Merge pull request #338 from Dragon-Seeker/1.21.2-NbtOptionalDeserial…

…izationFix [1.21.2] NBT Optional deserialization fix
wisp-forest · Dec 11, 2024 · 621e4b9 · 621e4b9
2 parents 96dfc05 + 28d3f2d
commit 621e4b9
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/src/main/java/io/wispforest/owo/serialization/format/nbt/IdentityHolder.java b/src/main/java/io/wispforest/owo/serialization/format/nbt/IdentityHolder.java
@@ -0,0 +1,19 @@
+package io.wispforest.owo.serialization.format.nbt;
+
+record IdentityHolder<T>(T t) {
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null || obj.getClass() != this.getClass()) return false;
+        return this.t == ((IdentityHolder<?>) obj).t;
+    }
+
+    @Override
+    public int hashCode() {
+        return System.identityHashCode(this.t);
+    }
+
+    @Override
+    public String toString() {
+        return "IdentityHolder[t=" + t + ']';
+    }
+}
diff --git a/src/main/java/io/wispforest/owo/serialization/format/nbt/NbtDeserializer.java b/src/main/java/io/wispforest/owo/serialization/format/nbt/NbtDeserializer.java
@@ -87,8 +87,15 @@ public byte[] readBytes(SerializationContext ctx) {
         return this.getAs(this.getValue(), NbtByteArray.class).getByteArray();
     }
 
+    private final Set<IdentityHolder<NbtElement>> encodedOptionals = Collections.newSetFromMap(new WeakHashMap<>());
+
     @Override
     public <V> Optional<V> readOptional(SerializationContext ctx, Endec<V> endec) {
+        var value = this.getValue();
+        if (this.encodedOptionals.contains(new IdentityHolder<>(value))) {
+            return Optional.of(endec.decode(ctx, this));
+        }
+
         var struct = this.struct();
         return struct.field("present", ctx, Endec.BOOLEAN)
                 ? Optional.of(struct.field("value", ctx, endec))
@@ -239,8 +246,10 @@ public Struct(NbtCompound compound) {
 
                 return defaultValueFactory.get();
             }
+            var element = this.compound.get(name);
+            if (defaultValueFactory != null) NbtDeserializer.this.encodedOptionals.add(new IdentityHolder<>(element));
             return NbtDeserializer.this.frame(
-                    () -> this.compound.get(name),
+                    () -> element,
                     () -> endec.decode(ctx, NbtDeserializer.this)
             );
         }

diff --git a/src/main/java/io/wispforest/owo/serialization/format/nbt/NbtSerializer.java b/src/main/java/io/wispforest/owo/serialization/format/nbt/NbtSerializer.java
@@ -101,7 +101,7 @@ public void writeBytes(SerializationContext ctx, byte[] bytes) {
         this.consume(new NbtByteArray(bytes));
     }
 
-    private final Set<NbtElement> encodedOptionals = Collections.newSetFromMap(new WeakHashMap<>());
+    private final Set<IdentityHolder<NbtElement>> encodedOptionals = Collections.newSetFromMap(new WeakHashMap<>());
 
     @Override
     public <V> void writeOptional(SerializationContext ctx, Endec<V> endec, Optional<V> optional) {
@@ -115,7 +115,7 @@ public <V> void writeOptional(SerializationContext ctx, Endec<V> endec, Optional
 
             var compound = encoded.require("optional representation");
 
-            encodedOptionals.add(compound);
+            encodedOptionals.add(new IdentityHolder<>(compound));
             frameData.setValue(compound);
         });
 
@@ -177,7 +177,7 @@ public <F> Struct field(String name, SerializationContext ctx, Endec<F> endec, F
 
                 var element = encoded.require("struct field");
 
-                if (mayOmit && NbtSerializer.this.encodedOptionals.contains(element)) {
+                if (mayOmit && NbtSerializer.this.encodedOptionals.contains(new IdentityHolder<>(element))) {
                     var nbtCompound = (NbtCompound) element;
 
                     if(!nbtCompound.getBoolean("present")) return;