Use the proper atlas for a material when loading its texture

Author: eclipseisoffline

rainbow/src/main/java/org/geysermc/rainbow/Rainbow.java

@@ -1,6 +1,9 @@
 package org.geysermc.rainbow;
 
 import com.mojang.logging.LogUtils;
+import net.minecraft.client.renderer.texture.TextureAtlas;
+import net.minecraft.client.resources.model.Material;
+import net.minecraft.data.AtlasIds;
 import net.minecraft.resources.Identifier;
 import org.slf4j.Logger;
 
@@ -25,4 +28,14 @@ public static Identifier decorateIdentifier(Identifier identifier, String type,
     public static Identifier decorateTextureIdentifier(Identifier identifier) {
         return decorateIdentifier(identifier, "textures", "png");
     }
+
+    public static Identifier getAtlasIdFromMaterial(Material material) {
+        Identifier atlasLocation = material.atlasLocation();
+        if (atlasLocation.equals(TextureAtlas.LOCATION_BLOCKS)) {
+            return AtlasIds.BLOCKS;
+        } else if (atlasLocation.equals(TextureAtlas.LOCATION_ITEMS)) {
+            return AtlasIds.ITEMS;
+        }
+        return atlasLocation;
+    }
 }

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

@@ -2,7 +2,9 @@
 
 import net.minecraft.client.renderer.item.ClientItem;
 import net.minecraft.client.resources.model.EquipmentClientInfo;
+import net.minecraft.client.resources.model.ModelManager;
 import net.minecraft.client.resources.model.ResolvedModel;
+import net.minecraft.data.AtlasIds;
 import net.minecraft.resources.Identifier;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.world.item.equipment.EquipmentAsset;
@@ -19,4 +21,12 @@ public interface AssetResolver {
     Optional<EquipmentClientInfo> getEquipmentInfo(ResourceKey<EquipmentAsset> key);
 
     Optional<TextureResource> getTexture(Identifier atlas, Identifier identifier);
+
+    default Optional<TextureResource> getTextureSafely(Identifier atlas, Identifier identifier) {
+        if (atlas.equals(ModelManager.BLOCK_OR_ITEM)) {
+            return getTexture(AtlasIds.BLOCKS, identifier)
+                    .or(() -> getTexture(AtlasIds.ITEMS, identifier));
+        }
+        return getTexture(atlas, identifier);
+    }
 }

rainbow/src/main/java/org/geysermc/rainbow/mapping/geometry/BedrockGeometryContext.java

@@ -3,7 +3,6 @@
 import net.minecraft.client.renderer.block.model.TextureSlots;
 import net.minecraft.client.resources.model.Material;
 import net.minecraft.client.resources.model.ResolvedModel;
-import net.minecraft.data.AtlasIds;
 import net.minecraft.resources.Identifier;
 import net.minecraft.world.item.ItemStack;
 import org.geysermc.rainbow.Rainbow;
@@ -37,7 +36,7 @@ public static BedrockGeometryContext create(Identifier bedrockIdentifier, Resolv
         if (layer0Texture != null) {
             geometry = Optional.empty();
             animation = Optional.empty();
-            icon = TextureHolder.createBuiltIn(AtlasIds.BLOCKS, layer0Texture.texture());
+            icon = TextureHolder.createBuiltIn(Rainbow.getAtlasIdFromMaterial(layer0Texture), layer0Texture.texture());
         } else {
             // Unknown model (doesn't use layer0), so we immediately assume the geometry is custom
             // This check should probably be done differently (actually check if the model is 2D or 3D)

rainbow/src/main/java/org/geysermc/rainbow/mapping/texture/BuiltInTextureHolder.java

@@ -23,7 +23,7 @@ public BuiltInTextureHolder(Identifier identifier, Identifier atlas, Identifier
     @Override
     public Optional<byte[]> load(AssetResolver assetResolver, ProblemReporter reporter) {
         return RainbowIO.safeIO(() -> {
-            try (TextureResource texture = assetResolver.getTexture(atlas, source).orElse(null)) {
+            try (TextureResource texture = assetResolver.getTextureSafely(atlas, source).orElse(null)) {
                 Objects.requireNonNull(texture);
                 try (NativeImage firstFrame = texture.getFirstFrame(false)) {
                     return NativeImageUtil.writeToByteArray(firstFrame);

rainbow/src/main/java/org/geysermc/rainbow/mapping/texture/StitchedTextures.java

@@ -7,8 +7,8 @@
 import net.minecraft.client.renderer.texture.TextureAtlasSprite;
 import net.minecraft.client.resources.model.Material;
 import net.minecraft.data.AtlasIds;
-import net.minecraft.resources.Identifier;
 import net.minecraft.util.Util;
+import org.geysermc.rainbow.Rainbow;
 import org.geysermc.rainbow.RainbowIO;
 import org.geysermc.rainbow.mapping.PackContext;
 import org.geysermc.rainbow.mixin.SpriteContentsAccessor;
@@ -37,7 +37,7 @@ public Optional<TextureAtlasSprite> getSprite(String key) {
 
     public static StitchedTextures stitchModelTextures(TextureSlots textures, PackContext context) {
         Map<String, Material> materials = ((TextureSlotsAccessor) textures).getResolvedValues();
-        SpriteLoader.Preparations preparations = prepareStitching(materials.values().stream().map(Material::texture), context);
+        SpriteLoader.Preparations preparations = prepareStitching(materials.values().stream(), context);
 
         Map<String, TextureAtlasSprite> sprites = new HashMap<>();
         for (Map.Entry<String, Material> material : materials.entrySet()) {
@@ -50,21 +50,21 @@ public static StitchedTextures stitchModelTextures(TextureSlots textures, PackCo
         return new StitchedTextures(Map.copyOf(sprites), () -> stitchTextureAtlas(preparations), preparations.width(), preparations.height());
     }
 
-    private static SpriteLoader.Preparations prepareStitching(Stream<Identifier> textures, PackContext context) {
-        // Atlas ID doesn't matter much here, but BLOCKS is the most appropriate
-        SpriteLoader spriteLoader = new SpriteLoader(AtlasIds.BLOCKS, MAX_TEXTURE_SIZE);
-        List<SpriteContents> sprites = textures.distinct()
-                .map(texture -> readSpriteContents(texture, context))
+    private static SpriteLoader.Preparations prepareStitching(Stream<Material> materials, PackContext context) {
+        // Atlas ID doesn't matter much here, but ITEMS is the most appropriate (though not always right)
+        SpriteLoader spriteLoader = new SpriteLoader(AtlasIds.ITEMS, MAX_TEXTURE_SIZE);
+        List<SpriteContents> sprites = materials.distinct()
+                .map(material -> readSpriteContents(material, context))
                 .<SpriteContents>mapMulti(Optional::ifPresent)
                 .toList();
         return  ((SpriteLoaderAccessor) spriteLoader).invokeStitch(sprites, 0, Util.backgroundExecutor());
     }
 
-    private static Optional<SpriteContents> readSpriteContents(Identifier identifier, PackContext context) {
+    private static Optional<SpriteContents> readSpriteContents(Material material, PackContext context) {
         return RainbowIO.safeIO(() -> {
-            try (TextureResource texture = context.assetResolver().getTexture(AtlasIds.BLOCKS, identifier).orElse(null)) {
+            try (TextureResource texture = context.assetResolver().getTextureSafely(Rainbow.getAtlasIdFromMaterial(material), material.texture()).orElse(null)) {
                 if (texture != null) {
-                    return new SpriteContents(identifier, texture.sizeOfFrame(), texture.getFirstFrame(true));
+                    return new SpriteContents(material.texture(), texture.sizeOfFrame(), texture.getFirstFrame(true));
                 }
             }
             return null;