4th Of July FlagView with LowLevelMesh, Dynamic Lights, Particles, Audio FX, and more!

ExampleAudioFXFiles.swift

import Foundation

enum ExampleAudioFXFiles {
    case cannonShot
    case fireworkExplode
    
    static let baseURL = URL(string: "https://matt54.github.io/Resources/")!
    var url: URL {
        return ExampleAudioFXFiles.baseURL.appendingPathComponent( "\(filenameWithExtension)" )
    }
    
    var filenameWithExtension: String {
        switch self {
        case .cannonShot:
            return "cannon_shot.wav"
        case .fireworkExplode:
            return "firework_explode.wav"
        }
    }
}

FireWorkEntity.swift

import Foundation
import RealityKit

struct FireworkDefinition {
    var startingPosition: SIMD3<Float> = SIMD3<Float>(0, -0.08, 0.15)
    var velocity: SIMD3<Float> = SIMD3<Float>(Float.random(in: -0.0015...0.0015),
                                              Float.random(in: 0.009...0.013),
                                              -0.0075)
    var timeToLive: Float = 0.3
    
    static let maxLightIntensity: Float = 10_000
    static let maxLightRadius: Float = 0.5
    static let acceleration: SIMD3<Float> = SIMD3<Float>(0.0, -0.00055, 0)
}

class FireWorkEntity: Entity {
    var fireworkDefinition: FireworkDefinition
    var litAmount: Float = 0
    var hasExploded: Bool = false
    var modelEntity: ModelEntity?
    
    let lifeRange: ClosedRange<Float> = 0.23...0.3
    
    var intensity: Float {
        FireworkDefinition.maxLightIntensity * litAmount
    }
    
    var attentionRadius: Float {
        FireworkDefinition.maxLightRadius * litAmount
    }
    
    required init() {
        fireworkDefinition = .init(timeToLive: Float.random(in: lifeRange))
        super.init()
        position = fireworkDefinition.startingPosition
        setModelComponent()
        setupParticles()
    }
    
    func setModelComponent() {
        let mesh = MeshResource.generateSphere(radius: 0.002)
        let material = UnlitMaterial(color: .white)
        let sphereModel = ModelEntity(mesh: mesh, materials: [material])
        sphereModel.position = .init(x: 0, y: 0, z: 0.03)
        addChild(sphereModel)
        self.modelEntity = sphereModel
    }
    
    func update() {
        if hasExploded {
            litAmount -= 0.01
            setLightComponent()
        } else {
            fireworkDefinition.velocity += FireworkDefinition.acceleration
            position += fireworkDefinition.velocity
            fireworkDefinition.timeToLive -= 0.01
            if fireworkDefinition.timeToLive <= 0 {
                explode()
            }
        }
    }
    
    func explode() {
        hasExploded = true
        litAmount = 1.0
        burstParticles()
    }
    
    func setupParticles() {
        var particleEmitter = ParticleEmitterComponent.Presets.fireworks
        particleEmitter.mainEmitter.lifeSpan = 0
        particleEmitter.isEmitting = false
        particleEmitter.spawnOccasion = .onBirth
        particleEmitter.emitterShape = .point
        particleEmitter.emitterShapeSize = SIMD3<Float>(0.0, -0.1, 0.1)
        particleEmitter.birthLocation = .surface
        
        components.set(particleEmitter)
    }
    
    func burstParticles() {
        if let modelEntity {
            removeChild(modelEntity)
        }
        
        if var particleEmitter = components[ParticleEmitterComponent.self] {
            particleEmitter.burstCount = 1
            particleEmitter.burst()
            components.set(particleEmitter)
        }
    }
    
    func setLightComponent() {
        let pointLightComponent = PointLightComponent(
            cgColor: .init(red: 1, green: 1, blue: 1, alpha: 1),
            intensity: intensity,
            attenuationRadius: attentionRadius
        )
        components.set(pointLightComponent)
    }
}

FlagParams.h

#include <simd/simd.h>

#ifndef FlagParams_h
#define FlagParams_h

struct FlagParams {
    simd_float2 size;
    simd_uint2 dimensions;
    float time;
    float windStrength;
};

#endif /* FlagParams_h */

FlagView.swift

import Metal
import RealityKit
import SwiftUI

struct FlagView: View {
    @State var rootEntity = Entity()
    @State private var flagMesh: LowLevelMesh?
    @State var flagAnimationTime: Float = 0
    @State private var timer: Timer?
    @State private var fireworkEntities = [FireWorkEntity]()
    @State var cannonRecoilAmount: Float = 0
    @State var cannonAudioResource: AudioFileResource?
    @State var fireworkAudioResource: AudioFileResource?
    
    var meshData: PlaneMeshData
    var flagAnimationRate: Float = 0.01
    let windStrength: Float = 1.0
    let maxCannonRecoilAmount: Float = 0.025
    let cannonParentEntity = Entity()
    let cannonRecoilRecoveryRate: Float = 0.05
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipeline: MTLComputePipelineState
    
    let flagAndTextZOffset: Float = -0.15
    let flagAndTextYOffset: Float = 0.025
    
    init(meshData: PlaneMeshData = .init(size: [0.19, 0.1], dimensions: [240, 160])) {
        self.meshData = meshData
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "updateFlagMesh")!
        self.computePipeline = try! device.makeComputePipelineState(function: updateFunction)
    }
    
    var body: some View {
        RealityView { content in
            await createAndAddFlag()
            await createAndAddCannon()
            await loadFireworkAudioResources()
            let textEntity = ForthOfJulyTextEntity()
            textEntity.position.z = flagAndTextZOffset
            textEntity.position.y += flagAndTextYOffset
            rootEntity.addChild(textEntity)
            content.add(rootEntity)
            updateMeshGeometry()
        }
        .gesture(
            SpatialTapGesture()
                .targetedToEntity(cannonParentEntity)
                .onEnded { _ in
                    fireCannon()
                }
        )
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
}

#Preview {
    FlagView()
}

// MARK: Animation Loop
extension FlagView {
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: 1/120.0, repeats: true) { _ in
            updateMeshGeometry()
            
            // Update cannon recoil
            if cannonRecoilAmount > 0 {
                cannonRecoilAmount -= cannonRecoilRecoveryRate
                cannonRecoilAmount = max(0, cannonRecoilAmount)
                cannonParentEntity.position.z = cannonRecoilAmount * maxCannonRecoilAmount
            }
            
            for fireworkEntity in fireworkEntities {
                fireworkEntity.update()
                
                // notice we use a negative value here, so the point light is long gone but we want to give the audio and particles some time to finish
                if fireworkEntity.hasExploded && fireworkEntity.litAmount <= -1.5 {
                    rootEntity.removeChild(fireworkEntity)
                }
                
                if fireworkEntity.litAmount == 1 {
                    if let fireworkAudioResource {
                        fireworkEntity.playAudio(fireworkAudioResource)
                    }
                }
            }
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
}

// MARK: Create and Update Flag Mesh
extension FlagView {
    func createAndAddFlag() async {
        let mesh = try! createMesh()
        let resource = try! await MeshResource(from: mesh)
        var material = PhysicallyBasedMaterial()
        let url = URL(string: "https://matt54.github.io/Resources/Flag_of_the_United_States.png")!
        let textureResource = try! await TextureResource.loadOnlineImage(url)
        material.baseColor = .init(texture: .init(textureResource))
        material.metallic = 1.0
        material.roughness = 1.0
        material.emissiveColor = .init(texture: .init(textureResource))
        material.emissiveIntensity = 1.0
        let modelComponent = ModelComponent(mesh: resource, materials: [material])
        
        let flagEntity = Entity()
        flagEntity.components.set(modelComponent)
        flagEntity.scale *= 1.5
        flagEntity.position.z = flagAndTextZOffset
        flagEntity.position.y += flagAndTextYOffset
        rootEntity.addChild(flagEntity)
        self.flagMesh = mesh
    }
    
    func createMesh() throws -> LowLevelMesh {
        let mesh = try! VertexData.initializeMesh(vertexCapacity: meshData.vertexCount,
                                                  indexCapacity: meshData.indexCount)
        
        // Set up indices once - they never change
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            var indexOffset = 0
            
            for y in 0..<(meshData.dimensions.y - 1) {
                for x in 0..<(meshData.dimensions.x - 1) {
                    let bottomLeft = UInt32(y * meshData.dimensions.x + x)
                    let bottomRight = bottomLeft + 1
                    let topLeft = bottomLeft + UInt32(meshData.dimensions.x)
                    let topRight = topLeft + 1
                    
                    // First triangle
                    indices[indexOffset] = bottomLeft
                    indices[indexOffset + 1] = bottomRight
                    indices[indexOffset + 2] = topLeft
                    
                    // Second triangle
                    indices[indexOffset + 3] = topLeft
                    indices[indexOffset + 4] = bottomRight
                    indices[indexOffset + 5] = topRight
                    
                    indexOffset += 6
                }
            }
        }
        
        return mesh
    }
    
    func updateMeshGeometry() {
        guard let mesh = flagMesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        flagAnimationTime += flagAnimationRate
        var params = FlagParams(
            size: meshData.size,
            dimensions: meshData.dimensions,
            time: flagAnimationTime,
            windStrength: windStrength
        )
        
        let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipeline)
        computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0)
        computeEncoder.setBytes(&params, length: MemoryLayout<FlagParams>.stride, index: 1)
        
        let threadgroupSize = MTLSize(width: 8, height: 8, depth: 1)
        let threadgroups = MTLSize(
            width: (Int(meshData.dimensions.x) + threadgroupSize.width - 1) / threadgroupSize.width,
            height: (Int(meshData.dimensions.y) + threadgroupSize.height - 1) / threadgroupSize.height,
            depth: 1
        )
        
        computeEncoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        computeEncoder.endEncoding()
        
        commandBuffer.commit()

        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: meshData.indexCount,
                topology: .triangle,
                bounds: meshData.boundingBox
            )
        ])
    }
}

// MARK: Cannon
extension FlagView {
    func createAndAddCannon() async {
        let url = URL(string: "https://matt54.github.io/Resources/cannon.usdz")!
        let (downloadedURL, _) = try! await URLSession.shared.download(from: url)
        let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
        let destinationURL = documentsDirectory.appendingPathComponent("downloadedModel.usdz")
        if FileManager.default.fileExists(atPath: destinationURL.path) {
            try! FileManager.default.removeItem(at: destinationURL)
        }
        try! FileManager.default.moveItem(at: downloadedURL, to: destinationURL)
        let cannon = try! await ModelEntity.init(contentsOf: destinationURL)
        try! FileManager.default.removeItem(at: destinationURL)
        cannon.transform.rotation = .init(angle: .pi, axis: .init(x: 0, y: 1, z: 0))
        cannon.scale *= 0.13
        cannon.position = .init(x: 0, y: -0.1, z: 0.22)
        
        cannon.components.set(InputTargetComponent(allowedInputTypes: .all))
        
        if let extents = cannon.model?.mesh.bounds.extents {
            cannon.components.set(CollisionComponent(shapes: [ShapeResource.generateBox(size: extents)], isStatic: true))
        }
        cannonParentEntity.addChild(cannon)
        rootEntity.addChild(cannonParentEntity)
        await loadCannonAudioResource()
    }
    
    func fireCannon() {
        let fireworkEntity = FireWorkEntity()
        rootEntity.addChild(fireworkEntity)
        fireworkEntities.append(fireworkEntity)
        cannonRecoilAmount = 1.0
        if let cannonAudioResource {
            cannonParentEntity.playAudio(cannonAudioResource)
        }
    }
}

// MARK: Audio Resource Loading
extension FlagView {
    func loadCannonAudioResource() async {
        let (downloadedURL, _) = try! await URLSession.shared.download(from: ExampleAudioFXFiles.cannonShot.url)
        let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
        let destinationURL = documentsDirectory.appendingPathComponent("cannon_shot.wav")
        if FileManager.default.fileExists(atPath: destinationURL.path) {
            try! FileManager.default.removeItem(at: destinationURL)
        }
        try! FileManager.default.moveItem(at: downloadedURL, to: destinationURL)
        cannonAudioResource = try! await AudioFileResource(contentsOf: destinationURL)
        try! FileManager.default.removeItem(at: destinationURL)
    }
    
    func loadFireworkAudioResources() async {
        let (downloadedURL, _) = try! await URLSession.shared.download(from: ExampleAudioFXFiles.fireworkExplode.url)
        let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
        let destinationURL = documentsDirectory.appendingPathComponent("firework_explosion.wav")
        if FileManager.default.fileExists(atPath: destinationURL.path) {
            try! FileManager.default.removeItem(at: destinationURL)
        }
        try! FileManager.default.moveItem(at: downloadedURL, to: destinationURL)
        fireworkAudioResource = try! await AudioFileResource(contentsOf: destinationURL)
        try! FileManager.default.removeItem(at: destinationURL)
    }
}

ForthOfJulyTextEntity.swift

import SwiftUI
import RealityKit

class ForthOfJulyTextEntity: Entity {
    required init() {
        super.init()
        try! setModelComponent()
        
        position = .init(x: 0, y: -0.1, z: 0)
        centerTextEntity()
        
    }
    
    func centerTextEntity() {
        if let modelComponent = components[ModelComponent.self] {
            let bounds = modelComponent.mesh.bounds
            let minX = bounds.min.x
            let maxX = bounds.max.x
            let minY = bounds.min.y
            let maxY = bounds.max.y
            transform.translation.x -= (maxX-minX) * 0.5
            transform.translation.y -= (maxY-minY) * 0.5
        }
    }
    
    func setModelComponent() throws {
        let resource = try getMeshResource()
        
        var material = PhysicallyBasedMaterial()
        material.baseColor.tint = .init(red: 1.0, green: 1.0, blue: 1.0, alpha: 1.0)
        
        var material2 = UnlitMaterial()
        material2.color.tint = .init(red: 0, green: 0.157, blue: 0.408, alpha: 1.0)

        let modelComponent = ModelComponent(mesh: resource, materials: [material, material2])

        components.set(modelComponent)
    }
    
    func getMeshResource() throws -> MeshResource {
        var extrusionOptions = MeshResource.ShapeExtrusionOptions()
        extrusionOptions.extrusionMethod = .linear(depth: 0.4)
        extrusionOptions.materialAssignment = .init(front: 0, back: 0, extrusion: 1, frontChamfer: 1, backChamfer: 1)
        extrusionOptions.chamferRadius = 0.01
        return try MeshResource(extruding: attributedString, extrusionOptions: extrusionOptions)
    }
    
    var attributedString: AttributedString {
        var textString = AttributedString("HAPPY 4TH OF JULY")
        textString.font = .systemFont(ofSize: 1.9)
        
        let paragraphStyle = NSMutableParagraphStyle()
        paragraphStyle.alignment = .center
        let centerAttributes = AttributeContainer([.paragraphStyle: paragraphStyle])
        textString.mergeAttributes(centerAttributes)
        
        return textString
    }
}

PlaneMeshData.swift

import Foundation
import RealityKit

struct PlaneMeshData: Equatable, Codable {
    var size: SIMD2<Float> = [0.2, 0.2]
    var dimensions: SIMD2<UInt32> = [16, 16]
    
    var vertexCount: Int {
        Int(dimensions.x * dimensions.y)
    }

    var indexCount: Int {
        Int(6 * (dimensions.x - 1) * (dimensions.y - 1))
    }
    
    var boundingBox: BoundingBox {
        BoundingBox(
            min: [-size.x/2, -size.y/2, 0],
            max: [size.x/2, size.y/2, 0]
        )
    }
}

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

updateFlagMesh.metal

#include <metal_stdlib>
using namespace metal;

#include "VertexData.h"
#include "FlagParams.h"

float calculateAmplitude(float xCoord01, float exponentialFalloff, float windStrength) {
    return (1.3 - exp(-exponentialFalloff * xCoord01)) * windStrength;
}

float3 calculateWaves(float time, float xPosition, float yPosition) {
    float wavePhase1 = -time * 3.0 + xPosition * 40.0;
    float wave1 = sin(wavePhase1) * 0.002;

    float wavePhase2 = -time * 4.5 + xPosition * 140.0 + yPosition * 6.0;
    float wave2 = sin(wavePhase2) * 0.0025;

    float wavePhase3 = -time * 2.0 + xPosition * 20.0 - yPosition * 30.0;
    float wave3 = sin(wavePhase3) * 0.003;
    
    return float3(wave1, wave2, wave3);
}

float calculateZPosition(float xCoord01, float yCoord01, constant FlagParams& params, float exponentialFalloff) {
    float xPosition = params.size.x * xCoord01 - params.size.x / 2;
    float yPosition = params.size.y * yCoord01 - params.size.y / 2;
    
    float amplitude = calculateAmplitude(xCoord01, exponentialFalloff, params.windStrength);
    float3 waves = calculateWaves(params.time, xPosition, yPosition);
    
    return (waves.x + waves.y + waves.z) * amplitude;
}

kernel void updateFlagMesh(device VertexData* vertices [[buffer(0)]],
                           constant FlagParams& params [[buffer(1)]],
                           uint2 id [[thread_position_in_grid]])
{
    uint x = id.x;
    uint y = id.y;
    
    if (x >= params.dimensions.x || y >= params.dimensions.y) return;

    uint vertexIndex = y * params.dimensions.x + x;

    float xCoord01 = float(x) / float(params.dimensions.x - 1);
    float yCoord01 = float(y) / float(params.dimensions.y - 1);

    float xPosition = params.size.x * xCoord01 - params.size.x / 2;
    float yPosition = params.size.y * yCoord01 - params.size.y / 2;

    float exponentialFalloff = 20.0;
    float zPosition = calculateZPosition(xCoord01, yCoord01, params, exponentialFalloff);

    vertices[vertexIndex].position = float3(xPosition, yPosition, zPosition);
    
    // Calculate normal based on neighbors
    float3 normal = float3(0, 0, 1);
    if (x > 0 && x < params.dimensions.x - 1 && y > 0 && y < params.dimensions.y - 1) {
        float leftXCoord = float(x - 1) / float(params.dimensions.x - 1);
        float rightXCoord = float(x + 1) / float(params.dimensions.x - 1);
        float upYCoord = float(y + 1) / float(params.dimensions.y - 1);
        float downYCoord = float(y - 1) / float(params.dimensions.y - 1);

        float leftZ = calculateZPosition(leftXCoord, yCoord01, params, exponentialFalloff);
        float rightZ = calculateZPosition(rightXCoord, yCoord01, params, exponentialFalloff);
        float upZ = calculateZPosition(xCoord01, upYCoord, params, exponentialFalloff);
        float downZ = calculateZPosition(xCoord01, downYCoord, params, exponentialFalloff);

        float3 tangentX = float3(params.size.x / float(params.dimensions.x - 1), 0, rightZ - leftZ);
        float3 tangentY = float3(0, params.size.y / float(params.dimensions.y - 1), upZ - downZ);

        normal = normalize(cross(tangentX, tangentY));
    }
    
    vertices[vertexIndex].normal = normal;

    float2 uv = float2(
        float(x) / float(params.dimensions.x - 1),
        float(y) / float(params.dimensions.y - 1)
    );
    vertices[vertexIndex].uv = uv;
}

VertexData+Extensions.swift

import Foundation
import RealityKit

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

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

VertexData.h

#include <simd/simd.h>

#ifndef VertexData_h
#define VertexData_h

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

#endif /* VertexData_h */