Cleanup: Use abstraction for attribute math node input

Since creating the attribute node, a helper function has been added to
automatically get the input attribute or a constant value, depending on
the "input type" values for the node. This commit replaces the specific
implementation of that behavior with the new helper function.

The versioning is necessary since the node now has a "storage" struct.

1 files changed, 23 insertions, 34 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 455b2a79394..a7b71fa614f 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc
@@ -30,9 +30,9 @@
 
 static bNodeSocketTemplate geo_node_attribute_math_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
-    {SOCK_STRING, N_("Attribute A")},
+    {SOCK_STRING, N_("A")},
     {SOCK_FLOAT, N_("A"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
-    {SOCK_STRING, N_("Attribute B")},
+    {SOCK_STRING, N_("B")},
     {SOCK_FLOAT, N_("B"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
     {SOCK_STRING, N_("Result")},
     {-1, ""},
@@ -45,27 +45,27 @@ static bNodeSocketTemplate geo_node_attribute_math_out[] = {
 
 static void geo_node_attribute_math_init(bNodeTree *UNUSED(tree), bNode *node)
 {
-  node->custom1 = NODE_MATH_ADD;
-  node->custom2 = GEO_NODE_USE_ATTRIBUTE_A | GEO_NODE_USE_ATTRIBUTE_B;
+  NodeAttributeMath *data = (NodeAttributeMath *)MEM_callocN(sizeof(NodeAttributeMath),
+                                                             "NodeAttributeMath");
+
+  data->operation = NODE_MATH_ADD;
+  data->input_type_a = GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE;
+  data->input_type_b = GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE;
+  node->storage = data;
 }
 
+namespace blender::nodes {
+
 static void geo_node_attribute_math_update(bNodeTree *UNUSED(ntree), bNode *node)
 {
-  bNodeSocket *sock_attribute_a = (bNodeSocket *)BLI_findlink(&node->inputs, 1);
-  bNodeSocket *sock_float_a = sock_attribute_a->next;
-  bNodeSocket *sock_attribute_b = sock_float_a->next;
-  bNodeSocket *sock_float_b = sock_attribute_b->next;
+  NodeAttributeMath *node_storage = (NodeAttributeMath *)node->storage;
 
-  GeometryNodeUseAttributeFlag flag = static_cast<GeometryNodeUseAttributeFlag>(node->custom2);
-
-  nodeSetSocketAvailability(sock_attribute_a, flag & GEO_NODE_USE_ATTRIBUTE_A);
-  nodeSetSocketAvailability(sock_attribute_b, flag & GEO_NODE_USE_ATTRIBUTE_B);
-  nodeSetSocketAvailability(sock_float_a, !(flag & GEO_NODE_USE_ATTRIBUTE_A));
-  nodeSetSocketAvailability(sock_float_b, !(flag & GEO_NODE_USE_ATTRIBUTE_B));
+  update_attribute_input_socket_availabilities(
+      *node, "A", (GeometryNodeAttributeInputMode)node_storage->input_type_a);
+  update_attribute_input_socket_availabilities(
+      *node, "B", (GeometryNodeAttributeInputMode)node_storage->input_type_b);
 }
 
-namespace blender::nodes {
-
 static void do_math_operation(const FloatReadAttribute &input_a,
                               const FloatReadAttribute &input_b,
                               FloatWriteAttribute result,
@@ -112,23 +112,10 @@ static void attribute_math_calc(GeometryComponent &component, const GeoNodeExecP
     return;
   }
 
-  GeometryNodeUseAttributeFlag flag = static_cast<GeometryNodeUseAttributeFlag>(node.custom2);
-
-  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);
-  };
-
-  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");
-
+  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. */
     return;
@@ -161,7 +148,9 @@ void register_node_type_geo_attribute_math()
   geo_node_type_base(&ntype, GEO_NODE_ATTRIBUTE_MATH, "Attribute Math", NODE_CLASS_ATTRIBUTE, 0);
   node_type_socket_templates(&ntype, geo_node_attribute_math_in, geo_node_attribute_math_out);
   ntype.geometry_node_execute = blender::nodes::geo_node_attribute_math_exec;
-  node_type_update(&ntype, geo_node_attribute_math_update);
+  node_type_update(&ntype, blender::nodes::geo_node_attribute_math_update);
   node_type_init(&ntype, geo_node_attribute_math_init);
+  node_type_storage(
+      &ntype, "NodeAttributeCompare", node_free_standard_storage, node_copy_standard_storage);
   nodeRegisterType(&ntype);
 }