Geometry Nodes: Add Attribute Curve Map Node

This node has the same functionality as the color and vector curve
mapping nodes in the shader editor. Here is works on every value for
the selected attribute, and it can also output a float value. Other
than that, the implementation is quite straightforward-- almost
completely boilerplate code.

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

4 files changed, 227 insertions, 0 deletions

diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index c116ceb8968..2bf419b5ea2 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -145,6 +145,7 @@ set(SRC
   geometry/nodes/node_geo_attribute_combine_xyz.cc
   geometry/nodes/node_geo_attribute_compare.cc
   geometry/nodes/node_geo_attribute_convert.cc
+  geometry/nodes/node_geo_attribute_curve_map.cc
   geometry/nodes/node_geo_attribute_fill.cc
   geometry/nodes/node_geo_attribute_map_range.cc
   geometry/nodes/node_geo_attribute_math.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 14c9e235bb6..7f77142b7e4 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -35,6 +35,7 @@ void register_node_type_geo_attribute_color_ramp(void);
 void register_node_type_geo_attribute_combine_xyz(void);
 void register_node_type_geo_attribute_compare(void);
 void register_node_type_geo_attribute_convert(void);
+void register_node_type_geo_attribute_curve_map(void);
 void register_node_type_geo_attribute_fill(void);
 void register_node_type_geo_attribute_map_range(void);
 void register_node_type_geo_attribute_math(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 2e4187229f4..f445e520c95 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -313,6 +313,7 @@ DefNode(GeometryNode, GEO_NODE_BOUNDING_BOX, 0, "BOUNDING_BOX", BoundBox, "Bound
 DefNode(GeometryNode, GEO_NODE_SWITCH, def_geo_switch, "SWITCH", Switch, "Switch", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_TRANSFER, def_geo_attribute_transfer, "ATTRIBUTE_TRANSFER", AttributeTransfer, "Attribute Transfer", "")
 DefNode(GeometryNode, GEO_NODE_CURVE_TO_MESH, 0, "CURVE_TO_MESH", CurveToMesh, "Curve to Mesh", "")
+DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_CURVE_MAP, def_geo_attribute_curve_map, "ATTRIBUTE_CURVE_MAP", AttributeCurveMap, "Attribute Curve Map", "")
 
 /* undefine macros */
 #undef DefNode
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_curve_map.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_curve_map.cc
new file mode 100644
index 00000000000..ba2e5b08352
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_curve_map.cc
@@ -0,0 +1,224 @@
+/*
+ * 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 "BLI_blenlib.h"
+
+#include "BKE_colortools.h"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "node_geometry_util.hh"
+
+static bNodeSocketTemplate geo_node_attribute_curve_map_in[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {SOCK_STRING, N_("Attribute")},
+    {SOCK_STRING, N_("Result")},
+    {-1, ""},
+};
+
+static bNodeSocketTemplate geo_node_attribute_curve_map_out[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {-1, ""},
+};
+
+static void geo_node_attribute_curve_map_layout(uiLayout *layout,
+                                                bContext *UNUSED(C),
+                                                PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
+  bNode *node = (bNode *)ptr->data;
+  NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)node->storage;
+  switch (data->data_type) {
+    case CD_PROP_FLOAT:
+      uiTemplateCurveMapping(layout, ptr, "curve_vec", 0, false, false, false, false);
+      break;
+    case CD_PROP_FLOAT3:
+      uiTemplateCurveMapping(layout, ptr, "curve_vec", 'v', false, false, false, false);
+      break;
+    case CD_PROP_COLOR:
+      uiTemplateCurveMapping(layout, ptr, "curve_rgb", 'c', false, false, false, false);
+      break;
+  }
+}
+
+static void geo_node_attribute_curve_map_free_storage(bNode *node)
+{
+  if (node->storage) {
+    NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)node->storage;
+    BKE_curvemapping_free(data->curve_vec);
+    BKE_curvemapping_free(data->curve_rgb);
+    MEM_freeN(node->storage);
+  }
+}
+
+static void geo_node_attribute_curve_map_copy_storage(bNodeTree *UNUSED(dest_ntree),
+                                                      bNode *dest_node,
+                                                      const bNode *src_node)
+{
+  dest_node->storage = MEM_dupallocN(src_node->storage);
+  NodeAttributeCurveMap *src_data = (NodeAttributeCurveMap *)src_node->storage;
+  NodeAttributeCurveMap *dest_data = (NodeAttributeCurveMap *)dest_node->storage;
+  dest_data->curve_vec = BKE_curvemapping_copy(src_data->curve_vec);
+  dest_data->curve_rgb = BKE_curvemapping_copy(src_data->curve_rgb);
+}
+
+static void geo_node_attribute_curve_map_init(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)MEM_callocN(sizeof(NodeAttributeCurveMap),
+                                                                     __func__);
+
+  data->data_type = CD_PROP_FLOAT;
+  data->curve_vec = BKE_curvemapping_add(4, -1.0f, -1.0f, 1.0f, 1.0f);
+  data->curve_vec->cur = 3;
+  data->curve_rgb = BKE_curvemapping_add(4, 0.0f, 0.0f, 1.0f, 1.0f);
+  node->storage = data;
+}
+
+static void geo_node_attribute_curve_map_update(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  /* Set the active curve when data type is changed. */
+  NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)node->storage;
+  if (data->data_type == CD_PROP_FLOAT) {
+    data->curve_vec->cur = 3;
+  }
+  else if (data->data_type == CD_PROP_FLOAT3) {
+    data->curve_vec->cur = 0;
+  }
+}
+
+namespace blender::nodes {
+
+static AttributeDomain get_result_domain(const GeometryComponent &component,
+                                         StringRef input_name,
+                                         StringRef result_name)
+{
+  /* Use the domain of the result attribute if it already exists. */
+  ReadAttributeLookup result_attribute = component.attribute_try_get_for_read(result_name);
+  if (result_attribute) {
+    return result_attribute.domain;
+  }
+
+  /* Otherwise use the input attribute's domain if it exists. */
+  ReadAttributeLookup input_attribute = component.attribute_try_get_for_read(input_name);
+  if (input_attribute) {
+    return input_attribute.domain;
+  }
+
+  return ATTR_DOMAIN_POINT;
+}
+
+static void execute_on_component(const GeoNodeExecParams &params, GeometryComponent &component)
+{
+  const bNode &bnode = params.node();
+  NodeAttributeCurveMap &node_storage = *(NodeAttributeCurveMap *)bnode.storage;
+  const std::string result_name = params.get_input<std::string>("Result");
+  const std::string input_name = params.get_input<std::string>("Attribute");
+
+  const CustomDataType result_type = (CustomDataType)node_storage.data_type;
+  const AttributeDomain result_domain = get_result_domain(component, input_name, result_name);
+
+  OutputAttribute attribute_result = component.attribute_try_get_for_output_only(
+      result_name, result_domain, result_type);
+  if (!attribute_result) {
+    return;
+  }
+
+  switch (result_type) {
+    case CD_PROP_FLOAT: {
+      const CurveMapping *cumap = (CurveMapping *)node_storage.curve_vec;
+      GVArray_Typed<float> attribute_in = component.attribute_get_for_read<float>(
+          input_name, result_domain, float(0.0f));
+      MutableSpan<float> results = attribute_result.as_span<float>();
+      for (const int i : IndexRange(attribute_in.size())) {
+        results[i] = BKE_curvemapping_evaluateF(cumap, 3, attribute_in[i]);
+      }
+      break;
+    }
+    case CD_PROP_FLOAT3: {
+      const CurveMapping *cumap = (CurveMapping *)node_storage.curve_vec;
+      GVArray_Typed<float3> attribute_in = component.attribute_get_for_read<float3>(
+          input_name, result_domain, float3(0.0f));
+      MutableSpan<float3> results = attribute_result.as_span<float3>();
+      for (const int i : IndexRange(attribute_in.size())) {
+        BKE_curvemapping_evaluate3F(cumap, results[i], attribute_in[i]);
+      }
+      break;
+    }
+    case CD_PROP_COLOR: {
+      const CurveMapping *cumap = (CurveMapping *)node_storage.curve_rgb;
+      GVArray_Typed<Color4f> attribute_in = component.attribute_get_for_read<Color4f>(
+          input_name, result_domain, Color4f(0.0f, 0.0f, 0.0f, 1.0f));
+      MutableSpan<Color4f> results = attribute_result.as_span<Color4f>();
+      for (const int i : IndexRange(attribute_in.size())) {
+        BKE_curvemapping_evaluateRGBF(cumap, results[i], attribute_in[i]);
+      }
+      break;
+    }
+    default: {
+      BLI_assert_unreachable();
+      break;
+    }
+  }
+
+  attribute_result.save();
+}
+
+static void geo_node_attribute_curve_map_exec(GeoNodeExecParams params)
+{
+  const bNode &bnode = params.node();
+  NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)bnode.storage;
+  BKE_curvemapping_init(data->curve_vec);
+  BKE_curvemapping_init(data->curve_rgb);
+
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+
+  geometry_set = geometry_set_realize_instances(geometry_set);
+
+  if (geometry_set.has<MeshComponent>()) {
+    execute_on_component(params, geometry_set.get_component_for_write<MeshComponent>());
+  }
+  if (geometry_set.has<PointCloudComponent>()) {
+    execute_on_component(params, geometry_set.get_component_for_write<PointCloudComponent>());
+  }
+  if (geometry_set.has<CurveComponent>()) {
+    execute_on_component(params, geometry_set.get_component_for_write<CurveComponent>());
+  }
+
+  params.set_output("Geometry", std::move(geometry_set));
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_attribute_curve_map()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_ATTRIBUTE_CURVE_MAP, "Attribute Curve Map", NODE_CLASS_ATTRIBUTE, 0);
+  node_type_socket_templates(
+      &ntype, geo_node_attribute_curve_map_in, geo_node_attribute_curve_map_out);
+  node_type_update(&ntype, geo_node_attribute_curve_map_update);
+  node_type_init(&ntype, geo_node_attribute_curve_map_init);
+  node_type_size_preset(&ntype, NODE_SIZE_LARGE);
+  node_type_storage(&ntype,
+                    "NodeAttributeCurveMap",
+                    geo_node_attribute_curve_map_free_storage,
+                    geo_node_attribute_curve_map_copy_storage);
+  ntype.geometry_node_execute = blender::nodes::geo_node_attribute_curve_map_exec;
+  ntype.draw_buttons = geo_node_attribute_curve_map_layout;
+  nodeRegisterType(&ntype);
+}