diff --git a/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleReader.java b/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleReader.java index 6528bef..72e9c90 100644 --- a/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleReader.java +++ b/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)) { - var result = reassembler.writeReassembled(allocator, RegionFile.CompressType.GZIP); + try (var reassembler = new RegionParser(entryData, allocator)) { + var result = reassembler.writeReassembled(RegionFile.CompressType.GZIP); try (var fc = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { result.readBytes(fc, 0, result.readableBytes()); } finally { diff --git a/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleWriter.java b/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleWriter.java index f864b61..ecc5290 100644 --- a/bundler/src/main/java/io/ib67/sfcraft/bundler/BundleWriter.java +++ b/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.netty.buffer.AdaptiveByteBufAllocator; +import io.ib67.sfcraft.bundler.region.RegionParser; 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 SORT = BundleWriter::preferRegion; + protected static final ThreadLocal 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> 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 result = rf.writeReassembled(allocator, RegionFile.CompressType.NONE); + var rf = new RegionParser(path, allocator)) { + 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); } diff --git a/bundler/src/main/java/io/ib67/sfcraft/bundler/Main.java b/bundler/src/main/java/io/ib67/sfcraft/bundler/Main.java index 2bff881..3654ee5 100644 --- a/bundler/src/main/java/io/ib67/sfcraft/bundler/Main.java +++ b/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", - "Should we reassemble NBT files for maximum compression rate", true); + var allowGunzip = opts.bool("reasm", "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()); + } + } } diff --git a/bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionFileReassembler.java b/bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionParser.java similarity index 73% rename from bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionFileReassembler.java rename to bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionParser.java index 1530d49..b3dd8a4 100644 --- a/bundler/src/main/java/io/ib67/sfcraft/bundler/region/RegionFileReassembler.java +++ b/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 { - if (Files.size(file) < 8192) throw new IOException("The region file " + file + " has been truncated"); - source = Unpooled.wrappedBuffer(Files.readAllBytes(file)); + public RegionParser(Path file, ByteBufAllocator allocator) throws IOException { + this.allocator = allocator; + 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)) { diff --git a/bundler/src/main/java/io/ib67/sfcraft/bundler/trainer/ZstdRegionTrainer.java b/bundler/src/main/java/io/ib67/sfcraft/bundler/trainer/ZstdRegionTrainer.java new file mode 100644 index 0000000..4415f1b --- /dev/null +++ b/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(); + } +} diff --git a/bundler/src/test/java/io/ib67/bundler/TestRegionReassembler.java b/bundler/src/test/java/io/ib67/bundler/TestRegionReassembler.java index f775946..dfde0c6 100644 --- a/bundler/src/test/java/io/ib67/bundler/TestRegionReassembler.java +++ b/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))) { - var result = reasm.writeReassembled(ByteBufAllocator.DEFAULT, RegionFile.CompressType.NONE); + try(var reasm = new RegionParser(Unpooled.wrappedBuffer(bytes), ByteBufAllocator.DEFAULT)) { + var result = reasm.writeReassembled(RegionFile.CompressType.NONE); } } } diff --git a/bundler/src/test/resources/region/r.9.-5.mca b/bundler/src/test/resources/region/r.9.-5.mca new file mode 100644 index 0000000..850e308 Binary files /dev/null and b/bundler/src/test/resources/region/r.9.-5.mca differ