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

8 files changed, 290 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index bcab456b36e..7978a0114ef 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1419,6 +1419,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_SWITCH 1043
 #define GEO_NODE_ATTRIBUTE_TRANSFER 1044
 #define GEO_NODE_CURVE_TO_MESH 1045
+#define GEO_NODE_ATTRIBUTE_CURVE_MAP 1046
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 6634c42a274..05eb7e43441 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -501,6 +501,12 @@ void ntreeBlendWrite(BlendWriter *writer, bNodeTree *ntree)
           ELEM(node->type, SH_NODE_CURVE_VEC, SH_NODE_CURVE_RGB)) {
         BKE_curvemapping_blend_write(writer, (const CurveMapping *)node->storage);
       }
+      else if ((ntree->type == NTREE_GEOMETRY) && (node->type == GEO_NODE_ATTRIBUTE_CURVE_MAP)) {
+        BLO_write_struct_by_name(writer, node->typeinfo->storagename, node->storage);
+        NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)node->storage;
+        BKE_curvemapping_blend_write(writer, (const CurveMapping *)data->curve_vec);
+        BKE_curvemapping_blend_write(writer, (const CurveMapping *)data->curve_rgb);
+      }
       else if (ntree->type == NTREE_SHADER && (node->type == SH_NODE_SCRIPT)) {
         NodeShaderScript *nss = (NodeShaderScript *)node->storage;
         if (nss->bytecode) {
@@ -676,6 +682,18 @@ void ntreeBlendReadData(BlendDataReader *reader, bNodeTree *ntree)
           BKE_curvemapping_blend_read(reader, (CurveMapping *)node->storage);
           break;
         }
+        case GEO_NODE_ATTRIBUTE_CURVE_MAP: {
+          NodeAttributeCurveMap *data = (NodeAttributeCurveMap *)node->storage;
+          BLO_read_data_address(reader, &data->curve_vec);
+          if (data->curve_vec) {
+            BKE_curvemapping_blend_read(reader, data->curve_vec);
+          }
+          BLO_read_data_address(reader, &data->curve_rgb);
+          if (data->curve_rgb) {
+            BKE_curvemapping_blend_read(reader, data->curve_rgb);
+          }
+          break;
+        }
         case SH_NODE_SCRIPT: {
           NodeShaderScript *nss = (NodeShaderScript *)node->storage;
           BLO_read_data_address(reader, &nss->bytecode);
@@ -4934,6 +4952,7 @@ static void registerGeometryNodes()
   register_node_type_geo_attribute_combine_xyz();
   register_node_type_geo_attribute_compare();
   register_node_type_geo_attribute_convert();
+  register_node_type_geo_attribute_curve_map();
   register_node_type_geo_attribute_fill();
   register_node_type_geo_attribute_map_range();
   register_node_type_geo_attribute_math();
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index 47f2c221279..56f8c9211c1 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -1188,6 +1188,14 @@ typedef struct NodeAttributeColorRamp {
   ColorBand color_ramp;
 } NodeAttributeColorRamp;
 
+typedef struct NodeAttributeCurveMap {
+  /* CustomDataType. */
+  uint8_t data_type;
+  char _pad[7];
+  CurveMapping *curve_vec;
+  CurveMapping *curve_rgb;
+} NodeAttributeCurveMap;
+
 typedef struct NodeInputVector {
   float vector[3];
 } NodeInputVector;
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 128c65aecd9..293c43a93e7 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -2182,6 +2182,17 @@ static const EnumPropertyItem *rna_GeometryNodeAttributeMapRange_type_itemf(
   return itemf_function_check(rna_enum_attribute_type_items, attribute_map_range_type_supported);
 }
 
+static bool attribute_curve_map_type_supported(const EnumPropertyItem *item)
+{
+  return ELEM(item->value, CD_PROP_FLOAT, CD_PROP_FLOAT3, CD_PROP_COLOR);
+}
+static const EnumPropertyItem *rna_GeometryNodeAttributeCurveMap_type_itemf(
+    bContext *UNUSED(C), PointerRNA *UNUSED(ptr), PropertyRNA *UNUSED(prop), bool *r_free)
+{
+  *r_free = true;
+  return itemf_function_check(rna_enum_attribute_type_items, attribute_curve_map_type_supported);
+}
+
 static StructRNA *rna_ShaderNode_register(Main *bmain,
                                           ReportList *reports,
                                           void *data,
@@ -9262,6 +9273,30 @@ static void def_geo_attribute_color_ramp(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
 }
 
+static void def_geo_attribute_curve_map(StructRNA *srna)
+{
+  PropertyRNA *prop;
+
+  RNA_def_struct_sdna_from(srna, "NodeAttributeCurveMap", "storage");
+
+  prop = RNA_def_property(srna, "data_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "data_type");
+  RNA_def_property_enum_items(prop, rna_enum_attribute_type_items);
+  RNA_def_property_enum_funcs(prop, NULL, NULL, "rna_GeometryNodeAttributeCurveMap_type_itemf");
+  RNA_def_property_ui_text(prop, "Data Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+
+  prop = RNA_def_property(srna, "curve_vec", PROP_POINTER, PROP_NONE);
+  RNA_def_property_struct_type(prop, "CurveMapping");
+  RNA_def_property_ui_text(prop, "Mapping", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
+
+  prop = RNA_def_property(srna, "curve_rgb", PROP_POINTER, PROP_NONE);
+  RNA_def_property_struct_type(prop, "CurveMapping");
+  RNA_def_property_ui_text(prop, "Mapping", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
+}
+
 static void def_geo_point_rotate(StructRNA *srna)
 {
   static const EnumPropertyItem type_items[] = {
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);
+}