[pypy-commit] pypy default: Merged in megaumi/pypy/zlib_zdict (pull request #286)

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: 
Changeset: r74045:c88db2931b91
Date: 2014-10-21 17:33 +0200
http://bitbucket.org/pypy/pypy/changeset/c88db2931b91/

Log:	Merged in megaumi/pypy/zlib_zdict (pull request #286)

	(vmukhame, ronan)

diff --git a/rpython/rlib/rzlib.py b/rpython/rlib/rzlib.py
--- a/rpython/rlib/rzlib.py
+++ b/rpython/rlib/rzlib.py
@@ -37,7 +37,7 @@
     Z_NO_FLUSH  Z_FINISH  Z_SYNC_FLUSH  Z_FULL_FLUSH
     MAX_WBITS  MAX_MEM_LEVEL
     Z_BEST_SPEED  Z_BEST_COMPRESSION  Z_DEFAULT_COMPRESSION
-    Z_FILTERED  Z_HUFFMAN_ONLY  Z_DEFAULT_STRATEGY
+    Z_FILTERED  Z_HUFFMAN_ONLY  Z_DEFAULT_STRATEGY Z_NEED_DICT
     '''.split()
 
 class SimpleCConfig:
@@ -165,6 +165,9 @@
     result = _inflateInit2_(stream, wbits, ZLIB_VERSION, size)
     return result
 
+_deflateSetDictionary = zlib_external('deflateSetDictionary', [z_stream_p, Bytefp, uInt], rffi.INT)
+_inflateSetDictionary = zlib_external('inflateSetDictionary', [z_stream_p, Bytefp, uInt], rffi.INT)
+
 # ____________________________________________________________
 
 CRC32_DEFAULT_START = 0
@@ -184,6 +187,23 @@
 
 ADLER32_DEFAULT_START = 1
 
+def deflateSetDictionary(stream, string):
+    bytes = rffi.get_nonmovingbuffer(string)
+    err = _deflateSetDictionary(stream, rffi.cast(Bytefp, bytes), len(string))
+    rffi.free_nonmovingbuffer(string, bytes)
+    if err == Z_STREAM_ERROR:
+        raise RZlibError("Parameter is invalid or the stream state is inconsistent")
+
+def inflateSetDictionary(stream, string):
+    bytes = rffi.get_nonmovingbuffer(string)
+    err = _inflateSetDictionary(stream, rffi.cast(Bytefp, bytes), len(string))
+    rffi.free_nonmovingbuffer(string, bytes)
+    if err == Z_STREAM_ERROR:
+        raise RZlibError("Parameter is invalid or the stream state is inconsistent")
+    elif err == Z_DATA_ERROR:
+        raise RZlibError("The given dictionary doesn't match the expected one")
+
+    
 def adler32(string, start=ADLER32_DEFAULT_START):
     """
     Compute the Adler-32 checksum of the string, possibly with the given
diff --git a/rpython/rlib/test/test_rzlib.py b/rpython/rlib/test/test_rzlib.py
--- a/rpython/rlib/test/test_rzlib.py
+++ b/rpython/rlib/test/test_rzlib.py
@@ -82,6 +82,39 @@
     rzlib.deflateEnd(stream)
 
 
+def test_deflate_set_dictionary():
+    text = 'abcabc'
+    zdict = 'abc'
+    stream = rzlib.deflateInit()
+    rzlib.deflateSetDictionary(stream, zdict)
+    bytes = rzlib.compress(stream, text, rzlib.Z_FINISH)
+    rzlib.deflateEnd(stream)
+    
+    stream2 = rzlib.inflateInit()
+
+    from rpython.rtyper.lltypesystem import lltype, rffi, rstr
+    from rpython.rtyper.annlowlevel import llstr
+    from rpython.rlib.rstring import StringBuilder
+    with lltype.scoped_alloc(rffi.CCHARP.TO, len(bytes)) as inbuf:
+        rstr.copy_string_to_raw(llstr(bytes), inbuf, 0, len(bytes))
+        stream2.c_next_in = rffi.cast(rzlib.Bytefp, inbuf)
+        rffi.setintfield(stream2, 'c_avail_in', len(bytes))
+        with lltype.scoped_alloc(rffi.CCHARP.TO, 100) as outbuf:
+            stream2.c_next_out = rffi.cast(rzlib.Bytefp, outbuf)
+            bufsize = 100
+            rffi.setintfield(stream2, 'c_avail_out', bufsize)
+            err = rzlib._inflate(stream2, rzlib.Z_SYNC_FLUSH)
+            assert err == rzlib.Z_NEED_DICT
+            rzlib.inflateSetDictionary(stream2, zdict)
+            rzlib._inflate(stream2, rzlib.Z_SYNC_FLUSH)
+            avail_out = rffi.cast(lltype.Signed, stream2.c_avail_out)
+            result = StringBuilder()
+            result.append_charpsize(outbuf, bufsize - avail_out)
+
+    rzlib.inflateEnd(stream2)
+    assert result.build() == text
+
+
 def test_compression():
     """
     Once we have got a deflate stream, rzlib.compress()