Geometry Nodes: new Field at Index node

This node allows accessing data of other elements in the context geometry.
It is similar to the Transfer Attribute node in Index mode. The main difference
is that this node does not require a geometry input, because the context
is used.

The node can e.g. be used to generalize what the Edge Vertices node is doing.
Instead of only being able to get the position of the vertices of an edge,
any field/attribute can be accessed on the vertices.

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

release/scripts/startup/nodeitems_builtins.py

@@ -757,6 +757,7 @@ geometry_node_categories = [
     ]),
     GeometryNodeCategory("GEO_UTILITIES", "Utilities", items=[
         NodeItem("GeometryNodeAccumulateField"),
+        NodeItem("GeometryNodeFieldAtIndex"),
         NodeItem("ShaderNodeMapRange"),
         NodeItem("ShaderNodeFloatCurve"),
         NodeItem("ShaderNodeClamp"),

source/blender/blenkernel/BKE_node.h

@@ -1628,6 +1628,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_INPUT_SCENE_TIME 1145
 #define GEO_NODE_ACCUMULATE_FIELD 1146
 #define GEO_NODE_INPUT_MESH_EDGE_ANGLE 1147
+#define GEO_NODE_FIELD_AT_INDEX 1148
 
 /** \} */
 

source/blender/blenkernel/intern/node.cc

@@ -4767,6 +4767,7 @@ static void registerGeometryNodes()
   register_node_type_geo_distribute_points_on_faces();
   register_node_type_geo_dual_mesh();
   register_node_type_geo_edge_split();
+  register_node_type_geo_field_at_index();
   register_node_type_geo_geometry_to_instance();
   register_node_type_geo_image_texture();
   register_node_type_geo_input_curve_handles();

source/blender/makesrna/intern/rna_nodetree.c

@@ -11377,6 +11377,24 @@ static void def_geo_realize_instances(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
 }
 
+static void def_geo_field_at_index(StructRNA *srna)
+{
+  PropertyRNA *prop;
+
+  prop = RNA_def_property(srna, "domain", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom1");
+  RNA_def_property_enum_items(prop, rna_enum_attribute_domain_items);
+  RNA_def_property_ui_text(prop, "Domain", "Domain the field is evaluated in");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
+
+  prop = RNA_def_property(srna, "data_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom2");
+  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_ui_text(prop, "Data Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
+}
+
 /* -------------------------------------------------------------------------- */
 
 static void rna_def_shader_node(BlenderRNA *brna)

source/blender/nodes/NOD_geometry.h

@@ -98,6 +98,7 @@ void register_node_type_geo_delete_geometry(void);
 void register_node_type_geo_distribute_points_on_faces(void);
 void register_node_type_geo_dual_mesh(void);
 void register_node_type_geo_edge_split(void);
+void register_node_type_geo_field_at_index(void);
 void register_node_type_geo_geometry_to_instance(void);
 void register_node_type_geo_image_texture(void);
 void register_node_type_geo_input_curve_handles(void);

source/blender/nodes/NOD_static_types.h

@@ -351,6 +351,7 @@ DefNode(GeometryNode, GEO_NODE_DELETE_GEOMETRY, def_geo_delete_geometry, "DELETE
 DefNode(GeometryNode, GEO_NODE_DISTRIBUTE_POINTS_ON_FACES, def_geo_distribute_points_on_faces, "DISTRIBUTE_POINTS_ON_FACES", DistributePointsOnFaces, "Distribute Points on Faces", "")
 DefNode(GeometryNode, GEO_NODE_ACCUMULATE_FIELD, def_geo_accumulate_field, "ACCUMULATE_FIELD", AccumulateField, "Accumulate Field", "")
 DefNode(GeometryNode, GEO_NODE_DUAL_MESH, 0, "DUAL_MESH", DualMesh, "Dual Mesh", "")
+DefNode(GeometryNode, GEO_NODE_FIELD_AT_INDEX, def_geo_field_at_index, "FIELD_AT_INDEX", FieldAtIndex, "Field at Index", "")
 DefNode(GeometryNode, GEO_NODE_FILL_CURVE, def_geo_curve_fill, "FILL_CURVE", FillCurve, "Fill Curve", "")
 DefNode(GeometryNode, GEO_NODE_FILLET_CURVE, def_geo_curve_fillet, "FILLET_CURVE", FilletCurve, "Fillet Curve", "")
 DefNode(GeometryNode, GEO_NODE_GEOMETRY_TO_INSTANCE, 0, "GEOMETRY_TO_INSTANCE", GeometryToInstance, "Geometry to Instance", "")

source/blender/nodes/geometry/CMakeLists.txt

@@ -116,6 +116,7 @@ set(SRC
   nodes/node_geo_distribute_points_on_faces.cc
   nodes/node_geo_dual_mesh.cc
   nodes/node_geo_edge_split.cc
+  nodes/node_geo_field_at_index.cc
   nodes/node_geo_geometry_to_instance.cc
   nodes/node_geo_image_texture.cc
   nodes/node_geo_input_curve_handles.cc

source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc

@@ -0,0 +1,193 @@
+/*
+ * 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 "node_geometry_util.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "BKE_attribute_math.hh"
+
+#include "BLI_task.hh"
+
+namespace blender::nodes::node_geo_field_at_index_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Int>(N_("Index")).min(0).supports_field();
+
+  b.add_input<decl::Float>(N_("Value"), "Value_Float").supports_field();
+  b.add_input<decl::Int>(N_("Value"), "Value_Int").supports_field();
+  b.add_input<decl::Vector>(N_("Value"), "Value_Vector").supports_field();
+  b.add_input<decl::Color>(N_("Value"), "Value_Color").supports_field();
+  b.add_input<decl::Bool>(N_("Value"), "Value_Bool").supports_field();
+
+  b.add_output<decl::Float>(N_("Value"), "Value_Float").field_source();
+  b.add_output<decl::Int>(N_("Value"), "Value_Int").field_source();
+  b.add_output<decl::Vector>(N_("Value"), "Value_Vector").field_source();
+  b.add_output<decl::Color>(N_("Value"), "Value_Color").field_source();
+  b.add_output<decl::Bool>(N_("Value"), "Value_Bool").field_source();
+}
+
+static void node_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
+  uiItemR(layout, ptr, "domain", 0, "", ICON_NONE);
+}
+
+static void node_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  node->custom1 = ATTR_DOMAIN_POINT;
+  node->custom2 = CD_PROP_FLOAT;
+}
+
+static void node_update(bNodeTree *ntree, bNode *node)
+{
+  const CustomDataType data_type = static_cast<CustomDataType>(node->custom2);
+
+  bNodeSocket *sock_index = static_cast<bNodeSocket *>(node->inputs.first);
+  bNodeSocket *sock_in_float = sock_index->next;
+  bNodeSocket *sock_in_int = sock_in_float->next;
+  bNodeSocket *sock_in_vector = sock_in_int->next;
+  bNodeSocket *sock_in_color = sock_in_vector->next;
+  bNodeSocket *sock_in_bool = sock_in_color->next;
+
+  bNodeSocket *sock_out_float = static_cast<bNodeSocket *>(node->outputs.first);
+  bNodeSocket *sock_out_int = sock_out_float->next;
+  bNodeSocket *sock_out_vector = sock_out_int->next;
+  bNodeSocket *sock_out_color = sock_out_vector->next;
+  bNodeSocket *sock_out_bool = sock_out_color->next;
+
+  nodeSetSocketAvailability(ntree, sock_in_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_in_int, data_type == CD_PROP_INT32);
+  nodeSetSocketAvailability(ntree, sock_in_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_in_color, data_type == CD_PROP_COLOR);
+  nodeSetSocketAvailability(ntree, sock_in_bool, data_type == CD_PROP_BOOL);
+
+  nodeSetSocketAvailability(ntree, sock_out_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, sock_out_int, data_type == CD_PROP_INT32);
+  nodeSetSocketAvailability(ntree, sock_out_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, sock_out_color, data_type == CD_PROP_COLOR);
+  nodeSetSocketAvailability(ntree, sock_out_bool, data_type == CD_PROP_BOOL);
+}
+
+class FieldAtIndex final : public GeometryFieldInput {
+ private:
+  Field<int> index_field_;
+  GField value_field_;
+  AttributeDomain value_field_domain_;
+
+ public:
+  FieldAtIndex(Field<int> index_field, GField value_field, AttributeDomain value_field_domain)
+      : GeometryFieldInput(value_field.cpp_type(), "Field at Index"),
+        index_field_(std::move(index_field)),
+        value_field_(std::move(value_field)),
+        value_field_domain_(value_field_domain)
+  {
+  }
+
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask mask) const final
+  {
+    const GeometryComponentFieldContext value_field_context{component, value_field_domain_};
+    FieldEvaluator value_evaluator{value_field_context,
+                                   component.attribute_domain_size(value_field_domain_)};
+    value_evaluator.add(value_field_);
+    value_evaluator.evaluate();
+    const GVArray &values = value_evaluator.get_evaluated(0);
+
+    const GeometryComponentFieldContext index_field_context{component, domain};
+    FieldEvaluator index_evaluator{index_field_context, &mask};
+    index_evaluator.add(index_field_);
+    index_evaluator.evaluate();
+    const VArray<int> &indices = index_evaluator.get_evaluated<int>(0);
+
+    GVArray output_array;
+    attribute_math::convert_to_static_type(*type_, [&](auto dummy) {
+      using T = decltype(dummy);
+      Array<T> dst_array(mask.min_array_size());
+      VArray<T> src_values = values.typed<T>();
+      threading::parallel_for(mask.index_range(), 1024, [&](const IndexRange range) {
+        for (const int i : mask.slice(range)) {
+          const int index = indices[i];
+          if (index >= 0 && index < src_values.size()) {
+            dst_array[i] = src_values[index];
+          }
+          else {
+            dst_array[i] = {};
+          }
+        }
+      });
+      output_array = VArray<T>::ForContainer(std::move(dst_array));
+    });
+
+    return output_array;
+  }
+};
+
+static StringRefNull identifier_suffix(CustomDataType data_type)
+{
+  switch (data_type) {
+    case CD_PROP_BOOL:
+      return "Bool";
+    case CD_PROP_FLOAT:
+      return "Float";
+    case CD_PROP_INT32:
+      return "Int";
+    case CD_PROP_COLOR:
+      return "Color";
+    case CD_PROP_FLOAT3:
+      return "Vector";
+    default:
+      BLI_assert_unreachable();
+      return "";
+  }
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  const bNode &node = params.node();
+  const AttributeDomain domain = static_cast<AttributeDomain>(node.custom1);
+  const CustomDataType data_type = static_cast<CustomDataType>(node.custom2);
+
+  Field<int> index_field = params.extract_input<Field<int>>("Index");
+  attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
+    using T = decltype(dummy);
+    static const std::string identifier = "Value_" + identifier_suffix(data_type);
+    Field<T> value_field = params.extract_input<Field<T>>(identifier);
+    Field<T> output_field{
+        std::make_shared<FieldAtIndex>(std::move(index_field), std::move(value_field), domain)};
+    params.set_output(identifier, std::move(output_field));
+  });
+}
+
+}  // namespace blender::nodes::node_geo_field_at_index_cc
+
+void register_node_type_geo_field_at_index()
+{
+  namespace file_ns = blender::nodes::node_geo_field_at_index_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_FIELD_AT_INDEX, "Field at Index", NODE_CLASS_CONVERTER);
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  ntype.declare = file_ns::node_declare;
+  ntype.draw_buttons = file_ns::node_layout;
+  ntype.initfunc = file_ns::node_init;
+  ntype.updatefunc = file_ns::node_update;
+  nodeRegisterType(&ntype);
+}