Implement mapping of range dispatch models (#7)

2025-12-19 14:59:16 +00:00 · 2025-10-13 17:35:57 +00:00
parent 0cecce88ee c8caa26ed0
commit 2a3428229e
15 changed files with 260 additions and 261 deletions
--- a/src/main/java/org/geysermc/rainbow/accessor/BlockModelWrapperLocationAccessor.java
+++ b/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);
-}
--- a/src/main/java/org/geysermc/rainbow/accessor/ResolvedModelAccessor.java
+++ b/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);
 }
--- a/src/main/java/org/geysermc/rainbow/accessor/SelectItemModelCasesAccessor.java
+++ b/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);
-}
--- a/src/main/java/org/geysermc/rainbow/mapping/BedrockItemMapper.java
+++ b/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,47 +72,65 @@ 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);
-        mapItem(model, stack, reporter.forChild(() -> "client item definition " + modelLocation + " "), base -> new GeyserSingleDefinition(base, Optional.of(modelLocation)), context);
+        getModels().rainbow$getClientItem(modelLocation).map(ClientItem::model)
+                .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));
-        reporter = reporter.forChild(() -> "item model " + vanillaModel + " with custom model data " + customModelData + " ");
-        if (vanillaModel instanceof RangeSelectItemModel rangeModel) {
-            RangeSelectItemModelAccessor accessor = (RangeSelectItemModelAccessor) rangeModel;
-            RangeSelectItemModelProperty property = accessor.getProperty();
+        ItemModel.Unbaked vanillaModel = getModels().rainbow$getClientItem(stack.get(DataComponents.ITEM_MODEL)).map(ClientItem::model).orElseThrow();
+        ProblemReporter childReporter = reporter.forChild(() -> "item model " + vanillaModel + " with custom model data " + customModelData + " ");
+        if (vanillaModel instanceof RangeSelectItemModel.Unbaked(RangeSelectItemModelProperty property, float scale, List<RangeSelectItemModel.Entry> entries, Optional<ItemModel.Unbaked> fallback)) {
            // 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);
-                    ItemModel model = modelIndex == -1 ? accessor.getFallback() : accessor.getModels()[modelIndex];
-                    mapItem(model, stack, reporter, base -> new GeyserLegacyDefinition(base, customModelData), context);
+                    float[] thresholds = ArrayUtils.toPrimitive(entries.stream()
+                            .map(RangeSelectItemModel.Entry::threshold)
+                            .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 {
-                reporter.report(() -> "range_dispatch model property is not custom model data, unable to apply custom model data");
+                return;
            }
-        } 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 -> {
-                ResourceLocation itemModelLocation = ((BlockModelWrapperLocationAccessor) modelWrapper).rainbow$getModelOrigin();
+            case BlockModelWrapper.Unbaked modelWrapper -> mapBlockModelWrapper(modelWrapper, context.child("plain model " + modelWrapper.model()));
+            case ConditionalItemModel.Unbaked conditional -> mapConditionalModel(conditional, context.child("condition model "));
+            case RangeSelectItemModel.Unbaked rangeSelect -> mapRangeSelectModel(rangeSelect, context.child("range select model "));
+            case SelectItemModel.Unbaked select -> mapSelectModel(select, context.child("select model "));
+            default -> context.reporter.report(() -> "unsupported item model " + getModelId(model));
+        }
+    }

-                ((ResolvedModelAccessor) Minecraft.getInstance().getModelManager()).rainbow$getResolvedModel(itemModelLocation)
+    private static void mapBlockModelWrapper(BlockModelWrapper.Unbaked model, MappingContext context) {
+        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
@@ -147,14 +167,9 @@ public class BedrockItemMapper {
                    }, () -> 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) {
-        ItemModelPropertyTest property = ((ConditionalItemModelAccessor) model).getProperty();
+    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 onFalse = ((ConditionalItemModelAccessor) model).getOnFalse();
+        ItemModel.Unbaked onTrue = model.onTrue();
+        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) {
-        SelectItemModelProperty<T> property = ((SelectItemModelAccessor<T>) model).getProperty();
-        Function<T, GeyserMatchPredicate.MatchPredicateData> dataConstructor = switch (property) {
+    private static void mapSelectModel(SelectItemModel.Unbaked model, MappingContext context) {
+        SelectItemModel.UnbakedSwitch<?, ?> unbakedSwitch = model.unbakedSwitch();
+        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) {
-                ItemModel gui = cases.get(ItemDisplayContext.GUI);
-                if (gui != null) {
-                    context.reporter.report(() -> "unsupported select model property display_context, only mapping \"gui\" case");
-                    mapItem(gui, context.child("select GUI display_context case (unsupported property) "));
+            if (unbakedSwitch.property() instanceof DisplayContext) {
+                context.reporter.report(() -> "unsupported select model property display_context, only mapping \"gui\" case, if it exists");
+                for (SelectItemModel.SwitchCase<?> switchCase : cases) {
+                    if (switchCase.values().contains(ItemDisplayContext.GUI)) {
+                        mapItem(switchCase.model(), context.child("select GUI display_context case (unsupported property) "));
                        return;
                    }
                }
-            context.reporter.report(() -> "unsupported select model property " + property + ", only mapping fallback");
-            mapItem(cases.defaultReturnValue(), context.child("select fallback case (unsupported property) "));
+            }
+            context.reporter.report(() -> "unsupported select model property " + unbakedSwitch.property() + ", only mapping fallback, if present");
+            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 + " ")));
-        mapItem(cases.defaultReturnValue(), context.child("select fallback case "));
+        cases.forEach(switchCase -> {
+            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,
--- a/src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserConditionPredicate.java
+++ b/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 ->
-                instance.group(
-                        ExtraCodecs.NON_NEGATIVE_INT.optionalFieldOf("index", 0).forGetter(CustomModelData::index)
-                ).apply(instance, CustomModelData::new)
-        );
+        public static final MapCodec<CustomModelData> CODEC = ExtraCodecs.NON_NEGATIVE_INT.optionalFieldOf("index", 0).xmap(CustomModelData::new, CustomModelData::index);

        @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 ->
-                instance.group(
-                        DataComponentType.CODEC.fieldOf("component").forGetter(HasComponent::component)
-                ).apply(instance, HasComponent::new)
-        );
+        public static final MapCodec<HasComponent> CODEC = DataComponentType.CODEC.fieldOf("component").xmap(HasComponent::new, HasComponent::component);

        @Override
        public Type type() {
--- a/src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserPredicate.java
+++ b/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);

--- a/src/main/java/org/geysermc/rainbow/mapping/geyser/predicate/GeyserRangeDispatchPredicate.java
+++ b/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;
+    }
+}
--- a/src/main/java/org/geysermc/rainbow/mixin/BlockModelWrapperMixin.java
+++ b/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);
-        }
-    }
-}
--- a/src/main/java/org/geysermc/rainbow/mixin/ConditionalItemModelAccessor.java
+++ b/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();
-}
--- a/src/main/java/org/geysermc/rainbow/mixin/LateBoundIdMapperAccessor.java
+++ b/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();
+}
--- a/src/main/java/org/geysermc/rainbow/mixin/ModelManagerMixin.java
+++ b/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));
+    }
 }
--- a/src/main/java/org/geysermc/rainbow/mixin/RangeSelectItemModelAccessor.java
+++ b/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();
--- a/src/main/java/org/geysermc/rainbow/mixin/SelectItemModelAccessor.java
+++ b/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();
-}
--- a/src/main/java/org/geysermc/rainbow/mixin/SelectItemModelMixin.java
+++ b/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);
-        }
-    }
-}
--- a/src/main/resources/rainbow.mixins.json
+++ b/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"
  ],