zstd: set pledged size if available

build.rs

@@ -1,6 +1,11 @@
 fn main() {
     // See note in Cargo.toml
     if cfg!(target_pointer_width = "32") {
+        #[cfg(any(
+            feature = "use-system-isal",
+            feature = "isal-static",
+            feature = "isal-shared"
+        ))]
         let msg = "feature set but ISA-L not supported on 32 bit systems.";
 
         #[cfg(feature = "use-system-isal")]

src/capi.rs

@@ -263,6 +263,7 @@ pub extern "C" fn compress(
             &mut decompressed,
             &mut compressed,
             level.map(|v: i32| v as i32),
+            Some(input_len),
         ),
         #[cfg(feature = "lz4")]
         Codec::Lz4 => lz4::compress(&mut decompressed, &mut compressed, level.map(|v| v as _)),
@@ -388,6 +389,7 @@ pub extern "C" fn compress_into(
             &mut decompressed,
             &mut compressed,
             level.map(|v: i32| v as i32),
+            Some(input_len),
         ),
         #[cfg(feature = "lz4")]
         Codec::Lz4 => lz4::compress(&mut decompressed, &mut compressed, level.map(|v| v as _)),
@@ -1063,7 +1065,7 @@ mod tests {
     #[test]
     fn test_zstd_roundtrip() {
         let mut expected = Cursor::new(vec![]);
-        zstd::compress(Cursor::new(DATA), &mut expected, Some(6)).unwrap();
+        zstd::compress(Cursor::new(DATA), &mut expected, Some(6), Some(DATA.len())).unwrap();
         let expected = expected.into_inner();
         roundtrip(Codec::Zstd, &expected, 6);
     }

src/lib.rs

@@ -172,7 +172,7 @@ mod tests {
     test_variant!(deflate, None);
 
     #[cfg(feature = "zstd")]
-    test_variant!(zstd, None);
+    test_variant!(zstd, None, None);
 
     #[cfg(feature = "zlib")]
     test_variant!(zlib, None);

src/zstd.rs

@@ -19,9 +19,17 @@ pub fn decompress<W: Write + ?Sized, R: Read>(input: R, output: &mut W) -> Resul
 
 /// Compress gzip data
 #[inline(always)]
-pub fn compress<W: Write + ?Sized, R: Read>(input: R, output: &mut W, level: Option<i32>) -> Result<usize, Error> {
+pub fn compress<W: Write + ?Sized, R: Read>(
+    input: R,
+    output: &mut W,
+    level: Option<i32>,
+    input_size: Option<usize>,
+) -> Result<usize, Error> {
     let level = level.unwrap_or_else(|| DEFAULT_COMPRESSION_LEVEL); // 0 will use zstd's default, currently 3
     let mut encoder = zstd::stream::read::Encoder::new(input, level)?;
+
+    encoder.set_pledged_src_size(input_size.map(|v| v as u64))?;
+
     let n_bytes = std::io::copy(&mut encoder, output)?;
     Ok(n_bytes as usize)
 }