Procedural tree RealityView with LowLevelMesh

Branch.swift

import Foundation

struct Branch {
    var segments: [BranchSegment] = []
    var dynamicSegment: BranchSegment? // the end of the branch that's currently growing
    var completedSegments: Int = 0
    var currentGrowthProgress: Float = 0.0
    var isActive: Bool = true
    var branchId: UUID = UUID()
    var parentBranchId: UUID?
    var generation: Int = 0 // 0 = main trunk, 1 = first branches, etc.
    var hasSpawnedChildren: Bool = false // Track if this branch has already spawned children
    
    // Branch-specific properties
    var baseRadius: Float
    var segmentHeight: Float
    var maxSegments: Int
    var taperRate: Float
    var directionVariation: Float
    var branchAtSegment: Int
    
    init(branchPoint: BranchPoint,
         segmentHeight: Float,
         maxSegments: Int,
         taperRate: Float,
         directionVariation: Float,
         branchAtSegment: Int,  // Branch decides when it branches
         parentBranchId: UUID? = nil,
         generation: Int = 0) {
        
        self.baseRadius = branchPoint.radius
        self.segmentHeight = segmentHeight
        self.maxSegments = maxSegments
        self.taperRate = taperRate
        self.directionVariation = directionVariation
        self.parentBranchId = parentBranchId
        self.generation = generation
        self.branchAtSegment = branchAtSegment
        
        createFirstSegment(from: branchPoint.position, direction: branchPoint.direction, radius: branchPoint.radius)
    }
    
    private mutating func createFirstSegment(from position: SIMD3<Float>, direction: SIMD3<Float>, radius: Float) {
        let endPos = position + direction * segmentHeight
        
        dynamicSegment = BranchSegment(
            startPosition: position,
            endPosition: endPos,
            radius: radius
        )
        currentGrowthProgress = 0.0
    }
    
    private mutating func createNextDynamicSegment() {
        guard let lastSegment = segments.last else {
            // This shouldn't happen in normal flow, but just in case
            return
        }
        
        let newDirection = generateNextDirection(from: lastSegment.direction)
        let startPos = lastSegment.endPosition
        let endPos = startPos + newDirection * segmentHeight
        let newRadius = lastSegment.radius * taperRate
        
        dynamicSegment = BranchSegment(
            startPosition: startPos,
            endPosition: endPos,
            radius: newRadius
        )
    }
    
    private func generateNextDirection(from currentDirection: SIMD3<Float>) -> SIMD3<Float> {
        let randomX = Float.random(in: -directionVariation...directionVariation)
        let randomZ = Float.random(in: -directionVariation...directionVariation)
        
        // Keep some upward bias, but reduce it for higher generations
        let upwardBias = max(0.1, 0.5 - Float(generation) * 0.1)
        let variation = SIMD3<Float>(randomX, upwardBias, randomZ)
        let newDirection = currentDirection + variation
        
        return normalize(newDirection)
    }
    
    func getCurrentDynamicEndPosition() -> SIMD3<Float>? {
        guard let dynamic = dynamicSegment else { return nil }
        return mix(dynamic.startPosition, dynamic.endPosition, t: currentGrowthProgress)
    }
    
    func shouldSpawnBranches(maxGenerations: Int) -> Bool {
        return !hasSpawnedChildren &&
               completedSegments == branchAtSegment &&
               generation < maxGenerations  // Global constraint
    }
    
    func getSpawnPoint() -> BranchPoint? {
        guard let dynamicSegment = dynamicSegment else { return nil }
        
        return BranchPoint(
            position: dynamicSegment.startPosition,
            direction: dynamicSegment.direction,
            radius: dynamicSegment.radius
        )
    }
    
    mutating func updateGrowth(progressRate: Float) -> Bool {
        guard isActive else { return false }
        
        if completedSegments >= maxSegments {
            // Branch is complete, stop growing
            isActive = false
            return false
        }
        
        currentGrowthProgress += progressRate
        
        if currentGrowthProgress >= 1.0 {
            // Complete the current segment
            if let segment = dynamicSegment {
                segments.append(segment)
                completedSegments += 1
                currentGrowthProgress = 0
                
                // Create next segment if we haven't reached max
                if completedSegments < maxSegments {
                    createNextDynamicSegment()
                } else {
                    // Branch is complete
                    dynamicSegment = nil
                    isActive = false
                }
            }
        }
        
        return true
    }
}

BranchingTreeMeshData.swift

import Foundation
import RealityKit

struct BranchingTreeMeshData {
    var branches: [Branch] = []
    var treeGrowthProgress: Float = 0
    
    // Is the tree still growing?
    var isActive: Bool = true
    
    // Configuration parameters
    var progressRate: Float = 0.0375                // How fast to animate growth
    var baseRadius: Float = 0.0075                  // Base radius of tree (real radius expands over time + tapers through generations)
    var treeGrowthRate: Float = 0.0002              // How fast should the radius expand
    var taperRate: Float = 0.925                    // Rate of radius reduction through the children
    var segmentHeight: Float = 0.03                 // Height between tree rings
    var radialSegments: Int = 16                    // How maybe vertices used per tree ring
    var maxSegments: Int = 9                        // Number of segments to grow to which effectively sets the branch length
    var baseDirectionVariation: Float = 0.25        // How much should branch segments change directions (increases through generations)
    var branchAtSegment: Int = 4                    // When to spawn branches (completedSegments == branchAtSegment)
    var maxGenerations: Int = 3                     // Maximum branching generations
    var branchCount: Int = 2                        // How many branches to spawn at each split
    var branchAngleSpread: Float = 0.375            // Branch angle spread (as fraction of π)
    var childRadiusScale: Float = 0.8               // Child branch radius multiplier
    var childSegmentScale: Float = 0.8              // Child branch segment height multiplier
    var childMaxSegmentReduction: Int = 0           // How much to reduce maxSegments per generation
    var generationVariationIncrease: Float = 0.25   // How much to increase directionVariation per generation
    
    init() {
        let branchPoint = BranchPoint(position: SIMD3<Float>(0, 0, 0), direction: SIMD3<Float>(0, 1, 0), radius: baseRadius)
        
        // Create the main trunk
        let mainBranch = Branch(
            branchPoint: branchPoint,
            segmentHeight: segmentHeight,
            maxSegments: maxSegments,
            taperRate: taperRate,
            directionVariation: baseDirectionVariation,
            branchAtSegment: branchAtSegment,
            generation: 0
        )
        branches.append(mainBranch)
    }
    
    var vertexCount: Int {
        var totalRings = 0
        var endCaps = 0
        
        for branch in branches {
            let hasStartRing = !branch.segments.isEmpty || branch.dynamicSegment != nil
            let totalBranchRings = (hasStartRing ? 1 : 0) + branch.completedSegments + (branch.dynamicSegment != nil ? 1 : 0)
            totalRings += totalBranchRings
            
            // Add 1 vertex for end cap center if branch has an end
            let hasEndCap = !branch.isActive || branch.dynamicSegment != nil
            if hasEndCap && totalBranchRings > 0 {
                endCaps += 1
            }
        }
        
        return (radialSegments + 1) * totalRings + endCaps
    }
    
    var indexCount: Int {
        var totalConnections = 0
        var endCapTriangles = 0
        
        for branch in branches {
            let hasStartRing = !branch.segments.isEmpty || branch.dynamicSegment != nil
            let totalRings = (hasStartRing ? 1 : 0) + branch.completedSegments + (branch.dynamicSegment != nil ? 1 : 0)
            
            if totalRings > 1 {
                totalConnections += totalRings - 1
            }
            
            // Add triangles for end cap
            let hasEndCap = !branch.isActive || branch.dynamicSegment != nil
            if hasEndCap && totalRings > 0 {
                endCapTriangles += radialSegments
            }
        }
        
        return radialSegments * totalConnections * 6 + endCapTriangles * 3
    }
    
    var boundingBox: BoundingBox {
        guard !branches.isEmpty else {
            return BoundingBox(min: SIMD3<Float>(-baseRadius, 0, -baseRadius),
                             max: SIMD3<Float>(baseRadius, segmentHeight, baseRadius))
        }
        
        var minBounds = SIMD3<Float>()
        var maxBounds = SIMD3<Float>()
        
        for branch in branches {
            for segment in branch.segments {
                let radius = segment.radius
                minBounds = min(minBounds, segment.startPosition - SIMD3<Float>(radius, radius, radius))
                minBounds = min(minBounds, segment.endPosition - SIMD3<Float>(radius, radius, radius))
                maxBounds = max(maxBounds, segment.startPosition + SIMD3<Float>(radius, radius, radius))
                maxBounds = max(maxBounds, segment.endPosition + SIMD3<Float>(radius, radius, radius))
            }
            
            if let dynamic = branch.dynamicSegment,
               let currentEndPosition = branch.getCurrentDynamicEndPosition() {
                let radius = dynamic.radius
                minBounds = min(minBounds, currentEndPosition - SIMD3<Float>(radius, radius, radius))
                maxBounds = max(maxBounds, currentEndPosition + SIMD3<Float>(radius, radius, radius))
            }
        }
        
        return BoundingBox(min: minBounds, max: maxBounds)
    }
    
    // Allocation capacities for mesh initialization
    var maxVertexCount: Int {
        let maxBranches = 50 // Reasonable upper limit
        let maxRingsPerBranch = maxSegments + 1
        return (radialSegments + 1) * maxBranches * maxRingsPerBranch
    }
    
    var maxIndexCount: Int {
        let maxBranches = 50
        return radialSegments * maxBranches * maxSegments * 6
    }
    
    func getExpandedRadius(baseRadius: Float, numberOfChildren: Int) -> Float {
        return baseRadius + baseRadius * treeGrowthProgress * Float(numberOfChildren)
    }
    
    func generateBranchSpawnData(from spawnPoint: BranchPoint) -> [BranchPoint] {
        let baseDirection = spawnPoint.direction
        let spawnPosition = spawnPoint.position
        let baseRadius = spawnPoint.radius
        
        var branchPoints: [BranchPoint] = []
        
        // Create coordinate system
        let up = baseDirection
        let tempRight = cross(up, SIMD3<Float>(0, 0, 1))
        let baseRight = length(tempRight) > 0.01 ? normalize(tempRight) : SIMD3<Float>(1, 0, 0)
        let baseForward = normalize(cross(baseRight, up))
        
        // Add random rotation around the up axis
        let randomRotation = Float.random(in: 0...(2 * Float.pi))
        let right = baseRight * cos(randomRotation) + baseForward * sin(randomRotation)
        let forward = normalize(cross(right, up))
        
        // Generate branches based on tree configuration
        for i in 0..<branchCount {
            let angle = Float.pi * branchAngleSpread * (Float(i) - Float(branchCount - 1) * 0.5)
            
            let direction: SIMD3<Float>
            if i % 2 == 0 {
                direction = normalize(up + right * tan(angle))
            } else {
                direction = normalize(up + forward * tan(angle))
            }
            
            let branchPoint = BranchPoint(
                position: spawnPosition,
                direction: direction,
                radius: baseRadius * childRadiusScale
            )
            branchPoints.append(branchPoint)
        }
        
        return branchPoints
    }
    
    mutating func updateGrowth() {
        var branchesToAdd: [Branch] = []
        
        var didGrow = false
        for i in 0..<branches.count {
            let hasChanged = branches[i].updateGrowth(progressRate: progressRate)
            if hasChanged { didGrow = true }
            
            if branches[i].shouldSpawnBranches(maxGenerations: maxGenerations) {
                if let spawnPoint = branches[i].getSpawnPoint() {
                    let branchPoints = generateBranchSpawnData(from: spawnPoint)
                    
                    for branchPoint in branchPoints {
                        let newBranch = Branch(
                            branchPoint: branchPoint,
                            segmentHeight: segmentHeight * childSegmentScale,
                            maxSegments: max(3, maxSegments - (branches[i].generation * childMaxSegmentReduction)),
                            taperRate: taperRate,
                            directionVariation: baseDirectionVariation + Float(branches[i].generation) * generationVariationIncrease,
                            branchAtSegment: 4,
                            parentBranchId: branches[i].branchId,
                            generation: branches[i].generation + 1
                        )
                        branchesToAdd.append(newBranch)
                    }
                    
                    branches[i].hasSpawnedChildren = true
                }
            }
        }
        
        branches.append(contentsOf: branchesToAdd)
        
        if didGrow {
            treeGrowthProgress += treeGrowthRate
        } else {
            isActive = false // so we can stop updating the mesh
        }
    }
}

BranchingTreeMeshView.swift

import SwiftUI
import RealityKit

struct BranchingTreeMeshView: View {
    @State private var meshData: BranchingTreeMeshData = .init()
    @State private var mesh: LowLevelMesh?
    @State private var timer: Timer?
    
    var body: some View {
        RealityView { content in
            let mesh = try! VertexData.initializeMesh(vertexCapacity: meshData.maxVertexCount,
                                                     indexCapacity: meshData.maxIndexCount)
            
            let meshResource = try! await MeshResource(from: mesh)
            
            var material = PhysicallyBasedMaterial()
            material.baseColor = .init(tint: .init(red: 0.6, green: 0.4, blue: 0.2, alpha: 1.0))
            material.metallic = 0.2
            material.roughness = 0.8
            
            let entity = ModelEntity(mesh: meshResource, materials: [material])
            entity.position.y = -0.25
            content.add(entity)
            self.mesh = mesh
            updateMeshGeometry()
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
    
    func startTimer() {
        // Start growth timer
        self.timer = Timer.scheduledTimer(withTimeInterval: 1/120, repeats: true) { _ in
            meshData.updateGrowth()
            updateMeshGeometry()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    private func updateMeshGeometry() {
        guard let mesh = mesh, meshData.isActive else { return }
        
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertices = rawBytes.bindMemory(to: VertexData.self)
            var vertexIndex = 0
            
            for branch in meshData.branches {
                // Generate start ring
                var startPosition: SIMD3<Float>
                var startDirection: SIMD3<Float>
                var startRadius: Float
                
                if let firstSegment = branch.segments.first {
                    startPosition = firstSegment.startPosition
                    startDirection = firstSegment.direction
                    startRadius = firstSegment.radius
                } else if let dynamicSegment = branch.dynamicSegment {
                    startPosition = dynamicSegment.startPosition
                    startDirection = dynamicSegment.direction
                    startRadius = dynamicSegment.radius
                } else {
                    continue
                }
                
                startRadius = meshData.getExpandedRadius(baseRadius: startRadius,
                                                         numberOfChildren: branch.segments.count)
                
                generateVerticesForRing(
                    vertices: vertices,
                    startVertexIndex: vertexIndex,
                    position: startPosition,
                    direction: startDirection,
                    radius: startRadius,
                    ringIndex: 0,
                    branch: branch,
                    meshData: meshData
                )
                vertexIndex += meshData.radialSegments + 1
                
                // Generate rings for completed segments
                for (segmentIndex, segment) in branch.segments.enumerated() {
                    let numberOfChildren = branch.segments.count - segmentIndex
                    let expandedRadius = meshData.getExpandedRadius(baseRadius: segment.radius,
                                                                    numberOfChildren: numberOfChildren)
                    
                    generateVerticesForRing(
                        vertices: vertices,
                        startVertexIndex: vertexIndex,
                        position: segment.endPosition,
                        direction: segment.direction,
                        radius: expandedRadius,
                        ringIndex: segmentIndex + 1,
                        branch: branch,
                        meshData: meshData
                    )
                    vertexIndex += meshData.radialSegments + 1
                }
                
                // Generate ring for dynamic segment
                var hasEndRing = false
                var endPosition: SIMD3<Float> = SIMD3<Float>()
                var endDirection: SIMD3<Float> = SIMD3<Float>()
                
                if let dynamicSegment = branch.dynamicSegment,
                   let currentEndPosition = branch.getCurrentDynamicEndPosition() {
                    
                    generateVerticesForRing(
                        vertices: vertices,
                        startVertexIndex: vertexIndex,
                        position: currentEndPosition,
                        direction: dynamicSegment.direction,
                        radius: dynamicSegment.radius,
                        ringIndex: branch.completedSegments + 1,
                        branch: branch,
                        meshData: meshData
                    )
                    vertexIndex += meshData.radialSegments + 1
                    hasEndRing = true
                    endPosition = currentEndPosition
                    endDirection = dynamicSegment.direction
                } else if let lastSegment = branch.segments.last {
                    // Branch is complete, use the last segment's end
                    hasEndRing = true
                    endPosition = lastSegment.endPosition
                    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 tipVertex = VertexData(
                        position: tipPosition,
                        normal: endDirection,
                        uv: SIMD2<Float>(0.5, 1.0) // Center UV
                    )
                    vertices[vertexIndex] = tipVertex
                    vertexIndex += 1
                }
            }
        }
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            var indexOffset = 0
            var vertexOffset = 0
            
            for branch in meshData.branches {
                let totalRings = 1 + branch.completedSegments + (branch.dynamicSegment != nil ? 1 : 0)
                let hasEndCap = !branch.isActive || branch.dynamicSegment != nil
                
                // Generate indices between consecutive rings
                for y in 0..<totalRings - 1 {
                    for x in 0..<meshData.radialSegments {
                        let a = vertexOffset + y * (meshData.radialSegments + 1) + x
                        let b = a + 1
                        let c = a + (meshData.radialSegments + 1)
                        let d = c + 1
                        
                        indices[indexOffset] = UInt32(a)
                        indices[indexOffset + 1] = UInt32(c)
                        indices[indexOffset + 2] = UInt32(b)
                        
                        indices[indexOffset + 3] = UInt32(b)
                        indices[indexOffset + 4] = UInt32(c)
                        indices[indexOffset + 5] = UInt32(d)
                        
                        indexOffset += 6
                    }
                }
                
                // Generate end cap indices (pointy tip)
                if hasEndCap {
                    let lastRingStart = vertexOffset + (totalRings - 1) * (meshData.radialSegments + 1)
                    let centerVertexIndex = vertexOffset + totalRings * (meshData.radialSegments + 1)
                    
                    for x in 0..<meshData.radialSegments {
                        let a = lastRingStart + x
                        let b = lastRingStart + (x + 1) % (meshData.radialSegments + 1)
                        let center = centerVertexIndex
                        
                        // Create triangle from ring edge to center (tip)
                        indices[indexOffset] = UInt32(a)
                        indices[indexOffset + 1] = UInt32(center)
                        indices[indexOffset + 2] = UInt32(b)
                        
                        indexOffset += 3
                    }
                    
                    vertexOffset += totalRings * (meshData.radialSegments + 1) + 1 // +1 for center vertex
                } else {
                    vertexOffset += totalRings * (meshData.radialSegments + 1)
                }
            }
        }
        
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: meshData.indexCount,
                topology: .triangle,
                bounds: meshData.boundingBox
            )
        ])
    }
    
    private func generateVerticesForRing(
        vertices: UnsafeMutableBufferPointer<VertexData>,
        startVertexIndex: Int,
        position: SIMD3<Float>,
        direction: SIMD3<Float>,
        radius: Float,
        ringIndex: Int,
        branch: Branch,
        meshData: BranchingTreeMeshData
    ) {
        // Create a coordinate system for this ring
        let up = direction
        let tempRight = cross(up, SIMD3<Float>(0, 0, 1))
        let right = length(tempRight) > 0.01 ? normalize(tempRight) : SIMD3<Float>(1, 0, 0)
        let forward = normalize(cross(right, up))
        
        let v = Float(ringIndex) / Float(branch.maxSegments) * 1.0
        
        // Generate vertices around the circumference
        for x in 0...meshData.radialSegments {
            let angle = Float(x) / Float(meshData.radialSegments) * 2 * Float.pi
            let localX = cos(angle) * radius
            let localZ = sin(angle) * radius
            
            // Transform local coordinates to world space
            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(meshData.radialSegments)
            let uv = SIMD2<Float>(u, v)
            
            let index = startVertexIndex + x
            vertices[index] = VertexData(position: vertexPosition,
                                         normal: normal,
                                         uv: uv)
        }
    }
}

#Preview {
    BranchingTreeMeshView()
}

BranchPoint.swift

import Foundation

struct BranchPoint {
    var position: SIMD3<Float>
    var direction: SIMD3<Float>
    var radius: Float
    
    init(position: SIMD3<Float>, direction: SIMD3<Float>, radius: Float) {
        self.position = position
        self.direction = direction
        self.radius = radius
    }
    
    // Convenience initializer from a segment's end point
    init(from segment: BranchSegment) {
        self.position = segment.endPosition
        self.direction = segment.direction
        self.radius = segment.radius
    }
    
    // Scale the radius (useful for child branches)
    func scaled(by factor: Float) -> BranchPoint {
        return BranchPoint(
            position: position,
            direction: direction,
            radius: radius * factor
        )
    }
}

BranchSegment.swift

import Foundation

struct BranchSegment {
    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)
    }
}

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