Implement mapping of range dispatch models

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/mapping/geyser/predicate/GeyserConditionPredicate.java

@@ -64,11 +64,7 @@ public MapCodec<? extends Property> codec() {
     }
 
     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 Type type() {
     }
 
     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() {

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/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 @@
 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 @@
 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 @@ private static Object setResolvedModels(CompletableFuture<?> instance, Operation
         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/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"
   ],