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

release/scripts/startup/nodeitems_builtins.py

@@ -488,6 +488,7 @@ geometry_node_categories = [
         NodeItem("GeometryNodeAttributeClamp"),
         NodeItem("GeometryNodeAttributeCompare"),
         NodeItem("GeometryNodeAttributeConvert"),
+        NodeItem("GeometryNodeAttributeCurveMap"),
         NodeItem("GeometryNodeAttributeFill"),
         NodeItem("GeometryNodeAttributeMix"),
         NodeItem("GeometryNodeAttributeProximity"),

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
 
 /** \} */
 

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();

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;

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[] = {

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

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);

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

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);
+}