Fix crash going to edit mode of particles with copy on write

We can not rely on edit->psys, it is not set for particle edit,
and there is some logic deeper inside which does different things
dependent on that.

We need to replace those checks with some some HAIR vs. PARTICLES
flag and always set psys pointer.

7 files changed, 81 insertions, 51 deletions

diff --git a/source/blender/draw/intern/draw_cache.c b/source/blender/draw/intern/draw_cache.c
index 6c6e9a732b3..28153505051 100644
--- a/source/blender/draw/intern/draw_cache.c
+++ b/source/blender/draw/intern/draw_cache.c
@@ -2886,19 +2886,28 @@ Gwn_Batch *DRW_cache_particles_get_dots(Object *object, ParticleSystem *psys)
 	return DRW_particles_batch_cache_get_dots(object, psys);
 }
 
-Gwn_Batch *DRW_cache_particles_get_edit_strands(Object *object, struct PTCacheEdit *edit)
+Gwn_Batch *DRW_cache_particles_get_edit_strands(
+        Object *object,
+        ParticleSystem *psys,
+        struct PTCacheEdit *edit)
 {
-	return DRW_particles_batch_cache_get_edit_strands(object, edit);
+	return DRW_particles_batch_cache_get_edit_strands(object, psys, edit);
 }
 
-Gwn_Batch *DRW_cache_particles_get_edit_inner_points(Object *object, struct PTCacheEdit *edit)
+Gwn_Batch *DRW_cache_particles_get_edit_inner_points(
+        Object *object,
+        ParticleSystem *psys,
+        struct PTCacheEdit *edit)
 {
-	return DRW_particles_batch_cache_get_edit_inner_points(object, edit);
+	return DRW_particles_batch_cache_get_edit_inner_points(object, psys, edit);
 }
 
-Gwn_Batch *DRW_cache_particles_get_edit_tip_points(Object *object, struct PTCacheEdit *edit)
+Gwn_Batch *DRW_cache_particles_get_edit_tip_points(
+        Object *object,
+        ParticleSystem *psys,
+        struct PTCacheEdit *edit)
 {
-	return DRW_particles_batch_cache_get_edit_tip_points(object, edit);
+	return DRW_particles_batch_cache_get_edit_tip_points(object, psys, edit);
 }
 
 Gwn_Batch *DRW_cache_particles_get_prim(int type)
diff --git a/source/blender/draw/intern/draw_cache.h b/source/blender/draw/intern/draw_cache.h
index 2dc07e40c42..22117d71534 100644
--- a/source/blender/draw/intern/draw_cache.h
+++ b/source/blender/draw/intern/draw_cache.h
@@ -169,9 +169,9 @@ struct Gwn_Batch *DRW_cache_lattice_vert_overlay_get(struct Object *ob);
 /* Particles */
 struct Gwn_Batch *DRW_cache_particles_get_hair(struct ParticleSystem *psys, struct ModifierData *md);
 struct Gwn_Batch *DRW_cache_particles_get_dots(struct Object *object, struct ParticleSystem *psys);
-struct Gwn_Batch *DRW_cache_particles_get_edit_strands(struct Object *object, struct PTCacheEdit *edit);
-struct Gwn_Batch *DRW_cache_particles_get_edit_inner_points(struct Object *object, struct PTCacheEdit *edit);
-struct Gwn_Batch *DRW_cache_particles_get_edit_tip_points(struct Object *object, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_cache_particles_get_edit_strands(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_cache_particles_get_edit_inner_points(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_cache_particles_get_edit_tip_points(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
 struct Gwn_Batch *DRW_cache_particles_get_prim(int type);
 
 /* Metaball */
diff --git a/source/blender/draw/intern/draw_cache_impl.h b/source/blender/draw/intern/draw_cache_impl.h
index 9d9668f3ffc..fd1a33452bc 100644
--- a/source/blender/draw/intern/draw_cache_impl.h
+++ b/source/blender/draw/intern/draw_cache_impl.h
@@ -125,8 +125,8 @@ void DRW_mesh_cache_sculpt_coords_ensure(struct Mesh *me);
 /* Particles */
 struct Gwn_Batch *DRW_particles_batch_cache_get_hair(struct ParticleSystem *psys, struct ModifierData *md);
 struct Gwn_Batch *DRW_particles_batch_cache_get_dots(struct Object *object, struct ParticleSystem *psys);
-struct Gwn_Batch *DRW_particles_batch_cache_get_edit_strands(struct Object *object, struct PTCacheEdit *edit);
-struct Gwn_Batch *DRW_particles_batch_cache_get_edit_inner_points(struct Object *object, struct PTCacheEdit *edit);
-struct Gwn_Batch *DRW_particles_batch_cache_get_edit_tip_points(struct Object *object, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_particles_batch_cache_get_edit_strands(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_particles_batch_cache_get_edit_inner_points(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+struct Gwn_Batch *DRW_particles_batch_cache_get_edit_tip_points(struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
 
 #endif /* __DRAW_CACHE_IMPL_H__ */
diff --git a/source/blender/draw/intern/draw_cache_impl_particles.c b/source/blender/draw/intern/draw_cache_impl_particles.c
index f1ae9d6ef8c..1c39cbc16cd 100644
--- a/source/blender/draw/intern/draw_cache_impl_particles.c
+++ b/source/blender/draw/intern/draw_cache_impl_particles.c
@@ -597,25 +597,12 @@ Gwn_Batch *DRW_particles_batch_cache_get_dots(Object *object, ParticleSystem *ps
 	return cache->hairs;
 }
 
-/* TODO(sergey): Avoid linear lookup. */
-static ParticleBatchCache *particle_batch_cache_get_edit(Object *object, PTCacheEdit *edit)
-{
-	ParticleSystem *psys_orig = edit->psys;
-	for (ParticleSystem *psys_eval = object->particlesystem.first;
-	     psys_eval != NULL;
-	     psys_eval = psys_eval->next)
-	{
-		if (STREQ(psys_orig->name, psys_eval->name)) {
-			return particle_batch_cache_get(psys_eval);
-		}
-	}
-	return NULL;
-}
-
-Gwn_Batch *DRW_particles_batch_cache_get_edit_strands(Object *object, PTCacheEdit *edit)
+Gwn_Batch *DRW_particles_batch_cache_get_edit_strands(
+        Object *UNUSED(object),
+        ParticleSystem *psys,
+        PTCacheEdit *edit)
 {
-	ParticleSystem *psys = edit->psys;
-	ParticleBatchCache *cache = particle_batch_cache_get_edit(object, edit);
+	ParticleBatchCache *cache = particle_batch_cache_get(psys);
 	if (cache->hairs != NULL) {
 		return cache->hairs;
 	}
@@ -691,9 +678,12 @@ static void particle_batch_cache_ensure_edit_inner_pos(
 	}
 }
 
-Gwn_Batch *DRW_particles_batch_cache_get_edit_inner_points(Object *object, PTCacheEdit *edit)
+Gwn_Batch *DRW_particles_batch_cache_get_edit_inner_points(
+        Object *UNUSED(object),
+        ParticleSystem *psys,
+        PTCacheEdit *edit)
 {
-	ParticleBatchCache *cache = particle_batch_cache_get_edit(object, edit);
+	ParticleBatchCache *cache = particle_batch_cache_get(psys);
 	if (cache->edit_inner_points != NULL) {
 		return cache->edit_inner_points;
 	}
@@ -752,9 +742,12 @@ static void particle_batch_cache_ensure_edit_tip_pos(
 	}
 }
 
-Gwn_Batch *DRW_particles_batch_cache_get_edit_tip_points(Object *object, PTCacheEdit *edit)
+Gwn_Batch *DRW_particles_batch_cache_get_edit_tip_points(
+        Object *UNUSED(object),
+        ParticleSystem *psys,
+        PTCacheEdit *edit)
 {
-	ParticleBatchCache *cache = particle_batch_cache_get_edit(object, edit);
+	ParticleBatchCache *cache = particle_batch_cache_get(psys);
 	if (cache->edit_tip_points != NULL) {
 		return cache->edit_tip_points;
 	}
diff --git a/source/blender/draw/modes/particle_mode.c b/source/blender/draw/modes/particle_mode.c
index d9363a18054..ef85fe1fdc4 100644
--- a/source/blender/draw/modes/particle_mode.c
+++ b/source/blender/draw/modes/particle_mode.c
@@ -139,22 +139,17 @@ static void particle_cache_init(void *vedata)
 	DRW_shgroup_uniform_float(stl->g_data->tip_points_group, "outlineWidth", &outline_width, 1);
 }
 
-/* TODO(sergey): Avoid linear lookup. */
-static void draw_update_ptcache_edit(Object *object_eval, PTCacheEdit *edit)
+static void draw_update_ptcache_edit(Object *object_eval,
+                                     ParticleSystem *psys,
+                                     PTCacheEdit *edit)
 {
 	if (edit->psys == NULL) {
 		return;
 	}
-	ParticleSystem *psys_eval;
-	for (psys_eval = object_eval->particlesystem.first;
-	     psys_eval != NULL;
-	     psys_eval = psys_eval->next)
-	{
-		if (STREQ(edit->psys->name, psys_eval->name)) {
-			break;
-		}
-	}
-	if (psys_eval->flag & PSYS_HAIR_UPDATED) {
+	/* NOTE: Get flag from particle system coming from drawing object.
+	 * this is where depsgraph will be setting flags to.
+	 */
+	if (psys->flag & PSYS_HAIR_UPDATED) {
 		const DRWContextState *draw_ctx = DRW_context_state_get();
 		Scene *scene_orig = (Scene *)DEG_get_original_id(&draw_ctx->scene->id);
 		Object *object_orig = DEG_get_original_object(object_eval);
@@ -165,25 +160,26 @@ static void draw_update_ptcache_edit(Object *object_eval, PTCacheEdit *edit)
 
 static void particle_edit_cache_populate(void *vedata,
                                          Object *object,
+                                         ParticleSystem *psys,
                                          PTCacheEdit *edit)
 {
 	PARTICLE_StorageList *stl = ((PARTICLE_Data *)vedata)->stl;
 	const DRWContextState *draw_ctx = DRW_context_state_get();
-	draw_update_ptcache_edit(object, edit);
+	draw_update_ptcache_edit(object, psys, edit);
 	ParticleEditSettings *pset = PE_settings(draw_ctx->scene);
 	{
 		struct Gwn_Batch *strands =
-		        DRW_cache_particles_get_edit_strands(object, edit);
+		        DRW_cache_particles_get_edit_strands(object, psys, edit);
 		DRW_shgroup_call_add(stl->g_data->strands_group, strands, NULL);
 	}
 	if (pset->selectmode == SCE_SELECT_POINT) {
 		struct Gwn_Batch *points =
-		        DRW_cache_particles_get_edit_inner_points(object, edit);
+		        DRW_cache_particles_get_edit_inner_points(object, psys, edit);
 		DRW_shgroup_call_add(stl->g_data->inner_points_group, points, NULL);
 	}
 	if (ELEM(pset->selectmode, SCE_SELECT_POINT, SCE_SELECT_END)) {
 		struct Gwn_Batch *points =
-		        DRW_cache_particles_get_edit_tip_points(object, edit);
+		        DRW_cache_particles_get_edit_tip_points(object, psys, edit);
 		DRW_shgroup_call_add(stl->g_data->tip_points_group, points, NULL);
 	}
 }
@@ -197,7 +193,20 @@ static void particle_cache_populate(void *vedata, Object *object)
 		printf("Particle edit struct is NULL, not supposed to happen.\n");
 		return;
 	}
-	particle_edit_cache_populate(vedata, object, edit);
+	ParticleSystem *psys = object->particlesystem.first;
+	ParticleSystem *psys_orig = object_orig->particlesystem.first;
+	while (psys_orig != NULL) {
+		if (PE_get_current_from_psys(psys_orig) == edit) {
+			break;
+		}
+		psys = psys->next;
+		psys_orig = psys_orig->next;
+	}
+	if (psys == NULL) {
+		printf("Error getting evaluated particle system for edit.\n");
+		return;
+	}
+	particle_edit_cache_populate(vedata, object, psys, edit);
 }
 
 /* Optional: Post-cache_populate callback */
diff --git a/source/blender/editors/include/ED_particle.h b/source/blender/editors/include/ED_particle.h
index ef0690546c0..8b522c91188 100644
--- a/source/blender/editors/include/ED_particle.h
+++ b/source/blender/editors/include/ED_particle.h
@@ -47,6 +47,7 @@ void PE_free_ptcache_edit(struct PTCacheEdit *edit);
 int PE_start_edit(struct PTCacheEdit *edit);
 
 /* access */
+struct PTCacheEdit *PE_get_current_from_psys(struct ParticleSystem *psys);
 struct PTCacheEdit *PE_get_current(struct Scene *scene, struct Object *ob);
 struct PTCacheEdit *PE_create_current(struct Depsgraph *depsgraph, struct Scene *scene, struct Object *ob);
 void PE_current_changed(struct Depsgraph *depsgraph, struct Scene *scene, struct Object *ob);
diff --git a/source/blender/editors/physics/particle_edit.c b/source/blender/editors/physics/particle_edit.c
index d42989424be..f65e598e204 100644
--- a/source/blender/editors/physics/particle_edit.c
+++ b/source/blender/editors/physics/particle_edit.c
@@ -191,6 +191,24 @@ static float pe_brush_size_get(const Scene *UNUSED(scene), ParticleBrushData *br
 	return brush->size * U.pixelsize;
 }
 
+PTCacheEdit *PE_get_current_from_psys(ParticleSystem *psys)
+{
+	if (psys->part && psys->part->type == PART_HAIR) {
+		if ((psys->flag & PSYS_HAIR_DYNAMICS) != 0 &&
+		    (psys->pointcache->flag & PTCACHE_BAKED) != 0)
+		{
+			return psys->pointcache->edit;
+		}
+		else {
+			return psys->edit;
+		}
+	}
+	else if (psys->pointcache->flag & PTCACHE_BAKED) {
+		return psys->pointcache->edit;
+	}
+	return NULL;
+}
+
 /* always gets at least the first particlesystem even if PSYS_CURRENT flag is not set
  *
  * note: this function runs on poll, therefor it can runs many times a second