Geometry Nodes: add utility to process all instances separately

This adds a new `GeometrySet::modify_geometry_sets` method that can be
used to update each sub-geometry-set separately without making any
instances real.

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

5 files changed, 104 insertions, 48 deletions

diff --git a/source/blender/blenkernel/BKE_geometry_set.hh b/source/blender/blenkernel/BKE_geometry_set.hh
index 571c6c6a0a0..08bfdbf2382 100644
--- a/source/blender/blenkernel/BKE_geometry_set.hh
+++ b/source/blender/blenkernel/BKE_geometry_set.hh
@@ -309,6 +309,10 @@ struct GeometrySet {
       bool include_instances,
       blender::Map<blender::bke::AttributeIDRef, AttributeKind> &r_attributes) const;
 
+  using ForeachSubGeometryCallback = blender::FunctionRef<void(GeometrySet &geometry_set)>;
+
+  void modify_geometry_sets(ForeachSubGeometryCallback callback);
+
   /* Utility methods for creation. */
   static GeometrySet create_with_mesh(
       Mesh *mesh, GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
@@ -479,7 +483,7 @@ class InstanceReference {
   Type type_ = Type::None;
   /** Depending on the type this is either null, an Object or Collection pointer. */
   void *data_ = nullptr;
-  std::shared_ptr<GeometrySet> geometry_set_;
+  std::unique_ptr<GeometrySet> geometry_set_;
 
  public:
   InstanceReference() = default;
@@ -494,8 +498,25 @@ class InstanceReference {
 
   InstanceReference(GeometrySet geometry_set)
       : type_(Type::GeometrySet),
-        geometry_set_(std::make_shared<GeometrySet>(std::move(geometry_set)))
+        geometry_set_(std::make_unique<GeometrySet>(std::move(geometry_set)))
+  {
+  }
+
+  InstanceReference(const InstanceReference &other) : type_(other.type_), data_(other.data_)
   {
+    if (other.geometry_set_) {
+      geometry_set_ = std::make_unique<GeometrySet>(*other.geometry_set_);
+    }
+  }
+
+  InstanceReference &operator=(const InstanceReference &other)
+  {
+    if (this == &other) {
+      return *this;
+    }
+    this->~InstanceReference();
+    new (this) InstanceReference(other);
+    return *this;
   }
 
   Type type() const
diff --git a/source/blender/blenkernel/intern/geometry_component_instances.cc b/source/blender/blenkernel/intern/geometry_component_instances.cc
index f1f60266545..4204d62e1a7 100644
--- a/source/blender/blenkernel/intern/geometry_component_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_component_instances.cc
@@ -125,44 +125,14 @@ blender::Span<int> InstancesComponent::instance_ids() const
 }
 
 /**
- * If references have a collection or object type, convert them into geometry instances. This
- * will join geometry components from nested instances if necessary. After that, the geometry
- * sets can be edited.
- */
-void InstancesComponent::ensure_geometry_instances()
-{
-  VectorSet<InstanceReference> new_references;
-  new_references.reserve(references_.size());
-  for (const InstanceReference &reference : references_) {
-    if (reference.type() == InstanceReference::Type::Object) {
-      GeometrySet geometry_set;
-      InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
-      const int handle = instances.add_reference(reference.object());
-      instances.add_instance(handle, float4x4::identity());
-      new_references.add_new(geometry_set);
-    }
-    else if (reference.type() == InstanceReference::Type::Collection) {
-      GeometrySet geometry_set;
-      InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
-      const int handle = instances.add_reference(reference.collection());
-      instances.add_instance(handle, float4x4::identity());
-      new_references.add_new(geometry_set);
-    }
-    else {
-      new_references.add_new(reference);
-    }
-  }
-  references_ = std::move(new_references);
-}
-
-/**
  * With write access to the instances component, the data in the instanced geometry sets can be
  * changed. This is a function on the component rather than each reference to ensure `const`
  * correctness for that reason.
  */
 GeometrySet &InstancesComponent::geometry_set_from_reference(const int reference_index)
 {
-  /* If this assert fails, it means #ensure_geometry_instances must be called first. */
+  /* If this assert fails, it means #ensure_geometry_instances must be called first or that the
+   * reference can't be converted to a geometry set. */
   BLI_assert(references_[reference_index].type() == InstanceReference::Type::GeometrySet);
 
   /* The const cast is okay because the instance's hash in the set
diff --git a/source/blender/blenkernel/intern/geometry_set.cc b/source/blender/blenkernel/intern/geometry_set.cc
index 1ebdde75f46..3abe2fd2213 100644
--- a/source/blender/blenkernel/intern/geometry_set.cc
+++ b/source/blender/blenkernel/intern/geometry_set.cc
@@ -458,6 +458,28 @@ void GeometrySet::gather_attributes_for_propagation(
   delete dummy_component;
 }
 
+/**
+ * Modify every (recursive) instance separately. This is often more efficient than realizing all
+ * instances just to change the same thing on all of them.
+ */
+void GeometrySet::modify_geometry_sets(ForeachSubGeometryCallback callback)
+{
+  callback(*this);
+  if (!this->has_instances()) {
+    return;
+  }
+  /* In the future this can be improved by deduplicating instance references across different
+   * instances. */
+  InstancesComponent &instances_component = this->get_component_for_write<InstancesComponent>();
+  instances_component.ensure_geometry_instances();
+  for (const int handle : instances_component.references().index_range()) {
+    if (instances_component.references()[handle].type() == InstanceReference::Type::GeometrySet) {
+      GeometrySet &instance_geometry = instances_component.geometry_set_from_reference(handle);
+      instance_geometry.modify_geometry_sets(callback);
+    }
+  }
+}
+
 /** \} */
 
 /* -------------------------------------------------------------------- */
diff --git a/source/blender/blenkernel/intern/geometry_set_instances.cc b/source/blender/blenkernel/intern/geometry_set_instances.cc
index 162f91a4a47..ad13342ad9e 100644
--- a/source/blender/blenkernel/intern/geometry_set_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_set_instances.cc
@@ -652,3 +652,58 @@ void InstancesComponent::foreach_referenced_geometry(
     }
   }
 }
+
+/**
+ * If references have a collection or object type, convert them into geometry instances
+ * recursively. After that, the geometry sets can be edited. There may still be instances of other
+ * types of they can't be converted to geometry sets.
+ */
+void InstancesComponent::ensure_geometry_instances()
+{
+  using namespace blender;
+  using namespace blender::bke;
+  VectorSet<InstanceReference> new_references;
+  new_references.reserve(references_.size());
+  for (const InstanceReference &reference : references_) {
+    switch (reference.type()) {
+      case InstanceReference::Type::None:
+      case InstanceReference::Type::GeometrySet: {
+        /* Those references can stay as their were. */
+        new_references.add_new(reference);
+        break;
+      }
+      case InstanceReference::Type::Object: {
+        /* Create a new reference that contains the geometry set of the object. We may want to
+         * treat e.g. lamps and similar object types separately here. */
+        const Object &object = reference.object();
+        GeometrySet object_geometry_set = object_get_geometry_set_for_read(object);
+        if (object_geometry_set.has_instances()) {
+          InstancesComponent &component =
+              object_geometry_set.get_component_for_write<InstancesComponent>();
+          component.ensure_geometry_instances();
+        }
+        new_references.add_new(std::move(object_geometry_set));
+        break;
+      }
+      case InstanceReference::Type::Collection: {
+        /* Create a new reference that contains a geometry set that contains all objects from the
+         * collection as instances. */
+        GeometrySet collection_geometry_set;
+        InstancesComponent &component =
+            collection_geometry_set.get_component_for_write<InstancesComponent>();
+        Collection &collection = reference.collection();
+        FOREACH_COLLECTION_OBJECT_RECURSIVE_BEGIN (&collection, object) {
+          const int handle = component.add_reference(*object);
+          component.add_instance(handle, object->obmat);
+          float4x4 &transform = component.instance_transforms().last();
+          sub_v3_v3(transform.values[3], collection.instance_offset);
+        }
+        FOREACH_COLLECTION_OBJECT_RECURSIVE_END;
+        component.ensure_geometry_instances();
+        new_references.add_new(std::move(collection_geometry_set));
+        break;
+      }
+    }
+  }
+  references_ = std::move(new_references);
+}
diff --git a/source/blender/nodes/geometry/nodes/node_geo_curve_fill.cc b/source/blender/nodes/geometry/nodes/node_geo_curve_fill.cc
index fcafdf93197..c30741cf786 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_curve_fill.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_curve_fill.cc
@@ -154,20 +154,8 @@ static void geo_node_curve_fill_exec(GeoNodeExecParams params)
   const NodeGeometryCurveFill &storage = *(const NodeGeometryCurveFill *)params.node().storage;
   const GeometryNodeCurveFillMode mode = (GeometryNodeCurveFillMode)storage.mode;
 
-  if (geometry_set.has_instances()) {
-    InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
-    instances.ensure_geometry_instances();
-
-    threading::parallel_for(IndexRange(instances.references_amount()), 16, [&](IndexRange range) {
-      for (int i : range) {
-        GeometrySet &geometry_set = instances.geometry_set_from_reference(i);
-        geometry_set = bke::geometry_set_realize_instances(geometry_set);
-        curve_fill_calculate(geometry_set, mode);
-      }
-    });
-  }
-
-  curve_fill_calculate(geometry_set, mode);
+  geometry_set.modify_geometry_sets(
+      [&](GeometrySet &geometry_set) { curve_fill_calculate(geometry_set, mode); });
 
   params.set_output("Mesh", std::move(geometry_set));
 }