Fix Remapping and Fixing of Blocks

src/main/java/com/nomiceu/nomilabs/mixin/ForgeRegistryMixin.java

@@ -0,0 +1,110 @@
+package com.nomiceu.nomilabs.mixin;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.event.RegistryEvent;
+import net.minecraftforge.registries.ForgeRegistry;
+import net.minecraftforge.registries.IForgeRegistryEntry;
+import net.minecraftforge.registries.IForgeRegistryInternal;
+import net.minecraftforge.registries.IForgeRegistryModifiable;
+
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
+
+import com.google.common.collect.Maps;
+import com.llamalad7.mixinextras.sugar.Local;
+import com.nomiceu.nomilabs.NomiLabs;
+import com.nomiceu.nomilabs.mixinhelper.RemappableForgeRegistry;
+import com.nomiceu.nomilabs.mixinhelper.RemappableSnapshot;
+
+/**
+ * This mixin saves a new map in Forge Registries: Remapped.<br>
+ * This is a map of remapped block ids to their new resource location, as specified by a remapper.<br>
+ * Note that this is only applied when the new resource location cannot be registered at the old id, as in that case, no
+ * saving of ids is required, as id loading will already load the new resource location.<br>
+ * <p>
+ * This mixin also fixes blocked lists not syncing between Forge Registries.
+ */
+@Mixin(value = ForgeRegistry.class, remap = false)
+public abstract class ForgeRegistryMixin<V extends IForgeRegistryEntry<V>>
+                                        implements IForgeRegistryInternal<V>, IForgeRegistryModifiable<V>,
+                                        RemappableForgeRegistry {
+
+    @Shadow
+    abstract void block(int id);
+
+    @Shadow
+    @Final
+    private Set<Integer> blocked;
+
+    @Unique
+    private final Map<Integer, ResourceLocation> remapped = Maps.newHashMap();
+
+    @Inject(method = "processMissingEvent",
+            at = @At(value = "INVOKE",
+                     target = "Lnet/minecraftforge/registries/ForgeRegistry;addAlias(Lnet/minecraft/util/ResourceLocation;Lnet/minecraft/util/ResourceLocation;)V"),
+            require = 1,
+            locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+    public void handleRemaps(ResourceLocation name, ForgeRegistry<V> pool,
+                             List<RegistryEvent.MissingMappings.Mapping<V>> mappings,
+                             Map<ResourceLocation, Integer> missing, Map<ResourceLocation, Integer[]> remaps,
+                             Collection<ResourceLocation> defaulted, Collection<ResourceLocation> failed,
+                             boolean injectNetworkDummies, CallbackInfo ci,
+                             @Local RegistryEvent.MissingMappings.Mapping<V> remap,
+                             @Local(ordinal = 2) int realId) {
+        if (remap.id != realId) {
+            block(remap.id);
+            remapped.put(remap.id, remap.getTarget().getRegistryName());
+            NomiLabs.LOGGER.warn(
+                    "[Forge Registry] Remap could not assign Id {} for Object {}! If this is of type BLOCK, without Data Fixers, after initial load, blocks will no longer be remapped!",
+                    remap.id, remap.getTarget().getRegistryName());
+        }
+    }
+
+    @Inject(method = "sync", at = @At("RETURN"))
+    void syncBlockedRemapped(ResourceLocation name, ForgeRegistry<V> from, CallbackInfo ci) {
+        blocked.clear();
+
+        var remFrom = (RemappableForgeRegistry) from;
+
+        remFrom.getBlocked().forEach(this::block);
+
+        remapped.clear();
+        remFrom.getRemapped().forEach(this::addRemapped);
+    }
+
+    @Inject(method = "makeSnapshot", at = @At("RETURN"))
+    public void addRemappedToSnapshot(CallbackInfoReturnable<ForgeRegistry.Snapshot> cir) {
+        ForgeRegistry.Snapshot ret = cir.getReturnValue();
+        ((RemappableSnapshot) ret).addAllRemapped(remapped);
+    }
+
+    @Override
+    @Unique
+    public void addRemapped(int id, ResourceLocation key) {
+        remapped.put(id, key);
+    }
+
+    @Override
+    @Unique
+    public Map<Integer, ResourceLocation> getRemapped() {
+        return remapped;
+    }
+
+    @Override
+    @Unique
+    public Set<Integer> getBlocked() {
+        return blocked;
+    }
+}

src/main/java/com/nomiceu/nomilabs/mixin/GameDataMixin.java

@@ -0,0 +1,38 @@
+package com.nomiceu.nomilabs.mixin;
+
+import java.util.Map;
+
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.registries.ForgeRegistry;
+import net.minecraftforge.registries.GameData;
+import net.minecraftforge.registries.RegistryManager;
+
+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;
+import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
+
+import com.llamalad7.mixinextras.sugar.Local;
+import com.nomiceu.nomilabs.mixinhelper.RemappableForgeRegistry;
+import com.nomiceu.nomilabs.mixinhelper.RemappableSnapshot;
+
+/**
+ * This Mixin allows for the remapped and block lists to be loaded from the in-world-save.
+ */
+@Mixin(value = GameData.class, remap = false)
+public class GameDataMixin {
+
+    @Inject(method = "loadPersistentDataToStagingRegistry",
+            at = @At(value = "INVOKE",
+                     target = "Lnet/minecraftforge/registries/ForgeRegistry;loadIds(Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Lnet/minecraftforge/registries/ForgeRegistry;Lnet/minecraft/util/ResourceLocation;)V"),
+            require = 1,
+            locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+    private static void loadRemappedToRegistry(RegistryManager pool, RegistryManager to,
+                                               Map<ResourceLocation, Integer[]> remaps,
+                                               Map<ResourceLocation, Integer> missing, ResourceLocation name,
+                                               ForgeRegistry.Snapshot snap, Class<?> regType, CallbackInfo ci,
+                                               @Local(ordinal = 1) ForgeRegistry<?> newRegistry) {
+        ((RemappableSnapshot) snap).loadToRegistry((RemappableForgeRegistry) newRegistry);
+    }
+}

src/main/java/com/nomiceu/nomilabs/mixin/SnapshotMixin.java

@@ -0,0 +1,73 @@
+package com.nomiceu.nomilabs.mixin;
+
+import java.util.Map;
+
+import net.minecraft.nbt.NBTTagCompound;
+import net.minecraft.nbt.NBTTagList;
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.common.util.Constants;
+import net.minecraftforge.registries.ForgeRegistry;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import com.google.common.collect.Maps;
+import com.nomiceu.nomilabs.mixinhelper.RemappableForgeRegistry;
+import com.nomiceu.nomilabs.mixinhelper.RemappableSnapshot;
+
+/**
+ * This mixin allows for saving of Forge Registry Remappings to snapshots.
+ */
+@Mixin(value = ForgeRegistry.Snapshot.class, remap = false)
+public class SnapshotMixin implements RemappableSnapshot {
+
+    @Unique
+    public Map<Integer, ResourceLocation> remapped = Maps.newHashMap();
+
+    @Unique
+    private static final String REMAPPED_KEY = "remapped";
+
+    @Inject(method = "write", at = @At("RETURN"))
+    public void saveRemapped(CallbackInfoReturnable<NBTTagCompound> cir) {
+        NBTTagCompound data = cir.getReturnValue();
+        NBTTagList remapList = new NBTTagList();
+        remapped.entrySet().stream().sorted(Map.Entry.comparingByKey()).forEachOrdered(e -> {
+            NBTTagCompound tag = new NBTTagCompound();
+            tag.setInteger("K", e.getKey());
+            tag.setString("V", e.getValue().toString());
+            remapList.appendTag(tag);
+        });
+        data.setTag(REMAPPED_KEY, remapList);
+    }
+
+    @Inject(method = "read", at = @At("RETURN"))
+    private static void readRemapped(NBTTagCompound nbt, CallbackInfoReturnable<ForgeRegistry.Snapshot> cir) {
+        ForgeRegistry.Snapshot ret = cir.getReturnValue();
+        NBTTagList list = nbt.getTagList(REMAPPED_KEY, Constants.NBT.TAG_COMPOUND);
+        list.forEach(e -> {
+            NBTTagCompound comp = (NBTTagCompound) e;
+            ((RemappableSnapshot) ret).addRemapped(comp.getInteger("K"), new ResourceLocation(comp.getString("V")));
+        });
+    }
+
+    @Override
+    @Unique
+    public void addRemapped(int id, ResourceLocation key) {
+        remapped.put(id, key);
+    }
+
+    @Override
+    @Unique
+    public void addAllRemapped(Map<Integer, ResourceLocation> map) {
+        remapped.putAll(map);
+    }
+
+    @Override
+    @Unique
+    public void loadToRegistry(RemappableForgeRegistry reg) {
+        remapped.forEach(reg::addRemapped);
+    }
+}

src/main/java/com/nomiceu/nomilabs/mixinhelper/RemappableForgeRegistry.java

@@ -0,0 +1,15 @@
+package com.nomiceu.nomilabs.mixinhelper;
+
+import java.util.Map;
+import java.util.Set;
+
+import net.minecraft.util.ResourceLocation;
+
+public interface RemappableForgeRegistry {
+
+    void addRemapped(int id, ResourceLocation key);
+
+    Set<Integer> getBlocked();
+
+    Map<Integer, ResourceLocation> getRemapped();
+}

src/main/java/com/nomiceu/nomilabs/mixinhelper/RemappableSnapshot.java

@@ -0,0 +1,14 @@
+package com.nomiceu.nomilabs.mixinhelper;
+
+import java.util.Map;
+
+import net.minecraft.util.ResourceLocation;
+
+public interface RemappableSnapshot {
+
+    void addRemapped(int id, ResourceLocation key);
+
+    void addAllRemapped(Map<Integer, ResourceLocation> map);
+
+    void loadToRegistry(RemappableForgeRegistry reg);
+}

src/main/java/com/nomiceu/nomilabs/remap/LabsRemapHelper.java

@@ -100,6 +100,11 @@ private static NBTTagCompound rewriteBlocksInSection(NBTTagCompound chunkSection
                 new NibbleArray(chunkSectionTag.getByteArray("Add")) : null;
         for (int i = 0; i < 4096; ++i) {
             int x = i & 0x0F, y = i >> 8 & 0x0F, z = i >> 4 & 0x0F;
+
+            // This is based off BlockStateContainer's setDataFromNBT
+            // There, the block id is shifted by 4, and extended is shifted by 12
+            // However, that is to allow the accommodation of 4 bits of metadata info
+            // Thus, here, the block id is not shifted, and extended is only shifted by 8.
             int id = extendedIds == null ? (blockIds[i] & 0xFF) :
                     ((blockIds[i] & 0xFF) | (extendedIds.get(x, y, z) << 8));
             var state = new BlockStateLike(id, (short) blockMetadata.get(x, y, z),

src/main/java/com/nomiceu/nomilabs/remap/datafixer/DataFixerHandler.java

@@ -26,13 +26,13 @@
 import net.minecraftforge.fml.common.FMLCommonHandler;
 import net.minecraftforge.fml.common.registry.ForgeRegistries;
 import net.minecraftforge.registries.ForgeRegistry;
-import net.minecraftforge.registries.GameData;
 
 import org.apache.commons.lang3.StringUtils;
 import org.apache.groovy.util.Arrays;
 
 import com.nomiceu.nomilabs.LabsValues;
 import com.nomiceu.nomilabs.NomiLabs;
+import com.nomiceu.nomilabs.mixinhelper.RemappableForgeRegistry;
 import com.nomiceu.nomilabs.remap.LabsRemapHelper;
 import com.nomiceu.nomilabs.remap.datafixer.fixes.BlockFixer;
 import com.nomiceu.nomilabs.remap.datafixer.fixes.ItemFixer;
@@ -67,8 +67,11 @@ public class DataFixerHandler {
 
     private static String savedLabsVersion;
 
-    /* Must be split up so that idToBlockMap is the old one (so we can use not registered resource locations) */
-    private static NBTTagList oldBlockRegistry;
+    /*
+     * Must be split up so that idToBlockMap has remapped info
+     * (Remapped Info from ForgeRegistry & Snapshot & Game Data Mixins)
+     * Remapped contains the old id mapped to the new resource location (as specified by LabsRemappers)
+     */
     private static Map<Integer, ResourceLocation> idToBlockMap;
     private static Map<ResourceLocation, Integer> blockToIdMap;
 
@@ -97,11 +100,9 @@ public static void onWorldLoad(SaveHandler save) {
         neededNewFixes = null;
         savedLabsVersion = null;
         fixAvailable = true;
-
-        // Clear Block Helper Maps, the ids can be different for each save
         idToBlockMap = null;
         blockToIdMap = null;
-        oldBlockRegistry = null;
+
         NomiLabs.LOGGER.info("Checking Data Fixers...");
 
         getInfoFromSave(save);
@@ -183,18 +184,9 @@ private static void getInfoFromSave(SaveHandler save) {
                         continue;
 
                     savedLabsVersion = compound.getString("ModVersion");
-                    break;
+                    return;
                 }
             }
-
-            if (!fml.hasKey("Registries", Constants.NBT.TAG_COMPOUND)) return;
-            NBTTagCompound registries = fml.getCompoundTag("Registries");
-
-            if (!registries.hasKey(GameData.BLOCKS.toString(), Constants.NBT.TAG_COMPOUND)) return;
-            NBTTagCompound blocks = registries.getCompoundTag(GameData.BLOCKS.toString());
-
-            if (!blocks.hasKey("ids", Constants.NBT.TAG_LIST)) return;
-            oldBlockRegistry = blocks.getTagList("ids", Constants.NBT.TAG_COMPOUND);
         } catch (IOException e) {
             NomiLabs.LOGGER.fatal("Failed to read level.dat.", e);
         }
@@ -229,24 +221,22 @@ private static void determineNeededFixesAndLog() {
     }
 
     /**
-     * The Id To Block Map is either loaded from the save's existing registry map, or is produced once needed,
-     * instead of in World Load or Post Init.<br>
+     * The Id to Block Map also includes Remapped IDs of Blocks.<br>
      * This means they are loaded after block id mismatch fixes, so the ids are correct to the world.
      */
     public static Map<Integer, ResourceLocation> getIdToBlockMap() {
         if (idToBlockMap != null) return idToBlockMap;
 
-        if (oldBlockRegistry == null || oldBlockRegistry.isEmpty()) {
-            ForgeRegistry<Block> registry = (ForgeRegistry<Block>) ForgeRegistries.BLOCKS;
-            idToBlockMap = registry.getKeys().stream()
-                    .collect(Collectors.toMap(registry::getID, Function.identity()));
-            NomiLabs.LOGGER.error(
-                    "Block Registry Save in level.dat was not found. Defaulting to Current Registry, some Remaps may not work correctly!");
-        } else {
-            idToBlockMap = oldBlockRegistry.tagList.stream()
-                    .map((tag) -> (NBTTagCompound) tag)
-                    .collect(Collectors.toMap((tag) -> tag.getInteger("V"),
-                            (tag) -> new ResourceLocation(tag.getString("K"))));
+        ForgeRegistry<Block> registry = (ForgeRegistry<Block>) ForgeRegistries.BLOCKS;
+        idToBlockMap = registry.getKeys().stream()
+                .collect(Collectors.toMap(registry::getID, Function.identity()));
+        var remReg = (RemappableForgeRegistry) registry;
+        if (!remReg.getRemapped().isEmpty()) {
+            NomiLabs.LOGGER.debug("Map Before Adding Remapped IDs:");
+            NomiLabs.LOGGER.debug(idToBlockMap);
+            NomiLabs.LOGGER.debug("Adding Block Remapped IDs:");
+            NomiLabs.LOGGER.debug(remReg.getRemapped());
+            idToBlockMap.putAll(remReg.getRemapped());
         }
 
         NomiLabs.LOGGER.debug("Generated Id to Block Map!");
@@ -311,6 +301,5 @@ public static void close() {
         neededNewFixes = null;
         idToBlockMap = null;
         blockToIdMap = null;
-        oldBlockRegistry = null;
     }
 }