Pain points

- Add Self parameter to TransformStack.
- Replace duck interface impl of TransformStack for PoseStack with
  lazily constructed wrapper object.
- Document Models.
- Replace PARTIAL_DIR with TRANSFORMED_PARTIAl and allow for arbitrary
  transformations with arbitrary keys.
  - Will be useful for transforming by different enums, or using
    alternative transforms by Direction.
- Add default impl of #delete to Instance.
- Store VisualizationContext in AbstractVisual.
- Reduce usage of @SuppressWarnings in Rotate and Transform.

src/main/java/com/jozufozu/flywheel/api/instance/Instance.java

@@ -4,4 +4,12 @@ public interface Instance {
 	InstanceType<?> type();
 
 	InstanceHandle handle();
+
+	default void delete() {
+		handle().setDeleted();
+	}
+
+	default void setChanged() {
+		handle().setChanged();
+	}
 }

src/main/java/com/jozufozu/flywheel/extension/PoseStackExtension.java

@@ -0,0 +1,16 @@
+package com.jozufozu.flywheel.extension;
+
+import com.jozufozu.flywheel.lib.transform.PoseTransformStack;
+import com.mojang.blaze3d.vertex.PoseStack;
+
+/**
+ * An extension interface for {@link PoseStack} that provides a {@link PoseTransformStack} wrapper.
+ * <br>
+ * Each PoseStack lazily creates and saves a wrapper instance. This wrapper is cached and reused for all future calls.
+ */
+public interface PoseStackExtension {
+	/**
+	 * @return The {@link PoseTransformStack} wrapper for this {@link PoseStack}.
+	 */
+	PoseTransformStack flywheel$transformStack();
+}

src/main/java/com/jozufozu/flywheel/lib/instance/AbstractInstance.java

@@ -23,11 +23,15 @@ public abstract class AbstractInstance implements Instance {
 		return handle;
 	}
 
+	@Override
 	public final void setChanged() {
+		// Override to mark final.
 		handle.setChanged();
 	}
 
+	@Override
 	public final void delete() {
+		// Override to mark final.
 		handle.setDeleted();
 	}
 }

src/main/java/com/jozufozu/flywheel/lib/model/Models.java

@@ -1,42 +1,93 @@
 package com.jozufozu.flywheel.lib.model;
 
+import java.util.function.BiConsumer;
+
 import com.jozufozu.flywheel.api.model.Model;
 import com.jozufozu.flywheel.lib.model.baked.BakedModelBuilder;
 import com.jozufozu.flywheel.lib.model.baked.BlockModelBuilder;
 import com.jozufozu.flywheel.lib.model.baked.PartialModel;
 import com.jozufozu.flywheel.lib.transform.TransformStack;
-import com.jozufozu.flywheel.lib.util.Pair;
 import com.mojang.blaze3d.vertex.PoseStack;
 
 import net.minecraft.core.Direction;
 import net.minecraft.world.level.block.state.BlockState;
 
+/**
+ * A collection of methods for creating models from various sources.
+ * <br>
+ * All Models returned from this class are cached, so calling the same
+ * method with the same parameters will return the same object.
+ */
 public final class Models {
 	private static final ModelCache<BlockState> BLOCK_STATE = new ModelCache<>(it -> new BlockModelBuilder(it).build());
 	private static final ModelCache<PartialModel> PARTIAL = new ModelCache<>(it -> new BakedModelBuilder(it.get()).build());
-	private static final ModelCache<Pair<PartialModel, Direction>> PARTIAL_DIR = new ModelCache<>(it -> new BakedModelBuilder(it.first().get()).poseStack(createRotation(it.second())).build());
+	private static final ModelCache<TransformedPartial<?>> TRANSFORMED_PARTIAL = new ModelCache<>(TransformedPartial::create);
 
 	private Models() {
 	}
 
+	/**
+	 * Get a usable model for a given blockstate.
+	 *
+	 * @param state The blockstate you wish to render.
+	 * @return A model corresponding to how the given blockstate would appear in the world.
+	 */
 	public static Model block(BlockState state) {
 		return BLOCK_STATE.get(state);
 	}
 
+	/**
+	 * Get a usable model for a given partial model.
+	 * @param partial The partial model you wish to render.
+	 * @return A model built from the baked model the partial model represents.
+	 */
 	public static Model partial(PartialModel partial) {
 		return PARTIAL.get(partial);
 	}
 
-	public static Model partial(PartialModel partial, Direction dir) {
-		return PARTIAL_DIR.get(Pair.of(partial, dir));
+	/**
+	 * Get a usable model for a given partial model, transformed in some way.
+	 * <br>
+	 * In general, you should try to avoid captures in the transformer function,
+	 * i.e. prefer static method references over lambdas.
+	 *
+	 * @param partial     The partial model you wish to render.
+	 * @param key         A key that will be used to cache the transformed model.
+	 * @param transformer A function that will transform the model in some way.
+	 * @param <T>         The type of the key.
+	 * @return A model built from the baked model the partial model represents, transformed by the given function.
+	 */
+	public static <T> Model partial(PartialModel partial, T key, BiConsumer<T, PoseStack> transformer) {
+		return TRANSFORMED_PARTIAL.get(new TransformedPartial<>(partial, key, transformer));
 	}
 
-	private static PoseStack createRotation(Direction facing) {
-		PoseStack stack = new PoseStack();
+	/**
+	 * Get a usable model for a given partial model, transformed to face a given direction.
+	 * <br>
+	 * {@link Direction#NORTH} is considered the default direction and the corresponding transform will be a no-op.
+	 *
+	 * @param partial The partial model you wish to render.
+	 * @param dir The direction you wish the model to be rotated to.
+	 * @return A model built from the baked model the partial model represents, transformed to face the given direction.
+	 */
+	public static Model partial(PartialModel partial, Direction dir) {
+		return partial(partial, dir, Models::rotateAboutCenterToFace);
+	}
+
+	private static void rotateAboutCenterToFace(Direction facing, PoseStack stack) {
 		TransformStack.of(stack)
 				.center()
 				.rotateToFace(facing.getOpposite())
 				.uncenter();
-		return stack;
+	}
+
+	private record TransformedPartial<T>(PartialModel partial, T key, BiConsumer<T, PoseStack> transformer) {
+		private Model create() {
+			var stack = new PoseStack();
+			transformer.accept(key, stack);
+			return new BakedModelBuilder(partial.get())
+					.poseStack(stack)
+					.build();
+		}
 	}
 }

src/main/java/com/jozufozu/flywheel/lib/model/baked/PartialModel.java

@@ -28,7 +28,9 @@ public class PartialModel {
 	protected BakedModel bakedModel;
 
 	public PartialModel(ResourceLocation modelLocation) {
-		if (tooLate) throw new RuntimeException("PartialModel '" + modelLocation + "' loaded after ModelRegistryEvent");
+		if (tooLate) {
+			throw new RuntimeException("PartialModel '" + modelLocation + "' loaded after ModelRegistryEvent");
+		}
 
 		this.modelLocation = modelLocation;
 		ALL.add(this);
@@ -44,8 +46,9 @@ public class PartialModel {
 
 	public static void onModelBake(ModelEvent.BakingCompleted event) {
 		var modelRegistry = event.getModels();
-		for (PartialModel partial : ALL)
+		for (PartialModel partial : ALL) {
 			partial.set(modelRegistry.get(partial.getLocation()));
+		}
 	}
 
 	@NotNull

src/main/java/com/jozufozu/flywheel/lib/transform/PoseTransformStack.java

@@ -0,0 +1,68 @@
+package com.jozufozu.flywheel.lib.transform;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.joml.Matrix3f;
+import org.joml.Matrix4f;
+import org.joml.Quaternionf;
+
+import com.mojang.blaze3d.vertex.PoseStack;
+
+/**
+ * A wrapper around {@link PoseStack} that implements {@link TransformStack}.
+ * <br>
+ * Only one instance of this class should exist per {@link PoseStack}.
+ */
+public class PoseTransformStack implements TransformStack<PoseTransformStack> {
+	private final PoseStack stack;
+
+	@ApiStatus.Internal
+	public PoseTransformStack(PoseStack stack) {
+		this.stack = stack;
+	}
+
+	@Override
+	public PoseTransformStack rotate(Quaternionf quaternion) {
+		stack.mulPose(quaternion);
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack scale(float factorX, float factorY, float factorZ) {
+		stack.scale(factorX, factorY, factorZ);
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack mulPose(Matrix4f pose) {
+		stack.last()
+				.pose()
+				.mul(pose);
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack mulNormal(Matrix3f normal) {
+		stack.last()
+				.normal()
+				.mul(normal);
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack pushPose() {
+		stack.pushPose();
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack popPose() {
+		stack.popPose();
+		return this;
+	}
+
+	@Override
+	public PoseTransformStack translate(double x, double y, double z) {
+		stack.translate(x, y, z);
+		return this;
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/transform/Rotate.java

@@ -8,38 +8,35 @@ import com.mojang.math.Axis;
 
 import net.minecraft.core.Direction;
 
-public interface Rotate<Self extends Rotate<Self>> {
-	Self rotate(Quaternionf quaternion);
+public interface Rotate<S extends Rotate<S>> {
+	S rotate(Quaternionf quaternion);
 
-	default Self rotate(AxisAngle4f axisAngle) {
+	default S rotate(AxisAngle4f axisAngle) {
 		return rotate(new Quaternionf(axisAngle));
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotate(float radians, Vector3fc axis) {
+	default S rotate(float radians, Vector3fc axis) {
 		if (radians == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate(new Quaternionf().setAngleAxis(radians, axis.x(), axis.y(), axis.z()));
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotate(float radians, Axis axis) {
+	default S rotate(float radians, Axis axis) {
 		if (radians == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate(axis.rotation(radians));
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotate(float radians, Direction axis) {
+	default S rotate(float radians, Direction axis) {
 		if (radians == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate(radians, axis.step());
 	}
 
-	default Self rotate(float radians, Direction.Axis axis) {
+	default S rotate(float radians, Direction.Axis axis) {
 		return switch (axis) {
 		case X -> rotateX(radians);
 		case Y -> rotateY(radians);
@@ -47,71 +44,71 @@ public interface Rotate<Self extends Rotate<Self>> {
 		};
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotateDegrees(float degrees, Vector3fc axis) {
+	default S rotateDegrees(float degrees, Vector3fc axis) {
 		if (degrees == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate((float) Math.toRadians(degrees), axis);
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotateDegrees(float degrees, Axis axis) {
+	default S rotateDegrees(float degrees, Axis axis) {
 		if (degrees == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate(axis.rotationDegrees(degrees));
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotateDegrees(float degrees, Direction axis) {
+	default S rotateDegrees(float degrees, Direction axis) {
 		if (degrees == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate((float) Math.toRadians(degrees), axis);
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotateDegrees(float degrees, Direction.Axis axis) {
+	default S rotateDegrees(float degrees, Direction.Axis axis) {
 		if (degrees == 0) {
-			return (Self) this;
+			return self();
 		}
 		return rotate((float) Math.toRadians(degrees), axis);
 	}
 
-	default Self rotateX(float radians) {
+	default S rotateX(float radians) {
 		return rotate(radians, Axis.XP);
 	}
 
-	default Self rotateY(float radians) {
+	default S rotateY(float radians) {
 		return rotate(radians, Axis.YP);
 	}
 
-	default Self rotateZ(float radians) {
+	default S rotateZ(float radians) {
 		return rotate(radians, Axis.ZP);
 	}
 
-	default Self rotateXDegrees(float degrees) {
+	default S rotateXDegrees(float degrees) {
 		return rotateDegrees(degrees, Axis.XP);
 	}
 
-	default Self rotateYDegrees(float degrees) {
+	default S rotateYDegrees(float degrees) {
 		return rotateDegrees(degrees, Axis.YP);
 	}
 
-	default Self rotateZDegrees(float degrees) {
+	default S rotateZDegrees(float degrees) {
 		return rotateDegrees(degrees, Axis.ZP);
 	}
 
-	@SuppressWarnings("unchecked")
-	default Self rotateToFace(Direction facing) {
+	default S rotateToFace(Direction facing) {
 		return switch (facing) {
 		case DOWN -> rotateXDegrees(-90);
 		case UP -> rotateXDegrees(90);
-		case NORTH -> (Self) this;
+			case NORTH -> self();
 		case SOUTH -> rotateYDegrees(180);
 		case WEST -> rotateYDegrees(90);
 		case EAST -> rotateYDegrees(270);
 		};
 	}
+
+	@SuppressWarnings("unchecked")
+	default S self() {
+		return (S) this;
+	}
 }

src/main/java/com/jozufozu/flywheel/lib/transform/Transform.java

@@ -15,8 +15,7 @@ public interface Transform<Self extends Transform<Self>> extends Scale<Self>, Ro
 	Self mulNormal(Matrix3f normal);
 
 	default Self transform(Matrix4f pose, Matrix3f normal) {
-		mulPose(pose);
-		return mulNormal(normal);
+		return mulPose(pose).mulNormal(normal);
 	}
 
 	default Self transform(PoseStack stack) {
@@ -24,24 +23,18 @@ public interface Transform<Self extends Transform<Self>> extends Scale<Self>, Ro
 		return transform(last.pose(), last.normal());
 	}
 
-	@SuppressWarnings("unchecked")
 	default Self rotateCentered(Quaternionf q) {
-		center().rotate(q)
+		return center().rotate(q)
 				.uncenter();
-		return (Self) this;
 	}
 
-	@SuppressWarnings("unchecked")
 	default Self rotateCentered(float radians, Axis axis) {
-		center().rotate(radians, axis)
+		return center().rotate(radians, axis)
 				.uncenter();
-		return (Self) this;
 	}
 
-	@SuppressWarnings("unchecked")
 	default Self rotateCentered(float radians, Direction axis) {
-		center().rotate(radians, axis)
+		return center().rotate(radians, axis)
 				.uncenter();
-		return (Self) this;
 	}
 }

src/main/java/com/jozufozu/flywheel/lib/transform/TransformStack.java

@@ -1,13 +1,14 @@
 package com.jozufozu.flywheel.lib.transform;
 
+import com.jozufozu.flywheel.extension.PoseStackExtension;
 import com.mojang.blaze3d.vertex.PoseStack;
 
-public interface TransformStack extends Transform<TransformStack> {
-	TransformStack pushPose();
+public interface TransformStack<Self extends TransformStack<Self>> extends Transform<Self> {
+	Self pushPose();
 
-	TransformStack popPose();
+	Self popPose();
 
-	static TransformStack of(PoseStack stack) {
-		return (TransformStack) stack;
+	static PoseTransformStack of(PoseStack stack) {
+		return ((PoseStackExtension) stack).flywheel$transformStack();
 	}
 }

src/main/java/com/jozufozu/flywheel/lib/visual/AbstractBlockEntityVisual.java

@@ -4,6 +4,7 @@ import org.joml.FrustumIntersection;
 
 import com.jozufozu.flywheel.api.visual.BlockEntityVisual;
 import com.jozufozu.flywheel.api.visual.DynamicVisual;
+import com.jozufozu.flywheel.api.visual.PlannedVisual;
 import com.jozufozu.flywheel.api.visual.TickableVisual;
 import com.jozufozu.flywheel.api.visual.VisualFrameContext;
 import com.jozufozu.flywheel.api.visualization.VisualManager;
@@ -18,13 +19,12 @@ import net.minecraft.world.level.block.state.BlockState;
 
 /**
  * The layer between a {@link BlockEntity} and the Flywheel backend.
- *
- * <br><br> {@link #updateLight()} is called after initialization.
- *
- * <br><br> There are a few additional features that overriding classes can opt in to:
+ * <br>
+ * <br> There are a few additional features that overriding classes can opt in to:
  * <ul>
  *     <li>{@link DynamicVisual}</li>
  *     <li>{@link TickableVisual}</li>
+ *     <li>{@link PlannedVisual}</li>
  * </ul>
  * See the interfaces' documentation for more information about each one.
  *
@@ -70,13 +70,15 @@ public abstract class AbstractBlockEntityVisual<T extends BlockEntity> extends A
 	}
 
 	/**
+	 * Check if this visual is within the given frustum.
 	 * @param frustum The current frustum.
-	 * @return {@code true} if this visual within the given frustum.
+	 * @return {@code true} if this visual is possibly visible.
 	 */
 	public boolean isVisible(FrustumIntersection frustum) {
 		float x = visualPos.getX() + 0.5f;
 		float y = visualPos.getY() + 0.5f;
 		float z = visualPos.getZ() + 0.5f;
+		// Default to checking a sphere exactly encompassing the block.
 		return frustum.testSphere(x, y, z, MoreMath.SQRT_3_OVER_2);
 	}
 

src/main/java/com/jozufozu/flywheel/lib/visual/AbstractEntityVisual.java

@@ -5,6 +5,7 @@ import org.joml.Vector3f;
 
 import com.jozufozu.flywheel.api.visual.DynamicVisual;
 import com.jozufozu.flywheel.api.visual.EntityVisual;
+import com.jozufozu.flywheel.api.visual.PlannedVisual;
 import com.jozufozu.flywheel.api.visual.TickableVisual;
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
 import com.jozufozu.flywheel.api.visualization.VisualizationManager;
@@ -19,11 +20,12 @@ import net.minecraft.world.phys.Vec3;
 
 /**
  * The layer between an {@link Entity} and the Flywheel backend.
- * *
- * <br><br> There are a few additional features that overriding classes can opt in to:
+ * <br>
+ * <br> There are a few additional features that overriding classes can opt in to:
  * <ul>
  *     <li>{@link DynamicVisual}</li>
  *     <li>{@link TickableVisual}</li>
+ *     <li>{@link PlannedVisual}</li>
  * </ul>
  * See the interfaces' documentation for more information about each one.
  *

src/main/java/com/jozufozu/flywheel/lib/visual/AbstractVisual.java

@@ -16,6 +16,12 @@ import net.minecraft.world.level.Level;
 import net.minecraft.world.level.LightLayer;
 
 public abstract class AbstractVisual implements Visual, LightListener {
+	/**
+	 * The visualization context used to construct this visual.
+	 * <br>
+	 * Useful for passing to child visuals.
+	 */
+	protected final VisualizationContext visualizationContext;
 	protected final InstancerProvider instancerProvider;
 	protected final Vec3i renderOrigin;
 	protected final Level level;
@@ -23,6 +29,7 @@ public abstract class AbstractVisual implements Visual, LightListener {
 	protected boolean deleted = false;
 
 	public AbstractVisual(VisualizationContext ctx, Level level) {
+		this.visualizationContext = ctx;
 		this.instancerProvider = ctx.instancerProvider();
 		this.renderOrigin = ctx.renderOrigin();
 		this.level = level;

src/main/java/com/jozufozu/flywheel/mixin/PoseStackMixin.java

@@ -1,40 +1,24 @@
 package com.jozufozu.flywheel.mixin;
 
-import org.joml.Quaternionf;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
 
-import com.jozufozu.flywheel.lib.transform.TransformStack;
+import com.jozufozu.flywheel.extension.PoseStackExtension;
+import com.jozufozu.flywheel.lib.transform.PoseTransformStack;
 import com.mojang.blaze3d.vertex.PoseStack;
 
 @Mixin(PoseStack.class)
-abstract class PoseStackMixin implements TransformStack {
-	@Override
-	public TransformStack pushPose() {
-		((PoseStack) (Object) this).pushPose();
-		return this;
-	}
+abstract class PoseStackMixin implements PoseStackExtension {
+	@Unique
+	private PoseTransformStack flywheel$wrapper;
 
 	@Override
-	public TransformStack popPose() {
-		((PoseStack) (Object) this).popPose();
-		return this;
-	}
-
-	@Override
-	public TransformStack scale(float factorX, float factorY, float factorZ) {
-		((PoseStack) (Object) this).scale(factorX, factorY, factorZ);
-		return this;
-	}
-
-	@Override
-	public TransformStack rotate(Quaternionf quaternion) {
-		((PoseStack) (Object) this).mulPose(quaternion);
-		return this;
-	}
-
-	@Override
-	public TransformStack translate(double x, double y, double z) {
-		((PoseStack) (Object) this).translate(x, y, z);
-		return this;
+	public PoseTransformStack flywheel$transformStack() {
+		if (flywheel$wrapper == null) {
+			// Thread safety: bless you if you're calling this from multiple threads, but as there is no state
+			// associated with the wrapper itself it's fine if we create multiple instances and one wins.
+			flywheel$wrapper = new PoseTransformStack((PoseStack) (Object) this);
+		}
+		return flywheel$wrapper;
 	}
 }