1 files changed, 38 insertions, 0 deletions
diff --git a/src/main/java/moe/nea/firmament/mixins/DFUEntityIdFixPatch.java b/src/main/java/moe/nea/firmament/mixins/DFUEntityIdFixPatch.java
new file mode 100644
index 0000000..8f16686
--- /dev/null
+++ b/src/main/java/moe/nea/firmament/mixins/DFUEntityIdFixPatch.java
@@ -0,0 +1,38 @@
+/*
+ * SPDX-FileCopyrightText: 2023 Linnea Gräf <nea@nea.moe>
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+package moe.nea.firmament.mixins;
+
+import com.mojang.datafixers.DataFix;
+import com.mojang.datafixers.TypeRewriteRule;
+import com.mojang.datafixers.schemas.Schema;
+import com.mojang.datafixers.util.Pair;
+import net.minecraft.datafixer.TypeReferences;
+import net.minecraft.datafixer.fix.EntityIdFix;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.util.Map;
+
+@Mixin(EntityIdFix.class)
+public abstract class DFUEntityIdFixPatch extends DataFix {
+    @Shadow
+    @Final
+    private static Map<String, String> RENAMED_ENTITIES;
+
+    public DFUEntityIdFixPatch(Schema outputSchema, boolean changesType) {
+        super(outputSchema, changesType);
+    }
+
+    @Inject(method = "makeRule", at = @At("RETURN"), cancellable = true)
+    public void onMakeRule(CallbackInfoReturnable<TypeRewriteRule> cir) {
+        cir.setReturnValue(TypeRewriteRule.seq(fixTypeEverywhere("EntityIdFix", getInputSchema().findChoiceType(TypeReferences.ENTITY), getOutputSchema().findChoiceType(TypeReferences.ENTITY), dynamicOps -> pair -> ((Pair) pair).mapFirst(string -> RENAMED_ENTITIES.getOrDefault(string, (String) string))), convertUnchecked("Fix Type", getInputSchema().getType(TypeReferences.ITEM_STACK), getOutputSchema().getType(TypeReferences.ITEM_STACK))));
+    }
+}