From 3cdea23beacec3c28ec428a8f7e8e339cf87310e Mon Sep 17 00:00:00 2001
From: "Mateusz \"Serafin\" Gajewski" <github@wendigo.pl>
Date: Wed, 4 Sep 2024 12:12:38 +0200
Subject: [PATCH] Add retained size to compressor

---
 src/main/java/io/airlift/compress/v3/Compressor.java        | 5 +++++
 .../java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java  | 6 ++++++
 .../java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java   | 2 +-
 src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java | 6 ++++++
 .../io/airlift/compress/v3/snappy/SnappyJavaCompressor.java | 6 ++++++
 .../io/airlift/compress/v3/snappy/SnappyRawCompressor.java  | 2 +-
 6 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/main/java/io/airlift/compress/v3/Compressor.java b/src/main/java/io/airlift/compress/v3/Compressor.java
index d47c4d2e..ecd00380 100644
--- a/src/main/java/io/airlift/compress/v3/Compressor.java
+++ b/src/main/java/io/airlift/compress/v3/Compressor.java
@@ -28,4 +28,9 @@ public interface Compressor
      * @return number of bytes written to the output
      */
     int compress(MemorySegment input, MemorySegment output);
+
+    default int getRetainedSizeInBytes(int inputLength)
+    {
+        return 0;
+    }
 }
diff --git a/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java b/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java
index 70888431..1b7ffd5e 100644
--- a/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java
+++ b/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java
@@ -69,6 +69,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return Lz4RawCompressor.computeTableSize(inputLength);
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");
diff --git a/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java b/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java
index 79f550c5..8ea20779 100644
--- a/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java
+++ b/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java
@@ -300,7 +300,7 @@ private static long encodeRunLength(
         return output;
     }
 
-    private static int computeTableSize(int inputSize)
+    static int computeTableSize(int inputSize)
     {
         // smallest power of 2 larger than inputSize
         int target = Integer.highestOneBit(inputSize - 1) << 1;
diff --git a/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java b/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java
index 29c7bc52..23e163d9 100644
--- a/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java
+++ b/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java
@@ -71,6 +71,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return MAX_TABLE_SIZE;
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");
diff --git a/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java b/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java
index 3411966f..e3ec79cc 100644
--- a/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java
+++ b/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java
@@ -75,6 +75,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return SnappyRawCompressor.getHashTableSize(inputLength);
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");
diff --git a/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java b/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java
index b2edf392..251cca9c 100644
--- a/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java
+++ b/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java
@@ -344,7 +344,7 @@ private static long emitCopy(Object outputBase, long output, long input, long ma
     }
 
     @SuppressWarnings("IllegalToken")
-    private static int getHashTableSize(int inputSize)
+    static int getHashTableSize(int inputSize)
     {
         // Use smaller hash table when input.size() is smaller, since we
         // fill the table, incurring O(hash table size) overhead for