diff --git a/src/main/java/com/mojang/datafixers/DataFix.java b/src/main/java/com/mojang/datafixers/DataFix.java index 36f51493..6c0aa26d 100644 --- a/src/main/java/com/mojang/datafixers/DataFix.java +++ b/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,18 +43,24 @@ protected TypeRewriteRule writeAndRead(final String name, final Type type, fi return writeFixAndRead(name, type, newType, Function.identity()); } + @SuppressWarnings("unchecked") protected TypeRewriteRule writeFixAndRead(final String name, final Type type, final Type newType, final Function, Dynamic> fix) { - return fixTypeEverywhere(name, type, newType, ops -> input -> { - final Optional> written = type.writeDynamic(ops, input).resultOrPartial(LOGGER::error); - if (!written.isPresent()) { + final AtomicReference> patchedType = new AtomicReference<>(); + final RewriteResult view = unchecked(name, type, newType, ops -> input -> { + final Optional> written = patchedType.getPlain().writeDynamic(ops, input).resultOrPartial(LOGGER::error); + if (written.isEmpty()) { throw new RuntimeException("Could not write the object in " + name); } final Optional, ?>> 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) type.all(rule, true, false).view().newType()); + return rule; } protected TypeRewriteRule fixTypeEverywhere(final String name, final Type type, final Type newType, final Function, Function> function) { @@ -61,7 +68,7 @@ protected TypeRewriteRule fixTypeEverywhere(final String name, final Type } protected TypeRewriteRule fixTypeEverywhere(final String name, final Type type, final Type newType, final Function, Function> 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 TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type type, final Function, Typed> function) { @@ -80,6 +87,10 @@ protected TypeRewriteRule fixTypeEverywhereTyped(final String name, final return fixTypeEverywhere(type, checked(name, type, newType, function, bitSet)); } + private static RewriteResult unchecked(final String name, final Type type, final Type newType, final Function, Function> function, final BitSet bitSet) { + return RewriteResult.create(View.create(name, type, newType, new NamedFunctionWrapper<>(name, function)), bitSet); + } + @SuppressWarnings("unchecked") public static RewriteResult checked(final String name, final Type type, final Type newType, final Function, Typed> function, final BitSet bitSet) { return RewriteResult.create(View.create(name, type, newType, new NamedFunctionWrapper<>(name, ops -> a -> { diff --git a/src/main/java/com/mojang/datafixers/DataFixerUpper.java b/src/main/java/com/mojang/datafixers/DataFixerUpper.java index 248e36a3..a4464616 100644 --- a/src/main/java/com/mojang/datafixers/DataFixerUpper.java +++ b/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 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 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;