scripts/startup/bl_ui/node_add_menu_geometry.py

@@ -307,6 +307,7 @@ class NODE_MT_geometry_node_GEO_MATERIAL(Menu):
     def draw(self, _context):
         layout = self.layout
         node_add_menu.add_node_type(layout, "GeometryNodeReplaceMaterial")
+        node_add_menu.add_node_type(layout, "GeometryNodeSampleMaterials")
         layout.separator()
         node_add_menu.add_node_type(layout, "GeometryNodeInputMaterialIndex")
         node_add_menu.add_node_type(layout, "GeometryNodeMaterialSelection")

source/blender/blenkernel/BKE_node.h

@@ -1575,6 +1575,7 @@ void BKE_nodetree_remove_layer_n(struct bNodeTree *ntree, struct Scene *scene, i
 #define GEO_NODE_SDF_VOLUME_SPHERE 1196
 #define GEO_NODE_MEAN_FILTER_SDF_VOLUME 1197
 #define GEO_NODE_OFFSET_SDF_VOLUME 1198
+#define GEO_NODE_SAMPLE_MATERIALS 1199
 
 /** \} */
 

source/blender/blenlib/BLI_vector.hh

@@ -486,6 +486,15 @@ class Vector {
     }
   }
 
+  int append_non_duplicates_and_get_index(const T &value)
+  {
+    const int index = this->first_index_of_try(value);
+    if (index == -1) {
+      return this->append_and_get_index(value);
+    }
+    return index;
+  }
+
   /**
    * Append the value and assume that vector has enough memory reserved. This invokes undefined
    * behavior when not enough capacity has been reserved beforehand. Only use this in performance

source/blender/nodes/NOD_static_types.h

@@ -397,6 +397,7 @@ DefNode(GeometryNode, GEO_NODE_REVERSE_CURVE, 0, "REVERSE_CURVE", ReverseCurve,
 DefNode(GeometryNode, GEO_NODE_ROTATE_INSTANCES, 0, "ROTATE_INSTANCES", RotateInstances, "Rotate Instances", "Rotate geometry instances in local or global space")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_CURVE, def_geo_curve_sample, "SAMPLE_CURVE", SampleCurve, "Sample Curve", "Retrieve data from a point on a curve at a certain distance from its start")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_INDEX, def_geo_sample_index, "SAMPLE_INDEX", SampleIndex, "Sample Index", "Retrieve values from specific geometry elements")
+DefNode(GeometryNode, GEO_NODE_SAMPLE_MATERIALS, 0, "SAMPLE_MATERIALS", SampleMaterials, "Sample Materials", "")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST_SURFACE, def_geo_sample_nearest_surface, "SAMPLE_NEAREST_SURFACE", SampleNearestSurface, "Sample Nearest Surface", "Calculate the interpolated value of a mesh attribute on the closest point of its surface")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_UV_SURFACE, def_geo_sample_uv_surface, "SAMPLE_UV_SURFACE", SampleUVSurface, "Sample UV Surface", "Calculate the interpolated values of a mesh attribute at a UV coordinate")

source/blender/nodes/geometry/CMakeLists.txt

@@ -148,6 +148,7 @@ set(SRC
   nodes/node_geo_remove_attribute.cc
   nodes/node_geo_rotate_instances.cc
   nodes/node_geo_sample_index.cc
+  nodes/node_geo_sample_materials.cc
   nodes/node_geo_sample_nearest.cc
   nodes/node_geo_sample_nearest_surface.cc
   nodes/node_geo_sample_uv_surface.cc

source/blender/nodes/geometry/node_geometry_register.cc

@@ -132,6 +132,7 @@ void register_geometry_nodes()
   register_node_type_geo_remove_attribute();
   register_node_type_geo_rotate_instances();
   register_node_type_geo_sample_index();
+  register_node_type_geo_sample_materials();
   register_node_type_geo_sample_nearest_surface();
   register_node_type_geo_sample_nearest();
   register_node_type_geo_sample_uv_surface();

source/blender/nodes/geometry/node_geometry_register.hh

@@ -129,6 +129,7 @@ void register_node_type_geo_realize_instances();
 void register_node_type_geo_remove_attribute();
 void register_node_type_geo_rotate_instances();
 void register_node_type_geo_sample_index();
+void register_node_type_geo_sample_materials();
 void register_node_type_geo_sample_nearest_surface();
 void register_node_type_geo_sample_nearest();
 void register_node_type_geo_sample_uv_surface();

source/blender/nodes/geometry/nodes/node_geo_sample_materials.cc

@@ -0,0 +1,110 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "DNA_mesh_types.h"
+
+#include "BKE_attribute.hh"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_sample_materials_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Geometry>(N_("Geometry"));
+  b.add_input<decl::Geometry>(N_("Material Source"));
+  b.add_input<decl::Bool>(N_("Selection")).default_value(true).hide_value().field_on({0});
+  b.add_input<decl::Int>(N_("Material Index")).field_on({0});
+  /* TODO: It's possible to avoid propagation from second geometry input? Or i miss something? */
+  b.add_output<decl::Geometry>(N_("Geometry")).propagate_all();
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry = params.extract_input<GeometrySet>("Geometry");
+  GeometrySet source_geometry = params.extract_input<GeometrySet>("Material Source");
+  Field<bool> selection_field = params.extract_input<Field<bool>>("Selection");
+  Field<int> index_field = params.extract_input<Field<int>>("Material Index");
+
+  const Mesh *source_mesh = source_geometry.get_mesh_for_read();
+
+  if (source_mesh == nullptr) {
+    params.set_output("Geometry", std::move(geometry));
+    return;
+  }
+
+  if (source_mesh->totcol == 0) {
+    params.set_output("Geometry", std::move(geometry));
+    return;
+  }
+
+  geometry.modify_geometry_sets([&](GeometrySet &geometry_set) {
+    const MeshComponent *component = geometry_set.get_component_for_read<MeshComponent>();
+    Mesh *mesh = geometry_set.get_mesh_for_write();
+    if (!mesh) {
+      return;
+    }
+
+    bke::GeometryFieldContext field_context{*component, ATTR_DOMAIN_FACE};
+    fn::FieldEvaluator evaluator{field_context, mesh->totpoly};
+    evaluator.set_selection(selection_field);
+    evaluator.add(index_field);
+    evaluator.evaluate();
+
+    const IndexMask selection = evaluator.get_evaluated_selection_as_mask();
+    if (selection.is_empty()) {
+      return;
+    }
+
+    const Span<Material *> original_materials(mesh->mat, mesh->totcol);
+
+    Vector<Material *> materials = original_materials;
+
+    const auto append_material = [source_mesh, &materials](const int index) -> int {
+      const int limited_index = math::clamp<int>(index, 0, source_mesh->totcol - 1);
+      Material *material = source_mesh->mat[limited_index];
+      return materials.append_non_duplicates_and_get_index(material);
+    };
+
+    bke::MutableAttributeAccessor attributes = mesh->attributes_for_write();
+    SpanAttributeWriter<int> dst_indices_attribute =
+        attributes.lookup_or_add_for_write_only_span<int>("material_index", ATTR_DOMAIN_FACE);
+    MutableSpan<int> dst_indices = dst_indices_attribute.span;
+
+    const VArray<int> &indices = evaluator.get_evaluated<int>(0);
+    devirtualize_varray(indices, [&](auto indices) {
+      selection.foreach_index(
+          [&](const int index) { dst_indices[index] = append_material(indices[index]); });
+    });
+    dst_indices_attribute.finish();
+
+    const int result_tot_material = materials.size();
+    const Span<Material *> new_materials = materials.as_span().drop_front(mesh->totcol);
+
+    if (new_materials.is_empty()) {
+      return;
+    }
+
+    Array<Material *> old_src_material = original_materials;
+    mesh->mat = static_cast<Material **>(
+        MEM_recallocN_id(mesh->mat, sizeof(void *) * result_tot_material, "matarray"));
+    MutableSpan<Material *> r_materials(mesh->mat, result_tot_material);
+    r_materials.drop_front(mesh->totcol).copy_from(new_materials);
+    mesh->totcol = result_tot_material;
+  });
+
+  params.set_output("Geometry", std::move(geometry));
+}
+
+}  // namespace blender::nodes::node_geo_sample_materials_cc
+
+void register_node_type_geo_sample_materials()
+{
+  namespace file_ns = blender::nodes::node_geo_sample_materials_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_SAMPLE_MATERIALS, "Sample Materials", NODE_CLASS_GEOMETRY);
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  ntype.declare = file_ns::node_declare;
+  nodeRegisterType(&ntype);
+}