RealityKit HeightMap Image to Terrain with LowLevelMesh and LowLevelTexture

convertHeightMapToColors.metal

#include <metal_stdlib>
using namespace metal;

// Function to get color based on height and color mode
float3 getColorForHeight(float height, int colorMode) {
    float3 finalColor;
    
    switch (colorMode) {
        case 0: {
            // Realistic terrain coloring
            float3 waterColor = float3(0.2, 0.4, 0.8);      // Deep blue for water/lowest areas
            float3 sandColor = float3(0.8, 0.7, 0.5);       // Sandy beige for beaches/low areas
            float3 grassColor = float3(0.3, 0.6, 0.2);      // Green for grasslands
            float3 forestColor = float3(0.2, 0.4, 0.1);     // Dark green for forests
            float3 rockColor = float3(0.5, 0.4, 0.3);       // Brown-gray for rocky areas
            float3 snowColor = float3(0.9, 0.9, 1.0);       // White-blue for snow peaks
            
            // Define elevation thresholds (0.0 to 1.0 range)
            float waterLevel = 0.15;
            float sandLevel = 0.25;
            float grassLevel = 0.45;
            float forestLevel = 0.65;
            float rockLevel = 0.85;
            
            // Smooth transitions between terrain types using mix function
            if (height < waterLevel) {
                finalColor = waterColor;
            } else if (height < sandLevel) {
                float t = (height - waterLevel) / (sandLevel - waterLevel);
                finalColor = mix(waterColor, sandColor, t);
            } else if (height < grassLevel) {
                float t = (height - sandLevel) / (grassLevel - sandLevel);
                finalColor = mix(sandColor, grassColor, t);
            } else if (height < forestLevel) {
                float t = (height - grassLevel) / (forestLevel - grassLevel);
                finalColor = mix(grassColor, forestColor, t);
            } else if (height < rockLevel) {
                float t = (height - forestLevel) / (rockLevel - forestLevel);
                finalColor = mix(forestColor, rockColor, t);
            } else {
                float t = (height - rockLevel) / (1.0 - rockLevel);
                finalColor = mix(rockColor, snowColor, t);
            }
            break;
        }
        
        case 1: {
            // Standard altitude heat map (blue to red gradient)
            float3 coldColor = float3(0.0, 0.0, 1.0);       // Blue for low elevations
            float3 coolColor = float3(0.0, 1.0, 1.0);       // Cyan
            float3 neutralColor = float3(0.0, 1.0, 0.0);    // Green
            float3 warmColor = float3(1.0, 1.0, 0.0);       // Yellow
            float3 hotColor = float3(1.0, 0.0, 0.0);        // Red for high elevations
            
            if (height < 0.25) {
                float t = height / 0.25;
                finalColor = mix(coldColor, coolColor, t);
            } else if (height < 0.5) {
                float t = (height - 0.25) / 0.25;
                finalColor = mix(coolColor, neutralColor, t);
            } else if (height < 0.75) {
                float t = (height - 0.5) / 0.25;
                finalColor = mix(neutralColor, warmColor, t);
            } else {
                float t = (height - 0.75) / 0.25;
                finalColor = mix(warmColor, hotColor, t);
            }
            break;
        }
        
        default: {
            // Fallback to grayscale
            finalColor = float3(height, height, height);
            break;
        }
    }
    
    return finalColor;
}

[[kernel]]
void convertHeightMapToColors(texture2d<float, access::read> imageTexture [[texture(0)]],
                              texture2d<float, access::write> outputTexture [[texture(1)]],
                              constant int& colorMode [[buffer(0)]],
                              constant float& progress [[buffer(1)]],
                              uint2 pixelCoords [[thread_position_in_grid]]) {
    // Skip out-of-bounds threads.
    if (pixelCoords.x >= imageTexture.get_width() || pixelCoords.y >= imageTexture.get_height()) { return; }
    
    // Read the RGB values from the input image texture
    float4 imageData = imageTexture.read(pixelCoords);
    
    // Use the red channel as height (assuming grayscale input)
    float height = dot(imageData.rgb, float3(0.299, 0.587, 0.114)) * 1.5 * progress;
    
    // Get the color for this height using the specified color mode
    float3 finalColor = getColorForHeight(height, colorMode);
    
    // Write the final color to the output texture
    outputTexture.write(float4(finalColor, 1.0), pixelCoords);
}

deriveNormalsFromHeightMap.metal

#include <metal_stdlib>
using namespace metal;

/// See: https://developer.apple.com/documentation/realitykit/creating-a-dynamic-height-map-with-low-level-texture
/// Derives normal directions from a height map, storing them in the texture's rgb channels.
[[kernel]]
void deriveNormalsFromHeightMap(texture2d<float, access::read> heightMapIn [[texture(0)]],
                                texture2d<float, access::write> heightMapOut [[texture(1)]],
                                constant float2 &cellSize [[buffer(2)]],
                                uint2 pixelCoords [[thread_position_in_grid]]) {
    // Get the dimensions of the height map.
    uint2 dimensions = uint2(heightMapIn.get_width(), heightMapIn.get_height());
    
    // Skip out-of-bounds threads.
    if (any(pixelCoords >= dimensions)) { return; }
    
    // The current pixel coordinate minus one in both dimensions, guaranteed to be in bounds.
    uint2 pixelCoordsMinusOne = max(pixelCoords, 1) - 1;
    // The current pixel coordinate plus one in both dimensions, guaranteed to be in bounds.
    uint2 pixelCoordsPlusOne = min(pixelCoords + 1, dimensions - 1);
    
    // Sample the current pixel along with its four neighbors.
    float height = heightMapIn.read(pixelCoords).a;
    float leftHeight = heightMapIn.read(uint2(pixelCoordsMinusOne.x, pixelCoords.y)).a;
    float rightHeight = heightMapIn.read(uint2(pixelCoordsPlusOne.x, pixelCoords.y)).a;
    float bottomHeight = heightMapIn.read(uint2(pixelCoords.x, pixelCoordsMinusOne.y)).a;
    float topHeight = heightMapIn.read(uint2(pixelCoords.x, pixelCoordsPlusOne.y)).a;
    
    // Compute the normal direction using central differences.
    float3 normal = normalize(float3((leftHeight - rightHeight) / (cellSize.x * 2),
                                     (bottomHeight - topHeight) / (cellSize.y * 2),
                                     1));
    
    // Write the normal direction to the height map.
    heightMapOut.write(float4(normal, height), pixelCoords);
}

HeightMapParams.h

#ifndef HeightMapParams_h
#define HeightMapParams_h

struct HeightMapParams {
    simd_float2 size;
    simd_uint2 dimensions;
};

#endif /* HeightMapParams_h */

HeightMapToColoredMeshView.swift

import SwiftUI
import RealityKit
import Metal

struct HeightMapToColoredMeshView: View {
    @State var heightMap: ExampleHeightMapImage
    @State var isTransitioning: Bool = false
    @State var colorMode: HeightMapColorMode = .terrain
    @State var meshData: HeightMapMeshData?
    @State var mesh: LowLevelMesh?
    
    // holds the rgb values of the original image
    @State var originalImageTexture: LowLevelTexture?
    
    // used to define the positions/normals of the mesh
    @State var heightMapTexture: LowLevelTexture?
    
    // the colors you see on the mesh
    @State var materialTexture: LowLevelTexture?
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let meshComputePipeline: MTLComputePipelineState
    let imageToGreyscalePipeline: MTLComputePipelineState
    let deriveNormalsPipeline: MTLComputePipelineState
    let convertHeightMapToColorPipeline: MTLComputePipelineState

    @State var timer: Timer?
    @State var heightProgress: Float = 0.0
    let heightProgressRate: Float = 0.01
    let timerUpdateDuration: TimeInterval = 1/120.0
    
    init(heightMap: ExampleHeightMapImage = ExampleHeightMapImage.mountEverest) {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "updateHeightMapMesh")!
        let heightFunction = library.makeFunction(name: "convertGreyscaleToAlpha")!
        let normalsFunction = library.makeFunction(name: "deriveNormalsFromHeightMap")!
        let colorFunction = library.makeFunction(name: "convertHeightMapToColors")!
        self.meshComputePipeline = try! device.makeComputePipelineState(function: updateFunction)
        self.imageToGreyscalePipeline = try! device.makeComputePipelineState(function: heightFunction)
        self.deriveNormalsPipeline = try! device.makeComputePipelineState(function: normalsFunction)
        self.convertHeightMapToColorPipeline = try! device.makeComputePipelineState(function: colorFunction)
        
        self.heightMap = heightMap
    }
    
    var body: some View {
        RealityView { content in
            let textureResource = try! await TextureResource.loadOnlineImage(heightMap.url)
            let originalImageTexture = try! copyTextureResourceToLowLevelTexture(from: textureResource)
            self.originalImageTexture = originalImageTexture
            
            let width = UInt32(textureResource.width)
            let height = UInt32(textureResource.height)
            let dimensions: SIMD2<UInt32> = SIMD2(x: width, y: height)
            let meshData = HeightMapMeshData(dimensions: SIMD2(x: width, y: height))

            self.meshData = meshData

            let textureDescriptor = LowLevelTexture.Descriptor(pixelFormat: .rgba32Float,
                                                               width: Int(dimensions.x),
                                                               height: Int(dimensions.y),
                                                               textureUsage: [.shaderRead, .shaderWrite])
            self.heightMapTexture = try! LowLevelTexture(descriptor: textureDescriptor)

            let meshMaterialTexture = try! LowLevelTexture(descriptor: textureDescriptor)
            self.materialTexture = meshMaterialTexture
            
            let mesh = try! VertexData.initializeMesh(vertexCapacity: meshData.vertexCount,
                                                      indexCapacity: meshData.indexCount)
            let resource = try! await MeshResource(from: mesh)
            
            var material = PhysicallyBasedMaterial()
            let textureResourceForMaterial = try! await TextureResource(from: meshMaterialTexture)
            material.baseColor = .init(texture: .init(textureResourceForMaterial))
            
            material.metallic = .init(floatLiteral: 0.5)  // Make it more metallic
            material.roughness = .init(floatLiteral: 0.25) // Make it smoother
            
            let modelComponent = ModelComponent(mesh: resource, materials: [material])
            
            let entity = Entity()
            entity.components.set(modelComponent)
            content.add(entity)
            
            entity.transform.rotation = .init(angle: -.pi / 2, axis: .init(x: 1, y: 0, z: 0))
            entity.position = .init(x: 0, y: -0.2, z: 0.0)
            
            self.mesh = mesh
            
            updateMeshAndTexture()
            startTimer()
        }
        .ornament(attachmentAnchor: .scene(.topBack), contentAlignment: .top) {
            VStack(spacing: 16) {
                Text("Height Map Visualization")
                    .font(.headline)
                
                Picker("Mountain", selection: $heightMap) {
                    ForEach(ExampleHeightMapImage.allCases, id: \.self) { mountain in
                        Text(mountain.displayName).tag(mountain)
                    }
                }
                .pickerStyle(.segmented)
                .onChange(of: heightMap) { _, newMountain in
                    changeMountain(to: newMountain)
                }
                
                Picker("Color Mode", selection: $colorMode) {
                    ForEach(HeightMapColorMode.allCases, id: \.self) { mode in
                        Text(mode.name).tag(mode)
                    }
                }
                .pickerStyle(.segmented)
                .onChange(of: colorMode) { _, newColorMode in
                    updateMaterialTextureColors()
                }
            }
            .padding()
            .background(Color.brown, in: RoundedRectangle(cornerRadius: 12))
            .padding(.top, 300)
        }
    }
    
    func changeMountain(to newMountain: ExampleHeightMapImage) {
        guard !isTransitioning else { return }
        
        isTransitioning = true
        
        // Stop current timer and start transition
        stopTimer()
        startTimer()
    }

    func loadNewMountainTexture() async throws {
        let textureResource = try await TextureResource.loadOnlineImage(heightMap.url)
        let newOriginalImageTexture = try copyTextureResourceToLowLevelTexture(from: textureResource)
        
        // Update the texture
        self.originalImageTexture = newOriginalImageTexture
        
        // Mark transition as complete so height can start rising
        isTransitioning = false
        
        startTimer()
    }
    
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: timerUpdateDuration, repeats: true) { timer in
            if isTransitioning && heightProgress > 0 {
                // lower the height to 0
                heightProgress -= heightProgressRate * 2
                if heightProgress <= 0 {
                    heightProgress = 0
                    // Load new texture when height reaches 0
                    stopTimer()
                    Task {
                        try! await loadNewMountainTexture()
                    }
                }
            } else {
                // increase height to 1
                heightProgress += heightProgressRate
                if heightProgress >= 1 {
                    heightProgress = 1
                    isTransitioning = false  // Clear transition flag if it was set
                    stopTimer()
                }
            }
            
            updateMeshAndTexture()
        }
    }
    
    func updateMeshAndTexture() {
        updateHeightMapFromImageTexture()
        updateTextureNormals()
        updateMaterialTextureColors()
        updateMeshGeometry()
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
    
    func updateHeightMapFromImageTexture() {
        guard let imageTexture = originalImageTexture,
              let heightMapTexture,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }

        computeEncoder.setComputePipelineState(imageToGreyscalePipeline)

        // input: rbg texture
        computeEncoder.setTexture(imageTexture.read(), index: 0)
        
        // output: alpha texture
        computeEncoder.setTexture(heightMapTexture.replace(using: commandBuffer), index: 1)
        
        var progress = heightProgress
        computeEncoder.setBytes(&progress, length: MemoryLayout<Float>.size, index: 0)
        
        let threadWidth = imageToGreyscalePipeline.threadExecutionWidth
        let threadHeight = imageToGreyscalePipeline.maxTotalThreadsPerThreadgroup / threadWidth
        let threadsPerThreadgroup = MTLSize(width: threadWidth, height: threadHeight, depth: 1)
        
        let threadGroupSize = MTLSizeMake(8, 8, 1)
        let threadGroups = MTLSizeMake(
            (heightMapTexture.descriptor.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (heightMapTexture.descriptor.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1
        )

        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadsPerThreadgroup)

        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
    
    func updateTextureNormals() {
        guard let heightMapTexture,
                let meshData,
                let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }

        computeEncoder.setComputePipelineState(deriveNormalsPipeline)

        // Pass a readable version of the height map texture to the compute shader.
        computeEncoder.setTexture(heightMapTexture.read(), index: 0)
        // Pass a writable version of the height map texture to the compute shader.
        computeEncoder.setTexture(heightMapTexture.replace(using: commandBuffer), index: 1)

        // Pass the cell size to the compute shader.
        var mutableCellSize = meshData.cellSize
        computeEncoder.setBytes(&mutableCellSize, length: MemoryLayout<SIMD2<Float>>.size, index: 2)

        let threadWidth = imageToGreyscalePipeline.threadExecutionWidth
        let threadHeight = imageToGreyscalePipeline.maxTotalThreadsPerThreadgroup / threadWidth
        let threadsPerThreadgroup = MTLSize(width: threadWidth, height: threadHeight, depth: 1)
        
        let threadGroupSize = MTLSizeMake(8, 8, 1)
        let threadGroups = MTLSizeMake(
            (heightMapTexture.descriptor.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (heightMapTexture.descriptor.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1
        )

        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadsPerThreadgroup)

        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
    
    func updateMaterialTextureColors() {
        guard let imageTexture = originalImageTexture,
              let materialTexture,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }

        computeEncoder.setComputePipelineState(convertHeightMapToColorPipeline)

        // input: rbg texture
        computeEncoder.setTexture(imageTexture.read(), index: 0)

        // output: rbg texture
        computeEncoder.setTexture(materialTexture.replace(using: commandBuffer), index: 1)
        
        var colorMode: Int32 = colorMode.value
        computeEncoder.setBytes(&colorMode, length: MemoryLayout<Int32>.size, index: 0)
        
        var progress = heightProgress
        computeEncoder.setBytes(&progress, length: MemoryLayout<Float>.size, index: 1)
        
        let threadWidth = convertHeightMapToColorPipeline.threadExecutionWidth
        let threadHeight = convertHeightMapToColorPipeline.maxTotalThreadsPerThreadgroup / threadWidth
        let threadsPerThreadgroup = MTLSize(width: threadWidth, height: threadHeight, depth: 1)
        
        let threadGroupSize = MTLSizeMake(8, 8, 1)
        let threadGroups = MTLSizeMake(
            (materialTexture.descriptor.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (materialTexture.descriptor.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1
        )

        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadsPerThreadgroup)

        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
    
    func updateMeshGeometry() {
        guard let mesh,
                let heightMapTexture,
                let meshData,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        var params = HeightMapParams(size: meshData.size, dimensions: meshData.dimensions)
        let vertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        let indexBuffer = mesh.replaceIndices(using: commandBuffer)
        
        computeEncoder.setComputePipelineState(meshComputePipeline)
        computeEncoder.setBuffer(vertexBuffer, offset: 0, index: 0)
        computeEncoder.setBuffer(indexBuffer, offset: 0, index: 1)
        computeEncoder.setBytes(&params, length: MemoryLayout<HeightMapParams>.stride, index: 2)
        computeEncoder.setTexture(heightMapTexture.read(), index: 3)
        
        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
            )
        ])
    }
    
    func copyTextureResourceToLowLevelTexture(from textureResource: TextureResource) throws -> LowLevelTexture {
        var descriptor = LowLevelTexture.Descriptor()
        descriptor.textureType = .type2D
        descriptor.pixelFormat = .rgba16Float
        descriptor.width = textureResource.width
        descriptor.height = textureResource.height
        descriptor.mipmapLevelCount = 1
        descriptor.textureUsage = [.shaderRead, .shaderWrite]
        
        let texture = try LowLevelTexture(descriptor: descriptor)
        try textureResource.copy(to: texture.read())
        
        return texture
    }
}

#Preview {
    HeightMapToColoredMeshView()
}

enum ExampleHeightMapImage: CaseIterable {
    case mountEverest
    case k2
    case kangchenjunga
    case lhotse
    case makalu

    static let baseURL = URL(string: "https://matt54.github.io/Resources/")!
    
    var url: URL {
        return ExampleHeightMapImage.baseURL.appendingPathComponent("\(filename)_2056.png")
    }

    var filename: String {
        switch self {
        case .mountEverest:
            return "MT_EVEREST"
        case .k2:
            return "K2"
        case .kangchenjunga:
            return "Kangchenjunga"
        case .lhotse:
            return "Lhotse"
        case .makalu:
            return "Makalu"
        }
    }
    
    var displayName: String {
        switch self {
        case .mountEverest:
            return "Mount Everest"
        case .k2:
            return "K2"
        case .kangchenjunga:
            return "Kangchenjunga"
        case .lhotse:
            return "Lhotse"
        case .makalu:
            return "Makalu"
        }
    }
}

enum HeightMapColorMode: CaseIterable {
    case terrain
    case heatMap
    
    var name: String {
        switch self {
        case .terrain:
            "Terrain"
        case .heatMap:
            "Heat Map"
        }
    }
    
    var value: Int32 {
        switch self {
        case .terrain:
            0
        case .heatMap:
            1
        }
    }
}

struct HeightMapMeshData {
    var size: SIMD2<Float> = [0.4, 0.4]
    var dimensions: SIMD2<UInt32> = [2056, 2056]
    
    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]
        )
    }
    
    var cellSize: SIMD2<Float> {
        SIMD2<Float>(size.x / Float(dimensions.x - 1),
                     size.y / Float(dimensions.y - 1))
    }
}

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

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

updateHeightMapMesh.metal

#include <metal_stdlib>
using namespace metal;

#include "VertexData.h"
#include "HeightMapParams.h"

/// See: https://developer.apple.com/documentation/realitykit/creating-a-dynamic-height-map-with-low-level-texture
kernel void updateHeightMapMesh(device VertexData* vertices [[buffer(0)]],
                                device uint* indices [[buffer(1)]],
                                constant HeightMapParams& params [[buffer(2)]],
                                texture2d<float, access::read> heightMap [[texture(3)]],
                                uint2 id [[thread_position_in_grid]])
{
    uint x = id.x;
    uint y = id.y;
    
    if (x >= params.dimensions.x || y >= params.dimensions.y) return;

    // Sample the height map pixel corresponding to this vertex.
    float4 heightMapData = heightMap.read(id);
    // Extract the normal direction and the height.
    float3 normal = heightMapData.rgb;
    float height = heightMapData.a;

    // Calculate normalized coordinates (0 to 1)
    float xCoord01 = float(x) / float(params.dimensions.x - 1);
    float yCoord01 = float(y) / float(params.dimensions.y - 1);
    
    // Apply height displacement to the Z position
    float xPosition = params.size.x * xCoord01 - params.size.x / 2;
    float yPosition = params.size.y * yCoord01 - params.size.y / 2;
    float zPosition = height;
    
    // Calculate vertex index
    uint vertexIndex = y * params.dimensions.x + x;

    vertices[vertexIndex].position = float3(xPosition, yPosition, zPosition);
    vertices[vertexIndex].normal = normal;
    
    float2 uv = float2(
        float(x) / float(params.dimensions.x - 1),
        1.0 - float(y) / float(params.dimensions.y - 1)  // Flip the V coordinate (texture was upside down otherwise)
    );
    vertices[vertexIndex].uv = uv;
    
    // Create indices for triangles
    if (x < params.dimensions.x - 1 && y < params.dimensions.y - 1) {
        uint indexBase = 6 * (y * (params.dimensions.x - 1) + x);
        uint bottomLeft = vertexIndex;
        uint bottomRight = bottomLeft + 1;
        uint topLeft = bottomLeft + params.dimensions.x;
        uint topRight = topLeft + 1;
        
        // First triangle
        indices[indexBase] = bottomLeft;
        indices[indexBase + 1] = bottomRight;
        indices[indexBase + 2] = topLeft;
        
        // Second triangle
        indices[indexBase + 3] = topLeft;
        indices[indexBase + 4] = bottomRight;
        indices[indexBase + 5] = topRight;
    }
}

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 */