Add WriteChunk::chunk_len method

This method returns the full length of the chunk when serialized, including the
chunk identifier and (written) chunk length.

Author: danieldk

src/chunks/io.rs

@@ -2,6 +2,7 @@ use std::convert::TryFrom;
 use std::fmt::{self, Display};
 use std::fs::File;
 use std::io::{BufReader, Read, Seek, Write};
+use std::mem;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
@@ -154,6 +155,12 @@ pub trait WriteChunk {
     /// Get the identifier of a chunk.
     fn chunk_identifier(&self) -> ChunkIdentifier;
 
+    /// Get the serialized length of a chunk.
+    ///
+    /// The `offset` of the chunk in the serialized data is required
+    /// because some chunks store arrays aligned.
+    fn chunk_len(&self, offset: u64) -> u64;
+
     fn write_chunk<W>(&self, write: &mut W) -> Result<()>
     where
         W: Write + Seek;
@@ -181,6 +188,14 @@ impl WriteChunk for Header {
         ChunkIdentifier::Header
     }
 
+    fn chunk_len(&self, _offset: u64) -> u64 {
+        // magic + model version (u32) + chunk ids len (u32), chunk ids (len * u32)
+        (MAGIC.len()
+            + mem::size_of_val(&MODEL_VERSION)
+            + mem::size_of::<u32>()
+            + self.chunk_identifiers.len() * mem::size_of::<u32>()) as u64
+    }
+
     fn write_chunk<W>(&self, write: &mut W) -> Result<()>
     where
         W: Write + Seek,
@@ -256,6 +271,17 @@ mod tests {
 
     use super::{ChunkIdentifier, Header, ReadChunk, WriteChunk};
 
+    #[test]
+    fn header_chunk_len_is_correct() {
+        let check_header =
+            Header::new(vec![ChunkIdentifier::SimpleVocab, ChunkIdentifier::NdArray]);
+        let mut cursor = Cursor::new(Vec::new());
+        check_header.write_chunk(&mut cursor).unwrap();
+
+        let data = cursor.into_inner();
+        assert_eq!(data.len() as u64, check_header.chunk_len(0));
+    }
+
     #[test]
     fn header_write_read_roundtrip() {
         let check_header =

src/chunks/metadata.rs

@@ -1,6 +1,7 @@
 //! Metadata chunks
 
 use std::io::{Read, Seek, Write};
+use std::mem;
 use std::ops::{Deref, DerefMut};
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
@@ -80,6 +81,11 @@ impl WriteChunk for Metadata {
         ChunkIdentifier::Metadata
     }
 
+    fn chunk_len(&self, _offset: u64) -> u64 {
+        // chunk identifier (u32) + metadata length (u64) + metadata
+        (mem::size_of::<u32>() + mem::size_of::<u64>() + self.to_string().len()) as u64
+    }
+
     fn write_chunk<W>(&self, write: &mut W) -> Result<()>
     where
         W: Write + Seek,
@@ -154,16 +160,24 @@ mod tests {
 
     #[test]
     fn metadata_correct_chunk_size() {
-        let check_metadata = test_metadata();
-        let mut cursor = Cursor::new(Vec::new());
-        check_metadata.write_chunk(&mut cursor).unwrap();
-        cursor.seek(SeekFrom::Start(0)).unwrap();
-
-        let chunk_size = read_chunk_size(&mut cursor);
-        assert_eq!(
-            cursor.read_to_end(&mut Vec::new()).unwrap(),
-            chunk_size as usize
-        );
+        for offset in 0..16u64 {
+            let check_metadata = test_metadata();
+            let mut cursor = Cursor::new(Vec::new());
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+            check_metadata.write_chunk(&mut cursor).unwrap();
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+
+            // Check remaining chunk size against size written into the chunk.
+            let chunk_size = read_chunk_size(&mut cursor);
+            assert_eq!(
+                cursor.read_to_end(&mut Vec::new()).unwrap() as u64,
+                chunk_size
+            );
+
+            // Check overall chunk size.
+            let data = cursor.into_inner();
+            assert_eq!(data.len() as u64 - offset, check_metadata.chunk_len(offset));
+        }
     }
 
     #[test]

src/chunks/norms.rs

@@ -2,6 +2,7 @@
 
 use std::convert::TryInto;
 use std::io::{Read, Seek, SeekFrom, Write};
+use std::mem;
 use std::mem::size_of;
 use std::ops::Deref;
 
@@ -89,6 +90,18 @@ impl WriteChunk for NdNorms {
         ChunkIdentifier::NdNorms
     }
 
+    fn chunk_len(&self, offset: u64) -> u64 {
+        let n_padding = padding::<f32>(offset + mem::size_of::<u32>() as u64);
+
+        // Chunk identifier (u32) + chunk len (u64) + len (u64) + type id (u32) + padding + vector.
+        (mem::size_of::<u32>()
+            + mem::size_of::<u64>()
+            + mem::size_of::<u64>()
+            + mem::size_of::<u32>()
+            + self.len() * mem::size_of::<f32>()) as u64
+            + n_padding
+    }
+
     fn write_chunk<W>(&self, write: &mut W) -> Result<()>
     where
         W: Write + Seek,
@@ -156,16 +169,26 @@ mod tests {
 
     #[test]
     fn ndnorms_correct_chunk_size() {
-        let check_arr = test_ndnorms();
-        let mut cursor = Cursor::new(Vec::new());
-        check_arr.write_chunk(&mut cursor).unwrap();
-        cursor.seek(SeekFrom::Start(0)).unwrap();
-
-        let chunk_size = read_chunk_size(&mut cursor);
-        assert_eq!(
-            cursor.read_to_end(&mut Vec::new()).unwrap(),
-            chunk_size as usize
-        );
+        for offset in 0..16u64 {
+            let check_arr = test_ndnorms();
+            let mut cursor = Cursor::new(Vec::new());
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+            check_arr.write_chunk(&mut cursor).unwrap();
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+
+            // Check size remained chunk against embedded chunk size.
+            let chunk_size = read_chunk_size(&mut cursor);
+            assert_eq!(
+                cursor.read_to_end(&mut Vec::new()).unwrap(),
+                chunk_size as usize
+            );
+
+            // Check overall chunk size.
+            assert_eq!(
+                cursor.into_inner().len() as u64 - offset,
+                check_arr.chunk_len(offset)
+            );
+        }
     }
 
     #[test]

src/chunks/storage/array.rs

@@ -1,5 +1,6 @@
 use std::convert::TryInto;
 use std::io::{Read, Seek, SeekFrom, Write};
+use std::mem;
 use std::mem::size_of;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
@@ -170,6 +171,10 @@ mod mmap {
             ChunkIdentifier::NdArray
         }
 
+        fn chunk_len(&self, offset: u64) -> u64 {
+            NdArray::chunk_len(self.view(), offset)
+        }
+
         fn write_chunk<W>(&self, write: &mut W) -> Result<()>
         where
             W: Write + Seek,
@@ -193,6 +198,19 @@ impl NdArray {
         NdArray { inner: arr }
     }
 
+    fn chunk_len(data: ArrayView2<f32>, offset: u64) -> u64 {
+        let n_padding = padding::<f32>(offset + mem::size_of::<u32>() as u64);
+
+        // Chunk identifier (u32) + chunk len (u64) + rows (u64) + cols (u32) + type id (u32) + padding + matrix.
+        (mem::size_of::<u32>()
+            + mem::size_of::<u64>()
+            + mem::size_of::<u64>()
+            + mem::size_of::<u32>()
+            + mem::size_of::<u32>()
+            + data.len() * mem::size_of::<f32>()) as u64
+            + n_padding
+    }
+
     fn write_ndarray_chunk<W>(data: ArrayView2<f32>, write: &mut W) -> Result<()>
     where
         W: Write + Seek,
@@ -346,6 +364,10 @@ impl WriteChunk for NdArray {
         ChunkIdentifier::NdArray
     }
 
+    fn chunk_len(&self, offset: u64) -> u64 {
+        Self::chunk_len(self.inner.view(), offset)
+    }
+
     fn write_chunk<W>(&self, write: &mut W) -> Result<()>
     where
         W: Write + Seek,
@@ -356,13 +378,13 @@ impl WriteChunk for NdArray {
 
 #[cfg(test)]
 mod tests {
-    use std::io::{Cursor, Read, Seek, SeekFrom};
+    use std::io::{Cursor, Seek, SeekFrom};
 
-    use byteorder::{LittleEndian, ReadBytesExt};
     use ndarray::Array2;
 
     use crate::chunks::io::{ReadChunk, WriteChunk};
     use crate::chunks::storage::{NdArray, Storage, StorageView};
+    use crate::storage::tests::test_storage_chunk_len;
 
     const N_ROWS: usize = 100;
     const N_COLS: usize = 100;
@@ -375,14 +397,6 @@ mod tests {
         NdArray::new(test_data)
     }
 
-    fn read_chunk_size(read: &mut impl Read) -> u64 {
-        // Skip identifier.
-        read.read_u32::<LittleEndian>().unwrap();
-
-        // Return chunk length.
-        read.read_u64::<LittleEndian>().unwrap()
-    }
-
     #[test]
     fn embeddings_returns_expected_embeddings() {
         const CHECK_INDICES: &[usize] = &[0, 50, 99, 0];
@@ -398,16 +412,7 @@ mod tests {
 
     #[test]
     fn ndarray_correct_chunk_size() {
-        let check_arr = test_ndarray();
-        let mut cursor = Cursor::new(Vec::new());
-        check_arr.write_chunk(&mut cursor).unwrap();
-        cursor.seek(SeekFrom::Start(0)).unwrap();
-
-        let chunk_size = read_chunk_size(&mut cursor);
-        assert_eq!(
-            cursor.read_to_end(&mut Vec::new()).unwrap(),
-            chunk_size as usize
-        );
+        test_storage_chunk_len(test_ndarray().into());
     }
 
     #[test]

src/chunks/storage/mod.rs

@@ -54,3 +54,40 @@ pub(crate) trait StorageViewMut: Storage {
     /// Get a view of the embedding matrix.
     fn view_mut(&mut self) -> ArrayViewMut2<f32>;
 }
+
+#[cfg(test)]
+mod tests {
+    use std::io::{Cursor, Read, Seek, SeekFrom};
+
+    use crate::chunks::io::WriteChunk;
+    use byteorder::{LittleEndian, ReadBytesExt};
+
+    use crate::storage::StorageWrap;
+
+    fn read_chunk_size(read: &mut impl Read) -> u64 {
+        // Skip identifier.
+        read.read_u32::<LittleEndian>().unwrap();
+
+        // Return chunk length.
+        read.read_u64::<LittleEndian>().unwrap()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn test_storage_chunk_len(check_storage: StorageWrap) {
+        for offset in 0..16u64 {
+            let mut cursor = Cursor::new(Vec::new());
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+            check_storage.write_chunk(&mut cursor).unwrap();
+            cursor.seek(SeekFrom::Start(offset)).unwrap();
+
+            let chunk_size = read_chunk_size(&mut cursor);
+            assert_eq!(
+                cursor.read_to_end(&mut Vec::new()).unwrap(),
+                chunk_size as usize
+            );
+
+            let data = cursor.into_inner();
+            assert_eq!(data.len() as u64 - offset, check_storage.chunk_len(offset));
+        }
+    }
+}