More progress converting the point distribute node to operations

I might not continue with this approach, it seems like building a
temporary vector of geometry components will be much simpler

source/blender/blenkernel/BKE_geometry_set.hh

@@ -476,3 +476,9 @@ class VolumeComponent : public GeometryComponent {
 
   static constexpr inline GeometryComponentType static_type = GeometryComponentType::Volume;
 };
+
+using ForeachGeometryCallbackConst = std::function<void(
+    const GeometryComponent &component, blender::Span<blender::float4x4> transforms)>;
+
+void BKE_foreach_geometry_component_recursive(const GeometrySet &geometry_set,
+                                              const ForeachGeometryCallbackConst &callback);

source/blender/blenkernel/intern/geometry_set.cc

@@ -14,19 +14,24 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BLI_listbase_wrapper.hh" /* TODO: Couldn't figure this out yet. */
+
 #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 "MEM_guardedalloc.h"
 
 using blender::float3;
 using blender::float4x4;
+using blender::ListBaseWrapper;
 using blender::MutableSpan;
 using blender::Span;
 using blender::StringRef;
@@ -566,6 +571,180 @@ bool InstancesComponent::is_empty() const
   return transforms_.size() == 0;
 }
 
+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);
+    }
+  }
+  // else if (object.type == OB_VOLUME) {
+  //   Volume *volume = BKE_modifier_get_volume...
+  // }
+
+  /* Return by value since there is no existing geometry set owned elsewhere to use. */
+  return new_geometry_set;
+}
+
+static void foreach_geometry_component_recursive(const GeometrySet &geometry_set,
+                                                 const ForeachGeometryCallbackConst &callback,
+                                                 const float4x4 &transform);
+
+static void foreach_collection_geometry_set_recursive(const Collection &collection,
+                                                      const ForeachGeometryCallbackConst &callback,
+                                                      const float4x4 &transform)
+{
+  LISTBASE_FOREACH (const CollectionObject *, collection_object, &collection.gobject) {
+    BLI_assert(collection_object->ob != nullptr);
+    const Object &object = *collection_object->ob;
+    GeometrySet instance_geometry_set = object_get_geometry_set_for_read(object);
+
+    /* TODO: This seems to work-- validate this. */
+    const float4x4 instance_transform = transform * object.obmat;
+    foreach_geometry_component_recursive(instance_geometry_set, callback, instance_transform);
+  }
+  LISTBASE_FOREACH (const CollectionChild *, collection_child, &collection.children) {
+    BLI_assert(collection_child->collection != nullptr);
+    const Collection &collection = *collection_child->collection;
+    foreach_collection_geometry_set_recursive(collection, callback, transform);
+  }
+}
+
+static void foreach_geometry_component_recursive(const GeometrySet &geometry_set,
+                                                 const ForeachGeometryCallbackConst &callback,
+                                                 const float4x4 &transform)
+{
+  if (geometry_set.has_mesh()) {
+    callback(*geometry_set.get_component_for_read<MeshComponent>(), {transform});
+  }
+  if (geometry_set.has_pointcloud()) {
+    callback(*geometry_set.get_component_for_read<PointCloudComponent>(), {transform});
+  }
+  if (geometry_set.has_volume()) {
+    callback(*geometry_set.get_component_for_read<VolumeComponent>(), {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 &transform = transforms[i];
+
+      if (data.type == INSTANCE_DATA_TYPE_OBJECT) {
+        BLI_assert(data.data.object != nullptr);
+        const Object &object = *data.data.object;
+        GeometrySet instance_geometry_set = object_get_geometry_set_for_read(object);
+        foreach_geometry_component_recursive(instance_geometry_set, callback, transform);
+      }
+      else if (data.type == INSTANCE_DATA_TYPE_COLLECTION) {
+        BLI_assert(data.data.collection != nullptr);
+        const Collection &collection = *data.data.collection;
+        foreach_collection_geometry_set_recursive(collection, callback, transform);
+      }
+    }
+  }
+}
+
+void BKE_foreach_geometry_component_recursive(const GeometrySet &geometry_set,
+                                              const ForeachGeometryCallbackConst &callback)
+{
+  float4x4 unit_transform;
+  unit_m4(unit_transform.values);
+
+  foreach_geometry_component_recursive(geometry_set, callback, unit_transform);
+}
+
+/* ============= API 2 =============== */
+
+using GeometrySetGroup = std::pair<GeometrySet, Vector<float4x4>>;
+
+static void collect_geometry_set_recursive(
+    const GeometrySet &geometry_set,
+    const float4x4 &transform,
+    Vector<std::pair<GeometrySet, Vector<float4x4>>> &r_sets);
+
+static void collect_collection_geometry_set_recursive(const Collection &collection,
+                                                      const float4x4 &transform,
+                                                      Vector<GeometrySetGroup> &r_sets)
+{
+  LISTBASE_FOREACH (const CollectionObject *, collection_object, &collection.gobject) {
+    BLI_assert(collection_object->ob != nullptr);
+    const Object &object = *collection_object->ob;
+    GeometrySet instance_geometry_set = object_get_geometry_set_for_read(object);
+
+    /* TODO: This seems to work-- validate this. */
+    const float4x4 instance_transform = transform * object.obmat;
+    collect_geometry_set_recursive(instance_geometry_set, instance_transform, r_sets);
+  }
+  LISTBASE_FOREACH (const CollectionChild *, collection_child, &collection.children) {
+    BLI_assert(collection_child->collection != nullptr);
+    const Collection &collection = *collection_child->collection;
+    collect_collection_geometry_set_recursive(collection, transform, r_sets);
+  }
+}
+
+static void collect_geometry_set_recursive(const GeometrySet &geometry_set,
+                                           const float4x4 &transform,
+                                           Vector<GeometrySetGroup> &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 &transform = transforms[i];
+
+      if (data.type == INSTANCE_DATA_TYPE_OBJECT) {
+        BLI_assert(data.data.object != nullptr);
+        const Object &object = *data.data.object;
+        GeometrySet instance_geometry_set = object_get_geometry_set_for_read(object);
+        collect_geometry_set_recursive(instance_geometry_set, transform, r_sets);
+      }
+      else if (data.type == INSTANCE_DATA_TYPE_COLLECTION) {
+        BLI_assert(data.data.collection != nullptr);
+        const Collection &collection = *data.data.collection;
+        collect_collection_geometry_set_recursive(collection, transform, r_sets);
+      }
+    }
+  }
+}
+
+Vector<GeometrySetGroup> BKE_geometry_set_gather_instanced(const GeometrySet &geometry_set)
+{
+  Vector<GeometrySetGroup> result_vector;
+
+  float4x4 unit_transform;
+  unit_m4(unit_transform.values);
+
+  collect_geometry_set_recursive(geometry_set, unit_transform, result_vector);
+
+  return result_vector;
+}
+
 /** \} */
 
 /* -------------------------------------------------------------------- */

source/blender/nodes/geometry/nodes/node_geo_object_info.cc

@@ -69,35 +69,15 @@ static void geo_node_object_info_exec(GeoNodeExecParams params)
     quat_to_eul(rotation, quaternion);
 
     if (object != self_object) {
-      if (object->type == OB_MESH) {
-        Mesh *mesh = BKE_modifier_get_evaluated_mesh_from_evaluated_object(object, false);
-        if (mesh != nullptr) {
-          BKE_mesh_wrapper_ensure_mdata(mesh);
+      InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
 
-          /* Make a copy because the life time of the other mesh might be shorter. */
-          Mesh *copied_mesh = BKE_mesh_copy_for_eval(mesh, false);
-
-          if (transform_space_relative) {
-            /* Transform into the local space of the object that is being modified. */
-            BKE_mesh_transform(copied_mesh, transform, true);
-          }
-
-          MeshComponent &mesh_component = geometry_set.get_component_for_write<MeshComponent>();
-          mesh_component.replace(copied_mesh);
-          mesh_component.copy_vertex_group_names_from_object(*object);
-        }
+      if (transform_space_relative) {
+        instances.add_instance(object, transform);
       }
-      if (object->type == OB_VOLUME) {
-        InstancesComponent &instances = geometry_set.get_component_for_write<InstancesComponent>();
-
-        if (transform_space_relative) {
-          instances.add_instance(object, transform);
-        }
-        else {
-          float unit_transform[4][4];
-          unit_m4(unit_transform);
-          instances.add_instance(object, unit_transform);
-        }
+      else {
+        float unit_transform[4][4];
+        unit_m4(unit_transform);
+        instances.add_instance(object, unit_transform);
       }
     }
   }

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

@@ -79,6 +79,7 @@ static Span<MLoopTri> get_mesh_looptris(const Mesh &mesh)
 }
 
 static void sample_mesh_surface(const Mesh &mesh,
+                                const float4x4 &transform,
                                 const float base_density,
                                 const FloatReadAttribute *density_factors,
                                 const int seed,
@@ -93,9 +94,10 @@ static void sample_mesh_surface(const Mesh &mesh,
     const int v0_index = mesh.mloop[looptri.tri[0]].v;
     const int v1_index = mesh.mloop[looptri.tri[1]].v;
     const int v2_index = mesh.mloop[looptri.tri[2]].v;
-    const float3 v0_pos = mesh.mvert[v0_index].co;
-    const float3 v1_pos = mesh.mvert[v1_index].co;
-    const float3 v2_pos = mesh.mvert[v2_index].co;
+
+    const float3 v0_pos = transform * float3(mesh.mvert[v0_index].co);
+    const float3 v1_pos = transform * float3(mesh.mvert[v1_index].co);
+    const float3 v2_pos = transform * float3(mesh.mvert[v2_index].co);
 
     float looptri_density_factor = 1.0f;
     if (density_factors != nullptr) {
@@ -194,8 +196,8 @@ BLI_NOINLINE static void update_elimination_mask_based_on_density_factors(
     const float v1_density_factor = std::max(0.0f, density_factors[v1_index]);
     const float v2_density_factor = std::max(0.0f, density_factors[v2_index]);
 
-    const float probablity = v0_density_factor * bary_coord.x + v1_density_factor * bary_coord.y +
-                             v2_density_factor * bary_coord.z;
+    const float probablity = attribute_math::mix3<float>(
+        bary_coord, v0_density_factor, v1_density_factor, v2_density_factor);
 
     const float hash = BLI_hash_int_01(bary_coord.hash());
     if (hash > probablity) {
@@ -386,77 +388,174 @@ BLI_NOINLINE static void add_remaining_point_attributes(const MeshComponent &mes
       *mesh_component.get_for_read(), component, bary_coords, looptri_indices);
 }
 
-static void sample_mesh_surface_with_minimum_distance(const Mesh &mesh,
-                                                      const float max_density,
-                                                      const float minimum_distance,
-                                                      const FloatReadAttribute &density_factors,
-                                                      const int seed,
-                                                      Vector<float3> &r_positions,
-                                                      Vector<float3> &r_bary_coords,
-                                                      Vector<int> &r_looptri_indices)
-{
-  sample_mesh_surface(
-      mesh, max_density, nullptr, seed, r_positions, r_bary_coords, r_looptri_indices);
-  Array<bool> elimination_mask(r_positions.size(), false);
-  update_elimination_mask_for_close_points(r_positions, minimum_distance, elimination_mask);
-  update_elimination_mask_based_on_density_factors(
-      mesh, density_factors, r_bary_coords, r_looptri_indices, elimination_mask);
-  eliminate_points_based_on_mask(elimination_mask, r_positions, r_bary_coords, r_looptri_indices);
-}
+struct AttributeInfo {
+  std::string name;
+  Vector<CustomDataType> data_types;
+  Vector<AttributeDomain> domains;
+};
+
+struct ScatterPointsOnMeshOp {
+  /* Input data. */
+  const GeometryNodePointDistributeMethod distribute_method;
+  const std::string &density_attribute_name;
+  const int seed;
+  const float density;
+
+  /* Output data. */
+  Vector<float3> &positions;
+  Vector<float3> &bary_coords;
+  Vector<int> &looptri_indices;
+  Set<AttributeInfo> &attributes;
+
+  void operator()(const GeometryComponent &component, blender::Span<blender::float4x4> transforms)
+  {
+    if (component.type() != GeometryComponentType::Mesh) {
+      return;
+    }
+
+    const MeshComponent &mesh_component = static_cast<const MeshComponent &>(component);
+    if (!mesh_component.has_mesh()) {
+      return;
+    }
+    const Mesh &mesh = *mesh_component.get_for_read();
+    for (const float4x4 &transform : transforms) {
+      switch (distribute_method) {
+        case GEO_NODE_POINT_DISTRIBUTE_RANDOM: {
+          const FloatReadAttribute density_factors = mesh_component.attribute_get_for_read<float>(
+              density_attribute_name, ATTR_DOMAIN_POINT, 1.0f);
+          sample_mesh_surface(mesh,
+                              transform,
+                              density,
+                              &density_factors,
+                              seed,
+                              positions,
+                              bary_coords,
+                              looptri_indices);
+          break;
+        }
+        case GEO_NODE_POINT_DISTRIBUTE_POISSON:
+          sample_mesh_surface(
+              mesh, transform, density, nullptr, seed, positions, bary_coords, looptri_indices);
+          break;
+      }
+    }
+  }
+};
+
+struct PoissonEliminateFromDensityOp {
+  /* Input data. */
+  const std::string &density_attribute_name;
+  const float density;
+  Span<float3> bary_coords;
+  Span<int> looptri_indices;
+
+  /* Output data. */
+  MutableSpan<bool> elimination_mask;
+
+  void operator()(const GeometryComponent &component, blender::Span<blender::float4x4> transforms)
+  {
+    if (component.type() != GeometryComponentType::Mesh) {
+      return;
+    }
+    const MeshComponent &mesh_component = static_cast<const MeshComponent &>(component);
+    if (!mesh_component.has_mesh()) {
+      return;
+    }
+    const Mesh &mesh = *mesh_component.get_for_read();
+    for (const float4x4 &transform : transforms) {
+      const FloatReadAttribute density_factors = mesh_component.attribute_get_for_read<float>(
+          density_attribute_name, ATTR_DOMAIN_POINT, 1.0f);
+      update_elimination_mask_based_on_density_factors(
+          mesh, density_factors, bary_coords, looptri_indices, elimination_mask);
+    }
+  }
+};
+
+struct AttributeInterpolateOp {
+  /* Input data. */
+  const std::string &density_attribute_name;
+  const float density;
+  Span<float3> bary_coords;
+  Span<int> looptri_indices;
+
+  int index_offset = 0;
+
+  /* Output data. */
+  MutableSpan<bool> elimination_mask;
+
+  void operator()(const GeometryComponent &component, blender::Span<blender::float4x4> transforms)
+  {
+    if (component.type() != GeometryComponentType::Mesh) {
+      return;
+    }
+    const MeshComponent &mesh_component = static_cast<const MeshComponent &>(component);
+    if (!mesh_component.has_mesh()) {
+      return;
+    }
+    const Mesh &mesh = *mesh_component.get_for_read();
+    for (const float4x4 &transform : transforms) {
+
+      const FloatReadAttribute density_factors = mesh_component.attribute_get_for_read<float>(
+          density_attribute_name, ATTR_DOMAIN_POINT, 1.0f);
+      update_elimination_mask_based_on_density_factors(
+          mesh, density_factors, bary_coords, looptri_indices, elimination_mask);
+    }
+  }
+};
 
 static void geo_node_point_distribute_exec(GeoNodeExecParams params)
 {
   GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
   GeometrySet geometry_set_out;
 
-  GeometryNodePointDistributeMethod distribute_method =
+  const GeometryNodePointDistributeMethod distribute_method =
       static_cast<GeometryNodePointDistributeMethod>(params.node().custom1);
 
-  if (!geometry_set.has_mesh()) {
+  if (!geometry_set.has_mesh() && !geometry_set.has_instances()) {
     params.set_output("Geometry", std::move(geometry_set_out));
     return;
   }
 
   const float density = params.extract_input<float>("Density Max");
-  const std::string density_attribute = params.extract_input<std::string>("Density Attribute");
+  const std::string density_attribute_name = params.extract_input<std::string>(
+      "Density Attribute");
 
   if (density <= 0.0f) {
     params.set_output("Geometry", std::move(geometry_set_out));
     return;
   }
 
-  const MeshComponent &mesh_component = *geometry_set.get_component_for_read<MeshComponent>();
-  const Mesh *mesh_in = mesh_component.get_for_read();
-
-  if (mesh_in == nullptr || mesh_in->mpoly == nullptr) {
-    params.set_output("Geometry", std::move(geometry_set_out));
-    return;
-  }
-
-  const FloatReadAttribute density_factors = mesh_component.attribute_get_for_read<float>(
-      density_attribute, ATTR_DOMAIN_POINT, 1.0f);
   const int seed = params.get_input<int>("Seed");
 
   Vector<float3> positions;
   Vector<float3> bary_coords;
   Vector<int> looptri_indices;
-  switch (distribute_method) {
-    case GEO_NODE_POINT_DISTRIBUTE_RANDOM:
-      sample_mesh_surface(
-          *mesh_in, density, &density_factors, seed, positions, bary_coords, looptri_indices);
-      break;
-    case GEO_NODE_POINT_DISTRIBUTE_POISSON:
-      const float minimum_distance = params.extract_input<float>("Distance Min");
-      sample_mesh_surface_with_minimum_distance(*mesh_in,
-                                                density,
-                                                minimum_distance,
-                                                density_factors,
-                                                seed,
-                                                positions,
-                                                bary_coords,
-                                                looptri_indices);
-      break;
+  Set<AttributeInfo> attributes;
+
+  ScatterPointsOnMeshOp scatter_points_op{distribute_method,
+                                          density_attribute_name,
+                                          seed,
+                                          density,
+                                          positions,
+                                          bary_coords,
+                                          looptri_indices,
+                                          attributes};
+  BKE_foreach_geometry_component_recursive(geometry_set, scatter_points_op);
+
+  /* Eliminate points based on the minimum distance for the poisson disk case. */
+  if (distribute_method == GEO_NODE_POINT_DISTRIBUTE_POISSON) {
+    Array<bool> elimination_mask(positions.size(), false);
+    const float minimum_distance = params.get_input<float>("Distance Min");
+
+    update_elimination_mask_for_close_points(positions, minimum_distance, elimination_mask);
+
+    PoissonEliminateFromDensityOp eliminate_density_op{
+        density_attribute_name, density, bary_coords, looptri_indices};
+    BKE_foreach_geometry_component_recursive(geometry_set, eliminate_density_op);
+
+    eliminate_points_based_on_mask(elimination_mask, positions, bary_coords, looptri_indices);
   }
+
   const int tot_points = positions.size();
 
   PointCloud *pointcloud = BKE_pointcloud_new_nomain(tot_points);
@@ -470,7 +569,7 @@ static void geo_node_point_distribute_exec(GeoNodeExecParams params)
       geometry_set_out.get_component_for_write<PointCloudComponent>();
   point_component.replace(pointcloud);
 
-  add_remaining_point_attributes(mesh_component, point_component, bary_coords, looptri_indices);
+  // add_remaining_point_attributes(mesh_component, point_component, bary_coords, looptri_indices);
 
   params.set_output("Geometry", std::move(geometry_set_out));
 }