Geometry Nodes: Make instances real on-demand

This commit makes the geometry output of the collection info usable.
The output is the geometry of a collection instance, but this commit
adds a utility to convert the instances to real geometry, used in the
background whenever it is needed, like copy on write.

The recursive nature of the "realize instances" code is essential,
because collection instances in the `InstancesComponent`, might have no
geometry sets of their own containing even more collection instances,
which might then contain object instances, etc.

Another consideration is that currently, every single instance contains
a reference to its data. This is inefficient since most of the time
there are many locations and only a few sets of unique data. So this
commit adds a `GeometryInstanceGroup` to support this future optimization.
The API for instances returns a vector of `GeometryInstanceGroup`.
This may be less efficient when there are many instances, but it makes
more complicated operations like point distribution that need to iterate
over input geometry multiple times much simpler.

Any code that needs to change data, like most of the attribute nodes,
can simply call `geometry_set_realize_instances(geometry_set)`,
which will move any geometry in the `InstancesComponent` to new "real"
geometry components.

Many nodes can support read-only access to instances in order to avoid
making them real, this will be addressed where needed in the near future.

Instances from the existing "dupli" system are not supported yet.

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

2 files changed, 172 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_geometry_set.hh b/source/blender/blenkernel/BKE_geometry_set.hh
index 02b3c88183a..5b429a61b3f 100644
--- a/source/blender/blenkernel/BKE_geometry_set.hh
+++ b/source/blender/blenkernel/BKE_geometry_set.hh
@@ -467,3 +467,25 @@ class VolumeComponent : public GeometryComponent {
 
   static constexpr inline GeometryComponentType static_type = GeometryComponentType::Volume;
 };
+
+/**
+ * Used to keep track of a group of instances using the same geometry data.
+ */
+struct GeometryInstanceGroup {
+  /**
+   * The geometry set instanced on each of the transforms. The components are not necessarily
+   * owned here. For example, they may be owned by the instanced object. This cannot be a
+   * reference because not all instanced data will necessarily have a #geometry_set_eval.
+   */
+  GeometrySet geometry_set;
+
+  /**
+   * As an optimization to avoid copying, the same geometry set can be associated with multiple
+   * instances. Each instance is stored as a transform matrix here. Again, these must be owned
+   * because they may be transformed from the original data. TODO: Validate that last statement.
+   */
+  blender::Vector<blender::float4x4> transforms;
+};
+
+blender::Vector<GeometryInstanceGroup> BKE_geometry_set_gather_instances(
+    const GeometrySet &geometry_set);
diff --git a/source/blender/blenkernel/intern/geometry_set.cc b/source/blender/blenkernel/intern/geometry_set.cc
index 7eb41a6c623..0bc248299ca 100644
--- a/source/blender/blenkernel/intern/geometry_set.cc
+++ b/source/blender/blenkernel/intern/geometry_set.cc
@@ -14,13 +14,19 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BLI_map.hh"
+
+#include "BKE_attribute.h"
+#include "BKE_attribute_access.hh"
 #include "BKE_geometry_set.hh"
 #include "BKE_lib_id.h"
 #include "BKE_mesh.h"
 #include "BKE_mesh_wrapper.h"
+#include "BKE_modifier.h"
 #include "BKE_pointcloud.h"
 #include "BKE_volume.h"
 
+#include "DNA_collection_types.h"
 #include "DNA_object_types.h"
 
 #include "BLI_rand.hh"
@@ -29,6 +35,7 @@
 
 using blender::float3;
 using blender::float4x4;
+using blender::Map;
 using blender::MutableSpan;
 using blender::Span;
 using blender::StringRef;
@@ -579,6 +586,149 @@ bool InstancesComponent::is_empty() const
   return transforms_.size() == 0;
 }
 
+/**
+ * \note This doesn't extract instances from the "dupli" system for non-geometry-nodes instances.
+ */
+static GeometrySet object_get_geometry_set_for_read(const Object &object)
+{
+  /* Objects evaluated with a nodes modifier will have a geometry set already. */
+  if (object.runtime.geometry_set_eval != nullptr) {
+    return *object.runtime.geometry_set_eval;
+  }
+
+  /* Otherwise, construct a new geometry set with the component based on the object type. */
+  GeometrySet new_geometry_set;
+
+  if (object.type == OB_MESH) {
+    Mesh *mesh = BKE_modifier_get_evaluated_mesh_from_evaluated_object(
+        &const_cast<Object &>(object), false);
+
+    if (mesh != nullptr) {
+      BKE_mesh_wrapper_ensure_mdata(mesh);
+
+      MeshComponent &mesh_component = new_geometry_set.get_component_for_write<MeshComponent>();
+      mesh_component.replace(mesh, GeometryOwnershipType::ReadOnly);
+      mesh_component.copy_vertex_group_names_from_object(object);
+    }
+  }
+
+  /* TODO: Cover the case of pointclouds without modifiers-- they may not be covered by the
+   * #geometry_set_eval case above. */
+
+  /* TODO: Add volume support. */
+
+  /* Return by value since there is not always an existing geometry set owned elsewhere to use. */
+  return new_geometry_set;
+}
+
+/** \} */
+
+/* -------------------------------------------------------------------- */
+/** \name Geometry Set Gather Recursive Instances
+ * \{ */
+
+static void geometry_set_collect_recursive(const GeometrySet &geometry_set,
+                                           const float4x4 &transform,
+                                           Vector<GeometryInstanceGroup> &r_sets);
+
+static void geometry_set_collect_recursive_collection(const Collection &collection,
+                                                      const float4x4 &transform,
+                                                      Vector<GeometryInstanceGroup> &r_sets);
+
+static void geometry_set_collect_recursive_collection_instance(
+    const Collection &collection, const float4x4 &transform, Vector<GeometryInstanceGroup> &r_sets)
+{
+  float4x4 offset_matrix;
+  unit_m4(offset_matrix.values);
+  sub_v3_v3(offset_matrix.values[3], collection.instance_offset);
+  const float4x4 instance_transform = transform * offset_matrix;
+  geometry_set_collect_recursive_collection(collection, instance_transform, r_sets);
+}
+
+static void geometry_set_collect_recursive_object(const Object &object,
+                                                  const float4x4 &transform,
+                                                  Vector<GeometryInstanceGroup> &r_sets)
+{
+  GeometrySet instance_geometry_set = object_get_geometry_set_for_read(object);
+  geometry_set_collect_recursive(instance_geometry_set, transform, r_sets);
+
+  if (object.type == OB_EMPTY) {
+    const Collection *collection_instance = object.instance_collection;
+    if (collection_instance != nullptr) {
+      geometry_set_collect_recursive_collection_instance(*collection_instance, transform, r_sets);
+    }
+  }
+}
+
+static void geometry_set_collect_recursive_collection(const Collection &collection,
+                                                      const float4x4 &transform,
+                                                      Vector<GeometryInstanceGroup> &r_sets)
+{
+  LISTBASE_FOREACH (const CollectionObject *, collection_object, &collection.gobject) {
+    BLI_assert(collection_object->ob != nullptr);
+    const Object &object = *collection_object->ob;
+    const float4x4 object_transform = transform * object.obmat;
+    geometry_set_collect_recursive_object(object, object_transform, r_sets);
+  }
+  LISTBASE_FOREACH (const CollectionChild *, collection_child, &collection.children) {
+    BLI_assert(collection_child->collection != nullptr);
+    const Collection &collection = *collection_child->collection;
+    geometry_set_collect_recursive_collection(collection, transform, r_sets);
+  }
+}
+
+static void geometry_set_collect_recursive(const GeometrySet &geometry_set,
+                                           const float4x4 &transform,
+                                           Vector<GeometryInstanceGroup> &r_sets)
+{
+  r_sets.append({geometry_set, {transform}});
+
+  if (geometry_set.has_instances()) {
+    const InstancesComponent &instances_component =
+        *geometry_set.get_component_for_read<InstancesComponent>();
+
+    Span<float4x4> transforms = instances_component.transforms();
+    Span<InstancedData> instances = instances_component.instanced_data();
+    for (const int i : instances.index_range()) {
+      const InstancedData &data = instances[i];
+      const float4x4 instance_transform = transform * transforms[i];
+
+      if (data.type == INSTANCE_DATA_TYPE_OBJECT) {
+        BLI_assert(data.data.object != nullptr);
+        const Object &object = *data.data.object;
+        geometry_set_collect_recursive_object(object, instance_transform, r_sets);
+      }
+      else if (data.type == INSTANCE_DATA_TYPE_COLLECTION) {
+        BLI_assert(data.data.collection != nullptr);
+        const Collection &collection = *data.data.collection;
+        geometry_set_collect_recursive_collection_instance(collection, instance_transform, r_sets);
+      }
+    }
+  }
+}
+
+/**
+ * Return flattened vector of the geometry component's recursive instances. I.e. all collection
+ * instances and object instances will be expanded into the instances of their geometry components.
+ * Even the instances in those geometry components' will be included.
+ *
+ * \note For convenience (to avoid duplication in the caller), the returned vector also contains
+ * the argument geometry set.
+ *
+ * \note This doesn't extract instances from the "dupli" system for non-geometry-nodes instances.
+ */
+Vector<GeometryInstanceGroup> BKE_geometry_set_gather_instances(const GeometrySet &geometry_set)
+{
+  Vector<GeometryInstanceGroup> result_vector;
+
+  float4x4 unit_transform;
+  unit_m4(unit_transform.values);
+
+  geometry_set_collect_recursive(geometry_set, unit_transform, result_vector);
+
+  return result_vector;
+}
+
 static blender::Array<int> generate_unique_instance_ids(Span<int> original_ids)
 {
   using namespace blender;