Grass + Wind RealityView using LowLevelMesh, Poisson Disk Sampling, and Metal Compute Shader

BladeSegment.swift

import Foundation

struct BladeSegment {
    var startPosition: SIMD3<Float>
    var endPosition: SIMD3<Float>
    var radius: Float
    
    init(startPosition: SIMD3<Float>, endPosition: SIMD3<Float>, radius: Float) {
        self.startPosition = startPosition
        self.endPosition = endPosition
        self.radius = radius
    }
    
    var direction: SIMD3<Float> {
        return normalize(endPosition - startPosition)
    }
    
    var length: Float {
        return distance(startPosition, endPosition)
    }
}

GrassData.swift

import RealityKit
import SwiftUI

struct GrassData {
    var settings: GrassSettings
    var materialIndex: Int = 0

    // Mesh buffer offsets (set once during initialization)
    var vertexOffset: Int = 0
    var indexOffset: Int = 0
    
    // Growth state
    var segments: [BladeSegment] = []
    var dynamicSegment: BladeSegment?
    var completedSegments: Int = 0
    var currentGrowthProgress: Float = 0.0
    var isGrowing: Bool = true
    
    var drawState: BladeDrawState = .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 BladeDrawState {
        var lastSegmentIndicesCompleted: Int = 0
        var lastSegmentVerticesCompleted: Int = 0
    }

    init(settings: GrassSettings = .init()) {
        self.settings = settings
        reset()
    }
    
    mutating func reset() {
        segments.removeAll()
        dynamicSegment = nil
        completedSegments = 0
        currentGrowthProgress = 0.0
        isGrowing = true
        drawState = .init()
        
        let baseSegment = BladeSegment(
            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 {
        let rings = completedSegments + 1 + (dynamicSegment != nil ? 1 : 0) // base + completed + dynamic
        let tipVertex = 1
        return (settings.radialSegments + 1) * rings + tipVertex
    }
    
    var indexCount: Int {
        let ringConnections = settings.radialSegments * (completedSegments + 1) * 6
        
        // Add end cap triangles if this blade has an end cap
        let hasEndCap = !isGrowing || dynamicSegment != nil
        let endCapTriangles = hasEndCap ? settings.radialSegments * 3 : 0
        
        return ringConnections + endCapTriangles
    }
    
    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) + 1
    }
    
    var maxIndexCount: Int {
        let ringConnections = settings.radialSegments * (settings.maxSegments + 1) * 6
        let endCapTriangles = settings.radialSegments * 3  // Always account for potential end cap
        return ringConnections + endCapTriangles
    }
    
    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 = BladeSegment(
            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
    }
}

GrassSettings.swift

import Foundation
import Metal

struct GrassSettings {
    var baseRadius: Float = 0.0025
    var segmentHeight: Float = 0.02
    var radialSegments: Int = 6
    var maxSegments: Int = 8
    var progressRate: Float = 0.2
    var directionVariation: Float = 0.1
    var taperRate: Float = 0.8
    var topology: MTLPrimitiveType = .triangle
}

GrassView.swift

import RealityKit
import SwiftUI

#Preview { GrassView() }
struct GrassView: View {
    @State var meshParts: [GrassData] = []
    @State var mesh: LowLevelMesh?
    @State var timer: Timer?
    @State var windTime: Float = 0.0
    @State var rootEntity = Entity()
    @State var grassEntity: Entity?
    @State var samplePoints: [SIMD2<Float>] = []
    
    // Poisson disk sampling parameters
    @State var grassAreaWidth: Float = 0.2
    @State var grassAreaHeight: Float = 0.2
    
    @State var soilPadding: Float = 0.1

    @State var populationDensity: PopulationDensity = .high
    var minimumSpacing: Float {
        populationDensity.minimumSpacing
    }
    
    // holding this so we don't recreate on resize
    @State var soilTextureResource: TextureResource?
    
    let numberOfMaterials: Int = 10
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipeline: MTLComputePipelineState
    
    enum PopulationDensity {
        case low
        case medium
        case high
        
        var minimumSpacing: Float {
            switch self {
            case .low:
                return 0.05
            case .medium:
                return 0.035
            case .high:
                return 0.025
            }
        }
    }
    
    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "updateGrassVertices")!
        self.computePipeline = try! device.makeComputePipelineState(function: updateFunction)
    }

    var body: some View {
        GeometryReader3D { proxy in
            RealityView { content in
                let extents = content.convert(proxy.frame(in: .local), from: .local, to: .scene).extents
                let yOffset = extents.y * 0.5
                await setupGrassMesh()
                await setupSoil()
                
                rootEntity.position = SIMD3<Float>(x: 0, y: -yOffset, z: 0)
                content.add(rootEntity)
            } update: { content in
                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)
            }
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
        .onChange(of: populationDensity) { oldValue, newValue in
            reset()
        }
        .ornament(attachmentAnchor: .scene(.top), contentAlignment: .top) {
            VStack {
                Button(action: {reset() },
                       label: {
                    HStack {
                        Image(systemName: "arrow.trianglehead.counterclockwise")
                        Text("Reset")
                    }
                })
                
                Picker("Population Density", selection: $populationDensity) {
                    Text("Low").tag(PopulationDensity.low)
                    Text("Medium").tag(PopulationDensity.medium)
                    Text("High").tag(PopulationDensity.high)
                }
                .pickerStyle(SegmentedPickerStyle())
                .frame(maxWidth: 300)
                
                VStack(spacing: 10) {
                    HStack {
                        Text("Width:")
                            .frame(width: 80, alignment: .leading)
                        Slider(value: $grassAreaWidth, in: 0.125...0.25, onEditingChanged: { editing in
                            if !editing {
                                reset(dimensionChanged: true)
                            }
                        })
                        Text("\(String(format: "%.3f", grassAreaWidth))")
                            .frame(width: 50, alignment: .trailing)
                    }
                    
                    HStack {
                        Text("Height:")
                            .frame(width: 80, alignment: .leading)
                        Slider(value: $grassAreaHeight, in: 0.125...0.25, onEditingChanged: { editing in
                            if !editing {
                                reset(dimensionChanged: true)
                            }
                        })
                        Text("\(String(format: "%.3f", grassAreaHeight))")
                            .frame(width: 50, alignment: .trailing)
                    }
                    
                    HStack {
                        Text("Padding:")
                            .frame(width: 80, alignment: .leading)
                        Slider(value: $soilPadding, in: 0.0...0.2, onEditingChanged: { editing in
                            if !editing {
                                reset(dimensionChanged: true)
                            }
                        })
                        Text("\(String(format: "%.3f", soilPadding))")
                            .frame(width: 50, alignment: .trailing)
                    }
                }
                .frame(width: 300)
                
                Text("Blades: \(samplePoints.count)")
                    .font(.title)
                
                Text("Min Spacing: \(String(format: "%.3f", minimumSpacing))")
                    .font(.title3)
            }
            .padding(20)
            .glassBackgroundEffect()
            .padding(.top, 200)
        }
    }
}

// MARK: - Reset
extension GrassView {
    func reset(dimensionChanged: Bool = false) {
        mesh = nil
        meshParts.removeAll()
        if let grassEntity {
            grassEntity.components.removeAll()
        }
        if dimensionChanged {
            rootEntity.children.removeAll()
        }
        Task {
            if dimensionChanged {
                await setupSoil()
            }
            await setupGrassMesh()
        }
    }
}

// MARK: - Poisson Disk Sampling Algorithm
extension GrassView {
    /// Generate sample points using Bridson's Algorithm (Fast Poisson Disk Sampling)
    func generatePoissonDiskSamples(width: Float, height: Float, minimumDistance: Float, maxAttempts: Int = 30) -> [SIMD2<Float>] {
        let cellSize = minimumDistance / sqrt(2.0)
        let gridWidth = Int(ceil(width / cellSize))
        let gridHeight = Int(ceil(height / cellSize))
        
        var grid = Array(repeating: Array<SIMD2<Float>?>(repeating: nil, count: gridWidth), count: gridHeight)
        var samples: [SIMD2<Float>] = []
        var activeList: [SIMD2<Float>] = []
        
        func gridCoordinates(for point: SIMD2<Float>) -> (x: Int, y: Int) {
            let x = Int(point.x / cellSize)
            let y = Int(point.y / cellSize)
            return (x, y)
        }
        
        func isInBounds(_ point: SIMD2<Float>) -> Bool {
            return point.x >= 0 && point.x < width && point.y >= 0 && point.y < height
        }
        
        func isFarEnough(from point: SIMD2<Float>) -> Bool {
            let coords = gridCoordinates(for: point)
            let gx = coords.x
            let gy = coords.y
            
            // Check neighboring cells in a 5x5 grid around the point
            for dx in -2...2 {
                for dy in -2...2 {
                    let nx = gx + dx
                    let ny = gy + dy
                    
                    if nx < 0 || ny < 0 || nx >= gridWidth || ny >= gridHeight {
                        continue
                    }
                    
                    if let existingPoint = grid[ny][nx] {
                        let distance = length(point - existingPoint)
                        if distance < minimumDistance {
                            return false
                        }
                    }
                }
            }
            return true
        }
        
        func generatePointAround(_ point: SIMD2<Float>) -> SIMD2<Float> {
            let angle = Float.random(in: 0..<(2 * Float.pi))
            let radius = Float.random(in: minimumDistance...(2 * minimumDistance))
            
            return SIMD2<Float>(
                point.x + cos(angle) * radius,
                point.y + sin(angle) * radius
            )
        }
        
        // Generate initial random point
        let initialPoint = SIMD2<Float>(
            Float.random(in: 0..<width),
            Float.random(in: 0..<height)
        )
        
        samples.append(initialPoint)
        activeList.append(initialPoint)
        
        let coords = gridCoordinates(for: initialPoint)
        grid[coords.y][coords.x] = initialPoint
        
        // Main algorithm loop
        while !activeList.isEmpty {
            let randomIndex = Int.random(in: 0..<activeList.count)
            let point = activeList[randomIndex]
            var foundValidPoint = false
            
            for _ in 0..<maxAttempts {
                let candidate = generatePointAround(point)
                
                if isInBounds(candidate) && isFarEnough(from: candidate) {
                    samples.append(candidate)
                    activeList.append(candidate)
                    
                    let coords = gridCoordinates(for: candidate)
                    grid[coords.y][coords.x] = candidate
                    foundValidPoint = true
                    break
                }
            }
            
            if !foundValidPoint {
                activeList.remove(at: randomIndex)
            }
        }
        
        return samples
    }
}

// MARK: Animation Timer
extension GrassView {
    func startTimer() {
        self.timer = Timer.scheduledTimer(withTimeInterval: 0.016, repeats: true) { _ in
            for index in 0..<meshParts.count {
                meshParts[index].updateGrowth()
                updateMeshGeometry(index: index)
            }
            updateGrassForWindEffect()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
}

// MARK: Repositioning on Resize
extension GrassView {
    func updateOffsets(yOffset: Float) {
        if let soilEntity = rootEntity.children.first(where: { $0.name == "soil" }) {
            soilEntity.position = SIMD3<Float>(x: 0, y: yOffset, z: 0)
        }
        if let grassEntity = rootEntity.children.first(where: { $0.name == "grass" }) {
            grassEntity.position = SIMD3<Float>(x: 0, y: yOffset, z: 0)
        }
    }
}

// MARK: Wind Offset Grass Mesh Vertices
extension GrassView {
    func updateGrassForWindEffect() {
        guard let mesh = mesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        windTime += 0.016
        
        // Read the current data AND get a buffer to write to
        let currentVertexBuffer = mesh.read(bufferIndex: 0, using: commandBuffer)
        let newVertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipeline)
        computeEncoder.setBuffer(currentVertexBuffer, offset: 0, index: 0)  // Read from original
        computeEncoder.setBuffer(newVertexBuffer, offset: 0, index: 1)      // Write to new
        
        // Wind parameters
        var windStrength: Float = 0.004
        computeEncoder.setBytes(&windStrength, length: MemoryLayout<Float>.stride, index: 2)
        computeEncoder.setBytes(&windTime, length: MemoryLayout<Float>.stride, index: 3)
        
        // Process full vertex capacity
        let totalVertices = mesh.vertexCapacity
        let threadgroupSize = MTLSize(width: 64, height: 1, depth: 1)
        let threadgroups = MTLSize(
            width: (totalVertices + threadgroupSize.width - 1) / threadgroupSize.width,
            height: 1,
            depth: 1
        )
        
        computeEncoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
}

// MARK: Setup Grass Mesh
extension GrassView {
    func setupGrassMesh() async {
        // Generate sample points using Poisson disk sampling
        samplePoints = generatePoissonDiskSamples(
            width: grassAreaWidth,
            height: grassAreaHeight,
            minimumDistance: minimumSpacing
        )
        
        var cumulativeVertexOffset = 0
        var cumulativeIndexOffset = 0
        
        // Create grass data for each sample point
        for index in 0..<samplePoints.count {
            var settings = GrassSettings()
            settings.maxSegments = Int.random(in: 5...10)
            var grassData = GrassData(settings: settings)
            
            // Set consistent material index
            grassData.materialIndex = index % numberOfMaterials
            
            // Set the offsets for this branch
            grassData.vertexOffset = cumulativeVertexOffset
            grassData.indexOffset = cumulativeIndexOffset
            
            // Update cumulative offsets for next branch
            cumulativeVertexOffset += grassData.maxVertexCount
            cumulativeIndexOffset += grassData.maxIndexCount
            
            self.meshParts.append(grassData)
        }
        
        var vertexCapacity = 0
        var indexCapacity = 0
        for data in self.meshParts {
            vertexCapacity += data.maxVertexCount
            indexCapacity += data.maxIndexCount
        }
        let mesh = try! VertexDataWithOriginalPosition.initializeMesh(vertexCapacity: vertexCapacity,
                                                                      indexCapacity: indexCapacity)
        self.mesh = mesh
        
        let meshResource = try! await MeshResource(from: mesh)
        
        var materials: [RealityKit.Material] = []
        
        for i in 0..<numberOfMaterials {
            var material = PhysicallyBasedMaterial()
            
            // Calculate hue based on blade index
            let hue = Float(i) / Float(numberOfMaterials) * 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 entity =  ModelEntity(mesh: meshResource, materials: materials)
        grassEntity = entity
        rootEntity.addChild(entity)
    }
}

// MARK: Setup Soil
extension GrassView {
    func setupSoil() async {
        var textureResource: TextureResource
        if let soilTextureResource {
            textureResource = soilTextureResource
        } else {
            let soilTextureURL = URL(string: "https://matt54.github.io/Resources/soil_texture_color.jpg")!
            textureResource = try! await TextureResource.loadOnlineImage(soilTextureURL)
            self.soilTextureResource = textureResource
        }
        
        let soilMesh = MeshResource.generatePlane(width: grassAreaWidth+soilPadding, depth: grassAreaHeight+soilPadding)
        
        var soilMaterial = PhysicallyBasedMaterial()
        soilMaterial.baseColor = .init(texture: .init(textureResource))
        let soilEntity = ModelEntity(mesh: soilMesh,
                                       materials: [soilMaterial])
        rootEntity.addChild(soilEntity)
    }
}

// MARK: Incrementally Growing Grass Mesh
extension GrassView {
    func updateMeshGeometry(index: Int) {
        guard index < samplePoints.count else { return }
        
        let meshData = meshParts[index]
        guard let mesh = mesh, meshData.isGrowing else { return }
        var drawState = meshData.drawState
        
        let startSegment: Int = drawState.lastSegmentVerticesCompleted
        let endSegment = meshData.completedSegments
        
        // Get the position from the Poisson disk sample
        let samplePoint = samplePoints[index]
        
        // Convert from sample coordinates to world coordinates (centered around origin)
        let bladeOffsetX = samplePoint.x - grassAreaWidth * 0.5
        let bladeOffsetZ = samplePoint.y - grassAreaHeight * 0.5
        
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertices = rawBytes.bindMemory(to: VertexDataWithOriginalPosition.self)
            
            if startSegment == 0 {
                // Ring 0: Base (always at origin with sample position)
                generateVerticesForRing(
                    vertices: vertices,
                    vertexOffset: meshData.vertexOffset,
                    ringIndex: 0,
                    position: SIMD3<Float>(bladeOffsetX, 0, bladeOffsetZ),
                    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 += bladeOffsetX
                    position.z += bladeOffsetZ
                    
                    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
                    )
                }
            }
            
            // Variables to track end cap information
            var hasEndRing = false
            var endPosition = SIMD3<Float>()
            var endDirection = SIMD3<Float>()
            
            // Ring for dynamic segment (constantly updating it's vertex positions)
            if let dynamicSegment = meshData.dynamicSegment,
               var currentEndPosition = meshData.getCurrentDynamicEndPosition() {
                currentEndPosition.x += bladeOffsetX
                currentEndPosition.z += bladeOffsetZ
                
                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
                )
                
                hasEndRing = true
                endPosition = currentEndPosition
                endDirection = dynamicSegment.direction
            } else if let lastSegment = meshData.segments.last {
                // Branch is complete, use the last segment's end
                hasEndRing = true
                var lastEndPosition = lastSegment.endPosition
                lastEndPosition.x += bladeOffsetX
                lastEndPosition.z += bladeOffsetZ
                endPosition = lastEndPosition
                endDirection = lastSegment.direction
            }
            
            // Generate center vertex for end cap (pointy tip)
            if hasEndRing {
                let tipExtension: Float = 0.0005 // How much to extend the tip beyond the end
                let tipPosition = endPosition + endDirection * tipExtension
                let totalRings = 1 + meshData.completedSegments + (meshData.dynamicSegment != nil ? 1 : 0)
                let centerVertexIndex = meshData.vertexOffset + totalRings * (meshData.settings.radialSegments + 1)
                
                let tipVertex = VertexDataWithOriginalPosition(
                    position: tipPosition,
                    originalPosition: tipPosition,
                    normal: endDirection,
                    uv: SIMD2<Float>(0.5, 1.0) // Center UV
                )
                vertices[centerVertexIndex] = tipVertex
            }
        }
        
        drawState.lastSegmentVerticesCompleted = endSegment
        
        let startSegmentIndices: Int = drawState.lastSegmentIndicesCompleted
        let endSegmentIndices = meshData.completedSegments+1
        
        // Just return if we aren't going to update the indices
        guard startSegmentIndices < endSegmentIndices else {
            // Make sure to stash vert changes
            self.meshParts[index].drawState = drawState
            return
        }

        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
                }
            }
            
            // Generate end cap indices (pointy tip)
            let totalRings = 1 + meshData.completedSegments + (meshData.dynamicSegment != nil ? 1 : 0)
            let hasEndCap = meshData.dynamicSegment != nil || !meshData.isGrowing
            
            if hasEndCap {
                let lastRingStart = meshData.vertexOffset + (totalRings - 1) * (meshData.settings.radialSegments + 1)
                let centerVertexIndex = meshData.vertexOffset + totalRings * (meshData.settings.radialSegments + 1)
                
                // Calculate the starting index for end cap indices
                let endCapStartIndex = indexOffset + 6 * meshData.settings.radialSegments * (totalRings - 1)
                
                for x in 0..<meshData.settings.radialSegments {
                    let a = lastRingStart + x
                    let b = lastRingStart + (x + 1) % (meshData.settings.radialSegments + 1)
                    let center = centerVertexIndex
                    
                    let triangleStartIndex = endCapStartIndex + x * 3
                    
                    // Create triangle from ring edge to center (tip)
                    indices[triangleStartIndex] = UInt32(a)
                    indices[triangleStartIndex + 1] = UInt32(center)
                    indices[triangleStartIndex + 2] = UInt32(b)
                }
            }
        }
        
        drawState.lastSegmentIndicesCompleted = endSegmentIndices

        let part = LowLevelMesh.Part(indexOffset: indexOffsetInBytes,
                                     indexCount:  meshData.indexCount,
                                     topology: meshData.settings.topology,
                                     materialIndex: meshData.materialIndex, // ← Use stored value
                                     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 = drawState
    }
    
    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<VertexDataWithOriginalPosition>,
        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] = VertexDataWithOriginalPosition(
                position: vertexPosition,
                originalPosition: vertexPosition,
                normal: normal,
                uv: uv
            )
        }
    }
}

TextureResource+loadOnlineImage.swift

import SwiftUI
import RealityKit

extension TextureResource {
    static func loadOnlineImage(_ url: URL) async throws -> TextureResource {
        let (data, _) = try await URLSession.shared.data(from: url)
        let image = UIImage(data: data)!
        let cgImage = image.cgImage!
        return try await TextureResource(image: cgImage, options: .init(semantic: nil))
    }
}

updateGrassVertices.metal

#include <metal_stdlib>
using namespace metal;

#include "VertexDataWithOriginalPosition.h"

kernel void updateGrassVertices(device const VertexDataWithOriginalPosition* inputVertices [[buffer(0)]],
                                device VertexDataWithOriginalPosition* outputVertices [[buffer(1)]],
                                constant float& windStrength [[buffer(2)]],
                                constant float& time [[buffer(3)]],
                                uint id [[thread_position_in_grid]])
{
    // Copy the input vertex
    VertexDataWithOriginalPosition inputVertex = inputVertices[id];
    VertexDataWithOriginalPosition outputVertex = inputVertex;
    float3 originalPosition = inputVertex.originalPosition;  // comes from UV1

    // Skip if this looks like uninitialized data
    if (length(originalPosition) < 0.0001) {
        outputVertices[id] = outputVertex;
        return;
    }

    // Define safe zone height - Y is up
    float safeZoneHeight = 0.01;
    float maxGrassHeight = 0.2;

    // Calculate height influence and wind displacement based on ORIGINAL position
    float heightInfluence = max(0.0, (originalPosition.y - safeZoneHeight) / (maxGrassHeight - safeZoneHeight));
    float windX = sin(time * 2.0 + originalPosition.x * 10.0 + originalPosition.y * 5.0) * windStrength * heightInfluence;
    float windZ = cos(time * 1.5 + originalPosition.y * 8.0 + originalPosition.x * 3.0) * windStrength * heightInfluence * 0.7;

    outputVertex.position.x = originalPosition.x + windX;
    outputVertex.position.z = originalPosition.z + windZ;
    outputVertex.position.y = originalPosition.y;
    outputVertex.originalPosition = originalPosition;
    outputVertices[id] = outputVertex;
}

VertexDataWithOriginalPosition.extensions.swift

import Foundation
import RealityKit

extension VertexDataWithOriginalPosition {
    static var vertexAttributes: [LowLevelMesh.Attribute] = [
        .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!),
        .init(semantic: .uv1, format: .float3, offset: MemoryLayout<Self>.offset(of: \.originalPosition)!),
        .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 = Self.vertexAttributes
        desc.vertexLayouts = Self.vertexLayouts
        desc.indexType = .uint32
        return desc
    }
    
    @MainActor static func initializeMesh(vertexCapacity: Int,
                                          indexCapacity: Int) throws -> LowLevelMesh {
        var desc = Self.descriptor
        desc.vertexCapacity = vertexCapacity
        desc.indexCapacity = indexCapacity
        return try LowLevelMesh(descriptor: desc)
    }
}

VertexDataWithOriginalPosition.h

#include <simd/simd.h>

#ifndef VertexDataWithOriginalPosition_h
#define VertexDataWithOriginalPosition_h

struct VertexDataWithOriginalPosition {
    simd_float3 position;
    simd_float3 originalPosition;
    simd_float3 normal;
    simd_float2 uv;
};

#endif /* VertexDataWithOriginalPosition_h */