Geometry Nodes: Expose material index attribute

The `material_index` attribute can adjust which material in the list
will be applied to each face of the mesh. There are two new things
about this attribute that haven't been exposed by the attribute API yet.
Each comes with limitations:
 1. Integer data type: Most attribute nodes are currently written to use
    float data types. This means that they can't write to this attribute
    because they can't change the type of a built-in attribute.
 2. Polygon domain: This is our first attribute using the polygon domain,
    meaning until some of the interpolations are implemented, some
    operations may not work as expected.

Currently the two nodes that work with this attribute are Attribute Fill
and Attribute Randomize.

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

3 files changed, 43 insertions, 5 deletions

diff --git a/source/blender/blenkernel/intern/attribute_access.cc b/source/blender/blenkernel/intern/attribute_access.cc
index f594ec54a7e..aa721a0b4ca 100644
--- a/source/blender/blenkernel/intern/attribute_access.cc
+++ b/source/blender/blenkernel/intern/attribute_access.cc
@@ -1292,6 +1292,29 @@ static void tag_normals_dirty_when_writing_position(GeometryComponent &component
   }
 }
 
+static int get_material_index(const MPoly &mpoly)
+{
+  return static_cast<int>(mpoly.mat_nr);
+}
+
+static void set_material_index(MPoly &mpoly, const int &index)
+{
+  mpoly.mat_nr = static_cast<short>(std::clamp(index, 0, SHRT_MAX));
+}
+
+static ReadAttributePtr make_material_index_read_attribute(const void *data, const int domain_size)
+{
+  return std::make_unique<DerivedArrayReadAttribute<MPoly, int, get_material_index>>(
+      ATTR_DOMAIN_POLYGON, Span<MPoly>((const MPoly *)data, domain_size));
+}
+
+static WriteAttributePtr make_material_index_write_attribute(void *data, const int domain_size)
+{
+  return std::make_unique<
+      DerivedArrayWriteAttribute<MPoly, int, get_material_index, set_material_index>>(
+      ATTR_DOMAIN_POLYGON, MutableSpan<MPoly>((MPoly *)data, domain_size));
+}
+
 template<typename T, AttributeDomain Domain>
 static ReadAttributePtr make_array_read_attribute(const void *data, const int domain_size)
 {
@@ -1355,6 +1378,19 @@ static ComponentAttributeProviders create_attribute_providers_for_mesh()
                                                  make_vertex_position_read_attribute,
                                                  make_vertex_position_write_attribute,
                                                  tag_normals_dirty_when_writing_position);
+
+  static BuiltinCustomDataLayerProvider material_index("material_index",
+                                                       ATTR_DOMAIN_POLYGON,
+                                                       CD_PROP_INT32,
+                                                       CD_MPOLY,
+                                                       BuiltinAttributeProvider::NonCreatable,
+                                                       BuiltinAttributeProvider::Writable,
+                                                       BuiltinAttributeProvider::NonDeletable,
+                                                       polygon_access,
+                                                       make_material_index_read_attribute,
+                                                       make_material_index_write_attribute,
+                                                       nullptr);
+
   static MeshUVsAttributeProvider uvs;
   static VertexGroupsAttributeProvider vertex_groups;
   static CustomDataAttributeProvider corner_custom_data(ATTR_DOMAIN_CORNER, corner_access);
@@ -1362,7 +1398,7 @@ static ComponentAttributeProviders create_attribute_providers_for_mesh()
   static CustomDataAttributeProvider edge_custom_data(ATTR_DOMAIN_EDGE, edge_access);
   static CustomDataAttributeProvider polygon_custom_data(ATTR_DOMAIN_POLYGON, polygon_access);
 
-  return ComponentAttributeProviders({&position},
+  return ComponentAttributeProviders({&position, &material_index},
                                      {&uvs,
                                       &corner_custom_data,
                                       &vertex_groups,
diff --git a/source/blender/blenkernel/intern/geometry_set_instances.cc b/source/blender/blenkernel/intern/geometry_set_instances.cc
index 8e9edc86bca..436fc0f1d1d 100644
--- a/source/blender/blenkernel/intern/geometry_set_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_set_instances.cc
@@ -350,9 +350,10 @@ static void join_instance_groups_mesh(Span<GeometryInstanceGroup> set_groups, Ge
   MeshComponent &dst_component = result.get_component_for_write<MeshComponent>();
   dst_component.replace(new_mesh);
 
-  /* The position attribute is handled above already. */
+  /* Don't copy attributes that are stored directly in the mesh data structs. */
   Map<std::string, AttributeInfo> attributes;
-  gather_attribute_info(attributes, GeometryComponentType::Mesh, set_groups, {"position"});
+  gather_attribute_info(
+      attributes, GeometryComponentType::Mesh, set_groups, {"position", "material_index"});
   join_attributes(set_groups,
                   GeometryComponentType::Mesh,
                   attributes,
diff --git a/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc b/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
index 3f6f28f2826..4e15f232934 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc
@@ -221,8 +221,9 @@ static void join_components(Span<const MeshComponent *> src_components, Geometry
   MeshComponent &dst_component = result.get_component_for_write<MeshComponent>();
   dst_component.replace(new_mesh);
 
-  /* The position attribute is handled above already. */
-  join_attributes(to_base_components(src_components), dst_component, {"position"});
+  /* Don't copy attributes that are stored directly in the mesh data structs. */
+  join_attributes(
+      to_base_components(src_components), dst_component, {"position", "material_index"});
 }
 
 static void join_components(Span<const PointCloudComponent *> src_components, GeometrySet &result)