minecraft-source/src/net/minecraft/util/datafix/fixes/TrappedChestBlockEntityFix.java

package net.minecraft.util.datafix.fixes;

import javax.annotation.Nullable;
import org.apache.logging.log4j.LogManager;
import com.mojang.datafixers.util.Pair;
import com.mojang.datafixers.types.templates.TaggedChoice;
import java.util.Iterator;
import it.unimi.dsi.fastutil.ints.IntSet;
import java.util.Optional;
import com.mojang.datafixers.OpticFinder;
import java.util.Objects;
import com.mojang.datafixers.Dynamic;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import com.mojang.datafixers.Typed;
import com.mojang.datafixers.types.Type;
import com.mojang.datafixers.DSL;
import com.mojang.datafixers.types.templates.List;
import com.mojang.datafixers.TypeRewriteRule;
import com.mojang.datafixers.schemas.Schema;
import org.apache.logging.log4j.Logger;
import com.mojang.datafixers.DataFix;

public class TrappedChestBlockEntityFix extends DataFix {
    private static final Logger LOGGER;
    
    public TrappedChestBlockEntityFix(final Schema schema, final boolean boolean2) {
        super(schema, boolean2);
    }
    
    public TypeRewriteRule makeRule() {
        final Type<?> type2 = this.getOutputSchema().getType(References.CHUNK);
        final Type<?> type3 = type2.findFieldType("Level");
        final Type<?> type4 = type3.findFieldType("TileEntities");
        if (!(type4 instanceof List.ListType)) {
            throw new IllegalStateException("Tile entity type is not a list type.");
        }
        final List.ListType<?> listType5 = type4;
        final OpticFinder<? extends java.util.List<?>> opticFinder6 = DSL.fieldFinder("TileEntities", (Type)listType5);
        final Type<?> type5 = this.getInputSchema().getType(References.CHUNK);
        final OpticFinder<?> opticFinder7 = type5.findField("Level");
        final OpticFinder<?> opticFinder8 = opticFinder7.type().findField("Sections");
        final Type<?> type6 = opticFinder8.type();
        if (!(type6 instanceof List.ListType)) {
            throw new IllegalStateException("Expecting sections to be a list.");
        }
        final Type<?> type7 = ((List.ListType)type6).getElement();
        final OpticFinder<?> opticFinder9 = DSL.typeFinder((Type)type7);
        final OpticFinder opticFinder10;
        final Optional<? extends Typed<?>> optional6;
        final OpticFinder opticFinder11;
        java.util.List<? extends Typed<?>> list7;
        IntSet intSet8;
        final Iterator<? extends Typed<?>> iterator;
        Typed<?> typed2;
        TrappedChestSection a11;
        int integer12;
        int integer13;
        Dynamic<?> dynamic9;
        int integer14;
        int integer15;
        TaggedChoice.TaggedChoiceType<String> taggedChoiceType12;
        final OpticFinder opticFinder12;
        final TaggedChoice.TaggedChoiceType taggedChoiceType13;
        Dynamic<?> dynamic10;
        final int n;
        int integer16;
        int integer17;
        final int n2;
        int integer18;
        final IntSet set;
        final TaggedChoice.TaggedChoiceType taggedChoiceType14;
        return TypeRewriteRule.seq(new AddNewChoices(this.getOutputSchema(), "AddTrappedChestFix", References.BLOCK_ENTITY).makeRule(), this.fixTypeEverywhereTyped("Trapped Chest fix", (Type)type5, typed -> typed.updateTyped((OpticFinder)opticFinder7, typed -> {
            optional6 = typed.getOptionalTyped(opticFinder10);
            if (!optional6.isPresent()) {
                return typed;
            }
            else {
                list7 = ((Typed)optional6.get()).getAllTyped(opticFinder11);
                intSet8 = (IntSet)new IntOpenHashSet();
                list7.iterator();
                while (iterator.hasNext()) {
                    typed2 = iterator.next();
                    a11 = new TrappedChestSection(typed2, this.getInputSchema());
                    if (a11.isSkippable()) {
                        continue;
                    }
                    else {
                        for (integer12 = 0; integer12 < 4096; ++integer12) {
                            integer13 = a11.getBlock(integer12);
                            if (a11.isTrappedChest(integer13)) {
                                intSet8.add(a11.getIndex() << 12 | integer12);
                            }
                        }
                    }
                }
                dynamic9 = typed.get(DSL.remainderFinder());
                integer14 = dynamic9.get("xPos").asInt(0);
                integer15 = dynamic9.get("zPos").asInt(0);
                taggedChoiceType12 = (TaggedChoice.TaggedChoiceType<String>)this.getInputSchema().findChoiceType(References.BLOCK_ENTITY);
                return typed.updateTyped(opticFinder12, typed -> typed.updateTyped(taggedChoiceType13.finder(), typed -> {
                    dynamic10 = ((Typed)typed).getOrCreate(DSL.remainderFinder());
                    integer16 = dynamic10.get("x").asInt(0) - (n << 4);
                    integer17 = dynamic10.get("y").asInt(0);
                    integer18 = dynamic10.get("z").asInt(0) - (n2 << 4);
                    if (set.contains(LeavesFix.getIndex(integer16, integer17, integer18))) {
                        return ((Typed)typed).update(taggedChoiceType14.finder(), pair -> pair.mapFirst(string -> {
                            if (!Objects.equals(string, "minecraft:chest")) {
                                TrappedChestBlockEntityFix.LOGGER.warn("Block Entity was expected to be a chest");
                            }
                            return "minecraft:trapped_chest";
                        }));
                    }
                    else {
                        return (Typed)typed;
                    }
                }));
            }
        })));
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
    
    public static final class TrappedChestSection extends LeavesFix.Section {
        @Nullable
        private IntSet chestIds;
        
        public TrappedChestSection(final Typed<?> typed, final Schema schema) {
            super(typed, schema);
        }
        
        @Override
        protected boolean skippable() {
            this.chestIds = (IntSet)new IntOpenHashSet();
            for (int integer2 = 0; integer2 < this.palette.size(); ++integer2) {
                final Dynamic<?> dynamic3 = this.palette.get(integer2);
                final String string4 = dynamic3.get("Name").asString("");
                if (Objects.equals(string4, "minecraft:trapped_chest")) {
                    this.chestIds.add(integer2);
                }
            }
            return this.chestIds.isEmpty();
        }
        
        public boolean isTrappedChest(final int integer) {
            return this.chestIds.contains(integer);
        }
    }
}