Add tint override to texture packs

9 files changed, 276 insertions, 72 deletions

diff --git a/src/main/java/moe/nea/firmament/init/EarlyRiser.java b/src/main/java/moe/nea/firmament/init/EarlyRiser.java
index 5eab563..9734e94 100644
--- a/src/main/java/moe/nea/firmament/init/EarlyRiser.java
+++ b/src/main/java/moe/nea/firmament/init/EarlyRiser.java
@@ -7,5 +7,6 @@ public class EarlyRiser implements Runnable {
         new ClientPlayerRiser().addTinkerers();
         new HandledScreenRiser().addTinkerers();
         new SectionBuilderRiser().addTinkerers();
+		new ItemColorsSodiumRiser().addTinkerers();
     }
 }
diff --git a/src/main/java/moe/nea/firmament/init/ItemColorsSodiumRiser.java b/src/main/java/moe/nea/firmament/init/ItemColorsSodiumRiser.java
new file mode 100644
index 0000000..80ee9aa
--- /dev/null
+++ b/src/main/java/moe/nea/firmament/init/ItemColorsSodiumRiser.java
@@ -0,0 +1,64 @@
+package moe.nea.firmament.init;
+
+import me.shedaniel.mm.api.ClassTinkerers;
+import moe.nea.firmament.util.ErrorUtil;
+import net.fabricmc.loader.api.FabricLoader;
+import net.minecraft.client.color.item.ItemColorProvider;
+import net.minecraft.client.color.item.ItemColors;
+import net.minecraft.item.ItemStack;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.InsnNode;
+import org.objectweb.asm.tree.MethodInsnNode;
+import org.objectweb.asm.tree.VarInsnNode;
+
+public class ItemColorsSodiumRiser extends RiserUtils {
+	@IntermediaryName(ItemColors.class)
+	String ItemColors;
+	@IntermediaryName(ItemColorProvider.class)
+	String ItemColorProvider;
+	@IntermediaryName(ItemStack.class)
+	String ItemStack;
+	String getColorProvider = "sodium$getColorProvider";
+	Type getColorProviderDesc = Type.getMethodType(getTypeForClassName(ItemColorProvider),
+	                                               getTypeForClassName(ItemStack));
+
+	@Override
+	public void addTinkerers() {
+		ClassTinkerers.addTransformation(ItemColors, this::addSodiumOverride, true);
+	}
+
+	private void addSodiumOverride(ClassNode classNode) {
+		var node = findMethod(classNode, getColorProvider, getColorProviderDesc);
+		if (node == null) {
+			if (!FabricLoader.getInstance().isModLoaded("sodium"))
+				ErrorUtil.INSTANCE.softError("Sodium is present, but sodium color override could not be injected.");
+			return;
+		}
+		var p = node.instructions.getFirst();
+		while (p != null) {
+			if (p.getOpcode() == Opcodes.ARETURN) {
+				node.instructions.insertBefore(
+					p,
+					mkOverrideSodiumCall()
+				);
+			}
+			p = p.getNext();
+		}
+	}
+
+	private InsnList mkOverrideSodiumCall() {
+		var insnList = new InsnList();
+		insnList.add(new VarInsnNode(Opcodes.ALOAD, 0));
+		insnList.add(new InsnNode(Opcodes.SWAP));
+		insnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL,
+		                                getTypeForClassName(ItemColors).getInternalName(),
+		                                "overrideSodium_firmament",
+		                                Type.getMethodType(getTypeForClassName(ItemColorProvider),
+		                                                   getTypeForClassName(ItemColorProvider)).getDescriptor(),
+		                                false));
+		return insnList;
+	}
+}
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBasic.java b/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBasic.java
index 7bf3732..3ed2177 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBasic.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBasic.java
@@ -2,6 +2,7 @@
 package moe.nea.firmament.mixins.custommodels;
 
 import moe.nea.firmament.features.texturepack.BakedModelExtra;
+import moe.nea.firmament.features.texturepack.TintOverrides;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.BasicBakedModel;
 import org.jetbrains.annotations.Nullable;
@@ -11,18 +12,31 @@ import org.spongepowered.asm.mixin.Unique;
 @Mixin(BasicBakedModel.class)
 public class BakedModelDataHolderBasic implements BakedModelExtra {
 
-    @Unique
-    private BakedModel headModel;
+	@Unique
+	private BakedModel headModel;
 
+	@Unique
+	@Nullable
+	private TintOverrides tintOverrides;
 
-    @Nullable
-    @Override
-    public BakedModel getHeadModel_firmament() {
-        return headModel;
-    }
+	@Nullable
+	@Override
+	public BakedModel getHeadModel_firmament() {
+		return headModel;
+	}
 
-    @Override
-    public void setHeadModel_firmament(@Nullable BakedModel headModel) {
-        this.headModel = headModel;
-    }
+	@Override
+	public void setHeadModel_firmament(@Nullable BakedModel headModel) {
+		this.headModel = headModel;
+	}
+
+	@Override
+	public @Nullable TintOverrides getTintOverrides_firmament() {
+		return tintOverrides;
+	}
+
+	@Override
+	public void setTintOverrides_firmament(@Nullable TintOverrides tintOverrides) {
+		this.tintOverrides = tintOverrides;
+	}
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBuiltin.java b/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBuiltin.java
index 76300a1..87aecb1 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBuiltin.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/BakedModelDataHolderBuiltin.java
@@ -2,6 +2,7 @@
 package moe.nea.firmament.mixins.custommodels;
 
 import moe.nea.firmament.features.texturepack.BakedModelExtra;
+import moe.nea.firmament.features.texturepack.TintOverrides;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.BuiltinBakedModel;
 import org.jetbrains.annotations.Nullable;
@@ -11,18 +12,32 @@ import org.spongepowered.asm.mixin.Unique;
 @Mixin(BuiltinBakedModel.class)
 public class BakedModelDataHolderBuiltin implements BakedModelExtra {
 
-    @Unique
-    private BakedModel headModel;
+	@Unique
+	@Nullable
+	private BakedModel headModel;
 
+	@Unique
+	@Nullable
+	private TintOverrides tintOverrides;
 
-    @Nullable
-    @Override
-    public BakedModel getHeadModel_firmament() {
-        return headModel;
-    }
+	@Override
+	public @Nullable TintOverrides getTintOverrides_firmament() {
+		return tintOverrides;
+	}
 
-    @Override
-    public void setHeadModel_firmament(@Nullable BakedModel headModel) {
-        this.headModel = headModel;
-    }
+	@Override
+	public void setTintOverrides_firmament(@Nullable TintOverrides tintOverrides) {
+		this.tintOverrides = tintOverrides;
+	}
+
+	@Nullable
+	@Override
+	public BakedModel getHeadModel_firmament() {
+		return headModel;
+	}
+
+	@Override
+	public void setHeadModel_firmament(@Nullable BakedModel headModel) {
+		this.headModel = headModel;
+	}
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemColorRemovalPatch.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemColorRemovalPatch.java
new file mode 100644
index 0000000..8c76c60
--- /dev/null
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemColorRemovalPatch.java
@@ -0,0 +1,39 @@
+package moe.nea.firmament.mixins.custommodels;
+
+import moe.nea.firmament.features.texturepack.TintOverrides;
+import moe.nea.firmament.init.ItemColorsSodiumRiser;
+import net.minecraft.client.color.item.ItemColorProvider;
+import net.minecraft.client.color.item.ItemColors;
+import net.minecraft.item.ItemStack;
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+@Mixin(ItemColors.class)
+public class ItemColorRemovalPatch {
+
+	/**
+	 * @see ItemColorsSodiumRiser
+	 */
+	private @Nullable ItemColorProvider overrideSodium_firmament(@Nullable ItemColorProvider original) {
+		var tintOverrides = TintOverrides.Companion.getCurrentOverrides();
+		if (!tintOverrides.hasOverrides()) return original;
+		return (stack, tintIndex) -> {
+			var override = tintOverrides.getOverride(tintIndex);
+			if (override != null) return override;
+			if (original != null) return original.getColor(stack, tintIndex);
+			return -1;
+		};
+	}
+
+
+	@Inject(method = "getColor", at = @At("HEAD"), cancellable = true)
+	private void overrideGetColorCall(ItemStack item, int tintIndex, CallbackInfoReturnable<Integer> cir) {
+		var tintOverrides = TintOverrides.Companion.getCurrentOverrides();
+		var override = tintOverrides.getOverride(tintIndex);
+		if (override != null)
+			cir.setReturnValue(override);
+	}
+}
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemModelGeneratorJsonUnbakedModelCopy.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemModelGeneratorJsonUnbakedModelCopy.java
index c897441..89d0411 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemModelGeneratorJsonUnbakedModelCopy.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemModelGeneratorJsonUnbakedModelCopy.java
@@ -12,8 +12,11 @@ import org.spongepowered.asm.mixin.injection.At;
 @Mixin(ItemModelGenerator.class)
 public class ItemModelGeneratorJsonUnbakedModelCopy {
     @ModifyReturnValue(method = "create", at = @At("RETURN"))
-    private JsonUnbakedModel copyHeadModel(JsonUnbakedModel original, @Local(argsOnly = true) JsonUnbakedModel oldModel) {
-        ((JsonUnbakedModelFirmExtra) original).setHeadModel_firmament(((JsonUnbakedModelFirmExtra) oldModel).getHeadModel_firmament());
+    private JsonUnbakedModel copyExtraModelData(JsonUnbakedModel original, @Local(argsOnly = true) JsonUnbakedModel oldModel) {
+        var extra = ((JsonUnbakedModelFirmExtra) original);
+		var oldExtra = ((JsonUnbakedModelFirmExtra) oldModel);
+		extra.setHeadModel_firmament(oldExtra.getHeadModel_firmament());
+		extra.setTintOverrides_firmament(oldExtra.getTintOverrides_firmament());
         return original;
     }
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java
new file mode 100644
index 0000000..5ed0fbe
--- /dev/null
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java
@@ -0,0 +1,33 @@
+package moe.nea.firmament.mixins.custommodels;
+
+import moe.nea.firmament.features.texturepack.BakedModelExtra;
+import moe.nea.firmament.features.texturepack.TintOverrides;
+import net.minecraft.client.render.VertexConsumerProvider;
+import net.minecraft.client.render.item.ItemRenderer;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.json.ModelTransformationMode;
+import net.minecraft.client.util.math.MatrixStack;
+import net.minecraft.item.ItemStack;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(ItemRenderer.class)
+public class ItemRendererTintContextPatch {
+	@Inject(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/client/render/model/json/ModelTransformationMode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;)V",
+		at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/BakedModel;getTransformation()Lnet/minecraft/client/render/model/json/ModelTransformation;"), allow = 1)
+	private void onStartRendering(ItemStack stack, ModelTransformationMode renderMode, boolean leftHanded, MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, int overlay, BakedModel model, CallbackInfo ci) {
+		if (model instanceof BakedModelExtra extra) {
+			TintOverrides.Companion.enter(extra.getTintOverrides_firmament());
+		}
+	}
+
+	@Inject(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/client/render/model/json/ModelTransformationMode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;)V",
+		at = @At(value = "INVOKE", target = "Lnet/minecraft/client/util/math/MatrixStack;pop()V"), allow = 1)
+	private void onEndRendering(ItemStack stack, ModelTransformationMode renderMode, boolean leftHanded, MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, int overlay, BakedModel model, CallbackInfo ci) {
+		if (model instanceof BakedModelExtra extra) {
+			TintOverrides.Companion.exit(extra.getTintOverrides_firmament());
+		}
+	}
+}
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java b/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
index 5fe3dec..20c69e2 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
@@ -1,4 +1,3 @@
-
 package moe.nea.firmament.mixins.custommodels;
 
 import com.google.gson.annotations.SerializedName;
@@ -6,12 +5,14 @@ import com.llamalad7.mixinextras.injector.ModifyReturnValue;
 import com.llamalad7.mixinextras.sugar.Local;
 import moe.nea.firmament.features.texturepack.BakedModelExtra;
 import moe.nea.firmament.features.texturepack.JsonUnbakedModelFirmExtra;
+import moe.nea.firmament.features.texturepack.TintOverrides;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.Baker;
 import net.minecraft.client.render.model.ModelRotation;
 import net.minecraft.client.render.model.UnbakedModel;
 import net.minecraft.client.render.model.json.JsonUnbakedModel;
 import net.minecraft.util.Identifier;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -23,46 +24,74 @@ import java.util.Objects;
 
 @Mixin(JsonUnbakedModel.class)
 public class JsonUnbakedModelDataHolder implements JsonUnbakedModelFirmExtra {
-    @Shadow
-    @Nullable
-    protected JsonUnbakedModel parent;
-    @Unique
-    @Nullable
-    public Identifier headModel;
+	@Shadow
+	@Nullable
+	protected JsonUnbakedModel parent;
+	@Unique
+	@Nullable
+	public Identifier headModel;
+	@Unique
+	@Nullable
+	public TintOverrides tintOverrides;
+	@Unique
+	@Nullable
+	public TintOverrides mergedTintOverrides;
+
+	@Override
+	public void setTintOverrides_firmament(@Nullable TintOverrides tintOverrides) {
+		this.tintOverrides = tintOverrides;
+		this.mergedTintOverrides = null;
+	}
+
+	@Override
+	public @NotNull TintOverrides getTintOverrides_firmament() {
+		if (mergedTintOverrides != null)
+			return mergedTintOverrides;
+		var mergedTintOverrides = parent == null ? new TintOverrides()
+			: ((JsonUnbakedModelFirmExtra) parent).getTintOverrides_firmament();
+		if (tintOverrides != null)
+			mergedTintOverrides = tintOverrides.mergeWithParent(mergedTintOverrides);
+		this.mergedTintOverrides = mergedTintOverrides;
+		return mergedTintOverrides;
+	}
 
-    @Override
-    public void setHeadModel_firmament(@Nullable Identifier identifier) {
-        this.headModel = identifier;
-    }
+	@Override
+	public void setHeadModel_firmament(@Nullable Identifier identifier) {
+		this.headModel = identifier;
+	}
 
-    @Override
-    public @Nullable Identifier getHeadModel_firmament() {
-        if (this.headModel != null) return this.headModel;
-        if (this.parent == null) return null;
-        return ((JsonUnbakedModelFirmExtra) this.parent).getHeadModel_firmament();
-    }
+	@Override
+	public @Nullable Identifier getHeadModel_firmament() {
+		if (this.headModel != null) return this.headModel;
+		if (this.parent == null) return null;
+		return ((JsonUnbakedModelFirmExtra) this.parent).getHeadModel_firmament();
+	}
 
-    @ModifyReturnValue(method = "getModelDependencies", at = @At("RETURN"))
-    private Collection<Identifier> addDependencies(Collection<Identifier> original) {
-        var headModel = getHeadModel_firmament();
-        if (headModel != null) {
-            original.add(headModel);
-        }
-        return original;
-    }
+	@ModifyReturnValue(method = "getModelDependencies", at = @At("RETURN"))
+	private Collection<Identifier> addDependencies(Collection<Identifier> original) {
+		var headModel = getHeadModel_firmament();
+		if (headModel != null) {
+			original.add(headModel);
+		}
+		return original;
+	}
 
-    @ModifyReturnValue(
-        method = "bake(Lnet/minecraft/client/render/model/Baker;Lnet/minecraft/client/render/model/json/JsonUnbakedModel;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Z)Lnet/minecraft/client/render/model/BakedModel;",
-        at = @At(value = "RETURN"))
-    private BakedModel bakeExtraInfo(BakedModel original, @Local(argsOnly = true) Baker baker) {
-        var headModel = getHeadModel_firmament();
-        if (headModel != null && original instanceof BakedModelExtra extra) {
-            UnbakedModel unbakedModel = baker.getOrLoadModel(headModel);
-            extra.setHeadModel_firmament(
-                Objects.equals(unbakedModel, parent)
-                    ? null
-                    : baker.bake(headModel, ModelRotation.X0_Y0));
-        }
-        return original;
-    }
+	@ModifyReturnValue(
+		method = "bake(Lnet/minecraft/client/render/model/Baker;Lnet/minecraft/client/render/model/json/JsonUnbakedModel;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Z)Lnet/minecraft/client/render/model/BakedModel;",
+		at = @At(value = "RETURN"))
+	private BakedModel bakeExtraInfo(BakedModel original, @Local(argsOnly = true) Baker baker) {
+		if (original instanceof BakedModelExtra extra) {
+			var headModel = getHeadModel_firmament();
+			if (headModel != null) {
+				UnbakedModel unbakedModel = baker.getOrLoadModel(headModel);
+				extra.setHeadModel_firmament(
+					Objects.equals(unbakedModel, parent)
+						? null
+						: baker.bake(headModel, ModelRotation.X0_Y0));
+			}
+			if (getTintOverrides_firmament().hasOverrides())
+				extra.setTintOverrides_firmament(getTintOverrides_firmament());
+		}
+		return original;
+	}
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/PatchJsonUnbakedModelDeserializer.java b/src/main/java/moe/nea/firmament/mixins/custommodels/PatchJsonUnbakedModelDeserializer.java
index cd6fb17..d6c25b5 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/PatchJsonUnbakedModelDeserializer.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/PatchJsonUnbakedModelDeserializer.java
@@ -6,6 +6,7 @@ import com.google.gson.JsonPrimitive;
 import com.llamalad7.mixinextras.injector.ModifyReturnValue;
 import com.llamalad7.mixinextras.sugar.Local;
 import moe.nea.firmament.features.texturepack.JsonUnbakedModelFirmExtra;
+import moe.nea.firmament.features.texturepack.TintOverrides;
 import net.minecraft.client.render.model.json.JsonUnbakedModel;
 import net.minecraft.util.Identifier;
 import org.spongepowered.asm.mixin.Mixin;
@@ -13,13 +14,18 @@ import org.spongepowered.asm.mixin.injection.At;
 
 @Mixin(JsonUnbakedModel.Deserializer.class)
 public class PatchJsonUnbakedModelDeserializer {
-    @ModifyReturnValue(method = "deserialize(Lcom/google/gson/JsonElement;Ljava/lang/reflect/Type;Lcom/google/gson/JsonDeserializationContext;)Lnet/minecraft/client/render/model/json/JsonUnbakedModel;",
-        at = @At("RETURN"))
-    private JsonUnbakedModel addHeadModel(JsonUnbakedModel original, @Local JsonObject jsonObject) {
-        var headModel = jsonObject.get("firmament:head_model");
-        if (headModel instanceof JsonPrimitive prim && prim.isString()) {
-            ((JsonUnbakedModelFirmExtra) original).setHeadModel_firmament(Identifier.of(prim.getAsString()));
-        }
-        return original;
-    }
+	@ModifyReturnValue(method = "deserialize(Lcom/google/gson/JsonElement;Ljava/lang/reflect/Type;Lcom/google/gson/JsonDeserializationContext;)Lnet/minecraft/client/render/model/json/JsonUnbakedModel;",
+		at = @At("RETURN"))
+	private JsonUnbakedModel addHeadModel(JsonUnbakedModel original, @Local JsonObject jsonObject) {
+		var headModel = jsonObject.get("firmament:head_model");
+		var extra = ((JsonUnbakedModelFirmExtra) original);
+		if (headModel instanceof JsonPrimitive prim && prim.isString()) {
+			extra.setHeadModel_firmament(Identifier.of(prim.getAsString()));
+		}
+		var tintOverrides = jsonObject.get("firmament:tint_overrides");
+		if (tintOverrides instanceof JsonObject object) {
+			extra.setTintOverrides_firmament(TintOverrides.Companion.parse(object));
+		}
+		return original;
+	}
 }