RealityKit burn fade sphere: LowLevelMesh, ShaderGraphMaterial, and compute shader

BurnFadeParams.h

#ifndef BurnFadeParams_h
#define BurnFadeParams_h

struct BurnFadeParams {
    float progress;
    float scale;
    float hueRotate;
    float edgeWidth;
    float emberRange;
};


#endif /* BurnFadeParams_h */

burnFadeVertices.metal

#include <metal_stdlib>
using namespace metal;

#include "HueRotate.metal"
#include "NoiseUtilities.metal"
#include "VertexDataWith4ChannelColor.h"
#include "BurnFadeParams.h"

kernel void burnFadeVertices(device const VertexDataWith4ChannelColor* inputVertices [[buffer(0)]],
                             device VertexDataWith4ChannelColor* outputVertices [[buffer(1)]],
                             constant BurnFadeParams& params [[buffer(2)]],
                             uint vid [[thread_position_in_grid]])
{
    VertexDataWith4ChannelColor inputVertex = inputVertices[vid];
    VertexDataWith4ChannelColor outputVertex = inputVertex;

    float3 scaledPosition = inputVertex.position * params.scale;
    
    // Create noise pattern that determines burn priority
    float primaryNoise = fractalNoise(scaledPosition);
    float secondaryNoise = fractalNoise(scaledPosition * 2.3 + float3(100.0, 50.0, 75.0));
    float detailNoise = noise(scaledPosition * 12.0);
    
    // Combine noise layers - this determines which areas burn first/last
    float burnPriority = primaryNoise * 0.6 + secondaryNoise * 0.3 + detailNoise * 0.1;
    
    // Map burnAmount to affect the full range
    float effectiveBurnAmount = params.progress;
    
    // Calculate how far this vertex is from being burned away
    float burnDistance = burnPriority - effectiveBurnAmount;
    
    // Create smooth transition for alpha
    float edgeWidth = params.edgeWidth;
    float alpha = smoothstep(-edgeWidth, edgeWidth, burnDistance);
    
    // Calculate ember/fire effect at burn edges
    float emberIntensity = 0.0;
    float emberRange = params.emberRange; // How far from the burn edge we show embers
    
    if (burnDistance > -emberRange && burnDistance < emberRange) {
        // We're near the burn edge - calculate ember intensity
        float distanceFromEdge = abs(burnDistance);
        emberIntensity = 1.0 - (distanceFromEdge / emberRange);
        emberIntensity = smoothstep(0.0, 1.0, emberIntensity);
        
        // Add flickering with noise
        float emberFlicker = noise(scaledPosition * 25.0 + float3(params.progress * 10.0, 0.0, 0.0));
        emberIntensity *= 0.7 + 0.3 * emberFlicker;
    }
    
    // Define fire colors
    float3 originalColor = float3(inputVertex.color.x, inputVertex.color.y, inputVertex.color.z); //float3(0,0,0);       // no color (since we add to
    float3 emberColor = float3(1.0, 0.4, 0.1);  // Bright orange
    float3 fireColor = float3(1.0, 0.2, 0.0);   // Red-orange
    float3 hotColor = float3(1.0, 0.8, 0.2);    // Yellow-white hot
    
    // Mix between different fire colors based on intensity
    float3 burnColor;
    if (emberIntensity < 0.3) {
        burnColor = mix(emberColor, fireColor, emberIntensity / 0.3);
    } else if (emberIntensity < 0.7) {
        burnColor = mix(fireColor, hotColor, (emberIntensity - 0.3) / 0.4);
    } else {
        burnColor = hotColor;
    }
    
    burnColor = rotateHue(burnColor, params.hueRotate);
    
    // Blend original color with burn effect
    float3 finalColor = mix(originalColor, burnColor, emberIntensity);
    
    // Make burn edges extra bright
    finalColor += emberIntensity * emberIntensity * float3(0.2, 0.2, 0.2);
    
    alpha = clamp(alpha, 0.0, 1.0);
    
    outputVertex.color.rgb = finalColor;
    outputVertex.color.w = alpha;
    outputVertices[vid] = outputVertex;
}

BurningSphereView.swift

import SwiftUI
import RealityKit
import Metal

struct BurnFadeSettings: Equatable {
    var burnAmount: Float = 0.0
    var burnScale: Float = 8.0
    var hueRotate: Float = 0.0
    var edgeWidth: Float = 0.08
    var emberRange: Float = 0.15
}

struct BurningSphereView: View {
    @State var mesh: LowLevelMesh?
    @State var burnSettings: BurnFadeSettings = .init()
    
    // Store original vertices to preserve color
    @State var originalVertices: [VertexDataWith4ChannelColor] = []
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipeline: MTLComputePipelineState
    
    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let function = library.makeFunction(name: "burnFadeVertices")!
        self.computePipeline = try! device.makeComputePipelineState(function: function)
    }
    
    var body: some View {
        VStack {
            RealityView { content in
                let mesh = try! generateIcosphereMesh(subdivisions: 5)
                let meshResource = try! await MeshResource(from: mesh)
                
                let material = try! await loadMaterial()
                
                let entity = ModelEntity(mesh: meshResource, materials: [material])
                content.add(entity)
                self.mesh = mesh
                
                storeOriginalVertices(from: mesh)
            }
            .onChange(of: burnSettings) { _, _ in
                updateMesh()
            }
            
            VStack {
                HStack {
                    Text("Burn Amount: \(burnSettings.burnAmount, specifier: "%.2f")")
                    Spacer()
                    Slider(value: $burnSettings.burnAmount, in: 0...1)
                        .frame(width: 300)
                }
                
                HStack {
                    Text("Burn Scale: \(burnSettings.burnScale, specifier: "%.2f")")
                    Spacer()
                    Slider(value: $burnSettings.burnScale, in: 1...40)
                        .frame(width: 300)
                }
                
                HStack {
                    Text("Burn Hue: \(burnSettings.hueRotate, specifier: "%.2f")")
                    Spacer()
                    Slider(value: $burnSettings.hueRotate, in: 0...(.pi*2))
                        .frame(width: 300)
                }
                HStack {
                    Text("Edge Width: \(burnSettings.edgeWidth, specifier: "%.3f")")
                    Spacer()
                    Slider(value: $burnSettings.edgeWidth, in: 0.01...0.2)
                        .frame(width: 300)
                }
                HStack {
                    Text("Ember Range: \(burnSettings.emberRange, specifier: "%.3f")")
                    Spacer()
                    Slider(value: $burnSettings.emberRange, in: 0.05...0.4)
                        .frame(width: 300)
                }
            }
            .frame(width: 500)
            .padding()
            .glassBackgroundEffect()
        }
    }
    
    func storeOriginalVertices(from mesh: LowLevelMesh) {
        let vertexCount = mesh.vertexCapacity
        originalVertices = Array(repeating: VertexDataWith4ChannelColor(
            position: SIMD3<Float>(0, 0, 0),
            normal: SIMD3<Float>(0, 0, 0),
            uv: SIMD2<Float>(0, 0),
            color: SIMD4<Float>(0, 0, 0, 0)
        ), count: vertexCount)
        
        mesh.withUnsafeBytes(bufferIndex: 0) { rawBytes in
            let vertexBuffer = rawBytes.bindMemory(to: VertexDataWith4ChannelColor.self)
            for i in 0..<vertexCount {
                originalVertices[i] = vertexBuffer[i]
            }
        }
    }
    
    func updateMesh() {
        guard let mesh = mesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder(),
              !originalVertices.isEmpty else { return }

        let originalVertexBuffer = device.makeBuffer(
            bytes: originalVertices,
            length: originalVertices.count * MemoryLayout<VertexDataWith4ChannelColor>.size,
            options: .storageModeShared
        )
        
        let newVertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipeline)
        computeEncoder.setBuffer(originalVertexBuffer, offset: 0, index: 0)  // Read from original
        computeEncoder.setBuffer(newVertexBuffer, offset: 0, index: 1)       // Write to new
        var params = BurnFadeParams(
            progress: burnSettings.burnAmount,
            scale: burnSettings.burnScale,
            hueRotate: burnSettings.hueRotate,
            edgeWidth: burnSettings.edgeWidth,
            emberRange: burnSettings.emberRange
        )
        computeEncoder.setBytes(&params, length: MemoryLayout<BurnFadeParams>.size, index: 2)
        
        // Use vertex count, not index count
        let vertexCount = mesh.vertexCapacity
        let threadsPerGrid = MTLSize(width: vertexCount, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        
        computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
}

#Preview {
    BurningSphereView()
}

// MARK: Get Online Shader Graph Material
extension BurningSphereView {
    func loadMaterial() async throws -> ShaderGraphMaterial {
        let baseURL = URL(string: "https://matt54.github.io/Resources/")!
        let fullURL = baseURL.appendingPathComponent("TextureCoordinatesColorMaterial.usda")
        let data = try Data(contentsOf: fullURL)
        let materialFilenameWithPath: String = "/Root/TextureCoordinatesColorMaterial"
        return try await ShaderGraphMaterial(named: materialFilenameWithPath, from: data)
    }
}

extension BurningSphereView {
    func generateIcosphereMesh(radius: Float = 0.1, subdivisions: Int) throws -> LowLevelMesh {
        // Define initial icosahedron vertices
        let t: Float = (1.0 + sqrt(5.0)) / 2.0
        
        var vertices: [SIMD3<Float>] = [
            SIMD3<Float>(-1, t, 0),
            SIMD3<Float>(1, t, 0),
            SIMD3<Float>(-1, -t, 0),
            SIMD3<Float>(1, -t, 0),
            
            SIMD3<Float>(0, -1, t),
            SIMD3<Float>(0, 1, t),
            SIMD3<Float>(0, -1, -t),
            SIMD3<Float>(0, 1, -t),
            
            SIMD3<Float>(t, 0, -1),
            SIMD3<Float>(t, 0, 1),
            SIMD3<Float>(-t, 0, -1),
            SIMD3<Float>(-t, 0, 1)
        ].map { normalize($0) * radius }
        
        var faces: [(Int, Int, Int)] = [
            (0, 11, 5), (0, 5, 1), (0, 1, 7), (0, 7, 10), (0, 10, 11),
            (1, 5, 9), (5, 11, 4), (11, 10, 2), (10, 7, 6), (7, 1, 8),
            (3, 9, 4), (3, 4, 2), (3, 2, 6), (3, 6, 8), (3, 8, 9),
            (4, 9, 5), (2, 4, 11), (6, 2, 10), (8, 6, 7), (9, 8, 1)
        ]
        
        var midpointCache: [Int: Int] = [:]
        func midpoint(_ v1: Int, _ v2: Int) -> Int {
            let key = v1 < v2 ? (v1 << 16) | v2 : (v2 << 16) | v1
            if let mid = midpointCache[key] {
                return mid
            }
            let midPoint = normalize((vertices[v1] + vertices[v2]) * 0.5) * radius
            
            vertices.append(midPoint)
            let index = vertices.count - 1
            midpointCache[key] = index
            return index
        }
        
        for _ in 0..<subdivisions {
            var newFaces: [(Int, Int, Int)] = []
            for (v1, v2, v3) in faces {
                let a = midpoint(v1, v2)
                let b = midpoint(v2, v3)
                let c = midpoint(v3, v1)
                newFaces.append((v1, a, c))
                newFaces.append((v2, b, a))
                newFaces.append((v3, c, b))
                newFaces.append((a, b, c))
            }
            faces = newFaces
        }
        
        let vertexCount = vertices.count
        let indexCount = faces.count * 3
        
        var desc = VertexDataWith4ChannelColor.descriptor
        desc.vertexCapacity = vertexCount
        desc.indexCapacity = indexCount
        
        let mesh = try LowLevelMesh(descriptor: desc)
        
        mesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertexBuffer = rawBytes.bindMemory(to: VertexDataWith4ChannelColor.self)
            for (i, position) in vertices.enumerated() {
                let normal = normalize(position)
                
                let u = (atan2(position.z, position.x) + Float.pi) / (2.0 * Float.pi)
                let v = (asin(position.y / radius) + Float.pi / 2.0) / Float.pi
                
                vertexBuffer[i] = VertexDataWith4ChannelColor(
                    position: position,
                    normal: normal,
                    uv: SIMD2<Float>(u, v),
                    color: SIMD4<Float>(0,0.0,0,1.0)
                )
            }
        }
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indexBuffer = rawIndices.bindMemory(to: UInt32.self)
            var index = 0
            for (v1, v2, v3) in faces {
                indexBuffer[index] = UInt32(v1)
                indexBuffer[index + 1] = UInt32(v2)
                indexBuffer[index + 2] = UInt32(v3)
                index += 3
            }
        }
        
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: indexCount,
                topology: .triangle,
                bounds: BoundingBox(min: [-radius, -radius, -radius], max: [radius, radius, radius])
            )
        ])
        
        return mesh
    }
}

HueRotate.metal

#include <metal_stdlib>
using namespace metal;

inline float3 rotateHue(float3 color, float hueRotation) {
    const float3 k = float3(0.57735, 0.57735, 0.57735);
    float cosAngle = cos(hueRotation);
    return color * cosAngle + cross(k, color) * sin(hueRotation) + k * dot(k, color) * (1.0 - cosAngle);
}

NoiseUtilities.metal

#include <metal_stdlib>
using namespace metal;

inline float hash(float3 p) {
    p = fract(p * 0.3183099 + 0.1);
    p *= 17.0;
    return fract(p.x * p.y * p.z * (p.x + p.y + p.z));
}

inline float noise(float3 p) {
    float3 i = floor(p);
    float3 f = fract(p);
    f = f * f * (3.0 - 2.0 * f);
    
    return mix(mix(mix(hash(i + float3(0,0,0)), hash(i + float3(1,0,0)), f.x),
                   mix(hash(i + float3(0,1,0)), hash(i + float3(1,1,0)), f.x), f.y),
               mix(mix(hash(i + float3(0,0,1)), hash(i + float3(1,0,1)), f.x),
                   mix(hash(i + float3(0,1,1)), hash(i + float3(1,1,1)), f.x), f.y), f.z);
}

inline float fractalNoise(float3 p) {
    float value = 0.0;
    float amplitude = 0.5;
    float frequency = 1.0;
    float maxValue = 0.0;
    
    for (int i = 0; i < 4; i++) {
        value += amplitude * noise(p * frequency);
        maxValue += amplitude;
        amplitude *= 0.5;
        frequency *= 2.0;
    }
    
    return value / maxValue;
}

VertexDataWith4ChannelColor.extensions.swift

import Foundation
import RealityKit

extension VertexDataWith4ChannelColor {
    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)!),
        .init(semantic: .uv2, format: .float4, offset: MemoryLayout<Self>.offset(of: \.color)!)
    ]
    
    static var vertexLayouts: [LowLevelMesh.Layout] = [
        .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride)
    ]

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

VertexDataWith4ChannelColor.h

#include <simd/simd.h>

#ifndef VertexDataWith4ChannelColor_h
#define VertexDataWith4ChannelColor_h

struct VertexDataWith4ChannelColor {
    simd_float3 position;
    simd_float3 normal;
    simd_float2 uv;
    simd_float4 color; // rgb + alpha
};

#endif /* VertexDataWith4ChannelColor_h */