Some data fixer example code

NewCodeTest.java

package net.minecraft.util.datafix;

import com.google.gson.JsonElement;
import com.google.gson.internal.Streams;
import com.google.gson.stream.JsonReader;
import com.mojang.datafixers.DSL;
import com.mojang.datafixers.DataFix;
import com.mojang.datafixers.Dynamic;
import com.mojang.datafixers.OpticFinder;
import com.mojang.datafixers.RewriteResult;
import com.mojang.datafixers.TypeRewriteRule;
import com.mojang.datafixers.Typed;
import com.mojang.datafixers.functions.PointFreeRule;
import com.mojang.datafixers.optics.Inj1;
import com.mojang.datafixers.optics.Optics;
import com.mojang.datafixers.optics.Proj2;
import com.mojang.datafixers.optics.ReForgetC;
import com.mojang.datafixers.types.JsonOps;
import com.mojang.datafixers.types.Type;
import com.mojang.datafixers.types.families.RecursiveTypeFamily;
import com.mojang.datafixers.types.templates.RecursivePoint;
import com.mojang.datafixers.types.templates.TypeTemplate;
import com.mojang.datafixers.util.Either;
import com.mojang.datafixers.util.Pair;

import java.io.StringReader;
import java.util.BitSet;
import java.util.Optional;
import java.util.function.Function;

import static com.mojang.datafixers.DSL.constType;
import static com.mojang.datafixers.DSL.fieldFinder;
import static com.mojang.datafixers.DSL.fields;
import static com.mojang.datafixers.DSL.intType;
import static com.mojang.datafixers.DSL.remainder;
import static com.mojang.datafixers.DSL.string;

public class NewCodeTest {
    public static void main(final String[] args) throws Exception {
        final PointFreeRule opSimple = PointFreeRule.orElse(PointFreeRule.CataFuseSame.INSTANCE, PointFreeRule.orElse(PointFreeRule.LensAppId.INSTANCE, PointFreeRule.LensComp.INSTANCE));
        final PointFreeRule opLeft = PointFreeRule.many(PointFreeRule.once(PointFreeRule.orElse(PointFreeRule.CompAssocLeft.INSTANCE, opSimple)));
        final PointFreeRule opRight = PointFreeRule.many(PointFreeRule.once(PointFreeRule.orElse(PointFreeRule.CompAssocRight.INSTANCE, opSimple)));
        final PointFreeRule op = PointFreeRule.seq(opLeft, () -> opRight);

        testSimpleOptics();
        testList(op);
    }

    public static void testSimpleOptics() {
        final Either<Integer, Integer> ev = Either.right(5);

        final Pair<Integer, Either<Integer, Integer>> v = Pair.of(1, ev);

        final Proj2<Integer, Either<Integer, Integer>, Either<Integer, Integer>> proj = Optics.proj2();
        final Inj1<Integer, Integer, Integer> inj = Optics.inj1();

        final Object nv = ReForgetC.unbox(proj.compose(inj).eval(new ReForgetC.Instance<Integer>()).apply(Optics.reForgetC("run", Either.left(Function.identity())))).run(v, 8);
        System.out.println(nv);
    }

    public static void testList(final PointFreeRule op) throws Exception {
        final JsonElement list;
        try (final JsonReader reader = new JsonReader(new StringReader("{value:3,next:{value:5,next:{value:11}}}"))) {
            reader.setLenient(true);
            list = Streams.parse(reader);
        }

        final TypeTemplate listF = fields(
            "value", constType(intType()),
            DSL.optionalFields(
                "next", DSL.id(0)
            )
        );

        final TypeTemplate listF2 = fields(
            "value", constType(intType()),
            DSL.optionalFields(
                "next", DSL.id(0),
                "next2", DSL.id(0)
            )
        );

        final RecursiveTypeFamily listT = new RecursiveTypeFamily("list", listF);
        final RecursiveTypeFamily listT2 = new RecursiveTypeFamily("list2", listF2);

        final RecursivePoint.RecursivePointType<?> listType = listT.apply(0);
        final RecursivePoint.RecursivePointType<?> listType2 = listT2.apply(0);

        final OpticFinder<Integer> valueField = fieldFinder("value", intType());
        final OpticFinder<?> nextField = fieldFinder("next", listType2);
        final OpticFinder<?> next2Field = fieldFinder("next2", listType2);

        final Type<?> next = DSL.optionalFields(
            "next", DSL.id(0)
        ).apply(listType2.family()).apply(0);

        final Type<?> next2 = DSL.optionalFields(
            "next", DSL.id(0),
            "next2", DSL.id(0)
        ).apply(listType2.family()).apply(0);

        final Type<?> end = remainder().apply(listType2.family()).apply(0);

        final OpticFinder<?> nextF = next.finder();
        final OpticFinder<?> next2F = next2.finder();
        final OpticFinder<?> endF = end.finder();

        final Pair<JsonElement, ? extends Optional<? extends Typed<?>>> read = listType.readTyped(new Dynamic<>(JsonOps.INSTANCE, list));
//        final Optional<? extends Any<?>> read2 = read.map(p -> p.getRight().update(valueField, string(), i -> "asd_" + (i + 3)));
        final TypeRewriteRule rule2 = fixTypeEverywhereT("rule2", intType(), string(), p -> new Typed<>(string(), JsonOps.INSTANCE, "fgh_" + p.getValue().toString()));

        final BitSet one = new BitSet();
        one.set(0);

        /*final TypeRewriteRule rule3 = fixTypeEverywhereT("rule3", listType, listType2, pp -> {
            // next -> next | next2
            final Any<?> t = pp.getT(nextF);
            return pp.set(endF, t.rename(n -> Objects.equals(n, "next") ? "next2" : n));
        }, one);*/

        final TypeRewriteRule rule4 = fixTypeEverywhereT("rule4", listType2, listType2, pp -> {
            final int currentV = pp.get(valueField);
            final int nextV = pp.getOptionalTyped(nextField).map(n -> n.get(valueField)).orElse(0);
            final int next2V = pp.getOptionalTyped(next2Field).map(n -> n.get(valueField)).orElse(0);
            return pp.set(valueField, currentV + nextV + next2V);
        }, one);

        final TypeRewriteRule rule5 = fixTypeEverywhereT("rule5", listType2, listType2, pp -> {
            return pp.set(valueField, pp.get(valueField) * 2);
        }, new BitSet());

        final TypeRewriteRule rule6 = fixTypeEverywhereT("rule6", listType2, listType2, pp -> {
            return pp.set(valueField, pp.get(valueField) - 10);
        }, new BitSet());

        System.out.println(read.getSecond());
//        System.out.println(read2);
        System.out.println("rule2: " + rule2.rewrite(listType));
//        System.out.println("rule3: " + rule3.rewrite(listType));
        System.out.println("rule4: " + rule4.rewrite(listType));
        System.out.println("rule5: " + rule5.rewrite(listType));
        System.out.println("rule6: " + rule5.rewrite(listType));
        final TypeRewriteRule seq = TypeRewriteRule.seq(rule4, rule5, rule6);
//        System.out.println("rule3+4+5+6: " + seq.rewrite(listType));

        System.out.println("optimized rule2: " + listType.rewrite(rule2, op));
//        System.out.println("optimized rule3: " + listType.rewrite(rule3, op));
        System.out.println("optimized rule4: " + listType.rewrite(rule4, op));
        System.out.println("optimized rule5: " + listType.rewrite(rule5, op));
//        System.out.println("optimized rule3+4+5+6: " + listType.rewrite(seq, op));
        System.out.println(test(op, list, listType, rule2));
//        System.out.println(listType.readAndWrite(JsonOps.INSTANCE, rule3.rewrite(listType).get().view().newType(), rule3, op, list));
//        System.out.println(listType.readAndWrite(JsonOps.INSTANCE, rule4.rewrite(listType).get().view().newType(), rule4, op, list));
//        System.out.println("unop apply: " + test(PointFreeRule.nop(), list, listType, seq));
//        System.out.println("op apply:   " + test(op, list, listType, seq));
    }

    public static Optional<JsonElement> test(final PointFreeRule op, final JsonElement list, final RecursivePoint.RecursivePointType<?> listType, final TypeRewriteRule rule2) {
        return listType.readAndWrite(JsonOps.INSTANCE, rule2.rewrite(listType).get().view().newType(), rule2, op, list);
    }

    private static <A, B> TypeRewriteRule fixTypeEverywhereT(final String name, final Type<A> type, final Type<B> newType, final Function<Typed<?>, Typed<?>> function) {
        return fixTypeEverywhereT(name, type, newType, function, new BitSet());
    }

    private static <A, B> TypeRewriteRule fixTypeEverywhereT(final String name, final Type<A> type, final Type<B> newType, final Function<Typed<?>, Typed<?>> function, final BitSet bitSet) {
        return fixTypeEverywhere(type, DataFix.checked(name, type, newType, function, bitSet));
    }

    private static <A, B> TypeRewriteRule fixTypeEverywhere(final Type<A> type, final RewriteResult<A, B> view) {
        return TypeRewriteRule.everywhere(TypeRewriteRule.ifSame(type, view), PointFreeRule.nop(), true, true);
    }
}