From 3cd9aca62b0ab1bd69fea68c23ca60e5953f2f15 Mon Sep 17 00:00:00 2001
From: Eryn Wells <eryn@erynwells.me>
Date: Thu, 22 Nov 2018 09:35:01 -0700
Subject: [PATCH] Schedule the vertex normal update kernel

---
 Terrain2/Renderer.swift                  |  2 +-
 Terrain2/Shaders/TerrainAlgorithms.metal |  7 +++++--
 Terrain2/Terrain.swift                   | 22 ++++++++++++++++++++--
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/Terrain2/Renderer.swift b/Terrain2/Renderer.swift
index 74752cf..f54c138 100644
--- a/Terrain2/Renderer.swift
+++ b/Terrain2/Renderer.swift
@@ -321,8 +321,8 @@ class Renderer: NSObject, MTKViewDelegate {
                         renderEncoder.setRenderPipelineState(normalPipelineState)
                         renderEncoder.setDepthStencilState(depthState)
 
-                        encodeVertexNormalsDrawCall(encoder: renderEncoder)
                         encodeFaceNormalsDrawCall(encoder: renderEncoder)
+                        encodeVertexNormalsDrawCall(encoder: renderEncoder)
 
                         renderEncoder.endEncoding()
                     }
diff --git a/Terrain2/Shaders/TerrainAlgorithms.metal b/Terrain2/Shaders/TerrainAlgorithms.metal
index 7d7c06d..b6a14a9 100644
--- a/Terrain2/Shaders/TerrainAlgorithms.metal
+++ b/Terrain2/Shaders/TerrainAlgorithms.metal
@@ -50,9 +50,12 @@ kernel void updateGeometryNormals(constant packed_float3 *meshPositions [[buffer
     faceMidpoints[tid] = (1.0 / 3.0) * (v1 + v2 + v3);
 }
 
-kernel void updateGeometryVertexNormals()
+kernel void updateGeometryVertexNormals(constant packed_float3 *meshPositions [[buffer(GeneratorBufferIndexMeshPositions)]],
+                                        constant packed_ushort3 *indexes [[buffer(GeneratorBufferIndexIndexes)]],
+                                        constant packed_float3 *faceNormals [[buffer(GeneratorBufferIndexFaceNormals)]],
+                                        device packed_float3 *vertexNormals [[buffer(GeneratorBufferIndexNormals)]],
+                                        uint2 tid [[thread_position_in_grid]])
 {
-    
 }
 
 #pragma mark - ZeroGenerator
diff --git a/Terrain2/Terrain.swift b/Terrain2/Terrain.swift
index 8f193f2..b34396c 100644
--- a/Terrain2/Terrain.swift
+++ b/Terrain2/Terrain.swift
@@ -82,6 +82,7 @@ class Terrain: NSObject {
 
     private let updateHeightsPipeline: MTLComputePipelineState
     private let updateSurfaceNormalsPipeline: MTLComputePipelineState
+    private let updateVertexNormalsPipeline: MTLComputePipelineState
 
     let dimensions: float2
     let segments: uint2
@@ -114,6 +115,7 @@ class Terrain: NSObject {
         do {
             updateHeightsPipeline = try Terrain.computePipeline(withFunctionNamed: "updateGeometryHeights", device: device, library: library)
             updateSurfaceNormalsPipeline = try Terrain.computePipeline(withFunctionNamed: "updateGeometryNormals", device: device, library: library)
+            updateVertexNormalsPipeline = try Terrain.computePipeline(withFunctionNamed: "updateGeometryVertexNormals", device: device, library: library)
         } catch {
             print("Unable to create compute pipelines for terrain geometry updates. Error: \(error)")
             return nil
@@ -158,7 +160,7 @@ class Terrain: NSObject {
         if let computeEncoder = commandBuffer.makeComputeCommandEncoder() {
             //print("Scheduling update geometry heights")
             computeEncoder.label = "Geometry Heights Encoder"
-            computeEncoder.pushDebugGroup("Update Geometry: Heights")
+            computeEncoder.pushDebugGroup("Update Heights")
             computeEncoder.setComputePipelineState(updateHeightsPipeline)
             computeEncoder.setTexture(generator.outTexture, index: GeneratorTextureIndex.in.rawValue)
             let vertexBuffer = mesh.vertexBuffers[BufferIndex.meshPositions.rawValue]
@@ -174,7 +176,7 @@ class Terrain: NSObject {
         if let computeEncoder = commandBuffer.makeComputeCommandEncoder() {
             //print("Scheduling update geometry normals")
             computeEncoder.label = "Surface Normals Encoder"
-            computeEncoder.pushDebugGroup("Update Geometry: Surface Normals")
+            computeEncoder.pushDebugGroup("Update Surface Normals")
             computeEncoder.setComputePipelineState(updateSurfaceNormalsPipeline)
             let indexBuffer = mesh.submeshes[0].indexBuffer
             computeEncoder.setBuffer(indexBuffer.buffer, offset: indexBuffer.offset, index: GeneratorBufferIndex.indexes.rawValue)
@@ -186,6 +188,22 @@ class Terrain: NSObject {
             computeEncoder.popDebugGroup()
             computeEncoder.endEncoding()
         }
+
+        if let computeEncoder = commandBuffer.makeComputeCommandEncoder() {
+            computeEncoder.label = "Vertex Normals Encoder"
+            computeEncoder.pushDebugGroup("Update Vertex Normals")
+            computeEncoder.setComputePipelineState(updateVertexNormalsPipeline)
+            let indexBuffer = mesh.submeshes[0].indexBuffer
+            computeEncoder.setBuffer(indexBuffer.buffer, offset: indexBuffer.offset, index: GeneratorBufferIndex.indexes.rawValue)
+            let positionsBuffer = mesh.vertexBuffers[BufferIndex.meshPositions.rawValue]
+            computeEncoder.setBuffer(positionsBuffer.buffer, offset: positionsBuffer.offset, index: GeneratorBufferIndex.meshPositions.rawValue)
+            computeEncoder.setBuffer(faceNormalsBuffer, offset: 0, index: GeneratorBufferIndex.faceNormals.rawValue)
+            let normalsBuffer = mesh.vertexBuffers[BufferIndex.normals.rawValue]
+            computeEncoder.setBuffer(normalsBuffer.buffer, offset: normalsBuffer.offset, index: GeneratorBufferIndex.normals.rawValue)
+            computeEncoder.dispatchThreads(MTLSize(width: 2 * Int(segments.x * segments.y), height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: 8, height: 8, depth: 1))
+            computeEncoder.popDebugGroup()
+            computeEncoder.endEncoding()
+        }
     }
 
     private func populateInitialFaceNormals() {