From 7da8f01f7523ea3a65f883d4ea128a368096345a Mon Sep 17 00:00:00 2001
From: Robert Wassens <robert.wassens@technolution.nl>
Date: Sat, 5 Dec 2020 01:39:37 +0100
Subject: [PATCH] make strip size configurable

---
 src/encoder/mod.rs     | 29 ++++++++++++++++++++++++++---
 tests/encode_images.rs | 16 ++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs
index e055fad..952c3e0 100644
--- a/src/encoder/mod.rs
+++ b/src/encoder/mod.rs
@@ -589,6 +589,9 @@ impl<'a, W: Write + Seek> Drop for DirectoryEncoder<'a, W> {
 /// // You can encode tags here
 /// image.encoder().write_tag(Tag::Artist, "Image-tiff").unwrap();
 ///
+/// // Smaller strip size can be configured before writing data
+/// image.rows_per_strip(2).unwrap();
+///
 /// let mut idx = 0;
 /// while image.next_strip_sample_count() > 0 {
 ///     let sample_count = image.next_strip_sample_count() as usize;
@@ -620,10 +623,9 @@ impl<'a, W: 'a + Write + Seek, T: ColorType> ImageEncoder<'a, W, T> {
         height: u32,
     ) -> TiffResult<ImageEncoder<'a, W, T>> {
         let row_samples = u64::from(width) * u64::try_from(<T>::BITS_PER_SAMPLE.len())?;
-        let row_bytes = row_samples * u64::from(<T::Inner>::BYTE_LEN);
 
-        // As per tiff spec each strip should be about 8k long
-        // let rows_per_strip = (8000 + row_bytes - 1) / row_bytes;
+        // By default write the image as single strip. This can be changed by calling
+        // rows_per_strip() before the first data is written.
         let rows_per_strip = height as u64;
 
         let strip_count = (u64::from(height) + rows_per_strip - 1) / rows_per_strip;
@@ -756,6 +758,27 @@ impl<'a, W: 'a + Write + Seek, T: ColorType> ImageEncoder<'a, W, T> {
         self.encoder.write_tag(Tag::YResolution, value).unwrap();
     }
 
+    /// Set image number of lines per strip
+    ///
+    /// This function must be called before calling write_data() or write_strip().
+    pub fn rows_per_strip(&mut self, value: u32) -> TiffResult<()> {
+        if self.strip_idx != 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "Cannot change strip size after data was written",
+            )
+            .into());
+        }
+        // Write tag as 32 bits
+        self.encoder.write_tag(Tag::RowsPerStrip, value)?;
+
+        let value: u64 = value as u64;
+        self.strip_count = (self.height as u64 + value - 1) / value;
+        self.rows_per_strip = value;
+
+        Ok(())
+    }
+
     fn finish_internal(&mut self) -> TiffResult<()> {
         self.encoder
             .write_tag(Tag::StripOffsets, &*self.strip_offsets)?;
diff --git a/tests/encode_images.rs b/tests/encode_images.rs
index e5c6a8c..d389122 100644
--- a/tests/encode_images.rs
+++ b/tests/encode_images.rs
@@ -360,3 +360,19 @@ fn test_multipage_image() {
         assert_eq!(img_decoder.dimensions().unwrap(), (3, 3));
     }
 }
+
+#[test]
+/// verify rows per strip setting
+fn test_rows_per_strip() {
+    let mut file = Cursor::new(Vec::new());
+    let mut img_encoder = TiffEncoder::new(&mut file).unwrap();
+
+    let mut image = img_encoder.new_image::<colortype::Gray8>(100, 100).unwrap();
+    assert_eq!(image.next_strip_sample_count(), 100 * 100);
+    image.rows_per_strip(2).unwrap();
+    assert_eq!(image.next_strip_sample_count(), 2 * 100);
+
+    let img2: Vec<u8> = vec![0; 2 * 100];
+    image.write_strip(&img2[..]).unwrap();
+    assert!(image.rows_per_strip(5).is_err());
+}