用 Metal 画一个三角形（Swift 函数式风格）

由于今年工作中用得语言换成 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 的子类。

自定义操作符

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

代码实现

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

最后执行出这种效果