cache decoded images during stitching

src/main/java/com/cricketcraft/chisel/mixins/MixinFallbackResourceManager.java

@@ -40,7 +40,7 @@ public abstract class MixinFallbackResourceManager {
         if (submap == null) return;
 
         IResource source = getResource(submap.source());
-        cir.setReturnValue(TextureSubmapResource.create(location, source, submap));
+        cir.setReturnValue(new TextureSubmapResource(source, submap, true));
     }
 
     @Inject(method = "getAllResources", at = @At("HEAD"), cancellable = true)
@@ -54,7 +54,7 @@ public abstract class MixinFallbackResourceManager {
         List<IResource> resources = new ArrayList<>(sources.size());
 
         for (IResource source : sources) {
-            resources.add(TextureSubmapResource.create(location, source, submap));
+            resources.add(new TextureSubmapResource(source, submap, false));
         }
 
         cir.setReturnValue(resources);

src/main/java/team/chisel/ctmlib/TextureSubmap.java

@@ -99,11 +99,18 @@ public class TextureSubmap implements IIcon, ISubmap {
         @SubscribeEvent
         public void onTextureStitchPre(TextureStitchEvent.Pre event) {
             if (event.map.getTextureType() != 0) return;
+            TextureSubmapResource.beginTextureStitch();
             for (TextureSubmap submap : submaps) {
                 submap.registerSubIcons(event.map);
             }
             submaps.clear();
         }
+
+        @SubscribeEvent
+        public void onTextureStitchPost(TextureStitchEvent.Post event) {
+            if (event.map.getTextureType() != 0) return;
+            TextureSubmapResource.endTextureStitch();
+        }
     }
 
     /**

src/main/java/team/chisel/ctmlib/TextureSubmapResource.java

@@ -6,9 +6,17 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 
+import javax.annotation.Nullable;
+import javax.imageio.IIOImage;
 import javax.imageio.ImageIO;
+import javax.imageio.ImageWriteParam;
+import javax.imageio.ImageWriter;
+import javax.imageio.stream.ImageOutputStream;
+import javax.imageio.stream.MemoryCacheImageOutputStream;
 
 import net.minecraft.client.resources.IResource;
 import net.minecraft.client.resources.data.AnimationFrame;
@@ -23,89 +31,23 @@ public final class TextureSubmapResource implements IResource {
     @Desugar
     public record SubmapLocation(ResourceLocation source, int columns, int rows, int cellX, int cellY) {}
 
+    @Desugar
+    private record SubmapCacheKey(ResourceLocation source, int columns, int rows) {}
+
+    @Nullable
+    private static HashMap<SubmapCacheKey, byte[][]> imageCache;
+
     private static final String PREFIX = "__chisel_submap/";
 
-    private final IResource source;
-    private final byte[] imageData;
-    private final int columns;
-    private final int rows;
+    private final IResource resource;
+    private final SubmapLocation submap;
+    private final boolean cacheable;
+    private byte[] imageData;
 
-    private TextureSubmapResource(IResource source, byte[] imageData, int columns, int rows) {
-        this.source = source;
-        this.imageData = imageData;
-        this.columns = columns;
-        this.rows = rows;
-    }
-
-    public static String getSubIconName(ResourceLocation source, int columns, int rows, int cellX, int cellY) {
-        return source.getResourceDomain() + ":"
-            + PREFIX
-            + columns
-            + "x"
-            + rows
-            + "/"
-            + cellX
-            + "_"
-            + cellY
-            + "/"
-            + source.getResourcePath();
-    }
-
-    public static IResource create(ResourceLocation location, IResource source, SubmapLocation submap)
-        throws IOException {
-        BufferedImage sourceImage;
-        try (InputStream stream = source.getInputStream()) {
-            sourceImage = ImageIO.read(stream);
-        }
-        if (sourceImage == null) {
-            throw new IOException("Unable to decode " + submap.source);
-        }
-
-        BufferedImage cropped = crop(sourceImage, submap);
-        ByteArrayOutputStream output = new ByteArrayOutputStream();
-
-        if (!ImageIO.write(cropped, "png", output)) {
-            throw new IOException("Unable to encode submap resource " + location);
-        }
-
-        return new TextureSubmapResource(source, output.toByteArray(), submap.columns, submap.rows);
-    }
-
-    private static BufferedImage crop(BufferedImage source, SubmapLocation submap) throws IOException {
-        int frameSize = source.getWidth();
-        int sourceHeight = source.getHeight();
-
-        if (frameSize < submap.columns || frameSize < submap.rows
-            || frameSize % submap.columns != 0
-            || frameSize % submap.rows != 0
-            || sourceHeight % frameSize != 0) {
-
-            throw new IOException(
-                "Texture size " + frameSize
-                    + "x"
-                    + sourceHeight
-                    + " cannot be split into submap "
-                    + submap.columns
-                    + "x"
-                    + submap.rows);
-        }
-
-        int cellWidth = frameSize / submap.columns;
-        int cellHeight = frameSize / submap.rows;
-        int frameCount = sourceHeight / frameSize;
-
-        BufferedImage result = new BufferedImage(cellWidth, cellHeight * frameCount, BufferedImage.TYPE_INT_ARGB);
-        int[] pixels = new int[cellWidth * cellHeight];
-
-        for (int frame = 0; frame < frameCount; frame++) {
-            int sourceX = submap.cellX * cellWidth;
-            int sourceY = frame * frameSize + submap.cellY * cellHeight;
-
-            source.getRGB(sourceX, sourceY, cellWidth, cellHeight, pixels, 0, cellWidth);
-            result.setRGB(0, frame * cellHeight, cellWidth, cellHeight, pixels, 0, cellWidth);
-        }
-
-        return result;
+    public TextureSubmapResource(IResource resource, SubmapLocation submap, boolean cacheable) {
+        this.resource = resource;
+        this.submap = submap;
+        this.cacheable = cacheable;
     }
 
     public static SubmapLocation parse(ResourceLocation location) {
@@ -146,22 +88,162 @@ public final class TextureSubmapResource implements IResource {
         }
     }
 
+    public static String getSubIconName(ResourceLocation source, int columns, int rows, int cellX, int cellY) {
+        return source.getResourceDomain() + ":"
+            + PREFIX
+            + columns
+            + "x"
+            + rows
+            + "/"
+            + cellX
+            + "_"
+            + cellY
+            + "/"
+            + source.getResourcePath();
+    }
+
+    public static void beginTextureStitch() {
+        imageCache = new HashMap<>();
+    }
+
+    public static void endTextureStitch() {
+        imageCache = null;
+    }
+
     @Override
     public InputStream getInputStream() {
+        if (imageData == null) {
+            try {
+                imageData = getCachedImageData(resource, submap);
+            } catch (IOException e) {
+                throw new RuntimeException("Failed to generate submap resource " + submap.source(), e);
+            }
+        }
+
         return new ByteArrayInputStream(imageData);
     }
 
+    private byte[] getCachedImageData(IResource resource, SubmapLocation submap) throws IOException {
+        if (!cacheable || imageCache == null) {
+            return generateCell(resource, submap);
+        }
+
+        SubmapCacheKey key = new SubmapCacheKey(submap.source(), submap.columns(), submap.rows());
+        byte[][] cells = imageCache.get(key);
+        if (cells == null) {
+            cells = generateSubmap(resource, submap);
+            imageCache.put(key, cells);
+        }
+
+        return cells[submap.cellX() * submap.rows() + submap.cellY()];
+    }
+
+    private static byte[][] generateSubmap(IResource resource, SubmapLocation submap) throws IOException {
+        BufferedImage source = read(resource, submap.source());
+        byte[][] cells = new byte[submap.columns() * submap.rows()][];
+
+        for (int x = 0; x < submap.columns(); x++) {
+            for (int y = 0; y < submap.rows(); y++) {
+                BufferedImage cropped = crop(source, submap.columns(), submap.rows(), x, y);
+                cells[x * submap.rows() + y] = encode(cropped);
+            }
+        }
+
+        return cells;
+    }
+
+    private static byte[] generateCell(IResource resource, SubmapLocation submap) throws IOException {
+        BufferedImage source = read(resource, submap.source());
+        BufferedImage cropped = crop(source, submap.columns(), submap.rows(), submap.cellX(), submap.cellY());
+        return encode(cropped);
+    }
+
+    private static BufferedImage read(IResource resource, ResourceLocation location) throws IOException {
+        BufferedImage image;
+        try (InputStream stream = resource.getInputStream()) {
+            image = ImageIO.read(stream);
+        }
+        if (image == null) {
+            throw new IOException("Unable to decode " + location);
+        }
+        return image;
+    }
+
+    private static byte[] encode(BufferedImage image) throws IOException {
+        ByteArrayOutputStream output = new ByteArrayOutputStream();
+        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("png");
+        ImageWriter writer = writers.next();
+
+        ImageWriteParam param = writer.getDefaultWriteParam();
+        if (param.canWriteCompressed()) {
+            // disable compression so we don't spend time on compressing and decompressing to just get the image bytes
+            param.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
+            param.setCompressionQuality(1.0f);
+        }
+
+        try (ImageOutputStream imageOutput = new MemoryCacheImageOutputStream(output)) {
+            writer.setOutput(imageOutput);
+            writer.write(null, new IIOImage(image, null, null), param);
+        } finally {
+            writer.dispose();
+        }
+
+        return output.toByteArray();
+    }
+
+    private static BufferedImage crop(BufferedImage source, int columns, int rows, int cellX, int cellY)
+        throws IOException {
+
+        int frameSize = source.getWidth();
+        int sourceHeight = source.getHeight();
+
+        if (frameSize < columns || frameSize < rows
+            || frameSize % columns != 0
+            || frameSize % rows != 0
+            || sourceHeight % frameSize != 0) {
+
+            throw new IOException(
+                "Texture size " + frameSize
+                    + "x"
+                    + sourceHeight
+                    + " cannot be split into submap "
+                    + columns
+                    + "x"
+                    + rows);
+        }
+
+        int cellWidth = frameSize / columns;
+        int cellHeight = frameSize / rows;
+        int frameCount = sourceHeight / frameSize;
+
+        BufferedImage result = new BufferedImage(cellWidth, cellHeight * frameCount, BufferedImage.TYPE_INT_ARGB);
+        int[] pixels = new int[cellWidth * cellHeight];
+
+        for (int frame = 0; frame < frameCount; frame++) {
+            int sourceX = cellX * cellWidth;
+            int sourceY = frame * frameSize + cellY * cellHeight;
+
+            source.getRGB(sourceX, sourceY, cellWidth, cellHeight, pixels, 0, cellWidth);
+            result.setRGB(0, frame * cellHeight, cellWidth, cellHeight, pixels, 0, cellWidth);
+        }
+
+        return result;
+    }
+
     @Override
     public boolean hasMetadata() {
-        return source.hasMetadata();
+        return resource.hasMetadata();
     }
 
     @Override
     public IMetadataSection getMetadata(String sectionName) {
-        IMetadataSection metadata = source.getMetadata(sectionName);
+        IMetadataSection metadata = resource.getMetadata(sectionName);
         if (!(metadata instanceof AnimationMetadataSection animation)) {
             return metadata;
         }
+        if (animation.getFrameWidth() < 0 && animation.getFrameHeight() < 0) {
+            return metadata;
+        }
 
         List<AnimationFrame> frames = new ArrayList<>(animation.getFrameCount());
         for (int i = 0; i < animation.getFrameCount(); i++) {
@@ -172,8 +254,8 @@ public final class TextureSubmapResource implements IResource {
             }
         }
 
-        int frameWidth = scaleMetadataDimension(animation.getFrameWidth(), columns);
-        int frameHeight = scaleMetadataDimension(animation.getFrameHeight(), rows);
+        int frameWidth = scaleMetadataDimension(animation.getFrameWidth(), submap.columns());
+        int frameHeight = scaleMetadataDimension(animation.getFrameHeight(), submap.rows());
         return new AnimationMetadataSection(frames, frameWidth, frameHeight, animation.getFrameTime());
     }