From 2e6223d90f3fddce013e17b271916d628df15425 Mon Sep 17 00:00:00 2001
From: Hans Goudey <hans@blender.org>
Date: Thu, 8 Feb 2024 20:55:34 +0100
Subject: [PATCH] Geometry Nodes: Move instance reference indices to a builtin
 attribute

This means the array can be shared between geometries when unchanged,
reducing memory usage and increasing performance. It also means that
handling the data can be done more generically using the attribute
system. In the future the transforms can become attributes too.

Two other changes are related here:
- The "almost unique ids" are cached with a shared cache so they
  are shared too.
- The reference indices are baked as an attribute now, making the process
  more generic and potentially easier to optimize in the future.

Pull Request: https://projects.blender.org/blender/blender/pulls/117951
---
 source/blender/blenkernel/BKE_instances.hh    | 18 ++---
 .../blenkernel/intern/attribute_access.cc     |  3 +
 .../blenkernel/intern/bake_items_serialize.cc | 25 ++++---
 .../intern/geometry_component_instances.cc    | 18 ++++-
 source/blender/blenkernel/intern/instances.cc | 70 ++++++++++---------
 .../geometry/intern/join_geometries.cc        |  4 +-
 source/blender/geometry/intern/randomize.cc   |  2 +-
 source/blender/geometry/intern/reorder.cc     |  4 --
 .../nodes/node_geo_duplicate_elements.cc      |  4 +-
 .../nodes/node_geo_instance_on_points.cc      |  4 +-
 .../nodes/node_geo_split_to_instances.cc      |  4 +-
 .../nodes/node_geo_string_to_curves.cc        |  2 +-
 12 files changed, 92 insertions(+), 66 deletions(-)

diff --git a/source/blender/blenkernel/BKE_instances.hh b/source/blender/blenkernel/BKE_instances.hh
index 87d2fc3efcc..7bd81aa0e08 100644
--- a/source/blender/blenkernel/BKE_instances.hh
+++ b/source/blender/blenkernel/BKE_instances.hh
@@ -19,13 +19,13 @@
  * which is then stored per instance. Many instances can use the same #InstanceReference.
  */
 
-#include <mutex>
 #include <optional>
 
 #include "BLI_array.hh"
 #include "BLI_function_ref.hh"
 #include "BLI_index_mask_fwd.hh"
 #include "BLI_math_matrix_types.hh"
+#include "BLI_shared_cache.hh"
 #include "BLI_vector.hh"
 
 #include "DNA_customdata_types.h"
@@ -99,19 +99,16 @@ class Instances {
    */
   Vector<InstanceReference> references_;
 
-  /** Indices into `references_`. Determines what data is instanced. */
-  Vector<int> reference_handles_;
   /** Transformation of the instances. */
   Vector<float4x4> transforms_;
 
+  CustomData attributes_;
+
   /* These almost unique ids are generated based on the `id` attribute, which might not contain
    * unique ids at all. They are *almost* unique, because under certain very unlikely
    * circumstances, they are not unique. Code using these ids should not crash when they are not
    * unique but can generally expect them to be unique. */
-  mutable std::mutex almost_unique_ids_mutex_;
-  mutable Array<int> almost_unique_ids_;
-
-  CustomData attributes_;
+  mutable SharedCache<Array<int>> almost_unique_ids_cache_;
 
  public:
   Instances();
@@ -161,7 +158,7 @@ class Instances {
   GeometrySet &geometry_set_from_reference(int reference_index);
 
   Span<int> reference_handles() const;
-  MutableSpan<int> reference_handles();
+  MutableSpan<int> reference_handles_for_write();
   MutableSpan<float4x4> transforms();
   Span<float4x4> transforms() const;
 
@@ -189,6 +186,11 @@ class Instances {
 
   bool owns_direct_data() const;
   void ensure_owns_direct_data();
+
+  void tag_reference_handles_changed()
+  {
+    almost_unique_ids_cache_.tag_dirty();
+  }
 };
 
 /* -------------------------------------------------------------------- */
diff --git a/source/blender/blenkernel/intern/attribute_access.cc b/source/blender/blenkernel/intern/attribute_access.cc
index 163e20708df..92764d2df09 100644
--- a/source/blender/blenkernel/intern/attribute_access.cc
+++ b/source/blender/blenkernel/intern/attribute_access.cc
@@ -129,6 +129,9 @@ bool allow_procedural_attribute_access(StringRef attribute_name)
   if (attribute_name.startswith(".uv")) {
     return false;
   }
+  if (attribute_name == ".reference_index") {
+    return false;
+  }
   if (attribute_name.startswith("." UV_VERTSEL_NAME ".")) {
     return false;
   }
diff --git a/source/blender/blenkernel/intern/bake_items_serialize.cc b/source/blender/blenkernel/intern/bake_items_serialize.cc
index ea5b57b3677..28edca63a09 100644
--- a/source/blender/blenkernel/intern/bake_items_serialize.cc
+++ b/source/blender/blenkernel/intern/bake_items_serialize.cc
@@ -724,19 +724,25 @@ static std::unique_ptr<Instances> try_load_instances(const DictionaryValue &io_g
     return {};
   }
 
-  const auto *io_handles = io_instances->lookup_dict("handles");
-  if (!io_handles) {
-    return {};
-  }
-  if (!read_blob_simple_gspan(blob_reader, *io_handles, instances->reference_handles())) {
-    return {};
-  }
-
   MutableAttributeAccessor attributes = instances->attributes_for_write();
   if (!load_attributes(*io_attributes, attributes, blob_reader, blob_sharing)) {
     return {};
   }
 
+  if (!attributes.contains(".reference_index")) {
+    /* Try reading the reference index attribute from the old bake format from before it was an
+     * attribute. */
+    const auto *io_handles = io_instances->lookup_dict("handles");
+    if (!io_handles) {
+      return {};
+    }
+    if (!read_blob_simple_gspan(
+            blob_reader, *io_handles, instances->reference_handles_for_write()))
+    {
+      return {};
+    }
+  }
+
   return instances;
 }
 
@@ -969,9 +975,6 @@ static std::shared_ptr<DictionaryValue> serialize_geometry_set(const GeometrySet
 
     io_instances->append(
         "transforms", write_blob_simple_gspan(blob_writer, blob_sharing, instances.transforms()));
-    io_instances->append(
-        "handles",
-        write_blob_simple_gspan(blob_writer, blob_sharing, instances.reference_handles()));
 
     auto io_attributes = serialize_attributes(
         instances.attributes(), blob_writer, blob_sharing, {"position"});
diff --git a/source/blender/blenkernel/intern/geometry_component_instances.cc b/source/blender/blenkernel/intern/geometry_component_instances.cc
index 961705a19dc..a417f29faf4 100644
--- a/source/blender/blenkernel/intern/geometry_component_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_component_instances.cc
@@ -168,6 +168,12 @@ class InstancePositionAttributeProvider final : public BuiltinAttributeProvider
   }
 };
 
+static void tag_component_reference_index_changed(void *owner)
+{
+  Instances &instances = *static_cast<Instances *>(owner);
+  instances.tag_reference_handles_changed();
+}
+
 static ComponentAttributeProviders create_attribute_providers_for_instances()
 {
   static InstancePositionAttributeProvider position;
@@ -200,10 +206,20 @@ static ComponentAttributeProviders create_attribute_providers_for_instances()
                                            instance_custom_data_access,
                                            nullptr);
 
+  /** Indices into `Instances::references_`. Determines what data is instanced. */
+  static BuiltinCustomDataLayerProvider reference_index(".reference_index",
+                                                        AttrDomain::Instance,
+                                                        CD_PROP_INT32,
+                                                        CD_PROP_INT32,
+                                                        BuiltinAttributeProvider::Creatable,
+                                                        BuiltinAttributeProvider::NonDeletable,
+                                                        instance_custom_data_access,
+                                                        tag_component_reference_index_changed);
+
   static CustomDataAttributeProvider instance_custom_data(AttrDomain::Instance,
                                                           instance_custom_data_access);
 
-  return ComponentAttributeProviders({&position, &id}, {&instance_custom_data});
+  return ComponentAttributeProviders({&position, &id, &reference_index}, {&instance_custom_data});
 }
 
 static AttributeAccessorFunctions get_instances_accessor_functions()
diff --git a/source/blender/blenkernel/intern/instances.cc b/source/blender/blenkernel/intern/instances.cc
index 2214bbe4da1..f43fba35002 100644
--- a/source/blender/blenkernel/intern/instances.cc
+++ b/source/blender/blenkernel/intern/instances.cc
@@ -51,19 +51,17 @@ Instances::Instances()
 
 Instances::Instances(Instances &&other)
     : references_(std::move(other.references_)),
-      reference_handles_(std::move(other.reference_handles_)),
       transforms_(std::move(other.transforms_)),
-      almost_unique_ids_(std::move(other.almost_unique_ids_)),
-      attributes_(other.attributes_)
+      attributes_(other.attributes_),
+      almost_unique_ids_cache_(std::move(other.almost_unique_ids_cache_))
 {
   CustomData_reset(&other.attributes_);
 }
 
 Instances::Instances(const Instances &other)
     : references_(other.references_),
-      reference_handles_(other.reference_handles_),
       transforms_(other.transforms_),
-      almost_unique_ids_(other.almost_unique_ids_)
+      almost_unique_ids_cache_(other.almost_unique_ids_cache_)
 {
   CustomData_copy(&other.attributes_, &attributes_, CD_MASK_ALL, other.instances_num());
 }
@@ -96,7 +94,6 @@ Instances &Instances::operator=(Instances &&other)
 void Instances::resize(int capacity)
 {
   const int old_size = this->instances_num();
-  reference_handles_.resize(capacity);
   transforms_.resize(capacity);
   CustomData_realloc(&attributes_, old_size, capacity, CD_SET_DEFAULT);
 }
@@ -106,19 +103,27 @@ void Instances::add_instance(const int instance_handle, const float4x4 &transfor
   BLI_assert(instance_handle >= 0);
   BLI_assert(instance_handle < references_.size());
   const int old_size = this->instances_num();
-  reference_handles_.append(instance_handle);
   transforms_.append(transform);
   CustomData_realloc(&attributes_, old_size, transforms_.size());
+  this->reference_handles_for_write().last() = instance_handle;
 }
 
 Span<int> Instances::reference_handles() const
 {
-  return reference_handles_;
+  return {static_cast<const int *>(
+              CustomData_get_layer_named(&attributes_, CD_PROP_INT32, ".reference_index")),
+          this->instances_num()};
 }
 
-MutableSpan<int> Instances::reference_handles()
+MutableSpan<int> Instances::reference_handles_for_write()
 {
-  return reference_handles_;
+  int *data = static_cast<int *>(CustomData_get_layer_named_for_write(
+      &attributes_, CD_PROP_INT32, ".reference_index", this->instances_num()));
+  if (!data) {
+    data = static_cast<int *>(CustomData_add_layer_named(
+        &attributes_, CD_PROP_INT32, CD_SET_DEFAULT, this->instances_num(), ".reference_index"));
+  }
+  return {data, this->instances_num()};
 }
 
 MutableSpan<float4x4> Instances::transforms()
@@ -178,10 +183,7 @@ void Instances::remove(const IndexMask &mask,
 
   Instances new_instances;
   new_instances.references_ = std::move(references_);
-  new_instances.reference_handles_.resize(new_size);
   new_instances.transforms_.resize(new_size);
-  array_utils::gather(
-      reference_handles_.as_span(), mask, new_instances.reference_handles_.as_mutable_span());
   array_utils::gather(transforms_.as_span(), mask, new_instances.transforms_.as_mutable_span());
 
   gather_attributes(this->attributes(),
@@ -212,6 +214,8 @@ void Instances::remove_unused_references()
     return;
   }
 
+  const Span<int> reference_handles = this->reference_handles();
+
   Array<bool> usage_by_handle(tot_references_before, false);
   std::mutex mutex;
 
@@ -221,7 +225,7 @@ void Instances::remove_unused_references()
     Array<bool> local_usage_by_handle(tot_references_before, false);
 
     for (const int i : range) {
-      const int handle = reference_handles_[i];
+      const int handle = reference_handles[i];
       BLI_assert(handle >= 0 && handle < tot_references_before);
       local_usage_by_handle[handle] = true;
     }
@@ -266,11 +270,14 @@ void Instances::remove_unused_references()
   }
 
   /* Update handles of instances. */
-  threading::parallel_for(IndexRange(tot_instances), 1000, [&](IndexRange range) {
-    for (const int i : range) {
-      reference_handles_[i] = handle_mapping[reference_handles_[i]];
-    }
-  });
+  {
+    const MutableSpan<int> reference_handles = this->reference_handles_for_write();
+    threading::parallel_for(IndexRange(tot_instances), 1000, [&](IndexRange range) {
+      for (const int i : range) {
+        reference_handles[i] = handle_mapping[reference_handles[i]];
+      }
+    });
+  }
 }
 
 int Instances::instances_num() const
@@ -357,21 +364,20 @@ static Array<int> generate_unique_instance_ids(Span<int> original_ids)
 
 Span<int> Instances::almost_unique_ids() const
 {
-  std::lock_guard lock(almost_unique_ids_mutex_);
-  bke::AttributeReader<int> instance_ids_attribute = this->attributes().lookup<int>("id");
-  if (instance_ids_attribute) {
-    Span<int> instance_ids = instance_ids_attribute.varray.get_internal_span();
-    if (almost_unique_ids_.size() != instance_ids.size()) {
-      almost_unique_ids_ = generate_unique_instance_ids(instance_ids);
+  almost_unique_ids_cache_.ensure([&](Array<int> &r_data) {
+    bke::AttributeReader<int> instance_ids_attribute = this->attributes().lookup<int>("id");
+    if (instance_ids_attribute) {
+      Span<int> instance_ids = instance_ids_attribute.varray.get_internal_span();
+      if (r_data.size() != instance_ids.size()) {
+        r_data = generate_unique_instance_ids(instance_ids);
+      }
     }
-  }
-  else {
-    almost_unique_ids_.reinitialize(this->instances_num());
-    for (const int i : almost_unique_ids_.index_range()) {
-      almost_unique_ids_[i] = i;
+    else {
+      r_data.reinitialize(this->instances_num());
+      array_utils::fill_index_range(r_data.as_mutable_span());
     }
-  }
-  return almost_unique_ids_;
+  });
+  return almost_unique_ids_cache_.data();
 }
 
 }  // namespace blender::bke
diff --git a/source/blender/geometry/intern/join_geometries.cc b/source/blender/geometry/intern/join_geometries.cc
index 3cc6175db0a..2acaf0d8909 100644
--- a/source/blender/geometry/intern/join_geometries.cc
+++ b/source/blender/geometry/intern/join_geometries.cc
@@ -110,7 +110,7 @@ static void join_instances(const Span<const GeometryComponent *> src_components,
   dst_instances->resize(offsets.total_size());
 
   MutableSpan<float4x4> all_transforms = dst_instances->transforms();
-  MutableSpan<int> all_handles = dst_instances->reference_handles();
+  MutableSpan<int> all_handles = dst_instances->reference_handles_for_write();
 
   for (const int i : src_components.index_range()) {
     const auto &src_component = static_cast<const bke::InstancesComponent &>(*src_components[i]);
@@ -131,7 +131,7 @@ static void join_instances(const Span<const GeometryComponent *> src_components,
 
   result.replace_instances(dst_instances.release());
   auto &dst_component = result.get_component_for_write<bke::InstancesComponent>();
-  join_attributes(src_components, dst_component, {"position"});
+  join_attributes(src_components, dst_component, {"position", ".reference_index"});
 }
 
 static void join_volumes(const Span<const GeometryComponent *> /*src_components*/,
diff --git a/source/blender/geometry/intern/randomize.cc b/source/blender/geometry/intern/randomize.cc
index f61c5af74bc..77a516b806c 100644
--- a/source/blender/geometry/intern/randomize.cc
+++ b/source/blender/geometry/intern/randomize.cc
@@ -253,7 +253,7 @@ void debug_randomize_instance_order(bke::Instances *instances)
     new_transforms[new_i] = old_transforms[old_i];
   }
 
-  instances->reference_handles().copy_from(new_reference_handles);
+  instances->reference_handles_for_write().copy_from(new_reference_handles);
   instances->transforms().copy_from(new_transforms);
 }
 
diff --git a/source/blender/geometry/intern/reorder.cc b/source/blender/geometry/intern/reorder.cc
index 0c5f14ff9b2..1007657e264 100644
--- a/source/blender/geometry/intern/reorder.cc
+++ b/source/blender/geometry/intern/reorder.cc
@@ -210,10 +210,6 @@ static void reorder_instaces_exec(const bke::Instances &src_instances,
                          old_by_new_map,
                          dst_instances.attributes_for_write());
 
-  const Span<int> old_reference_handles = src_instances.reference_handles();
-  MutableSpan<int> new_reference_handles = dst_instances.reference_handles();
-  array_utils::gather(old_reference_handles, old_by_new_map, new_reference_handles);
-
   const Span<float4x4> old_transforms = src_instances.transforms();
   MutableSpan<float4x4> new_transforms = dst_instances.transforms();
   array_utils::gather(old_transforms, old_by_new_map, new_transforms);
diff --git a/source/blender/nodes/geometry/nodes/node_geo_duplicate_elements.cc b/source/blender/nodes/geometry/nodes/node_geo_duplicate_elements.cc
index f52ff22a86d..319a2553da9 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_duplicate_elements.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_duplicate_elements.cc
@@ -956,13 +956,13 @@ static void duplicate_instances(GeometrySet &geometry_set,
     const int new_handle = dst_instances->add_reference(reference);
     const float4x4 transform = src_instances.transforms()[i_selection];
     dst_instances->transforms().slice(range).fill(transform);
-    dst_instances->reference_handles().slice(range).fill(new_handle);
+    dst_instances->reference_handles_for_write().slice(range).fill(new_handle);
   }
 
   bke::gather_attributes_to_groups(src_instances.attributes(),
                                    AttrDomain::Instance,
                                    propagation_info,
-                                   {"id"},
+                                   {"id", ".reference_index"},
                                    duplicates,
                                    selection,
                                    dst_instances->attributes_for_write());
diff --git a/source/blender/nodes/geometry/nodes/node_geo_instance_on_points.cc b/source/blender/nodes/geometry/nodes/node_geo_instance_on_points.cc
index a186cbf6baf..d43059578e9 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_instance_on_points.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_instance_on_points.cc
@@ -85,7 +85,8 @@ static void add_instances_from_component(
   const int select_len = selection.index_range().size();
   dst_component.resize(start_len + select_len);
 
-  MutableSpan<int> dst_handles = dst_component.reference_handles().slice(start_len, select_len);
+  MutableSpan<int> dst_handles = dst_component.reference_handles_for_write().slice(start_len,
+                                                                                   select_len);
   MutableSpan<float4x4> dst_transforms = dst_component.transforms().slice(start_len, select_len);
 
   const VArraySpan positions = *src_attributes.lookup<float3>("position");
@@ -213,6 +214,7 @@ static void node_geo_exec(GeoNodeExecParams params)
                                                    propagation_info,
                                                    attributes_to_propagate);
     attributes_to_propagate.remove("position");
+    attributes_to_propagate.remove(".reference_index");
 
     for (const GeometryComponent::Type type : types) {
       if (geometry_set.has(type)) {
diff --git a/source/blender/nodes/geometry/nodes/node_geo_split_to_instances.cc b/source/blender/nodes/geometry/nodes/node_geo_split_to_instances.cc
index 0ea24934dfe..2fd9da1203f 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_split_to_instances.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_split_to_instances.cc
@@ -264,8 +264,6 @@ static void split_instance_groups(const InstancesComponent &component,
       }
 
       array_utils::gather(src_instances.transforms(), mask, group_instances->transforms());
-      array_utils::gather(
-          src_instances.reference_handles(), mask, group_instances->reference_handles());
       bke::gather_attributes(src_instances.attributes(),
                              AttrDomain::Instance,
                              propagation_info,
@@ -343,7 +341,7 @@ static void node_geo_exec(GeoNodeExecParams params)
   }
 
   dst_instances->transforms().fill(float4x4::identity());
-  array_utils::fill_index_range(dst_instances->reference_handles());
+  array_utils::fill_index_range(dst_instances->reference_handles_for_write());
 
   for (auto item : geometry_by_group_id.items()) {
     std::unique_ptr<GeometrySet> &group_geometry = item.value;
diff --git a/source/blender/nodes/geometry/nodes/node_geo_string_to_curves.cc b/source/blender/nodes/geometry/nodes/node_geo_string_to_curves.cc
index 36229517354..730ce16a1d9 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_string_to_curves.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_string_to_curves.cc
@@ -319,7 +319,7 @@ static void add_instances_from_handles(bke::Instances &instances,
                                        const TextLayout &layout)
 {
   instances.resize(layout.positions.size());
-  MutableSpan<int> handles = instances.reference_handles();
+  MutableSpan<int> handles = instances.reference_handles_for_write();
   MutableSpan<float4x4> transforms = instances.transforms();
 
   threading::parallel_for(IndexRange(layout.positions.size()), 256, [&](IndexRange range) {