RealityRenderer Test (visionOS)

RealityRendererTest.swift

//
//  RealityRendererView.swift
//  RealityRendererTest
//
//  Created by Arthur Schiller on 11.01.24.
//

import SwiftUI
import Metal
import MetalKit
import RealityKit

struct RealityRendererView: View {
    var body: some View {
        MetalViewRepresentable()
            .background(Color.red)
    }
}

#Preview(windowStyle: .automatic) {
    RealityRendererView()
}

struct MetalViewRepresentable: UIViewRepresentable {
    func makeUIView(context: Context) -> MetalView {
        return MetalView()
    }

    func updateUIView(_ uiView: MetalView, context: Context) {}
}

class MetalView: UIView {
    override class var layerClass: AnyClass {
        return CAMetalLayer.self
    }

    var metalLayer: CAMetalLayer {
        return layer as! CAMetalLayer
    }

    var device: MTLDevice!
    var displayLink: CADisplayLink!
    
    var commandQueue: MTLCommandQueue!
    
    var offscreenTexture: MTLTexture!
    var imageTexture: MTLTexture!
    
    var blitTestImage: Bool = false
    
    private var realityRenderer: RealityKit.RealityRenderer?
    private var lockFlag = false
    private var previousTimeStamp: TimeInterval?
    
    private let camera = PerspectiveCamera()
    private let modelEntity = ModelEntity(
        mesh: .generateBox(size: 0.3, cornerRadius: 0.1),
        materials: [
            UnlitMaterial(color: .blue)
        ]
    )

    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)
        commonInit()
    }

    override init(frame: CGRect) {
        super.init(frame: frame)
        commonInit()
    }

    private func commonInit() {
        device = MTLCreateSystemDefaultDevice()
        metalLayer.device = device
        metalLayer.pixelFormat = .bgra8Unorm
        metalLayer.framebufferOnly = false

        commandQueue = device.makeCommandQueue()
        createCommandQueue(device: device)
        startDisplayLink()
        
        do {
            let realityRenderer = try RealityKit.RealityRenderer()
            realityRenderer.cameraSettings.colorBackground = .color(
                .init(srgbRed: 0.5, green: 0.2, blue: 0.3, alpha: 1)
            )
//            realityRenderer.cameraSettings.colorBackground = .outputTexture()
            
            realityRenderer.entities.append(modelEntity)
            realityRenderer.entities.append(camera)
            
            realityRenderer.activeCamera = camera
            
            self.realityRenderer = realityRenderer
        } catch {
            print(error)
        }
    }
    
    private func setupOffscreenTexture(withDrawable drawable: CAMetalDrawable) {
        let texture = drawable.texture
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: texture.width,
            height: texture.height,
            mipmapped: false
        )
        
        textureDescriptor.usage = [.shaderRead, .shaderWrite, .renderTarget]
        offscreenTexture = device.makeTexture(descriptor: textureDescriptor)
    }
    
    private func loadTextureFromAssetCatalog(named imageName: String, device: MTLDevice) -> MTLTexture? {
        let textureLoader = MTKTextureLoader(device: device)

        do {
            let texture = try textureLoader.newTexture(name: imageName, scaleFactor: 1.0, bundle: nil, options: nil)
            return texture
        } catch {
            print("Error loading texture from asset catalog: \(error)")
            return nil
        }
    }
        
    //MARK: Builders
    private func createCommandQueue(device: MTLDevice) {
        commandQueue = device.makeCommandQueue()
    }

    private func startDisplayLink() {
        displayLink = CADisplayLink(target: self, selector: #selector(drawFrame))
        displayLink.add(to: .current, forMode: .default)
    }
    
    @objc private func drawFrame() {
        guard
            !lockFlag,
            let drawable = metalLayer.nextDrawable()
        else {
            return
        }
        
        print("")
        
        lockFlag = true
        
        modelEntity.position.z = -0.5
        camera.look(at: modelEntity.position, from: .zero, relativeTo: nil)
        
        if imageTexture == nil {
            guard let texture = loadTextureFromAssetCatalog(named: "testTexture", device: device) else {
                fatalError()
            }
            imageTexture = texture
        }
        
        if offscreenTexture == nil || offscreenTexture.width != drawable.texture.width || offscreenTexture.height != drawable.texture.height {
            setupOffscreenTexture(withDrawable: drawable)
        }
        
        let commandBuffer = commandQueue.makeCommandBuffer()!
        let timestamp = displayLink.timestamp
        let deltaTime: TimeInterval = {
            guard let previousTimeStamp else {
                return 0
            }
            return timestamp - previousTimeStamp
        }()
        
        let event = device.makeEvent()!
        
        print("Delta time: \(deltaTime)")
        
        do {
            if let realityRenderer {
                let descriptor = RealityKit.RealityRenderer.CameraOutput.Descriptor.singleProjection(
                    colorTexture: offscreenTexture // can I render this directly into the drawable?
                )

                try realityRenderer.updateAndRender(
                    deltaTime: deltaTime,
                    cameraOutput: .init(descriptor),
                    whenScheduled: { _ in
                        print("Render scheduled for timestamp: \(timestamp)")
                    },
                    onComplete: { renderer in
                        print("Render complete for timestamp: \(timestamp)")
                    },
                    actionsBeforeRender: [
                    ],
                    actionsAfterRender: [
                        .signal(event, value: 1)
                    ]
                )
            }
        } catch {
            print("Error: \(error)")
        }
        
        commandBuffer.encodeWaitForEvent(event, value: 1)
        
        let blitEncoder = commandBuffer.makeBlitCommandEncoder()!
        
        // Blit the texture to the offscreenTexture's texture -> just for testing, not important
        if blitTestImage {
            blitEncoder.copy(from: imageTexture, sourceSlice: 0, sourceLevel: 0, sourceOrigin: MTLOrigin(x: 0, y: 0, z: 0), sourceSize: MTLSize(width: imageTexture.width, height: imageTexture.height, depth: 1), to: offscreenTexture, destinationSlice: 0, destinationLevel: 0, destinationOrigin: MTLOrigin(x: 0, y: 0, z: 0))
        }
        
        // Blit the offscreenTexture to the drawable texture
        blitEncoder.copy(from: offscreenTexture, sourceSlice: 0, sourceLevel: 0, sourceOrigin: MTLOrigin(x: 0, y: 0, z: 0), sourceSize: MTLSize(width: offscreenTexture.width, height: offscreenTexture.height, depth: 1), to: drawable.texture, destinationSlice: 0, destinationLevel: 0, destinationOrigin: MTLOrigin(x: 0, y: 0, z: 0))
        blitEncoder.endEncoding()
        
        commandBuffer.present(drawable)
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
        
        print("Render Completed for timestamp: \(timestamp)")
        print("---- ✅ ----")
        
        previousTimeStamp = timestamp

        lockFlag = false
    }
}