Implement Zlib decompression with inflater caching

Author: Glavo

src/main/java/org/glavo/nbt/internal/input/InputContext.java

@@ -22,10 +22,9 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.IdentityHashMap;
 import java.util.Map;
-import java.util.Objects;
-import java.util.function.Supplier;
 
 public final class InputContext implements Closeable {
     // Used for reading UTF-8 strings
@@ -37,7 +36,7 @@ public final class InputContext implements Closeable {
     public final InputSource source;
     public final MinecraftEdition edition;
 
-    public final DataReader rawReader;
+    public final RawDataReader rawReader;
 
     public final StringCache stringCache = DEFAULT_CACHE;
     @Nullable StringBuilder charsBuffer;
@@ -69,26 +68,46 @@ public void releaseDecompressBuffer(InputBuffer buffer) {
     }
 
     @Override
+    @SuppressWarnings({"rawtypes", "unchecked"})
     public void close() throws IOException {
         source.close();
+        if (cacheMap != null) {
+            for (var entry : new ArrayList<>(cacheMap.entrySet())) {
+                ((CacheKey) entry.getKey()).close(entry.getValue());
+            }
+
+            cacheMap.clear();
+        }
     }
 
     @SuppressWarnings("unchecked")
-    public static final class CacheKey<T> {
-
-        public @Nullable T get(InputContext context) {
-            return context.cacheMap != null ? (T) context.cacheMap.remove(this) : null;
-        }
+    public static abstract class CacheKey<T> {
+
+        public T get(InputContext context) {
+            if (context.cacheMap != null) {
+                T value = (T) context.cacheMap.remove(this);
+                if (value != null) {
+                    return value;
+                }
+            }
 
-        public T getOrCreate(InputContext context, Supplier<? extends T> supplier) {
-            return Objects.requireNonNullElseGet(get(context), supplier);
+            return create(context);
         }
 
-        public void put(InputContext context, T value) {
+        public void release(InputContext context, T value) {
             if (context.cacheMap == null) {
                 context.cacheMap = new IdentityHashMap<>();
             }
-            context.cacheMap.put(this, value);
+
+            T oldValue = (T) context.cacheMap.put(this, value);
+            if (oldValue != null) {
+                close(oldValue);
+            }
+        }
+
+        protected abstract T create(InputContext context);
+
+        public void close(T value) {
         }
     }
 }

src/main/java/org/glavo/nbt/internal/input/RawDataReader.java

@@ -18,7 +18,7 @@
 import java.io.EOFException;
 import java.io.IOException;
 
-final class RawDataReader extends DataReader {
+public final class RawDataReader extends DataReader {
 
     /// Used for the default reader;
     RawDataReader(InputContext context, InputBuffer buffer) {

src/main/java/org/glavo/nbt/internal/input/ZlibDataReader.java

@@ -18,15 +18,27 @@
 import java.io.EOFException;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.zip.DataFormatException;
 import java.util.zip.Inflater;
 
 public final class ZlibDataReader extends DataReader {
-    static final InputContext.CacheKey<Inflater> INFLATER_CACHE_KEY = new InputContext.CacheKey<>();
+    static final InputContext.CacheKey<Inflater> INFLATER_CACHE_KEY = new InputContext.CacheKey<>() {
+
+        @Override
+        protected Inflater create(InputContext context) {
+            return new Inflater();
+        }
+
+        @Override
+        public void close(Inflater value) {
+            value.end();
+        }
+    };
     private final Inflater inflater;
 
     ZlibDataReader(InputContext context, long compressedSize) {
         super(context, context.getDecompressBuffer());
-        this.inflater = INFLATER_CACHE_KEY.getOrCreate(context, Inflater::new);
+        this.inflater = INFLATER_CACHE_KEY.get(context);
         this.remainingInput = compressedSize;
     }
 
@@ -36,23 +48,46 @@ public void ensureBufferRemaining(int required) throws IOException {
             return;
         }
 
+        if (inflater.finished() || inflater.needsDictionary()) {
+            throw new EOFException("Inflater finished or needs dictionary");
+        }
+
         buffer.ensureCapacity(required);
-        ByteBuffer byteBuffer = this.buffer.getByteBuffer();
-
-        // TODO
-
-//        do {
-//            if (inflater.finished() || inflater.needsDictionary()) {
-//                throw new EOFException();
-//            }
-//
-//            if (inflater.needsInput()) {
-//                if (context.rawReader.buffer.remaining() == 0) {
-//                    context.rawReader.ensureBufferRemaining(1);
-//                }
-//
-//                inflater.setInput(context.rawReader.buffer.getByteBuffer());
-//            }
-//        } while ((n = inflater.inflate(b, off, len)) == 0);
+        ByteBuffer output = this.buffer.getByteBuffer();
+        output.compact();
+
+        do {
+            if (inflater.finished() || inflater.needsDictionary()) {
+                throw new EOFException();
+            }
+
+            if (inflater.needsInput()) {
+                if (context.rawReader.buffer.remaining() == 0) {
+                    context.rawReader.ensureBufferRemaining(1);
+                }
+                inflater.setInput(context.rawReader.buffer.getByteBuffer());
+            }
+
+            try {
+                inflater.inflate(output);
+            } catch (DataFormatException exception) {
+                throw new IOException(exception);
+            }
+
+        } while (output.position() < required);
+
+        inflater.setInput(EMPTY_BYTE_ARRAY);
+        output.flip();
+    }
+
+    @Override
+    public void close() throws IOException {
+        inflater.end();
+        context.releaseDecompressBuffer(buffer);
+
+        inflater.reset();
+        INFLATER_CACHE_KEY.release(context, inflater);
     }
+
+    private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
 }