Small fixes for recursive data fixers

src/main/java/com/mojang/datafixers/DataFix.java

@@ -14,6 +14,7 @@
 import java.util.BitSet;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
 public abstract class DataFix {
@@ -42,26 +43,32 @@ protected TypeRewriteRule writeAndRead(final String name, final Type<?> type, fi
         return writeFixAndRead(name, type, newType, Function.identity());
     }
 
+    @SuppressWarnings("unchecked")
     protected <A, B> TypeRewriteRule writeFixAndRead(final String name, final Type<A> type, final Type<B> newType, final Function<Dynamic<?>, Dynamic<?>> fix) {
-        return fixTypeEverywhere(name, type, newType, ops -> input -> {
-            final Optional<? extends Dynamic<?>> written = type.writeDynamic(ops, input).resultOrPartial(LOGGER::error);
-            if (!written.isPresent()) {
+        final AtomicReference<Type<A>> patchedType = new AtomicReference<>();
+        final RewriteResult<A, B> view = unchecked(name, type, newType, ops -> input -> {
+            final Optional<? extends Dynamic<?>> written = patchedType.getPlain().writeDynamic(ops, input).resultOrPartial(LOGGER::error);
+            if (written.isEmpty()) {
                 throw new RuntimeException("Could not write the object in " + name);
             }
             final Optional<? extends Pair<Typed<B>, ?>> read = newType.readTyped(fix.apply(written.get())).resultOrPartial(LOGGER::error);
-            if (!read.isPresent()) {
+            if (read.isEmpty()) {
                 throw new RuntimeException("Could not read the new object in " + name);
             }
             return read.get().getFirst().getValue();
-        });
+        }, new BitSet());
+        final TypeRewriteRule rule = fixTypeEverywhere(type, view);
+        // Replace the input type within itself recursively, as this is what is actually passed to the fixer
+        patchedType.setPlain((Type<A>) type.all(rule, true, false).view().newType());
+        return rule;
     }
 
     protected <A, B> TypeRewriteRule fixTypeEverywhere(final String name, final Type<A> type, final Type<B> newType, final Function<DynamicOps<?>, Function<A, B>> function) {
         return fixTypeEverywhere(name, type, newType, function, new BitSet());
     }
 
     protected <A, B> TypeRewriteRule fixTypeEverywhere(final String name, final Type<A> type, final Type<B> newType, final Function<DynamicOps<?>, Function<A, B>> function, final BitSet bitSet) {
-        return fixTypeEverywhere(type, RewriteResult.create(View.create(name, type, newType, new NamedFunctionWrapper<>(name, function)), bitSet));
+        return fixTypeEverywhere(type, unchecked(name, type, newType, function, bitSet));
     }
 
     protected <A> TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type<A> type, final Function<Typed<?>, Typed<?>> function) {
@@ -80,6 +87,10 @@ protected <A, B> TypeRewriteRule fixTypeEverywhereTyped(final String name, final
         return fixTypeEverywhere(type, checked(name, type, newType, function, bitSet));
     }
 
+    private static <A, B> RewriteResult<A, B> unchecked(final String name, final Type<A> type, final Type<B> newType, final Function<DynamicOps<?>, Function<A, B>> function, final BitSet bitSet) {
+        return RewriteResult.create(View.create(name, type, newType, new NamedFunctionWrapper<>(name, function)), bitSet);
+    }
+
     @SuppressWarnings("unchecked")
     public static <A, B> RewriteResult<A, B> checked(final String name, final Type<A> type, final Type<B> newType, final Function<Typed<?>, Typed<?>> function, final BitSet bitSet) {
         return RewriteResult.create(View.create(name, type, newType, new NamedFunctionWrapper<>(name, ops -> a -> {

src/main/java/com/mojang/datafixers/DataFixerUpper.java

@@ -88,7 +88,7 @@ public Schema getSchema(final int key) {
     }
 
     protected Type<?> getType(final DSL.TypeReference type, final int version) {
-        return getSchema(DataFixUtils.makeKey(version)).getType(type);
+        return getSchema(DataFixUtils.makeKey(version)).getTypeRaw(type);
     }
 
     protected static int getLowestSchemaSameVersion(final Int2ObjectSortedMap<Schema> schemas, final int versionKey) {
@@ -107,20 +107,20 @@ private int getLowestFixSameVersion(final int versionKey) {
         return fixerVersions.subSet(0, versionKey + 1).lastInt();
     }
 
-    protected TypeRewriteRule getRule(final int version, final int dataVersion) {
-        if (version >= dataVersion) {
+    protected TypeRewriteRule getRule(final int version, final int newVersion) {
+        if (version >= newVersion) {
             return TypeRewriteRule.nop();
         }
 
-        final int expandedVersion = getLowestFixSameVersion(DataFixUtils.makeKey(version));
-        final int expandedDataVersion = DataFixUtils.makeKey(dataVersion);
-
-        final long key = (long) expandedVersion << 32 | expandedDataVersion;
+        final long key = (long) version << 32 | newVersion;
         return rules.computeIfAbsent(key, k -> {
+            final int expandedVersion = getLowestFixSameVersion(DataFixUtils.makeKey(version));
+
             final List<TypeRewriteRule> rules = Lists.newArrayList();
             for (final DataFix fix : globalList) {
-                final int fixVersion = fix.getVersionKey();
-                if (fixVersion > expandedVersion && fixVersion <= expandedDataVersion) {
+                final int expandedFixVersion = fix.getVersionKey();
+                final int fixVersion = DataFixUtils.getVersion(expandedFixVersion);
+                if (expandedFixVersion > expandedVersion && fixVersion <= newVersion) {
                     final TypeRewriteRule fixRule = fix.getRule();
                     if (fixRule == TypeRewriteRule.nop()) {
                         continue;