Skip to main content

tiff_writer/
writer.rs

1//! Main TiffWriter: orchestrates multi-IFD streaming writes.
2
3use std::io::{Seek, SeekFrom, Write};
4
5use tiff_core::{ByteOrder, Tag};
6
7use crate::builder::ImageBuilder;
8use crate::compress;
9use crate::encoder;
10use crate::error::{Error, Result};
11use crate::sample::TiffWriteSample;
12
13const CLASSIC_TIFF_LIMIT: u64 = u32::MAX as u64;
14
15fn checked_len_u64(len: usize, context: &str) -> Result<u64> {
16    u64::try_from(len).map_err(|_| Error::Other(format!("{context} length exceeds u64::MAX")))
17}
18
19fn checked_add_u64(lhs: u64, rhs: u64, context: &str) -> Result<u64> {
20    lhs.checked_add(rhs)
21        .ok_or_else(|| Error::Other(format!("{context} overflow")))
22}
23
24fn classic_offset_u32(offset: u64) -> Result<u32> {
25    u32::try_from(offset).map_err(|_| Error::ClassicOffsetOverflow { offset })
26}
27
28fn classic_byte_count_u32(byte_count: u64) -> Result<u32> {
29    u32::try_from(byte_count).map_err(|_| Error::ClassicByteCountOverflow { byte_count })
30}
31
32/// TIFF format variant.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum TiffVariant {
35    Classic,
36    BigTiff,
37    /// Exact auto-selection.
38    ///
39    /// The writer records block data first and chooses Classic TIFF vs
40    /// BigTIFF from the finalized file layout during `finish()`.
41    Auto,
42}
43
44/// Configuration for the TIFF writer.
45#[derive(Debug, Clone)]
46pub struct WriteOptions {
47    pub byte_order: ByteOrder,
48    pub variant: TiffVariant,
49}
50
51impl Default for WriteOptions {
52    fn default() -> Self {
53        Self {
54            byte_order: ByteOrder::LittleEndian,
55            variant: TiffVariant::Auto,
56        }
57    }
58}
59
60impl WriteOptions {
61    /// Construct exact auto-selection options.
62    ///
63    /// The `estimated_bytes` parameter is retained for source compatibility,
64    /// but the writer now decides from the exact finalized layout instead of
65    /// an upfront size heuristic.
66    pub fn auto(_estimated_bytes: u64) -> Self {
67        Self {
68            byte_order: ByteOrder::LittleEndian,
69            variant: TiffVariant::Auto,
70        }
71    }
72}
73
74/// Handle identifying a specific image within the writer.
75#[derive(Debug, Clone)]
76pub struct ImageHandle {
77    pub(crate) index: usize,
78}
79
80/// Write state for one IFD.
81struct IfdState {
82    builder: ImageBuilder,
83    block_records: Vec<Option<(u64, u64)>>,
84}
85
86/// A streaming TIFF/BigTIFF file writer.
87pub struct TiffWriter<W: Write + Seek> {
88    sink: W,
89    byte_order: ByteOrder,
90    requested_variant: TiffVariant,
91    header_offset: u64,
92    images: Vec<IfdState>,
93    finalized: bool,
94}
95
96impl<W: Write + Seek> TiffWriter<W> {
97    /// Create a new TIFF writer.
98    ///
99    /// Image data is streamed immediately. The final IFD chain and the header
100    /// are emitted during `finish()`, which allows `TiffVariant::Auto` to
101    /// choose Classic TIFF vs BigTIFF from the exact completed layout.
102    pub fn new(mut sink: W, options: WriteOptions) -> Result<Self> {
103        let header_offset = sink.stream_position()?;
104        let reserved_header_len = match options.variant {
105            TiffVariant::Classic => encoder::header_len(false),
106            TiffVariant::BigTiff | TiffVariant::Auto => encoder::header_len(true),
107        };
108        sink.write_all(&[0; encoder::BIGTIFF_HEADER_LEN as usize][..reserved_header_len as usize])?;
109
110        Ok(Self {
111            sink,
112            byte_order: options.byte_order,
113            requested_variant: options.variant,
114            header_offset,
115            images: Vec::new(),
116            finalized: false,
117        })
118    }
119
120    /// Add an image (IFD) to the file.
121    pub fn add_image(&mut self, builder: ImageBuilder) -> Result<ImageHandle> {
122        if self.finalized {
123            return Err(Error::AlreadyFinalized);
124        }
125        builder.validate()?;
126
127        let index = self.images.len();
128        self.images.push(IfdState {
129            block_records: vec![None; builder.block_count()],
130            builder,
131        });
132
133        Ok(ImageHandle { index })
134    }
135
136    /// Write raw bytes between the reserved header area and the image data.
137    ///
138    /// When `TiffVariant::Auto` is used, the writer reserves the 16-byte
139    /// BigTIFF header footprint up front so the finalized header can switch
140    /// variants without relocating block data.
141    pub fn write_header_prefix(&mut self, bytes: &[u8]) -> Result<()> {
142        if self.finalized {
143            return Err(Error::AlreadyFinalized);
144        }
145        if !self.images.is_empty() {
146            return Err(Error::Other(
147                "header prefix bytes must be written before adding images".into(),
148            ));
149        }
150
151        self.sink.seek(SeekFrom::End(0))?;
152        let prefix_end = checked_add_u64(
153            self.sink.stream_position()?,
154            checked_len_u64(bytes.len(), "header prefix")?,
155            "header prefix size",
156        )?;
157        if matches!(self.requested_variant, TiffVariant::Classic) {
158            classic_offset_u32(prefix_end)?;
159        }
160        self.sink.write_all(bytes)?;
161        Ok(())
162    }
163
164    /// Write a single strip or tile for the given image.
165    pub fn write_block<T: TiffWriteSample>(
166        &mut self,
167        handle: &ImageHandle,
168        block_index: usize,
169        samples: &[T],
170    ) -> Result<()> {
171        if self.finalized {
172            return Err(Error::AlreadyFinalized);
173        }
174        let state = self
175            .images
176            .get(handle.index)
177            .ok_or(Error::Other("invalid image handle".into()))?;
178
179        let total_blocks = state.builder.block_count();
180        if block_index >= total_blocks {
181            return Err(Error::BlockIndexOutOfRange {
182                index: block_index,
183                total: total_blocks,
184            });
185        }
186
187        let expected = state.builder.block_sample_count(block_index);
188        if samples.len() != expected {
189            return Err(Error::BlockSizeMismatch {
190                index: block_index,
191                expected,
192                actual: samples.len(),
193            });
194        }
195
196        let compressed = if matches!(state.builder.compression, tiff_core::Compression::Lerc) {
197            let opts = state.builder.lerc_options.unwrap_or_default();
198            let block_width = state.builder.block_row_width() as u32;
199            let block_height = state.builder.block_height(block_index);
200            let depth = state.builder.block_samples_per_pixel() as u32;
201            compress::compress_block_lerc(
202                samples,
203                block_width,
204                block_height,
205                depth,
206                &opts,
207                block_index,
208            )?
209        } else {
210            compress::compress_block(
211                samples,
212                compress::BlockEncodingOptions {
213                    byte_order: self.byte_order,
214                    compression: state.builder.compression,
215                    predictor: state.builder.predictor,
216                    samples_per_pixel: state.builder.block_samples_per_pixel(),
217                    row_width_pixels: state.builder.block_row_width(),
218                    jpeg_options: state.builder.jpeg_options.as_ref(),
219                },
220                block_index,
221            )?
222        };
223
224        self.write_block_raw(handle, block_index, &compressed)
225    }
226
227    /// Write a pre-compressed block (bypass the compression pipeline).
228    pub fn write_block_raw(
229        &mut self,
230        handle: &ImageHandle,
231        block_index: usize,
232        compressed_bytes: &[u8],
233    ) -> Result<()> {
234        if self.finalized {
235            return Err(Error::AlreadyFinalized);
236        }
237
238        let state = self
239            .images
240            .get(handle.index)
241            .ok_or(Error::Other("invalid image handle".into()))?;
242        let total = state.builder.block_count();
243        if block_index >= total {
244            return Err(Error::BlockIndexOutOfRange {
245                index: block_index,
246                total,
247            });
248        }
249
250        let offset = self.sink.seek(SeekFrom::End(0))?;
251        let byte_count = checked_len_u64(compressed_bytes.len(), "block payload")?;
252        if matches!(self.requested_variant, TiffVariant::Classic) {
253            classic_offset_u32(offset)?;
254            classic_byte_count_u32(byte_count)?;
255        }
256
257        self.sink.write_all(compressed_bytes)?;
258
259        let state = self
260            .images
261            .get_mut(handle.index)
262            .ok_or(Error::Other("invalid image handle".into()))?;
263        state.block_records[block_index] = Some((offset, byte_count));
264        Ok(())
265    }
266
267    /// Write a block whose on-disk bytes include a prefix and/or suffix that
268    /// must not be reflected in the TIFF block offset/byte-count arrays.
269    pub fn write_block_raw_segmented(
270        &mut self,
271        handle: &ImageHandle,
272        block_index: usize,
273        prefix: &[u8],
274        payload: &[u8],
275        suffix: &[u8],
276    ) -> Result<()> {
277        if self.finalized {
278            return Err(Error::AlreadyFinalized);
279        }
280
281        let state = self
282            .images
283            .get(handle.index)
284            .ok_or(Error::Other("invalid image handle".into()))?;
285        let total = state.builder.block_count();
286        if block_index >= total {
287            return Err(Error::BlockIndexOutOfRange {
288                index: block_index,
289                total,
290            });
291        }
292
293        let start = self.sink.seek(SeekFrom::End(0))?;
294        let prefix_len = checked_len_u64(prefix.len(), "block prefix")?;
295        let byte_count = checked_len_u64(payload.len(), "block payload")?;
296        let suffix_len = checked_len_u64(suffix.len(), "block suffix")?;
297        let offset = checked_add_u64(start, prefix_len, "block offset")?;
298        let end = checked_add_u64(
299            checked_add_u64(offset, byte_count, "segmented block size")?,
300            suffix_len,
301            "segmented block size",
302        )?;
303        if matches!(self.requested_variant, TiffVariant::Classic) {
304            classic_offset_u32(offset)?;
305            classic_byte_count_u32(byte_count)?;
306            classic_offset_u32(end)?;
307        }
308
309        self.sink.write_all(prefix)?;
310        self.sink.write_all(payload)?;
311        self.sink.write_all(suffix)?;
312
313        let state = self
314            .images
315            .get_mut(handle.index)
316            .ok_or(Error::Other("invalid image handle".into()))?;
317        state.block_records[block_index] = Some((offset, byte_count));
318        Ok(())
319    }
320
321    fn choose_is_bigtiff(&mut self) -> Result<bool> {
322        match self.requested_variant {
323            TiffVariant::Classic => {
324                self.ensure_classic_layout()?;
325                Ok(false)
326            }
327            TiffVariant::BigTiff => Ok(true),
328            TiffVariant::Auto => Ok(!self.classic_layout_fits()?),
329        }
330    }
331
332    fn classic_layout_fits(&mut self) -> Result<bool> {
333        for state in &self.images {
334            for &(offset, byte_count) in state.block_records.iter().flatten() {
335                if offset > CLASSIC_TIFF_LIMIT || byte_count > CLASSIC_TIFF_LIMIT {
336                    return Ok(false);
337                }
338            }
339        }
340
341        let mut current = self.sink.seek(SeekFrom::End(0))?;
342        for state in &self.images {
343            let tags = state.builder.build_tags(false);
344            current = checked_add_u64(
345                current,
346                encoder::estimate_ifd_size(self.byte_order, false, &tags),
347                "classic IFD layout",
348            )?;
349            if current > CLASSIC_TIFF_LIMIT {
350                return Ok(false);
351            }
352        }
353
354        Ok(true)
355    }
356
357    fn ensure_classic_layout(&mut self) -> Result<()> {
358        for state in &self.images {
359            for &(offset, byte_count) in state.block_records.iter().flatten() {
360                classic_offset_u32(offset)?;
361                classic_byte_count_u32(byte_count)?;
362            }
363        }
364
365        let mut current = self.sink.seek(SeekFrom::End(0))?;
366        for state in &self.images {
367            let tags = state.builder.build_tags(false);
368            current = checked_add_u64(
369                current,
370                encoder::estimate_ifd_size(self.byte_order, false, &tags),
371                "classic IFD layout",
372            )?;
373            classic_offset_u32(current)?;
374        }
375
376        Ok(())
377    }
378
379    fn write_final_ifds(
380        &mut self,
381        is_bigtiff: bool,
382    ) -> Result<Vec<(Vec<Tag>, encoder::IfdWriteResult)>> {
383        let mut results = Vec::with_capacity(self.images.len());
384        for state in &self.images {
385            let tags = state.builder.build_tags(is_bigtiff);
386            let (offsets_tag_code, byte_counts_tag_code) = state.builder.offset_tag_codes();
387            let ifd_result = encoder::write_ifd(
388                &mut self.sink,
389                self.byte_order,
390                is_bigtiff,
391                &tags,
392                offsets_tag_code,
393                byte_counts_tag_code,
394                state.builder.block_count(),
395            )?;
396            results.push((tags, ifd_result));
397        }
398        Ok(results)
399    }
400
401    /// Finalize the TIFF file. Emits the IFD chain and patches the header.
402    pub fn finish(mut self) -> Result<W> {
403        if self.finalized {
404            return Err(Error::AlreadyFinalized);
405        }
406        self.finalized = true;
407
408        for state in &self.images {
409            let total = state.builder.block_count();
410            let written = state
411                .block_records
412                .iter()
413                .filter(|record| record.is_some())
414                .count();
415            if written != total {
416                return Err(Error::IncompleteImage { written, total });
417            }
418        }
419
420        let is_bigtiff = self.choose_is_bigtiff()?;
421
422        self.sink.seek(SeekFrom::Start(self.header_offset))?;
423        encoder::write_header(&mut self.sink, self.byte_order, is_bigtiff)?;
424
425        self.sink.seek(SeekFrom::End(0))?;
426        let ifd_results = self.write_final_ifds(is_bigtiff)?;
427
428        for (img_idx, state) in self.images.iter().enumerate() {
429            let offsets: Vec<u64> = state
430                .block_records
431                .iter()
432                .map(|record| record.unwrap().0)
433                .collect();
434            let byte_counts: Vec<u64> = state
435                .block_records
436                .iter()
437                .map(|record| record.unwrap().1)
438                .collect();
439
440            let (tags, ifd_result) = &ifd_results[img_idx];
441            let (offsets_tag_code, byte_counts_tag_code) = state.builder.offset_tag_codes();
442
443            if offsets.len() == 1 {
444                if let Some(off) = encoder::find_inline_tag_value_offset(
445                    ifd_result.ifd_offset,
446                    is_bigtiff,
447                    tags,
448                    offsets_tag_code,
449                ) {
450                    self.sink.seek(SeekFrom::Start(off))?;
451                    if is_bigtiff {
452                        self.sink
453                            .write_all(&self.byte_order.write_u64(offsets[0]))?;
454                    } else {
455                        self.sink.write_all(
456                            &self.byte_order.write_u32(classic_offset_u32(offsets[0])?),
457                        )?;
458                    }
459                }
460                if let Some(off) = encoder::find_inline_tag_value_offset(
461                    ifd_result.ifd_offset,
462                    is_bigtiff,
463                    tags,
464                    byte_counts_tag_code,
465                ) {
466                    self.sink.seek(SeekFrom::Start(off))?;
467                    if is_bigtiff {
468                        self.sink
469                            .write_all(&self.byte_order.write_u64(byte_counts[0]))?;
470                    } else {
471                        self.sink.write_all(
472                            &self
473                                .byte_order
474                                .write_u32(classic_byte_count_u32(byte_counts[0])?),
475                        )?;
476                    }
477                }
478            } else {
479                if let Some(off) = ifd_result.offsets_tag_data_offset {
480                    encoder::patch_block_offsets(
481                        &mut self.sink,
482                        self.byte_order,
483                        is_bigtiff,
484                        off,
485                        &offsets,
486                    )?;
487                }
488                if let Some(off) = ifd_result.byte_counts_tag_data_offset {
489                    encoder::patch_block_byte_counts(
490                        &mut self.sink,
491                        self.byte_order,
492                        is_bigtiff,
493                        off,
494                        &byte_counts,
495                    )?;
496                }
497            }
498
499            if img_idx == 0 {
500                encoder::patch_first_ifd(
501                    &mut self.sink,
502                    self.header_offset,
503                    self.byte_order,
504                    is_bigtiff,
505                    ifd_result.ifd_offset,
506                )?;
507            } else {
508                let prev = &ifd_results[img_idx - 1].1;
509                encoder::patch_next_ifd(
510                    &mut self.sink,
511                    self.byte_order,
512                    is_bigtiff,
513                    prev.next_ifd_pointer_offset,
514                    ifd_result.ifd_offset,
515                )?;
516            }
517        }
518
519        self.sink.seek(SeekFrom::End(0))?;
520        Ok(self.sink)
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use std::io::{self, Cursor, Seek, SeekFrom, Write};
527
528    use super::*;
529    use crate::builder::ImageBuilder;
530
531    #[derive(Default)]
532    struct CountingSink {
533        pos: u64,
534        len: u64,
535    }
536
537    impl Write for CountingSink {
538        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
539            self.pos += buf.len() as u64;
540            self.len = self.len.max(self.pos);
541            Ok(buf.len())
542        }
543
544        fn flush(&mut self) -> io::Result<()> {
545            Ok(())
546        }
547    }
548
549    impl Seek for CountingSink {
550        fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
551            let next = match pos {
552                SeekFrom::Start(offset) => offset as i128,
553                SeekFrom::End(delta) => self.len as i128 + delta as i128,
554                SeekFrom::Current(delta) => self.pos as i128 + delta as i128,
555            };
556            if next < 0 {
557                return Err(io::Error::new(
558                    io::ErrorKind::InvalidInput,
559                    "negative seek in CountingSink",
560                ));
561            }
562            self.pos = next as u64;
563            self.len = self.len.max(self.pos);
564            Ok(self.pos)
565        }
566    }
567
568    #[test]
569    fn auto_promotes_to_bigtiff_from_the_final_layout() {
570        let mut writer = TiffWriter::new(CountingSink::default(), WriteOptions::default()).unwrap();
571        let handle = writer
572            .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
573            .unwrap();
574
575        writer
576            .sink
577            .seek(SeekFrom::Start(CLASSIC_TIFF_LIMIT + 1))
578            .unwrap();
579        writer.write_block_raw(&handle, 0, &[1]).unwrap();
580
581        assert!(writer.choose_is_bigtiff().unwrap());
582    }
583
584    #[test]
585    fn auto_keeps_classic_for_small_layouts() {
586        let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
587        let handle = writer
588            .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
589            .unwrap();
590        writer.write_block(&handle, 0, &[7u8]).unwrap();
591
592        assert!(!writer.choose_is_bigtiff().unwrap());
593    }
594
595    #[test]
596    fn write_block_raw_validates_before_mutating_sink() {
597        let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
598        let handle = writer
599            .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
600            .unwrap();
601
602        let len_before = writer.sink.get_ref().len();
603        let err = writer.write_block_raw(&handle, 1, &[1, 2, 3]).unwrap_err();
604
605        assert!(matches!(err, Error::BlockIndexOutOfRange { .. }));
606        assert_eq!(writer.sink.get_ref().len(), len_before);
607    }
608
609    #[test]
610    fn write_block_raw_segmented_validates_before_mutating_sink() {
611        let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
612        let handle = writer
613            .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
614            .unwrap();
615
616        let len_before = writer.sink.get_ref().len();
617        let err = writer
618            .write_block_raw_segmented(&handle, 1, &[1, 2], &[3, 4], &[5, 6])
619            .unwrap_err();
620
621        assert!(matches!(err, Error::BlockIndexOutOfRange { .. }));
622        assert_eq!(writer.sink.get_ref().len(), len_before);
623    }
624}