Geometry Nodes: Field on Domain Node

As described in T98943, this commit adds a node that can
evaluate a field on a separate domain in a larger field context.
This is potentially useful in many cases, to avoid relying on
a separate capture attribute node, which can make it easier
to build reusable fields that don't need geometry inputs.

Internally, the node just evaluates the input field in the larger
field context and then uses the generic domain interpolation,
so the code is simple. One future optimization might be using
the input selection to only evaluate part of the input field, but
then the selection has to be interpolated as well, and that might
not always be worth it.

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

8 files changed, 171 insertions, 0 deletions

diff --git a/release/scripts/startup/nodeitems_builtins.py b/release/scripts/startup/nodeitems_builtins.py
index 7ab919839dc..cddd96ff3cf 100644
--- a/release/scripts/startup/nodeitems_builtins.py
+++ b/release/scripts/startup/nodeitems_builtins.py
@@ -695,6 +695,7 @@ geometry_node_categories = [
     GeometryNodeCategory("GEO_UTILITIES", "Utilities", items=[
         NodeItem("GeometryNodeAccumulateField"),
         NodeItem("GeometryNodeFieldAtIndex"),
+        NodeItem("GeometryNodeFieldOnDomain"),
         NodeItem("ShaderNodeMapRange"),
         NodeItem("ShaderNodeFloatCurve"),
         NodeItem("ShaderNodeClamp"),
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index f531d693b6c..4a84f8fc3c4 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1497,6 +1497,7 @@ struct TexResult;
 #define GEO_NODE_INPUT_INSTANCE_SCALE 1160
 #define GEO_NODE_VOLUME_CUBE 1161
 #define GEO_NODE_POINTS 1162
+#define GEO_NODE_FIELD_ON_DOMAIN 1163
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 2e7755bca0e..b6ab1e0174b 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -4756,6 +4756,7 @@ static void registerGeometryNodes()
   register_node_type_geo_edge_split();
   register_node_type_geo_extrude_mesh();
   register_node_type_geo_field_at_index();
+  register_node_type_geo_field_on_domain();
   register_node_type_geo_flip_faces();
   register_node_type_geo_geometry_to_instance();
   register_node_type_geo_image_texture();
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 65a8d8bf24a..1a8f47605a3 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -10702,6 +10702,25 @@ static void def_geo_field_at_index(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
 }
 
+static void def_geo_field_on_domain(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_GeometryNodeAttributeType_type_with_socket_itemf");
+  RNA_def_property_ui_text(prop, "Data Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNode_socket_update");
+}
+
 static void def_geo_scale_elements(StructRNA *srna)
 {
   PropertyRNA *prop;
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 391322a95c6..56d1472e840 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -54,6 +54,7 @@ void register_node_type_geo_dual_mesh(void);
 void register_node_type_geo_edge_split(void);
 void register_node_type_geo_extrude_mesh(void);
 void register_node_type_geo_field_at_index(void);
+void register_node_type_geo_field_on_domain(void);
 void register_node_type_geo_flip_faces(void);
 void register_node_type_geo_geometry_to_instance(void);
 void register_node_type_geo_image_texture(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 090594d2308..20b3a61fd0b 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -308,6 +308,7 @@ DefNode(GeometryNode, GEO_NODE_ACCUMULATE_FIELD, def_geo_accumulate_field, "ACCU
 DefNode(GeometryNode, GEO_NODE_DUAL_MESH, 0, "DUAL_MESH", DualMesh, "Dual Mesh", "")
 DefNode(GeometryNode, GEO_NODE_EXTRUDE_MESH, def_geo_extrude_mesh, "EXTRUDE_MESH", ExtrudeMesh, "Extrude Mesh", "")
 DefNode(GeometryNode, GEO_NODE_FIELD_AT_INDEX, def_geo_field_at_index, "FIELD_AT_INDEX", FieldAtIndex, "Field at Index", "")
+DefNode(GeometryNode, GEO_NODE_FIELD_ON_DOMAIN, def_geo_field_on_domain, "FIELD_ON_DOMAIN", FieldOnDomain, "Field on Domain", "")
 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_FLIP_FACES, 0, "FLIP_FACES", FlipFaces, "Flip Faces", "")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index 8334a673539..1fad6ce5b34 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/source/blender/nodes/geometry/CMakeLists.txt
@@ -64,6 +64,7 @@ set(SRC
   nodes/node_geo_edge_split.cc
   nodes/node_geo_extrude_mesh.cc
   nodes/node_geo_field_at_index.cc
+  nodes/node_geo_field_on_domain.cc
   nodes/node_geo_flip_faces.cc
   nodes/node_geo_geometry_to_instance.cc
   nodes/node_geo_image_texture.cc
diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc b/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc
new file mode 100644
index 00000000000..e6906f4fb21
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc
@@ -0,0 +1,146 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#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_on_domain_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  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 eCustomDataType data_type = static_cast<eCustomDataType>(node->custom2);
+
+  bNodeSocket *sock_in_float = static_cast<bNodeSocket *>(node->inputs.first);
+  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 FieldOnDomain final : public GeometryFieldInput {
+ private:
+  GField src_field_;
+  eAttrDomain src_domain_;
+
+ public:
+  FieldOnDomain(GField field, eAttrDomain domain)
+      : GeometryFieldInput(field.cpp_type(), "Field on Domain"),
+        src_field_(std::move(field)),
+        src_domain_(domain)
+  {
+  }
+
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const eAttrDomain domain,
+                                 IndexMask /* mask */) const final
+  {
+    const GeometryComponentFieldContext context{component, src_domain_};
+    FieldEvaluator value_evaluator{context, component.attribute_domain_num(src_domain_)};
+    value_evaluator.add(src_field_);
+    value_evaluator.evaluate();
+    const GVArray &values = value_evaluator.get_evaluated(0);
+
+    return component.attribute_try_adapt_domain(values, src_domain_, domain);
+  }
+};
+
+static StringRefNull identifier_suffix(eCustomDataType 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 eAttrDomain domain = static_cast<eAttrDomain>(node.custom1);
+  const eCustomDataType data_type = static_cast<eCustomDataType>(node.custom2);
+
+  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> src_field = params.extract_input<Field<T>>(identifier);
+    Field<T> dst_field{std::make_shared<FieldOnDomain>(std::move(src_field), domain)};
+    params.set_output(identifier, std::move(dst_field));
+  });
+}
+
+}  // namespace blender::nodes::node_geo_field_on_domain_cc
+
+void register_node_type_geo_field_on_domain()
+{
+  namespace file_ns = blender::nodes::node_geo_field_on_domain_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_FIELD_ON_DOMAIN, "Field on Domain", 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);
+}