Truncate dictionary to window size.

2 files changed, 22 insertions, 15 deletions

diff --git a/python/brotlimodule.cc b/python/brotlimodule.cc
index 99a4f8d..936d1a4 100644
--- a/python/brotlimodule.cc
+++ b/python/brotlimodule.cc
@@ -105,8 +105,8 @@ PyDoc_STRVAR(compress__doc__,
 "  lgblock (int, optional): Base 2 logarithm of the maximum input block size.\n"
 "    Range is 16 to 24. If set to 0, the value will be set based on the\n"
 "    quality. Defaults to 0.\n"
-"  dictionary (bytes, optional): Custom dictionary. Should be shorter than\n"
-"     sliding window size.\n"
+"  dictionary (bytes, optional): Custom dictionary. Only last sliding window\n"
+"     size bytes will be used.\n"
 "\n"
 "Returns:\n"
 "  The compressed byte string.\n"
@@ -158,10 +158,16 @@ static PyObject* brotli_compress(PyObject *self, PyObject *args, PyObject *keywd
     ok = BrotliCompressBuffer(params, length, input,
                               &output_length, output);
   } else {
+    uint8_t *custom_dictionary_start = custom_dictionary;
     BrotliMemIn in(input, length);
     BrotliMemOut out(output, output_length);
+    size_t sliding_window_size = ((size_t)1) << params.lgwin;
+    if (custom_dictionary_length > sliding_window_size) {
+      custom_dictionary_start += custom_dictionary_length - sliding_window_size;
+      custom_dictionary_length = sliding_window_size;
+    }
     ok = BrotliCompressWithCustomDictionary(custom_dictionary_length,
-        custom_dictionary, params, &in, &out);
+        custom_dictionary_start, params, &in, &out);
     output_length = out.position();
   }
 
diff --git a/python/tests/custom_dictionary_test.py b/python/tests/custom_dictionary_test.py
index 8fdebb7..afbf07a 100644
--- a/python/tests/custom_dictionary_test.py
+++ b/python/tests/custom_dictionary_test.py
@@ -21,15 +21,16 @@ testdata/plrabn12.txt
 os.chdir(os.path.abspath("../../tests"))
 for filename in INPUTS.splitlines():
     for quality in (1, 6, 9, 11):
-        filename = os.path.abspath(filename)
-        print('Roundtrip testing file "%s" at quality %d with auto-custom-dictionary' %
-              (os.path.basename(filename), quality))
-        compressed = os.path.splitext(filename)[0] + ".custom_bro"
-        uncompressed = os.path.splitext(filename)[0] + ".custom_unbro"
-        check_call([PYTHON, BRO, "-f", "-q", str(quality), "-i", filename,
-                    "-o", compressed, "--lgwin", "24",
-                    "--custom-dictionary", filename], env=TEST_ENV)
-        check_call([PYTHON, BRO, "-f", "-d", "-i", compressed, "-o",
-                    uncompressed, "--custom-dictionary", filename], env=TEST_ENV)
-        if diff_q(filename, uncompressed) != 0:
-            sys.exit(1)
+        for lgwin in (10, 15, 20, 24):
+            filename = os.path.abspath(filename)
+            print('Roundtrip testing file "%s" at quality %d with lg(win)=%d and auto-custom-dictionary' %
+                  (os.path.basename(filename), quality, lgwin))
+            compressed = os.path.splitext(filename)[0] + ".custom_bro"
+            uncompressed = os.path.splitext(filename)[0] + ".custom_unbro"
+            check_call([PYTHON, BRO, "-f", "-q", str(quality), "-i", filename,
+                        "-o", compressed, "--lgwin", str(lgwin),
+                        "--custom-dictionary", filename], env=TEST_ENV)
+            check_call([PYTHON, BRO, "-f", "-d", "-i", compressed, "-o",
+                        uncompressed, "--custom-dictionary", filename], env=TEST_ENV)
+            if diff_q(filename, uncompressed) != 0:
+                sys.exit(1)