Merge-in SharedDictionary feature (#916)

Co-authored-by: Eugene Kliuchnikov <eustas@chromium.org>

14 files changed, 427 insertions, 0 deletions

diff --git a/java/org/brotli/wrapper/dec/BUILD b/java/org/brotli/wrapper/dec/BUILD
index 4fd51d0..6190504 100644
--- a/java/org/brotli/wrapper/dec/BUILD
+++ b/java/org/brotli/wrapper/dec/BUILD
@@ -24,6 +24,7 @@ java_library(
         ":dec",
         "//org/brotli/integration:brotli_jni_test_base",
         "//org/brotli/integration:bundle_helper",
+        "//org/brotli/wrapper/enc",
         "@maven//:junit_junit",
     ],
 )
@@ -82,3 +83,16 @@ java_test(
     test_class = "org.brotli.wrapper.dec.DecoderTest",
     runtime_deps = [":test_lib"],
 )
+
+java_test(
+    name = "CornerCasesTest",
+    size = "large",
+    data = [
+        ":brotli_jni",  # Bazel JNI workaround
+    ],
+    jvm_flags = [
+        "-DBROTLI_JNI_LIBRARY=$(location :brotli_jni)",
+    ],
+    test_class = "org.brotli.wrapper.dec.CornerCasesTest",
+    runtime_deps = [":test_lib"],
+)
diff --git a/java/org/brotli/wrapper/dec/BrotliDecoderChannel.java b/java/org/brotli/wrapper/dec/BrotliDecoderChannel.java
index 6dfe8f2..30f6f1d 100644
--- a/java/org/brotli/wrapper/dec/BrotliDecoderChannel.java
+++ b/java/org/brotli/wrapper/dec/BrotliDecoderChannel.java
@@ -36,6 +36,11 @@ public class BrotliDecoderChannel extends Decoder implements ReadableByteChannel
   }
 
   @Override
+  public void attachDictionary(ByteBuffer dictionary) throws IOException {
+    super.attachDictionary(dictionary);
+  }
+
+  @Override
   public boolean isOpen() {
     synchronized (mutex) {
       return !closed;
diff --git a/java/org/brotli/wrapper/dec/BrotliInputStream.java b/java/org/brotli/wrapper/dec/BrotliInputStream.java
index 6e2e6e5..13a8880 100644
--- a/java/org/brotli/wrapper/dec/BrotliInputStream.java
+++ b/java/org/brotli/wrapper/dec/BrotliInputStream.java
@@ -8,6 +8,7 @@ package org.brotli.wrapper.dec;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
 
 /**
@@ -34,6 +35,10 @@ public class BrotliInputStream extends InputStream {
     this(source, DEFAULT_BUFFER_SIZE);
   }
 
+  public void attachDictionary(ByteBuffer dictionary) throws IOException {
+    decoder.attachDictionary(dictionary);
+  }
+
   public void enableEagerOutput() {
     decoder.enableEagerOutput();
   }
diff --git a/java/org/brotli/wrapper/dec/CornerCasesTest.java b/java/org/brotli/wrapper/dec/CornerCasesTest.java
new file mode 100644
index 0000000..ab02f23
--- /dev/null
+++ b/java/org/brotli/wrapper/dec/CornerCasesTest.java
@@ -0,0 +1,33 @@
+/* Copyright 2021 Google Inc. All Rights Reserved.
+
+   Distributed under MIT license.
+   See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
+*/
+
+package org.brotli.wrapper.dec;
+
+import static org.junit.Assert.assertEquals;
+
+import org.brotli.integration.BrotliJniTestBase;
+import org.brotli.wrapper.enc.Encoder;
+import java.io.IOException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link org.brotli.wrapper.enc.Encoder}. */
+@RunWith(JUnit4.class)
+public class CornerCasesTest extends BrotliJniTestBase {
+  @Test
+  public void testPowerOfTwoInput() throws IOException {
+    // 24 == max window bits to ensure ring-buffer size is not greater than input.
+    int len = 1 << 24;
+    byte[] data = new byte[len];
+    for (int i = 0; i < len; ++i) {
+      data[i] = (byte) Integer.bitCount(i);
+    }
+    byte[] encoded = Encoder.compress(data);
+    byte[] decoded = Decoder.decompress(encoded);
+    assertEquals(len, decoded.length);
+  }
+}
diff --git a/java/org/brotli/wrapper/dec/Decoder.java b/java/org/brotli/wrapper/dec/Decoder.java
index 018317d..de69d3a 100644
--- a/java/org/brotli/wrapper/dec/Decoder.java
+++ b/java/org/brotli/wrapper/dec/Decoder.java
@@ -50,6 +50,12 @@ public class Decoder {
     throw new IOException(message);
   }
 
+  void attachDictionary(ByteBuffer dictionary) throws IOException {
+    if (!decoder.attachDictionary(dictionary)) {
+      fail("failed to attach dictionary");
+    }
+  }
+
   public void enableEagerOutput() {
     this.eager = true;
   }
diff --git a/java/org/brotli/wrapper/dec/DecoderJNI.java b/java/org/brotli/wrapper/dec/DecoderJNI.java
index 2319b1e..fc5225d 100644
--- a/java/org/brotli/wrapper/dec/DecoderJNI.java
+++ b/java/org/brotli/wrapper/dec/DecoderJNI.java
@@ -17,6 +17,7 @@ public class DecoderJNI {
   private static native void nativePush(long[] context, int length);
   private static native ByteBuffer nativePull(long[] context);
   private static native void nativeDestroy(long[] context);
+  private static native boolean nativeAttachDictionary(long[] context, ByteBuffer dictionary);
 
   public enum Status {
     ERROR,
@@ -40,6 +41,19 @@ public class DecoderJNI {
       }
     }
 
+    public boolean attachDictionary(ByteBuffer dictionary) {
+      if (!dictionary.isDirect()) {
+        throw new IllegalArgumentException("only direct buffers allowed");
+      }
+      if (context[0] == 0) {
+        throw new IllegalStateException("brotli decoder is already destroyed");
+      }
+      if (!fresh) {
+        throw new IllegalStateException("decoding is already started");
+      }
+      return nativeAttachDictionary(context, dictionary);
+    }
+
     public void push(int length) {
       if (length < 0) {
         throw new IllegalArgumentException("negative block length");
diff --git a/java/org/brotli/wrapper/dec/decoder_jni.cc b/java/org/brotli/wrapper/dec/decoder_jni.cc
index 268a10b..e10ffe3 100644
--- a/java/org/brotli/wrapper/dec/decoder_jni.cc
+++ b/java/org/brotli/wrapper/dec/decoder_jni.cc
@@ -15,6 +15,9 @@ namespace {
 typedef struct DecoderHandle {
   BrotliDecoderState* state;
 
+  jobject dictionary_refs[15];
+  size_t dictionary_count;
+
   uint8_t* input_start;
   size_t input_offset;
   size_t input_length;
@@ -54,6 +57,10 @@ Java_org_brotli_wrapper_dec_DecoderJNI_nativeCreate(
   ok = !!handle;
 
   if (ok) {
+    for (int i = 0; i < 15; ++i) {
+      handle->dictionary_refs[i] = nullptr;
+    }
+    handle->dictionary_count = 0;
     handle->input_offset = 0;
     handle->input_length = 0;
     handle->input_start = nullptr;
@@ -193,10 +200,53 @@ Java_org_brotli_wrapper_dec_DecoderJNI_nativeDestroy(
   env->GetLongArrayRegion(ctx, 0, 3, context);
   DecoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
   BrotliDecoderDestroyInstance(handle->state);
+  for (size_t i = 0; i < handle->dictionary_count; ++i) {
+    env->DeleteGlobalRef(handle->dictionary_refs[i]);
+  }
   delete[] handle->input_start;
   delete handle;
 }
 
+JNIEXPORT jboolean JNICALL
+Java_org_brotli_wrapper_dec_DecoderJNI_nativeAttachDictionary(
+    JNIEnv* env, jobject /*jobj*/, jlongArray ctx, jobject dictionary) {
+  jlong context[3];
+  env->GetLongArrayRegion(ctx, 0, 3, context);
+  DecoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
+  jobject ref = nullptr;
+  uint8_t* address = nullptr;
+  jlong capacity = 0;
+
+  bool ok = true;
+  if (ok && !dictionary) {
+    ok = false;
+  }
+  if (ok && handle->dictionary_count >= 15) {
+    ok = false;
+  }
+  if (ok) {
+    ref = env->NewGlobalRef(dictionary);
+    ok = !!ref;
+  }
+  if (ok) {
+    handle->dictionary_refs[handle->dictionary_count] = ref;
+    handle->dictionary_count++;
+    address = static_cast<uint8_t*>(env->GetDirectBufferAddress(ref));
+    ok = !!address;
+  }
+  if (ok) {
+    capacity = env->GetDirectBufferCapacity(ref);
+    ok = (capacity > 0) && (capacity < (1 << 30));
+  }
+  if (ok) {
+    size_t size = static_cast<size_t>(capacity);
+    ok = !!BrotliDecoderAttachDictionary(handle->state,
+        BROTLI_SHARED_DICTIONARY_RAW, size, address);
+  }
+
+  return static_cast<jboolean>(ok);
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/java/org/brotli/wrapper/enc/BUILD b/java/org/brotli/wrapper/enc/BUILD
index 22154d2..0f8d2e1 100644
--- a/java/org/brotli/wrapper/enc/BUILD
+++ b/java/org/brotli/wrapper/enc/BUILD
@@ -18,6 +18,10 @@ java_library(
         ["*.java"],
         exclude = ["*Test*.java"],
     ),
+    deps = [
+        "//org/brotli/common:shared_dictionary",
+        "//org/brotli/enc:prepared_dictionary",
+    ],
     resource_jars = ["//:license"],
 )
 
@@ -27,8 +31,11 @@ java_library(
     srcs = glob(["*Test*.java"]),
     deps = [
         ":enc",
+        "//org/brotli/common:shared_dictionary",
+        "//org/brotli/enc:prepared_dictionary",
         "//org/brotli/integration:brotli_jni_test_base",
         "//org/brotli/integration:bundle_helper",
+        "//org/brotli/wrapper/common",
         "//org/brotli/wrapper/dec",
         "@maven//:junit_junit",
     ],
diff --git a/java/org/brotli/wrapper/enc/BrotliEncoderChannel.java b/java/org/brotli/wrapper/enc/BrotliEncoderChannel.java
index 047c356..2ee2ba6 100644
--- a/java/org/brotli/wrapper/enc/BrotliEncoderChannel.java
+++ b/java/org/brotli/wrapper/enc/BrotliEncoderChannel.java
@@ -6,6 +6,7 @@
 
 package org.brotli.wrapper.enc;
 
+import org.brotli.enc.PreparedDictionary;
 import java.io.IOException;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
@@ -43,6 +44,11 @@ public class BrotliEncoderChannel extends Encoder implements WritableByteChannel
   }
 
   @Override
+  public void attachDictionary(PreparedDictionary dictionary) throws IOException {
+    super.attachDictionary(dictionary);
+  }
+
+  @Override
   public boolean isOpen() {
     synchronized (mutex) {
       return !closed;
diff --git a/java/org/brotli/wrapper/enc/BrotliOutputStream.java b/java/org/brotli/wrapper/enc/BrotliOutputStream.java
index 5bd3957..09cbd5a 100644
--- a/java/org/brotli/wrapper/enc/BrotliOutputStream.java
+++ b/java/org/brotli/wrapper/enc/BrotliOutputStream.java
@@ -6,6 +6,7 @@
 
 package org.brotli.wrapper.enc;
 
+import org.brotli.enc.PreparedDictionary;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.channels.Channels;
@@ -40,6 +41,10 @@ public class BrotliOutputStream extends OutputStream {
     this(destination, new Encoder.Parameters());
   }
 
+  public void attachDictionary(PreparedDictionary dictionary) throws IOException {
+    encoder.attachDictionary(dictionary);
+  }
+
   @Override
   public void close() throws IOException {
     encoder.close();
diff --git a/java/org/brotli/wrapper/enc/Encoder.java b/java/org/brotli/wrapper/enc/Encoder.java
index af579a3..2aa9890 100644
--- a/java/org/brotli/wrapper/enc/Encoder.java
+++ b/java/org/brotli/wrapper/enc/Encoder.java
@@ -6,17 +6,20 @@
 
 package org.brotli.wrapper.enc;
 
+import org.brotli.enc.PreparedDictionary;
 import java.io.IOException;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.channels.WritableByteChannel;
 import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Base class for OutputStream / Channel implementations.
  */
 public class Encoder {
   private final WritableByteChannel destination;
+  private final List<PreparedDictionary> dictionaries;
   private final EncoderJNI.Wrapper encoder;
   private ByteBuffer buffer;
   final ByteBuffer inputBuffer;
@@ -111,6 +114,7 @@ public class Encoder {
     if (destination == null) {
       throw new NullPointerException("destination can not be null");
     }
+    this.dictionaries = new ArrayList<PreparedDictionary>();
     this.destination = destination;
     this.encoder = new EncoderJNI.Wrapper(inputBufferSize, params.quality, params.lgwin, params.mode);
     this.inputBuffer = this.encoder.getInputBuffer();
@@ -125,6 +129,14 @@ public class Encoder {
     throw new IOException(message);
   }
 
+  public void attachDictionary(PreparedDictionary dictionary) throws IOException {
+    if (!encoder.attachDictionary(dictionary.getData())) {
+      fail("failed to attach dictionary");
+    }
+    // Reference to native prepared dictionary wrapper should be held till the end of encoding.
+    dictionaries.add(dictionary);
+  }
+
   /**
    * @param force repeat pushing until all output is consumed
    * @return true if all encoder output is consumed
@@ -239,4 +251,15 @@ public class Encoder {
   public static byte[] compress(byte[] data) throws IOException {
     return compress(data, new Parameters());
   }
+
+  /**
+   * Prepares raw or serialized dictionary for being used by encoder.
+   *
+   * @param dictionary raw / serialized dictionary data; MUST be direct
+   * @param sharedDictionaryType dictionary data type
+   */
+  public static PreparedDictionary prepareDictionary(ByteBuffer dictionary,
+      int sharedDictionaryType) {
+    return EncoderJNI.prepareDictionary(dictionary, sharedDictionaryType);
+  }
 }
diff --git a/java/org/brotli/wrapper/enc/EncoderJNI.java b/java/org/brotli/wrapper/enc/EncoderJNI.java
index d0c5fac..1cbd7c1 100644
--- a/java/org/brotli/wrapper/enc/EncoderJNI.java
+++ b/java/org/brotli/wrapper/enc/EncoderJNI.java
@@ -6,6 +6,7 @@
 
 package org.brotli.wrapper.enc;
 
+import org.brotli.enc.PreparedDictionary;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 
@@ -17,6 +18,9 @@ class EncoderJNI {
   private static native void nativePush(long[] context, int length);
   private static native ByteBuffer nativePull(long[] context);
   private static native void nativeDestroy(long[] context);
+  private static native boolean nativeAttachDictionary(long[] context, ByteBuffer dictionary);
+  private static native ByteBuffer nativePrepareDictionary(ByteBuffer dictionary, long type);
+  private static native void nativeDestroyDictionary(ByteBuffer dictionary);
 
   enum Operation {
     PROCESS,
@@ -24,6 +28,47 @@ class EncoderJNI {
     FINISH
   }
 
+  private static class PreparedDictionaryImpl implements PreparedDictionary {
+    private ByteBuffer data;
+
+    private PreparedDictionaryImpl(ByteBuffer data) {
+      this.data = data;
+    }
+
+    @Override
+    public ByteBuffer getData() {
+      return data;
+    }
+
+    @Override
+    protected void finalize() throws Throwable {
+      try {
+        ByteBuffer data = this.data;
+        this.data = null;
+        nativeDestroyDictionary(data);
+      } finally {
+        super.finalize();
+      }
+    }
+  }
+
+  /**
+   * Prepares raw or serialized dictionary for being used by encoder.
+   *
+   * @param dictionary raw / serialized dictionary data; MUST be direct
+   * @param sharedDictionaryType dictionary data type
+   */
+  static PreparedDictionary prepareDictionary(ByteBuffer dictionary, int sharedDictionaryType) {
+    if (!dictionary.isDirect()) {
+      throw new IllegalArgumentException("only direct buffers allowed");
+    }
+    ByteBuffer dictionaryData = nativePrepareDictionary(dictionary, sharedDictionaryType);
+    if (dictionaryData == null) {
+      throw new IllegalStateException("OOM");
+    }
+    return new PreparedDictionaryImpl(dictionaryData);
+  }
+
   static class Wrapper {
     protected final long[] context = new long[5];
     private final ByteBuffer inputBuffer;
@@ -48,6 +93,19 @@ class EncoderJNI {
       this.context[4] = 0;
     }
 
+    boolean attachDictionary(ByteBuffer dictionary) {
+      if (!dictionary.isDirect()) {
+        throw new IllegalArgumentException("only direct buffers allowed");
+      }
+      if (context[0] == 0) {
+        throw new IllegalStateException("brotli decoder is already destroyed");
+      }
+      if (!fresh) {
+        throw new IllegalStateException("decoding is already started");
+      }
+      return nativeAttachDictionary(context, dictionary);
+    }
+
     void push(Operation op, int length) {
       if (length < 0) {
         throw new IllegalArgumentException("negative block length");
diff --git a/java/org/brotli/wrapper/enc/UseCompoundDictionaryTest.java b/java/org/brotli/wrapper/enc/UseCompoundDictionaryTest.java
new file mode 100644
index 0000000..cbfd0d9
--- /dev/null
+++ b/java/org/brotli/wrapper/enc/UseCompoundDictionaryTest.java
@@ -0,0 +1,114 @@
+package org.brotli.wrapper.enc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import org.brotli.common.SharedDictionaryType;
+import org.brotli.enc.PreparedDictionary;
+import org.brotli.enc.PreparedDictionaryGenerator;
+import org.brotli.integration.BrotliJniTestBase;
+import org.brotli.integration.BundleHelper;
+import org.brotli.wrapper.common.BrotliCommon;
+import org.brotli.wrapper.dec.BrotliInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.List;
+import junit.framework.TestCase;
+import junit.framework.TestSuite;
+import org.junit.runner.RunWith;
+import org.junit.runners.AllTests;
+
+/** Tests for compression / decompression aided with LZ77 dictionary. */
+@RunWith(AllTests.class)
+public class UseCompoundDictionaryTest extends BrotliJniTestBase {
+
+  static InputStream getBundle() throws IOException {
+    return new FileInputStream(System.getProperty("TEST_BUNDLE"));
+  }
+
+  /** Creates a test suite. */
+  public static TestSuite suite() throws IOException {
+    TestSuite suite = new TestSuite();
+    InputStream bundle = getBundle();
+    try {
+      List<String> entries = BundleHelper.listEntries(bundle);
+      for (String entry : entries) {
+        suite.addTest(new UseCompoundDictionaryTestCase(entry));
+      }
+    } finally {
+      bundle.close();
+    }
+    return suite;
+  }
+
+  /** Test case with a unique name. */
+  static class UseCompoundDictionaryTestCase extends TestCase {
+    final String entryName;
+    UseCompoundDictionaryTestCase(String entryName) {
+      super("UseCompoundDictionaryTest." + entryName);
+      this.entryName = entryName;
+    }
+
+    @Override
+    protected void runTest() throws Throwable {
+      UseCompoundDictionaryTest.run(entryName);
+    }
+  }
+
+  private static PreparedDictionary prepareRawDictionary(String entryName, ByteBuffer data) {
+    if (entryName.endsWith("E.coli.txt")) {
+      // Default prepared dictionary parameters doesn't work well for DNA data.
+      // Should work well with 8-byte hash.
+      return PreparedDictionaryGenerator.generate(data, 17, 3, 64, 5);
+    } else {
+       return Encoder.prepareDictionary(data, SharedDictionaryType.RAW);
+    }
+  }
+
+  private static void run(String entryName) throws Throwable {
+    InputStream bundle = getBundle();
+    byte[] original;
+    try {
+      original = BundleHelper.readEntry(bundle, entryName);
+    } finally {
+      bundle.close();
+    }
+
+    if (original == null) {
+      throw new RuntimeException("Can't read bundle entry: " + entryName);
+    }
+
+    ByteBuffer dictionary = BrotliCommon.makeNative(original);
+    PreparedDictionary preparedDictionary = prepareRawDictionary(entryName, dictionary);
+
+    ByteArrayOutputStream dst = new ByteArrayOutputStream();
+    BrotliOutputStream encoder =
+        new BrotliOutputStream(dst, new Encoder.Parameters().setQuality(9).setWindow(23), 1 << 23);
+    encoder.attachDictionary(preparedDictionary);
+    try {
+      encoder.write(original);
+    } finally {
+      encoder.close();
+    }
+
+    byte[] compressed = dst.toByteArray();
+
+    // Just copy self from LZ77 dictionary -> ultimate compression ratio.
+    assertTrue(compressed.length < 80 + original.length / 65536);
+
+    BrotliInputStream decoder =
+        new BrotliInputStream(new ByteArrayInputStream(compressed), 1 << 23);
+    decoder.attachDictionary(dictionary);
+    try {
+      long originalCrc = BundleHelper.fingerprintStream(new ByteArrayInputStream(original));
+      long crc = BundleHelper.fingerprintStream(decoder);
+      assertEquals(originalCrc, crc);
+    } finally {
+      decoder.close();
+    }
+  }
+}
\ No newline at end of file
diff --git a/java/org/brotli/wrapper/enc/encoder_jni.cc b/java/org/brotli/wrapper/enc/encoder_jni.cc
index e69f719..00fbfd8 100644
--- a/java/org/brotli/wrapper/enc/encoder_jni.cc
+++ b/java/org/brotli/wrapper/enc/encoder_jni.cc
@@ -15,6 +15,9 @@ namespace {
 typedef struct EncoderHandle {
   BrotliEncoderState* state;
 
+  jobject dictionary_refs[15];
+  size_t dictionary_count;
+
   uint8_t* input_start;
   size_t input_offset;
   size_t input_last;
@@ -53,6 +56,10 @@ Java_org_brotli_wrapper_enc_EncoderJNI_nativeCreate(
   ok = !!handle;
 
   if (ok) {
+    for (int i = 0; i < 15; ++i) {
+      handle->dictionary_refs[i] = nullptr;
+    }
+    handle->dictionary_count = 0;
     handle->input_offset = 0;
     handle->input_last = 0;
     handle->input_start = nullptr;
@@ -190,10 +197,90 @@ Java_org_brotli_wrapper_enc_EncoderJNI_nativeDestroy(
   env->GetLongArrayRegion(ctx, 0, 2, context);
   EncoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
   BrotliEncoderDestroyInstance(handle->state);
+  for (size_t i = 0; i < handle->dictionary_count; ++i) {
+    env->DeleteGlobalRef(handle->dictionary_refs[i]);
+  }
   delete[] handle->input_start;
   delete handle;
 }
 
+JNIEXPORT jboolean JNICALL
+Java_org_brotli_wrapper_enc_EncoderJNI_nativeAttachDictionary(
+    JNIEnv* env, jobject /*jobj*/, jlongArray ctx, jobject dictionary) {
+  jlong context[2];
+  env->GetLongArrayRegion(ctx, 0, 2, context);
+  EncoderHandle* handle = getHandle(reinterpret_cast<void*>(context[0]));
+  jobject ref = nullptr;
+  uint8_t* address = nullptr;
+
+  bool ok = true;
+  if (ok && !dictionary) {
+    ok = false;
+  }
+  if (ok && handle->dictionary_count >= 15) {
+    ok = false;
+  }
+  if (ok) {
+    ref = env->NewGlobalRef(dictionary);
+    ok = !!ref;
+  }
+  if (ok) {
+    handle->dictionary_refs[handle->dictionary_count] = ref;
+    handle->dictionary_count++;
+    address = static_cast<uint8_t*>(env->GetDirectBufferAddress(ref));
+    ok = !!address;
+  }
+  if (ok) {
+    ok = !!BrotliEncoderAttachPreparedDictionary(handle->state,
+        reinterpret_cast<BrotliEncoderPreparedDictionary*>(address));
+  }
+
+  return static_cast<jboolean>(ok);
+}
+
+JNIEXPORT void JNICALL
+Java_org_brotli_wrapper_enc_EncoderJNI_nativeDestroyDictionary(
+    JNIEnv* env, jobject /*jobj*/, jobject dictionary) {
+  if (!dictionary) {
+    return;
+  }
+  uint8_t* address =
+      static_cast<uint8_t*>(env->GetDirectBufferAddress(dictionary));
+  if (!address) {
+    return;
+  }
+  BrotliEncoderDestroyPreparedDictionary(
+      reinterpret_cast<BrotliEncoderPreparedDictionary*>(address));
+}
+
+JNIEXPORT jobject JNICALL
+Java_org_brotli_wrapper_enc_EncoderJNI_nativePrepareDictionary(
+    JNIEnv* env, jobject /*jobj*/, jobject dictionary, jlong type) {
+  if (!dictionary) {
+    return nullptr;
+  }
+  uint8_t* address =
+      static_cast<uint8_t*>(env->GetDirectBufferAddress(dictionary));
+  if (!address) {
+    return nullptr;
+  }
+  jlong capacity = env->GetDirectBufferCapacity(dictionary);
+  if ((capacity <= 0) || (capacity >= (1 << 30))) {
+    return nullptr;
+  }
+  BrotliSharedDictionaryType dictionary_type =
+      static_cast<BrotliSharedDictionaryType>(type);
+  size_t size = static_cast<size_t>(capacity);
+  BrotliEncoderPreparedDictionary* prepared_dictionary =
+      BrotliEncoderPrepareDictionary(dictionary_type, size, address,
+        BROTLI_MAX_QUALITY, nullptr, nullptr, nullptr);
+  if (!prepared_dictionary) {
+    return nullptr;
+  }
+  /* Size is 4 - just enough to check magic bytes. */
+  return env->NewDirectByteBuffer(prepared_dictionary, 4);
+}
+
 #ifdef __cplusplus
 }
 #endif