Geometry Nodes: refactor instances component

The main goal of this refactor is to not store Object/Collection
pointers for every individual instance. Instead instances now
store a handle for the referenced data. The actual Object/Collection
pointers are stored in a new `InstanceReference` class.

This refactor also allows for some better optimizations further down
the line, because one does not have to search through all instances
anymore to find what data is instanced.

Furthermore, this refactor makes it easier to support instancing
`GeometrySet` or any other data that has to be owned by the
`InstancesComponent`.

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

5 files changed, 57 insertions, 55 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_collection_info.cc b/source/blender/nodes/geometry/nodes/node_geo_collection_info.cc
index 0ad495aa4db..5800d46b70d 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_collection_info.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_collection_info.cc
@@ -58,10 +58,6 @@ static void geo_node_collection_info_exec(GeoNodeExecParams params)
   const bool transform_space_relative = (node_storage->transform_space ==
                                          GEO_NODE_TRANSFORM_SPACE_RELATIVE);
 
-  InstancedData instance;
-  instance.type = INSTANCE_DATA_TYPE_COLLECTION;
-  instance.data.collection = collection;
-
   InstancesComponent &instances = geometry_set_out.get_component_for_write<InstancesComponent>();
 
   float transform_mat[4][4];
@@ -73,7 +69,9 @@ static void geo_node_collection_info_exec(GeoNodeExecParams params)
 
     mul_m4_m4_pre(transform_mat, self_object->imat);
   }
-  instances.add_instance(instance, transform_mat, -1);
+
+  const int handle = instances.add_reference(*collection);
+  instances.add_instance(handle, transform_mat, -1);
 
   params.set_output("Geometry", geometry_set_out);
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc b/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
index edc2fb58e01..68bb3614751 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
@@ -243,13 +243,30 @@ static void join_components(Span<const PointCloudComponent *> src_components, Ge
 static void join_components(Span<const InstancesComponent *> src_components, GeometrySet &result)
 {
   InstancesComponent &dst_component = result.get_component_for_write<InstancesComponent>();
-  for (const InstancesComponent *component : src_components) {
-    const int size = component->instances_amount();
-    Span<InstancedData> instanced_data = component->instanced_data();
-    Span<float4x4> transforms = component->transforms();
-    Span<int> ids = component->ids();
-    for (const int i : IndexRange(size)) {
-      dst_component.add_instance(instanced_data[i], transforms[i], ids[i]);
+
+  int tot_instances = 0;
+  for (const InstancesComponent *src_component : src_components) {
+    tot_instances += src_component->instances_amount();
+  }
+  dst_component.reserve(tot_instances);
+
+  for (const InstancesComponent *src_component : src_components) {
+    Span<InstanceReference> src_references = src_component->references();
+    Array<int> handle_map(src_references.size());
+    for (const int src_handle : src_references.index_range()) {
+      handle_map[src_handle] = dst_component.add_reference(src_references[src_handle]);
+    }
+
+    Span<float4x4> src_transforms = src_component->instance_transforms();
+    Span<int> src_ids = src_component->instance_ids();
+    Span<int> src_reference_handles = src_component->instance_reference_handles();
+
+    for (const int i : src_transforms.index_range()) {
+      const int src_handle = src_reference_handles[i];
+      const int dst_handle = handle_map[src_handle];
+      const float4x4 &transform = src_transforms[i];
+      const int id = src_ids[i];
+      dst_component.add_instance(dst_handle, transform, id);
     }
   }
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_object_info.cc b/source/blender/nodes/geometry/nodes/node_geo_object_info.cc
index bd42b4c11d6..de099b8062f 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_object_info.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_object_info.cc
@@ -73,14 +73,15 @@ static void geo_node_object_info_exec(GeoNodeExecParams params)
 
     if (object != self_object) {
       InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
+      const int handle = instances.add_reference(*object);
 
       if (transform_space_relative) {
-        instances.add_instance(object, transform);
+        instances.add_instance(handle, transform);
       }
       else {
         float unit_transform[4][4];
         unit_m4(unit_transform);
-        instances.add_instance(object, unit_transform);
+        instances.add_instance(handle, unit_transform);
       }
     }
   }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc b/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
index f199ed71067..7929a5a1546 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
@@ -65,8 +65,8 @@ static void geo_node_point_instance_update(bNodeTree *UNUSED(tree), bNode *node)
       seed_socket, type == GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION && !use_whole_collection);
 }
 
-static void get_instanced_data__object(const GeoNodeExecParams &params,
-                                       MutableSpan<std::optional<InstancedData>> r_instances_data)
+static void get_instance_references__object(const GeoNodeExecParams &params,
+                                            MutableSpan<InstanceReference> r_references)
 {
   bke::PersistentObjectHandle object_handle = params.get_input<bke::PersistentObjectHandle>(
       "Object");
@@ -75,17 +75,13 @@ static void get_instanced_data__object(const GeoNodeExecParams &params,
     object = nullptr;
   }
   if (object != nullptr) {
-    InstancedData instance;
-    instance.type = INSTANCE_DATA_TYPE_OBJECT;
-    instance.data.object = object;
-    r_instances_data.fill(instance);
+    r_references.fill(*object);
   }
 }
 
-static void get_instanced_data__collection(
-    const GeoNodeExecParams &params,
-    const GeometryComponent &component,
-    MutableSpan<std::optional<InstancedData>> r_instances_data)
+static void get_instance_references__collection(const GeoNodeExecParams &params,
+                                                const GeometryComponent &component,
+                                                MutableSpan<InstanceReference> r_references)
 {
   const bNode &node = params.node();
   NodeGeometryPointInstance *node_storage = (NodeGeometryPointInstance *)node.storage;
@@ -106,62 +102,51 @@ static void get_instanced_data__collection(
   const bool use_whole_collection = (node_storage->flag &
                                      GEO_NODE_POINT_INSTANCE_WHOLE_COLLECTION) != 0;
   if (use_whole_collection) {
-    InstancedData instance;
-    instance.type = INSTANCE_DATA_TYPE_COLLECTION;
-    instance.data.collection = collection;
-    r_instances_data.fill(instance);
+    r_references.fill(*collection);
   }
   else {
-    Vector<InstancedData> possible_instances;
+    Vector<InstanceReference> possible_references;
     /* Direct child objects are instanced as objects. */
     LISTBASE_FOREACH (CollectionObject *, cob, &collection->gobject) {
-      Object *object = cob->ob;
-      InstancedData instance;
-      instance.type = INSTANCE_DATA_TYPE_OBJECT;
-      instance.data.object = object;
-      possible_instances.append(instance);
+      possible_references.append(*cob->ob);
     }
     /* Direct child collections are instanced as collections. */
     LISTBASE_FOREACH (CollectionChild *, child, &collection->children) {
-      Collection *child_collection = child->collection;
-      InstancedData instance;
-      instance.type = INSTANCE_DATA_TYPE_COLLECTION;
-      instance.data.collection = child_collection;
-      possible_instances.append(instance);
+      possible_references.append(*child->collection);
     }
 
-    if (!possible_instances.is_empty()) {
+    if (!possible_references.is_empty()) {
       const int seed = params.get_input<int>("Seed");
       Array<uint32_t> ids = get_geometry_element_ids_as_uints(component, ATTR_DOMAIN_POINT);
-      for (const int i : r_instances_data.index_range()) {
-        const int index = BLI_hash_int_2d(ids[i], seed) % possible_instances.size();
-        r_instances_data[i] = possible_instances[index];
+      for (const int i : r_references.index_range()) {
+        const int index = BLI_hash_int_2d(ids[i], seed) % possible_references.size();
+        r_references[i] = possible_references[index];
       }
     }
   }
 }
 
-static Array<std::optional<InstancedData>> get_instanced_data(const GeoNodeExecParams &params,
-                                                              const GeometryComponent &component,
-                                                              const int amount)
+static Array<InstanceReference> get_instance_references(const GeoNodeExecParams &params,
+                                                        const GeometryComponent &component,
+                                                        const int amount)
 {
   const bNode &node = params.node();
   NodeGeometryPointInstance *node_storage = (NodeGeometryPointInstance *)node.storage;
   const GeometryNodePointInstanceType type = (GeometryNodePointInstanceType)
                                                  node_storage->instance_type;
-  Array<std::optional<InstancedData>> instances_data(amount);
+  Array<InstanceReference> references(amount);
 
   switch (type) {
     case GEO_NODE_POINT_INSTANCE_TYPE_OBJECT: {
-      get_instanced_data__object(params, instances_data);
+      get_instance_references__object(params, references);
       break;
     }
     case GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION: {
-      get_instanced_data__collection(params, component, instances_data);
+      get_instance_references__collection(params, component, references);
       break;
     }
   }
-  return instances_data;
+  return references;
 }
 
 static void add_instances_from_geometry_component(InstancesComponent &instances,
@@ -171,8 +156,7 @@ static void add_instances_from_geometry_component(InstancesComponent &instances,
   const AttributeDomain domain = ATTR_DOMAIN_POINT;
 
   const int domain_size = src_geometry.attribute_domain_size(domain);
-  Array<std::optional<InstancedData>> instances_data = get_instanced_data(
-      params, src_geometry, domain_size);
+  Array<InstanceReference> references = get_instance_references(params, src_geometry, domain_size);
 
   GVArray_Typed<float3> positions = src_geometry.attribute_get_for_read<float3>(
       "position", domain, {0, 0, 0});
@@ -183,9 +167,11 @@ static void add_instances_from_geometry_component(InstancesComponent &instances,
   GVArray_Typed<int> ids = src_geometry.attribute_get_for_read<int>("id", domain, -1);
 
   for (const int i : IndexRange(domain_size)) {
-    if (instances_data[i].has_value()) {
+    const InstanceReference &reference = references[i];
+    if (reference.type() != InstanceReference::Type::None) {
       const float4x4 matrix = float4x4::from_loc_eul_scale(positions[i], rotations[i], scales[i]);
-      instances.add_instance(*instances_data[i], matrix, ids[i]);
+      const int handle = instances.add_reference(reference);
+      instances.add_instance(handle, matrix, ids[i]);
     }
   }
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_transform.cc b/source/blender/nodes/geometry/nodes/node_geo_transform.cc
index cba2c53907f..ce52dc90668 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_transform.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_transform.cc
@@ -101,7 +101,7 @@ static void transform_instances(InstancesComponent &instances,
                                 const float3 rotation,
                                 const float3 scale)
 {
-  MutableSpan<float4x4> transforms = instances.transforms();
+  MutableSpan<float4x4> transforms = instances.instance_transforms();
 
   /* Use only translation if rotation and scale don't apply. */
   if (use_translate(rotation, scale)) {