Refactor object engine data storage

Main idea is to make specific engine types be a subclass of generic
ObjectEngineData structure.

This required following changes:

- Have extra size argument to engine data allocation function.

  Not sure whether there is less error-prone way of doing this.

- Add init() callback to engine data allocation function.

Additionally, added some extra checks to Eevee's engine data getters, so we do
not silently cast lamp data to lightprobe data.

Reviewers: dfelinto, fclem

Differential Revision: https://developer.blender.org/D3027

7 files changed, 127 insertions, 73 deletions

diff --git a/source/blender/blenkernel/intern/object.c b/source/blender/blenkernel/intern/object.c
index 9d752b59abe..8c85a0413b8 100644
--- a/source/blender/blenkernel/intern/object.c
+++ b/source/blender/blenkernel/intern/object.c
@@ -458,11 +458,8 @@ void BKE_object_free(Object *ob)
 	GPU_lamp_free(ob);
 
 	for (ObjectEngineData *oed = ob->drawdata.first; oed; oed = oed->next) {
-		if (oed->storage) {
-			if (oed->free) {
-				oed->free(oed->storage);
-			}
-			MEM_freeN(oed->storage);
+		if (oed->free != NULL) {
+			oed->free(oed);
 		}
 	}
 	BLI_freelistN(&ob->drawdata);
diff --git a/source/blender/draw/engines/eevee/eevee_data.c b/source/blender/draw/engines/eevee/eevee_data.c
index 0a052b12e93..c6e5f645f10 100644
--- a/source/blender/draw/engines/eevee/eevee_data.c
+++ b/source/blender/draw/engines/eevee/eevee_data.c
@@ -65,13 +65,6 @@ static void eevee_view_layer_data_free(void *storage)
 	DRW_TEXTURE_FREE_SAFE(sldata->irradiance_rt);
 }
 
-static void eevee_lightprobe_data_free(void *storage)
-{
-	EEVEE_LightProbeEngineData *ped = (EEVEE_LightProbeEngineData *)storage;
-
-	BLI_freelistN(&ped->captured_object_list);
-}
-
 EEVEE_ViewLayerData *EEVEE_view_layer_data_get(void)
 {
 	return (EEVEE_ViewLayerData *)DRW_view_layer_engine_data_get(
@@ -90,61 +83,95 @@ EEVEE_ViewLayerData *EEVEE_view_layer_data_ensure(void)
 	return *sldata;
 }
 
+/* Object data. */
+
+static void eevee_object_data_init(ObjectEngineData *engine_data)
+{
+	EEVEE_ObjectEngineData *eevee_data = (EEVEE_ObjectEngineData *)engine_data;
+	eevee_data->shadow_caster_id = -1;
+}
+
 EEVEE_ObjectEngineData *EEVEE_object_data_get(Object *ob)
 {
+	if (ELEM(ob->type, OB_LIGHTPROBE, OB_LAMP)) {
+		return NULL;
+	}
 	return (EEVEE_ObjectEngineData *)DRW_object_engine_data_get(
 	        ob, &draw_engine_eevee_type);
 }
 
 EEVEE_ObjectEngineData *EEVEE_object_data_ensure(Object *ob)
 {
-	EEVEE_ObjectEngineData **oedata = (EEVEE_ObjectEngineData **)DRW_object_engine_data_ensure(
-	        ob, &draw_engine_eevee_type, NULL);
+	BLI_assert(!ELEM(ob->type, OB_LIGHTPROBE, OB_LAMP));
+	return (EEVEE_ObjectEngineData *)DRW_object_engine_data_ensure(
+	        ob,
+	        &draw_engine_eevee_type,
+	        sizeof(EEVEE_ObjectEngineData),
+	        eevee_object_data_init,
+	        NULL);
+}
 
-	if (*oedata == NULL) {
-		*oedata = MEM_callocN(sizeof(**oedata), "EEVEE_ObjectEngineData");
-		(*oedata)->shadow_caster_id = -1;
-	}
+/* Light probe data. */
+
+static void eevee_lightprobe_data_init(ObjectEngineData *engine_data)
+{
+	EEVEE_LightProbeEngineData *ped = (EEVEE_LightProbeEngineData *)engine_data;
+	ped->need_full_update = true;
+	ped->need_update = true;
+}
+
+static void eevee_lightprobe_data_free(ObjectEngineData *engine_data)
+{
+	EEVEE_LightProbeEngineData *ped = (EEVEE_LightProbeEngineData *)engine_data;
 
-	return *oedata;
+	BLI_freelistN(&ped->captured_object_list);
 }
 
 EEVEE_LightProbeEngineData *EEVEE_lightprobe_data_get(Object *ob)
 {
+	if (ob->type != OB_LIGHTPROBE) {
+		return NULL;
+	}
 	return (EEVEE_LightProbeEngineData *)DRW_object_engine_data_get(
 	        ob, &draw_engine_eevee_type);
 }
 
 EEVEE_LightProbeEngineData *EEVEE_lightprobe_data_ensure(Object *ob)
 {
-	EEVEE_LightProbeEngineData **pedata = (EEVEE_LightProbeEngineData **)DRW_object_engine_data_ensure(
-	        ob, &draw_engine_eevee_type, &eevee_lightprobe_data_free);
+	BLI_assert(ob->type == OB_LIGHTPROBE);
+	return (EEVEE_LightProbeEngineData *)DRW_object_engine_data_ensure(
+	        ob,
+	        &draw_engine_eevee_type,
+	        sizeof(EEVEE_LightProbeEngineData),
+	        &eevee_lightprobe_data_init,
+	        &eevee_lightprobe_data_free);
+}
 
-	if (*pedata == NULL) {
-		*pedata = MEM_callocN(sizeof(**pedata), "EEVEE_LightProbeEngineData");
-		(*pedata)->need_full_update = true;
-		(*pedata)->need_update = true;
-	}
+/* Lamp data. */
 
-	return *pedata;
+static void eevee_lamp_data_init(ObjectEngineData *engine_data)
+{
+	EEVEE_LampEngineData *led = (EEVEE_LampEngineData *)engine_data;
+	led->need_update = true;
+	led->prev_cube_shadow_id = -1;
 }
 
 EEVEE_LampEngineData *EEVEE_lamp_data_get(Object *ob)
 {
+	if (ob->type != OB_LAMP) {
+		return NULL;
+	}
 	return (EEVEE_LampEngineData *)DRW_object_engine_data_get(
 	        ob, &draw_engine_eevee_type);
 }
 
 EEVEE_LampEngineData *EEVEE_lamp_data_ensure(Object *ob)
 {
-	EEVEE_LampEngineData **ledata = (EEVEE_LampEngineData **)DRW_object_engine_data_ensure(
-	        ob, &draw_engine_eevee_type, NULL);
-
-	if (*ledata == NULL) {
-		*ledata = MEM_callocN(sizeof(**ledata), "EEVEE_LampEngineData");
-		(*ledata)->need_update = true;
-		(*ledata)->prev_cube_shadow_id = -1;
-	}
-
-	return *ledata;
+	BLI_assert(ob->type == OB_LAMP);
+	return (EEVEE_LampEngineData *)DRW_object_engine_data_ensure(
+	        ob,
+	        &draw_engine_eevee_type,
+	        sizeof(EEVEE_LampEngineData),
+	        eevee_lamp_data_init,
+	        NULL);
 }
diff --git a/source/blender/draw/engines/eevee/eevee_private.h b/source/blender/draw/engines/eevee/eevee_private.h
index c91bfb1405f..648048d44fb 100644
--- a/source/blender/draw/engines/eevee/eevee_private.h
+++ b/source/blender/draw/engines/eevee/eevee_private.h
@@ -647,6 +647,8 @@ typedef struct EEVEE_ShadowCascadeData {
  * It works with even if the object is in multiple layers
  * because we don't get the same "Object *" for each layer. */
 typedef struct EEVEE_LampEngineData {
+	ObjectEngineData engine_data;
+
 	bool need_update;
 	/* This needs to be out of the union to avoid undefined behaviour. */
 	short prev_cube_shadow_id;
@@ -658,6 +660,8 @@ typedef struct EEVEE_LampEngineData {
 } EEVEE_LampEngineData;
 
 typedef struct EEVEE_LightProbeEngineData {
+	ObjectEngineData engine_data;
+
 	/* NOTE: need_full_update is set by dependency graph when the probe or it's
 	 * object is updated. This triggers full probe update, including it's
 	 * "progressive" GI refresh.
@@ -686,6 +690,8 @@ typedef struct EEVEE_LightProbeEngineData {
 } EEVEE_LightProbeEngineData;
 
 typedef struct EEVEE_ObjectEngineData {
+	ObjectEngineData engine_data;
+
 	bool need_update;
 	unsigned int shadow_caster_id;
 } EEVEE_ObjectEngineData;
diff --git a/source/blender/draw/intern/DRW_render.h b/source/blender/draw/intern/DRW_render.h
index 1f16f63ba14..f5943f2130b 100644
--- a/source/blender/draw/intern/DRW_render.h
+++ b/source/blender/draw/intern/DRW_render.h
@@ -391,9 +391,13 @@ void *DRW_view_layer_engine_data_get(DrawEngineType *engine_type);
 void **DRW_view_layer_engine_data_ensure(DrawEngineType *engine_type, void (*callback)(void *storage));
 
 /* Objects */
-void *DRW_object_engine_data_get(Object *ob, DrawEngineType *engine_type);
-void **DRW_object_engine_data_ensure(
-        Object *ob, DrawEngineType *engine_type, void (*callback)(void *storage));
+ObjectEngineData *DRW_object_engine_data_get(Object *ob, DrawEngineType *engine_type);
+ObjectEngineData *DRW_object_engine_data_ensure(
+        Object *ob,
+        DrawEngineType *engine_type,
+        size_t size,
+        ObjectEngineDataInitCb init_cb,
+        ObjectEngineDataFreeCb free_cb);
 struct LampEngineData *DRW_lamp_engine_data_ensure(Object *ob, struct RenderEngineType *engine_type);
 void DRW_lamp_engine_data_free(struct LampEngineData *led);
 
diff --git a/source/blender/draw/intern/draw_manager.c b/source/blender/draw/intern/draw_manager.c
index b1718413d4d..a9cb3ad77f5 100644
--- a/source/blender/draw/intern/draw_manager.c
+++ b/source/blender/draw/intern/draw_manager.c
@@ -2765,33 +2765,40 @@ void **DRW_view_layer_engine_data_ensure(DrawEngineType *engine_type, void (*cal
 /** \name Objects (DRW_object)
  * \{ */
 
-void *DRW_object_engine_data_get(Object *ob, DrawEngineType *engine_type)
+ObjectEngineData *DRW_object_engine_data_get(Object *ob, DrawEngineType *engine_type)
 {
 	for (ObjectEngineData *oed = ob->drawdata.first; oed; oed = oed->next) {
 		if (oed->engine_type == engine_type) {
-			return oed->storage;
+			return oed;
 		}
 	}
 	return NULL;
 }
 
-void **DRW_object_engine_data_ensure(
-        Object *ob, DrawEngineType *engine_type, void (*callback)(void *storage))
+ObjectEngineData *DRW_object_engine_data_ensure(
+        Object *ob,
+        DrawEngineType *engine_type,
+        size_t size,
+        ObjectEngineDataInitCb init_cb,
+        ObjectEngineDataFreeCb free_cb)
 {
-	ObjectEngineData *oed;
-
-	for (oed = ob->drawdata.first; oed; oed = oed->next) {
-		if (oed->engine_type == engine_type) {
-			return &oed->storage;
-		}
+	BLI_assert(size >= sizeof(ObjectEngineData));
+	/* Try to re-use existing data. */
+	ObjectEngineData *oed = DRW_object_engine_data_get(ob, engine_type);
+	if (oed != NULL) {
+		return oed;
 	}
-
-	oed = MEM_callocN(sizeof(ObjectEngineData), "ObjectEngineData");
+	/* Allocate new data. */
+	oed = MEM_callocN(size, "ObjectEngineData");
 	oed->engine_type = engine_type;
-	oed->free = callback;
+	oed->free = free_cb;
+	/* Perform user-side initialization, if needed. */
+	if (init_cb != NULL) {
+		init_cb(oed);
+	}
+	/* Register in the list. */
 	BLI_addtail(&ob->drawdata, oed);
-
-	return &oed->storage;
+	return oed;
 }
 
 /* XXX There is definitly some overlap between this and DRW_object_engine_data_ensure.
diff --git a/source/blender/draw/modes/object_mode.c b/source/blender/draw/modes/object_mode.c
index 5ad9e19079e..10cbf631561 100644
--- a/source/blender/draw/modes/object_mode.c
+++ b/source/blender/draw/modes/object_mode.c
@@ -1124,15 +1124,22 @@ static void DRW_shgroup_lamp(OBJECT_StorageList *stl, Object *ob, ViewLayer *vie
 	int theme_id = DRW_object_wire_theme_get(ob, view_layer, &color);
 	static float zero = 0.0f;
 
-	float **la_mats = (float **)DRW_object_engine_data_ensure(ob, &draw_engine_object_type, NULL);
-	if (*la_mats == NULL) {
-		/* we need 2 matrices */
-		*la_mats = MEM_mallocN(sizeof(float) * 16 * 2, "Lamp Object Mode Matrices");
-	}
-
-	float (*shapemat)[4], (*spotblendmat)[4];
-	shapemat = (float (*)[4])(*la_mats);
-	spotblendmat = (float (*)[4])(*la_mats + 16);
+	typedef struct LampEngineData {
+		ObjectEngineData engine_data;
+		float shape_mat[4][4];
+		float spot_blend_mat[4][4];
+	} LampEngineData;
+
+	LampEngineData *lamp_engine_data =
+	        (LampEngineData *)DRW_object_engine_data_ensure(
+	                ob,
+	                &draw_engine_object_type,
+	                sizeof(LampEngineData),
+	                NULL,
+	                NULL);
+
+	float (*shapemat)[4] = lamp_engine_data->shape_mat;
+	float (*spotblendmat)[4] = lamp_engine_data->spot_blend_mat;
 
 	/* Don't draw the center if it's selected or active */
 	if (theme_id == TH_GROUP)
@@ -1481,13 +1488,13 @@ static void DRW_shgroup_lightprobe(OBJECT_StorageList *stl, OBJECT_PassList *psl
 	bool do_outlines = ((ob->base_flag & BASE_SELECTED) != 0);
 	DRW_object_wire_theme_get(ob, view_layer, &color);
 
-	OBJECT_LightProbeEngineData *prb_data;
-	OBJECT_LightProbeEngineData **prb_data_pt = (OBJECT_LightProbeEngineData **)DRW_object_engine_data_ensure(ob, &draw_engine_object_type, NULL);
-	if (*prb_data_pt == NULL) {
-		*prb_data_pt = MEM_mallocN(sizeof(OBJECT_LightProbeEngineData), "Probe Clip distances Matrices");
-	}
-
-	prb_data = *prb_data_pt;
+	OBJECT_LightProbeEngineData *prb_data =
+	        (OBJECT_LightProbeEngineData *)DRW_object_engine_data_ensure(
+	                ob,
+	                &draw_engine_object_type,
+	                sizeof(OBJECT_LightProbeEngineData),
+	                NULL,
+	                NULL);
 
 	if ((DRW_state_is_select() || do_outlines) && ((prb->flag & LIGHTPROBE_FLAG_SHOW_DATA) != 0)) {
 
diff --git a/source/blender/makesdna/DNA_object_types.h b/source/blender/makesdna/DNA_object_types.h
index 7364dcbc167..d62ec746309 100644
--- a/source/blender/makesdna/DNA_object_types.h
+++ b/source/blender/makesdna/DNA_object_types.h
@@ -75,11 +75,17 @@ typedef struct bFaceMap {
 } bFaceMap;
 
 /* Object Runtime display data */
+struct ObjectEngineData;
+typedef void (*ObjectEngineDataInitCb)(struct ObjectEngineData *engine_data);
+typedef void (*ObjectEngineDataFreeCb)(struct ObjectEngineData *engine_data);
+
+#
+#
 typedef struct ObjectEngineData {
 	struct ObjectEngineData *next, *prev;
 	struct DrawEngineType *engine_type;
-	void *storage;
-	void (*free)(void *storage);
+	/* Only nested data, NOT the engine data itself. */
+	ObjectEngineDataFreeCb free;
 } ObjectEngineData;
 
 #define MAX_VGROUP_NAME 64