DebugDraw

DebugDrawPass.swift

import Substrate
import SubstrateMath
import DebugDraw
import SubstrateUtilities
import ShaderReflection

final class DebugDrawPass : ReflectableDrawRenderPass {
    
    public typealias Reflection = DebugDrawPassReflection
    
    static let vertexDescriptor : VertexDescriptor = {
        var descriptor = VertexDescriptor()
        
        // position
        descriptor.attributes[0].bufferIndex = 0
        descriptor.attributes[0].offset = 0
        descriptor.attributes[0].format = .float3
        
        // colour
        descriptor.attributes[1].bufferIndex = 0
        descriptor.attributes[1].offset = 12
        descriptor.attributes[1].format = .uchar4Normalized
        
        descriptor.layouts[0].stepFunction = .perVertex
        descriptor.layouts[0].stepRate = 1
        descriptor.layouts[0].stride = MemoryLayout<DebugDraw.DebugVertex>.size
        
        return descriptor
    }()
    
    static let vertexDescriptorPoint : VertexDescriptor = {
        var descriptor = DebugDrawPass.vertexDescriptor
        
        // size
        descriptor.attributes[2].bufferIndex = 1
        descriptor.attributes[2].offset = 0
        descriptor.attributes[2].format = .float
        
        descriptor.layouts[1].stepFunction = .perVertex
        descriptor.layouts[1].stepRate = 1
        descriptor.layouts[1].stride = MemoryLayout<Float>.size
        
        return descriptor
    }()
    
    static let depthStencilNoDepth : DepthStencilDescriptor = {
        var depthStencilDescriptor = DepthStencilDescriptor()
        depthStencilDescriptor.depthCompareFunction = .always
        depthStencilDescriptor.isDepthWriteEnabled = false
        return depthStencilDescriptor
    }()
    
    static let depthStencilWithDepth : DepthStencilDescriptor = {
        var depthStencilDescriptor = DepthStencilDescriptor()
        depthStencilDescriptor.depthCompareFunction = .greater
        depthStencilDescriptor.isDepthWriteEnabled = true
        return depthStencilDescriptor
    }()
    
    static let pipelineDescriptor : RenderPipelineDescriptor = {
        var descriptor = RenderPipelineDescriptor(attachmentCount: 1)
        
        var blendDescriptor = BlendDescriptor()
        
        blendDescriptor.alphaBlendOperation = .add
        blendDescriptor.rgbBlendOperation = .add
        blendDescriptor.sourceRGBBlendFactor = .sourceAlpha
        blendDescriptor.sourceAlphaBlendFactor = .sourceAlpha
        blendDescriptor.destinationRGBBlendFactor = .oneMinusSourceAlpha
        blendDescriptor.destinationAlphaBlendFactor = .oneMinusSourceAlpha
        
        descriptor.blendStates[0] = blendDescriptor
        
        descriptor.vertexDescriptor = DebugDrawPass.vertexDescriptor
        
        return descriptor
    }()
    
    
    let renderTargetDescriptor: RenderTargetDescriptor
    
    var name: String = "Debug Draw"
    
    let drawData : DebugDraw.DrawData
    let projectionMatrix: Matrix4x4f
    
    init(drawData: DebugDraw.DrawData, projectionMatrix: Matrix4x4f, renderTargetDescriptor: RenderTargetDescriptor) {
        self.drawData = drawData
        self.projectionMatrix = projectionMatrix
        self.renderTargetDescriptor = renderTargetDescriptor
    }
    
    func execute(renderCommandEncoder rce: TypedRenderCommandEncoder<Reflection>) async {
        rce.setCullMode(.none)
        
        rce.pipeline.descriptor = DebugDrawPass.pipelineDescriptor
        rce.pipeline.descriptor.fillMode = .fill
        
        rce.pushConstants.worldToProjection = projectionMatrix
        
        let bufferSize = self.drawData.depthDisabledData.requiredBufferCapacity + self.drawData.depthEnabledData.requiredBufferCapacity
        
        guard bufferSize > 0 else { return }
        
        let buffer = Buffer(length: bufferSize)
        
        rce.setVertexBuffer(buffer, offset: 0, index: 0)
        rce.setVertexBuffer(buffer, offset: 0, index: 1)
        
        var bufferOffset = 0
        // Non-points
        do {
            
            rce.pipeline.vertexFunction =  .vertex
            rce.pipeline.fragmentFunction =  .fragment
            
            if self.renderTargetDescriptor.depthAttachment != nil {
                rce.depthStencil = DebugDrawPass.depthStencilWithDepth
            }
            await self.drawData.depthEnabledData.drawNonPoints(encoder: rce, buffer: buffer, bufferOffset: &bufferOffset)
            
            rce.depthStencil = DebugDrawPass.depthStencilNoDepth
            await self.drawData.depthDisabledData.drawNonPoints(encoder: rce, buffer: buffer, bufferOffset: &bufferOffset)
        }
        
        // Points
        do {
            rce.pipeline.vertexFunction =  .vertexPoint
            rce.pipeline.fragmentFunction =  .fragmentPoint
            
            rce.pipeline.vertexDescriptor = DebugDrawPass.vertexDescriptorPoint
            
            rce.depthStencil = DebugDrawPass.depthStencilWithDepth
            await self.drawData.depthEnabledData.drawPoints(encoder: rce, buffer: buffer, bufferOffset: &bufferOffset)
            
            rce.depthStencil = DebugDrawPass.depthStencilNoDepth
            await self.drawData.depthDisabledData.drawPoints(encoder: rce, buffer: buffer, bufferOffset: &bufferOffset)
        }
    }
}

extension DebugDraw.DebugDrawData {
    var requiredBufferCapacity : Int {
        return MemoryLayout<DebugDraw.DebugVertex>.size * (self.points.count + self.lines.count + self.wireframeTriangles.count + self.filledTriangles.count) + MemoryLayout<Float>.size * self.pointSizes.count
    }
    
    func drawPoints<R>(encoder: TypedRenderCommandEncoder<R>, buffer: Buffer, bufferOffset: inout Int) async {
        let pointsSize = self.points.count * MemoryLayout<DebugDraw.DebugVertex>.size
        if pointsSize > 0 {
            buffer.withMutableContents(range: bufferOffset..<(bufferOffset + self.points.count * MemoryLayout<DebugDraw.DebugVertex>.size)) { contents, _ in
                contents.copyMemory(from: UnsafeRawBufferPointer(start: self.points.buffer, count: self.points.count * MemoryLayout<DebugDraw.DebugVertex>.size))
            }
            encoder.setVertexBufferOffset(bufferOffset, index: 0)
            bufferOffset += pointsSize
            
            let pointSizesSize = self.pointSizes.count * MemoryLayout<Float>.size
            buffer.withMutableContents(range: bufferOffset..<bufferOffset + pointSizesSize) { contents, _ in
                contents.copyMemory(from: UnsafeRawBufferPointer(start: self.pointSizes.buffer, count: pointSizesSize))
            }
            encoder.setVertexBufferOffset(bufferOffset, index: 1)
            bufferOffset += pointSizesSize
            
            await encoder.drawPrimitives(type: .point, vertexStart: 0, vertexCount: self.points.count)
        }
    }
    
    func drawNonPoints<R>(encoder: TypedRenderCommandEncoder<R>, buffer: Buffer, bufferOffset: inout Int) async {
        let linesSize = self.lines.count * MemoryLayout<DebugDraw.DebugVertex>.size
        if linesSize > 0 {
            buffer.withMutableContents(range: bufferOffset..<(self.lines.count * MemoryLayout<DebugDraw.DebugVertex>.stride)) { contents, _ in
                contents.copyMemory(from: UnsafeRawBufferPointer(start: self.lines.buffer, count: self.lines.count * MemoryLayout<DebugDraw.DebugVertex>.size))
            }
            encoder.setVertexBufferOffset(bufferOffset, index: 0)
            bufferOffset += linesSize
            
            await encoder.drawPrimitives(type: .line, vertexStart: 0, vertexCount: self.lines.count)
        }
        
        let filledTrianglesSize = self.filledTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size
        if filledTrianglesSize > 0 {
            encoder.pipeline.descriptor.fillMode = .fill
            
            buffer.withMutableContents(range: bufferOffset..<(bufferOffset + self.filledTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size)) { contents, _ in
                contents.copyMemory(from: UnsafeRawBufferPointer(start: self.filledTriangles.buffer, count: self.filledTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size))
            }
            
            encoder.setVertexBufferOffset(bufferOffset, index: 0)
            bufferOffset += filledTrianglesSize
            
            await encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: self.filledTriangles.count)
        }
        
        let wireframeTrianglesSize = self.wireframeTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size
        if wireframeTrianglesSize > 0 {
            encoder.pipeline.descriptor.fillMode = .lines
            
            buffer.withMutableContents(range: bufferOffset..<(bufferOffset + self.wireframeTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size)) { contents, _ in
                contents.copyMemory(from: UnsafeRawBufferPointer(start: self.wireframeTriangles.buffer, count: self.wireframeTriangles.count * MemoryLayout<DebugDraw.DebugVertex>.size))
            }
            encoder.setVertexBufferOffset(bufferOffset, index: 0)
            bufferOffset += wireframeTrianglesSize
            
            await encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: self.wireframeTriangles.count)
        }
    }
}