Geometry Nodes: Add fields version of material nodes

This commit moves the old material nodes to a "legacy" folder and adds
versions of the nodes that work with fields.

The "Select by Material" node is a field node now, so it doesn't have
a geometry output. This is an improvement because there are fewer links
to connect, and it's more easily usable in different situations.
It's also called "Material Selection", since it's more of an input
than an action now.

It's sometimes necessary to use the attribute capture node to get a
more predictable interpolation to mesh faces. This is because the
selection field input is always evaluated on the face domain, so
attribute inputs are interpolated before they are booleans, so they
cannot use the new interpolations from rB5841f8656d9580d7b9.

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

10 files changed, 265 insertions, 21 deletions

diff --git a/release/scripts/startup/nodeitems_builtins.py b/release/scripts/startup/nodeitems_builtins.py
index c0af91b7ad6..569c5291576 100644
--- a/release/scripts/startup/nodeitems_builtins.py
+++ b/release/scripts/startup/nodeitems_builtins.py
@@ -564,6 +564,8 @@ geometry_node_categories = [
         NodeItem("GeometryNodeLegacyMaterialAssign", poll=geometry_nodes_fields_legacy_poll),
         NodeItem("GeometryNodeLegacySelectByMaterial", poll=geometry_nodes_fields_legacy_poll),
 
+        NodeItem("GeometryNodeMaterialAssign", poll=geometry_nodes_fields_poll),
+        NodeItem("GeometryNodeMaterialSelection", poll=geometry_nodes_fields_poll),
         NodeItem("GeometryNodeMaterialReplace"),
     ]),
     GeometryNodeCategory("GEO_MESH", "Mesh", items=[
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index d4bc0245a61..a57281e4478 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1491,6 +1491,8 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_INPUT_INDEX 1078
 #define GEO_NODE_INPUT_NORMAL 1079
 #define GEO_NODE_ATTRIBUTE_CAPTURE 1080
+#define GEO_NODE_MATERIAL_SELECTION 1081
+#define GEO_NODE_MATERIAL_ASSIGN 1082
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 61b90b44b03..0d691854c1d 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -5153,6 +5153,9 @@ static void registerGeometryNodes()
 {
   register_node_type_geo_group();
 
+  register_node_type_geo_legacy_material_assign();
+  register_node_type_geo_legacy_select_by_material();
+
   register_node_type_geo_align_rotation_to_vector();
   register_node_type_geo_attribute_clamp();
   register_node_type_geo_attribute_color_ramp();
@@ -5225,7 +5228,7 @@ static void registerGeometryNodes()
   register_node_type_geo_raycast();
   register_node_type_geo_sample_texture();
   register_node_type_geo_select_by_handle_type();
-  register_node_type_geo_select_by_material();
+  register_node_type_geo_material_selection();
   register_node_type_geo_separate_components();
   register_node_type_geo_set_position();
   register_node_type_geo_subdivision_surface();
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 614be1d7229..b741461f820 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -141,6 +141,9 @@ set(SRC
   function/nodes/node_fn_random_float.cc
   function/node_function_util.cc
 
+  geometry/nodes/legacy/node_geo_material_assign.cc
+  geometry/nodes/legacy/node_geo_select_by_material.cc
+
   geometry/nodes/node_geo_align_rotation_to_vector.cc
   geometry/nodes/node_geo_attribute_capture.cc
   geometry/nodes/node_geo_attribute_clamp.cc
@@ -195,6 +198,7 @@ set(SRC
   geometry/nodes/node_geo_join_geometry.cc
   geometry/nodes/node_geo_material_assign.cc
   geometry/nodes/node_geo_material_replace.cc
+  geometry/nodes/node_geo_material_selection.cc
   geometry/nodes/node_geo_mesh_primitive_circle.cc
   geometry/nodes/node_geo_mesh_primitive_cone.cc
   geometry/nodes/node_geo_mesh_primitive_cube.cc
@@ -214,7 +218,6 @@ set(SRC
   geometry/nodes/node_geo_point_translate.cc
   geometry/nodes/node_geo_points_to_volume.cc
   geometry/nodes/node_geo_raycast.cc
-  geometry/nodes/node_geo_select_by_material.cc
   geometry/nodes/node_geo_separate_components.cc
   geometry/nodes/node_geo_set_position.cc
   geometry/nodes/node_geo_subdivision_surface.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index e9fb4ad123c..a713da45f0b 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -29,6 +29,9 @@ void register_node_tree_type_geo(void);
 void register_node_type_geo_group(void);
 void register_node_type_geo_custom_group(bNodeType *ntype);
 
+void register_node_type_geo_legacy_material_assign(void);
+void register_node_type_geo_legacy_select_by_material(void);
+
 void register_node_type_geo_align_rotation_to_vector(void);
 void register_node_type_geo_attribute_clamp(void);
 void register_node_type_geo_attribute_color_ramp(void);
@@ -80,6 +83,7 @@ void register_node_type_geo_is_viewport(void);
 void register_node_type_geo_join_geometry(void);
 void register_node_type_geo_material_assign(void);
 void register_node_type_geo_material_replace(void);
+void register_node_type_geo_material_selection(void);
 void register_node_type_geo_mesh_primitive_circle(void);
 void register_node_type_geo_mesh_primitive_cone(void);
 void register_node_type_geo_mesh_primitive_cube(void);
@@ -101,7 +105,6 @@ void register_node_type_geo_points_to_volume(void);
 void register_node_type_geo_raycast(void);
 void register_node_type_geo_sample_texture(void);
 void register_node_type_geo_select_by_handle_type(void);
-void register_node_type_geo_select_by_material(void);
 void register_node_type_geo_separate_components(void);
 void register_node_type_geo_set_position(void);
 void register_node_type_geo_subdivision_surface(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 215dca63fa3..51d59821d3c 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -332,7 +332,9 @@ DefNode(GeometryNode, GEO_NODE_INPUT_NORMAL, 0, "INPUT_NORMAL", InputNormal, "No
 DefNode(GeometryNode, GEO_NODE_INPUT_POSITION, 0, "POSITION", InputPosition, "Position", "")
 DefNode(GeometryNode, GEO_NODE_IS_VIEWPORT, 0, "IS_VIEWPORT", IsViewport, "Is Viewport", "")
 DefNode(GeometryNode, GEO_NODE_JOIN_GEOMETRY, 0, "JOIN_GEOMETRY", JoinGeometry, "Join Geometry", "")
+DefNode(GeometryNode, GEO_NODE_MATERIAL_ASSIGN, 0, "MATERIAL_ASSIGN", MaterialAssign, "Material Assign", "")
 DefNode(GeometryNode, GEO_NODE_MATERIAL_REPLACE, 0, "MATERIAL_REPLACE", MaterialReplace, "Material Replace", "")
+DefNode(GeometryNode, GEO_NODE_MATERIAL_SELECTION, 0, "MATERIAL_SELECTION", MaterialSelection, "Material Selection", "")
 DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_CIRCLE, def_geo_mesh_circle, "MESH_PRIMITIVE_CIRCLE", MeshCircle, "Mesh Circle", "")
 DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_CONE, def_geo_mesh_cone, "MESH_PRIMITIVE_CONE", MeshCone, "Cone", "")
 DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_CUBE, 0, "MESH_PRIMITIVE_CUBE", MeshCube, "Cube", "")
diff --git a/source/blender/nodes/geometry/nodes/legacy/node_geo_material_assign.cc b/source/blender/nodes/geometry/nodes/legacy/node_geo_material_assign.cc
new file mode 100644
index 00000000000..7d3481c1067
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/legacy/node_geo_material_assign.cc
@@ -0,0 +1,95 @@
+/*
+ * 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 "DNA_mesh_types.h"
+#include "DNA_meshdata_types.h"
+
+#include "BKE_material.h"
+
+namespace blender::nodes {
+
+static void geo_node_legacy_material_assign_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Geometry>("Geometry");
+  b.add_input<decl::Material>("Material").hide_label(true);
+  b.add_input<decl::String>("Selection");
+  b.add_output<decl::Geometry>("Geometry");
+}
+
+static void assign_material_to_faces(Mesh &mesh, const VArray<bool> &face_mask, Material *material)
+{
+  int new_material_index = -1;
+  for (const int i : IndexRange(mesh.totcol)) {
+    Material *other_material = mesh.mat[i];
+    if (other_material == material) {
+      new_material_index = i;
+      break;
+    }
+  }
+  if (new_material_index == -1) {
+    /* Append a new material index. */
+    new_material_index = mesh.totcol;
+    BKE_id_material_eval_assign(&mesh.id, new_material_index + 1, material);
+  }
+
+  mesh.mpoly = (MPoly *)CustomData_duplicate_referenced_layer(&mesh.pdata, CD_MPOLY, mesh.totpoly);
+  for (const int i : IndexRange(mesh.totpoly)) {
+    if (face_mask[i]) {
+      MPoly &poly = mesh.mpoly[i];
+      poly.mat_nr = new_material_index;
+    }
+  }
+}
+
+static void geo_node_legacy_material_assign_exec(GeoNodeExecParams params)
+{
+  Material *material = params.extract_input<Material *>("Material");
+  const std::string mask_name = params.extract_input<std::string>("Selection");
+
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+
+  geometry_set = geometry_set_realize_instances(geometry_set);
+
+  if (geometry_set.has<MeshComponent>()) {
+    MeshComponent &mesh_component = geometry_set.get_component_for_write<MeshComponent>();
+    Mesh *mesh = mesh_component.get_for_write();
+    if (mesh != nullptr) {
+      GVArray_Typed<bool> face_mask = mesh_component.attribute_get_for_read<bool>(
+          mask_name, ATTR_DOMAIN_FACE, true);
+      assign_material_to_faces(*mesh, face_mask, material);
+    }
+  }
+
+  params.set_output("Geometry", std::move(geometry_set));
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_legacy_material_assign()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_LEGACY_MATERIAL_ASSIGN, "Material Assign", NODE_CLASS_GEOMETRY, 0);
+  ntype.declare = blender::nodes::geo_node_legacy_material_assign_declare;
+  ntype.geometry_node_execute = blender::nodes::geo_node_legacy_material_assign_exec;
+  nodeRegisterType(&ntype);
+}
diff --git a/source/blender/nodes/geometry/nodes/node_geo_select_by_material.cc b/source/blender/nodes/geometry/nodes/legacy/node_geo_select_by_material.cc
index 1d1b9712fe0..eabdd2bcd5a 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_select_by_material.cc
+++ b/source/blender/nodes/geometry/nodes/legacy/node_geo_select_by_material.cc
@@ -28,7 +28,7 @@
 
 namespace blender::nodes {
 
-static void geo_node_select_by_material_declare(NodeDeclarationBuilder &b)
+static void geo_node_legacy_select_by_material_declare(NodeDeclarationBuilder &b)
 {
   b.add_input<decl::Geometry>("Geometry");
   b.add_input<decl::Material>("Material").hide_label();
@@ -54,7 +54,7 @@ static void select_mesh_by_material(const Mesh &mesh,
   });
 }
 
-static void geo_node_select_by_material_exec(GeoNodeExecParams params)
+static void geo_node_legacy_select_by_material_exec(GeoNodeExecParams params)
 {
   Material *material = params.extract_input<Material *>("Material");
   const std::string selection_name = params.extract_input<std::string>("Selection");
@@ -80,13 +80,13 @@ static void geo_node_select_by_material_exec(GeoNodeExecParams params)
 
 }  // namespace blender::nodes
 
-void register_node_type_geo_select_by_material()
+void register_node_type_geo_legacy_select_by_material()
 {
   static bNodeType ntype;
 
   geo_node_type_base(
       &ntype, GEO_NODE_LEGACY_SELECT_BY_MATERIAL, "Select by Material", NODE_CLASS_GEOMETRY, 0);
-  ntype.declare = blender::nodes::geo_node_select_by_material_declare;
-  ntype.geometry_node_execute = blender::nodes::geo_node_select_by_material_exec;
+  ntype.declare = blender::nodes::geo_node_legacy_select_by_material_declare;
+  ntype.geometry_node_execute = blender::nodes::geo_node_legacy_select_by_material_exec;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_material_assign.cc b/source/blender/nodes/geometry/nodes/node_geo_material_assign.cc
index 9e99ae29b00..bdcf06fe416 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_material_assign.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_material_assign.cc
@@ -30,11 +30,11 @@ static void geo_node_material_assign_declare(NodeDeclarationBuilder &b)
 {
   b.add_input<decl::Geometry>("Geometry");
   b.add_input<decl::Material>("Material").hide_label();
-  b.add_input<decl::String>("Selection");
+  b.add_input<decl::Bool>("Selection").default_value(true);
   b.add_output<decl::Geometry>("Geometry");
 }
 
-static void assign_material_to_faces(Mesh &mesh, const VArray<bool> &face_mask, Material *material)
+static void assign_material_to_faces(Mesh &mesh, const IndexMask selection, Material *material)
 {
   int new_material_index = -1;
   for (const int i : IndexRange(mesh.totcol)) {
@@ -51,18 +51,16 @@ static void assign_material_to_faces(Mesh &mesh, const VArray<bool> &face_mask,
   }
 
   mesh.mpoly = (MPoly *)CustomData_duplicate_referenced_layer(&mesh.pdata, CD_MPOLY, mesh.totpoly);
-  for (const int i : IndexRange(mesh.totpoly)) {
-    if (face_mask[i]) {
-      MPoly &poly = mesh.mpoly[i];
-      poly.mat_nr = new_material_index;
-    }
+  for (const int i : selection) {
+    MPoly &poly = mesh.mpoly[i];
+    poly.mat_nr = new_material_index;
   }
 }
 
 static void geo_node_material_assign_exec(GeoNodeExecParams params)
 {
   Material *material = params.extract_input<Material *>("Material");
-  const std::string mask_name = params.extract_input<std::string>("Selection");
+  const Field<bool> selection_field = params.extract_input<Field<bool>>("Selection");
 
   GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
 
@@ -72,9 +70,15 @@ static void geo_node_material_assign_exec(GeoNodeExecParams params)
     MeshComponent &mesh_component = geometry_set.get_component_for_write<MeshComponent>();
     Mesh *mesh = mesh_component.get_for_write();
     if (mesh != nullptr) {
-      GVArray_Typed<bool> face_mask = mesh_component.attribute_get_for_read<bool>(
-          mask_name, ATTR_DOMAIN_FACE, true);
-      assign_material_to_faces(*mesh, face_mask, material);
+
+      GeometryComponentFieldContext field_context{mesh_component, ATTR_DOMAIN_FACE};
+
+      fn::FieldEvaluator selection_evaluator{field_context, mesh->totpoly};
+      selection_evaluator.add(selection_field);
+      selection_evaluator.evaluate();
+      const IndexMask selection = selection_evaluator.get_evaluated_as_mask(0);
+
+      assign_material_to_faces(*mesh, selection, material);
     }
   }
 
@@ -87,8 +91,7 @@ void register_node_type_geo_material_assign()
 {
   static bNodeType ntype;
 
-  geo_node_type_base(
-      &ntype, GEO_NODE_LEGACY_MATERIAL_ASSIGN, "Material Assign", NODE_CLASS_GEOMETRY, 0);
+  geo_node_type_base(&ntype, GEO_NODE_MATERIAL_ASSIGN, "Material Assign", NODE_CLASS_GEOMETRY, 0);
   ntype.declare = blender::nodes::geo_node_material_assign_declare;
   ntype.geometry_node_execute = blender::nodes::geo_node_material_assign_exec;
   nodeRegisterType(&ntype);
diff --git a/source/blender/nodes/geometry/nodes/node_geo_material_selection.cc b/source/blender/nodes/geometry/nodes/node_geo_material_selection.cc
new file mode 100644
index 00000000000..22c24e34314
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_material_selection.cc
@@ -0,0 +1,131 @@
+/*
+ * 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 "DNA_mesh_types.h"
+#include "DNA_meshdata_types.h"
+
+#include "BLI_task.hh"
+
+#include "BKE_material.h"
+
+namespace blender::nodes {
+
+static void geo_node_material_selection_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Material>("Material").hide_label(true);
+  b.add_output<decl::Bool>("Selection");
+}
+
+static void select_mesh_by_material(const Mesh &mesh,
+                                    const Material *material,
+                                    const IndexMask mask,
+                                    const MutableSpan<bool> r_selection)
+{
+  BLI_assert(mesh.totpoly >= r_selection.size());
+  Vector<int> material_indices;
+  for (const int i : IndexRange(mesh.totcol)) {
+    if (mesh.mat[i] == material) {
+      material_indices.append(i);
+    }
+  }
+  threading::parallel_for(mask.index_range(), 1024, [&](IndexRange range) {
+    for (const int i : range) {
+      const int face_index = mask[i];
+      r_selection[i] = material_indices.contains(mesh.mpoly[face_index].mat_nr);
+    }
+  });
+}
+
+class MaterialSelectionFieldInput final : public fn::FieldInput {
+  Material *material_;
+
+ public:
+  MaterialSelectionFieldInput(Material *material)
+      : fn::FieldInput(CPPType::get<bool>(), "Material Selection"), material_(material)
+  {
+  }
+
+  const GVArray *get_varray_for_context(const fn::FieldContext &context,
+                                        IndexMask mask,
+                                        ResourceScope &scope) const final
+  {
+    if (const GeometryComponentFieldContext *geometry_context =
+            dynamic_cast<const GeometryComponentFieldContext *>(&context)) {
+      const GeometryComponent &component = geometry_context->geometry_component();
+      const AttributeDomain domain = geometry_context->domain();
+      if (component.type() != GEO_COMPONENT_TYPE_MESH) {
+        return nullptr;
+      }
+      const MeshComponent &mesh_component = static_cast<const MeshComponent &>(component);
+      const Mesh *mesh = mesh_component.get_for_read();
+      if (mesh == nullptr) {
+        return nullptr;
+      }
+
+      if (domain == ATTR_DOMAIN_FACE) {
+        Array<bool> selection(mask.min_array_size());
+        select_mesh_by_material(*mesh, material_, mask, selection);
+        return &scope.construct<fn::GVArray_For_ArrayContainer<Array<bool>>>(std::move(selection));
+      }
+
+      Array<bool> selection(mesh->totpoly);
+      select_mesh_by_material(*mesh, material_, IndexMask(mesh->totpoly), selection);
+      GVArrayPtr face_selection = std::make_unique<fn::GVArray_For_ArrayContainer<Array<bool>>>(
+          std::move(selection));
+      GVArrayPtr final_selection = mesh_component.attribute_try_adapt_domain(
+          std::move(face_selection), ATTR_DOMAIN_FACE, domain);
+      return scope.add_value(std::move(final_selection)).get();
+    }
+
+    return nullptr;
+  }
+
+  uint64_t hash() const override
+  {
+    /* Some random constant hash. */
+    return 91619626;
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const override
+  {
+    return dynamic_cast<const MaterialSelectionFieldInput *>(&other) != nullptr;
+  }
+};
+
+static void geo_node_material_selection_exec(GeoNodeExecParams params)
+{
+  Material *material = params.extract_input<Material *>("Material");
+  Field<bool> material_field{std::make_shared<MaterialSelectionFieldInput>(material)};
+  params.set_output("Selection", std::move(material_field));
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_material_selection()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_MATERIAL_SELECTION, "Material Selection", NODE_CLASS_GEOMETRY, 0);
+  ntype.declare = blender::nodes::geo_node_material_selection_declare;
+  ntype.geometry_node_execute = blender::nodes::geo_node_material_selection_exec;
+  nodeRegisterType(&ntype);
+}