Refactor NBTReader to DataReader and add InputContext

Author: Glavo

…/glavo/nbt/internal/input/NBTReader.java …glavo/nbt/internal/input/DataReader.javasrc/main/java/org/glavo/nbt/internal/input/NBTReader.java renamed to src/main/java/org/glavo/nbt/internal/input/DataReader.java src/main/java/org/glavo/nbt/internal/input/NBTReader.java renamed to src/main/java/org/glavo/nbt/internal/input/DataReader.java

@@ -16,59 +16,34 @@
 package org.glavo.nbt.internal.input;
 
 import org.glavo.nbt.MinecraftEdition;
-import org.glavo.nbt.internal.IOUtils;
-import org.glavo.nbt.internal.StringCache;
-import org.jetbrains.annotations.Nullable;
 
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 
-public final class NBTReader implements Closeable {
-
-    // Used for reading UTF-8 strings
-    private static final StringCache CACHE = new StringCache(
-            "data", "Data", "DataVersion"
-            // TODO: More tag names
-    );
-
-    private final InputSource source;
-    private final InputBuffer buffer;
-    private final MinecraftEdition edition;
-    private final long initialPosition;
-
-    /// Used for reading UTF-8 strings
-    private @Nullable StringBuilder charsBuffer;
-
-    public NBTReader(InputSource source, MinecraftEdition edition) throws IOException {
-        this.source = source;
-        this.buffer = InputBuffer.allocate(IOUtils.DEFAULT_BUFFER_SIZE, source.supportDirectBuffer(), edition.byteOrder());
-        this.edition = edition;
-
-        try {
-            this.initialPosition = source.position();
-        } catch (Throwable e) {
-            try {
-                source.close();
-            } catch (IOException e2) {
-                e.addSuppressed(e2);
-            }
-            throw e;
-        }
-    }
+public abstract class DataReader implements Closeable {
+    final InputContext context;
+    final InputBuffer buffer;
+    long remainingInput = -1L;
 
-    public MinecraftEdition getEdition() {
-        return edition;
+    protected DataReader(InputContext context, InputBuffer buffer) {
+        this.context = context;
+        this.buffer = buffer;
     }
 
+    public abstract void ensureBufferRemaining(int required) throws IOException;
+
     @Override
     public void close() throws IOException {
-        source.close();
-    }
-
-    private void fillBuffer(int required) throws IOException {
-        source.fillBuffer(buffer, required);
+        if (remainingInput > 0) {
+            int bufferRemaining = context.rawReader.buffer.remaining();
+            if (bufferRemaining > 0) {
+                int bytesDrop = (int) Math.min(bufferRemaining, remainingInput);
+                context.rawReader.buffer.drop(bytesDrop);
+            }
+            this.remainingInput -= bufferRemaining;
+        }
     }
 
     public byte[] readByteArray() throws IOException {
@@ -77,7 +52,7 @@ public byte[] readByteArray() throws IOException {
             throw new IOException("Array length too large");
         }
 
-        fillBuffer(len);
+        ensureBufferRemaining(len);
         return buffer.getByteArray(len);
     }
 
@@ -87,7 +62,7 @@ public int[] readIntArray() throws IOException {
             throw new IOException("Array length too large");
         }
 
-        fillBuffer(len * Integer.BYTES);
+        ensureBufferRemaining(len * Integer.BYTES);
         return buffer.getIntArray(len);
     }
 
@@ -97,13 +72,13 @@ public long[] readLongArray() throws IOException {
             throw new IOException("Array length too large");
         }
 
-        fillBuffer(len * Long.BYTES);
+        ensureBufferRemaining(len * Long.BYTES);
         return buffer.getLongArray(len);
     }
 
     /// Read a byte from the input stream.
     public byte readByte() throws IOException {
-        fillBuffer(Byte.BYTES);
+        ensureBufferRemaining(Byte.BYTES);
         return buffer.getByte();
     }
 
@@ -114,7 +89,7 @@ public int readUnsignedByte() throws IOException {
 
     /// Read a short from the input stream.
     public short readShort() throws IOException {
-        fillBuffer(Short.BYTES);
+        ensureBufferRemaining(Short.BYTES);
         return buffer.getShort();
     }
 
@@ -125,7 +100,7 @@ public int readUnsignedShort() throws IOException {
 
     /// Read an int from the input stream.
     public int readInt() throws IOException {
-        fillBuffer(Integer.BYTES);
+        ensureBufferRemaining(Integer.BYTES);
         return buffer.getInt();
     }
 
@@ -136,24 +111,24 @@ public long readUnsignedInt() throws IOException {
 
     /// Read a long from the input stream.
     public long readLong() throws IOException {
-        fillBuffer(Long.BYTES);
+        ensureBufferRemaining(Long.BYTES);
         return buffer.getLong();
     }
 
     /// Read a float from the input stream.
     public float readFloat() throws IOException {
-        fillBuffer(Float.BYTES);
+        ensureBufferRemaining(Float.BYTES);
         return buffer.getFloat();
     }
 
     /// Read a double from the input stream.
     public double readDouble() throws IOException {
-        fillBuffer(Double.BYTES);
+        ensureBufferRemaining(Double.BYTES);
         return buffer.getDouble();
     }
 
     private String getUTF8(ByteBuffer buffer, int offset, int length) {
-        String cached = CACHE.get(buffer, offset, length);
+        String cached = context.stringCache.get(buffer, offset, length);
         if (cached != null) {
             return cached;
         }
@@ -178,7 +153,7 @@ public String readString() throws IOException {
             return "";
         }
 
-        fillBuffer(len);
+        ensureBufferRemaining(len);
 
         ByteBuffer bytes = buffer.bytesBuffer();
         int offset = bytes.position();
@@ -187,7 +162,7 @@ public String readString() throws IOException {
         bytes.position(limit);
 
         // For Minecraft Bedrock Edition, the string is encoded in standard UTF-8
-        if (edition == MinecraftEdition.BEDROCK_EDITION) {
+        if (context.edition == MinecraftEdition.BEDROCK_EDITION) {
             return getUTF8(bytes, offset, len);
         }
 
@@ -207,10 +182,13 @@ public String readString() throws IOException {
         }
 
         // Slow path
-        if (charsBuffer != null) {
-            charsBuffer.setLength(0);
+        StringBuilder charsBuffer;
+        if (context.charsBuffer != null) {
+            charsBuffer = context.charsBuffer;
+            context.charsBuffer.setLength(0);
         } else {
             charsBuffer = new StringBuilder(len);
+            context.charsBuffer = charsBuffer;
         }
 
         int c, char2, char3;
@@ -262,5 +240,4 @@ public String readString() throws IOException {
         }
         return charsBuffer.toString();
     }
-
 }

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

@@ -15,6 +15,7 @@
  */
 package org.glavo.nbt.internal.input;
 
+import java.nio.BufferOverflowException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -43,6 +44,26 @@ public ByteOrder order() {
         return bytesBuffer.order();
     }
 
+    public void drop() {
+        bytesBuffer.position(0);
+        bytesBuffer.limit(0);
+    }
+
+    public void drop(int n) {
+        int remaining = this.remaining();
+        if (n < remaining) {
+            bytesBuffer.position(bytesBuffer.position() + n);
+        } else if (n == remaining) {
+            drop();
+        } else {
+            throw new BufferOverflowException();
+        }
+    }
+
+    public int remaining() {
+        return bytesBuffer.remaining();
+    }
+
     public byte getByte() {
         return bytesBuffer.get();
     }

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2026 Glavo
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.glavo.nbt.internal.input;
+
+import org.glavo.nbt.MinecraftEdition;
+import org.glavo.nbt.internal.IOUtils;
+import org.glavo.nbt.internal.StringCache;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+public final class InputContext implements Closeable {
+    // Used for reading UTF-8 strings
+    private static final StringCache DEFAULT_CACHE = new StringCache(
+            "data", "Data", "DataVersion"
+            // TODO: More tag names
+    );
+
+    public final InputSource source;
+    public final MinecraftEdition edition;
+
+    @Nullable StringBuilder charsBuffer;
+
+    public final DataReader rawReader;
+
+    public final StringCache stringCache = DEFAULT_CACHE;
+
+    public InputContext(InputSource source, MinecraftEdition edition) {
+        this.source = source;
+        this.edition = edition;
+        this.rawReader = new RawDataReader(
+                this,
+                InputBuffer.allocate(IOUtils.DEFAULT_BUFFER_SIZE, source.supportDirectBuffer(), edition.byteOrder()));
+    }
+
+    @Override
+    public void close() throws IOException {
+        source.close();
+    }
+}

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2026 Glavo
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.glavo.nbt.internal.input;
+
+import java.io.EOFException;
+import java.io.IOException;
+
+final class RawDataReader extends DataReader {
+
+    /// Used for the default reader;
+    RawDataReader(InputContext context, InputBuffer buffer) {
+        super(context, buffer);
+    }
+
+    public RawDataReader(InputContext context, long limit) {
+        super(context, context.rawReader.buffer);
+        this.remainingInput = limit;
+    }
+
+    @Override
+    public void ensureBufferRemaining(int required) throws IOException {
+        if (remainingInput >= 0 && remainingInput < required) {
+            throw new EOFException("Not enough data to read, required: " + required + ", remaining: " + remainingInput);
+        }
+
+        context.source.fillBuffer(buffer, required);
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2026 Glavo
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.glavo.nbt.internal.input;
+
+import java.io.IOException;
+
+public final class ZlibDataReader extends DataReader {
+    ZlibDataReader(InputContext context, InputBuffer buffer) {
+        super(context, buffer);
+    }
+
+    @Override
+    public void ensureBufferRemaining(int required) throws IOException {
+
+    }
+}

src/main/java/org/glavo/nbt/tag/ByteArrayTag.java

@@ -15,7 +15,7 @@
  */
 package org.glavo.nbt.tag;
 
-import org.glavo.nbt.internal.input.NBTReader;
+import org.glavo.nbt.internal.input.DataReader;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -103,7 +103,7 @@ public Byte next() {
     }
 
     @Override
-    protected void readContent(NBTReader reader) throws IOException {
+    protected void readContent(DataReader reader) throws IOException {
         value = reader.readByteArray();
     }
 

src/main/java/org/glavo/nbt/tag/ByteTag.java

@@ -15,7 +15,7 @@
  */
 package org.glavo.nbt.tag;
 
-import org.glavo.nbt.internal.input.NBTReader;
+import org.glavo.nbt.internal.input.DataReader;
 
 import java.io.IOException;
 
@@ -75,7 +75,7 @@ public void setBoolean(boolean value) {
     }
 
     @Override
-    protected void readContent(NBTReader reader) throws IOException {
+    protected void readContent(DataReader reader) throws IOException {
         set(reader.readByte());
     }