[1.21] Reimplement the ability to wrap ItemGroups and adjust RegistryAccess to fix #252

src/main/java/io/wispforest/owo/itemgroup/json/WrapperGroup.java

@@ -1,17 +1,16 @@
 package io.wispforest.owo.itemgroup.json;
 
-import com.mojang.serialization.Lifecycle;
 import io.wispforest.owo.itemgroup.Icon;
 import io.wispforest.owo.itemgroup.OwoItemGroup;
 import io.wispforest.owo.itemgroup.gui.ItemGroupButton;
 import io.wispforest.owo.itemgroup.gui.ItemGroupTab;
 import io.wispforest.owo.mixin.itemgroup.ItemGroupAccessor;
-import io.wispforest.owo.mixin.ui.SimpleRegistryAccessor;
+import io.wispforest.owo.util.pond.OwoSimpleRegistryExtensions;
 import net.minecraft.item.ItemGroup;
 import net.minecraft.registry.Registries;
 import net.minecraft.registry.RegistryKey;
 import net.minecraft.registry.RegistryKeys;
-import net.minecraft.registry.SimpleRegistry;
+import net.minecraft.registry.entry.RegistryEntryInfo;
 import net.minecraft.util.Identifier;
 import org.jetbrains.annotations.ApiStatus;
 
@@ -34,9 +33,7 @@ public WrapperGroup(ItemGroup parent, Identifier parentId, List<ItemGroupTab> ta
 
         int parentRawId = Registries.ITEM_GROUP.getRawId(parent);
 
-        // TODO: set doesn't exist anymore. figure out what to do.
-//        ((SimpleRegistryAccessor<ItemGroup>) Registries.ITEM_GROUP).owo$getValueToEntry().remove(parent);
-//        ((SimpleRegistry<ItemGroup>) Registries.ITEM_GROUP).set(parentRawId, RegistryKey.of(RegistryKeys.ITEM_GROUP, parentId), this, Lifecycle.stable());
+        ((OwoSimpleRegistryExtensions<ItemGroup>) Registries.ITEM_GROUP).owo$set(parentRawId, RegistryKey.of(RegistryKeys.ITEM_GROUP, parentId), this, RegistryEntryInfo.DEFAULT);
 
         ((ItemGroupAccessor) this).owo$setDisplayName(parent.getDisplayName());
         ((ItemGroupAccessor) this).owo$setColumn(parent.getColumn());

src/main/java/io/wispforest/owo/mixin/registry/ReferenceAccessor.java

@@ -0,0 +1,15 @@
+package io.wispforest.owo.mixin.registry;
+
+import net.minecraft.registry.RegistryKey;
+import net.minecraft.registry.entry.RegistryEntry;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+@Mixin(RegistryEntry.Reference.class)
+public interface ReferenceAccessor<T> {
+    @Invoker("setRegistryKey")
+    void owo$setRegistryKey(RegistryKey<T> registryKey);
+
+    @Invoker("setValue")
+    void owo$setValue(T value);
+}

src/main/java/io/wispforest/owo/mixin/registry/SimpleRegistryMixin.java

@@ -0,0 +1,76 @@
+package io.wispforest.owo.mixin.registry;
+
+import com.mojang.serialization.Lifecycle;
+import io.wispforest.owo.util.OwoFreezer;
+import io.wispforest.owo.util.pond.OwoSimpleRegistryExtensions;
+import it.unimi.dsi.fastutil.objects.ObjectList;
+import it.unimi.dsi.fastutil.objects.Reference2IntMap;
+import net.fabricmc.fabric.api.event.registry.RegistryEntryAddedCallback;
+import net.minecraft.registry.MutableRegistry;
+import net.minecraft.registry.RegistryKey;
+import net.minecraft.registry.SimpleRegistry;
+import net.minecraft.registry.entry.RegistryEntry;
+import net.minecraft.registry.entry.RegistryEntryInfo;
+import net.minecraft.util.Identifier;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+@Mixin(SimpleRegistry.class)
+public abstract class SimpleRegistryMixin<T> implements MutableRegistry<T>, OwoSimpleRegistryExtensions<T> {
+
+    @Shadow private Map<T, RegistryEntry.Reference<T>> intrusiveValueToEntry;
+    @Shadow @Final private Map<RegistryKey<T>, RegistryEntry.Reference<T>> keyToEntry;
+    @Shadow @Final private Map<Identifier, RegistryEntry.Reference<T>> idToEntry;
+    @Shadow @Final private Map<T, RegistryEntry.Reference<T>> valueToEntry;
+    @Shadow @Final private ObjectList<RegistryEntry.Reference<T>> rawIdToEntry;
+    @Shadow @Final private Reference2IntMap<T> entryToRawId;
+    @Shadow @Final private Map<RegistryKey<T>, RegistryEntryInfo> keyToEntryInfo;
+    @Shadow private Lifecycle lifecycle;
+
+    //--
+
+    /**
+     * Copy of the {@link SimpleRegistry#add} function but uses {@link List#set} instead of {@link List#add} for {@link SimpleRegistry#rawIdToEntry}
+     */
+    public RegistryEntry.Reference<T> owo$set(int id, RegistryKey<T> arg, T object, RegistryEntryInfo arg2) {
+        this.valueToEntry.remove(object);
+
+        OwoFreezer.checkRegister("Registry Set Calls"); //this.assertNotFrozen(arg);
+
+        Objects.requireNonNull(arg);
+        Objects.requireNonNull(object);
+
+        RegistryEntry.Reference<T> reference;
+
+        if (this.intrusiveValueToEntry != null) {
+            reference = this.intrusiveValueToEntry.remove(object);
+
+            if (reference == null) {
+                throw new AssertionError("Missing intrusive holder for " + arg + ":" + object);
+            }
+
+            ((ReferenceAccessor<T>) reference).owo$setRegistryKey(arg);
+        } else {
+            reference = this.keyToEntry.computeIfAbsent(arg, k -> RegistryEntry.Reference.standAlone(this.getEntryOwner(), k));
+            ((ReferenceAccessor<T>) reference).owo$setValue((T)object);
+        }
+
+        this.keyToEntry.put(arg, reference);
+        this.idToEntry.put(arg.getValue(), reference);
+        this.valueToEntry.put(object, reference);
+        this.rawIdToEntry.set(id, reference);
+        this.entryToRawId.put(object, id);
+        this.keyToEntryInfo.put(arg, arg2);
+        this.lifecycle = this.lifecycle.add(arg2.lifecycle());
+
+        // TODO: SHOULD WE BE REFIREING THE EVENT?
+        RegistryEntryAddedCallback.event(this).invoker().onEntryAdded(id, arg.getValue(), (T)object);
+
+        return reference;
+    }
+}

src/main/java/io/wispforest/owo/util/RegistryAccess.java

@@ -1,56 +1,28 @@
 package io.wispforest.owo.util;
 
 import net.minecraft.registry.Registry;
-import net.minecraft.registry.RegistryKey;
 import net.minecraft.util.Identifier;
 import net.minecraft.registry.entry.RegistryEntry;
-import net.minecraft.registry.SimpleRegistry;
 import org.jetbrains.annotations.Nullable;
 
+@Deprecated
 public final class RegistryAccess {
 
     private RegistryAccess() {}
 
     /**
-     * Gets a {@link RegistryEntry} from its id
-     *
-     * @param registry The registry to operate on. Must be a {@link SimpleRegistry} at some point in the hierarchy
-     * @param id       The id to use
-     * @param <T>      The type of the registry and returned entry
-     * @return The entry, or {@code null} if it's not present
+     * @deprecated Use {@link Registry#getEntry(Identifier)}
      */
     @Nullable
-    @SuppressWarnings("unchecked")
     public static <T> RegistryEntry<T> getEntry(Registry<T> registry, Identifier id) {
-        checkSimple(registry);
-        return ((AccessibleRegistry<T>) registry).getEntry(id);
+        return registry.getEntry(id).orElse(null);
     }
 
     /**
-     * Gets a {@link RegistryEntry} from its value
-     *
-     * @param registry The registry to operate on. Must be a {@link SimpleRegistry} at some point in the hierarchy
-     * @param value    The value to use
-     * @param <T>      The type of the registry and returned entry
-     * @return The entry, or {@code null} if it's not present
+     * @deprecated Use {@link Registry#getEntry(T)}
      */
     @Nullable
-    @SuppressWarnings("unchecked")
     public static <T> RegistryEntry<T> getEntry(Registry<T> registry, T value) {
-        checkSimple(registry);
-        return ((AccessibleRegistry<T>) registry).getEntry(value);
+        return registry.getEntry(value);
     }
-
-    private static void checkSimple(Registry<?> registry) {
-        if (registry instanceof SimpleRegistry<?>) return;
-        throw new IllegalArgumentException("[RegistryAccess] Tried to operate on Registry of class '"
-                + registry.getClass() + "', but only 'SimpleRegistry' and descendants are supported");
-    }
-
-    public interface AccessibleRegistry<T> {
-        @Nullable RegistryEntry<T> getEntry(Identifier id);
-
-        @Nullable RegistryEntry<T> getEntry(T value);
-    }
-
 }

src/main/java/io/wispforest/owo/util/pond/OwoSimpleRegistryExtensions.java

@@ -0,0 +1,12 @@
+package io.wispforest.owo.util.pond;
+
+import net.minecraft.registry.RegistryKey;
+import net.minecraft.registry.entry.RegistryEntry;
+import net.minecraft.registry.entry.RegistryEntryInfo;
+import org.jetbrains.annotations.ApiStatus;
+
+public interface OwoSimpleRegistryExtensions<T> {
+
+    @ApiStatus.Internal
+    RegistryEntry.Reference<T> owo$set(int id, RegistryKey<T> arg, T object, RegistryEntryInfo arg2);
+}

src/main/resources/owo.mixins.json

@@ -20,7 +20,6 @@
     "ServerPlayerEntityMixin",
     "ServerPlayerInteractionManagerMixin",
     "SetComponentsLootFunctionAccessor",
-    "SimpleRegistryMixin",
     "TagGroupLoaderMixin",
     "itemgroup.ItemGroupAccessor",
     "itemgroup.ItemMixin",
@@ -30,13 +29,14 @@
     "offline.WorldSaveHandlerMixin",
     "recipe_remainders.CraftingResultSlotMixin",
     "recipe_remainders.RecipeManagerMixin",
+    "registry.SimpleRegistryMixin",
+    "registry.ReferenceAccessor",
     "text.LanguageMixin",
     "text.TextCodecsMixin",
     "text.TranslatableTextContentMixin",
     "text.stapi.SystemDelegatedLanguageFixin",
     "tweaks.EulaReaderMixin",
     "tweaks.LevelInfoMixin",
-    "ui.SimpleRegistryAccessor",
     "ui.SlotAccessor",
     "ui.SlotMixin",
     "ui.access.BlockEntityAccessor"