// Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. package com.mojang.datafixers; import com.mojang.datafixers.schemas.Schema; import com.mojang.datafixers.types.Type; import com.mojang.datafixers.util.Pair; import com.mojang.serialization.Dynamic; import com.mojang.serialization.DynamicOps; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import javax.annotation.Nullable; import java.util.BitSet; import java.util.Objects; import java.util.Optional; import java.util.function.Function; public abstract class DataFix { private static final Logger LOGGER = LogManager.getLogger(); private final Schema outputSchema; private final boolean changesType; @Nullable private TypeRewriteRule rule; public DataFix(final Schema outputSchema, final boolean changesType) { this.outputSchema = outputSchema; this.changesType = changesType; } protected <A> TypeRewriteRule fixTypeEverywhere(final String name, final Type<A> type, final Function<DynamicOps<?>, Function<A, A>> function) { return fixTypeEverywhere(name, type, type, function, new BitSet()); } @SuppressWarnings("unchecked") protected <A, B> TypeRewriteRule convertUnchecked(final String name, final Type<A> type, final Type<B> newType) { return fixTypeEverywhere(name, type, newType, ops -> (Function<A, B>) Function.identity(), new BitSet()); } protected TypeRewriteRule writeAndRead(final String name, final Type<?> type, final Type<?> newType) { return writeFixAndRead(name, type, newType, Function.identity()); } 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()) { 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()) { throw new RuntimeException("Could not read the new object in " + name); } return read.get().getFirst().getValue(); }); } 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)); } protected <A> TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type<A> type, final Function<Typed<?>, Typed<?>> function) { return fixTypeEverywhereTyped(name, type, function, new BitSet()); } protected <A> TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type<A> type, final Function<Typed<?>, Typed<?>> function, final BitSet bitSet) { return fixTypeEverywhereTyped(name, type, type, function, bitSet); } protected <A, B> TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type<A> type, final Type<B> newType, final Function<Typed<?>, Typed<?>> function) { return fixTypeEverywhereTyped(name, type, newType, function, new BitSet()); } protected <A, B> TypeRewriteRule fixTypeEverywhereTyped(final String name, final Type<A> type, final Type<B> newType, final Function<Typed<?>, Typed<?>> function, final BitSet bitSet) { return fixTypeEverywhere(type, checked(name, type, newType, 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 -> { final Typed<?> result = function.apply(new Typed<>(type, ops, a)); if (!newType.equals(result.type, true, false)) { throw new IllegalStateException(String.format("Dynamic type check failed: %s not equal to %s", newType, result.type)); } return (B) result.value; })), bitSet); } protected <A, B> TypeRewriteRule fixTypeEverywhere(final Type<A> type, final RewriteResult<A, B> view) { return TypeRewriteRule.checkOnce(TypeRewriteRule.everywhere(TypeRewriteRule.ifSame(type, view), DataFixerUpper.OPTIMIZATION_RULE, true, true), this::onFail); } protected void onFail(final Type<?> type) { LOGGER.info("Not matched: " + this + " " + type); } public final int getVersionKey() { return getOutputSchema().getVersionKey(); } public TypeRewriteRule getRule() { if (rule == null) { rule = makeRule(); } return rule; } protected abstract TypeRewriteRule makeRule(); protected Schema getInputSchema() { if (changesType) { return outputSchema.getParent(); } return getOutputSchema(); } protected Schema getOutputSchema() { return outputSchema; } private static final class NamedFunctionWrapper<A, B> implements Function<DynamicOps<?>, Function<A, B>> { private final String name; private final Function<DynamicOps<?>, Function<A, B>> delegate; public NamedFunctionWrapper(final String name, final Function<DynamicOps<?>, Function<A, B>> delegate) { this.name = name; this.delegate = delegate; } @Override public Function<A, B> apply(final DynamicOps<?> ops) { return delegate.apply(ops); } @Override public boolean equals(final Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } final NamedFunctionWrapper<?, ?> that = (NamedFunctionWrapper<?, ?>) o; return Objects.equals(name, that.name); } @Override public int hashCode() { return Objects.hash(name); } } }