用 Metal 画一个三角形(Swift 函数式风格)

4,398 阅读4分钟

由于今年工作中用得语言换成 Rust/OCaml/ReScript 啦,所以导致我现在写代码更倾向于写函数式风格的代码。
顺便试试 Swift 在函数式方面能达到啥好玩的程度。主要是我不会 Swift,仅仅为了好玩。

创建工程

随便创建个工程,小玩具就不打算跑在手机上了,因为我的设备是 ARM 芯片的,所以直接创建个 Mac 项目,记得勾上包含测试。

构建 MTKView 子类

现在来创建个 MTKView 的子类,其实我现在已经不接受这种所谓的面向对象,开发者用这种方式,就要写太多篇幅来描述一个上下文结构跟函数就能实现的动作。

import MetalKit

class MetalView: MTKView {
    required init(coder: NSCoder) {
        super.init(coder: coder)
        device = MTLCreateSystemDefaultDevice()
        render()
    }
}

extension MetalView {
    func render() {
        // TODO: 具体实现
    }
}

我们这里给 MetalView extension 了一个 render 函数,里面是后续要写得具体实现。

普通的方式画一个三角形

先用常见的方式来画一个三角形

class MetalView: MTKView {
    required init(coder: NSCoder) {
        super.init(coder: coder)
        device = MTLCreateSystemDefaultDevice()
        render()
    }
}

extension MetalView {
    func render() {
        guard let device = device else { fatalError("Failed to find default device.") }
        let vertexData: [Float] = [
            -1.0, -1.0, 0.0, 1.0,
             1.0, -1.0, 0.0, 1.0,
             0.0,  1.0, 0.0, 1.0
        ]

        let dataSize = vertexData.count * MemoryLayout<Float>.size
        let vertexBuffer = device.makeBuffer(bytes: vertexData, length: dataSize, options: [])
        let library = device.makeDefaultLibrary()
        let renderPassDesc = MTLRenderPassDescriptor()
        let renderPipelineDesc = MTLRenderPipelineDescriptor()
        if let currentDrawable = currentDrawable, let library = library {
            renderPassDesc.colorAttachments[0].texture = currentDrawable.texture
            renderPassDesc.colorAttachments[0].clearColor = MTLClearColor(red: 0.0, green: 0.5, blue: 0.5, alpha: 1.0)
            renderPassDesc.colorAttachments[0].loadAction = .clear
            renderPipelineDesc.vertexFunction = library.makeFunction(name: "vertexFn")
            renderPipelineDesc.fragmentFunction = library.makeFunction(name: "fragmentFn")
            renderPipelineDesc.colorAttachments[0].pixelFormat = .bgra8Unorm
            let commandQueue = device.makeCommandQueue()
            guard let commandQueue = commandQueue else { fatalError("Failed to make command queue.") }
            let commandBuffer = commandQueue.makeCommandBuffer()
            guard let commandBuffer = commandBuffer else { fatalError("Failed to make command buffer.") }
            let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDesc)
            guard let encoder = encoder else { fatalError("Failed to make render command encoder.") }
            if let renderPipelineState = try? device.makeRenderPipelineState(descriptor: renderPipelineDesc) {
                encoder.setRenderPipelineState(renderPipelineState)
                encoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
                encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 3, instanceCount: 1)
                encoder.endEncoding()
                commandBuffer.present(currentDrawable)
                commandBuffer.commit()
            }
        }
    }
}

然后是我们需要注册的 Shader 两个函数

#include <metal_stdlib>

using namespace metal;

struct Vertex {
    float4 position [[position]];
};

vertex Vertex vertexFn(constant Vertex *vertices [[buffer(0)]], uint vid [[vertex_id]]) {
    return vertices[vid];
}

fragment float4 fragmentFn(Vertex vert [[stage_in]]) {
    return float4(0.7, 1, 1, 1);
}

在运行之前需要把 StoryBoard 控制器上的 View 改成我们写得这个 MTKView 的子类。

story board 设置

自定义操作符

函数式当然不是指可以定义操作符,但是没有这些操作符,感觉没有魂灵,所以先定义个管道符

代码实现

precedencegroup SingleForwardPipe {
    associativity: left
    higherThan: BitwiseShiftPrecedence
}

infix operator |> : SingleForwardPipe

func |> <T, U>(_ value: T, _ fn: ((T) -> U)) -> U {
    fn(value)
}

测试管道符

因为创建项目的时候,勾上了 include Tests,直接写点测试代码,执行测试。

final class using_metalTests: XCTestCase {
    // ...

    func testPipeOperator() throws {
        let add = { (a: Int) in
            return { (b: Int) in
                return a + b
            }
        }
        assert(10 |> add(11) == 21)
        let doSth = { 10 }
        assert(() |> doSth == 10)
    }
}

目前随便写个测试通过嘞。

Functional Programming

现在需要把上面的逻辑分割成小函数,事实上,因为 Cocoa 的基础是建立在面向对象上的,我们还是没法完全摆脱面向对象,目前先小范围应用它。

生成 MTLBuffer

先理一下逻辑,代码开始是创建顶点数据,生成 buffer

fileprivate let makeBuffer = { (device: MTLDevice) in
    let vertexData: [Float] = [
        -1.0, -1.0, 0.0, 1.0,
         1.0, -1.0, 0.0, 1.0,
         0.0,  1.0, 0.0, 1.0
    ]
    
    let dataSize = vertexData.count * MemoryLayout<Float>.size
    return device.makeBuffer(bytes: vertexData, length: dataSize, options: [])
}

创建 MTLLibrary

接着是创建 MTLLibrary 来注册两个 shader 方法,还创建了一个 MTLRenderPipelineDescriptor 对象用于创建 MTLRenderPipelineState,但是创建的 MTLLibrary 对象是一个 Optional 的,所以其实得有两步,总之先提取它再说吧

fileprivate let makeLib = { (device: MTLDevice) in device.makeDefaultLibrary() }

抽象 map 函数

根据我们有限的函数式编程经验,像 Optional 这种对象大概率有一个 map 函数,所以我们自家实现一个,同时还要写成柯里化的(建议自动柯里化语法糖入常),因为这里有逃逸闭包,所以要加上 @escaping

func map<T, U>(_ transform: @escaping (T) throws -> U) rethrows -> (T?) -> U? {
    return { (o: T?) in
        return try? o.map(transform)
    }
}

处理 MTLRenderPipelineState

这里最终目的就是 new 了一个 MTLRenderPipelineState,顺带处理把程序的一些上下文给渲染管线描述器(MTLRenderPipelineDescriptor),譬如我们用到的着色器(Shader)函数,像素格式。 最后一行直接 try! 不处理错误啦,反正出问题直接会抛出来的

fileprivate let makeState = { (device: MTLDevice) in
    return { (lib: MTLLibrary) in
        let renderPipelineDesc = MTLRenderPipelineDescriptor()
        renderPipelineDesc.vertexFunction = lib.makeFunction(name: "vertexFn")
        renderPipelineDesc.fragmentFunction = lib.makeFunction(name: "fragmentFn")
        renderPipelineDesc.colorAttachments[0].pixelFormat = .bgra8Unorm
        return (try! device.makeRenderPipelineState(descriptor: renderPipelineDesc))
    }
}

暂时收尾

已经不想再抽取函数啦,其实还能更细粒度地处理,因为函数式有个纯函数跟副作用的概念,像 Haskell 里是可以用 Monad 来处理副作用的情况,这个主题留给后续吧。先把 render 改造一下

fileprivate let render = { (device: MTLDevice, currentDrawable: CAMetalDrawable?) in
    return { state in
        let renderPassDesc = MTLRenderPassDescriptor()
        if let currentDrawable = currentDrawable {
            renderPassDesc.colorAttachments[0].texture = currentDrawable.texture
            renderPassDesc.colorAttachments[0].clearColor = MTLClearColor(red: 0.0, green: 0.5, blue: 0.5, alpha: 1.0)
            renderPassDesc.colorAttachments[0].loadAction = .clear
            let commandQueue = device.makeCommandQueue()
            guard let commandQueue = commandQueue else { fatalError("Failed to make command queue.") }
            let commandBuffer = commandQueue.makeCommandBuffer()
            guard let commandBuffer = commandBuffer else { fatalError("Failed to make command buffer.") }
            let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDesc)
            guard let encoder = encoder else { fatalError("Failed to make render command encoder.") }
            encoder.setRenderPipelineState(state)
            encoder.setVertexBuffer(device |> makeBuffer, offset: 0, index: 0)
            encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 3, instanceCount: 1)
            encoder.endEncoding()
            commandBuffer.present(currentDrawable)
            commandBuffer.commit()
        }
    }
}

然后再调用,于是就变成下面这副鸟样子

class MetalView: MTKView {
    required init(coder: NSCoder) {
        super.init(coder: coder)
        device = MTLCreateSystemDefaultDevice()
        device |> map {
            makeLib($0)
            |> map(makeState($0))
            |> map(render($0, self.currentDrawable))
        }
    }
}

最后执行出这种效果

执行结果