RealityKit Multi-Branch LowLevelMesh with Incremental Growth

IncrementallyUpdatingMultiPartData.swift

import RealityKit
import SwiftUI

struct IncrementallyUpdatingMultiPartData {
    var settings: IncrementallyUpdatingMultiPartSettings

    // Mesh buffer offsets (set once during initialization)
    var vertexOffset: Int = 0
    var indexOffset: Int = 0
    
    // Growth state
    var segments: [BranchSegment] = []
    var dynamicSegment: BranchSegment?
    var completedSegments: Int = 0
    var currentGrowthProgress: Float = 0.0
    var isGrowing: Bool = true
    
    var drawState: BranchDrawState = .init()
    
    // We are only appending segments and moving the dynamic segment's vertex positions
    // So storing this state allows us to not recalculate all of the index
    // and vertex values that are unchanging
    struct BranchDrawState {
        var lastSegmentIndicesCompleted: Int = 0
        var lastSegmentVerticesCompleted: Int = 0
    }

    init(settings: IncrementallyUpdatingMultiPartSettings = .init()) {
        self.settings = settings
        reset()
    }
    
    mutating func reset() {
        segments.removeAll()
        dynamicSegment = nil
        completedSegments = 0
        currentGrowthProgress = 0.0
        isGrowing = true
        drawState = .init()
        
        let baseSegment = BranchSegment(
            startPosition: SIMD3<Float>(0, 0, 0),
            endPosition: SIMD3<Float>(0, settings.segmentHeight, 0),
            radius: settings.baseRadius
        )
        segments.append(baseSegment)
        completedSegments = 1
        initializeDynamicSegment()
    }
    
    // Computed properties for mesh generation
    var vertexCount: Int {
        (settings.radialSegments + 1) * (completedSegments + 2)
    }
    
    var indexCount: Int {
        settings.radialSegments * (completedSegments + 1) * 6
    }
    
    var boundingBox: BoundingBox {
        guard !segments.isEmpty else {
            return BoundingBox(min: SIMD3<Float>(-settings.baseRadius, 0, -settings.baseRadius),
                               max: SIMD3<Float>(settings.baseRadius, settings.segmentHeight, settings.baseRadius))
        }
        
        var minBounds = segments[0].startPosition
        var maxBounds = segments[0].startPosition
        
        for segment in segments {
            minBounds = min(minBounds, segment.startPosition - SIMD3<Float>(segment.radius, segment.radius, segment.radius))
            minBounds = min(minBounds, segment.endPosition - SIMD3<Float>(segment.radius, segment.radius, segment.radius))
            maxBounds = max(maxBounds, segment.startPosition + SIMD3<Float>(segment.radius, segment.radius, segment.radius))
            maxBounds = max(maxBounds, segment.endPosition + SIMD3<Float>(segment.radius, segment.radius, segment.radius))
        }
        
        // Include the dynamic segment in bounds calculation
        if let dynamic = dynamicSegment {
            let currentEndPosition = mix(dynamic.startPosition, dynamic.endPosition, t: currentGrowthProgress)
            
            minBounds = min(minBounds, currentEndPosition - SIMD3<Float>(dynamic.radius, dynamic.radius, dynamic.radius))
            maxBounds = max(maxBounds, currentEndPosition + SIMD3<Float>(dynamic.radius, dynamic.radius, dynamic.radius))
        }
        
        return BoundingBox(min: minBounds, max: maxBounds)
    }
    
    // Allocation capacities for mesh initialization
    var maxVertexCount: Int {
        (settings.radialSegments + 1) * (settings.maxSegments + 2)
    }
    
    var maxIndexCount: Int {
        settings.radialSegments * (settings.maxSegments + 1) * 6
    }
    
    private func generateNextDirection(from currentDirection: SIMD3<Float>) -> SIMD3<Float> {
        // Generate a random variation in direction
        let randomX = Float.random(in: -settings.directionVariation...settings.directionVariation)
        let randomZ = Float.random(in: -settings.directionVariation...settings.directionVariation)
        
        // Keep some upward bias
        let variation = SIMD3<Float>(randomX, 0.1, randomZ)
        let newDirection = currentDirection + variation
        
        return normalize(newDirection)
    }

    // Initialize the dynamic segment when growth starts
    private mutating func initializeDynamicSegment() {
        guard let lastSegment = segments.last else { return }
        
        // Generate new direction from the last segment's direction
        let newDirection = generateNextDirection(from: lastSegment.direction)
        let startPos = lastSegment.endPosition
        let endPos = startPos + newDirection * settings.segmentHeight
        let newRadius = lastSegment.radius * settings.taperRate
        
        dynamicSegment = BranchSegment(
            startPosition: startPos,
            endPosition: endPos,
            radius: newRadius
        )
    }
    
    // Get the current end position of the dynamic segment based on progress
    func getCurrentDynamicEndPosition() -> SIMD3<Float>? {
        guard let dynamic = dynamicSegment else { return nil }
        return mix(dynamic.startPosition, dynamic.endPosition, t: currentGrowthProgress)
    }
    
    // Growth state management - unified approach
    mutating func updateGrowth() {
        guard isGrowing else { return }
        currentGrowthProgress += settings.progressRate
        
        if currentGrowthProgress >= 1.0 && completedSegments < settings.maxSegments {
            // Complete the current growing segment by adding it to segments array
            if let segment = dynamicSegment {
                segments.append(segment)
                completedSegments += 1
                currentGrowthProgress = 0
                initializeDynamicSegment() // Create next growing segment
            }
        } else if completedSegments >= settings.maxSegments {
            isGrowing = false
        }
    }
    
    // Helper function for linear interpolation
    private func mix(_ a: SIMD3<Float>, _ b: SIMD3<Float>, t: Float) -> SIMD3<Float> {
        return a + (b - a) * t
    }
}

IncrementallyUpdatingMultiPartSettings.swift

import Foundation
import Metal

struct IncrementallyUpdatingMultiPartSettings {
    var baseRadius: Float = 0.0025
    var segmentHeight: Float = 0.02
    var radialSegments: Int = 16
    var maxSegments: Int = 16
    var progressRate: Float = 0.2
    var directionVariation: Float = 0.05
    var taperRate: Float = 0.9
    var topology: MTLPrimitiveType = .triangle
}

IncrementallyUpdatingMultiPartView.swift

import RealityKit
import SwiftUI

struct IncrementallyUpdatingMultiPartView: View {
    @State var meshParts: [IncrementallyUpdatingMultiPartData] = []
    @State var mesh: LowLevelMesh?
    @State var timer: Timer?
    let numberOfBranches: Int = 50
    
    var body: some View {
        GeometryReader3D { proxy in
            RealityView { content in
                var cumulativeVertexOffset = 0
                var cumulativeIndexOffset = 0
                
                for _ in 0..<numberOfBranches {
                    var settings = IncrementallyUpdatingMultiPartSettings()
                    settings.maxSegments = Int.random(in: 5...10)
                    var branchData = IncrementallyUpdatingMultiPartData(settings: settings)
                    
                    // Set the offsets for this branch
                    branchData.vertexOffset = cumulativeVertexOffset
                    branchData.indexOffset = cumulativeIndexOffset
                    
                    // Update cumulative offsets for next branch
                    cumulativeVertexOffset += branchData.maxVertexCount
                    cumulativeIndexOffset += branchData.maxIndexCount
                    
                    self.meshParts.append(branchData)
                }
                
                var vertexCapacity = 0
                var indexCapacity = 0
                for data in self.meshParts {
                    vertexCapacity += data.maxVertexCount
                    indexCapacity += data.maxIndexCount
                }
                let mesh = try! VertexData.initializeMesh(vertexCapacity: vertexCapacity,
                                                          indexCapacity: indexCapacity)
                self.mesh = mesh
                
                let meshResource = try! await MeshResource(from: mesh)
                
                var materials: [RealityKit.Material] = []
                
                for i in 0..<numberOfBranches {
                    var material = PhysicallyBasedMaterial()
                    
                    // Calculate hue based on branch index
                    let hue = Float(i) / Float(numberOfBranches) * 0.125 + 0.3
                    let color = UIColor(hue: CGFloat(hue), saturation: 1.0, brightness: 1.0, alpha: 1.0)
                    
                    material.baseColor = .init(tint: color)
                    material.metallic = 0.0
                    material.roughness = 0.0
                    material.faceCulling = .none
                    
                    materials.append(material)
                }
                
                let rootEntity = Entity()

                let extents = content.convert(proxy.frame(in: .local), from: .local, to: .scene).extents
                let yOffset = extents.y * 0.5
                rootEntity.position = SIMD3<Float>(x: 0, y: -yOffset, z: 0)
                content.add(rootEntity)
                
                let branchEntity = ModelEntity(mesh: meshResource, materials: materials)
                rootEntity.addChild(branchEntity)
            }
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
        .ornament(attachmentAnchor: .scene(.topTrailingFront), contentAlignment: .topLeading) {
            Button(action: {
                for index in 0..<meshParts.count {
                    meshParts[index].reset()
                }
            }, label: {
                Image(systemName: "arrow.trianglehead.counterclockwise")
            })
        }
    }
    
    func startTimer() {
        self.timer = Timer.scheduledTimer(withTimeInterval: 0.016, repeats: true) { _ in
            for index in 0..<meshParts.count {
                meshParts[index].updateGrowth()
                updateMeshGeometry(index: index)
            }
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    func updateMeshGeometry(index: Int) {
        let meshData = meshParts[index]
        guard let mesh = mesh, meshData.isGrowing else { return }
        var branchDrawState = meshData.drawState
        
        let startSegment: Int = branchDrawState.lastSegmentVerticesCompleted
        let endSegment = meshData.completedSegments
        
        // Grid configuration
        let branchesPerRow = 10
        let spacing: Float = 0.025
        let branchRadius = meshData.settings.baseRadius
        let branchDiameter = branchRadius * 2
        
        // Calculate grid position for this branch
        let row = index / branchesPerRow
        let column = index % branchesPerRow
        
        // Calculate grid dimensions
        let gridWidth = Float(branchesPerRow) * branchDiameter + Float(branchesPerRow - 1) * spacing
        let numberOfRows = (numberOfBranches + branchesPerRow - 1) / branchesPerRow // Ceiling division
        let gridDepth = Float(numberOfRows) * branchDiameter + Float(numberOfRows - 1) * spacing
        
        // Position this branch in the grid (centered around origin)
        let branchOffsetX = -gridWidth * 0.5 + branchRadius + Float(column) * (branchDiameter + spacing)
        let branchOffsetZ = -gridDepth * 0.5 + branchRadius + Float(row) * (branchDiameter + spacing)
        
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertices = rawBytes.bindMemory(to: VertexData.self)
            
            if startSegment == 0 {
                // Ring 0: Base (always at origin with grid position)
                generateVerticesForRing(
                    vertices: vertices,
                    vertexOffset: meshData.vertexOffset,
                    ringIndex: 0,
                    position: SIMD3<Float>(branchOffsetX, 0, branchOffsetZ),
                    direction: SIMD3<Float>(0, 1, 0), // Always start straight up
                    radius: meshData.settings.baseRadius,
                    radialSegments: meshData.settings.radialSegments,
                    maxSegments: meshData.settings.maxSegments
                )
            }
            
            if startSegment < endSegment {
                for segmentIndex in startSegment..<endSegment {
                    let segment = meshData.segments[segmentIndex]
                    
                    var position = segment.endPosition
                    position.x += branchOffsetX
                    position.z += branchOffsetZ
                    
                    generateVerticesForRing(
                        vertices: vertices,
                        vertexOffset: meshData.vertexOffset,
                        ringIndex: segmentIndex + 1,
                        position: position,
                        direction: segment.direction,
                        radius: segment.radius,
                        radialSegments: meshData.settings.radialSegments,
                        maxSegments: meshData.settings.maxSegments
                    )
                }
            }
            
            // Ring for dynamic segment (constantly updating it's vertex positions)
            if let dynamicSegment = meshData.dynamicSegment,
               var currentEndPosition = meshData.getCurrentDynamicEndPosition() {
                currentEndPosition.x += branchOffsetX
                currentEndPosition.z += branchOffsetZ
                
                generateVerticesForRing(
                    vertices: vertices,
                    vertexOffset: meshData.vertexOffset,
                    ringIndex: meshData.completedSegments + 1,
                    position: currentEndPosition,
                    direction: dynamicSegment.direction,
                    radius: dynamicSegment.radius,
                    radialSegments: meshData.settings.radialSegments,
                    maxSegments: meshData.settings.maxSegments
                )
            }
        }
        
        branchDrawState.lastSegmentVerticesCompleted = endSegment
        
        let startSegmentIndices: Int = branchDrawState.lastSegmentIndicesCompleted
        let endSegmentIndices = meshData.completedSegments+1

        let indexOffset = meshData.indexOffset
        let indexOffsetInBytes = indexOffsetInBytes(fromIndexCount: meshData.indexOffset)
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            
            if startSegmentIndices < endSegmentIndices {
                var indicesIndex = 6 * meshData.settings.radialSegments * startSegmentIndices + indexOffset
                for segment in startSegmentIndices..<endSegmentIndices {
                    generateIndicesForRing(indices: indices,
                                           index: indicesIndex,
                                           radialSegments: meshData.settings.radialSegments,
                                           segmentNumber: segment,
                                           vertexOffset: meshData.vertexOffset)
                    indicesIndex += 6 * meshData.settings.radialSegments
                }
            }
        }
        
        branchDrawState.lastSegmentIndicesCompleted = endSegmentIndices
        
        
        
        // Transform the index to make colors change gradually along rows
        let rowCount = (numberOfBranches + branchesPerRow - 1) / branchesPerRow
        let transformedMaterialIndex = column * rowCount + row

        let part = LowLevelMesh.Part(indexOffset: indexOffsetInBytes,
                                     indexCount:  meshData.indexCount,
                                     topology: meshData.settings.topology,
                                     materialIndex: transformedMaterialIndex,
                                     bounds: meshData.boundingBox)
        
        if let index = mesh.parts.firstIndex(where: { $0.indexOffset == indexOffsetInBytes }) {
            mesh.parts[index] = part
        } else {
            mesh.parts.append(part)
        }
        
        self.meshParts[index].drawState = branchDrawState
    }
    
    func indexOffsetInBytes(fromIndexCount indexCount: Int) -> Int {
        return indexCount * MemoryLayout<UInt32>.stride
    }
    
    func generateIndicesForRing(
        indices: UnsafeMutableBufferPointer<UInt32>,
        index: Int,
        radialSegments: Int,
        segmentNumber: Int,
        vertexOffset: Int
    ) {
        var localIndex = index
        for x in 0..<radialSegments {
            let a = vertexOffset + segmentNumber * (radialSegments + 1) + x
            let b = a + 1
            let c = a + (radialSegments + 1)
            let d = c + 1
            
            indices[localIndex] = UInt32(a)
            indices[localIndex + 1] = UInt32(c)
            indices[localIndex + 2] = UInt32(b)
            
            indices[localIndex + 3] = UInt32(b)
            indices[localIndex + 4] = UInt32(c)
            indices[localIndex + 5] = UInt32(d)
            
            localIndex += 6
        }
    }
    
    func generateVerticesForRing(
        vertices: UnsafeMutableBufferPointer<VertexData>,
        vertexOffset: Int,
        ringIndex: Int,
        position: SIMD3<Float>,
        direction: SIMD3<Float>,
        radius: Float,
        radialSegments: Int,
        maxSegments: Int
    ) {
        // Create a coordinate system for this ring
        let up = direction
        let right = normalize(cross(up, SIMD3<Float>(0, 0, 1)))
        let forward = normalize(cross(right, up))
        
        // Generate vertices around the circumference
        for x in 0...radialSegments {
            let angle = Float(x) / Float(radialSegments) * 2 * Float.pi
            let localX = cos(angle) * radius
            let localZ = sin(angle) * radius
            
            // Transform local coordinates to world space using segment's coordinate system
            let localOffset = right * localX + forward * localZ
            let vertexPosition = position + localOffset
            
            // Normal points outward from the cylinder axis
            let normal = normalize(localOffset)
            
            let u = Float(x) / Float(radialSegments)
            let v = Float(ringIndex) / Float(maxSegments)
            let uv = SIMD2<Float>(u, v)
            
            let index = vertexOffset + ringIndex * (radialSegments + 1) + x
            vertices[index] = VertexData(position: vertexPosition,
                                         normal: normal,
                                         uv: uv)
        }
    }
}

#Preview { IncrementallyUpdatingMultiPartView() }

VertexData+Extensions.swift

import Foundation
import RealityKit

extension VertexData {
    static var vertexAttributes: [LowLevelMesh.Attribute] = [
        .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!),
        .init(semantic: .normal, format: .float3, offset: MemoryLayout<Self>.offset(of: \.normal)!),
        .init(semantic: .uv0, format: .float2, offset: MemoryLayout<Self>.offset(of: \.uv)!)
    ]
    
    static var vertexLayouts: [LowLevelMesh.Layout] = [
        .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride)
    ]

    static var descriptor: LowLevelMesh.Descriptor {
        var desc = LowLevelMesh.Descriptor()
        desc.vertexAttributes = VertexData.vertexAttributes
        desc.vertexLayouts = VertexData.vertexLayouts
        desc.indexType = .uint32
        return desc
    }
    
    @MainActor static func initializeMesh(vertexCapacity: Int,
                                          indexCapacity: Int) throws -> LowLevelMesh {
        var desc = VertexData.descriptor
        desc.vertexCapacity = vertexCapacity
        desc.indexCapacity = indexCapacity
        return try LowLevelMesh(descriptor: desc)
    }
}

VertexData.h

#include <simd/simd.h>

#ifndef VertexData_h
#define VertexData_h

struct VertexData {
    simd_float3 position;
    simd_float3 normal;
    simd_float2 uv;
};

#endif /* VertexData_h */