modify frozen and vanilla registries too

CleanroomMC · Nov 26, 2023 · fdc2b78 · fdc2b78
1 parent a41c385
commit fdc2b78
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 16 deletions.
diff --git a/src/main/java/com/cleanroommc/groovyscript/api/IReloadableForgeRegistry.java b/src/main/java/com/cleanroommc/groovyscript/api/IReloadableForgeRegistry.java
@@ -1,6 +1,7 @@
 package com.cleanroommc.groovyscript.api;
 
 import net.minecraft.util.ResourceLocation;
+import org.jetbrains.annotations.ApiStatus;
 
 /**
  * Access interface for reloadable forge registries
@@ -12,4 +13,10 @@ public interface IReloadableForgeRegistry<V> {
     void removeEntry(ResourceLocation name);
 
     void onReload();
+
+    @ApiStatus.Internal
+    void groovyscript$putDummy(V dummy, V realEntry, ResourceLocation name, int id, Object owner);
+
+    @ApiStatus.Internal
+    void groovyscript$forceAdd(V entry, int id, Object owner);
 }
diff --git a/src/main/java/com/cleanroommc/groovyscript/core/mixin/ForgeRegistryMixin.java b/src/main/java/com/cleanroommc/groovyscript/core/mixin/ForgeRegistryMixin.java
@@ -2,13 +2,14 @@
 
 import com.cleanroommc.groovyscript.api.IReloadableForgeRegistry;
 import com.cleanroommc.groovyscript.registry.ReloadableRegistryManager;
+import com.cleanroommc.groovyscript.registry.VirtualizedForgeRegistryEntry;
 import com.google.common.collect.BiMap;
 import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
 import net.minecraft.util.ResourceLocation;
 import net.minecraftforge.registries.ForgeRegistry;
 import net.minecraftforge.registries.IForgeRegistry;
 import net.minecraftforge.registries.IForgeRegistryEntry;
-import org.apache.commons.lang3.tuple.Triple;
+import net.minecraftforge.registries.RegistryManager;
 import org.jetbrains.annotations.Nullable;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
@@ -17,6 +18,7 @@
 
 import java.util.BitSet;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 @Mixin(value = ForgeRegistry.class, remap = false)
@@ -50,8 +52,17 @@ public abstract class ForgeRegistryMixin<V extends IForgeRegistryEntry<V>> imple
     @Final
     private BitSet availabilityMap;
 
+    @Shadow
+    public abstract V remove(ResourceLocation key);
+
+    @Shadow
+    @Final
+    private Class<V> superType;
+    @Shadow
+    @Final
+    private RegistryManager stage;
     @Unique
-    private Set<Triple<V, Integer, Object>> backups;
+    private Set<VirtualizedForgeRegistryEntry<V>> backups;
     @Unique
     private Set<V> scripted;
     @Unique
@@ -60,66 +71,88 @@ public abstract class ForgeRegistryMixin<V extends IForgeRegistryEntry<V>> imple
     @Final
     private final Set<ResourceLocation> dummies = new ObjectOpenHashSet<>();
 
+    @Unique
+    private IReloadableForgeRegistry<V> vanilla;
+    @Unique
+    private IReloadableForgeRegistry<V> frozen;
+
     @Override
     public V registerEntry(V registryEntry) {
+        if (stage != RegistryManager.ACTIVE) throw new IllegalStateException("Do not modify VANILLA or FROZEN registry directly!");
         if (registryEntry != null) {
             ResourceLocation rl = registryEntry.getRegistryName();
             if (rl != null) {
                 groovyscript$removeDummy(rl);
             }
         }
-        V newEntry = getValue(add(-1, registryEntry, null));
+        int id = add(-1, registryEntry, null);
+        V newEntry = getValue(id);
         if (newEntry == registryEntry) {
             if (this.scripted == null) {
                 this.scripted = new ObjectOpenHashSet<>();
             }
             this.scripted.add(registryEntry);
+            Object owner = this.owners.inverse().get(registryEntry);
+            grs$do(reg -> reg.groovyscript$forceAdd(registryEntry, id, owner));
         }
         return newEntry;
     }
 
     @Override
     public void removeEntry(ResourceLocation name) {
+        if (stage != RegistryManager.ACTIVE) throw new IllegalStateException("Do not modify VANILLA or FROZEN registry directly!");
+        if (this.dummies.contains(name)) return;
         V entry = this.names.remove(name);
         if (entry != null) {
             if (this.backups == null) {
                 this.backups = new ObjectOpenHashSet<>();
             }
             Integer id = this.ids.inverse().remove(entry);
             Object ownerOverride = this.owners.inverse().remove(entry);
+            if (id == null) throw new IllegalStateException();
             if (this.scripted == null || !this.scripted.contains(entry)) {
-                this.backups.add(Triple.of(entry, id, ownerOverride));
+                this.backups.add(new VirtualizedForgeRegistryEntry<>(entry, id, ownerOverride));
             }
-            groovyscript$putDummy(entry, name, id, ownerOverride);
+            V dummy = groovyscript$putDummy(entry, name, id, ownerOverride);
+            grs$do(reg -> reg.groovyscript$putDummy(dummy, entry, name, id, ownerOverride));
         }
     }
 
     @Override
     public void onReload() {
+        if (stage != RegistryManager.ACTIVE) throw new IllegalStateException("Do not modify VANILLA or FROZEN registry directly!");
         unfreeze();
         if (this.scripted != null) {
             for (V entry : this.scripted) {
                 ResourceLocation rl = this.names.inverse().remove(entry);
                 Integer id = this.ids.inverse().remove(entry);
                 Object owner = this.owners.inverse().remove(entry);
-                groovyscript$putDummy(entry, rl, id, owner);
+                V dummy = groovyscript$putDummy(entry, rl, id, owner);
+                grs$do(reg -> reg.groovyscript$putDummy(dummy, entry, rl, id, owner));
             }
             this.scripted = null;
         }
         if (this.backups != null) {
-            for (Triple<V, Integer, Object> entry : this.backups) {
-                this.names.put(entry.getLeft().getRegistryName(), entry.getLeft());
-                this.ids.put(entry.getMiddle(), entry.getLeft());
-                this.owners.put(entry.getRight(), entry.getLeft());
-                this.dummies.remove(entry.getLeft().getRegistryName());
+            for (VirtualizedForgeRegistryEntry<V> entry : this.backups) {
+                this.names.put(entry.getValue().getRegistryName(), entry.getValue());
+                this.ids.put(entry.getId(), entry.getValue());
+                this.owners.put(entry.getOverride(), entry.getValue());
+                this.dummies.remove(entry.getValue().getRegistryName());
+                grs$do(reg -> reg.groovyscript$forceAdd(entry.getValue(), entry.getId(), entry.getOverride()));
             }
             this.backups = null;
         }
     }
 
-    public void groovyscript$putDummy(V entry, ResourceLocation rl, Integer id, Object owner) {
-        if (entry == null || rl == null || id == null) return;
+    public V groovyscript$putDummy(V entry, ResourceLocation rl, Integer id, Object owner) {
+        if (entry == null || rl == null || id == null) throw new IllegalArgumentException();
         V dummy = groovyscript$getDummy(rl);
+        groovyscript$putDummy(dummy, entry, rl, id, owner);
+        return dummy;
+    }
+
+    @Override
+    public void groovyscript$putDummy(V dummy, V entry, ResourceLocation rl, int id, Object owner) {
         if (dummy != null) {
             this.names.put(rl, dummy);
             this.ids.put(id, dummy);
@@ -152,5 +185,21 @@ public void onReload() {
         if (value != null) value.setRegistryName(rl);
         return value;
     }
+
+    @Override
+    public void groovyscript$forceAdd(V entry, int id, Object owner) {
+        names.put(entry.getRegistryName(), entry);
+        ids.put(id, entry);
+        if (owner != null) owners.put(owner, entry);
+        availabilityMap.set(id);
+    }
+
+    @SuppressWarnings("unchecked")
+    private void grs$do(Consumer<IReloadableForgeRegistry<V>> consumer) {
+        if (frozen == null) frozen = (IReloadableForgeRegistry<V>) RegistryManager.FROZEN.getRegistry(superType);
+        if (vanilla == null) vanilla = (IReloadableForgeRegistry<V>) RegistryManager.VANILLA.getRegistry(superType);
+        if (frozen != null) consumer.accept(frozen);
+        if (vanilla != null) consumer.accept(vanilla);
+    }
 }
 
diff --git a/src/main/java/com/cleanroommc/groovyscript/registry/ReloadableRegistryManager.java b/src/main/java/com/cleanroommc/groovyscript/registry/ReloadableRegistryManager.java
@@ -121,7 +121,7 @@ public static <V extends IForgeRegistryEntry<V>> Supplier<V> getDummySupplier(Cl
      */
     @ApiStatus.Internal
     @SideOnly(Side.CLIENT)
-    public static long reloadJei(boolean msgPlayer) {
+    public static void reloadJei(boolean msgPlayer) {
         if (ModSupport.JEI.isLoaded()) {
             JeiProxyAccessor jeiProxy = (JeiProxyAccessor) JustEnoughItems.getProxy();
             long time = System.currentTimeMillis();
@@ -130,9 +130,7 @@ public static long reloadJei(boolean msgPlayer) {
             if (msgPlayer) {
                 Minecraft.getMinecraft().player.sendMessage(new TextComponentString("Reloading JEI took " + time + "ms"));
             }
-            return time;
         }
-        return 0;
     }
 
     private static void reloadForgeRegistries() {

diff --git a/src/main/java/com/cleanroommc/groovyscript/registry/VirtualizedForgeRegistryEntry.java b/src/main/java/com/cleanroommc/groovyscript/registry/VirtualizedForgeRegistryEntry.java
@@ -0,0 +1,48 @@
+package com.cleanroommc.groovyscript.registry;
+
+import net.minecraftforge.registries.IForgeRegistryEntry;
+
+import java.util.Objects;
+
+public class VirtualizedForgeRegistryEntry<T extends IForgeRegistryEntry<T>> {
+
+    private final T value;
+    private final int id;
+    private final Object override;
+
+    public VirtualizedForgeRegistryEntry(T value, int id, Object override) {
+        this.value = value;
+        this.id = id;
+        this.override = override;
+    }
+
+    public T getValue() {
+        return value;
+    }
+
+    public int getId() {
+        return id;
+    }
+
+    public Object getOverride() {
+        return override;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        VirtualizedForgeRegistryEntry<?> that = (VirtualizedForgeRegistryEntry<?>) o;
+        return id == that.id && Objects.equals(value, that.value);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(value, id);
+    }
+
+    @Override
+    public String toString() {
+        return value.getRegistryName().toString();
+    }
+}