Kind of update dynamic buffer source, but question is, do we still need it

src/main/kotlin/ru/dbotthepony/mc/otm/client/render/DynamicBufferSource.kt

@@ -2,6 +2,7 @@ package ru.dbotthepony.mc.otm.client.render
 
 import com.google.common.collect.ImmutableList
 import com.mojang.blaze3d.vertex.BufferBuilder
+import com.mojang.blaze3d.vertex.ByteBufferBuilder
 import com.mojang.blaze3d.vertex.VertexConsumer
 import com.mojang.blaze3d.vertex.VertexSorting
 import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
@@ -9,6 +10,10 @@ import it.unimi.dsi.fastutil.objects.ReferenceArraySet
 import net.minecraft.client.renderer.MultiBufferSource
 import net.minecraft.client.renderer.RenderType
 import net.minecraft.client.renderer.Sheets
+import net.minecraft.client.resources.model.ModelBakery
+import net.neoforged.fml.ModLoader
+import net.neoforged.neoforge.client.event.RegisterRenderBuffersEvent
+import java.io.Closeable
 
 private fun equals(existing: ImmutableList<RenderType>?, types: Array<out RenderType>): Boolean {
 	if (types.isEmpty()) {
@@ -63,7 +68,11 @@ private fun equals(existing: ImmutableList<RenderType>?, types: ImmutableList<Re
  *
  * Allows to batch OTM's geometry
  */
-class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInitialBufferSize: Int = Int.MAX_VALUE, val vertexSorting: VertexSorting = VertexSorting.DISTANCE_TO_ORIGIN) : MultiBufferSource {
+class DynamicBufferSource(
+	val minimalInitialBufferSize: Int = 0,
+	val maximalInitialBufferSize: Int = Int.MAX_VALUE,
+	val vertexSorting: VertexSorting = VertexSorting.DISTANCE_TO_ORIGIN
+) : MultiBufferSource, Closeable {
 	init {
 		require(minimalInitialBufferSize >= 0) { "Invalid minimal buffer size $minimalInitialBufferSize" }
 		require(maximalInitialBufferSize >= minimalInitialBufferSize) { "Maximal buffer size $maximalInitialBufferSize must be greater or equal to minimal buffer size $minimalInitialBufferSize" }
@@ -73,16 +82,17 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 		val type: RenderType,
 		var after: ImmutableList<RenderType>? = null,
 		immutableAfter: Boolean = false,
-		chained: Boolean = true
+		chained: Boolean = true,
+		private val byteBuffer: ByteBufferBuilder = ByteBufferBuilder(type.bufferSize().coerceIn(minimalInitialBufferSize, maximalInitialBufferSize))
 	) {
 		val immutableAfter: Boolean
 		val chained: Boolean
 
 		init {
 			if (
-				type.name == "forge_text" ||
+				type.name == "neoforge_text" ||
 				type.name == "text_intensity" ||
-				type.name == "forge_text_see_through" ||
+				type.name == "neoforge_text_see_through" ||
 				type.name.contains("font_") ||
 				type.name.contains("text_") ||
 				type.name.contains("_text") ||
@@ -95,54 +105,87 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 				this.immutableAfter = immutableAfter
 				this.chained = chained
 			}
+
+			buffers[type] = this
+			bufferList.add(this)
 		}
 
 		val dependants = ArrayList<State>(0)
 		var priority = -1
-		var dirty: Boolean = false
+
-		val builder by lazy(LazyThreadSafetyMode.NONE) {
+		private var buffer: BufferBuilder? = null
-			BufferBuilder(type.bufferSize().coerceAtLeast(minimalInitialBufferSize).coerceAtMost(maximalInitialBufferSize))
+
+		fun begin(): BufferBuilder {
+			var buffer = buffer
+
+			if (buffer != null) {
+				return buffer
+			}
+
+			buffer = BufferBuilder(byteBuffer, type.mode, type.format)
+			this.buffer = buffer
+			return buffer
+		}
+
+		fun finish() {
+			val built = buffer?.build()
+
+			if (built != null) {
+				if (type.sortOnUpload()) {
+					built.sortQuads(byteBuffer, vertexSorting)
+				}
+
+				type.draw(built)
+			}
+
+			this.buffer = null
+		}
+
+		fun finishChain() {
+			finish()
+			dependants.forEach { it.finishChain() }
+		}
+
+		fun close() {
+			byteBuffer.close()
 		}
 	}
 
-	@Suppress("BooleanLiteralArgument")
+	private val buffers = Reference2ObjectOpenHashMap<RenderType, State>()
-	private val bufferList = ArrayList<State>().also {
+	private val bufferList = ArrayList<State>()
-		var next = State(Sheets.solidBlockSheet(), null, true, false)
+
-		it.add(next)
+	init {
-		next = State(Sheets.cutoutBlockSheet(), ImmutableList.of(next.type), true, false)
+		var next = State(Sheets.solidBlockSheet(), null, immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(Sheets.cutoutBlockSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.bannerSheet(), ImmutableList.of(next.type), true, false)
+		next = State(Sheets.bannerSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(Sheets.translucentCullBlockSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.translucentCullBlockSheet(), ImmutableList.of(next.type), true, false)
+
-		it.add(next)
+		next = State(Sheets.shieldSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.shieldSheet(), ImmutableList.of(next.type), true, false)
+		next = State(Sheets.bedSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(Sheets.shulkerBoxSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.bedSheet(), ImmutableList.of(next.type), true, false)
+		next = State(Sheets.signSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(Sheets.hangingSignSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.shulkerBoxSheet(), ImmutableList.of(next.type), true, false)
+		next = State(Sheets.chestSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(RenderType.armorEntityGlint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.signSheet(), ImmutableList.of(next.type), true, false)
+		next = State(RenderType.glint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(RenderType.glintTranslucent(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(Sheets.chestSheet(), ImmutableList.of(next.type), true, false)
+		next = State(RenderType.entityGlint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+		next = State(RenderType.entityGlintDirect(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(RenderType.translucentNoCrumbling(), ImmutableList.of(next.type), true, false)
+		next = State(RenderType.waterMask(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		it.add(next)
+
-		next = State(RenderType.armorGlint(), ImmutableList.of(next.type), true, false)
+		ModelBakery.DESTROY_TYPES.forEach {
-		it.add(next)
+			next = State(it, ImmutableList.of(next.type), immutableAfter = true, chained = false)
-		next = State(RenderType.armorEntityGlint(), ImmutableList.of(next.type), true, false)
+		}
-		it.add(next)
+
-		next = State(RenderType.glint(), ImmutableList.of(next.type), true, false)
+		val modBuffers = LinkedHashMap<RenderType, ByteBufferBuilder>()
-		it.add(next)
+		ModLoader.postEvent(RegisterRenderBuffersEvent(modBuffers))
-		next = State(RenderType.glintDirect(), ImmutableList.of(next.type), true, false)
+
-		it.add(next)
+		for ((type, buffer) in modBuffers) {
-		next = State(RenderType.glintTranslucent(), ImmutableList.of(next.type), true, false)
+			// TODO: determine whenever other mods render types should be chained
-		it.add(next)
+			next = State(type, ImmutableList.of(next.type), immutableAfter = false, chained = true, buffer)
-		next = State(RenderType.entityGlint(), ImmutableList.of(next.type), true, false)
+		}
-		it.add(next)
+
-		next = State(RenderType.entityGlintDirect(), ImmutableList.of(next.type), true, false)
+		sortBufferList()
-		it.add(next)
-		next = State(RenderType.waterMask(), ImmutableList.of(next.type), true, false)
-		it.add(next)
 	}
 
 	private fun determineHeight(input: State, seen: MutableSet<State>) {
@@ -194,10 +237,6 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 		}
 	}
 
-	private val buffers = Reference2ObjectOpenHashMap<RenderType, State>().also {
-		for (v in bufferList) it[v.type] = v
-	}
-
 	init {
 		sortBufferList()
 	}
@@ -207,20 +246,9 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 
 		if (getState == null) {
 			getState = State(renderType)
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			getState.dirty = true
-			return getState.builder.also {
-				it.begin(renderType.mode(), renderType.format())
-			}
 		}
 
-		if (!getState.dirty) {
+		return getState.begin()
-			getState.dirty = true
-			getState.builder.begin(renderType.mode(), renderType.format())
-		}
-
-		return getState.builder
 	}
 
 	fun getBuffer(renderType: RenderType, after: RenderType): VertexConsumer {
@@ -228,13 +256,8 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 
 		if (getState == null) {
 			getState = State(renderType, ImmutableList.of(after))
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			getState.dirty = true
 			sortBufferList()
-			return getState.builder.also {
+			return getState.begin()
-				it.begin(renderType.mode(), renderType.format())
-			}
 		}
 
 		if (!getState.immutableAfter && getState.after?.getOrNull(0) !== after) {
@@ -242,12 +265,7 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 			sortBufferList()
 		}
 
-		if (!getState.dirty) {
+		return getState.begin()
-			getState.dirty = true
-			getState.builder.begin(renderType.mode(), renderType.format())
-		}
-
-		return getState.builder
 	}
 
 	fun getBuffer(renderType: RenderType, vararg after: RenderType): VertexConsumer {
@@ -255,13 +273,8 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 
 		if (getState == null) {
 			getState = State(renderType, ImmutableList.copyOf(after))
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			getState.dirty = true
 			sortBufferList()
-			return getState.builder.also {
+			return getState.begin()
-				it.begin(renderType.mode(), renderType.format())
-			}
 		}
 
 		if (!getState.immutableAfter && !equals(getState.after, after)) {
@@ -269,123 +282,46 @@ class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInit
 			sortBufferList()
 		}
 
-		if (!getState.dirty) {
+		return getState.begin()
-			getState.dirty = true
-			getState.builder.begin(renderType.mode(), renderType.format())
-		}
-
-		return getState.builder
-	}
-
-	fun getBuffer(renderType: RenderType, after: ImmutableList<RenderType>): VertexConsumer {
-		var getState = buffers[renderType]
-
-		if (getState == null) {
-			getState = State(renderType, after)
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			getState.dirty = true
-			sortBufferList()
-			return getState.builder.also {
-				it.begin(renderType.mode(), renderType.format())
-			}
-		}
-
-		if (!getState.immutableAfter && !equals(getState.after, after)) {
-			getState.after = after
-			sortBufferList()
-		}
-
-		if (!getState.dirty) {
-			getState.dirty = true
-			getState.builder.begin(renderType.mode(), renderType.format())
-		}
-
-		return getState.builder
 	}
 
 	fun registerType(renderType: RenderType, after: RenderType) {
-		var getState = buffers[renderType]
+		return registerType(renderType, ImmutableList.of(after))
-
-		if (getState == null) {
-			getState = State(renderType, ImmutableList.of(after), immutableAfter = true)
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			sortBufferList()
-		} else if (!getState.immutableAfter && getState.after?.getOrNull(0) !== after) {
-			getState.after = ImmutableList.of(after)
-			sortBufferList()
-		}
 	}
 
 	fun registerType(renderType: RenderType, vararg after: RenderType) {
-		var getState = buffers[renderType]
+		return registerType(renderType, ImmutableList.copyOf(after))
-
-		if (getState == null) {
-			getState = State(renderType, ImmutableList.copyOf(after), immutableAfter = true)
-			buffers[renderType] = getState
-			bufferList.add(getState)
-			sortBufferList()
-		} else if (!getState.immutableAfter && !equals(getState.after, after)) {
-			getState.after = ImmutableList.copyOf(after)
-			sortBufferList()
-		}
 	}
 
-	fun registerType(renderType: RenderType, after: ImmutableList<RenderType>) {
+	fun registerType(renderType: RenderType, after: List<RenderType>) {
+		val afterList = ImmutableList.copyOf(after)
 		var getState = buffers[renderType]
 
 		if (getState == null) {
-			getState = State(renderType, after, immutableAfter = true)
+			getState = State(renderType, afterList, immutableAfter = true)
 			buffers[renderType] = getState
 			bufferList.add(getState)
 			sortBufferList()
-		} else if (!getState.immutableAfter && !equals(getState.after, after)) {
+		} else if (!getState.immutableAfter && !equals(getState.after, afterList)) {
-			getState.after = after
+			getState.after = afterList
 			sortBufferList()
 		}
 	}
 
 	fun endBatch() {
-		for (state in bufferList) {
+		bufferList.forEach { it.finish() }
-			if (state.dirty) {
-				state.dirty = false
-				state.type.end(state.builder, vertexSorting)
-			}
-		}
 	}
 
 	fun endBatch(type: RenderType) {
-		val state = buffers[type] ?: return
+		buffers[type]?.finish()
-
-		if (state.dirty) {
-			state.dirty = false
-			type.end(state.builder, vertexSorting)
-		}
 	}
 
 	fun endBatchChain(type: RenderType) {
-		val state = buffers[type] ?: return
+		buffers[type]?.finishChain()
-
-		if (state.dirty) {
-			state.dirty = false
-			type.end(state.builder, vertexSorting)
-		}
-
-		for (ustate in state.dependants) {
-			endBatchChain(ustate)
-		}
 	}
 
-	private fun endBatchChain(state: State) {
+	override fun close() {
-		if (state.dirty) {
+		bufferList.forEach { it.close() }
-			state.dirty = false
-			state.type.end(state.builder, vertexSorting)
-		}
-
-		for (ustate in state.dependants) {
-			endBatchChain(ustate)
-		}
 	}
 
 	companion object {