Geometry Nodes: simplify looping over attributes in geometry set

This adds three new methods:
* `InstancesComponent::foreach_reference_as_geometry(...)`
* `GeometrySet::attribute_foreach(...)`
* `GeometrySet::gather_attributes_for_propagation(...)`

The goal is that these iteration primitives can be used in places
where we use more specialized iterators currently.

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

source/blender/blenkernel/BKE_attribute_access.hh

@@ -122,6 +122,11 @@ struct AttributeMetaData {
   }
 };
 
+struct AttributeKind {
+  AttributeDomain domain;
+  CustomDataType data_type;
+};
+
 /**
  * Base class for the attribute initializer types described below.
  */

source/blender/blenkernel/BKE_geometry_set.hh

@@ -25,6 +25,7 @@
 
 #include "BLI_float3.hh"
 #include "BLI_float4x4.hh"
+#include "BLI_function_ref.hh"
 #include "BLI_hash.hh"
 #include "BLI_map.hh"
 #include "BLI_set.hh"
@@ -293,6 +294,21 @@ struct GeometrySet {
   bool owns_direct_data() const;
   void ensure_owns_direct_data();
 
+  using AttributeForeachCallback =
+      blender::FunctionRef<void(const blender::bke::AttributeIDRef &attribute_id,
+                                const AttributeMetaData &meta_data,
+                                const GeometryComponent &component)>;
+
+  void attribute_foreach(blender::Span<GeometryComponentType> component_types,
+                         bool include_instances,
+                         AttributeForeachCallback callback) const;
+
+  void gather_attributes_for_propagation(
+      blender::Span<GeometryComponentType> component_types,
+      GeometryComponentType dst_component_type,
+      bool include_instances,
+      blender::Map<blender::bke::AttributeIDRef, AttributeKind> &r_attributes) const;
+
   /* Utility methods for creation. */
   static GeometrySet create_with_mesh(
       Mesh *mesh, GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
@@ -597,6 +613,9 @@ class InstancesComponent : public GeometryComponent {
 
   int attribute_domain_size(const AttributeDomain domain) const final;
 
+  void foreach_referenced_geometry(
+      blender::FunctionRef<void(const GeometrySet &geometry_set)> callback) const;
+
   bool is_empty() const final;
 
   bool owns_direct_data() const override;

source/blender/blenkernel/BKE_geometry_set_instances.hh

@@ -49,11 +49,6 @@ void geometry_set_gather_instances(const GeometrySet &geometry_set,
 GeometrySet geometry_set_realize_mesh_for_modifier(const GeometrySet &geometry_set);
 GeometrySet geometry_set_realize_instances(const GeometrySet &geometry_set);
 
-struct AttributeKind {
-  CustomDataType data_type;
-  AttributeDomain domain;
-};
-
 /**
  * Add information about all the attributes on every component of the type. The resulting info
  * will contain the highest complexity data type and the highest priority domain among every

source/blender/blenkernel/intern/geometry_set.cc

@@ -375,6 +375,76 @@ CurveEval *GeometrySet::get_curve_for_write()
   return component.get_for_write();
 }
 
+void GeometrySet::attribute_foreach(const Span<GeometryComponentType> component_types,
+                                    const bool include_instances,
+                                    const AttributeForeachCallback callback) const
+{
+  using namespace blender;
+  using namespace blender::bke;
+  for (const GeometryComponentType component_type : component_types) {
+    if (!this->has(component_type)) {
+      continue;
+    }
+    const GeometryComponent &component = *this->get_component_for_read(component_type);
+    component.attribute_foreach(
+        [&](const AttributeIDRef &attribute_id, const AttributeMetaData &meta_data) {
+          callback(attribute_id, meta_data, component);
+          return true;
+        });
+  }
+  if (include_instances && this->has_instances()) {
+    const InstancesComponent &instances = *this->get_component_for_read<InstancesComponent>();
+    instances.foreach_referenced_geometry([&](const GeometrySet &instance_geometry_set) {
+      instance_geometry_set.attribute_foreach(component_types, include_instances, callback);
+    });
+  }
+}
+
+void GeometrySet::gather_attributes_for_propagation(
+    const Span<GeometryComponentType> component_types,
+    const GeometryComponentType dst_component_type,
+    bool include_instances,
+    blender::Map<blender::bke::AttributeIDRef, AttributeKind> &r_attributes) const
+{
+  using namespace blender;
+  using namespace blender::bke;
+  /* Only needed right now to check if an attribute is built-in on this component type.
+   * TODO: Get rid of the dummy component. */
+  const GeometryComponent *dummy_component = GeometryComponent::create(dst_component_type);
+  this->attribute_foreach(
+      component_types,
+      include_instances,
+      [&](const AttributeIDRef &attribute_id,
+          const AttributeMetaData &meta_data,
+          const GeometryComponent &component) {
+        if (component.attribute_is_builtin(attribute_id)) {
+          if (!dummy_component->attribute_is_builtin(attribute_id)) {
+            /* Don't propagate built-in attributes that are not built-in on the destination
+             * component. */
+            return;
+          }
+        }
+        if (attribute_id.is_anonymous()) {
+          if (!BKE_anonymous_attribute_id_has_strong_references(&attribute_id.anonymous_id())) {
+            /* Don't propagate anonymous attributes that are not used anymore. */
+            return;
+          }
+        }
+        auto add_info = [&](AttributeKind *attribute_kind) {
+          attribute_kind->domain = meta_data.domain;
+          attribute_kind->data_type = meta_data.data_type;
+        };
+        auto modify_info = [&](AttributeKind *attribute_kind) {
+          attribute_kind->domain = bke::attribute_domain_highest_priority(
+              {attribute_kind->domain, meta_data.domain});
+          attribute_kind->data_type = bke::attribute_data_type_highest_complexity(
+              {attribute_kind->data_type, meta_data.data_type});
+        };
+        r_attributes.add_or_modify(attribute_id, add_info, modify_info);
+      });
+  delete dummy_component;
+}
+
 /** \} */
 
 /* -------------------------------------------------------------------- */

source/blender/blenkernel/intern/geometry_set_instances.cc

@@ -14,6 +14,7 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BKE_collection.h"
 #include "BKE_geometry_set_instances.hh"
 #include "BKE_material.h"
 #include "BKE_mesh.h"
@@ -23,6 +24,7 @@
 #include "BKE_spline.hh"
 
 #include "DNA_collection_types.h"
+#include "DNA_layer_types.h"
 #include "DNA_mesh_types.h"
 #include "DNA_meshdata_types.h"
 #include "DNA_object_types.h"
@@ -745,3 +747,36 @@ GeometrySet geometry_set_realize_instances(const GeometrySet &geometry_set)
 }
 
 }  // namespace blender::bke
+
+void InstancesComponent::foreach_referenced_geometry(
+    blender::FunctionRef<void(const GeometrySet &geometry_set)> callback) const
+{
+  using namespace blender::bke;
+  for (const InstanceReference &reference : references_) {
+    switch (reference.type()) {
+      case InstanceReference::Type::Object: {
+        const Object &object = reference.object();
+        const GeometrySet object_geometry_set = object_get_geometry_set_for_read(object);
+        callback(object_geometry_set);
+        break;
+      }
+      case InstanceReference::Type::Collection: {
+        Collection &collection = reference.collection();
+        FOREACH_COLLECTION_OBJECT_RECURSIVE_BEGIN (&collection, object) {
+          const GeometrySet object_geometry_set = object_get_geometry_set_for_read(*object);
+          callback(object_geometry_set);
+        }
+        FOREACH_COLLECTION_OBJECT_RECURSIVE_END;
+        break;
+      }
+      case InstanceReference::Type::GeometrySet: {
+        const GeometrySet &instance_geometry_set = reference.geometry_set();
+        callback(instance_geometry_set);
+        break;
+      }
+      case InstanceReference::Type::None: {
+        break;
+      }
+    }
+  }
+}

source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc

@@ -36,7 +36,6 @@
 
 #include "node_geometry_util.hh"
 
-using blender::bke::AttributeKind;
 using blender::bke::GeometryInstanceGroup;
 
 namespace blender::nodes {