Implement mapping of range dispatch models (#7)

src/main/java/org/geysermc/rainbow/accessor/BlockModelWrapperLocationAccessor.java

@@ -1,11 +0,0 @@
-package org.geysermc.rainbow.accessor;
-
-import net.minecraft.resources.ResourceLocation;
-
-// Implemented on BlockModelWrapper, since this class doesn't store its model after baking, we have to store it manually
-public interface BlockModelWrapperLocationAccessor {
-
-    ResourceLocation rainbow$getModelOrigin();
-
-    void rainbow$setModelOrigin(ResourceLocation model);
-}

src/main/java/org/geysermc/rainbow/accessor/ResolvedModelAccessor.java

@@ -1,12 +1,16 @@
 package org.geysermc.rainbow.accessor;
 
+import net.minecraft.client.renderer.item.ClientItem;
 import net.minecraft.client.resources.model.ResolvedModel;
 import net.minecraft.resources.ResourceLocation;
 
 import java.util.Optional;
 
-// Implemented on ModelManager, since this class doesn't keep the resolved models after baking, we have to store it manually
+// Implemented on ModelManager, since this class doesn't keep the resolved models or unbaked client items after baking, we have to store them manually.
+// This comes with some extra memory usage, but Rainbow should only be used to convert packs, so it should be fine
 public interface ResolvedModelAccessor {
 
     Optional<ResolvedModel> rainbow$getResolvedModel(ResourceLocation location);
+
+    Optional<ClientItem> rainbow$getClientItem(ResourceLocation location);
 }

src/main/java/org/geysermc/rainbow/accessor/SelectItemModelCasesAccessor.java

@@ -1,12 +0,0 @@
-package org.geysermc.rainbow.accessor;
-
-import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
-import net.minecraft.client.renderer.item.ItemModel;
-
-// Implemented on BlockModelWrapper, since this class doesn't store the cases it has in a nice format after baking, we have to store it manually
-public interface SelectItemModelCasesAccessor<T> {
-
-    Object2ObjectMap<T, ItemModel> rainbow$getCases();
-
-    void rainbow$setCases(Object2ObjectMap<T, ItemModel> cases);
-}

src/main/java/org/geysermc/rainbow/mapping/BedrockItemMapper.java

@@ -1,10 +1,11 @@
 package org.geysermc.rainbow.mapping;
 
-import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.renderer.item.BlockModelWrapper;
+import net.minecraft.client.renderer.item.ClientItem;
 import net.minecraft.client.renderer.item.ConditionalItemModel;
 import net.minecraft.client.renderer.item.ItemModel;
+import net.minecraft.client.renderer.item.ItemModels;
 import net.minecraft.client.renderer.item.RangeSelectItemModel;
 import net.minecraft.client.renderer.item.SelectItemModel;
 import net.minecraft.client.renderer.item.properties.conditional.Broken;
@@ -13,11 +14,13 @@ import net.minecraft.client.renderer.item.properties.conditional.Damaged;
 import net.minecraft.client.renderer.item.properties.conditional.FishingRodCast;
 import net.minecraft.client.renderer.item.properties.conditional.HasComponent;
 import net.minecraft.client.renderer.item.properties.conditional.ItemModelPropertyTest;
+import net.minecraft.client.renderer.item.properties.numeric.BundleFullness;
+import net.minecraft.client.renderer.item.properties.numeric.Count;
+import net.minecraft.client.renderer.item.properties.numeric.Damage;
 import net.minecraft.client.renderer.item.properties.numeric.RangeSelectItemModelProperty;
 import net.minecraft.client.renderer.item.properties.select.Charge;
 import net.minecraft.client.renderer.item.properties.select.ContextDimension;
 import net.minecraft.client.renderer.item.properties.select.DisplayContext;
-import net.minecraft.client.renderer.item.properties.select.SelectItemModelProperty;
 import net.minecraft.client.renderer.item.properties.select.TrimMaterialProperty;
 import net.minecraft.client.resources.model.Material;
 import net.minecraft.client.resources.model.ResolvedModel;
@@ -34,9 +37,8 @@ import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.component.ItemAttributeModifiers;
 import net.minecraft.world.item.equipment.trim.TrimMaterial;
 import net.minecraft.world.level.Level;
-import org.geysermc.rainbow.accessor.BlockModelWrapperLocationAccessor;
+import org.apache.commons.lang3.ArrayUtils;
 import org.geysermc.rainbow.accessor.ResolvedModelAccessor;
-import org.geysermc.rainbow.accessor.SelectItemModelCasesAccessor;
 import org.geysermc.rainbow.mapping.animation.AnimationMapper;
 import org.geysermc.rainbow.mapping.animation.BedrockAnimationContext;
 import org.geysermc.rainbow.mapping.attachable.AttachableMapper;
@@ -50,9 +52,9 @@ import org.geysermc.rainbow.mapping.geyser.GeyserSingleDefinition;
 import org.geysermc.rainbow.mapping.geyser.predicate.GeyserConditionPredicate;
 import org.geysermc.rainbow.mapping.geyser.predicate.GeyserMatchPredicate;
 import org.geysermc.rainbow.mapping.geyser.predicate.GeyserPredicate;
-import org.geysermc.rainbow.mixin.ConditionalItemModelAccessor;
+import org.geysermc.rainbow.mapping.geyser.predicate.GeyserRangeDispatchPredicate;
+import org.geysermc.rainbow.mixin.LateBoundIdMapperAccessor;
 import org.geysermc.rainbow.mixin.RangeSelectItemModelAccessor;
-import org.geysermc.rainbow.mixin.SelectItemModelAccessor;
 import org.geysermc.rainbow.mixin.TextureSlotsAccessor;
 import org.geysermc.rainbow.pack.BedrockItem;
 import org.geysermc.rainbow.pack.BedrockTextures;
@@ -70,91 +72,104 @@ public class BedrockItemMapper {
             .map(ResourceLocation::withDefaultNamespace)
             .toList();
 
+    private static ResolvedModelAccessor getModels() {
+        return (ResolvedModelAccessor) Minecraft.getInstance().getModelManager();
+    }
+
+    private static ResourceLocation getModelId(ItemModel.Unbaked model) {
+        //noinspection unchecked
+        return ((LateBoundIdMapperAccessor<ResourceLocation, ?>) ItemModels.ID_MAPPER).getIdToValue().inverse().get(model.type());
+    }
+
     public static void tryMapStack(ItemStack stack, ResourceLocation modelLocation, ProblemReporter reporter, PackContext context) {
-        ItemModel model = Minecraft.getInstance().getModelManager().getItemModel(modelLocation);
+        getModels().rainbow$getClientItem(modelLocation).map(ClientItem::model)
-        mapItem(model, stack, reporter.forChild(() -> "client item definition " + modelLocation + " "), base -> new GeyserSingleDefinition(base, Optional.of(modelLocation)), context);
+                .ifPresentOrElse(model -> mapItem(model, stack, reporter.forChild(() -> "client item definition " + modelLocation + " "), base -> new GeyserSingleDefinition(base, Optional.of(modelLocation)), context),
+                        () -> reporter.report(() -> "missing client item definition " + modelLocation));
     }
 
     public static void tryMapStack(ItemStack stack, int customModelData, ProblemReporter reporter, PackContext context) {
-        ItemModel vanillaModel = Minecraft.getInstance().getModelManager().getItemModel(stack.get(DataComponents.ITEM_MODEL));
+        ItemModel.Unbaked vanillaModel = getModels().rainbow$getClientItem(stack.get(DataComponents.ITEM_MODEL)).map(ClientItem::model).orElseThrow();
-        reporter = reporter.forChild(() -> "item model " + vanillaModel + " with custom model data " + customModelData + " ");
+        ProblemReporter childReporter = reporter.forChild(() -> "item model " + vanillaModel + " with custom model data " + customModelData + " ");
-        if (vanillaModel instanceof RangeSelectItemModel rangeModel) {
+        if (vanillaModel instanceof RangeSelectItemModel.Unbaked(RangeSelectItemModelProperty property, float scale, List<RangeSelectItemModel.Entry> entries, Optional<ItemModel.Unbaked> fallback)) {
-            RangeSelectItemModelAccessor accessor = (RangeSelectItemModelAccessor) rangeModel;
-            RangeSelectItemModelProperty property = accessor.getProperty();
             // WHY, Mojang?
             if (property instanceof net.minecraft.client.renderer.item.properties.numeric.CustomModelDataProperty(int index)) {
                 if (index == 0) {
-                    float scaledCustomModelData = customModelData * accessor.getScale();
+                    float scaledCustomModelData = customModelData * scale;
 
-                    int modelIndex = RangeSelectItemModelAccessor.invokeLastIndexLessOrEqual(accessor.getThresholds(), scaledCustomModelData);
+                    float[] thresholds = ArrayUtils.toPrimitive(entries.stream()
-                    ItemModel model = modelIndex == -1 ? accessor.getFallback() : accessor.getModels()[modelIndex];
+                            .map(RangeSelectItemModel.Entry::threshold)
-                    mapItem(model, stack, reporter, base -> new GeyserLegacyDefinition(base, customModelData), context);
+                            .toArray(Float[]::new));
+                    int modelIndex = RangeSelectItemModelAccessor.invokeLastIndexLessOrEqual(thresholds, scaledCustomModelData);
+                    Optional<ItemModel.Unbaked> model = modelIndex == -1 ? fallback : Optional.of(entries.get(modelIndex).model());
+                    model.ifPresentOrElse(present -> mapItem(present, stack, childReporter, base -> new GeyserLegacyDefinition(base, customModelData), context),
+                            () -> childReporter.report(() -> "custom model data index lookup returned -1, and no fallback is present"));
                 } else {
-                    reporter.report(() -> "range_dispatch custom model data property index is not zero, unable to apply custom model data");
+                    childReporter.report(() -> "range_dispatch custom model data property index is not zero, unable to apply custom model data");
                 }
-            } else {
+                return;
-                reporter.report(() -> "range_dispatch model property is not custom model data, unable to apply custom model data");
             }
-        } else {
-            reporter.report(() -> "item model is not range_dispatch, unable to apply custom model data");
         }
+        childReporter.report(() -> "item model is not range_dispatch, unable to apply custom model data");
     }
 
-    public static void mapItem(ItemModel model, ItemStack stack, ProblemReporter reporter,
+    public static void mapItem(ItemModel.Unbaked model, ItemStack stack, ProblemReporter reporter,
                                Function<GeyserBaseDefinition, GeyserItemDefinition> definitionCreator, PackContext packContext) {
         mapItem(model, new MappingContext(List.of(), stack, reporter, definitionCreator, packContext));
     }
 
-    private static void mapItem(ItemModel model, MappingContext context) {
+    private static void mapItem(ItemModel.Unbaked model, MappingContext context) {
         switch (model) {
-            case BlockModelWrapper modelWrapper -> {
+            case BlockModelWrapper.Unbaked modelWrapper -> mapBlockModelWrapper(modelWrapper, context.child("plain model " + modelWrapper.model()));
-                ResourceLocation itemModelLocation = ((BlockModelWrapperLocationAccessor) modelWrapper).rainbow$getModelOrigin();
+            case ConditionalItemModel.Unbaked conditional -> mapConditionalModel(conditional, context.child("condition model "));
-
+            case RangeSelectItemModel.Unbaked rangeSelect -> mapRangeSelectModel(rangeSelect, context.child("range select model "));
-                ((ResolvedModelAccessor) Minecraft.getInstance().getModelManager()).rainbow$getResolvedModel(itemModelLocation)
+            case SelectItemModel.Unbaked select -> mapSelectModel(select, context.child("select model "));
-                        .ifPresentOrElse(itemModel -> {
+            default -> context.reporter.report(() -> "unsupported item model " + getModelId(model));
-                            ResolvedModel parentModel = itemModel.parent();
-                            // debugName() returns the resource location of the model as a string
-                            boolean handheld = parentModel != null && HANDHELD_MODELS.contains(ResourceLocation.parse(parentModel.debugName()));
-
-                            ResourceLocation bedrockIdentifier;
-                            if (itemModelLocation.getNamespace().equals(ResourceLocation.DEFAULT_NAMESPACE)) {
-                                bedrockIdentifier = ResourceLocation.fromNamespaceAndPath("geyser_mc", itemModelLocation.getPath());
-                            } else {
-                                bedrockIdentifier = itemModelLocation;
-                            }
-
-                            Material layer0Texture = itemModel.getTopTextureSlots().getMaterial("layer0");
-                            Optional<ResourceLocation> texture;
-                            Optional<ResolvedModel> customGeometry;
-                            if (layer0Texture != null) {
-                                texture = Optional.of(layer0Texture.texture());
-                                customGeometry = Optional.empty();
-                            } else {
-                                // We can't stitch multiple textures together yet, so we just grab the first one we see
-                                // This will only work properly for models with just one texture
-                                texture = ((TextureSlotsAccessor) itemModel.getTopTextureSlots()).getResolvedValues().values().stream()
-                                        .map(Material::texture)
-                                        .findAny();
-                                // Unknown texture (doesn't use layer0), so we immediately assume the geometry is custom
-                                // This check should probably be done differently
-                                customGeometry = Optional.of(itemModel);
-                            }
-
-                            texture.ifPresentOrElse(itemTexture -> {
-                                // Not a problem, but just report to get the model printed in the report file
-                                context.reporter.report(() -> "creating mapping for block model " + itemModelLocation);
-                                context.create(bedrockIdentifier, itemTexture, handheld, customGeometry);
-                            }, () -> context.reporter.report(() -> "not mapping block model " + itemModelLocation + " because it has no texture"));
-                        }, () -> context.reporter.report(() -> "missing block model " + itemModelLocation));
-            }
-            case ConditionalItemModel conditional -> mapConditionalModel(conditional, context.child("condition model "));
-            case SelectItemModel<?> select -> mapSelectModel(select, context.child("select model "));
-            default -> context.reporter.report(() -> "unsupported item model " + model.getClass()); // TODO intermediary stuff
         }
     }
 
-    private static void mapConditionalModel(ConditionalItemModel model, MappingContext context) {
+    private static void mapBlockModelWrapper(BlockModelWrapper.Unbaked model, MappingContext context) {
-        ItemModelPropertyTest property = ((ConditionalItemModelAccessor) model).getProperty();
+        ResourceLocation itemModelLocation = model.model();
+
+        getModels().rainbow$getResolvedModel(itemModelLocation)
+                .ifPresentOrElse(itemModel -> {
+                    ResolvedModel parentModel = itemModel.parent();
+                    // debugName() returns the resource location of the model as a string
+                    boolean handheld = parentModel != null && HANDHELD_MODELS.contains(ResourceLocation.parse(parentModel.debugName()));
+
+                    ResourceLocation bedrockIdentifier;
+                    if (itemModelLocation.getNamespace().equals(ResourceLocation.DEFAULT_NAMESPACE)) {
+                        bedrockIdentifier = ResourceLocation.fromNamespaceAndPath("geyser_mc", itemModelLocation.getPath());
+                    } else {
+                        bedrockIdentifier = itemModelLocation;
+                    }
+
+                    Material layer0Texture = itemModel.getTopTextureSlots().getMaterial("layer0");
+                    Optional<ResourceLocation> texture;
+                    Optional<ResolvedModel> customGeometry;
+                    if (layer0Texture != null) {
+                        texture = Optional.of(layer0Texture.texture());
+                        customGeometry = Optional.empty();
+                    } else {
+                        // We can't stitch multiple textures together yet, so we just grab the first one we see
+                        // This will only work properly for models with just one texture
+                        texture = ((TextureSlotsAccessor) itemModel.getTopTextureSlots()).getResolvedValues().values().stream()
+                                .map(Material::texture)
+                                .findAny();
+                        // Unknown texture (doesn't use layer0), so we immediately assume the geometry is custom
+                        // This check should probably be done differently
+                        customGeometry = Optional.of(itemModel);
+                    }
+
+                    texture.ifPresentOrElse(itemTexture -> {
+                        // Not a problem, but just report to get the model printed in the report file
+                        context.reporter.report(() -> "creating mapping for block model " + itemModelLocation);
+                        context.create(bedrockIdentifier, itemTexture, handheld, customGeometry);
+                    }, () -> context.reporter.report(() -> "not mapping block model " + itemModelLocation + " because it has no texture"));
+                }, () -> context.reporter.report(() -> "missing block model " + itemModelLocation));
+    }
+
+    private static void mapConditionalModel(ConditionalItemModel.Unbaked model, MappingContext context) {
+        ItemModelPropertyTest property = model.property();
         GeyserConditionPredicate.Property predicateProperty = switch (property) {
             case Broken ignored -> GeyserConditionPredicate.BROKEN;
             case Damaged ignored -> GeyserConditionPredicate.DAMAGED;
@@ -163,8 +178,8 @@ public class BedrockItemMapper {
             case FishingRodCast ignored -> GeyserConditionPredicate.FISHING_ROD_CAST;
             default -> null;
         };
-        ItemModel onTrue = ((ConditionalItemModelAccessor) model).getOnTrue();
+        ItemModel.Unbaked onTrue = model.onTrue();
-        ItemModel onFalse = ((ConditionalItemModelAccessor) model).getOnFalse();
+        ItemModel.Unbaked onFalse = model.onFalse();
 
         if (predicateProperty == null) {
             context.reporter.report(() -> "unsupported conditional model property " + property + ", only mapping on_false");
@@ -176,10 +191,32 @@ public class BedrockItemMapper {
         mapItem(onFalse, context.with(new GeyserConditionPredicate(predicateProperty, false), "condition on false "));
     }
 
+    private static void mapRangeSelectModel(RangeSelectItemModel.Unbaked model, MappingContext context) {
+        RangeSelectItemModelProperty property = model.property();
+        GeyserRangeDispatchPredicate.Property predicateProperty = switch (property) {
+            case BundleFullness ignored -> GeyserRangeDispatchPredicate.BUNDLE_FULLNESS;
+            case Count count -> new GeyserRangeDispatchPredicate.Count(count.normalize());
+            // Mojang, why? :(
+            case net.minecraft.client.renderer.item.properties.numeric.CustomModelDataProperty customModelData -> new GeyserRangeDispatchPredicate.CustomModelData(customModelData.index());
+            case Damage damage -> new GeyserRangeDispatchPredicate.Damage(damage.normalize());
+            default -> null;
+        };
+
+        if (predicateProperty == null) {
+            context.reporter.report(() -> "unsupported range dispatch model property " + property + ", only mapping fallback, if it is present");
+        } else {
+            for (RangeSelectItemModel.Entry entry : model.entries()) {
+                mapItem(entry.model(), context.with(new GeyserRangeDispatchPredicate(predicateProperty, entry.threshold(), model.scale()), "threshold " + entry.threshold()));
+            }
+        }
+
+        model.fallback().ifPresent(fallback -> mapItem(fallback, context.child("range dispatch fallback")));
+    }
+
     @SuppressWarnings("unchecked")
-    private static <T> void mapSelectModel(SelectItemModel<T> model, MappingContext context) {
+    private static void mapSelectModel(SelectItemModel.Unbaked model, MappingContext context) {
-        SelectItemModelProperty<T> property = ((SelectItemModelAccessor<T>) model).getProperty();
+        SelectItemModel.UnbakedSwitch<?, ?> unbakedSwitch = model.unbakedSwitch();
-        Function<T, GeyserMatchPredicate.MatchPredicateData> dataConstructor = switch (property) {
+        Function<Object, GeyserMatchPredicate.MatchPredicateData> dataConstructor = switch (unbakedSwitch.property()) {
             case Charge ignored -> chargeType -> new GeyserMatchPredicate.ChargeType((CrossbowItem.ChargeType) chargeType);
             case TrimMaterialProperty ignored -> material -> new GeyserMatchPredicate.TrimMaterialData((ResourceKey<TrimMaterial>) material);
             case ContextDimension ignored -> dimension -> new GeyserMatchPredicate.ContextDimension((ResourceKey<Level>) dimension);
@@ -188,24 +225,29 @@ public class BedrockItemMapper {
             default -> null;
         };
 
-        Object2ObjectMap<T, ItemModel> cases = ((SelectItemModelCasesAccessor<T>) model).rainbow$getCases();
+        List<? extends SelectItemModel.SwitchCase<?>> cases = unbakedSwitch.cases();
 
         if (dataConstructor == null) {
-            if (property instanceof DisplayContext) {
+            if (unbakedSwitch.property() instanceof DisplayContext) {
-                ItemModel gui = cases.get(ItemDisplayContext.GUI);
+                context.reporter.report(() -> "unsupported select model property display_context, only mapping \"gui\" case, if it exists");
-                if (gui != null) {
+                for (SelectItemModel.SwitchCase<?> switchCase : cases) {
-                    context.reporter.report(() -> "unsupported select model property display_context, only mapping \"gui\" case");
+                    if (switchCase.values().contains(ItemDisplayContext.GUI)) {
-                    mapItem(gui, context.child("select GUI display_context case (unsupported property) "));
+                        mapItem(switchCase.model(), context.child("select GUI display_context case (unsupported property) "));
-                    return;
+                        return;
+                    }
                 }
             }
-            context.reporter.report(() -> "unsupported select model property " + property + ", only mapping fallback");
+            context.reporter.report(() -> "unsupported select model property " + unbakedSwitch.property() + ", only mapping fallback, if present");
-            mapItem(cases.defaultReturnValue(), context.child("select fallback case (unsupported property) "));
+            model.fallback().ifPresent(fallback -> mapItem(fallback, context.child("select fallback case (unsupported property) ")));
             return;
         }
 
-        cases.forEach((key, value) -> mapItem(value, context.with(new GeyserMatchPredicate(dataConstructor.apply(key)), "select case " + key + " ")));
+        cases.forEach(switchCase -> {
-        mapItem(cases.defaultReturnValue(), context.child("select fallback case "));
+            switchCase.values().forEach(value -> {
+                mapItem(switchCase.model(), context.with(new GeyserMatchPredicate(dataConstructor.apply(value)), "select case " + value + " "));
+            });
+        });
+        model.fallback().ifPresent(fallback -> mapItem(fallback, context.child("select fallback case ")));
     }
 
     private record MappingContext(List<GeyserPredicate> predicateStack, ItemStack stack, ProblemReporter reporter,

src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserConditionPredicate.java

@@ -64,11 +64,7 @@ public record GeyserConditionPredicate(Property property, boolean expected) impl
     }
 
     public record CustomModelData(int index) implements Property {
-        public static final MapCodec<CustomModelData> CODEC = RecordCodecBuilder.mapCodec(instance ->
+        public static final MapCodec<CustomModelData> CODEC = ExtraCodecs.NON_NEGATIVE_INT.optionalFieldOf("index", 0).xmap(CustomModelData::new, CustomModelData::index);
-                instance.group(
-                        ExtraCodecs.NON_NEGATIVE_INT.optionalFieldOf("index", 0).forGetter(CustomModelData::index)
-                ).apply(instance, CustomModelData::new)
-        );
 
         @Override
         public Type type() {
@@ -77,11 +73,7 @@ public record GeyserConditionPredicate(Property property, boolean expected) impl
     }
 
     public record HasComponent(DataComponentType<?> component) implements Property {
-        public static final MapCodec<HasComponent> CODEC = RecordCodecBuilder.mapCodec(instance ->
+        public static final MapCodec<HasComponent> CODEC = DataComponentType.CODEC.fieldOf("component").xmap(HasComponent::new, HasComponent::component);
-                instance.group(
-                        DataComponentType.CODEC.fieldOf("component").forGetter(HasComponent::component)
-                ).apply(instance, HasComponent::new)
-        );
 
         @Override
         public Type type() {

src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserPredicate.java

@@ -17,7 +17,8 @@ public interface GeyserPredicate {
 
     enum Type implements StringRepresentable {
         CONDITION("condition", GeyserConditionPredicate.CODEC),
-        MATCH("match", GeyserMatchPredicate.CODEC);
+        MATCH("match", GeyserMatchPredicate.CODEC),
+        RANGE_DISPATCH("range_dispatch", GeyserRangeDispatchPredicate.CODEC);
 
         public static final Codec<Type> CODEC = StringRepresentable.fromEnum(Type::values);
 

src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserRangeDispatchPredicate.java

@@ -0,0 +1,93 @@
+package org.geysermc.rainbow.mapping.geyser.predicate;
+
+import com.google.common.base.Suppliers;
+import com.mojang.serialization.Codec;
+import com.mojang.serialization.MapCodec;
+import com.mojang.serialization.codecs.RecordCodecBuilder;
+import net.minecraft.util.ExtraCodecs;
+import net.minecraft.util.StringRepresentable;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.function.Supplier;
+
+public record GeyserRangeDispatchPredicate(Property property, float threshold, float scale) implements GeyserPredicate {
+
+    public static final MapCodec<GeyserRangeDispatchPredicate> CODEC = RecordCodecBuilder.mapCodec(instance ->
+            instance.group(
+                    Property.CODEC.forGetter(GeyserRangeDispatchPredicate::property),
+                    Codec.FLOAT.fieldOf("threshold").forGetter(GeyserRangeDispatchPredicate::threshold),
+                    Codec.FLOAT.fieldOf("scale").forGetter(GeyserRangeDispatchPredicate::scale)
+            ).apply(instance, GeyserRangeDispatchPredicate::new)
+    );
+
+    public static final Property BUNDLE_FULLNESS = unit(Property.Type.BUNDLE_FULLNESS);
+
+    @Override
+    public Type type() {
+        return null;
+    }
+
+    public interface Property {
+
+        MapCodec<Property> CODEC = Type.CODEC.dispatchMap("property", Property::type, Type::codec);
+
+        Type type();
+
+        enum Type implements StringRepresentable {
+            BUNDLE_FULLNESS("bundle_fullness", () -> MapCodec.unit(GeyserRangeDispatchPredicate.BUNDLE_FULLNESS)),
+            DAMAGE("damage", () -> Damage.CODEC),
+            COUNT("count", () -> Count.CODEC),
+            CUSTOM_MODEL_DATA("custom_model_data", () -> CustomModelData.CODEC);
+
+            public static final Codec<Type> CODEC = StringRepresentable.fromEnum(Type::values);
+
+            private final String name;
+            private final Supplier<MapCodec<? extends Property>> codec;
+
+            Type(String name, Supplier<MapCodec<? extends Property>> codec) {
+                this.name = name;
+                this.codec = Suppliers.memoize(codec::get);
+            }
+
+            public MapCodec<? extends Property> codec() {
+                return codec.get();
+            }
+
+            @Override
+            public @NotNull String getSerializedName() {
+                return name;
+            }
+        }
+    }
+
+    public record Damage(boolean normalize) implements Property {
+        public static final MapCodec<Damage> CODEC = Codec.BOOL.fieldOf("normalize").xmap(Damage::new, Damage::normalize);
+
+        @Override
+        public Type type() {
+            return Type.DAMAGE;
+        }
+    }
+
+    public record Count(boolean normalize) implements Property {
+        public static final MapCodec<Count> CODEC = Codec.BOOL.fieldOf("normalize").xmap(Count::new, Count::normalize);
+
+        @Override
+        public Type type() {
+            return Type.COUNT;
+        }
+    }
+
+    public record CustomModelData(int index) implements Property {
+        public static final MapCodec<CustomModelData> CODEC = ExtraCodecs.NON_NEGATIVE_INT.optionalFieldOf("index", 0).xmap(CustomModelData::new, CustomModelData::index);
+
+        @Override
+        public Type type() {
+            return Type.CUSTOM_MODEL_DATA;
+        }
+    }
+
+    private static Property unit(Property.Type type) {
+        return () -> type;
+    }
+}

src/main/java/org/geysermc/rainbow/mixin/BlockModelWrapperMixin.java

@@ -1,43 +0,0 @@
-package org.geysermc.rainbow.mixin;
-
-import net.minecraft.client.renderer.item.BlockModelWrapper;
-import net.minecraft.client.renderer.item.ItemModel;
-import net.minecraft.resources.ResourceLocation;
-import org.geysermc.rainbow.accessor.BlockModelWrapperLocationAccessor;
-import org.spongepowered.asm.mixin.Final;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.Shadow;
-import org.spongepowered.asm.mixin.Unique;
-import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
-
-@Mixin(BlockModelWrapper.class)
-public abstract class BlockModelWrapperMixin implements ItemModel, BlockModelWrapperLocationAccessor {
-
-    @Unique
-    private ResourceLocation modelOrigin;
-
-    @Override
-    public ResourceLocation rainbow$getModelOrigin() {
-        return modelOrigin;
-    }
-
-    @Override
-    public void rainbow$setModelOrigin(ResourceLocation model) {
-        modelOrigin = model;
-    }
-
-    @Mixin(BlockModelWrapper.Unbaked.class)
-    public abstract static class UnbakedMixin implements Unbaked {
-
-        @Shadow
-        @Final
-        private ResourceLocation model;
-
-        @Inject(method = "bake", at = @At("TAIL"))
-        public void setModelOrigin(BakingContext context, CallbackInfoReturnable<ItemModel> callbackInfoReturnable) {
-            ((BlockModelWrapperLocationAccessor) callbackInfoReturnable.getReturnValue()).rainbow$setModelOrigin(model);
-        }
-    }
-}

src/main/java/org/geysermc/rainbow/mixin/ConditionalItemModelAccessor.java

@@ -1,20 +0,0 @@
-package org.geysermc.rainbow.mixin;
-
-import net.minecraft.client.renderer.item.ConditionalItemModel;
-import net.minecraft.client.renderer.item.ItemModel;
-import net.minecraft.client.renderer.item.properties.conditional.ItemModelPropertyTest;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.gen.Accessor;
-
-@Mixin(ConditionalItemModel.class)
-public interface ConditionalItemModelAccessor {
-
-    @Accessor
-    ItemModelPropertyTest getProperty();
-
-    @Accessor
-    ItemModel getOnTrue();
-
-    @Accessor
-    ItemModel getOnFalse();
-}

src/main/java/org/geysermc/rainbow/mixin/LateBoundIdMapperAccessor.java

@@ -0,0 +1,13 @@
+package org.geysermc.rainbow.mixin;
+
+import com.google.common.collect.BiMap;
+import net.minecraft.util.ExtraCodecs;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ExtraCodecs.LateBoundIdMapper.class)
+public interface LateBoundIdMapperAccessor<I, V> {
+
+    @Accessor
+    BiMap<I, V> getIdToValue();
+}

src/main/java/org/geysermc/rainbow/mixin/ModelManagerMixin.java

@@ -3,6 +3,8 @@ package org.geysermc.rainbow.mixin;
 import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
 import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 import net.minecraft.client.Minecraft;
+import net.minecraft.client.renderer.item.ClientItem;
+import net.minecraft.client.resources.model.ClientItemInfoLoader;
 import net.minecraft.client.resources.model.ModelManager;
 import net.minecraft.client.resources.model.ResolvedModel;
 import net.minecraft.resources.ResourceLocation;
@@ -21,6 +23,8 @@ import java.util.concurrent.CompletableFuture;
 public abstract class ModelManagerMixin implements PreparableReloadListener, AutoCloseable, ResolvedModelAccessor {
     @Unique
     private Map<ResourceLocation, ResolvedModel> unbakedResolvedModels;
+    @Unique
+    private Map<ResourceLocation, ClientItem> clientItems;
 
     @WrapOperation(method = "method_65753", at = @At(value = "INVOKE", target = "Ljava/util/concurrent/CompletableFuture;join()Ljava/lang/Object;", ordinal = 1))
     private static Object setResolvedModels(CompletableFuture<?> instance, Operation<Object> original) {
@@ -35,8 +39,21 @@ public abstract class ModelManagerMixin implements PreparableReloadListener, Aut
         return resolved;
     }
 
+    @WrapOperation(method = "method_65753", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/resources/model/ClientItemInfoLoader$LoadedClientInfos;contents()Ljava/util/Map;"))
+    private static Map<ResourceLocation, ClientItem> setClientItems(ClientItemInfoLoader.LoadedClientInfos instance, Operation<Map<ResourceLocation, ClientItem>> original) {
+        // Same note as above for not using "this"
+        ModelManagerMixin thiz = ((ModelManagerMixin) (Object) Minecraft.getInstance().getModelManager());
+        thiz.clientItems = original.call(instance);
+        return thiz.clientItems;
+    }
+
     @Override
     public Optional<ResolvedModel> rainbow$getResolvedModel(ResourceLocation location) {
         return unbakedResolvedModels == null ? Optional.empty() : Optional.ofNullable(unbakedResolvedModels.get(location));
     }
+
+    @Override
+    public Optional<ClientItem> rainbow$getClientItem(ResourceLocation location) {
+        return clientItems == null ? Optional.empty() : Optional.ofNullable(clientItems.get(location));
+    }
 }

src/main/java/org/geysermc/rainbow/mixin/RangeSelectItemModelAccessor.java

@@ -1,30 +1,12 @@
 package org.geysermc.rainbow.mixin;
 
-import net.minecraft.client.renderer.item.ItemModel;
 import net.minecraft.client.renderer.item.RangeSelectItemModel;
-import net.minecraft.client.renderer.item.properties.numeric.RangeSelectItemModelProperty;
 import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.gen.Accessor;
 import org.spongepowered.asm.mixin.gen.Invoker;
 
 @Mixin(RangeSelectItemModel.class)
 public interface RangeSelectItemModelAccessor {
 
-    @Accessor
-    RangeSelectItemModelProperty getProperty();
-
-    @Accessor
-    float getScale();
-
-    @Accessor
-    float[] getThresholds();
-
-    @Accessor
-    ItemModel[] getModels();
-
-    @Accessor
-    ItemModel getFallback();
-
     @Invoker
     static int invokeLastIndexLessOrEqual(float[] thresholds, float value) {
         throw new AssertionError();

src/main/java/org/geysermc/rainbow/mixin/SelectItemModelAccessor.java

@@ -1,13 +0,0 @@
-package org.geysermc.rainbow.mixin;
-
-import net.minecraft.client.renderer.item.SelectItemModel;
-import net.minecraft.client.renderer.item.properties.select.SelectItemModelProperty;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.gen.Accessor;
-
-@Mixin(SelectItemModel.class)
-public interface SelectItemModelAccessor<T> {
-
-    @Accessor
-    SelectItemModelProperty<T> getProperty();
-}

src/main/java/org/geysermc/rainbow/mixin/SelectItemModelMixin.java

@@ -1,41 +0,0 @@
-package org.geysermc.rainbow.mixin;
-
-import com.llamalad7.mixinextras.sugar.Local;
-import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
-import net.minecraft.client.renderer.item.ItemModel;
-import net.minecraft.client.renderer.item.SelectItemModel;
-import net.minecraft.client.renderer.item.properties.select.SelectItemModelProperty;
-import org.geysermc.rainbow.accessor.SelectItemModelCasesAccessor;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.Unique;
-import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
-
-@Mixin(SelectItemModel.class)
-public abstract class SelectItemModelMixin<T> implements ItemModel, SelectItemModelCasesAccessor<T> {
-
-    @Unique
-    private Object2ObjectMap<T, ItemModel> cases;
-
-    @Override
-    public Object2ObjectMap<T, ItemModel> rainbow$getCases() {
-        return cases;
-    }
-
-    @Override
-    public void rainbow$setCases(Object2ObjectMap<T, ItemModel> cases) {
-        this.cases = cases;
-    }
-
-    @Mixin(SelectItemModel.UnbakedSwitch.class)
-    public abstract static class UnbakedSwitchMixin<P extends SelectItemModelProperty<T>, T> {
-
-        @Inject(method = "bake", at = @At("TAIL"))
-        public void setCases(BakingContext bakingContext, ItemModel model, CallbackInfoReturnable<ItemModel> callbackInfoReturnable,
-                             @Local Object2ObjectMap<T, ItemModel> cases) {
-            //noinspection unchecked
-            ((SelectItemModelCasesAccessor<T>) callbackInfoReturnable.getReturnValue()).rainbow$setCases(cases);
-        }
-    }
-}

src/main/resources/rainbow.mixins.json

@@ -5,18 +5,13 @@
   "compatibilityLevel": "JAVA_21",
   "mixins": [],
   "client": [
-    "BlockModelWrapperMixin",
-    "BlockModelWrapperMixin$UnbakedMixin",
-    "ConditionalItemModelAccessor",
     "EntityRenderDispatcherAccessor",
     "GuiItemRenderStateMixin",
+    "LateBoundIdMapperAccessor",
     "ModelManagerMixin",
     "PictureInPictureRendererAccessor",
     "PictureInPictureRendererMixin",
     "RangeSelectItemModelAccessor",
-    "SelectItemModelAccessor",
-    "SelectItemModelMixin",
-    "SelectItemModelMixin$UnbakedSwitchMixin",
     "SplashRendererAccessor",
     "TextureSlotsAccessor"
   ],