Code/net/minecraft/util/datafix/fixes/SavedDataFeaturePoolElementFix.java

package net.minecraft.util.datafix.fixes;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Sets;
import com.mojang.datafixers.DataFix;
import com.mojang.datafixers.TypeRewriteRule;
import com.mojang.datafixers.schemas.Schema;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.Dynamic;
import com.mojang.serialization.OptionalDynamic;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Stream;

public class SavedDataFeaturePoolElementFix extends DataFix {
    private static final Pattern INDEX_PATTERN = Pattern.compile("\\[(\\d+)\\]");
    private static final Set<String> PIECE_TYPE = Sets.newHashSet(
        "minecraft:jigsaw", "minecraft:nvi", "minecraft:pcp", "minecraft:bastionremnant", "minecraft:runtime"
    );
    private static final Set<String> FEATURES = Sets.newHashSet("minecraft:tree", "minecraft:flower", "minecraft:block_pile", "minecraft:random_patch");

    public SavedDataFeaturePoolElementFix(Schema p_145646_) {
        super(p_145646_, false);
    }

    @Override
    public TypeRewriteRule makeRule() {
        return this.writeFixAndRead(
            "SavedDataFeaturePoolElementFix",
            this.getInputSchema().getType(References.STRUCTURE_FEATURE),
            this.getOutputSchema().getType(References.STRUCTURE_FEATURE),
            SavedDataFeaturePoolElementFix::fixTag
        );
    }

    private static <T> Dynamic<T> fixTag(Dynamic<T> p_145663_) {
        return p_145663_.update("Children", SavedDataFeaturePoolElementFix::updateChildren);
    }

    private static <T> Dynamic<T> updateChildren(Dynamic<T> p_145665_) {
        return p_145665_.asStreamOpt().map(SavedDataFeaturePoolElementFix::updateChildren).map(p_145665_::createList).result().orElse(p_145665_);
    }

    private static Stream<? extends Dynamic<?>> updateChildren(Stream<? extends Dynamic<?>> p_145661_) {
        return p_145661_.map(
            p_145667_ -> {
                String s = p_145667_.get("id").asString("");
                if (!PIECE_TYPE.contains(s)) {
                    return p_145667_;
                } else {
                    OptionalDynamic<?> optionaldynamic = p_145667_.get("pool_element");
                    return !optionaldynamic.get("element_type").asString("").equals("minecraft:feature_pool_element")
                        ? p_145667_
                        : p_145667_.update("pool_element", p_145669_ -> p_145669_.update("feature", SavedDataFeaturePoolElementFix::fixFeature));
                }
            }
        );
    }

    private static <T> OptionalDynamic<T> get(Dynamic<T> p_145650_, String... p_145651_) {
        if (p_145651_.length == 0) {
            throw new IllegalArgumentException("Missing path");
        } else {
            OptionalDynamic<T> optionaldynamic = p_145650_.get(p_145651_[0]);

            for (int i = 1; i < p_145651_.length; i++) {
                String s = p_145651_[i];
                Matcher matcher = INDEX_PATTERN.matcher(s);
                if (matcher.matches()) {
                    int j = Integer.parseInt(matcher.group(1));
                    List<? extends Dynamic<T>> list = optionaldynamic.asList(Function.identity());
                    if (j >= 0 && j < list.size()) {
                        optionaldynamic = new OptionalDynamic<>(p_145650_.getOps(), DataResult.success((Dynamic<T>)list.get(j)));
                    } else {
                        optionaldynamic = new OptionalDynamic<>(p_145650_.getOps(), DataResult.error(() -> "Missing id:" + j));
                    }
                } else {
                    optionaldynamic = optionaldynamic.get(s);
                }
            }

            return optionaldynamic;
        }
    }

    @VisibleForTesting
    protected static Dynamic<?> fixFeature(Dynamic<?> p_145648_) {
        Optional<String> optional = getReplacement(
            get(p_145648_, "type").asString(""),
            get(p_145648_, "name").asString(""),
            get(p_145648_, "config", "state_provider", "type").asString(""),
            get(p_145648_, "config", "state_provider", "state", "Name").asString(""),
            get(p_145648_, "config", "state_provider", "entries", "[0]", "data", "Name").asString(""),
            get(p_145648_, "config", "foliage_placer", "type").asString(""),
            get(p_145648_, "config", "leaves_provider", "state", "Name").asString("")
        );
        return optional.isPresent() ? p_145648_.createString(optional.get()) : p_145648_;
    }

    private static Optional<String> getReplacement(
        String p_145653_, String p_145654_, String p_145655_, String p_145656_, String p_145657_, String p_145658_, String p_145659_
    ) {
        String s;
        if (!p_145653_.isEmpty()) {
            s = p_145653_;
        } else {
            if (p_145654_.isEmpty()) {
                return Optional.empty();
            }

            if ("minecraft:normal_tree".equals(p_145654_)) {
                s = "minecraft:tree";
            } else {
                s = p_145654_;
            }
        }

        if (FEATURES.contains(s)) {
            if ("minecraft:random_patch".equals(s)) {
                if ("minecraft:simple_state_provider".equals(p_145655_)) {
                    if ("minecraft:sweet_berry_bush".equals(p_145656_)) {
                        return Optional.of("minecraft:patch_berry_bush");
                    }

                    if ("minecraft:cactus".equals(p_145656_)) {
                        return Optional.of("minecraft:patch_cactus");
                    }
                } else if ("minecraft:weighted_state_provider".equals(p_145655_) && ("minecraft:grass".equals(p_145657_) || "minecraft:fern".equals(p_145657_))
                    )
                 {
                    return Optional.of("minecraft:patch_taiga_grass");
                }
            } else if ("minecraft:block_pile".equals(s)) {
                if (!"minecraft:simple_state_provider".equals(p_145655_) && !"minecraft:rotated_block_provider".equals(p_145655_)) {
                    if ("minecraft:weighted_state_provider".equals(p_145655_)) {
                        if ("minecraft:packed_ice".equals(p_145657_) || "minecraft:blue_ice".equals(p_145657_)) {
                            return Optional.of("minecraft:pile_ice");
                        }

                        if ("minecraft:jack_o_lantern".equals(p_145657_) || "minecraft:pumpkin".equals(p_145657_)) {
                            return Optional.of("minecraft:pile_pumpkin");
                        }
                    }
                } else {
                    if ("minecraft:hay_block".equals(p_145656_)) {
                        return Optional.of("minecraft:pile_hay");
                    }

                    if ("minecraft:melon".equals(p_145656_)) {
                        return Optional.of("minecraft:pile_melon");
                    }

                    if ("minecraft:snow".equals(p_145656_)) {
                        return Optional.of("minecraft:pile_snow");
                    }
                }
            } else {
                if ("minecraft:flower".equals(s)) {
                    return Optional.of("minecraft:flower_plain");
                }

                if ("minecraft:tree".equals(s)) {
                    if ("minecraft:acacia_foliage_placer".equals(p_145658_)) {
                        return Optional.of("minecraft:acacia");
                    }

                    if ("minecraft:blob_foliage_placer".equals(p_145658_) && "minecraft:oak_leaves".equals(p_145659_)) {
                        return Optional.of("minecraft:oak");
                    }

                    if ("minecraft:pine_foliage_placer".equals(p_145658_)) {
                        return Optional.of("minecraft:pine");
                    }

                    if ("minecraft:spruce_foliage_placer".equals(p_145658_)) {
                        return Optional.of("minecraft:spruce");
                    }
                }
            }
        }

        return Optional.empty();
    }
}