1 files changed, 53 insertions, 0 deletions

diff --git a/source/blender/blenkernel/intern/geometry_component_instances.cc b/source/blender/blenkernel/intern/geometry_component_instances.cc
index c4e1fe2f8e9..9479d012cb8 100644
--- a/source/blender/blenkernel/intern/geometry_component_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_component_instances.cc
@@ -24,6 +24,7 @@
 #include "DNA_collection_types.h"
 
 #include "BKE_geometry_set.hh"
+#include "BKE_geometry_set_instances.hh"
 
 #include "attribute_access_intern.hh"
 
@@ -32,6 +33,7 @@ using blender::Map;
 using blender::MutableSpan;
 using blender::Set;
 using blender::Span;
+using blender::VectorSet;
 
 /* -------------------------------------------------------------------- */
 /** \name Geometry Component Implementation
@@ -120,6 +122,52 @@ 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. */
+  BLI_assert(references_[reference_index].type() == InstanceReference::Type::GeometrySet);
+
+  /* The const cast is okay because the instance's hash in the set
+   * is not changed by adjusting the data inside the geometry set. */
+  return const_cast<GeometrySet &>(references_[reference_index].geometry_set());
+}
+
+/**
  * Returns a handle for the given reference.
  * If the reference exists already, the handle of the existing reference is returned.
  * Otherwise a new handle is added.
@@ -139,6 +187,11 @@ int InstancesComponent::instances_amount() const
   return instance_transforms_.size();
 }
 
+int InstancesComponent::references_amount() const
+{
+  return references_.size();
+}
+
 bool InstancesComponent::is_empty() const
 {
   return this->instance_reference_handles_.size() == 0;