add Store Local Attribute node

8 files changed, 233 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index d80feec94ef..8a9cb9bd6e0 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1491,6 +1491,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_CURVE_SELECT_HANDLES 1074
 #define GEO_NODE_GEOMETRY_EXPANDER 1075
 #define GEO_NODE_EXTRUDE 1076
+#define GEO_NODE_ATTRIBUTE_STORE_LOCAL 1077
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 93d40ec70aa..7098ecc812e 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -3935,6 +3935,9 @@ static std::string expander_output_to_name(const GeometryExpanderOutput &expande
 {
   switch (expander_output.type) {
     case GEOMETRY_EXPANDER_OUTPUT_TYPE_LOCAL: {
+      if (expander_output.local_socket_identifier[0] == '\0') {
+        return expander_output.local_node_name;
+      }
       return expander_output.local_node_name + StringRef(" ▶ ") +
              expander_output.local_socket_identifier;
     }
@@ -5177,6 +5180,7 @@ static void registerGeometryNodes()
   register_node_type_geo_attribute_randomize();
   register_node_type_geo_attribute_remove();
   register_node_type_geo_attribute_separate_xyz();
+  register_node_type_geo_attribute_store_local();
   register_node_type_geo_attribute_transfer();
   register_node_type_geo_attribute_vector_math();
   register_node_type_geo_attribute_vector_rotate();
diff --git a/source/blender/editors/space_node/node_edit.cc b/source/blender/editors/space_node/node_edit.cc
index c71805bddf9..8e8ec9fd51c 100644
--- a/source/blender/editors/space_node/node_edit.cc
+++ b/source/blender/editors/space_node/node_edit.cc
@@ -2990,6 +2990,25 @@ void NODE_OT_cryptomatte_layer_remove(wmOperatorType *ot)
 
 /* ****************** Geometry Expander Add Output  ******************* */
 
+static eNodeSocketDatatype custom_data_type_to_socket_type(const CustomDataType type)
+{
+  switch (type) {
+    case CD_PROP_FLOAT:
+      return SOCK_FLOAT;
+    case CD_PROP_FLOAT2:
+    case CD_PROP_FLOAT3:
+      return SOCK_VECTOR;
+    case CD_PROP_COLOR:
+      return SOCK_RGBA;
+    case CD_PROP_INT32:
+      return SOCK_INT;
+    case CD_PROP_BOOL:
+      return SOCK_BOOLEAN;
+    default:
+      return SOCK_FLOAT;
+  }
+}
+
 static void foreach_available_attribute(
     bNodeTree *ntree,
     bNode *UNUSED(current_node),
@@ -3018,6 +3037,15 @@ static void foreach_available_attribute(
         callback(attribute);
       }
     }
+    if (node->type == GEO_NODE_ATTRIBUTE_STORE_LOCAL) {
+      GeometryExpanderOutput attribute;
+      attribute.type = GEOMETRY_EXPANDER_OUTPUT_TYPE_LOCAL;
+      attribute.socket_type = custom_data_type_to_socket_type((CustomDataType)node->custom1);
+      attribute.array_source = GEOMETRY_EXPANDER_ARRAY_SOURCE_MESH_VERTICES;
+      STRNCPY(attribute.local_node_name, node->name);
+      STRNCPY(attribute.local_socket_identifier, "");
+      callback(attribute);
+    }
   }
   {
     GeometryExpanderOutput attribute;
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index ab9fa208229..359e88f182f 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -9187,6 +9187,26 @@ static void def_geo_attribute_fill(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
 }
 
+static void def_geo_attribute_store_local(StructRNA *srna)
+{
+  PropertyRNA *prop;
+
+  prop = RNA_def_property(srna, "data_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom1");
+  RNA_def_property_enum_items(prop, rna_enum_attribute_type_items);
+  RNA_def_property_enum_funcs(prop, NULL, NULL, "rna_GeometryNodeAttributeFill_type_itemf");
+  RNA_def_property_enum_default(prop, CD_PROP_FLOAT);
+  RNA_def_property_ui_text(prop, "Data Type", "Type of data stored in attribute");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
+
+  prop = RNA_def_property(srna, "domain", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom2");
+  RNA_def_property_enum_items(prop, rna_enum_attribute_domain_with_auto_items);
+  RNA_def_property_enum_default(prop, ATTR_DOMAIN_AUTO);
+  RNA_def_property_ui_text(prop, "Domain", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
+}
+
 static void def_geo_attribute_convert(StructRNA *srna)
 {
   PropertyRNA *prop;
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 7dd5c0f0ae1..35577f68c9c 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -156,6 +156,7 @@ set(SRC
   geometry/nodes/node_geo_attribute_remove.cc
   geometry/nodes/node_geo_attribute_sample_texture.cc
   geometry/nodes/node_geo_attribute_separate_xyz.cc
+  geometry/nodes/node_geo_attribute_store_local.cc
   geometry/nodes/node_geo_attribute_transfer.cc
   geometry/nodes/node_geo_attribute_vector_math.cc
   geometry/nodes/node_geo_attribute_vector_rotate.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index b2ea7e66246..c9bfe9758d2 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -44,6 +44,7 @@ void register_node_type_geo_attribute_proximity(void);
 void register_node_type_geo_attribute_randomize(void);
 void register_node_type_geo_attribute_remove(void);
 void register_node_type_geo_attribute_separate_xyz(void);
+void register_node_type_geo_attribute_store_local(void);
 void register_node_type_geo_attribute_transfer(void);
 void register_node_type_geo_attribute_vector_math(void);
 void register_node_type_geo_attribute_vector_rotate(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 2b80f915ab3..413aab4c9f7 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -284,6 +284,7 @@ DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_RANDOMIZE, def_geo_attribute_randomize,
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_REMOVE, 0, "ATTRIBUTE_REMOVE", AttributeRemove, "Attribute Remove", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_SAMPLE_TEXTURE, 0, "ATTRIBUTE_SAMPLE_TEXTURE", AttributeSampleTexture, "Attribute Sample Texture", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_SEPARATE_XYZ, def_geo_attribute_separate_xyz, "ATTRIBUTE_SEPARATE_XYZ", AttributeSeparateXYZ, "Attribute Separate XYZ", "")
+DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_STORE_LOCAL, def_geo_attribute_store_local, "ATTRIBUTE_STORE_LOCAL", AttributeStoreLocal, "Attribute Store Local", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_TRANSFER, def_geo_attribute_transfer, "ATTRIBUTE_TRANSFER", AttributeTransfer, "Attribute Transfer", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_VECTOR_MATH, def_geo_attribute_vector_math, "ATTRIBUTE_VECTOR_MATH", AttributeVectorMath, "Attribute Vector Math", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_VECTOR_ROTATE, def_geo_attribute_vector_rotate, "ATTRIBUTE_VECTOR_ROTATE", AttributeVectorRotate, "Attribute Vector Rotate", "")
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_store_local.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_store_local.cc
new file mode 100644
index 00000000000..58dd9db2f2c
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_store_local.cc
@@ -0,0 +1,177 @@
+/*
+ * 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 "UI_interface.h"
+#include "UI_resources.h"
+
+#include "node_geometry_util.hh"
+
+static bNodeSocketTemplate geo_node_attribute_store_local_in[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {SOCK_VECTOR, N_("Value"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_FLOAT, N_("Value"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_RGBA, N_("Value"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_BOOLEAN, N_("Value"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_INT, N_("Value"), 0, 0, 0, 0, -10000000.0f, 10000000.0f},
+    {-1, ""},
+};
+
+static bNodeSocketTemplate geo_node_attribute_store_local_out[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {-1, ""},
+};
+
+static void geo_node_attribute_store_local_layout(uiLayout *layout,
+                                                  bContext *UNUSED(C),
+                                                  PointerRNA *ptr)
+{
+  uiLayoutSetPropSep(layout, true);
+  uiLayoutSetPropDecorate(layout, false);
+  uiItemR(layout, ptr, "domain", 0, "", ICON_NONE);
+  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
+}
+
+static void geo_node_attribute_store_local_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  node->custom1 = CD_PROP_FLOAT;
+  node->custom2 = ATTR_DOMAIN_POINT;
+}
+
+static void geo_node_attribute_store_local_update(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  bNodeSocket *socket_value_vector = (bNodeSocket *)BLI_findlink(&node->inputs, 1);
+  bNodeSocket *socket_value_float = socket_value_vector->next;
+  bNodeSocket *socket_value_color4f = socket_value_float->next;
+  bNodeSocket *socket_value_boolean = socket_value_color4f->next;
+  bNodeSocket *socket_value_int32 = socket_value_boolean->next;
+
+  const CustomDataType data_type = static_cast<CustomDataType>(node->custom1);
+
+  nodeSetSocketAvailability(socket_value_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(socket_value_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(socket_value_color4f, data_type == CD_PROP_COLOR);
+  nodeSetSocketAvailability(socket_value_boolean, data_type == CD_PROP_BOOL);
+  nodeSetSocketAvailability(socket_value_int32, data_type == CD_PROP_INT32);
+}
+
+namespace blender::nodes {
+
+static AttributeDomain get_result_domain(const GeometryComponent &component, const StringRef name)
+{
+  /* Use the domain of the result attribute if it already exists. */
+  std::optional<AttributeMetaData> result_info = component.attribute_get_meta_data(name);
+  if (result_info) {
+    return result_info->domain;
+  }
+  return ATTR_DOMAIN_POINT;
+}
+
+static void store_local_attribute(GeometryComponent &component, const GeoNodeExecParams &params)
+{
+  const bNode &node = params.node();
+  const bNodeTree &tree = params.ntree();
+  const std::string attribute_name = get_local_attribute_name(tree.id.name, node.name, "");
+
+  const CustomDataType data_type = static_cast<CustomDataType>(node.custom1);
+  const AttributeDomain domain = static_cast<AttributeDomain>(node.custom2);
+  const AttributeDomain result_domain = (domain == ATTR_DOMAIN_AUTO) ?
+                                            get_result_domain(component, attribute_name) :
+                                            domain;
+
+  OutputAttribute attribute = component.attribute_try_get_for_output_only(
+      attribute_name, result_domain, data_type);
+  if (!attribute) {
+    return;
+  }
+
+  switch (data_type) {
+    case CD_PROP_FLOAT: {
+      const Array<float> values = params.get_input<Array<float>>("Value_001");
+      fn::GVArray_For_RepeatedGSpan values_repeated{attribute->size(), values.as_span()};
+      fn::GVArray_GSpan values_span{values_repeated};
+      attribute->set_all(values_span.data());
+      break;
+    }
+    case CD_PROP_FLOAT3: {
+      const Array<float3> values = params.get_input<Array<float3>>("Value");
+      fn::GVArray_For_RepeatedGSpan values_repeated{attribute->size(), values.as_span()};
+      fn::GVArray_GSpan values_span{values_repeated};
+      attribute->set_all(values_span.data());
+      break;
+    }
+    case CD_PROP_COLOR: {
+      const Array<ColorGeometry4f> values = params.get_input<Array<ColorGeometry4f>>("Value_002");
+      fn::GVArray_For_RepeatedGSpan values_repeated{attribute->size(), values.as_span()};
+      fn::GVArray_GSpan values_span{values_repeated};
+      attribute->set_all(values_span.data());
+      break;
+    }
+    case CD_PROP_BOOL: {
+      const Array<bool> values = params.get_input<Array<bool>>("Value_003");
+      fn::GVArray_For_RepeatedGSpan values_repeated{attribute->size(), values.as_span()};
+      fn::GVArray_GSpan values_span{values_repeated};
+      attribute->set_all(values_span.data());
+      break;
+    }
+    case CD_PROP_INT32: {
+      const Array<int> values = params.get_input<Array<int>>("Value_004");
+      fn::GVArray_For_RepeatedGSpan values_repeated{attribute->size(), values.as_span()};
+      fn::GVArray_GSpan values_span{values_repeated};
+      attribute->set_all(values_span.data());
+      break;
+    }
+    default:
+      break;
+  }
+
+  attribute.save();
+}
+
+static void geo_node_attribute_store_local_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+
+  geometry_set = geometry_set_realize_instances(geometry_set);
+
+  if (geometry_set.has<MeshComponent>()) {
+    store_local_attribute(geometry_set.get_component_for_write<MeshComponent>(), params);
+  }
+  if (geometry_set.has<PointCloudComponent>()) {
+    store_local_attribute(geometry_set.get_component_for_write<PointCloudComponent>(), params);
+  }
+  if (geometry_set.has<CurveComponent>()) {
+    store_local_attribute(geometry_set.get_component_for_write<CurveComponent>(), params);
+  }
+
+  params.set_output("Geometry", geometry_set);
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_attribute_store_local()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_ATTRIBUTE_STORE_LOCAL, "Store Local Attribute", NODE_CLASS_ATTRIBUTE, 0);
+  node_type_socket_templates(
+      &ntype, geo_node_attribute_store_local_in, geo_node_attribute_store_local_out);
+  node_type_init(&ntype, geo_node_attribute_store_local_init);
+  node_type_update(&ntype, geo_node_attribute_store_local_update);
+  ntype.geometry_node_execute = blender::nodes::geo_node_attribute_store_local_exec;
+  ntype.draw_buttons = geo_node_attribute_store_local_layout;
+  nodeRegisterType(&ntype);
+}