From 53e208720e25624a149429952dd2000c82e59593 Mon Sep 17 00:00:00 2001
From: Jacob Greenfield <jag625@lehigh.edu>
Date: Sat, 11 Nov 2023 21:13:47 +0000
Subject: [PATCH] Allocate seqprod state during registration, so that each
 compression context gets its own state allocation

---
 .../github/luben/zstd/SequenceProducer.java   | 27 ++++++++++++-------
 .../github/luben/zstd/ZstdCompressCtx.java    | 18 ++++++++++---
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/main/java/com/github/luben/zstd/SequenceProducer.java b/src/main/java/com/github/luben/zstd/SequenceProducer.java
index 306cf1e..37a3336 100644
--- a/src/main/java/com/github/luben/zstd/SequenceProducer.java
+++ b/src/main/java/com/github/luben/zstd/SequenceProducer.java
@@ -1,21 +1,30 @@
 package com.github.luben.zstd;
 
 /**
- * Interface for an extenal sequence producer. The class implementing 
- * SequenceProducer must be provided by the user. To use, SequenceProducer should 
+ * Interface for an extenal sequence producer. The class implementing
+ * SequenceProducer must be provided by the user. To use, SequenceProducer should
  * be passed to `ZstdCompressCtx.registerSequenceProducer()`.
  */
 public interface SequenceProducer {
   /**
-   * Returns a function pointer of type `ZSTD_sequenceProducer_F` that will be
-   * passed to zstd as the `sequenceProducer` function. The returned pointer
-   * must be valid for the duration of this SeqeunceProducers's lifetime.
+   * Returns a pointer to the sequence producer function. This method is called
+   * once in {@link ZstdCompressCtx registerSequenceProducer}.
+   * @return A function pointer of type `ZSTD_sequenceProducer_F *`
    */
-  public long getProducerFunction();
+  public long getFunctionPointer();
 
   /**
-   * Returns a `void *` pointer to the `sequenceProducerState`. The returned
-   * pointer must be valid for the duration of this SeqeunceProducers's lifetime.
+   * Allocate the sequence producer state. The returned pointer will be passed
+   * to the sequence producer function. This method is called once in
+   * {@link ZstdCompressCtx#registerSequenceProducer(SequenceProducer)}.
+   * @return A `void *` pointer to the sequence producer state
    */
-  public long getStatePointer();
+  public long createState();
+
+  /**
+   * Free the sequence producer state. This method is called when closing the
+   * {@link ZstdCompressCtx} or registering a different sequence producer.
+   * @param statePointer the state pointer to be freed
+   */
+  public void freeState(long statePointer);
 }
\ No newline at end of file
diff --git a/src/main/java/com/github/luben/zstd/ZstdCompressCtx.java b/src/main/java/com/github/luben/zstd/ZstdCompressCtx.java
index 6fe7255..49b4377 100644
--- a/src/main/java/com/github/luben/zstd/ZstdCompressCtx.java
+++ b/src/main/java/com/github/luben/zstd/ZstdCompressCtx.java
@@ -16,7 +16,9 @@ public class ZstdCompressCtx extends AutoCloseBase {
 
     private ZstdDictCompress compression_dict = null;
 
-    private SequenceProducer sequence_producer = null;
+    private SequenceProducer seqprod = null;
+
+    private long seqprod_state = 0;
 
     private static native long init();
 
@@ -38,6 +40,10 @@ void doClose() {
         if (nativePtr != 0) {
             free(nativePtr);
             nativePtr = 0;
+            if (seqprod != null) {
+                seqprod.freeState(seqprod_state);
+                seqprod = null;
+            }
         }
     }
 
@@ -277,12 +283,18 @@ public ZstdCompressCtx setLong(int windowLog) {
     public void registerSequenceProducer(SequenceProducer producer) {
         ensureOpen();
         acquireSharedLock();
+        if (this.seqprod != null) {
+            this.seqprod.freeState(seqprod_state);
+        }
+
         if (producer == null) {
+            seqprod_state = 0;
             Zstd.registerSequenceProducer(nativePtr, 0, 0);
         } else {
-            Zstd.registerSequenceProducer(nativePtr, producer.getStatePointer(), producer.getProducerFunction());
+            seqprod_state = producer.createState();
+            Zstd.registerSequenceProducer(nativePtr, seqprod_state, producer.getFunctionPointer());
         }
-        this.sequence_producer = producer;
+        this.seqprod = producer;
         releaseSharedLock();
     }