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

7 files changed, 216 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index a76fe2f83e0..8f408ecd675 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/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
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 8c6bdb13708..b3bf80cd448 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/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();
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index ecbeadf1fa4..775c7a68f65 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/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)
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 96a5e1d87a6..3b9dd5a1f7c 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/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);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index e3566f81e2d..7434e0071ee 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/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", "")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index d692bd050c8..b9e762a341a 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/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
diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc b/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc
new file mode 100644
index 00000000000..9512323834c
--- /dev/null
+++ b/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);
+}