refactor: optimize region file handling and add dictionary training support

Key changes:
- Renamed RegionFileReassembler to RegionParser for better semantics
- Refactored RegionParser to use ByteBufAllocator injection and improved memory management
- Added thread-local buffer reuse in BundleWriter to reduce allocations
- Implemented Zstd dictionary training functionality for region files
- Added region file statistics collection for training optimization
- Improved error handling and resource management

Detailed changes:

BundleReader/BundleWriter:
- Updated to use RegionParser instead of RegionFileReassembler
- Modified RegionParser constructor to accept ByteBufAllocator
- Changed writeReassembled() method signature to use injected allocator
- Added thread-local buffer reuse to reduce memory allocation overhead
- Improved logging with better formatted output

Main class:
- Added dictionary training support with --train, --train-sample-size, and --train-dict-size options
- Implemented region file statistics collection with --stat option
- Added proper validation for training parameters
- Enhanced help text and error messages

RegionParser:
- Improved file reading using FileChannel for better performance
- Added size validation to prevent processing overly large files
- Implemented readChunkDecompress() method for training purposes
- Better resource management with proper buffer allocation/release

New features:
- ZstdRegionTrainer class for training zstd dictionaries on region files
- Statistics collection for chunk size distribution analysis
- Dictionary training workflow integrated into main application

Performance improvements:
- Reduced memory allocations through buffer reuse
- Better resource management with proper close() implementations
- Improved error handling for corrupted region files

Testing:
- Updated test cases to reflect RegionParser rename

bundler/src/main/java/io/ib67/sfcraft/bundler/BundleReader.java

@@ -2,7 +2,7 @@ package io.ib67.sfcraft.bundler;
 
 import com.github.luben.zstd.ZstdInputStreamNoFinalizer;
 import io.ib67.sfcraft.bundler.region.RegionFile;
-import io.ib67.sfcraft.bundler.region.RegionFileReassembler;
+import io.ib67.sfcraft.bundler.region.RegionParser;
 import io.netty.buffer.ByteBufAllocator;
 import lombok.Builder;
 
@@ -61,8 +61,8 @@ public class BundleReader {
                         toRead -= read;
                         entryData.writeBytes(buffer, 0, read);
                     }
-                    try (var reassembler = new RegionFileReassembler(entryData)) {
+                    try (var reassembler = new RegionParser(entryData, allocator)) {
-                        var result = reassembler.writeReassembled(allocator, RegionFile.CompressType.GZIP);
+                        var result = reassembler.writeReassembled(RegionFile.CompressType.GZIP);
                         try (var fc = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE)) {
                             result.readBytes(fc, 0, result.readableBytes());
                         } finally {

bundler/src/main/java/io/ib67/sfcraft/bundler/BundleWriter.java

@@ -4,8 +4,7 @@ import com.github.luben.zstd.ZstdDictCompress;
 import com.github.luben.zstd.ZstdOutputStreamNoFinalizer;
 import io.ib67.kiwi.routine.Uni;
 import io.ib67.sfcraft.bundler.region.RegionFile;
-import io.ib67.sfcraft.bundler.region.RegionFileReassembler;
+import io.ib67.sfcraft.bundler.region.RegionParser;
-import io.netty.buffer.AdaptiveByteBufAllocator;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufOutputStream;
 
@@ -21,12 +20,12 @@ import java.util.function.UnaryOperator;
 
 public class BundleWriter implements Closeable {
     public static Comparator<Path> SORT = BundleWriter::preferRegion;
+    protected static final ThreadLocal<byte[]> buffer = ThreadLocal.withInitial(() -> new byte[4096]);
     protected final ByteBufAllocator allocator;
     protected final boolean allowReassemble;
     protected final int compressionLevel;
     protected final Path relativeRoot;
     protected final EntryOutputStream outputStream;
-    protected final byte[] buffer;
 
     @SneakyThrows
     @Builder
@@ -48,7 +47,6 @@ public class BundleWriter implements Closeable {
         }
         if (maxWorkers > 1) zstdOut.setWorkers(maxWorkers);
         this.outputStream = new EntryOutputStream(zstdOut);
-        this.buffer = new byte[4096];
     }
 
     public static CompletableFuture<List<Path>> createBundleParallelized(
@@ -153,16 +151,16 @@ public class BundleWriter implements Closeable {
     @SneakyThrows
     private void writeReassembleInMem(Path path) {
         if (!"region".equals(path.getParent().toString()) && !path.toString().endsWith(".mca")) {
-            System.out.println("Mismatch "+path+", parent: "+path.getParent());
+            System.out.println("Mismatch " + path + ", parent: " + path.getParent());
             writePlain(path);
             return;
         }
-        System.out.println("Reassembling "+path);
+        System.out.println("Reassembling " + path);
         var uncompressed = allocator.buffer();
         var rPath = relativeRoot.toAbsolutePath().relativize(path.toAbsolutePath());
         try (var out = new EntryOutputStream(new ByteBufOutputStream(uncompressed));
-             var rf = new RegionFileReassembler(path)) {
+             var rf = new RegionParser(path, allocator)) {
-            var result = rf.writeReassembled(allocator, RegionFile.CompressType.NONE);
+            var result = rf.writeReassembled(RegionFile.CompressType.NONE);
             var length = result.readableBytes();
             out.writeEntryHeader(rPath.toString(), BundleEntry.ATTR_REASSEMBLE, length, System.currentTimeMillis());
             result.readBytes(out, length);
@@ -178,11 +176,11 @@ public class BundleWriter implements Closeable {
 
     private void writePlain(Path path) throws IOException {
         var rPath = relativeRoot.toAbsolutePath().relativize(path.toAbsolutePath());
-        var buffer = this.buffer;
         var outputStream = this.outputStream;
         outputStream.writeEntryHeader(rPath.toString(), (short) 0, Files.size(path), 0);
         var read = 0;
         try (var fs = Files.newInputStream(path)) {
+            var buffer = BundleWriter.buffer.get();
             while ((read = fs.read(buffer)) > 0) {
                 outputStream.write(buffer, 0, read);
             }

bundler/src/main/java/io/ib67/sfcraft/bundler/Main.java

@@ -1,14 +1,23 @@
 package io.ib67.sfcraft.bundler;
 
 import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdDictTrainer;
 import com.github.luben.zstd.ZstdOutputStreamNoFinalizer;
 import io.ib67.kiwi.ArgOpts;
+import io.ib67.sfcraft.bundler.region.RegionParser;
+import io.ib67.sfcraft.bundler.trainer.ZstdRegionTrainer;
+import io.netty.buffer.ByteBufAllocator;
+import it.unimi.dsi.fastutil.ints.Int2IntMap;
+import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
 import lombok.SneakyThrows;
 
 import java.io.BufferedOutputStream;
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Comparator;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ForkJoinPool;
 import java.util.function.UnaryOperator;
@@ -21,18 +30,41 @@ public class Main {
         var in = opts.string("in", "Input. Can be directory or jpack bundle", null);
         var out = opts.string("out", "Output. Can be jpack bundle or directory", null);
         var help = opts.bool("help", false);
-        var allowGunzip = opts.bool("reasm",
+        var allowGunzip = opts.bool("reasm", "Should we reassemble NBT files for maximum compression rate", true);
-                "Should we reassemble NBT files for maximum compression rate", true);
+
+        // useless parallel
         var parallel = opts.bool("parallel", "Use parallelized bundler. This conflicts with compress options", false);
         var parallelThreshold = opts.integer("threshold", "The threshold of parallel bundling", 512);
+
+        // zstd options
         var parallelCompression = opts.integer("parallel-compression", "The worker threads used for parallel compression. Do not confuse this with --parallel, which focus on bundling.\n" +
                 "This option will be ignored when --parallel is present.", 1);
         var compressionLevel = opts.integer("level", "Zstd compression level", Zstd.defaultCompressionLevel());
+
+        // dictionary
         var dictionaryPath = opts.string("dict", "Zstd dictionary file", null);
+        var train = opts.string("train", "Train zstd dictionary based on your map files. The value should be a directory which contains a lot of .mca files", null);
+        var trainSampleSize = opts.integer("train-sample-size", "Train sample size (in total)", -1);
+        var trainDictSize = opts.integer("train-dict-size", "Train dict size", -1);
+        var stat = opts.bool("stat", "Count region and chunk sizes. You should set --train for this option to find region files", false);
         if (help) {
             opts.printHelp(System.out);
             return;
         }
+        if (stat) {
+            Objects.requireNonNull(train, "You should also set --train for this to work.");
+            calculateStat(Path.of(train));
+            return;
+        }
+        if (train != null) {
+            if (trainDictSize < 0 || trainSampleSize < 0) {
+                System.err.println("You should also set train-sample-size and train-dict-size.");
+                System.err.println("TIP: use --stat to gather some useful information.");
+                return;
+            }
+            train(train, trainSampleSize, trainDictSize);
+            return;
+        }
         var dict = (dictionaryPath != null) ? Files.readAllBytes(Path.of(dictionaryPath)) : null;
         if (in == null || out == null) {
             System.out.println("Both in and out should be provided");
@@ -48,6 +80,7 @@ public class Main {
             System.out.println("The output is already exists. Please move it to other places");
             return;
         }
+
         if (in.endsWith(".jbp.zst")) {
             BundleReader.builder().build().extract(inPath, outPath);
         } else if (out.endsWith(".jbp.zst")) {
@@ -82,4 +115,64 @@ public class Main {
             System.out.println("Invalid arguments. You should provide a pair of jpack and directory to act as in/output.");
         }
     }
+
+    private static void calculateStat(Path inPath) throws IOException {
+        try (var stream = Files.walk(inPath)) {
+            var regions = stream
+                    .filter(Files::isRegularFile)
+                    .filter(it -> it.toString().endsWith("mca"))
+                    .toList();
+            var intFreqMap = new Int2IntOpenHashMap();
+            for (Path region : regions) {
+                try (var parser = new RegionParser(region, ByteBufAllocator.DEFAULT)) {
+                    for (long packedEntry : parser.readAllValidChunks()) {
+                        var entry = RegionParser.readEntry(packedEntry);
+                        var occupiedSectors = entry & 0xFF;
+                        intFreqMap.merge(occupiedSectors, 1, Integer::sum);
+                    }
+                } catch (IOException e) {
+                    System.err.println("Skipping truncated region " + region + ": " + e.getMessage());
+                }
+            }
+            var entrySet = intFreqMap.int2IntEntrySet().stream()
+                    .sorted(Comparator.comparingInt(Int2IntMap.Entry::getIntValue).reversed())
+                    .toList();
+            var overallAppears = entrySet
+                    .stream().mapToInt(Int2IntMap.Entry::getIntValue)
+                    .sum();
+            var all = entrySet.stream().mapToLong(it -> (long) it.getIntKey() * 4096 * it.getIntValue()).sum();
+            System.out.println("All samples size in total: " + all+"B");
+            System.out.println("Region chunk size stat (first 32):");
+            for (int i = 0; i < Math.min(32, entrySet.size()); i++) {
+                var entry = entrySet.get(i);
+                var freq = String.format("%.2f", ((double) entry.getIntValue() / (double) overallAppears));
+                System.out.println(" " + i + ". Sectors: " + entry.getIntKey() + "(" + entry.getIntKey() * 4096 + "KiB)" +
+                        ", frequency: " + freq + ", count: " + entry.getIntValue());
+            }
+        }
+    }
+
+    private static void train(String train, int sampleSize, int dictSize) throws IOException {
+        var path = Path.of(train);
+        try (var stream = Files.walk(path);
+             var trainer = new ZstdRegionTrainer(
+                     new ZstdDictTrainer(sampleSize, dictSize),
+                     ByteBufAllocator.DEFAULT)
+        ) {
+            var regions = stream
+                    .filter(Files::isRegularFile)
+                    .filter(it -> it.toString().endsWith("mca"))
+                    .toList();
+
+            for (var region : regions) {
+                try {
+                    trainer.feed(region);
+                } catch (IOException e) {
+                    System.err.println("Failed to train on " + region + ": " + e.getMessage() + ". Maybe sample is full..");
+                    break;
+                }
+            }
+            Files.write(Path.of("train_"+System.currentTimeMillis()+".train"), trainer.train());
+        }
+    }
 }

bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionParser.java

@@ -6,23 +6,32 @@ import it.unimi.dsi.fastutil.longs.LongList;
 import net.jpountz.lz4.LZ4BlockInputStream;
 
 import java.io.*;
-import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
 import java.util.zip.DeflaterInputStream;
 import java.util.zip.GZIPInputStream;
 
-public class RegionFileReassembler implements Closeable {
+public class RegionParser implements Closeable {
     protected final ByteBuf source;
     protected final byte[] sectorBuffer = new byte[4096];
+    protected final ByteBufAllocator allocator;
 
-    public RegionFileReassembler(Path file) throws IOException {
+    public RegionParser(Path file, ByteBufAllocator allocator) throws IOException {
-        if (Files.size(file) < 8192) throw new IOException("The region file " + file + " has been truncated");
+        this.allocator = allocator;
-        source = Unpooled.wrappedBuffer(Files.readAllBytes(file));
+        var size = (int) Files.size(file);
+        if (size < 8192) throw new IOException("The region file " + file + " has been truncated");
+        if(size > 8388608 * 2) throw new IOException("The region file " + file + " is too large"); //todo bug
+        source = allocator.buffer(size);
+        try (var fc = FileChannel.open(file, StandardOpenOption.READ)) {
+            while (source.writeBytes(fc, 0L, (int) size) > 0) ;
+        }
     }
 
-    public RegionFileReassembler(ByteBuf source) {
+    public RegionParser(ByteBuf source, ByteBufAllocator allocator) {
         this.source = source;
+        this.allocator = allocator;
         source.retain();
     }
 
@@ -53,12 +62,23 @@ public class RegionFileReassembler implements Closeable {
         if (compressType < 0) {
             // mcc todo
         }
-        buf.writeBytes(source, sectorLen-1);
+        buf.writeBytes(source, sectorLen - 1);
         return compressType;
     }
 
-    public ByteBuf writeReassembled(ByteBufAllocator alloc, RegionFile.CompressType newCompressType) throws IOException {
+    public void readChunkDecompress(int entry, ByteBuf buf) throws IOException {
+        var compressed = allocator.buffer();
+        var compressType = readChunk(entry, compressed);
+        try (var decompressor = getDecompressorByType(compressType, compressed)) {
+            while (buf.writeBytes(decompressor, 4096) != -1) ;
+        } finally {
+            compressed.release();
+        }
+    }
+
+    public ByteBuf writeReassembled(RegionFile.CompressType newCompressType) throws IOException {
         var chunks = readAllValidChunks();
+        var alloc = allocator;
         var compressed = alloc.buffer();
         var uncompressed = alloc.buffer();
         try (var regionFile = new RegionFile(alloc)) {

bundler/src/main/java/io/ib67/sfcraft/bundler/trainer/ZstdRegionTrainer.java

@@ -0,0 +1,50 @@
+package io.ib67.sfcraft.bundler.trainer;
+
+import com.github.luben.zstd.ZstdDictTrainer;
+import io.ib67.sfcraft.bundler.region.RegionParser;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.Unpooled;
+import lombok.SneakyThrows;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Objects;
+
+public class ZstdRegionTrainer implements Closeable {
+    protected final ZstdDictTrainer trainer;
+    protected final ByteBufAllocator allocator;
+    protected final ByteBuf decompressBuffer;
+
+    public ZstdRegionTrainer(ZstdDictTrainer trainer, ByteBufAllocator allocator) {
+        this.trainer = Objects.requireNonNull(trainer);
+        this.allocator = allocator;
+        this.decompressBuffer = Unpooled.buffer();
+    }
+
+    public void feed(Path pathRegion) throws IOException {
+        if(Files.size(pathRegion) < 8192) return;;
+        try (var regionParser = new RegionParser(pathRegion, allocator)) {
+            var validChunks = regionParser.readAllValidChunks();
+            var buffer = decompressBuffer;
+            for (var idx = 0; idx < validChunks.size(); idx++) {
+                regionParser.readChunkDecompress(RegionParser.readEntry(validChunks.getLong(idx)), buffer);
+                if(!trainer.addSample(buffer.array())){
+                    throw new IOException("Cannot add sample idx "+idx+" from "+pathRegion+", size: "+buffer.readableBytes());
+                }
+                buffer.clear();
+            }
+        }
+    }
+
+    public byte[] train() {
+        return trainer.trainSamples();
+    }
+
+    @Override
+    public void close() throws IOException {
+        decompressBuffer.release();
+    }
+}

bundler/src/test/java/io/ib67/bundler/TestRegionReassembler.java

@@ -1,7 +1,7 @@
 package io.ib67.bundler;
 
 import io.ib67.sfcraft.bundler.region.RegionFile;
-import io.ib67.sfcraft.bundler.region.RegionFileReassembler;
+import io.ib67.sfcraft.bundler.region.RegionParser;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.Unpooled;
 import org.junit.jupiter.api.Test;
@@ -15,8 +15,8 @@ public class TestRegionReassembler {
         try(var in = TestRegionReassembler.class.getClassLoader().getResourceAsStream("region/r.9.-5.mca")){
             bytes = in.readAllBytes();
         }
-        try(var reasm = new RegionFileReassembler(Unpooled.wrappedBuffer(bytes))) {
+        try(var reasm = new RegionParser(Unpooled.wrappedBuffer(bytes), ByteBufAllocator.DEFAULT)) {
-            var result = reasm.writeReassembled(ByteBufAllocator.DEFAULT, RegionFile.CompressType.NONE);
+            var result = reasm.writeReassembled(RegionFile.CompressType.NONE);
         }
     }
 }