From 75f1fc219ae6a49df1bd06cb4c37f26ef7d2a28a Mon Sep 17 00:00:00 2001
From: tomerr90 <37086811+tomerr90@users.noreply.github.com>
Date: Fri, 24 Jan 2025 22:50:15 +0000
Subject: [PATCH] Add support for compression level in ZstdDictTrainer

---
 src/main/java/com/github/luben/zstd/ZstdDictTrainer.java | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java b/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java
index 61757798..07557cbb 100644
--- a/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java
+++ b/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java
@@ -10,12 +10,18 @@ public class ZstdDictTrainer {
     private final List<Integer> sampleSizes;
     private final int dictSize;
     private long filledSize;
+    private int level;
 
     public ZstdDictTrainer(int sampleSize, int dictSize) {
+        this(sampleSize, dictSize, Zstd.defaultCompressionLevel());
+    }
+    
+    public ZstdDictTrainer(int sampleSize, int dictSize, int level) {
         trainingSamples = ByteBuffer.allocateDirect(sampleSize);
         sampleSizes =  new ArrayList<Integer>();
         this.allocatedSize = sampleSize;
         this.dictSize = dictSize;
+        this.level = level;
     }
 
     public synchronized boolean addSample(byte[] sample) {
@@ -34,7 +40,7 @@ public ByteBuffer trainSamplesDirect() throws ZstdException {
 
     public synchronized ByteBuffer trainSamplesDirect(boolean legacy) throws ZstdException {
         ByteBuffer dictBuffer = ByteBuffer.allocateDirect(dictSize);
-        long l = Zstd.trainFromBufferDirect(trainingSamples, copyToIntArray(sampleSizes), dictBuffer, legacy);
+        long l = Zstd.trainFromBufferDirect(trainingSamples, copyToIntArray(sampleSizes), dictBuffer, legacy, level);
         if (Zstd.isError(l)) {
             dictBuffer.limit(0);
             throw new ZstdException(l);