save registry IDs to hard drive

src/main/java/net/fabricmc/fabric/mixin/registry/MixinIdRegistry.java

@@ -25,6 +25,7 @@ import net.fabricmc.fabric.registry.RemapException;
 import net.fabricmc.fabric.registry.RemappableRegistry;
 import net.minecraft.class_3513;
 import net.minecraft.util.Identifier;
+import net.minecraft.util.registry.DefaultMappedRegistry;
 import net.minecraft.util.registry.IdRegistry;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -37,11 +38,13 @@ import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 @Mixin(IdRegistry.class)
 public abstract class MixinIdRegistry<T> implements RemappableRegistry, ListenableRegistry<T>, RegistryListener<T> {
 	@Shadow
-	protected static final Logger ID_LOGGER = LogManager.getLogger();
+	protected static Logger ID_LOGGER;
 	@Shadow
 	protected class_3513<T> idStore;
 	@Shadow
 	protected BiMap<Identifier, T> objectMap;
+	@Shadow
+	private int nextId;
 
 	private Object2IntMap<Identifier> initialIdMap;
 	private RegistryListener[] listeners;
@@ -69,11 +72,17 @@ public abstract class MixinIdRegistry<T> implements RemappableRegistry, Listenab
 	}
 
 	@Override
-	public void remap(Object2IntMap<Identifier> idMap) throws RemapException {
-		//noinspection unchecked
+	public void remap(Object2IntMap<Identifier> idMap, boolean reallocateMissingEntries) throws RemapException {
+		//noinspection unchecked, ConstantConditions
 		IdRegistry<Object> registry = (IdRegistry<Object>) (Object) this;
 
-		if (!idMap.keySet().equals(registry.keys())) {
+		Object defaultValue = null;
+		//noinspection ConstantConditions
+		if (registry instanceof DefaultMappedRegistry) {
+			defaultValue = registry.get(((DefaultMappedRegistry) registry).method_10137());
+		}
+
+		if (!reallocateMissingEntries && !idMap.keySet().equals(registry.keys())) {
 			throw new RemapException("Source and destination keys differ!");
 		}
 
@@ -85,6 +94,25 @@ public abstract class MixinIdRegistry<T> implements RemappableRegistry, Listenab
 			}
 		}
 
+		if (reallocateMissingEntries) {
+			int maxValue = 0;
+
+			Object2IntMap<Identifier> idMapOld = idMap;
+			idMap = new Object2IntOpenHashMap<>();
+			for (Identifier id : idMapOld.keySet()) {
+				int v = idMapOld.getInt(id);
+				idMap.put(id, v);
+				if (v > maxValue) maxValue = v;
+			}
+
+			for (Identifier id : registry.keys()) {
+				if (!idMap.containsKey(id)) {
+					ID_LOGGER.warn("Adding " + id + " to registry.");
+					idMap.put(id, ++maxValue);
+				}
+			}
+		}
+
 		if (listeners != null) {
 			for (RegistryListener listener : listeners) {
 				listener.beforeCleared(registry);
@@ -93,11 +121,24 @@ public abstract class MixinIdRegistry<T> implements RemappableRegistry, Listenab
 
 		// We don't really need to clear anything but idStore yet.
 		idStore.method_15229();
+		nextId = 0;
 
-		for (Identifier identifier : objectMap.keySet()) {
+		for (Identifier identifier : idMap.keySet()) {
 			int id = idMap.getInt(identifier);
 			T object = objectMap.get(identifier);
+			if (object == null) {
+				ID_LOGGER.warn(identifier + " missing from registry, but requested!");
+				continue;
+				
+				//noinspection unchecked, ConstantConditions
+				// object = (T) defaultValue;
+				// objectMap.put(identifier, object);
+			}
+
 			idStore.method_15230(object, id);
+			if (nextId <= id) {
+				nextId = id + 1;
+			}
 
 			if (listeners != null) {
 				for (RegistryListener listener : listeners) {
@@ -110,7 +151,7 @@ public abstract class MixinIdRegistry<T> implements RemappableRegistry, Listenab
 	@Override
 	public void unmap() throws RemapException {
 		if (initialIdMap != null) {
-			remap(initialIdMap);
+			remap(initialIdMap, true);
 			initialIdMap = null;
 		}
 	}

src/main/java/net/fabricmc/fabric/mixin/registry/MixinWorldSaveHandler.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) 2016, 2017, 2018 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.registry;
+
+import net.fabricmc.fabric.registry.RegistrySyncManager;
+import net.fabricmc.fabric.registry.RemapException;
+import net.minecraft.nbt.TagCompound;
+import net.minecraft.nbt.TagStorageHelper;
+import net.minecraft.world.WorldSaveHandlerOld;
+import net.minecraft.world.level.LevelProperties;
+import org.apache.logging.log4j.Logger;
+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.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+
+@Mixin(value = WorldSaveHandlerOld.class)
+public class MixinWorldSaveHandler {
+	private static final int ID_REGISTRY_BACKUPS = 3;
+
+	@Shadow
+	private static Logger LOGGER;
+	@Shadow
+	public File worldDir;
+	private TagCompound lastSavedIdMap = null;
+
+	private boolean readWorldIdMap(File file) {
+		try {
+			if (file.exists()) {
+				FileInputStream fileInputStream = new FileInputStream(file);
+				TagCompound tag = TagStorageHelper.readCompoundTagCompressed(fileInputStream);
+				fileInputStream.close();
+				if (tag != null) {
+					RegistrySyncManager.apply(tag, true);
+					return true;
+				}
+			}
+
+			return false;
+		} catch (IOException e) {
+			return false;
+		} catch (RemapException e) {
+			e.printStackTrace();
+			return false;
+		}
+	}
+
+	private File getWorldIdMapFile(int i) {
+		return new File(new File(worldDir, "data"), "fabricRegistry" + ".dat" + (i == 0 ? "" : ("." + i)));
+	}
+
+	// TODO: stop double save on client?
+	@Inject(method = "readProperties", at = @At("HEAD"))
+	public void readWorldProperties(CallbackInfoReturnable<LevelProperties> callbackInfo) {
+		// Load
+		for (int i = 0; i < ID_REGISTRY_BACKUPS; i++) {
+			LOGGER.info("Loading Fabric registry [file " + (i + 1) + "/" + (ID_REGISTRY_BACKUPS + 1) + "]");
+			if (readWorldIdMap(getWorldIdMapFile(i))) {
+				break;
+			}
+		}
+
+		TagCompound newIdMap = RegistrySyncManager.toTag(false);
+		if (!newIdMap.equals(lastSavedIdMap)) {
+			for (int i = ID_REGISTRY_BACKUPS - 1; i >= 0; i--) {
+				File file = getWorldIdMapFile(i);
+				if (file.exists()) {
+					if (i == ID_REGISTRY_BACKUPS - 1) {
+						file.delete();
+					} else {
+						File target = getWorldIdMapFile(i + 1);
+						file.renameTo(target);
+					}
+				}
+			}
+
+			try {
+				FileOutputStream fileOutputStream = new FileOutputStream(getWorldIdMapFile(0));
+				TagStorageHelper.writeCompoundTagCompressed(newIdMap, fileOutputStream);
+				fileOutputStream.close();
+			} catch (IOException e) {
+				e.printStackTrace();
+			}
+
+			lastSavedIdMap = newIdMap;
+		}
+	}
+}

src/main/java/net/fabricmc/fabric/registry/RegistrySyncManager.java

@@ -46,7 +46,6 @@ public final class RegistrySyncManager {
 
 	public static CPacketCustomPayload createPacket() {
 		PacketByteBuf buf = new PacketByteBuf(Unpooled.buffer());
-		buf.writeVarInt(1);
 		buf.writeTagCompound(toTag(true));
 
 		CPacketCustomPayload packet = new CPacketCustomPayload(ID, buf);
@@ -54,14 +53,10 @@ public final class RegistrySyncManager {
 	}
 
 	public static void receivePacket(PacketContext context, PacketByteBuf buf) {
-		int version = buf.readVarInt();
-		if (version != 1) {
-			// TODO: log error
-		}
-
 		TagCompound compound = buf.readTagCompound();
+
 		try {
-			apply(compound);
+			apply(compound, false);
 		} catch (RemapException e) {
 			// TODO: log error properly
 			e.printStackTrace();
@@ -89,10 +84,16 @@ public final class RegistrySyncManager {
 			}
 		}
 
-		return mainTag;
+		TagCompound tag = new TagCompound();
+		tag.setInt("version", 1);
+		tag.setTag("registries", mainTag);
+
+		return tag;
 	}
 
-	public static void apply(TagCompound mainTag) throws RemapException {
+	public static void apply(TagCompound tag, boolean reallocateMissingEntries) throws RemapException {
+		TagCompound mainTag = tag.getTagCompound("registries");
+
 		for (Identifier registryId : Registry.REGISTRIES.keys()) {
 			if (!mainTag.hasKey(registryId.toString())) {
 				continue;
@@ -105,7 +106,7 @@ public final class RegistrySyncManager {
 				for (String key : registryTag.getKeys()) {
 					idMap.put(new Identifier(key), registryTag.getInt(key));
 				}
-				((RemappableRegistry) registry).remap(idMap);
+				((RemappableRegistry) registry).remap(idMap, reallocateMissingEntries);
 			}
 		}
 	}

src/main/java/net/fabricmc/fabric/registry/RemappableRegistry.java

@@ -20,6 +20,6 @@ import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import net.minecraft.util.Identifier;
 
 public interface RemappableRegistry {
-	void remap(Object2IntMap<Identifier> idMap) throws RemapException;
+	void remap(Object2IntMap<Identifier> idMap, boolean reallocateMissingEntries) throws RemapException;
 	void unmap() throws RemapException;
 }

src/main/resources/net.fabricmc.fabric.mixins.common.json

@@ -10,6 +10,7 @@
     "registry.MixinIdList",
     "registry.MixinIdRegistry",
     "registry.MixinServerPlayNetworkHandler",
+    "registry.MixinWorldSaveHandler",
     "resources.MixinMinecraftServer"
   ],
   "injectors": {