Nicer wrapper for cfzlib (#149)

Cargo.lock

@@ -83,6 +83,14 @@ dependencies = [
  "vec_map 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "cloudflare-zlib"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "cloudflare-zlib-sys 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "cloudflare-zlib-sys"
 version = "0.2.0"
@@ -257,7 +265,7 @@ dependencies = [
  "bit-vec 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "byteorder 1.2.7 (registry+https://github.com/rust-lang/crates.io-index)",
  "clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "cloudflare-zlib-sys 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "cloudflare-zlib 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "crc 1.8.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "image 0.20.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "itertools 0.7.9 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -414,6 +422,7 @@ dependencies = [
 "checksum cc 1.0.25 (registry+https://github.com/rust-lang/crates.io-index)" = "f159dfd43363c4d08055a07703eb7a3406b0dac4d0584d96965a3262db3c9d16"
 "checksum cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)" = "082bb9b28e00d3c9d39cc03e64ce4cea0f1bb9b3fde493f0cbc008472d22bdf4"
 "checksum clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b957d88f4b6a63b9d70d5f454ac8011819c6efa7727858f458ab71c756ce2d3e"
+"checksum cloudflare-zlib 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "398f5d6b4431f230c89858e5389803bd0135bc764fbf9b7b50d4be0ca2657a6c"
 "checksum cloudflare-zlib-sys 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "7e195cb274a0d6ee87e718838a09baecd7cbc9f6075dac256a84cb5842739c06"
 "checksum crc 1.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "d663548de7f5cca343f1e0a48d14dcfb0e9eb4e079ec58883b7251539fa10aeb"
 "checksum crossbeam-deque 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f739f8c5363aca78cfb059edf753d8f0d36908c348f3d8d1503f03d8b75d9cf3"

Cargo.toml

@@ -49,7 +49,7 @@ optional = true
 version = "2.0.0"
 
 [target.'cfg(any(target_arch = "x86_64", target_arch = "aarch64"))'.dependencies]
-cloudflare-zlib-sys = "^0.2.0"
+cloudflare-zlib = "^0.2.0"
 
 [dependencies.image]
 default-features = false

src/atomicmin.rs

@@ -21,6 +21,11 @@ impl AtomicMin {
         }
     }
 
+    /// Unset value is usize_max
+    pub fn as_atomic_usize(&self) -> &AtomicUsize {
+        &self.val
+    }
+
     pub fn set_min(&self, new_val: usize) {
         let mut current_val = self.val.load(Relaxed);
         loop {

src/deflate/cfzlib.rs

@@ -0,0 +1,42 @@
+use atomicmin::AtomicMin;
+use PngResult;
+use PngError;
+use cloudflare_zlib::*;
+pub use cloudflare_zlib::is_supported;
+
+impl From<ZError> for PngError {
+    fn from(err: ZError) -> Self {
+        match err {
+            ZError::DeflatedDataTooLarge(n) => PngError::DeflatedDataTooLong(n),
+            other => PngError::Other(other.to_string().into()),
+        }
+    }
+}
+
+pub fn cfzlib_deflate(
+    data: &[u8],
+    level: u8,
+    strategy: u8,
+    window_bits: u8,
+    max_size: &AtomicMin,
+) -> PngResult<Vec<u8>> {
+    let mut stream = Deflate::new(level.into(), strategy.into(), window_bits.into())?;
+    stream.reserve(max_size.get().unwrap_or(data.len()/2));
+    let max_size = max_size.as_atomic_usize();
+    // max size is generally checked after each split,
+    // so splitting the buffer into pieces gices more checks
+    // = better chance of hitting it sooner.
+    let (first, rest) = data.split_at(data.len()/2);
+    stream.compress_with_limit(first, max_size)?;
+    let (rest1, rest2) = rest.split_at(rest.len()/2);
+    stream.compress_with_limit(rest1, max_size)?;
+    stream.compress_with_limit(rest2, max_size)?;
+    Ok(stream.finish()?)
+}
+
+#[test]
+fn compress_test() {
+    let vec = cfzlib_deflate(b"azxcvbnm", Z_BEST_COMPRESSION as u8, Z_DEFAULT_STRATEGY as u8, 15, &AtomicMin::new(None)).unwrap();
+    let res = ::deflate::inflate(&vec).unwrap();
+    assert_eq!(&res, b"azxcvbnm");
+}

src/deflate/mod.rs

@@ -8,6 +8,9 @@ use PngResult;
 #[doc(hidden)]
 pub mod miniz_stream;
 
+#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
+mod cfzlib;
+
 /// Decompress a data stream using the DEFLATE algorithm
 pub fn inflate(data: &[u8]) -> PngResult<Vec<u8>> {
     miniz_oxide::inflate::decompress_to_vec_zlib(data)
@@ -18,86 +21,14 @@ pub fn inflate(data: &[u8]) -> PngResult<Vec<u8>> {
 pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8, max_size: &AtomicMin) -> PngResult<Vec<u8>> {
     #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
     {
-        if is_cfzlib_supported() {
+        if cfzlib::is_supported() {
-            return cfzlib_deflate(data, zc, zs, zw, max_size);
+            return cfzlib::cfzlib_deflate(data, zc, zs, zw, max_size);
         }
     }
 
     miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into(), max_size)
 }
 
-#[cfg(target_arch = "x86_64")]
-fn is_cfzlib_supported() -> bool {
-    if is_x86_feature_detected!("sse4.2") && is_x86_feature_detected!("pclmulqdq") {
-        return true;
-    }
-    false
-}
-
-#[cfg(target_arch = "aarch64")]
-fn is_cfzlib_supported() -> bool {
-    if is_arm_feature_detected!("neon") && is_arm_feature_detected!("crc") {
-        return true;
-    }
-    false
-}
-
-#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
-pub fn cfzlib_deflate(
-    data: &[u8],
-    level: u8,
-    strategy: u8,
-    window_bits: u8,
-    max_size: &AtomicMin,
-) -> PngResult<Vec<u8>> {
-    use cloudflare_zlib_sys::*;
-    use std::mem;
-
-    assert!(data.len() < u32::max_value() as usize);
-    unsafe {
-        let mut stream = mem::zeroed();
-        if Z_OK != deflateInit2(
-            &mut stream,
-            level.into(),
-            Z_DEFLATED,
-            window_bits.into(),
-            MAX_MEM_LEVEL,
-            strategy.into(),
-        ) {
-            return Err(PngError::new("deflateInit2"));
-        }
-
-        let upper_bound = deflateBound(&mut stream, data.len() as uLong) as usize;
-        let max_size = max_size.get().unwrap_or(upper_bound).min(upper_bound);
-        // it's important to have the capacity pre-allocated,
-        // as unsafe set_len is called later
-        let mut out = Vec::with_capacity(max_size);
-
-        stream.next_in = data.as_ptr() as *mut _;
-        stream.total_in = data.len() as uLong;
-        stream.avail_in = data.len() as uInt;
-        stream.next_out = out.as_mut_ptr();
-        stream.avail_out = out.capacity() as uInt;
-        match deflate(&mut stream, Z_FINISH) {
-            Z_STREAM_END => {}
-            Z_OK | Z_BUF_ERROR => {
-                deflateEnd(&mut stream);
-                return Err(PngError::DeflatedDataTooLong(stream.total_out as usize));
-            }
-            _ => {
-                deflateEnd(&mut stream);
-                return Err(PngError::new("deflate"));
-            }
-        }
-        if Z_OK != deflateEnd(&mut stream) {
-            return Err(PngError::new("deflateEnd"));
-        }
-        debug_assert!(stream.total_out as usize <= out.capacity());
-        out.set_len(stream.total_out as usize);
-        Ok(out)
-    }
-}
-
 pub fn zopfli_deflate(data: &[u8]) -> PngResult<Vec<u8>> {
     let mut output = Vec::with_capacity(max(1024, data.len() / 20));
     let options = zopfli::Options::default();

src/lib.rs

@@ -1,7 +1,7 @@
 extern crate bit_vec;
 extern crate byteorder;
 #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
-extern crate cloudflare_zlib_sys;
+extern crate cloudflare_zlib;
 extern crate crc;
 extern crate image;
 extern crate rgb;