From 56c20c2c200c2b9c0f4af175e213a766346da531 Mon Sep 17 00:00:00 2001
From: Alan Race <alan.race@uni-marburg.de>
Date: Wed, 27 Apr 2022 17:38:44 +0200
Subject: [PATCH] Added ability to select which image (IFD) to load

---
 src/decoder/mod.rs | 67 +++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 6 deletions(-)

diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs
index f0be812..14cf693 100644
--- a/src/decoder/mod.rs
+++ b/src/decoder/mod.rs
@@ -323,8 +323,10 @@ where
     reader: SmartReader<R>,
     bigtiff: bool,
     limits: Limits,
-    next_ifd: Option<u64>,
     current_chunk: usize,
+    next_ifd: Option<u64>,
+    next_ifd_index: usize,
+    ifds: Vec<u64>,
     seen_ifds: HashSet<u64>,
     image: Image,
 }
@@ -598,12 +600,16 @@ impl<R: Read + Seek> Decoder<R> {
 
         let mut seen_ifds = HashSet::new();
         seen_ifds.insert(*next_ifd.as_ref().unwrap());
+        let ifds = vec![*next_ifd.as_ref().unwrap()];
 
         let mut decoder = Decoder {
             reader,
             bigtiff,
             limits: Default::default(),
             next_ifd,
+            next_ifd_index: 1,
+            ifds,
+            seen_ifds,
             image: Image {
                 ifd: None,
                 width: 0,
@@ -622,7 +628,6 @@ impl<R: Read + Seek> Decoder<R> {
                 chunk_bytes: Vec::new(),
             },
             current_chunk: 0,
-            seen_ifds,
         };
         decoder.next_image()?;
         Ok(decoder)
@@ -645,6 +650,48 @@ impl<R: Read + Seek> Decoder<R> {
         &self.image
     }
 
+    /// Loads the IFD at the specified index in the list, if one exists
+    pub fn image_at(&mut self, ifd_index: usize) -> TiffResult<()> {
+        if ifd_index >= self.next_ifd_index {
+            // We possibly need to load in the next IFD
+            if self.next_ifd.is_none() {
+                return Err(TiffError::FormatError(
+                    TiffFormatError::ImageFileDirectoryNotFound,
+                ));
+            }
+
+            loop {
+                // Follow the list until we find the one we want
+                let (ifd, next_ifd) = self.next_ifd()?;
+
+                if next_ifd.is_none() {
+                    break;
+                }
+
+                if ifd_index < self.ifds.len() {
+                    break;
+                }
+            }
+        }
+
+        if ifd_index < self.ifds.len() {
+            let (ifd, next_ifd) = Self::read_ifd(
+                &mut self.reader,
+                self.bigtiff,
+                *self.ifds.get(ifd_index).unwrap(),
+            )?;
+
+            self.current_chunk = 0;
+            self.image = Image::from_reader(&mut self.reader, ifd, &self.limits, self.bigtiff)?;
+
+            Ok(())
+        } else {
+            Err(TiffError::FormatError(
+                TiffFormatError::ImageFileDirectoryNotFound,
+            ))
+        }
+    }
+
     /// Reset the decoder.
     #[deprecated = "Never should have been public. Only use Decoder::new()"]
     pub fn init(self) -> TiffResult<Decoder<R>> {
@@ -652,10 +699,7 @@ impl<R: Read + Seek> Decoder<R> {
         Self::new(reader.into_inner())
     }
 
-    /// Reads in the next image.
-    /// If there is no further image in the TIFF file a format error is returned.
-    /// To determine whether there are more images call `TIFFDecoder::more_images` instead.
-    pub fn next_image(&mut self) -> TiffResult<()> {
+    fn next_ifd(&mut self) -> TiffResult<(Directory, Option<u64>)> {
         if self.next_ifd.is_none() {
             return Err(TiffError::FormatError(
                 TiffFormatError::ImageFileDirectoryNotFound,
@@ -673,8 +717,19 @@ impl<R: Read + Seek> Decoder<R> {
                 return Err(TiffError::FormatError(TiffFormatError::CycleInOffsets));
             }
             self.next_ifd = Some(next);
+            self.next_ifd_index += 1;
+            self.ifds.push(next);
         }
 
+        Ok((ifd, next_ifd))
+    }
+
+    /// Reads in the next image.
+    /// If there is no further image in the TIFF file a format error is returned.
+    /// To determine whether there are more images call `TIFFDecoder::more_images` instead.
+    pub fn next_image(&mut self) -> TiffResult<()> {
+        let (ifd, next_ifd) = self.next_ifd()?;
+
         self.current_chunk = 0;
         self.image = Image::from_reader(&mut self.reader, ifd, &self.limits, self.bigtiff)?;
         Ok(())