Geometry Nodes: Accumulate Fields Node

This function node creates a running total of a given Vector, Float, or
Int field.

Inputs:
  - Value: The field to be accumulated
  - Group Index: The values of this input are used to aggregate the input
    into separate 'bins', creating multiple accumulations.
Outputs:
  - Leading and Trailing: Returns the running totals starting
   at either the first value of each accumulations or 0 respectively.
  - Total: Returns the total accumulation at all positions of the field.

There's currently plenty of duplicate work happening when multiple outputs
are used that could be optimized by a future refactor to field inputs.

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

4 files changed, 434 insertions, 0 deletions

diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index c6efe6efb62..2735dc445e2 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -49,6 +49,7 @@ void register_node_type_geo_legacy_select_by_material(void);
 void register_node_type_geo_legacy_subdivision_surface(void);
 void register_node_type_geo_legacy_volume_to_mesh(void);
 
+void register_node_type_geo_accumulate_field(void);
 void register_node_type_geo_align_rotation_to_vector(void);
 void register_node_type_geo_attribute_capture(void);
 void register_node_type_geo_attribute_clamp(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 0c6e42deb27..108a37d0176 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -347,6 +347,7 @@ DefNode(GeometryNode, GEO_NODE_CURVE_TO_MESH, 0, "CURVE_TO_MESH", CurveToMesh, "
 DefNode(GeometryNode, GEO_NODE_CURVE_TO_POINTS, def_geo_curve_to_points, "CURVE_TO_POINTS", CurveToPoints, "Curve to Points", "")
 DefNode(GeometryNode, GEO_NODE_DELETE_GEOMETRY, def_geo_delete_geometry, "DELETE_GEOMETRY", DeleteGeometry, "Delete Geometry", "")
 DefNode(GeometryNode, GEO_NODE_DISTRIBUTE_POINTS_ON_FACES, def_geo_distribute_points_on_faces, "DISTRIBUTE_POINTS_ON_FACES", DistributePointsOnFaces, "Distribute Points on Faces", "")
+DefNode(GeometryNode, GEO_NODE_ACCUMULATE_FIELD, def_geo_accumulate_field, "ACCUMULATE_FIELD", AccumulateField, "Accumulate Field", "")
 DefNode(GeometryNode, GEO_NODE_DUAL_MESH, 0, "DUAL_MESH", DualMesh, "Dual Mesh", "")
 DefNode(GeometryNode, GEO_NODE_FILL_CURVE, def_geo_curve_fill, "FILL_CURVE", FillCurve, "Fill Curve", "")
 DefNode(GeometryNode, GEO_NODE_FILLET_CURVE, def_geo_curve_fillet, "FILLET_CURVE", FilletCurve, "Fillet Curve", "")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index c24f8849532..5d45fe9021a 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/source/blender/nodes/geometry/CMakeLists.txt
@@ -80,6 +80,7 @@ set(SRC
   nodes/legacy/node_geo_legacy_subdivision_surface.cc
   nodes/legacy/node_geo_legacy_volume_to_mesh.cc
 
+  nodes/node_geo_accumulate_field.cc
   nodes/node_geo_attribute_capture.cc
   nodes/node_geo_attribute_domain_size.cc
   nodes/node_geo_attribute_remove.cc
diff --git a/source/blender/nodes/geometry/nodes/node_geo_accumulate_field.cc b/source/blender/nodes/geometry/nodes/node_geo_accumulate_field.cc
new file mode 100644
index 00000000000..887a801aa2e
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_accumulate_field.cc
@@ -0,0 +1,431 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "BKE_attribute_math.hh"
+
+#include "NOD_socket_search_link.hh"
+
+#include "node_geometry_util.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+namespace blender::nodes::node_geo_accumulate_field_cc {
+
+NODE_STORAGE_FUNCS(NodeAccumulateField)
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  std::string value_in_description = "The values to be accumulated";
+  std::string leading_out_description =
+      "The running total of values in the corresponding group, starting at the first value";
+  std::string trailing_out_description =
+      "The running total of values in the corresponding group, starting at zero";
+  std::string total_out_description = "The total of all of the values in the corresponding group";
+
+  b.add_input<decl::Vector>(N_("Value"), "Value Vector")
+      .default_value({1.0f, 1.0f, 1.0f})
+      .supports_field()
+      .description(N_(value_in_description));
+  b.add_input<decl::Float>(N_("Value"), "Value Float")
+      .default_value(1.0f)
+      .supports_field()
+      .description(N_(value_in_description));
+  b.add_input<decl::Int>(N_("Value"), "Value Int")
+      .default_value(1)
+      .supports_field()
+      .description(N_(value_in_description));
+  b.add_input<decl::Int>(N_("Group Index"))
+      .supports_field()
+      .description(
+          N_("An index used to group values together for multiple separate accumulations"));
+
+  b.add_output<decl::Vector>(N_("Leading"), "Leading Vector")
+      .field_source()
+      .description(N_(leading_out_description));
+  b.add_output<decl::Float>(N_("Leading"), "Leading Float")
+      .field_source()
+      .description(N_(leading_out_description));
+  b.add_output<decl::Int>(N_("Leading"), "Leading Int")
+      .field_source()
+      .description(N_(leading_out_description));
+
+  b.add_output<decl::Vector>(N_("Trailing"), "Trailing Vector")
+      .field_source()
+      .description(N_(trailing_out_description));
+  b.add_output<decl::Float>(N_("Trailing"), "Trailing Float")
+      .field_source()
+      .description(N_(trailing_out_description));
+  b.add_output<decl::Int>(N_("Trailing"), "Trailing Int")
+      .field_source()
+      .description(N_(trailing_out_description));
+
+  b.add_output<decl::Vector>(N_("Total"), "Total Vector")
+      .field_source()
+      .description(N_(total_out_description));
+  b.add_output<decl::Float>(N_("Total"), "Total Float")
+      .field_source()
+      .description(N_(total_out_description));
+  b.add_output<decl::Int>(N_("Total"), "Total Int")
+      .field_source()
+      .description(N_(total_out_description));
+}
+
+static void node_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
+  uiItemR(layout, ptr, "domain", 0, "", ICON_NONE);
+}
+
+static void node_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  NodeAccumulateField *data = MEM_cnew<NodeAccumulateField>(__func__);
+  data->data_type = CD_PROP_FLOAT;
+  data->domain = ATTR_DOMAIN_POINT;
+  node->storage = data;
+}
+
+static void node_update(bNodeTree *ntree, bNode *node)
+{
+  const NodeAccumulateField &storage = node_storage(*node);
+  const CustomDataType data_type = static_cast<CustomDataType>(storage.data_type);
+
+  bNodeSocket *sock_in_vector = (bNodeSocket *)node->inputs.first;
+  bNodeSocket *sock_in_float = sock_in_vector->next;
+  bNodeSocket *sock_in_int = sock_in_float->next;
+
+  bNodeSocket *sock_out_vector = (bNodeSocket *)node->outputs.first;
+  bNodeSocket *sock_out_float = sock_out_vector->next;
+  bNodeSocket *sock_out_int = sock_out_float->next;
+
+  bNodeSocket *sock_out_first_vector = sock_out_int->next;
+  bNodeSocket *sock_out_first_float = sock_out_first_vector->next;
+  bNodeSocket *sock_out_first_int = sock_out_first_float->next;
+  bNodeSocket *sock_out_total_vector = sock_out_first_int->next;
+  bNodeSocket *sock_out_total_float = sock_out_total_vector->next;
+  bNodeSocket *sock_out_total_int = sock_out_total_float->next;
+
+  nodeSetSocketAvailability(ntree, sock_in_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_in_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_in_int, data_type == CD_PROP_INT32);
+
+  nodeSetSocketAvailability(ntree, sock_out_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_out_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_out_int, data_type == CD_PROP_INT32);
+
+  nodeSetSocketAvailability(ntree, sock_out_first_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_out_first_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_out_first_int, data_type == CD_PROP_INT32);
+
+  nodeSetSocketAvailability(ntree, sock_out_total_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_out_total_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_out_total_int, data_type == CD_PROP_INT32);
+}
+
+enum class AccumulationMode { Leading = 0, Trailing = 1 };
+
+static std::optional<CustomDataType> node_type_from_other_socket(const bNodeSocket &socket)
+{
+  switch (socket.type) {
+    case SOCK_FLOAT:
+      return CD_PROP_FLOAT;
+    case SOCK_BOOLEAN:
+    case SOCK_INT:
+      return CD_PROP_INT32;
+    case SOCK_VECTOR:
+    case SOCK_RGBA:
+      return CD_PROP_FLOAT3;
+    default:
+      return {};
+  }
+}
+
+static void node_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const std::optional<CustomDataType> type = node_type_from_other_socket(params.other_socket());
+  if (!type) {
+    return;
+  }
+  if (params.in_out() == SOCK_OUT) {
+    params.add_item(
+        IFACE_("Leading"),
+        [type](LinkSearchOpParams &params) {
+          bNode &node = params.add_node("GeometryNodeAccumulateField");
+          node_storage(node).data_type = *type;
+          params.update_and_connect_available_socket(node, "Leading");
+        },
+        0);
+    params.add_item(
+        IFACE_("Trailing"),
+        [type](LinkSearchOpParams &params) {
+          bNode &node = params.add_node("GeometryNodeAccumulateField");
+          node_storage(node).data_type = *type;
+          params.update_and_connect_available_socket(node, "Trailing");
+        },
+        -1);
+    params.add_item(
+        IFACE_("Total"),
+        [type](LinkSearchOpParams &params) {
+          bNode &node = params.add_node("GeometryNodeAccumulateField");
+          node_storage(node).data_type = *type;
+          params.update_and_connect_available_socket(node, "Total");
+        },
+        -2);
+  }
+  else {
+    params.add_item(
+        IFACE_("Value"),
+        [type](LinkSearchOpParams &params) {
+          bNode &node = params.add_node("GeometryNodeAccumulateField");
+          node_storage(node).data_type = *type;
+          params.update_and_connect_available_socket(node, "Value");
+        },
+        0);
+
+    params.add_item(
+        IFACE_("Group Index"),
+        [type](LinkSearchOpParams &params) {
+          bNode &node = params.add_node("GeometryNodeAccumulateField");
+          node_storage(node).data_type = *type;
+          params.update_and_connect_available_socket(node, "Group Index");
+        },
+        -1);
+  }
+}
+
+template<typename T> class AccumulateFieldInput final : public GeometryFieldInput {
+ private:
+  Field<T> input_;
+  Field<int> group_index_;
+  AttributeDomain source_domain_;
+  AccumulationMode accumulation_mode_;
+
+ public:
+  AccumulateFieldInput(const AttributeDomain source_domain,
+                       Field<T> input,
+                       Field<int> group_index,
+                       AccumulationMode accumulation_mode)
+      : GeometryFieldInput(CPPType::get<T>(), "Accumulation"),
+        input_(input),
+        group_index_(group_index),
+        source_domain_(source_domain),
+        accumulation_mode_(accumulation_mode)
+  {
+  }
+
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask UNUSED(mask)) const final
+  {
+    const GeometryComponentFieldContext field_context{component, source_domain_};
+    const int domain_size = component.attribute_domain_size(field_context.domain());
+
+    fn::FieldEvaluator evaluator{field_context, domain_size};
+    evaluator.add(input_);
+    evaluator.add(group_index_);
+    evaluator.evaluate();
+    const VArray<T> &values = evaluator.get_evaluated<T>(0);
+    const VArray<int> &group_indices = evaluator.get_evaluated<int>(1);
+
+    Array<T> accumulations_out(domain_size);
+
+    if (group_indices.is_single()) {
+      T accumulation = T();
+      if (accumulation_mode_ == AccumulationMode::Leading) {
+        for (const int i : values.index_range()) {
+          accumulation = values[i] + accumulation;
+          accumulations_out[i] = accumulation;
+        }
+      }
+      else {
+        for (const int i : values.index_range()) {
+          accumulations_out[i] = accumulation;
+          accumulation = values[i] + accumulation;
+        }
+      }
+    }
+    else {
+      Map<int, T> accumulations;
+      if (accumulation_mode_ == AccumulationMode::Leading) {
+        for (const int i : values.index_range()) {
+          T &accumulation_value = accumulations.lookup_or_add_default(group_indices[i]);
+          accumulation_value += values[i];
+          accumulations_out[i] = accumulation_value;
+        }
+      }
+      else {
+        for (const int i : values.index_range()) {
+          T &accumulation_value = accumulations.lookup_or_add_default(group_indices[i]);
+          accumulations_out[i] = accumulation_value;
+          accumulation_value += values[i];
+        }
+      }
+    }
+
+    return component.attribute_try_adapt_domain<T>(
+        VArray<T>::ForContainer(std::move(accumulations_out)), source_domain_, domain);
+  }
+
+  uint64_t hash() const override
+  {
+    return get_default_hash_4(input_, group_index_, source_domain_, accumulation_mode_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const override
+  {
+    if (const AccumulateFieldInput *other_accumulate = dynamic_cast<const AccumulateFieldInput *>(
+            &other)) {
+      return input_ == other_accumulate->input_ &&
+             group_index_ == other_accumulate->group_index_ &&
+             source_domain_ == other_accumulate->source_domain_ &&
+             accumulation_mode_ == other_accumulate->accumulation_mode_;
+    }
+    return false;
+  }
+};
+
+template<typename T> class TotalFieldInput final : public GeometryFieldInput {
+ private:
+  Field<T> input_;
+  Field<int> group_index_;
+  AttributeDomain source_domain_;
+
+ public:
+  TotalFieldInput(const AttributeDomain source_domain, Field<T> input, Field<int> group_index)
+      : GeometryFieldInput(CPPType::get<T>(), "Total Value"),
+        input_(input),
+        group_index_(group_index),
+        source_domain_(source_domain)
+  {
+  }
+
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask UNUSED(mask)) const final
+  {
+    const GeometryComponentFieldContext field_context{component, source_domain_};
+    const int domain_size = component.attribute_domain_size(field_context.domain());
+
+    fn::FieldEvaluator evaluator{field_context, domain_size};
+    evaluator.add(input_);
+    evaluator.add(group_index_);
+    evaluator.evaluate();
+    const VArray<T> &values = evaluator.get_evaluated<T>(0);
+    const VArray<int> &group_indices = evaluator.get_evaluated<int>(1);
+
+    if (group_indices.is_single()) {
+      T accumulation = T();
+      for (const int i : values.index_range()) {
+        accumulation = values[i] + accumulation;
+      }
+      return VArray<T>::ForSingle(accumulation, domain_size);
+    }
+
+    Array<T> accumulations_out(domain_size);
+    Map<int, T> accumulations;
+    for (const int i : values.index_range()) {
+      T &value = accumulations.lookup_or_add_default(group_indices[i]);
+      value = value + values[i];
+    }
+    for (const int i : values.index_range()) {
+      accumulations_out[i] = accumulations.lookup(group_indices[i]);
+    }
+
+    return component.attribute_try_adapt_domain<T>(
+        VArray<T>::ForContainer(std::move(accumulations_out)), source_domain_, domain);
+  }
+
+  uint64_t hash() const override
+  {
+    return get_default_hash_3(input_, group_index_, source_domain_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const override
+  {
+    if (const TotalFieldInput *other_field = dynamic_cast<const TotalFieldInput *>(&other)) {
+      return input_ == other_field->input_ && group_index_ == other_field->group_index_ &&
+             source_domain_ == other_field->source_domain_;
+    }
+    return false;
+  }
+};
+
+template<typename T> std::string identifier_suffix()
+{
+  if constexpr (std::is_same_v<T, int>) {
+    return "Int";
+  }
+  if constexpr (std::is_same_v<T, float>) {
+    return "Float";
+  }
+  if constexpr (std::is_same_v<T, float3>) {
+    return "Vector";
+  }
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  const NodeAccumulateField &storage = node_storage(params.node());
+  const CustomDataType data_type = static_cast<CustomDataType>(storage.data_type);
+  const AttributeDomain source_domain = static_cast<AttributeDomain>(storage.domain);
+
+  Field<int> group_index_field = params.extract_input<Field<int>>("Group Index");
+  attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
+    using T = decltype(dummy);
+    if constexpr (std::is_same_v<T, int> || std::is_same_v<T, float> ||
+                  std::is_same_v<T, float3>) {
+      const std::string suffix = " " + identifier_suffix<T>();
+      Field<T> input_field = params.extract_input<Field<T>>("Value" + suffix);
+      if (params.output_is_required("Leading" + suffix)) {
+        params.set_output(
+            "Leading" + suffix,
+            Field<T>{std::make_shared<AccumulateFieldInput<T>>(
+                source_domain, input_field, group_index_field, AccumulationMode::Leading)});
+      }
+      if (params.output_is_required("Trailing" + suffix)) {
+        params.set_output(
+            "Trailing" + suffix,
+            Field<T>{std::make_shared<AccumulateFieldInput<T>>(
+                source_domain, input_field, group_index_field, AccumulationMode::Trailing)});
+      }
+      if (params.output_is_required("Total" + suffix)) {
+        params.set_output("Total" + suffix,
+                          Field<T>{std::make_shared<TotalFieldInput<T>>(
+                              source_domain, input_field, group_index_field)});
+      }
+    }
+  });
+}
+}  // namespace blender::nodes::node_geo_accumulate_field_cc
+
+void register_node_type_geo_accumulate_field()
+{
+  namespace file_ns = blender::nodes::node_geo_accumulate_field_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_ACCUMULATE_FIELD, "Accumulate Field", NODE_CLASS_CONVERTER, 0);
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  node_type_init(&ntype, file_ns::node_init);
+  node_type_update(&ntype, file_ns::node_update);
+  ntype.draw_buttons = file_ns::node_layout;
+  ntype.declare = file_ns::node_declare;
+  ntype.gather_link_search_ops = file_ns::node_gather_link_searches;
+  node_type_storage(
+      &ntype, "NodeAccumulateField", node_free_standard_storage, node_copy_standard_storage);
+  nodeRegisterType(&ntype);
+}