RealityKit LowLevelMesh Sphere to Rectangular Plane Morph with Metal Compute Shader

MorphingSpherePlaneParams.h

#ifndef MorphingSpherePlaneParams_h
#define MorphingSpherePlaneParams_h

struct MorphingSpherePlaneParams {
    int latitudeBands;
    int longitudeBands;
    float radius;
    float morphAmount;
};

#endif /* MorphingSpherePlaneParams_h */

SphereToRectangularPlaneMorphingView.swift

import SwiftUI
import RealityKit
import Metal

#Preview { SphereToRectangularPlaneMorphingView() }
struct SphereToRectangularPlaneMorphingView: View {
    @State var mesh: LowLevelMesh?
    @State var timer: Timer?
    @State var frameDuration: TimeInterval = 0.0
    @State var lastUpdateTime = CACurrentMediaTime()
    @State var morphAmount: Float = 0.0
    @State var isMorphForward: Bool = true
    @State var isDwelling: Bool = false
    @State var dwellCounter: Float = 0.0

    let rootEntity: Entity = Entity()
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipeline: MTLComputePipelineState
    
    var radius: Float = 0.1
    var latitudeBands = 120
    var longitudeBands = 120
    var morphSpeed: Float = 0.00375
    var morphRange: ClosedRange<Float> = 0...1
    var dwellDuration: Float = 180
    var useTexture: Bool = true
    var topology: MTLPrimitiveType = .triangle

    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "updateSphereToRectangularPlaneGeometry")!
        self.computePipeline = try! device.makeComputePipelineState(function: updateFunction)
    }
    
    var body: some View {
        RealityView { content in
            let mesh = try! createMesh()
            let meshResource = try! await MeshResource(from: mesh)
            let material = try! await getMaterial()
            let modelComponent = ModelComponent(mesh: meshResource, materials: [material])
            rootEntity.components.set(modelComponent)
            content.add(rootEntity)
            self.mesh = mesh
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
}

// MARK: Animation / Timer
extension SphereToRectangularPlaneMorphingView {
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: 1/120.0, repeats: true) { _ in
            let currentTime = CACurrentMediaTime()
            frameDuration = currentTime - lastUpdateTime
            lastUpdateTime = currentTime
            updateMesh()
            stepMorphAmount()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    func stepMorphAmount() {
        // If we're dwelling, count frames
        if isDwelling {
            dwellCounter += 1.0
            if dwellCounter >= dwellDuration {
                isDwelling = false
                dwellCounter = 0.0
                // Switch direction after dwelling
                isMorphForward.toggle()
            }
            return
        }
        
        // Normal morphing behavior
        if isMorphForward {
            morphAmount += morphSpeed
            if morphAmount >= morphRange.upperBound {
                morphAmount = morphRange.upperBound
                isDwelling = true
                dwellCounter = 0.0
            }
        } else {
            morphAmount -= morphSpeed
            if morphAmount <= morphRange.lowerBound {
                morphAmount = morphRange.lowerBound
                isDwelling = true
                dwellCounter = 0.0
            }
        }
    }
}

// MARK: Material
extension SphereToRectangularPlaneMorphingView {
    func getMaterial() async throws -> RealityKit.Material {
        var material = PhysicallyBasedMaterial()
        material.baseColor.tint = .yellow
        material.roughness.scale = 0.5
        material.metallic.scale = 1.0
        material.blending = .transparent(opacity: 1.0)
        material.faceCulling = .none
        
        if useTexture,
            let url = URL(string: "https://matt54.github.io/Resources/square_uv_grid.png") {
            // download and apply image as texture
            let (downloadedURL, _) = try await URLSession.shared.download(from: url)
            let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
            let destinationURL = documentsDirectory.appendingPathComponent("square_uv_grid.png")
            if FileManager.default.fileExists(atPath: destinationURL.path) {
                try FileManager.default.removeItem(at: destinationURL)
            }
            try FileManager.default.moveItem(at: downloadedURL, to: destinationURL)
            if let textureResource = try? await TextureResource(contentsOf: destinationURL) {
                material.baseColor = .init(texture: .init(textureResource))
                material.emissiveColor = .init(texture: .init(textureResource))
                material.emissiveIntensity = 0.25
            }
        }
        return material
    }
}

// MARK: Mesh
extension SphereToRectangularPlaneMorphingView {
    var vertexCapacity: Int {
        return (latitudeBands + 1) * (longitudeBands + 1)
    }
    var indexCount: Int {
        return latitudeBands * longitudeBands * 6
    }
    
    func createMesh() throws -> LowLevelMesh {
        var desc = VertexData.descriptor
        desc.vertexCapacity = vertexCapacity
        desc.indexCapacity = indexCount
        
        let mesh = try LowLevelMesh(descriptor: desc)
        
        mesh.withUnsafeMutableIndices { rawIndices in
            let indices = rawIndices.bindMemory(to: UInt32.self)
            var index = 0
            
            for latNumber in 0..<latitudeBands {
                for longNumber in 0..<longitudeBands {
                    let first = (latNumber * (longitudeBands + 1)) + longNumber
                    let second = first + longitudeBands + 1
                    
                    indices[index] = UInt32(first)
                    indices[index + 1] = UInt32(second)
                    indices[index + 2] = UInt32(first + 1)
                    
                    indices[index + 3] = UInt32(second)
                    indices[index + 4] = UInt32(second + 1)
                    indices[index + 5] = UInt32(first + 1)
                    
                    index += 6
                }
            }
        }
        
        return mesh
    }
    
    func updateMesh() {
        guard let mesh = mesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipeline)
        computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0)

        var params = MorphingSpherePlaneParams(
            latitudeBands: Int32(latitudeBands),
            longitudeBands: Int32(longitudeBands),
            radius:  radius,
            morphAmount: morphAmount
        )
        computeEncoder.setBytes(&params, length: MemoryLayout<MorphingSpherePlaneParams>.size, index: 1)
        
        let threadsPerGrid = MTLSize(width: vertexCapacity, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        
        computeEncoder.endEncoding()
        commandBuffer.commit()
        
        let meshBounds = BoundingBox(min: [-radius, -radius, -radius],
                                     max: [radius, radius,  radius])
        mesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: indexCount,
                topology: topology,
                bounds: meshBounds
            )
        ])
    }
}

updateSphereToRectangularPlaneGeometry.metal

#include <metal_stdlib>
using namespace metal;

#include "Shared/VertexData.h"
#include "MorphingSpherePlaneParams.h"

kernel void updateSphereToRectangularPlaneGeometry(device VertexData* vertices [[buffer(0)]],
                                                   constant MorphingSpherePlaneParams& params [[buffer(1)]],
                                                   uint id [[thread_position_in_grid]])
{
    int x = id % (params.longitudeBands + 1);
    int y = id / (params.longitudeBands + 1);
    
    if (x > params.longitudeBands || y > params.latitudeBands) return;
    
    float lat = float(y) / float(params.latitudeBands);
    float lon = float(x) / float(params.longitudeBands);
    float theta = (1.0 - lat) * M_PI_F;
    float phi = lon * 2 * M_PI_F;
    
    float sinTheta = sin(theta);
    float cosTheta = cos(theta);
    float sinPhi = sin(phi);
    float cosPhi = cos(phi);
    
    float3 spherePosition = float3(cosPhi * sinTheta, cosTheta, sinPhi * sinTheta);
    float2 uv = float2(lon, lat);
    float3 planePosition = float3(0, (uv.y - 0.5)  * 2.0,  (uv.x - 0.5)  * -2.0  * M_PI_F);
    float3 interpolatedPosition = mix(spherePosition, planePosition, params.morphAmount);

    // Constant Y-axis rotation matrix (-90 degrees)
    // Applied to position and normal
    float3x3 rotationMatrix = float3x3(
        float3(0, 0, -1),
        float3(0, 1,  0),
        float3(1, 0,  0)
    );

    // Calculate proper normals for both sphere and plane, then interpolate
    float3 sphereNormal = normalize(spherePosition);
    float3 planeNormal = float3(0, -1, 0); // Plane faces upward
    float3 interpolatedNormal = mix(sphereNormal, planeNormal, params.morphAmount);
    float3 rotatedNormal = rotationMatrix * interpolatedNormal;
    float3 normal = normalize(rotatedNormal);

    vertices[id].position = rotationMatrix * interpolatedPosition * params.radius;
    vertices[id].normal = normal;
    uv.x *= -1; // flip texture
    vertices[id].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