Merge pull request #281 from image-rs/write-directory

Adjust Directory for ease of writing

src/decoder/ifd.rs

@@ -350,9 +350,9 @@ impl ::std::fmt::Debug for Entry {
 
 impl Entry {
     pub fn new(type_: Type, count: u32, offset: [u8; 4]) -> Entry {
-        let mut offset = offset.to_vec();
-        offset.append(&mut vec![0; 4]);
-        Entry::new_u64(type_, count.into(), offset[..].try_into().unwrap())
+        let mut entry_off = [0u8; 8];
+        entry_off[..4].copy_from_slice(&offset);
+        Entry::new_u64(type_, count.into(), entry_off)
     }
 
     pub fn new_u64(type_: Type, count: u64, offset: [u8; 8]) -> Entry {
@@ -363,6 +363,10 @@ impl Entry {
         }
     }
 
+    pub fn field_type(&self) -> Type {
+        self.type_
+    }
+
     pub fn count(&self) -> u64 {
         self.count
     }

src/directory.rs

@@ -27,10 +27,10 @@ pub struct Directory {
 }
 
 impl Directory {
-    pub fn empty(next: IfdPointer) -> Self {
+    pub fn empty() -> Self {
         Directory {
             entries: BTreeMap::new(),
-            next_ifd: NonZeroU64::new(next.0),
+            next_ifd: None,
         }
     }
 
@@ -64,15 +64,12 @@ impl Directory {
         self.extend_inner(iter.into_iter().by_ref())
     }
 
-    /// Get the length as a 2-byte integer.
-    ///
-    /// This is *almost* naturally bounded by the tag type being a `u16` itself. The
-    /// bounds are upheld when extending the iteration, i.e. this always corresponds to the
-    /// underlying data.
-    pub fn len(&self) -> u16 {
+    /// Get the number of entries.
+    pub fn len(&self) -> usize {
         // The keys are `u16`. Since IFDs are required to have at least one entry this would have
-        // been a trivial thing to do in the specification by storing it minus one but alas.
-        self.entries.len() as u16
+        // been a trivial thing to do in the specification by storing it minus one but alas. In
+        // BigTIFF the count is stored as 8-bit anyways.
+        self.entries.len()
     }
 
     /// Get the pointer to the next IFD, if it was defined.
@@ -80,16 +77,15 @@ impl Directory {
         self.next_ifd.map(|n| IfdPointer(n.get()))
     }
 
+    pub fn set_next(&mut self, next: Option<IfdPointer>) {
+        self.next_ifd = next.and_then(|n| NonZeroU64::new(n.0));
+    }
+
     fn extend_inner(&mut self, iter: &mut dyn Iterator<Item = (Tag, Entry)>) {
         for (tag, entry) in iter {
-            let is_full = self.entries.len() == u16::MAX as usize;
             // If the tag is already present, it will be overwritten.
             let map_entry = self.entries.entry(tag.to_u16());
 
-            if is_full {
-                // Only allowed to modify if the entry is already present.
-                map_entry.and_modify(|place| *place = entry);
-            } else {
             match map_entry {
                 std::collections::btree_map::Entry::Vacant(vacant_entry) => {
                     vacant_entry.insert(entry);
@@ -101,7 +97,6 @@ impl Directory {
         }
     }
 }
-}
 
 impl fmt::Debug for Directory {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -117,40 +112,12 @@ impl fmt::Debug for Directory {
 
 #[cfg(test)]
 mod tests {
-    use super::{Directory, IfdPointer};
+    use super::Directory;
     use crate::{decoder::ifd::Entry, tags::Tag};
 
-    #[test]
-    fn directory_upholds_len_invariant() {
-        let mut dir = Directory::empty(IfdPointer(0));
-        assert_eq!(dir.len(), 0);
-
-        let bogus_entry = Entry::new_u64(crate::tags::Type::BYTE, 0x0, [0; 8]);
-        dir.extend((0..=u16::MAX).map(|i| {
-            let tag = Tag::Unknown(i);
-            let entry = bogus_entry.clone();
-            (tag, entry)
-        }));
-
-        assert_eq!(dir.len(), u16::MAX);
-
-        // Ensure the other tags are still writable.
-        dir.extend([(
-            Tag::Unknown(0),
-            Entry::new_u64(crate::tags::Type::BYTE, 0x42, [0; 8]),
-        )]);
-
-        assert_eq!(
-            dir.get(Tag::Unknown(0))
-                .expect("tag 0 should be present after overwriting")
-                .count(),
-            0x42
-        );
-    }
-
     #[test]
     fn directory_multiple_entries() {
-        let mut dir = Directory::empty(IfdPointer(0));
+        let mut dir = Directory::empty();
         assert_eq!(dir.len(), 0);
 
         dir.extend((0..=u16::MAX).map(|i| {
@@ -171,7 +138,7 @@ mod tests {
 
     #[test]
     fn iteration_order() {
-        let mut dir = Directory::empty(IfdPointer(0));
+        let mut dir = Directory::empty();
         assert_eq!(dir.len(), 0);
 
         let fake_entry = Entry::new_u64(crate::tags::Type::BYTE, 0, [0; 8]);