Geometry Nodes: Support all operations in the "Attribute Math" node

This adds the ability to use all the math operations in the regular
utility "Math" node. The code is quite similar to the attribute vector
math node, with the simplification that the result is always a float.

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

1 files changed, 145 insertions, 27 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 f8ec9124db3..f3fc45fc1be 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
@@ -34,6 +34,8 @@ static bNodeSocketTemplate geo_node_attribute_math_in[] = {
     {SOCK_FLOAT, N_("A"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
     {SOCK_STRING, N_("B")},
     {SOCK_FLOAT, N_("B"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_STRING, N_("C")},
+    {SOCK_FLOAT, N_("C"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
     {SOCK_STRING, N_("Result")},
     {-1, ""},
 };
@@ -51,45 +53,132 @@ static void geo_node_attribute_math_init(bNodeTree *UNUSED(tree), bNode *node)
   data->operation = NODE_MATH_ADD;
   data->input_type_a = GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE;
   data->input_type_b = GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE;
+  data->input_type_c = GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE;
   node->storage = data;
 }
 
+static bool operation_use_input_c(const NodeMathOperation operation)
+{
+  return ELEM(operation,
+              NODE_MATH_MULTIPLY_ADD,
+              NODE_MATH_SMOOTH_MIN,
+              NODE_MATH_SMOOTH_MAX,
+              NODE_MATH_WRAP,
+              NODE_MATH_COMPARE);
+}
+
+static bool operation_use_input_b(const NodeMathOperation operation)
+{
+  switch (operation) {
+    case NODE_MATH_ADD:
+    case NODE_MATH_SUBTRACT:
+    case NODE_MATH_MULTIPLY:
+    case NODE_MATH_DIVIDE:
+    case NODE_MATH_POWER:
+    case NODE_MATH_LOGARITHM:
+    case NODE_MATH_MINIMUM:
+    case NODE_MATH_MAXIMUM:
+    case NODE_MATH_LESS_THAN:
+    case NODE_MATH_GREATER_THAN:
+    case NODE_MATH_MODULO:
+    case NODE_MATH_ARCTAN2:
+    case NODE_MATH_SNAP:
+    case NODE_MATH_WRAP:
+    case NODE_MATH_COMPARE:
+    case NODE_MATH_MULTIPLY_ADD:
+    case NODE_MATH_PINGPONG:
+    case NODE_MATH_SMOOTH_MIN:
+    case NODE_MATH_SMOOTH_MAX:
+      return true;
+    case NODE_MATH_SINE:
+    case NODE_MATH_COSINE:
+    case NODE_MATH_TANGENT:
+    case NODE_MATH_ARCSINE:
+    case NODE_MATH_ARCCOSINE:
+    case NODE_MATH_ARCTANGENT:
+    case NODE_MATH_ROUND:
+    case NODE_MATH_ABSOLUTE:
+    case NODE_MATH_FLOOR:
+    case NODE_MATH_CEIL:
+    case NODE_MATH_FRACTION:
+    case NODE_MATH_SQRT:
+    case NODE_MATH_INV_SQRT:
+    case NODE_MATH_SIGN:
+    case NODE_MATH_EXPONENT:
+    case NODE_MATH_RADIANS:
+    case NODE_MATH_DEGREES:
+    case NODE_MATH_SINH:
+    case NODE_MATH_COSH:
+    case NODE_MATH_TANH:
+    case NODE_MATH_TRUNC:
+      return false;
+  }
+  BLI_assert(false);
+  return false;
+}
+
 namespace blender::nodes {
 
 static void geo_node_attribute_math_update(bNodeTree *UNUSED(ntree), bNode *node)
 {
-  NodeAttributeMath *node_storage = (NodeAttributeMath *)node->storage;
+  NodeAttributeMath &node_storage = *(NodeAttributeMath *)node->storage;
+  NodeMathOperation operation = static_cast<NodeMathOperation>(node_storage.operation);
 
   update_attribute_input_socket_availabilities(
-      *node, "A", (GeometryNodeAttributeInputMode)node_storage->input_type_a);
+      *node, "A", (GeometryNodeAttributeInputMode)node_storage.input_type_a);
   update_attribute_input_socket_availabilities(
-      *node, "B", (GeometryNodeAttributeInputMode)node_storage->input_type_b);
+      *node,
+      "B",
+      (GeometryNodeAttributeInputMode)node_storage.input_type_b,
+      operation_use_input_b(operation));
+  update_attribute_input_socket_availabilities(
+      *node,
+      "C",
+      (GeometryNodeAttributeInputMode)node_storage.input_type_c,
+      operation_use_input_c(operation));
 }
 
-static void do_math_operation(const FloatReadAttribute &input_a,
-                              const FloatReadAttribute &input_b,
-                              FloatWriteAttribute result,
-                              const int operation)
+static void do_math_operation(Span<float> span_a,
+                              Span<float> span_b,
+                              Span<float> span_c,
+                              MutableSpan<float> span_result,
+                              const NodeMathOperation operation)
 {
-  const int size = input_a.size();
-
-  Span<float> span_a = input_a.get_span();
-  Span<float> span_b = input_b.get_span();
-  MutableSpan<float> span_result = result.get_span_for_write_only();
+  bool success = try_dispatch_float_math_fl_fl_fl_to_fl(
+      operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
+        for (const int i : IndexRange(span_result.size())) {
+          span_result[i] = math_function(span_a[i], span_b[i], span_c[i]);
+        }
+      });
+  BLI_assert(success);
+  UNUSED_VARS_NDEBUG(success);
+}
 
+static void do_math_operation(Span<float> span_a,
+                              Span<float> span_b,
+                              MutableSpan<float> span_result,
+                              const NodeMathOperation operation)
+{
   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 = span_a[i];
-          const float in2 = span_b[i];
-          const float out = math_function(in1, in2);
-          span_result[i] = out;
+        for (const int i : IndexRange(span_result.size())) {
+          span_result[i] = math_function(span_a[i], span_b[i]);
         }
       });
+  BLI_assert(success);
+  UNUSED_VARS_NDEBUG(success);
+}
 
-  result.apply_span();
-
-  /* The operation is not supported by this node currently. */
+static void do_math_operation(Span<float> span_input,
+                              MutableSpan<float> span_result,
+                              const NodeMathOperation operation)
+{
+  bool success = try_dispatch_float_math_fl_to_fl(
+      operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
+        for (const int i : IndexRange(span_result.size())) {
+          span_result[i] = math_function(span_input[i]);
+        }
+      });
   BLI_assert(success);
   UNUSED_VARS_NDEBUG(success);
 }
@@ -98,7 +187,7 @@ static void attribute_math_calc(GeometryComponent &component, const GeoNodeExecP
 {
   const bNode &node = params.node();
   const NodeAttributeMath *node_storage = (const NodeAttributeMath *)node.storage;
-  const int operation = node_storage->operation;
+  const NodeMathOperation operation = static_cast<NodeMathOperation>(node_storage->operation);
 
   /* The result type of this node is always float. */
   const CustomDataType result_type = CD_PROP_FLOAT;
@@ -115,15 +204,44 @@ static void attribute_math_calc(GeometryComponent &component, const GeoNodeExecP
 
   ReadAttributePtr attribute_a = params.get_input_attribute(
       "A", component, result_domain, result_type, nullptr);
-  ReadAttributePtr attribute_b = params.get_input_attribute(
-      "B", component, result_domain, result_type, nullptr);
-  if (!attribute_a || !attribute_b) {
-    /* Attribute wasn't found. */
+  if (!attribute_a) {
     return;
   }
 
-  do_math_operation(*attribute_a, *attribute_b, *attribute_result, operation);
-  attribute_result.save();
+  /* Note that passing the data with `get_span<float>()` works
+   * because the attributes were accessed with #CD_PROP_FLOAT. */
+  if (operation_use_input_b(operation)) {
+    ReadAttributePtr attribute_b = params.get_input_attribute(
+        "B", component, result_domain, result_type, nullptr);
+    if (!attribute_b) {
+      return;
+    }
+    if (operation_use_input_c(operation)) {
+      ReadAttributePtr attribute_c = params.get_input_attribute(
+          "C", component, result_domain, result_type, nullptr);
+      if (!attribute_c) {
+        return;
+      }
+      do_math_operation(attribute_a->get_span<float>(),
+                        attribute_b->get_span<float>(),
+                        attribute_c->get_span<float>(),
+                        attribute_result->get_span_for_write_only<float>(),
+                        operation);
+    }
+    else {
+      do_math_operation(attribute_a->get_span<float>(),
+                        attribute_b->get_span<float>(),
+                        attribute_result->get_span_for_write_only<float>(),
+                        operation);
+    }
+  }
+  else {
+    do_math_operation(attribute_a->get_span<float>(),
+                      attribute_result->get_span_for_write_only<float>(),
+                      operation);
+  }
+
+  attribute_result.apply_span_and_save();
 }
 
 static void geo_node_attribute_math_exec(GeoNodeExecParams params)