Fix spawn egg data fixer

2 files changed, 52 insertions, 23 deletions

diff --git a/src/main/java/moe/nea/firmament/mixins/MixinEntityIdFix.java b/src/main/java/moe/nea/firmament/mixins/MixinEntityIdFix.java
index 1d5e312..c1c9868 100644
--- a/src/main/java/moe/nea/firmament/mixins/MixinEntityIdFix.java
+++ b/src/main/java/moe/nea/firmament/mixins/MixinEntityIdFix.java
@@ -1,4 +1,32 @@
 package moe.nea.firmament.mixins;
 
-public class MixinEntityIdFix {
+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 MixinEntityIdFix extends DataFix {
+    @Shadow
+    @Final
+    private static Map<String, String> RENAMED_ENTITIES;
+
+    public MixinEntityIdFix(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))));
+    }
 }
diff --git a/src/main/resources/firmament.mixins.json b/src/main/resources/firmament.mixins.json
index aa35abf..1f63bd0 100644
--- a/src/main/resources/firmament.mixins.json
+++ b/src/main/resources/firmament.mixins.json
@@ -1,25 +1,26 @@
 {
-    "required": true,
-    "plugin": "moe.nea.firmament.init.MixinPlugin",
-    "package": "moe.nea.firmament.mixins",
-    "compatibilityLevel": "JAVA_16",
-    "client": [
-        "MixinDownloadingTerrainScreen",
-        "MixinMessageHandler",
-        "MixinMinecraft",
-        "MixinWorldRenderer",
-        "accessor.AccessorHandledScreen",
-        "devenv.DisableCommonPacketWarnings"
-    ],
-    "mixins": [
-        "MixinClientPacketHandler",
-        "MixinClientPlayerEntity",
-        "MixinHandledScreen",
-        "MixinPlayerInventory",
-        "devenv.DisableInvalidFishingHook",
-        "devenv.MixinScoreboard"
-    ],
-    "injectors": {
-        "defaultRequire": 1
+  "required": true,
+  "plugin": "moe.nea.firmament.init.MixinPlugin",
+  "package": "moe.nea.firmament.mixins",
+  "compatibilityLevel": "JAVA_16",
+  "client": [
+    "MixinDownloadingTerrainScreen",
+    "MixinMessageHandler",
+    "MixinMinecraft",
+    "MixinWorldRenderer",
+    "accessor.AccessorHandledScreen",
+    "devenv.DisableCommonPacketWarnings"
+  ],
+  "mixins": [
+    "MixinClientPacketHandler",
+    "MixinClientPlayerEntity",
+    "MixinEntityIdFix",
+    "MixinHandledScreen",
+    "MixinPlayerInventory",
+    "devenv.DisableInvalidFishingHook",
+    "devenv.MixinScoreboard"
+  ],
+  "injectors": {
+    "defaultRequire": 1
   }
 }