refactor: introduce a custom render pass abstraction

2 files changed, 169 insertions, 55 deletions

diff --git a/src/main/kotlin/features/misc/CustomCapes.kt b/src/main/kotlin/features/misc/CustomCapes.kt
index 5004c15..d353b38 100644
--- a/src/main/kotlin/features/misc/CustomCapes.kt
+++ b/src/main/kotlin/features/misc/CustomCapes.kt
@@ -26,6 +26,7 @@ import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.gui.config.ManagedConfig
 import moe.nea.firmament.util.MC
 import moe.nea.firmament.util.TimeMark
+import moe.nea.firmament.util.mc.CustomRenderPassHelper
 
 object CustomCapes : FirmamentFeature {
 	override val identifier: String
@@ -72,63 +73,24 @@ object CustomCapes : FirmamentFeature {
 			matrixStack: MatrixStack,
 			model: (VertexConsumer) -> Unit
 		) {
-			// TODO: figure out how exactly the rotation abstraction works
 			val animationValue = (startTime.passedTime() / animationSpeed).mod(1F)
-			val commandEncoder = RenderSystem.getDevice().createCommandEncoder()
-			val animationBufSource = Std140Builder.intoBuffer(ByteBuffer.allocateDirect(16).order(ByteOrder.nativeOrder()))
-				.putFloat(animationValue.toFloat())
-				.align(16)
-				.get()
-			BufferAllocator(2048).use { allocator ->
-				val bufferBuilder = BufferBuilder(allocator, renderLayer.drawMode, renderLayer.vertexFormat)
-				model(bufferBuilder)
-				bufferBuilder.end().use { buffer ->
-					val vertexBuffer = renderLayer.vertexFormat.uploadImmediateVertexBuffer(buffer.buffer)
-					val indexBufferConstructor = RenderSystem.getSequentialBuffer(renderLayer.drawMode)
-					val indexBuffer = indexBufferConstructor.getIndexBuffer(buffer.drawParameters.indexCount)
-					val templateTexture = MC.textureManager.getTexture(template)
-					val backgroundTexture = MC.textureManager.getTexture(background)
-					val foregroundTexture = MC.textureManager.getTexture(overlay)
-					val dynamicTransforms = RenderSystem.getDynamicUniforms()
-						.write(
-							RenderSystem.getModelViewMatrix(),
-							Vector4f(1.0F, 1.0F, 1.0F, 1.0F),
-							RenderSystem.getModelOffset(),
-							RenderSystem.getTextureMatrix(),
-							RenderSystem.getShaderLineWidth()
-						)
-					val framebuffer = MC.instance.framebuffer
-					val colorAttachment =
-						RenderSystem.outputColorTextureOverride ?: framebuffer.getColorAttachmentView()
-					val depthAttachment =
-						(RenderSystem.outputDepthTextureOverride
-							?: framebuffer.getDepthAttachmentView()).takeIf { framebuffer.useDepthAttachment }
-
-					RenderSystem.getDevice()
-						.createBuffer({ "Firm Cape Animation" }, GpuBuffer.USAGE_UNIFORM or GpuBuffer.USAGE_MAP_READ, animationBufSource)
-						.use { animationUniformBuf ->
-							commandEncoder.createRenderPass(
-								{ "FirmamentCustomCape" },
-								colorAttachment,
-								OptionalInt.empty(),
-								depthAttachment,
-								OptionalDouble.empty(),
-							).use { renderPass ->
-								// TODO: account for lighting
-								renderPass.setPipeline(CustomRenderPipelines.PARALLAX_CAPE_SHADER)
-								renderPass.bindSampler("Sampler0", templateTexture.glTextureView)
-								renderPass.bindSampler("Sampler1", backgroundTexture.glTextureView)
-								renderPass.bindSampler("Sampler3", foregroundTexture.glTextureView)
-								RenderSystem.bindDefaultUniforms(renderPass)
-								renderPass.setUniform("DynamicTransforms", dynamicTransforms)
-								renderPass.setUniform("Animation", animationUniformBuf)
-								renderPass.setIndexBuffer(indexBuffer, indexBufferConstructor.indexType)
-								renderPass.setVertexBuffer(0, vertexBuffer)
-								renderPass.drawIndexed(0, 0, buffer.drawParameters.indexCount, 1)
-							}
-
-						}
+			CustomRenderPassHelper(
+				{ "Firmament Cape Renderer" },
+				renderLayer.drawMode,
+				renderLayer.vertexFormat,
+				MC.instance.framebuffer,
+				true,
+			).use { renderPass ->
+				renderPass.setPipeline(CustomRenderPipelines.PARALLAX_CAPE_SHADER)
+				renderPass.setAllDefaultUniforms()
+				renderPass.setUniform("Animation", 4) {
+					it.putFloat(animationValue.toFloat())
 				}
+				renderPass.bindSampler("Sampler0", template)
+				renderPass.bindSampler("Sampler1", background)
+				renderPass.bindSampler("Sampler3", overlay)
+				renderPass.uploadVertices(2048, model)
+				renderPass.draw()
 			}
 		}
 	}
diff --git a/src/main/kotlin/util/mc/CustomRenderPassHelper.kt b/src/main/kotlin/util/mc/CustomRenderPassHelper.kt
new file mode 100644
index 0000000..1bd1148
--- /dev/null
+++ b/src/main/kotlin/util/mc/CustomRenderPassHelper.kt
@@ -0,0 +1,152 @@
+package moe.nea.firmament.util.mc
+
+import com.mojang.blaze3d.buffers.GpuBuffer
+import com.mojang.blaze3d.buffers.GpuBufferSlice
+import com.mojang.blaze3d.buffers.Std140Builder
+import com.mojang.blaze3d.pipeline.RenderPipeline
+import com.mojang.blaze3d.systems.RenderPass
+import com.mojang.blaze3d.systems.RenderSystem
+import com.mojang.blaze3d.vertex.VertexFormat
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+import java.util.OptionalDouble
+import java.util.OptionalInt
+import org.joml.Vector4f
+import net.minecraft.client.gl.Framebuffer
+import net.minecraft.client.render.BufferBuilder
+import net.minecraft.client.render.BuiltBuffer
+import net.minecraft.client.texture.AbstractTexture
+import net.minecraft.client.util.BufferAllocator
+import net.minecraft.util.Identifier
+import net.minecraft.util.math.MathHelper
+import moe.nea.firmament.util.ErrorUtil
+import moe.nea.firmament.util.MC
+
+
+class CustomRenderPassHelper(
+	val labelSupplier: () -> String,
+	val drawMode: VertexFormat.DrawMode,
+	val vertexFormat: VertexFormat,
+	val frameBuffer: Framebuffer,
+	val hasDepth: Boolean,
+) : AutoCloseable {
+	private val scope = mutableListOf<AutoCloseable>()
+	private val preparations = mutableListOf<(RenderPass) -> Unit>()
+	val device = RenderSystem.getDevice()
+	private var hasPipelineAction = false
+	val commandEncoder = device.createCommandEncoder()
+	fun setPipeline(pipeline: RenderPipeline) {
+		ErrorUtil.softCheck("Already has a pipeline", !hasPipelineAction)
+		hasPipelineAction = true
+		queueAction {
+			it.setPipeline(pipeline)
+		}
+	}
+
+	fun bindSampler(name: String, texture: Identifier) {
+		bindSampler(name, MC.textureManager.getTexture(texture))
+	}
+
+	fun bindSampler(name: String, texture: AbstractTexture) {
+		queueAction { it.bindSampler(name, texture.glTextureView) }
+	}
+
+	fun setAllDefaultUniforms() {
+		queueAction {
+			RenderSystem.bindDefaultUniforms(it)
+		}
+		setUniform(
+			"DynamicTransforms", RenderSystem.getDynamicUniforms()
+				.write(
+					RenderSystem.getModelViewMatrix(),
+					Vector4f(1.0F, 1.0F, 1.0F, 1.0F),
+					RenderSystem.getModelOffset(),
+					RenderSystem.getTextureMatrix(),
+					RenderSystem.getShaderLineWidth()
+				)
+		)
+	}
+
+	fun setUniform(name: String, slice: GpuBufferSlice) = queueAction { it.setUniform(name, slice) }
+	fun setUniform(name: String, slice: GpuBuffer) = queueAction { it.setUniform(name, slice) }
+
+	fun setUniform(name: String, size: Int, labelSupplier: () -> String = { name }, init: (Std140Builder) -> Unit) {
+		val buffer = createUniformBuffer(labelSupplier, allocateByteBuf(size, init))
+		setUniform(name, buffer)
+	}
+
+	var vertices: BuiltBuffer? = null
+
+	fun uploadVertices(size: Int, init: (BufferBuilder) -> Unit) {
+		uploadVertices(
+			BufferBuilder(queueClose(BufferAllocator(size)), drawMode, vertexFormat)
+				.also(init)
+				.end()
+		)
+	}
+
+	fun uploadVertices(buffer: BuiltBuffer) {
+		queueClose(buffer)
+		ErrorUtil.softCheck("Vertices have already been uploaded", vertices == null)
+		vertices = buffer
+		val vertexBuffer = vertexFormat.uploadImmediateVertexBuffer(buffer.buffer)
+		val indexBufferConstructor = RenderSystem.getSequentialBuffer(drawMode)
+		val indexBuffer = indexBufferConstructor.getIndexBuffer(buffer.drawParameters.indexCount)
+		queueAction {
+			it.setIndexBuffer(indexBuffer, indexBufferConstructor.indexType)
+			it.setVertexBuffer(0, vertexBuffer)
+		}
+	}
+
+	fun createUniformBuffer(labelSupplier: () -> String, buffer: ByteBuffer): GpuBuffer {
+		return queueClose(
+			device.createBuffer(
+				labelSupplier::invoke,
+				GpuBuffer.USAGE_UNIFORM or GpuBuffer.USAGE_MAP_READ,
+				buffer
+			)
+		)
+	}
+
+	fun allocateByteBuf(size: Int, init: (Std140Builder) -> Unit): ByteBuffer {
+		return Std140Builder.intoBuffer(
+			ByteBuffer
+				.allocateDirect(MathHelper.roundUpToMultiple(size, 16))
+				.order(ByteOrder.nativeOrder())
+		).also(init).get()
+	}
+
+	fun queueAction(action: (RenderPass) -> Unit) {
+		preparations.add(action)
+	}
+
+	fun <T : AutoCloseable> queueClose(t: T): T = t.also { scope.add(it) }
+	override fun close() {
+		scope.reversed().forEach { it.close() }
+	}
+
+	object DrawToken
+
+	fun draw(): DrawToken {
+		val vertexData = (ErrorUtil.notNullOr(vertices, "No vertex data uploaded") { return DrawToken })
+		ErrorUtil.softCheck("Missing a pipeline", hasPipelineAction)
+		val renderPass = queueClose(
+			commandEncoder.createRenderPass(
+				labelSupplier::invoke,
+				RenderSystem.outputColorTextureOverride ?: frameBuffer.getColorAttachmentView(),
+				OptionalInt.empty(),
+				(RenderSystem.outputDepthTextureOverride
+					?: frameBuffer.getDepthAttachmentView()).takeIf { frameBuffer.useDepthAttachment && hasDepth },
+				OptionalDouble.empty()
+			)
+		)
+		preparations.forEach { it(renderPass) }
+		renderPass.drawIndexed(
+			0,
+			0,
+			vertexData.drawParameters.indexCount,
+			1
+		)
+		return DrawToken
+	}
+}