RealityKit Multi-Part LowLevelMesh with Material Indexing

MultiPartMeshView.swift

import RealityKit
import SwiftUI

struct MultiPartMeshView: View {
    let radius: Float = 0.25
    
    var body: some View {
        RealityView { content in
            if let mesh = try? createMesh(),
               let resource = try? MeshResource(from: mesh)
            {
                updateMesh(mesh, radius: radius)
                
                // Two materials we will choose from using material index
                let materials = [UnlitMaterial(color: .red), UnlitMaterial(color: .green)]
                
                let entity = ModelEntity(mesh: resource, materials: materials)
                content.add(entity)
            }
        }
    }
    
    func createMesh() throws -> LowLevelMesh {
        var desc = VertexDefinition.descriptor
        desc.vertexCapacity = 6
        desc.indexCapacity = 6
        let mesh = try LowLevelMesh(descriptor: desc)
        return mesh
    }
    
    func updateMesh(_ mesh: LowLevelMesh, radius: Float) {
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertices = rawBytes.bindMemory(to: VertexDefinition.self)
            
            // Triangle 1 (left triangle)
            vertices[0] = VertexDefinition(position: [-radius, -radius, 0]) // bottom-left
            vertices[1] = VertexDefinition(position: [0, -radius, 0]) // bottom-right
            vertices[2] = VertexDefinition(position: [0, radius, 0]) // top
            
            // Triangle 2 (right triangle)
            vertices[3] = VertexDefinition(position: [0, -radius, 0]) // bottom-left
            vertices[4] = VertexDefinition(position: [radius, -radius, 0]) // bottom-right
            vertices[5] = VertexDefinition(position: [0, radius, 0]) // top
        }

        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            // Triangle 1 indices
            indices[0] = 0 // bottom-left
            indices[1] = 1 // bottom-right
            indices[2] = 2 // top
            
            // Triangle 2 indices
            indices[3] = 3 // bottom-left
            indices[4] = 4 // bottom-right
            indices[5] = 5 // top
        }

        let meshBounds = BoundingBox(min: [-radius, -radius, 0], max: [radius, radius, 0])
        
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: 3,
                topology: .triangle,
                materialIndex: 0,
                bounds: meshBounds
            ),
            LowLevelMesh.Part(
                indexOffset: indexOffsetInBytes(fromIndexCount: 3),
                indexCount: 3,
                topology: .triangle,
                materialIndex: 1,
                bounds: meshBounds
            )
        ])
    }
    
    func indexOffsetInBytes(fromIndexCount indexCount: Int) -> Int {
        return indexCount * MemoryLayout<UInt32>.stride
    }
}

#Preview {
    MultiPartMeshView()
}

fileprivate struct VertexDefinition {
    var position: SIMD3<Float> = .zero
    
    static var vertexAttributes: [LowLevelMesh.Attribute] = [
        .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!),
    ]

    static var vertexLayouts: [LowLevelMesh.Layout] = [
        .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride)
    ]

    static var descriptor: LowLevelMesh.Descriptor {
        var desc = LowLevelMesh.Descriptor()
        desc.vertexAttributes = VertexDefinition.vertexAttributes
        desc.vertexLayouts = VertexDefinition.vertexLayouts
        desc.indexType = .uint32
        return desc
    }
}