properly integrate 'RegistryOps' with SerializationAttributes.REGISTRIES

wisp-forest · Apr 24, 2024 · ead2902 · ead2902
1 parent 2f12212
commit ead2902
Show file tree

Hide file tree

Showing 7 changed files with 101 additions and 14 deletions.
diff --git a/src/main/java/io/wispforest/owo/mixin/RegistryOpsAccessor.java b/src/main/java/io/wispforest/owo/mixin/RegistryOpsAccessor.java
@@ -0,0 +1,15 @@
+package io.wispforest.owo.mixin;
+
+import net.minecraft.registry.RegistryOps;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+@Mixin(RegistryOps.class)
+public interface RegistryOpsAccessor {
+    @Invoker("caching")
+    static RegistryOps.RegistryInfoGetter owo$caching(RegistryOps.RegistryInfoGetter registryInfoGetter) {throw new UnsupportedOperationException();}
+
+    @Accessor("registryInfoGetter")
+    RegistryOps.RegistryInfoGetter owo$infoGetter();
+}
diff --git a/src/main/java/io/wispforest/owo/serialization/Endec.java b/src/main/java/io/wispforest/owo/serialization/Endec.java
@@ -6,14 +6,16 @@
 import com.mojang.serialization.DataResult;
 import com.mojang.serialization.DynamicOps;
 import io.netty.buffer.ByteBuf;
+import io.wispforest.owo.mixin.RegistryOpsAccessor;
 import io.wispforest.owo.serialization.endec.*;
 import io.wispforest.owo.serialization.format.bytebuf.ByteBufDeserializer;
 import io.wispforest.owo.serialization.format.bytebuf.ByteBufSerializer;
 import io.wispforest.owo.serialization.format.edm.*;
-import net.fabricmc.fabric.api.networking.v1.PacketByteBufs;
 import net.minecraft.network.PacketByteBuf;
 import net.minecraft.network.RegistryByteBuf;
 import net.minecraft.network.codec.PacketCodec;
+import net.minecraft.registry.RegistryOps;
+import org.apache.commons.lang3.ArrayUtils;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
@@ -219,15 +221,15 @@ static <T> Endec<T> ofCodec(Codec<T> codec) {
                 (serializer, value) -> {
                     DynamicOps<EdmElement<?>> ops = EdmOps.INSTANCE;
                     if (serializer.hasAttribute(SerializationAttributes.REGISTRIES)) {
-                        ops = serializer.getAttributeValue(SerializationAttributes.REGISTRIES).getOps(ops);
+                        ops = RegistryOps.of(ops, serializer.getAttributeValue(SerializationAttributes.REGISTRIES).infoGetter());
                     }
 
                     EdmEndec.INSTANCE.encode(serializer, codec.encodeStart(ops, value).getOrThrow(IllegalStateException::new));
                 },
                 deserializer -> {
                     DynamicOps<EdmElement<?>> ops = EdmOps.INSTANCE;
                     if (deserializer.hasAttribute(SerializationAttributes.REGISTRIES)) {
-                        ops = deserializer.getAttributeValue(SerializationAttributes.REGISTRIES).getOps(ops);
+                        ops = RegistryOps.of(ops, deserializer.getAttributeValue(SerializationAttributes.REGISTRIES).infoGetter());
                     }
 
                     return codec.parse(ops, EdmEndec.INSTANCE.decode(deserializer)).getOrThrow(IllegalStateException::new);
@@ -448,7 +450,11 @@ default Codec<T> codec(SerializationAttribute.Instance... assumedAttributes) {
             @Override
             public <D> DataResult<Pair<T, D>> decode(DynamicOps<D> ops, D input) {
                 try {
-                    return DataResult.success(new Pair<>(Endec.this.decode(LenientEdmDeserializer.of(ops.convertTo(EdmOps.INSTANCE, input)).withAttributes(assumedAttributes)), input));
+                    var attributes = ops instanceof RegistryOps<D> registryOps
+                            ? ArrayUtils.add(assumedAttributes, RegistriesAttribute.infoGetterOnly(((RegistryOpsAccessor) registryOps).owo$infoGetter()))
+                            : assumedAttributes;
+
+                    return DataResult.success(new Pair<>(Endec.this.decode(LenientEdmDeserializer.of(ops.convertTo(EdmOps.INSTANCE, input)).withAttributes(attributes)), input));
                 } catch (Exception e) {
                     return DataResult.error(e::getMessage);
                 }
@@ -457,7 +463,11 @@ public <D> DataResult<Pair<T, D>> decode(DynamicOps<D> ops, D input) {
             @Override
             public <D> DataResult<D> encode(T input, DynamicOps<D> ops, D prefix) {
                 try {
-                    return DataResult.success(EdmOps.INSTANCE.convertTo(ops, Endec.this.encodeFully(() -> EdmSerializer.of().withAttributes(assumedAttributes), input)));
+                    var attributes = ops instanceof RegistryOps<D> registryOps
+                            ? ArrayUtils.add(assumedAttributes, RegistriesAttribute.infoGetterOnly(((RegistryOpsAccessor) registryOps).owo$infoGetter()))
+                            : assumedAttributes;
+
+                    return DataResult.success(EdmOps.INSTANCE.convertTo(ops, Endec.this.encodeFully(() -> EdmSerializer.of().withAttributes(attributes), input)));
                 } catch (Exception e) {
                     return DataResult.error(e::getMessage);
                 }
@@ -472,7 +482,7 @@ default <B extends PacketByteBuf> PacketCodec<B, T> packetCodec() {
             public T decode(B buf) {
                 Deserializer<ByteBuf> deserializer = ByteBufDeserializer.of(buf);
                 if (buf instanceof RegistryByteBuf registryByteBuf) {
-                    deserializer = deserializer.withAttributes(SerializationAttributes.REGISTRIES.instance(registryByteBuf.getRegistryManager()));
+                    deserializer = deserializer.withAttributes(RegistriesAttribute.of(registryByteBuf.getRegistryManager()));
                 }
 
                 return Endec.this.decode(deserializer);
@@ -482,7 +492,7 @@ public T decode(B buf) {
             public void encode(B buf, T value) {
                 Serializer<ByteBuf> serializer = ByteBufSerializer.of(buf);
                 if (buf instanceof RegistryByteBuf registryByteBuf) {
-                    serializer = serializer.withAttributes(SerializationAttributes.REGISTRIES.instance(registryByteBuf.getRegistryManager()));
+                    serializer = serializer.withAttributes(RegistriesAttribute.of(registryByteBuf.getRegistryManager()));
                 }
 
                 Endec.this.encode(serializer, value);

diff --git a/src/main/java/io/wispforest/owo/serialization/RegistriesAttribute.java b/src/main/java/io/wispforest/owo/serialization/RegistriesAttribute.java
@@ -0,0 +1,62 @@
+package io.wispforest.owo.serialization;
+
+import net.minecraft.registry.*;
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Optional;
+
+public final class RegistriesAttribute implements SerializationAttribute.Instance {
+
+    private final RegistryOps.RegistryInfoGetter infoGetter;
+    private final @Nullable DynamicRegistryManager registryManager;
+
+    private RegistriesAttribute(RegistryOps.RegistryInfoGetter infoGetter, @Nullable DynamicRegistryManager registryManager) {
+        this.infoGetter = infoGetter;
+        this.registryManager = registryManager;
+    }
+
+    public static RegistriesAttribute of(DynamicRegistryManager registryManager) {
+        return new RegistriesAttribute(
+                new RegistryOps.RegistryInfoGetter() {
+                    @Override
+                    public <T> Optional<RegistryOps.RegistryInfo<T>> getRegistryInfo(RegistryKey<? extends Registry<? extends T>> registryRef) {
+                        return registryManager.getOptionalWrapper(registryRef).map(RegistryOps.RegistryInfo::fromWrapper);
+                    }
+                },
+                registryManager
+        );
+    }
+
+    @ApiStatus.Internal
+    public static RegistriesAttribute infoGetterOnly(RegistryOps.RegistryInfoGetter lookup) {
+        return new RegistriesAttribute(lookup, null);
+    }
+
+    public RegistryOps.RegistryInfoGetter infoGetter() {
+        return this.infoGetter;
+    }
+
+    public boolean hasRegistryManager() {
+        return this.registryManager != null;
+    }
+
+    public @NotNull DynamicRegistryManager registryManager() {
+        if (!this.hasRegistryManager()) {
+            throw new IllegalStateException("This instance of RegistriesAttribute does not supply a DynamicRegistryManager");
+        }
+
+        return this.registryManager;
+    }
+
+    @Override
+    public SerializationAttribute attribute() {
+        return SerializationAttributes.REGISTRIES;
+    }
+
+    @Override
+    public Object value() {
+        return this;
+    }
+}
diff --git a/src/main/java/io/wispforest/owo/serialization/SerializationAttribute.java b/src/main/java/io/wispforest/owo/serialization/SerializationAttribute.java
@@ -3,7 +3,9 @@
 public sealed abstract class SerializationAttribute permits SerializationAttribute.Marker, SerializationAttribute.WithValue {
 
     public final String name;
-    protected SerializationAttribute(String name) {this.name = name;}
+    protected SerializationAttribute(String name) {
+        this.name = name;
+    }
 
     public static SerializationAttribute.Marker marker(String name) {
         return new Marker(name);

diff --git a/src/main/java/io/wispforest/owo/serialization/SerializationAttributes.java b/src/main/java/io/wispforest/owo/serialization/SerializationAttributes.java
@@ -1,7 +1,5 @@
 package io.wispforest.owo.serialization;
 
-import net.minecraft.registry.DynamicRegistryManager;
-
 public final class SerializationAttributes {
 
     /**
@@ -22,8 +20,7 @@ public final class SerializationAttributes {
      */
     public static final SerializationAttribute.Marker HUMAN_READABLE = SerializationAttribute.marker("human_readable");
 
-    public static final SerializationAttribute.WithValue<DynamicRegistryManager> REGISTRIES = SerializationAttribute.withValue("registries");
+    public static final SerializationAttribute.WithValue<RegistriesAttribute> REGISTRIES = SerializationAttribute.withValue("registries");
 
     private SerializationAttributes() {}
-
 }
diff --git a/src/main/java/io/wispforest/owo/serialization/endec/ReflectiveEndecBuilder.java b/src/main/java/io/wispforest/owo/serialization/endec/ReflectiveEndecBuilder.java
@@ -247,7 +247,7 @@ private static Endec<?> createSealedSerializer(Class<?> commonClass) {
                 (serializer, particleEffect) -> {
                     var buffer = new RegistryByteBuf(
                             PacketByteBufs.create(),
-                            serializer.requireAttributeValue(SerializationAttributes.REGISTRIES)
+                            serializer.requireAttributeValue(SerializationAttributes.REGISTRIES).registryManager()
                     );
 
                     buffer.writeInt(Registries.PARTICLE_TYPE.getRawId(particleEffect.getType()));
@@ -257,7 +257,7 @@ private static Endec<?> createSealedSerializer(Class<?> commonClass) {
                 }, deserializer -> {
                     var buffer = new RegistryByteBuf(
                             BuiltInEndecs.PACKET_BYTE_BUF.decode(deserializer),
-                            deserializer.requireAttributeValue(SerializationAttributes.REGISTRIES)
+                            deserializer.requireAttributeValue(SerializationAttributes.REGISTRIES).registryManager()
                     );
 
                     return Registries.PARTICLE_TYPE.get(buffer.readInt()).getPacketCodec().decode(buffer);

diff --git a/src/main/resources/owo.mixins.json b/src/main/resources/owo.mixins.json
@@ -9,6 +9,7 @@
     "Copenhagen",
     "NbtCompoundMixin",
     "PacketByteBufMixin",
+    "RegistryOpsAccessor",
     "ScreenHandlerInvoker",
     "ScreenHandlerMixin",
     "ServerCommonNetworkHandlerAccessor",