Better buffering

- Abstract ModelUtil.getBufferBuilder and .getBufferBuilderFromTemplate
 - BakedModelBuilder and WorldModelBuilder as parameter objects for getBufferBuilder
 - WorldModelBuilder supports IModelData

src/main/java/com/jozufozu/flywheel/core/PartialModel.java

@@ -4,6 +4,8 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+import javax.annotation.Nonnull;
+
 import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraftforge.client.event.ModelBakeEvent;
@@ -49,6 +51,12 @@ public class PartialModel {
 			partial.set(modelRegistry.get(partial.getLocation()));
 	}
 
+	@Nonnull
+	public String getName() {
+		return getLocation()
+				.toString();
+	}
+
 	protected void set(BakedModel bakedModel) {
 		this.bakedModel = bakedModel;
 	}

src/main/java/com/jozufozu/flywheel/core/model/BakedModelBuilder.java

@@ -0,0 +1,47 @@
+package com.jozufozu.flywheel.core.model;
+
+import java.util.Random;
+
+import com.jozufozu.flywheel.core.virtual.VirtualEmptyBlockGetter;
+import com.jozufozu.flywheel.core.virtual.VirtualEmptyModelData;
+import com.mojang.blaze3d.vertex.PoseStack;
+import com.mojang.blaze3d.vertex.VertexConsumer;
+
+import net.minecraft.client.renderer.block.ModelBlockRenderer;
+import net.minecraft.client.renderer.texture.OverlayTexture;
+import net.minecraft.client.resources.model.BakedModel;
+import net.minecraft.core.BlockPos;
+import net.minecraft.world.level.BlockAndTintGetter;
+import net.minecraft.world.level.block.Blocks;
+import net.minecraft.world.level.block.state.BlockState;
+
+public final class BakedModelBuilder implements Bufferable {
+	private final BakedModel model;
+	private BlockAndTintGetter renderWorld = VirtualEmptyBlockGetter.INSTANCE;
+	private BlockState referenceState = Blocks.AIR.defaultBlockState();
+	private PoseStack poseStack = new PoseStack();
+
+	public BakedModelBuilder(BakedModel model) {
+		this.model = model;
+	}
+
+	public BakedModelBuilder withRenderWorld(BlockAndTintGetter renderWorld) {
+		this.renderWorld = renderWorld;
+		return this;
+	}
+
+	public BakedModelBuilder withReferenceState(BlockState referenceState) {
+		this.referenceState = referenceState;
+		return this;
+	}
+
+	public BakedModelBuilder withPoseStack(PoseStack poseStack) {
+		this.poseStack = poseStack;
+		return this;
+	}
+
+	@Override
+	public void bufferInto(ModelBlockRenderer blockRenderer, VertexConsumer consumer, Random random) {
+		blockRenderer.tesselateBlock(renderWorld, model, referenceState, BlockPos.ZERO, poseStack, consumer, false, random, 42, OverlayTexture.NO_OVERLAY, VirtualEmptyModelData.INSTANCE);
+	}
+}

src/main/java/com/jozufozu/flywheel/core/model/BlockModel.java

@@ -7,7 +7,6 @@ import com.mojang.blaze3d.vertex.PoseStack;
 
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.resources.model.BakedModel;
-import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
 
 /**
@@ -35,11 +34,18 @@ public class BlockModel implements Model {
 	}
 
 	public BlockModel(PartialModel model, PoseStack ms) {
-		this(Formats.BLOCK.createReader(ModelUtil.getBufferBuilder(model.get(), Blocks.AIR.defaultBlockState(), ms)), model.getLocation().toString());
+		this(ModelUtil.bakedModel(model.get())
+				.withPoseStack(ms), model.getName());
 	}
 
 	public BlockModel(BakedModel model, BlockState referenceState, PoseStack ms) {
-		this(Formats.BLOCK.createReader(ModelUtil.getBufferBuilder(model, referenceState, ms)), referenceState.toString());
+		this(ModelUtil.bakedModel(model)
+				.withReferenceState(referenceState)
+				.withPoseStack(ms), referenceState.toString());
+	}
+
+	public BlockModel(BakedModelBuilder builder, String name) {
+		this(Formats.BLOCK.createReader(builder.build()), name);
 	}
 
 	public BlockModel(VertexList reader, String name) {

src/main/java/com/jozufozu/flywheel/core/model/Bufferable.java

@@ -0,0 +1,18 @@
+package com.jozufozu.flywheel.core.model;
+
+import java.util.Random;
+
+import com.mojang.blaze3d.vertex.VertexConsumer;
+
+import net.minecraft.client.renderer.block.ModelBlockRenderer;
+
+/**
+ * An interface for objects that can "rendered" into a BufferBuilder.
+ */
+public interface Bufferable {
+	void bufferInto(ModelBlockRenderer renderer, VertexConsumer consumer, Random random);
+
+	default ShadeSeparatedBufferBuilder build() {
+		return ModelUtil.getBufferBuilder(this);
+	}
+}

src/main/java/com/jozufozu/flywheel/core/model/ModelUtil.java

@@ -1,13 +1,10 @@
 package com.jozufozu.flywheel.core.model;
 
 import java.lang.reflect.Field;
-import java.util.Collection;
 import java.util.EnumMap;
 import java.util.Random;
 
 import com.jozufozu.flywheel.Flywheel;
-import com.jozufozu.flywheel.core.virtual.VirtualEmptyBlockGetter;
-import com.jozufozu.flywheel.core.virtual.VirtualEmptyModelData;
 import com.jozufozu.flywheel.util.transform.TransformStack;
 import com.mojang.blaze3d.vertex.BufferBuilder;
 import com.mojang.blaze3d.vertex.DefaultVertexFormat;
@@ -15,20 +12,11 @@ import com.mojang.blaze3d.vertex.PoseStack;
 import com.mojang.blaze3d.vertex.VertexFormat;
 
 import net.minecraft.client.Minecraft;
-import net.minecraft.client.renderer.ItemBlockRenderTypes;
 import net.minecraft.client.renderer.RenderType;
 import net.minecraft.client.renderer.block.BlockRenderDispatcher;
 import net.minecraft.client.renderer.block.ModelBlockRenderer;
-import net.minecraft.client.renderer.texture.OverlayTexture;
 import net.minecraft.client.resources.model.BakedModel;
-import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
-import net.minecraft.world.level.BlockAndTintGetter;
-import net.minecraft.world.level.block.RenderShape;
-import net.minecraft.world.level.block.state.BlockState;
-import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
-import net.minecraftforge.client.ForgeHooksClient;
-import net.minecraftforge.client.model.data.EmptyModelData;
 import net.minecraftforge.fml.util.ObfuscationReflectionHelper;
 
 public class ModelUtil {
@@ -56,72 +44,29 @@ public class ModelUtil {
 		return dispatcher;
 	}
 
-	public static ShadeSeparatedBufferBuilder getBufferBuilder(BakedModel model, BlockState referenceState, PoseStack poseStack) {
+	public static BakedModelBuilder bakedModel(BakedModel model) {
-		return getBufferBuilder(VirtualEmptyBlockGetter.INSTANCE, model, referenceState, poseStack);
+		return new BakedModelBuilder(model);
 	}
 
-	public static ShadeSeparatedBufferBuilder getBufferBuilder(BlockAndTintGetter renderWorld, BakedModel model, BlockState referenceState, PoseStack poseStack) {
+	public static WorldModelBuilder worldLayer(RenderType layer) {
+		return new WorldModelBuilder(layer);
+	}
+
+	public static ShadeSeparatedBufferBuilder getBufferBuilder(Bufferable object) {
 		ModelBlockRenderer blockRenderer = VANILLA_RENDERER.getModelRenderer();
 		ThreadLocalObjects objects = THREAD_LOCAL_OBJECTS.get();
 
-		ShadeSeparatingVertexConsumer shadeSeparatingWrapper = objects.shadeSeparatingWrapper;
 		ShadeSeparatedBufferBuilder builder = new ShadeSeparatedBufferBuilder(512);
-		BufferBuilder unshadedBuilder = objects.unshadedBuilder;
 
 		builder.begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.BLOCK);
-		unshadedBuilder.begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.BLOCK);
+		objects.unshadedBuilder.begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.BLOCK);
-		shadeSeparatingWrapper.prepare(builder, unshadedBuilder);
+		objects.shadeSeparatingWrapper.prepare(builder, objects.unshadedBuilder);
-		blockRenderer.tesselateBlock(renderWorld, model, referenceState, BlockPos.ZERO, poseStack, shadeSeparatingWrapper,
-				false, objects.random, 42, OverlayTexture.NO_OVERLAY, VirtualEmptyModelData.INSTANCE);
-		shadeSeparatingWrapper.clear();
-		unshadedBuilder.end();
-		builder.appendUnshadedVertices(unshadedBuilder);
-		builder.end();
 
-		return builder;
+		object.bufferInto(blockRenderer, objects.shadeSeparatingWrapper, objects.random);
-	}
 
-	public static ShadeSeparatedBufferBuilder getBufferBuilderFromTemplate(BlockAndTintGetter renderWorld, RenderType layer, Collection<StructureTemplate.StructureBlockInfo> blocks) {
+		objects.shadeSeparatingWrapper.clear();
-		return getBufferBuilderFromTemplate(renderWorld, layer, blocks, new PoseStack());
+		objects.unshadedBuilder.end();
-	}
+		builder.appendUnshadedVertices(objects.unshadedBuilder);
-
-	public static ShadeSeparatedBufferBuilder getBufferBuilderFromTemplate(BlockAndTintGetter renderWorld, RenderType layer, Collection<StructureTemplate.StructureBlockInfo> blocks, PoseStack poseStack) {
-		ModelBlockRenderer modelRenderer = VANILLA_RENDERER.getModelRenderer();
-		ThreadLocalObjects objects = THREAD_LOCAL_OBJECTS.get();
-
-		Random random = objects.random;
-		ShadeSeparatingVertexConsumer shadeSeparatingWrapper = objects.shadeSeparatingWrapper;
-		ShadeSeparatedBufferBuilder builder = new ShadeSeparatedBufferBuilder(512);
-		BufferBuilder unshadedBuilder = objects.unshadedBuilder;
-
-		builder.begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.BLOCK);
-		unshadedBuilder.begin(VertexFormat.Mode.QUADS, DefaultVertexFormat.BLOCK);
-		shadeSeparatingWrapper.prepare(builder, unshadedBuilder);
-
-		ForgeHooksClient.setRenderType(layer);
-		ModelBlockRenderer.enableCaching();
-		for (StructureTemplate.StructureBlockInfo info : blocks) {
-			BlockState state = info.state;
-
-			if (state.getRenderShape() != RenderShape.MODEL)
-				continue;
-			if (!ItemBlockRenderTypes.canRenderInLayer(state, layer))
-				continue;
-
-			BlockPos pos = info.pos;
-
-			poseStack.pushPose();
-			poseStack.translate(pos.getX(), pos.getY(), pos.getZ());
-			modelRenderer.tesselateBlock(renderWorld, VANILLA_RENDERER.getBlockModel(state), state, pos, poseStack, shadeSeparatingWrapper,
-					true, random, 42, OverlayTexture.NO_OVERLAY, EmptyModelData.INSTANCE);
-			poseStack.popPose();
-		}
-		ModelBlockRenderer.clearCache();
-		ForgeHooksClient.setRenderType(null);
-
-		shadeSeparatingWrapper.clear();
-		unshadedBuilder.end();
-		builder.appendUnshadedVertices(unshadedBuilder);
 		builder.end();
 
 		return builder;
@@ -153,4 +98,5 @@ public class ModelUtil {
 		public final ShadeSeparatingVertexConsumer shadeSeparatingWrapper = new ShadeSeparatingVertexConsumer();
 		public final BufferBuilder unshadedBuilder = new BufferBuilder(512);
 	}
+
 }

src/main/java/com/jozufozu/flywheel/core/model/WorldModel.java

@@ -19,7 +19,10 @@ public class WorldModel implements Model {
 	 * It is expected that {@code renderWorld.getShade(...)} returns a constant.
 	 */
 	public WorldModel(BlockAndTintGetter renderWorld, RenderType layer, Collection<StructureTemplate.StructureBlockInfo> blocks, String name) {
-		reader = Formats.BLOCK.createReader(ModelUtil.getBufferBuilderFromTemplate(renderWorld, layer, blocks));
+		reader = Formats.BLOCK.createReader(ModelUtil.worldLayer(layer)
+				.withBlocks(blocks)
+				.withRenderWorld(renderWorld)
+				.build());
 		this.name = name;
 	}
 

src/main/java/com/jozufozu/flywheel/core/model/WorldModelBuilder.java

@@ -0,0 +1,79 @@
+package com.jozufozu.flywheel.core.model;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Random;
+
+import com.jozufozu.flywheel.core.virtual.VirtualEmptyBlockGetter;
+import com.mojang.blaze3d.vertex.PoseStack;
+import com.mojang.blaze3d.vertex.VertexConsumer;
+
+import net.minecraft.client.renderer.ItemBlockRenderTypes;
+import net.minecraft.client.renderer.RenderType;
+import net.minecraft.client.renderer.block.ModelBlockRenderer;
+import net.minecraft.client.renderer.texture.OverlayTexture;
+import net.minecraft.core.BlockPos;
+import net.minecraft.world.level.BlockAndTintGetter;
+import net.minecraft.world.level.block.RenderShape;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
+import net.minecraftforge.client.ForgeHooksClient;
+import net.minecraftforge.client.model.data.EmptyModelData;
+import net.minecraftforge.client.model.data.IModelData;
+
+public final class WorldModelBuilder implements Bufferable {
+	private final RenderType layer;
+
+	private PoseStack poseStack = new PoseStack();
+	private Map<BlockPos, IModelData> modelData = Collections.emptyMap();
+	private BlockAndTintGetter renderWorld = VirtualEmptyBlockGetter.INSTANCE;
+	private Collection<StructureTemplate.StructureBlockInfo> blocks = Collections.emptyList();
+
+	public WorldModelBuilder(RenderType layer) {
+		this.layer = layer;
+	}
+
+	@Override
+	public void bufferInto(ModelBlockRenderer modelRenderer, VertexConsumer consumer, Random random) {
+		ForgeHooksClient.setRenderType(this.layer);
+		ModelBlockRenderer.enableCaching();
+		for (StructureTemplate.StructureBlockInfo info : this.blocks) {
+			BlockState state = info.state;
+
+			if (state.getRenderShape() != RenderShape.MODEL) continue;
+			if (!ItemBlockRenderTypes.canRenderInLayer(state, this.layer)) continue;
+
+			BlockPos pos = info.pos;
+
+			IModelData data = this.modelData.getOrDefault(pos, EmptyModelData.INSTANCE);
+
+			poseStack.pushPose();
+			poseStack.translate(pos.getX(), pos.getY(), pos.getZ());
+			modelRenderer.tesselateBlock(this.renderWorld, ModelUtil.VANILLA_RENDERER.getBlockModel(state), state, pos, poseStack, consumer, true, random, 42, OverlayTexture.NO_OVERLAY, data);
+			poseStack.popPose();
+		}
+		ModelBlockRenderer.clearCache();
+		ForgeHooksClient.setRenderType(null);
+	}
+
+	public WorldModelBuilder withRenderWorld(BlockAndTintGetter renderWorld) {
+		this.renderWorld = renderWorld;
+		return this;
+	}
+
+	public WorldModelBuilder withBlocks(Collection<StructureTemplate.StructureBlockInfo> blocks) {
+		this.blocks = blocks;
+		return this;
+	}
+
+	public WorldModelBuilder withModelData(Map<BlockPos, IModelData> modelData) {
+		this.modelData = modelData;
+		return this;
+	}
+
+	public WorldModelBuilder withPoseStack(PoseStack poseStack) {
+		this.poseStack = poseStack;
+		return this;
+	}
+}