RealityKit Growing/Shrinking Branch using LowLevelMesh

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)
    }
}

DirectionalBranchMeshData.swift

import RealityKit
import SwiftUI

struct DirectionalBranchMeshData {
    var settings: DirectionalBranchSettings
    
    var segments: [BranchSegment] = []
    var dynamicSegment: BranchSegment?
    var completedSegments: Int = 0
    var currentGrowthProgress: Float = 0.0
    var isGrowingUp: Bool = true
    
    init(settings: DirectionalBranchSettings = .init()) {
        self.settings = settings
        // Initialize with base segment pointing straight up
        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()
    }
    
    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)
    }
    
    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)
    }
    
    mutating func updateGrowth() {
        if isGrowingUp {
            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 {
                // Switch to shrinking mode
                isGrowingUp = false
                // Convert the last completed segment to a shrinking segment
                if !segments.isEmpty {
                    dynamicSegment = segments.removeLast()
                    completedSegments -= 1
                    currentGrowthProgress = 1.0 // Start shrinking from full size
                }
            }
        } else {
            // Shrinking mode - progress goes from 1.0 to 0.0
            currentGrowthProgress -= settings.progressRate
            
            if currentGrowthProgress <= 0.0 {
                if completedSegments > 0 {
                    // Move to the next segment to shrink
                    dynamicSegment = segments.removeLast()
                    completedSegments -= 1
                    currentGrowthProgress = 1.0 // Start shrinking this segment from full size
                } else {
                    // Done shrinking, start growing again
                    isGrowingUp = true
                    currentGrowthProgress = 0.0
                    initializeDynamicSegment() // Create new growing segment
                }
            }
        }
    }
    
    private func mix(_ a: SIMD3<Float>, _ b: SIMD3<Float>, t: Float) -> SIMD3<Float> {
        return a + (b - a) * t
    }
}

DirectionalBranchMeshView.swift

import SwiftUI
import RealityKit

struct DirectionalBranchMeshView: View {
    @State var meshData: DirectionalBranchMeshData = .init()
    @State var mesh: LowLevelMesh?
    @State var timer: Timer?
    
    init(meshData: DirectionalBranchMeshData = .init()) {
        self._meshData = State(initialValue: meshData)
    }
    
    var body: some View {
        GeometryReader3D { proxy in
            RealityView { content in
                let mesh = try! VertexData.initializeMesh(vertexCapacity: meshData.maxVertexCount,
                                                          indexCapacity: meshData.maxIndexCount)
                self.mesh = mesh
                
                let meshResource = try! await MeshResource(from: mesh)
                
                var material = PhysicallyBasedMaterial()
                material.baseColor = .init(tint: .gray)
                material.metallic = 1.0
                material.roughness = 0.25
                material.faceCulling = .none

                let entity = ModelEntity(mesh: meshResource, materials: [material])
                let extents = content.convert(proxy.frame(in: .local), from: .local, to: .scene).extents
                let yOffset = extents.y * 0.5
                entity.position = SIMD3<Float>(x: 0, y: -yOffset, z: 0)
                content.add(entity)
                
                updateMeshGeometry()
            }
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
    
    func startTimer() {
        self.timer = Timer.scheduledTimer(withTimeInterval: 0.016, repeats: true) { _ in
            meshData.updateGrowth()
            updateMeshGeometry()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    func updateMeshGeometry() {
        guard let mesh = mesh else { return }
        
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertices = rawBytes.bindMemory(to: VertexData.self)
            
            // Ring 0: Base (always at origin)
            generateVerticesForRing(
                vertices: vertices,
                ringIndex: 0,
                position: SIMD3<Float>(0, 0, 0),
                direction: SIMD3<Float>(0, 1, 0), // Always start straight up
                radius: meshData.settings.baseRadius,
                meshData: meshData
            )
            
            // Rings for completed segments
            for segmentIndex in 0..<meshData.completedSegments {
                let segment = meshData.segments[segmentIndex]
                
                generateVerticesForRing(
                    vertices: vertices,
                    ringIndex: segmentIndex + 1,
                    position: segment.endPosition,
                    direction: segment.direction,
                    radius: segment.radius,
                    meshData: meshData
                )
            }
            
            // Ring for dynamic segment
            if let dynamicSegment = meshData.dynamicSegment,
               let currentEndPosition = meshData.getCurrentDynamicEndPosition() {
                
                generateVerticesForRing(
                    vertices: vertices,
                    ringIndex: meshData.completedSegments + 1,
                    position: currentEndPosition,
                    direction: dynamicSegment.direction,
                    radius: dynamicSegment.radius,
                    meshData: meshData
                )
            }
        }
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            var index = 0

            for y in 0..<meshData.completedSegments + 1 {
                for x in 0..<meshData.settings.radialSegments {
                    let a = y * (meshData.settings.radialSegments + 1) + x
                    let b = a + 1
                    let c = a + (meshData.settings.radialSegments + 1)
                    let d = c + 1
                    
                    indices[index] = UInt32(a)
                    indices[index + 1] = UInt32(c)
                    indices[index + 2] = UInt32(b)
                    
                    indices[index + 3] = UInt32(b)
                    indices[index + 4] = UInt32(c)
                    indices[index + 5] = UInt32(d)
                    
                    index += 6
                }
            }
        }
        
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: meshData.indexCount,
                topology: meshData.settings.topology,
                bounds: meshData.boundingBox
            )
        ])
    }
    
    func generateVerticesForRing(
        vertices: UnsafeMutableBufferPointer<VertexData>,
        ringIndex: Int,
        position: SIMD3<Float>,
        direction: SIMD3<Float>,
        radius: Float,
        meshData: DirectionalBranchMeshData
    ) {
        // 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...meshData.settings.radialSegments {
            let angle = Float(x) / Float(meshData.settings.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
            let normal = normalize(localOffset)
            
            let u = Float(x) / Float(meshData.settings.radialSegments)
            let v = Float(ringIndex) / Float(meshData.settings.maxSegments)
            let uv = SIMD2<Float>(u, v)
            
            let index = ringIndex * (meshData.settings.radialSegments + 1) + x
            vertices[index] = VertexData(position: vertexPosition,
                                         normal: normal,
                                         uv: uv)
        }
    }
}

#Preview { DirectionalBranchMeshView() }

DirectionalBranchSettings.swift

import Foundation
import Metal

struct DirectionalBranchSettings {
    var baseRadius: Float = 0.03
    var segmentHeight: Float = 0.02
    var radialSegments: Int = 64
    var maxSegments: Int = 16
    var progressRate: Float = 0.2
    var directionVariation: Float = 0.25
    var taperRate: Float = 0.85
    var topology: MTLPrimitiveType = .triangle
    
    static var standard: DirectionalBranchSettings = .init()
    
    static var lineExample: DirectionalBranchSettings {
        var settings = DirectionalBranchSettings()
        settings.segmentHeight = 0.04
        settings.radialSegments = 128
        settings.maxSegments = 8
        settings.progressRate = 0.02
        settings.taperRate = 0.9
        settings.topology = .line
        return settings
    }
}

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

// In this case you could define this as a Swift struct, but I reuse this in my project with other examples that use Metal
#include <simd/simd.h>

#ifndef VertexData_h
#define VertexData_h

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

#endif /* VertexData_h */