simplify math nodegeometry-nodes-deduplicate-float-math

1 files changed, 47 insertions, 196 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
index 810bc6734f6..d8c2b75a0a9 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
@@ -26,6 +26,8 @@
 #include "DNA_mesh_types.h"
 #include "DNA_pointcloud_types.h"
 
+#include "NOD_math_functions.hh"
+
 static bNodeSocketTemplate geo_node_attribute_math_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
     {SOCK_STRING, N_("Attribute A")},
@@ -64,221 +66,70 @@ static void geo_node_random_attribute_update(bNodeTree *UNUSED(ntree), bNode *no
 
 namespace blender::nodes {
 
-/**
- * Do implicit conversion from other attribute types to float.
- */
-static Array<float> attribute_ensure_float_type(ReadAttributePtr attribute)
-{
-  Array<float> array(attribute->size());
-  if (attribute->cpp_type().is<float>()) {
-    FloatReadAttribute float_attribute = std::move(attribute);
-    for (const int i : IndexRange(float_attribute.size())) {
-      array[i] = float_attribute[i];
-    }
-  }
-  else if (attribute->cpp_type().is<float3>()) {
-    Float3ReadAttribute float3_attribute = std::move(attribute);
-    for (const int i : IndexRange(float3_attribute.size())) {
-      array[i] = float3_attribute[i].length();
-    }
-  }
-#if 0 /* More CPP types to support in the future. */
-  else if (attribute->cpp_type().is<Color4f>()) {
-  }
-  else if (attribute->cpp_type().is<float2>()) {
-  }
-  else if (attribute->cpp_type().is<int>()) {
-  }
-#endif
-
-  return array;
-}
-
-static void math_operation_attribute_float(Array<float> input_a,
-                                           float input_b,
-                                           FloatWriteAttribute float_attribute,
-                                           const int operation)
+static void do_math_operation(const FloatReadAttribute &input_a,
+                              const FloatReadAttribute &input_b,
+                              FloatWriteAttribute result,
+                              const int operation)
 {
-  switch (operation) {
-    case NODE_MATH_ADD:
-      for (const int i : IndexRange(input_a.size())) {
-        float_attribute.set(i, input_a[i] + input_b);
-      }
-      break;
-    case NODE_MATH_SUBTRACT:
-      for (const int i : IndexRange(input_a.size())) {
-        float_attribute.set(i, input_a[i] - input_b);
-      }
-      break;
-    case NODE_MATH_MULTIPLY:
-      for (const int i : IndexRange(input_a.size())) {
-        float_attribute.set(i, input_a[i] * input_b);
-      }
-      break;
-    case NODE_MATH_DIVIDE:
-      /* Protect against division by zero. */
-      if (input_b == 0.0f) {
-        for (const int i : IndexRange(input_a.size())) {
-          float_attribute.set(i, 0.0f);
-        }
-      }
-      else {
-        for (const int i : IndexRange(input_a.size())) {
-          float_attribute.set(i, input_a[i] / input_b);
+  const int size = input_a.size();
+
+  bool success = try_dispatch_float_math_fl_fl_to_fl(
+      operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
+        for (const int i : IndexRange(size)) {
+          const float in1 = input_a[i];
+          const float in2 = input_b[i];
+          const float out = math_function(in1, in2);
+          result.set(i, out);
         }
-      }
-      break;
-  }
-}
+      });
 
-static void math_operation_float_attribute(float input_a,
-                                           Array<float> input_b,
-                                           FloatWriteAttribute float_attribute,
-                                           const int operation)
-{
-  switch (operation) {
-    case NODE_MATH_ADD:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a + input_b[i]);
-      }
-      break;
-    case NODE_MATH_SUBTRACT:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a - input_b[i]);
-      }
-      break;
-    case NODE_MATH_MULTIPLY:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a * input_b[i]);
-      }
-      break;
-    case NODE_MATH_DIVIDE:
-      for (const int i : IndexRange(float_attribute.size())) {
-        /* Protect against division by zero. */
-        float_attribute.set(i, safe_divide(input_a, input_b[i]));
-      }
-      break;
-  }
-}
-
-static void math_operation_attribute_attribute(Array<float> input_a,
-                                               Array<float> input_b,
-                                               FloatWriteAttribute float_attribute,
-                                               const int operation)
-{
-  switch (operation) {
-    case NODE_MATH_ADD:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a[i] + input_b[i]);
-      }
-      break;
-    case NODE_MATH_SUBTRACT:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a[i] - input_b[i]);
-      }
-      break;
-    case NODE_MATH_MULTIPLY:
-      for (const int i : IndexRange(float_attribute.size())) {
-        float_attribute.set(i, input_a[i] * input_b[i]);
-      }
-      break;
-    case NODE_MATH_DIVIDE:
-      for (const int i : IndexRange(float_attribute.size())) {
-        /* Protect against division by zero. */
-        float_attribute.set(i, safe_divide(input_a[i], input_b[i]));
-      }
-      break;
-  }
+  /* The operation is not supported by this node currently. */
+  BLI_assert(success);
+  UNUSED_VARS_NDEBUG(success);
 }
 
 static void attribute_math_calc(GeometryComponent &component, const GeoNodeExecParams &params)
 {
   const bNode &node = params.node();
-  const int operation = params.node().custom1;
-  const std::string result_name = params.get_input<std::string>("Result");
+  const int operation = node.custom1;
 
   /* The result type of this node is always float. */
-  const CustomDataType result_type = bke::cpp_type_to_custom_data_type(CPPType::get<float>());
-
-  /* Use a single float for none or one of the inputs, which leads to three cases,
-   * because the order of the inputs in the operations can matter. */
-  GeometryNodeUseAttributeFlag flag = static_cast<GeometryNodeUseAttributeFlag>(node.custom2);
-  if ((flag & GEO_NODE_USE_ATTRIBUTE_A) && (flag & GEO_NODE_USE_ATTRIBUTE_B)) {
-    /* Two attributes. */
-    const std::string attribute_a_name = params.get_input<std::string>("Attribute A");
-    const std::string attribute_b_name = params.get_input<std::string>("Attribute B");
-    ReadAttributePtr attribute_a = component.attribute_try_get_for_read(attribute_a_name);
-    ReadAttributePtr attribute_b = component.attribute_try_get_for_read(attribute_b_name);
-    if (!attribute_a || !attribute_b) {
-      return;
-    }
-
-    /* Don't handle domain interpolation for now. */
-    AttributeDomain domain_a = attribute_a->domain();
-    AttributeDomain domain_b = attribute_b->domain();
-    if (domain_a != domain_b) {
-      return;
-    }
+  const CustomDataType result_type = CD_PROP_FLOAT;
+  /* The result domain is always point for now. */
+  const AttributeDomain result_domain = ATTR_DOMAIN_POINT;
 
-    BLI_assert(attribute_a->size() == attribute_b->size());
-
-    Array<float> input_a = attribute_ensure_float_type(std::move(attribute_a));
-    Array<float> input_b = attribute_ensure_float_type(std::move(attribute_b));
-
-    WriteAttributePtr result = component.attribute_try_ensure_for_write(
-        result_name, domain_a, result_type);
-    if (!result) {
-      return;
-    }
-
-    math_operation_attribute_attribute(input_a, input_b, std::move(result), operation);
+  /* Get result attribute first, in case it has to overwrite one of the existing attributes. */
+  const std::string result_name = params.get_input<std::string>("Result");
+  WriteAttributePtr attribute_result = component.attribute_try_ensure_for_write(
+      result_name, result_domain, result_type);
+  if (!attribute_result) {
+    return;
   }
-  else if (flag & GEO_NODE_USE_ATTRIBUTE_A) {
-    /* Attribute and float. */
-    const std::string attribute_a_name = params.get_input<std::string>("Attribute A");
-    ReadAttributePtr attribute_a = component.attribute_try_get_for_read(attribute_a_name);
-    if (!attribute_a) {
-      return;
-    }
-
-    AttributeDomain domain = attribute_a->domain();
-    Array<float> input_a = attribute_ensure_float_type(std::move(attribute_a));
-    const float input_b = params.get_input<float>("B");
 
-    WriteAttributePtr result = component.attribute_try_ensure_for_write(
-        result_name, domain, result_type);
-    if (!result) {
-      return;
-    }
+  GeometryNodeUseAttributeFlag flag = static_cast<GeometryNodeUseAttributeFlag>(node.custom2);
 
-    math_operation_attribute_float(input_a, input_b, std::move(result), operation);
-  }
-  else if (flag & GEO_NODE_USE_ATTRIBUTE_B) {
-    /* Float and attribute. */
-    const std::string attribute_b_name = params.get_input<std::string>("Attribute B");
-    ReadAttributePtr attribute_b = component.attribute_try_get_for_read(attribute_b_name);
-    if (!attribute_b) {
-      return;
+  auto get_input_attribute = [&](GeometryNodeUseAttributeFlag use_flag,
+                                 StringRef attribute_socket_identifier,
+                                 StringRef value_socket_identifier) {
+    if (flag & use_flag) {
+      const std::string attribute_name = params.get_input<std::string>(
+          attribute_socket_identifier);
+      return component.attribute_try_get_for_read(attribute_name, result_domain, result_type);
     }
+    const float value = params.get_input<float>(value_socket_identifier);
+    return component.attribute_get_constant_for_read(result_domain, result_type, &value);
+  };
 
-    AttributeDomain domain = attribute_b->domain();
-    Array<float> input_b = attribute_ensure_float_type(std::move(attribute_b));
-    const float input_a = params.get_input<float>("A");
-
-    WriteAttributePtr result = component.attribute_try_ensure_for_write(
-        result_name, domain, result_type);
-    if (!result) {
-      return;
-    }
+  ReadAttributePtr attribute_a = get_input_attribute(GEO_NODE_USE_ATTRIBUTE_A, "Attribute A", "A");
+  ReadAttributePtr attribute_b = get_input_attribute(GEO_NODE_USE_ATTRIBUTE_B, "Attribute B", "B");
 
-    math_operation_float_attribute(input_a, input_b, std::move(result), operation);
-  }
-  else {
-    /* There was no attribute provided, so just return. Otherwise choosing which domain to use
-     * for a new attribute would be completely arbitrary, and we don't want to make those sort
-     * of decisions in the attribute math node. */
+  if (!attribute_a || !attribute_b) {
+    /* Attribute wasn't found. */
     return;
   }
+
+  do_math_operation(
+      std::move(attribute_a), std::move(attribute_b), std::move(attribute_result), operation);
 }
 
 static void geo_node_attribute_math_exec(GeoNodeExecParams params)