RealityKit LowLevelMesh pancake effect (smashing z position of vertices)

PancakeEffectView.swift

import SwiftUI
import RealityKit
import Metal

struct PancakeEffectView: View {
    @State var entity: ModelEntity?
    @State var lowLevelMesh: LowLevelMesh?
    @State var originalVerticesBuffer: MTLBuffer?
    @State var timer: Timer?
    @State var isForward: Bool = true
    @State var modulationProgress: Float = 0.0
    @State var dwellCounter: Int = 0
    @State var isDwelling: Bool = false
    
    let timerUpdateDuration: TimeInterval = 1/120.0
    let dwellDuration: Int = 30
    let modulationRate: Float = 0.0025
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipelineState: MTLComputePipelineState
    
    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        let library = device.makeDefaultLibrary()!
        let kernelFunction = library.makeFunction(name: "smashMeshKernel")!
        self.computePipelineState = try! device.makeComputePipelineState(function: kernelFunction)
    }
    
    var body: some View {
        RealityView { content in
            let model = try! await loadModelEntity()
            content.add(model)
            let lowLevelMesh = createMesh(from: model)!
            
//            model.model?.materials = [goldMaterial]

            // Store original vertex data
            lowLevelMesh.withUnsafeBytes(bufferIndex: 0) { buffer in
                let vertices = buffer.bindMemory(to: VertexData.self)
                let originalVertices = Array(vertices)
                
                // Create GPU buffer for original vertices
                let bufferSize = originalVertices.count * MemoryLayout<VertexData>.stride
                self.originalVerticesBuffer = device.makeBuffer(bytes: originalVertices, length: bufferSize, options: [])
            }

            model.model?.mesh = try! await MeshResource(from: lowLevelMesh)
            self.entity = model
            self.lowLevelMesh = lowLevelMesh
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
    
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: timerUpdateDuration, repeats: true) { timer in
            if isDwelling {
                // Count dwell time
                dwellCounter += 1
                if dwellCounter >= dwellDuration {
                    // Finished dwelling, switch direction and resume morphing
                    isDwelling = false
                    dwellCounter = 0
                    isForward.toggle()
                }
            } else {
                // Update Modulation Progress
                if isForward {
                    modulationProgress += modulationRate
                } else {
                    modulationProgress -= modulationRate
                }
                
                // Handle bounds and start dwelling
                if modulationProgress >= 1.0 {
                    modulationProgress = 1.0
                    isDwelling = true
                    dwellCounter = 0
                } else if modulationProgress <= 0.0 {
                    modulationProgress = 0.0
                    isDwelling = true
                    dwellCounter = 0
                }
            }
            
            updateMesh()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }

    func createMesh(from modelEntity: ModelEntity) -> LowLevelMesh? {
        guard let meshResource = modelEntity.model?.mesh,
              let meshPart = meshResource.contents.models.first?.parts.first else {
            return nil
        }
        
        let positions = meshPart[MeshBuffers.positions]?.elements ?? []
        let normals = meshPart[MeshBuffers.normals]?.elements ?? []
        let textureCoordinates = meshPart[MeshBuffers.textureCoordinates]?.elements ?? []
        let triangleIndices = meshPart.triangleIndices?.elements ?? []
        
        var descriptor = VertexData.descriptor
        descriptor.vertexCapacity = positions.count
        descriptor.indexCapacity = triangleIndices.count
        
        guard let lowLevelMesh = try? LowLevelMesh(descriptor: descriptor) else {
            return nil
        }
        
        // Copy vertex data
        lowLevelMesh.withUnsafeMutableBytes(bufferIndex: 0) { buffer in
            let vertices = buffer.bindMemory(to: (SIMD3<Float>, SIMD3<Float>, SIMD2<Float>).self)
            for i in 0..<positions.count {
                vertices[i] = (positions[i], normals[i], textureCoordinates[i])
            }
        }
        
        // Copy index data
        lowLevelMesh.withUnsafeMutableIndices { buffer in
            let indices = buffer.bindMemory(to: UInt32.self)
            for (index, triangleIndex) in triangleIndices.enumerated() {
                indices[index] = UInt32(triangleIndex)
            }
        }
        
        // Set up parts
        let bounds = meshResource.bounds
        lowLevelMesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: triangleIndices.count,
                topology: .triangle,
                bounds: bounds
            )
        ])
        
        return lowLevelMesh
    }
    
    func updateMesh() {
        guard let mesh = lowLevelMesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder(),
                let entity = entity,
              let model = entity.model,
        let originalBuffer = originalVerticesBuffer else { return }
                                                                            
        let bounds = model.mesh.bounds
        var minZ = bounds.min.z
        var maxZ = bounds.max.z

        let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)

        computeEncoder.setComputePipelineState(computePipelineState)
        computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0)
        computeEncoder.setBuffer(originalBuffer, offset: 0, index: 1)
        computeEncoder.setBytes(&modulationProgress, length: MemoryLayout<Float>.size, index: 2)
        computeEncoder.setBytes(&minZ, length: MemoryLayout<Float>.size, index: 3)
        computeEncoder.setBytes(&maxZ, length: MemoryLayout<Float>.size, index: 4)

        let threadsPerGrid = MTLSize(width: mesh.vertexCapacity, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        
        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
    
    func loadModelEntity(url: URL = URL(string: "https://matt54.github.io/Resources/Anubis_Statue_1.usdz")!) async throws -> ModelEntity {
        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 entity = try await ModelEntity.init(contentsOf: destinationURL)
        try FileManager.default.removeItem(at: destinationURL)
        return entity
    }
    
    var goldMaterial: PhysicallyBasedMaterial {
        var material = PhysicallyBasedMaterial()
        material.baseColor.tint = .init(red: 0.75, green: 0.75, blue: 0.5, alpha: 1.0)
        material.roughness = 0.0
        material.metallic = 1.0
        material.faceCulling = .none
        return material
    }
}

#Preview {
    PancakeEffectView()
}

smashMeshKernel.metal

#include <metal_stdlib>
using namespace metal;

#include "VertexData.h"

kernel void smashMeshKernel(device VertexData* vertices [[buffer(0)]],
                            device const VertexData* originalVertices [[buffer(1)]],
                            constant float& smashProgress [[buffer(2)]],
                            constant float& minZ [[buffer(3)]],
                            constant float& maxZ [[buffer(4)]],
                            uint vid [[thread_position_in_grid]])
{
    vertices[vid] = originalVertices[vid];
    
    float originalZ = vertices[vid].position.z;
    float3 originalNormal = vertices[vid].normal;
    
    // Calculate the normalized position of this vertex within the z range (0.0 to 1.0)
    float normalizedZ = (originalZ - minZ) / (maxZ - minZ);
    
    // Create a smashing effect that progressively crushes vertices towards minZ
    // Vertices closer to maxZ get smashed first and more severely
    float smashFactor = smoothstep(0.0, 1.0, smashProgress + (1.0 - normalizedZ) * 0.5);
    
    // Limit the maximum smash to prevent complete collapse (keep 2% of original height)
    smashFactor = min(smashFactor, 0.98);
    
    float targetZ = mix(originalZ, minZ, smashFactor);
    
    vertices[vid].position.z = targetZ;
    
    // This normal calculation is unrealistic but may work fine in context
    // assuming that you are looking at the model from the front
    float3 flattenedNormal = float3(0.0, 0.0, 1.0); // Points up in Z direction
    
    // Interpolate between original normal and flattened normal based on smash factor
    vertices[vid].normal = normalize(mix(originalNormal, flattenedNormal, smashProgress));
}

VertexData.h

#include <simd/simd.h>

#ifndef VertexData_h
#define VertexData_h

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

#endif /* PlaneVertex_h */