Geometry Nodes: Refactor / fix point separate node

The point separate node was failing in situations where one of the
outputs was empty. In addition, the code was not structured very well.
This new implementation stores less temporary information, is more
geometry component type agnostic, and is more self-descriptive.
It also solves the problems mentioned above.

Fixes T86573

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

1 files changed, 68 insertions, 120 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_point_separate.cc b/source/blender/nodes/geometry/nodes/node_geo_point_separate.cc
index 53d0fe1e112..522dea4aa0e 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_point_separate.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_point_separate.cc
@@ -14,12 +14,11 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BKE_attribute_math.hh"
 #include "BKE_mesh.h"
-#include "BKE_persistent_data_handle.hh"
 #include "BKE_pointcloud.h"
 
 #include "DNA_mesh_types.h"
-#include "DNA_meshdata_types.h"
 #include "DNA_pointcloud_types.h"
 
 #include "node_geometry_util.hh"
@@ -38,166 +37,115 @@ static bNodeSocketTemplate geo_node_point_instance_out[] = {
 
 namespace blender::nodes {
 
-static void fill_new_attribute_from_input(const ReadAttribute &input_attribute,
-                                          WriteAttribute &out_attribute_a,
-                                          WriteAttribute &out_attribute_b,
-                                          Span<bool> a_or_b)
+template<typename T>
+static void copy_data_based_on_mask(Span<T> data,
+                                    Span<bool> masks,
+                                    const bool invert,
+                                    MutableSpan<T> out_data)
 {
-  fn::GSpan in_span = input_attribute.get_span();
-  int i_a = 0;
-  int i_b = 0;
-  for (int i_in = 0; i_in < in_span.size(); i_in++) {
-    const bool move_to_b = a_or_b[i_in];
-    if (move_to_b) {
-      out_attribute_b.set(i_b, in_span[i_in]);
-      i_b++;
-    }
-    else {
-      out_attribute_a.set(i_a, in_span[i_in]);
-      i_a++;
+  int offset = 0;
+  for (const int i : data.index_range()) {
+    if (masks[i] != invert) {
+      out_data[offset] = data[i];
+      offset++;
     }
   }
 }
 
-/**
- * Move the original attribute values to the two output components.
- *
- * \note This assumes a consistent ordering of indices before and after the split,
- * which is true for points and a simple vertex array.
- */
-static void move_split_attributes(const GeometryComponent &in_component,
-                                  GeometryComponent &out_component_a,
-                                  GeometryComponent &out_component_b,
-                                  Span<bool> a_or_b)
+static void copy_attributes_based_on_mask(const GeometryComponent &in_component,
+                                          GeometryComponent &result_component,
+                                          Span<bool> masks,
+                                          const bool invert)
 {
-  Set<std::string> attribute_names = in_component.attribute_names();
-
-  for (const std::string &name : attribute_names) {
+  for (const std::string &name : in_component.attribute_names()) {
     ReadAttributePtr attribute = in_component.attribute_try_get_for_read(name);
-    BLI_assert(attribute);
+    const CustomDataType data_type = attribute->custom_data_type();
 
-    /* Since this node only creates points and vertices, don't copy other attributes. */
+    /* Only copy point attributes. Theoretically this could interpolate attributes on other
+     * domains to the point domain, but that would conflict with attributes that are built-in
+     * on other domains, which causes creating the attributes to fail. */
     if (attribute->domain() != ATTR_DOMAIN_POINT) {
       continue;
     }
 
-    const CustomDataType data_type = bke::cpp_type_to_custom_data_type(attribute->cpp_type());
-    const AttributeDomain domain = attribute->domain();
+    OutputAttributePtr result_attribute = result_component.attribute_try_get_for_output(
+        name, ATTR_DOMAIN_POINT, data_type);
 
-    /* Don't try to create the attribute on the new component if it already exists (i.e. has been
-     * initialized by someone else). */
-    if (!out_component_a.attribute_exists(name)) {
-      if (!out_component_a.attribute_try_create(name, domain, data_type)) {
-        continue;
-      }
-    }
-    if (!out_component_b.attribute_exists(name)) {
-      if (!out_component_b.attribute_try_create(name, domain, data_type)) {
-        continue;
-      }
-    }
+    attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
+      using T = decltype(dummy);
+      Span<T> span = attribute->get_span<T>();
+      MutableSpan<T> out_span = result_attribute->get_span_for_write_only<T>();
+      copy_data_based_on_mask(span, masks, invert, out_span);
+    });
 
-    WriteAttributePtr out_attribute_a = out_component_a.attribute_try_get_for_write(name);
-    WriteAttributePtr out_attribute_b = out_component_b.attribute_try_get_for_write(name);
-    if (!out_attribute_a || !out_attribute_b) {
-      BLI_assert(false);
-      continue;
-    }
-
-    fill_new_attribute_from_input(*attribute, *out_attribute_a, *out_attribute_b, a_or_b);
+    result_attribute.apply_span_and_save();
   }
 }
 
-/**
- * Find total in each new set and find which of the output sets each point will belong to.
- */
-static Array<bool> count_point_splits(const GeometryComponent &component,
-                                      const GeoNodeExecParams &params,
-                                      int *r_a_total,
-                                      int *r_b_total)
+static void create_component_points(GeometryComponent &component, const int total)
 {
-  const BooleanReadAttribute mask_attribute = params.get_input_attribute<bool>(
-      "Mask", component, ATTR_DOMAIN_POINT, false);
-  Array<bool> masks = mask_attribute.get_span();
-  const int in_total = masks.size();
-
-  *r_b_total = 0;
-  for (const bool mask : masks) {
-    if (mask) {
-      *r_b_total += 1;
-    }
+  switch (component.type()) {
+    case GEO_COMPONENT_TYPE_MESH:
+      static_cast<MeshComponent &>(component).replace(BKE_mesh_new_nomain(total, 0, 0, 0, 0));
+      break;
+    case GEO_COMPONENT_TYPE_POINT_CLOUD:
+      static_cast<PointCloudComponent &>(component).replace(BKE_pointcloud_new_nomain(total));
+      break;
+    default:
+      BLI_assert(false);
+      break;
   }
-  *r_a_total = in_total - *r_b_total;
-
-  return masks;
 }
 
-static void separate_mesh(const MeshComponent &in_component,
-                          const GeoNodeExecParams &params,
-                          MeshComponent &out_component_a,
-                          MeshComponent &out_component_b)
+static void separate_points_from_component(const GeometryComponent &in_component,
+                                           GeometryComponent &out_component,
+                                           const StringRef mask_name,
+                                           const bool invert)
 {
-  const int size = in_component.attribute_domain_size(ATTR_DOMAIN_POINT);
-  if (size == 0) {
+  if (!in_component.attribute_domain_supported(ATTR_DOMAIN_POINT) ||
+      in_component.attribute_domain_size(ATTR_DOMAIN_POINT) == 0) {
     return;
   }
 
-  int a_total;
-  int b_total;
-  Array<bool> a_or_b = count_point_splits(in_component, params, &a_total, &b_total);
-
-  out_component_a.replace(BKE_mesh_new_nomain(a_total, 0, 0, 0, 0));
-  out_component_b.replace(BKE_mesh_new_nomain(b_total, 0, 0, 0, 0));
-
-  move_split_attributes(in_component, out_component_a, out_component_b, a_or_b);
-}
+  const BooleanReadAttribute mask_attribute = in_component.attribute_get_for_read<bool>(
+      mask_name, ATTR_DOMAIN_POINT, false);
+  Span<bool> masks = mask_attribute.get_span();
 
-static void separate_point_cloud(const PointCloudComponent &in_component,
-                                 const GeoNodeExecParams &params,
-                                 PointCloudComponent &out_component_a,
-                                 PointCloudComponent &out_component_b)
-{
-  const int size = in_component.attribute_domain_size(ATTR_DOMAIN_POINT);
-  if (size == 0) {
+  const int total = masks.count(!invert);
+  if (total == 0) {
     return;
   }
 
-  int a_total;
-  int b_total;
-  Array<bool> a_or_b = count_point_splits(in_component, params, &a_total, &b_total);
+  create_component_points(out_component, total);
 
-  out_component_a.replace(BKE_pointcloud_new_nomain(a_total));
-  out_component_b.replace(BKE_pointcloud_new_nomain(b_total));
+  copy_attributes_based_on_mask(in_component, out_component, masks, invert);
+}
 
-  move_split_attributes(in_component, out_component_a, out_component_b, a_or_b);
+static GeometrySet separate_geometry_set(const GeometrySet &set_in,
+                                         const StringRef mask_name,
+                                         const bool invert)
+{
+  GeometrySet set_out;
+  for (const GeometryComponent *component : set_in.get_components_for_read()) {
+    GeometryComponent &out_component = set_out.get_component_for_write(component->type());
+    separate_points_from_component(*component, out_component, mask_name, invert);
+  }
+  return set_out;
 }
 
 static void geo_node_point_separate_exec(GeoNodeExecParams params)
 {
+  const std::string mask_attribute_name = params.extract_input<std::string>("Mask");
   GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
-  GeometrySet out_set_a(geometry_set);
-  GeometrySet out_set_b;
 
   /* TODO: This is not necessary-- the input geometry set can be read only,
    * but it must be rewritten to handle instance groups. */
   geometry_set = geometry_set_realize_instances(geometry_set);
 
-  if (geometry_set.has<PointCloudComponent>()) {
-    separate_point_cloud(*geometry_set.get_component_for_read<PointCloudComponent>(),
-                         params,
-                         out_set_a.get_component_for_write<PointCloudComponent>(),
-                         out_set_b.get_component_for_write<PointCloudComponent>());
-  }
-  if (geometry_set.has<MeshComponent>()) {
-    separate_mesh(*geometry_set.get_component_for_read<MeshComponent>(),
-                  params,
-                  out_set_a.get_component_for_write<MeshComponent>(),
-                  out_set_b.get_component_for_write<MeshComponent>());
-  }
-
-  params.set_output("Geometry 1", std::move(out_set_a));
-  params.set_output("Geometry 2", std::move(out_set_b));
+  params.set_output("Geometry 1", separate_geometry_set(geometry_set, mask_attribute_name, true));
+  params.set_output("Geometry 2", separate_geometry_set(geometry_set, mask_attribute_name, false));
 }
+
 }  // namespace blender::nodes
 
 void register_node_type_geo_point_separate()