Fix: pass through entry partial results for map and list codecs

Mojang · Mar 19, 2024 · b2fc9d6 · b2fc9d6
1 parent 0e55329
commit b2fc9d6
Show file tree

Hide file tree

Showing 4 changed files with 152 additions and 15 deletions.
diff --git a/src/main/java/com/mojang/serialization/DataResult.java b/src/main/java/com/mojang/serialization/DataResult.java
@@ -89,6 +89,10 @@ public Optional<R> resultOrPartial(final Consumer<String> onError) {
         );
     }
 
+    public Optional<R> resultOrPartial() {
+        return result.map(Optional::of, error -> error.partialResult);
+    }
+
     public <E extends Throwable> R getOrThrow(final Function<String, E> exceptionSupplier) throws E {
         final Optional<PartialResult<R>> error = result.right();
         if (error.isPresent()) {

diff --git a/src/main/java/com/mojang/serialization/codecs/BaseMapCodec.java b/src/main/java/com/mojang/serialization/codecs/BaseMapCodec.java
@@ -2,7 +2,6 @@
 // Licensed under the MIT license.
 package com.mojang.serialization.codecs;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.mojang.datafixers.util.Pair;
 import com.mojang.datafixers.util.Unit;
@@ -14,6 +13,8 @@
 import com.mojang.serialization.RecordBuilder;
 
 import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Stream;
 
 public interface BaseMapCodec<K, V> {
     Codec<K> keyCodec();
@@ -22,27 +23,27 @@ public interface BaseMapCodec<K, V> {
 
     default <T> DataResult<Map<K, V>> decode(final DynamicOps<T> ops, final MapLike<T> input) {
         final ImmutableMap.Builder<K, V> read = ImmutableMap.builder();
-        final ImmutableList.Builder<Pair<T, T>> failed = ImmutableList.builder();
+        final Stream.Builder<Pair<T, T>> failed = Stream.builder();
 
         final DataResult<Unit> result = input.entries().reduce(
             DataResult.success(Unit.INSTANCE, Lifecycle.stable()),
             (r, pair) -> {
-                final DataResult<K> k = keyCodec().parse(ops, pair.getFirst());
-                final DataResult<V> v = elementCodec().parse(ops, pair.getSecond());
+                final DataResult<K> key = keyCodec().parse(ops, pair.getFirst());
+                final DataResult<V> value = elementCodec().parse(ops, pair.getSecond());
 
-                final DataResult<Pair<K, V>> entry = k.apply2stable(Pair::of, v);
-                entry.error().ifPresent(e -> failed.add(pair));
+                final DataResult<Pair<K, V>> entryResult = key.apply2stable(Pair::of, value);
+                entryResult.resultOrPartial().ifPresent(entry -> read.put(entry.getFirst(), entry.getSecond()));
+                if (entryResult.error().isPresent()) {
+                    failed.add(pair);
+                }
 
-                return r.apply2stable((u, p) -> {
-                    read.put(p.getFirst(), p.getSecond());
-                    return u;
-                }, entry);
+                return r.apply2stable((u, p) -> u, entryResult);
             },
             (r1, r2) -> r1.apply2stable((u1, u2) -> u1, r2)
         );
 
         final Map<K, V> elements = read.build();
-        final T errors = ops.createMap(failed.build().stream());
+        final T errors = ops.createMap(failed.build());
 
         return result.map(unit -> elements).setPartial(elements).mapError(e -> e + " missed input: " + errors);
     }

diff --git a/src/main/java/com/mojang/serialization/codecs/ListCodec.java b/src/main/java/com/mojang/serialization/codecs/ListCodec.java
@@ -44,10 +44,8 @@ public <T> DataResult<Pair<List<A>, T>> decode(final DynamicOps<T> ops, final T
             stream.accept(t -> {
                 final DataResult<Pair<A, T>> element = elementCodec.decode(ops, t);
                 element.error().ifPresent(e -> failed.add(t));
-                result.setPlain(result.getPlain().apply2stable((r, v) -> {
-                    read.add(v.getFirst());
-                    return r;
-                }, element));
+                element.resultOrPartial().ifPresent(pair -> read.add(pair.getFirst()));
+                result.setPlain(result.getPlain().apply2stable((r, v) -> r, element));
             });
 
             final ImmutableList<A> elements = read.build();

diff --git a/src/test/java/com/mojang/serialization/CodecTests.java b/src/test/java/com/mojang/serialization/CodecTests.java
@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+package com.mojang.serialization;
+
+import com.google.common.collect.ImmutableMap;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class CodecTests {
+    private static <T> Object toJava(final Codec<T> codec, final T value) {
+        return codec.encodeStart(JavaOps.INSTANCE, value).getOrThrow(AssertionError::new);
+    }
+
+    private static <T> T fromJava(final Codec<T> codec, final Object value) {
+        return codec.parse(JavaOps.INSTANCE, value).getOrThrow(AssertionError::new);
+    }
+
+    private static <T> T fromJavaOrPartial(final Codec<T> codec, final Object value) {
+        return codec.parse(JavaOps.INSTANCE, value).getPartialOrThrow(AssertionError::new);
+    }
+
+    private static void assertFromJavaFails(final Codec<?> codec, final Object value) {
+        final DataResult<?> result = codec.parse(JavaOps.INSTANCE, value);
+        assertTrue("Expected data result error, but got: " + result.result(), result.result().isEmpty());
+    }
+
+    private static <T> void assertRoundTrip(final Codec<T> codec, final T value, final Object java) {
+        assertEquals(
+            java,
+            toJava(codec, value)
+        );
+        assertEquals(
+            value,
+            fromJava(codec, java)
+        );
+    }
+
+    @Test
+    public void unboundedMap_simple() {
+        assertRoundTrip(
+            Codec.unboundedMap(Codec.STRING, Codec.INT),
+            Map.of(
+                "foo", 1,
+                "bar", 2
+            ),
+            Map.of(
+                "foo", 1,
+                "bar", 2
+            )
+        );
+    }
+
+    @Test
+    public void unboundedMap_invalidEntry() {
+        final Codec<Map<String, Integer>> codec = Codec.unboundedMap(Codec.STRING, Codec.INT);
+        assertFromJavaFails(codec, Map.of(
+            "foo", 1,
+            "bar", "garbage",
+            "baz", 3
+        ));
+    }
+
+    @Test
+    public void unboundedMap_invalidEntryPartial() {
+        final Codec<Map<String, Integer>> codec = Codec.unboundedMap(Codec.STRING, Codec.INT);
+        assertEquals(
+            Map.of(
+                "foo", 1,
+                "baz", 3
+            ),
+            fromJavaOrPartial(codec, ImmutableMap.of(
+                "foo", 1,
+                "bar", "garbage",
+                "baz", 3
+            ))
+        );
+    }
+
+    @Test
+    public void unboundedMap_invalidEntryNestedPartial() {
+        final Codec<Map<String, Map<String, Integer>>> codec = Codec.unboundedMap(Codec.STRING, Codec.unboundedMap(Codec.STRING, Codec.INT));
+        assertEquals(
+            Map.of(
+                "foo", Map.of(
+                    "foo", 1
+                ),
+                "bar", Map.of(
+                    "foo", 1,
+                    "baz", 3
+                )
+            ),
+            fromJavaOrPartial(codec, ImmutableMap.of(
+                "foo", Map.of(
+                    "foo", 1
+                ),
+                "bar", Map.of(
+                    "foo", 1,
+                    "bar", "garbage",
+                    "baz", 3
+                )
+            ))
+        );
+    }
+
+    @Test
+    public void list_roundTrip() {
+        assertRoundTrip(
+            Codec.STRING.listOf(),
+            List.of("foo", "bar", "baz"),
+            List.of("foo", "bar", "baz")
+        );
+    }
+
+    @Test
+    public void list_invalidValues() {
+        final Codec<List<String>> codec = Codec.STRING.listOf();
+        assertFromJavaFails(codec, List.of("foo", 2, "baz", false));
+
+        assertEquals(
+            List.of("foo", "bar"),
+            fromJavaOrPartial(codec, List.of("foo", "bar", 2, false))
+        );
+
+        assertEquals(
+            List.of("foo", "baz"),
+            fromJavaOrPartial(codec, List.of("foo", 2, "baz", false))
+        );
+    }
+}