Cleanup top-level modules

- All chunk modules become submodules of `chunk`.
- Modules for reading foreign formats become submodules of `compat`.

Author: danieldk

src/chunks/io.rs

@@ -0,0 +1,263 @@
+use std::fmt::{self, Display};
+use std::fs::File;
+use std::io::{BufReader, Read, Seek, Write};
+
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+
+use crate::io::{Error, ErrorKind, Result};
+
+const MODEL_VERSION: u32 = 0;
+
+const MAGIC: [u8; 4] = [b'F', b'i', b'F', b'u'];
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[repr(u32)]
+pub enum ChunkIdentifier {
+    Header = 0,
+    SimpleVocab = 1,
+    NdArray = 2,
+    FinalfusionSubwordVocab = 3,
+    QuantizedArray = 4,
+    Metadata = 5,
+    NdNorms = 6,
+    FastTextSubwordVocab = 7,
+}
+
+impl ChunkIdentifier {
+    pub fn try_from(identifier: u32) -> Option<Self> {
+        use self::ChunkIdentifier::*;
+
+        match identifier {
+            1 => Some(SimpleVocab),
+            2 => Some(NdArray),
+            3 => Some(FinalfusionSubwordVocab),
+            4 => Some(QuantizedArray),
+            5 => Some(Metadata),
+            6 => Some(NdNorms),
+            7 => Some(FastTextSubwordVocab),
+            _ => None,
+        }
+    }
+
+    /// Read and ensure that the chunk has the given identifier.
+    pub fn ensure_chunk_type<R>(read: &mut R, identifier: ChunkIdentifier) -> Result<()>
+    where
+        R: Read,
+    {
+        let chunk_id = read
+            .read_u32::<LittleEndian>()
+            .map_err(|e| ErrorKind::io_error("Cannot read chunk identifier", e))?;
+        let chunk_id = ChunkIdentifier::try_from(chunk_id)
+            .ok_or_else(|| ErrorKind::Format(format!("Unknown chunk identifier: {}", chunk_id)))
+            .map_err(Error::from)?;
+        if chunk_id != identifier {
+            return Err(ErrorKind::Format(format!(
+                "Invalid chunk identifier, expected: {}, got: {}",
+                identifier, chunk_id
+            ))
+            .into());
+        }
+
+        Ok(())
+    }
+}
+
+impl Display for ChunkIdentifier {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use self::ChunkIdentifier::*;
+
+        match self {
+            Header => write!(f, "Header"),
+            SimpleVocab => write!(f, "SimpleVocab"),
+            NdArray => write!(f, "NdArray"),
+            FastTextSubwordVocab => write!(f, "FastTextSubwordVocab"),
+            FinalfusionSubwordVocab => write!(f, "FinalfusionSubwordVocab"),
+            QuantizedArray => write!(f, "QuantizedArray"),
+            Metadata => write!(f, "Metadata"),
+            NdNorms => write!(f, "NdNorms"),
+        }
+    }
+}
+
+pub trait TypeId {
+    /// Read and ensure that the data type is equal to `Self`.
+    fn ensure_data_type<R>(read: &mut R) -> Result<()>
+    where
+        R: Read;
+
+    fn type_id() -> u32;
+}
+
+macro_rules! typeid_impl {
+    ($type:ty, $id:expr) => {
+        impl TypeId for $type {
+            fn ensure_data_type<R>(read: &mut R) -> Result<()>
+            where
+                R: Read,
+            {
+                let type_id = read
+                    .read_u32::<LittleEndian>()
+                    .map_err(|e| ErrorKind::io_error("Cannot read type identifier", e))?;
+                if type_id != Self::type_id() {
+                    return Err(ErrorKind::Format(format!(
+                        "Invalid type, expected: {}, got: {}",
+                        Self::type_id(),
+                        type_id
+                    ))
+                    .into());
+                }
+
+                Ok(())
+            }
+
+            fn type_id() -> u32 {
+                $id
+            }
+        }
+    };
+}
+
+typeid_impl!(f32, 10);
+typeid_impl!(u8, 1);
+
+pub trait ReadChunk
+where
+    Self: Sized,
+{
+    fn read_chunk<R>(read: &mut R) -> Result<Self>
+    where
+        R: Read + Seek;
+}
+
+/// Memory-mappable chunks.
+pub trait MmapChunk
+where
+    Self: Sized,
+{
+    /// Memory map a chunk.
+    ///
+    /// The given `File` object should be positioned at the start of the chunk.
+    fn mmap_chunk(read: &mut BufReader<File>) -> Result<Self>;
+}
+
+pub trait WriteChunk {
+    /// Get the identifier of a chunk.
+    fn chunk_identifier(&self) -> ChunkIdentifier;
+
+    fn write_chunk<W>(&self, write: &mut W) -> Result<()>
+    where
+        W: Write + Seek;
+}
+
+#[derive(Debug, Eq, PartialEq)]
+pub(crate) struct Header {
+    chunk_identifiers: Vec<ChunkIdentifier>,
+}
+
+impl Header {
+    pub fn new(chunk_identifiers: impl Into<Vec<ChunkIdentifier>>) -> Self {
+        Header {
+            chunk_identifiers: chunk_identifiers.into(),
+        }
+    }
+
+    pub fn chunk_identifiers(&self) -> &[ChunkIdentifier] {
+        &self.chunk_identifiers
+    }
+}
+
+impl WriteChunk for Header {
+    fn chunk_identifier(&self) -> ChunkIdentifier {
+        ChunkIdentifier::Header
+    }
+
+    fn write_chunk<W>(&self, write: &mut W) -> Result<()>
+    where
+        W: Write + Seek,
+    {
+        write
+            .write_all(&MAGIC)
+            .map_err(|e| ErrorKind::io_error("Cannot write magic", e))?;
+        write
+            .write_u32::<LittleEndian>(MODEL_VERSION)
+            .map_err(|e| ErrorKind::io_error("Cannot write model version", e))?;
+        write
+            .write_u32::<LittleEndian>(self.chunk_identifiers.len() as u32)
+            .map_err(|e| ErrorKind::io_error("Cannot write chunk identifiers length", e))?;
+
+        for &identifier in &self.chunk_identifiers {
+            write
+                .write_u32::<LittleEndian>(identifier as u32)
+                .map_err(|e| ErrorKind::io_error("Cannot write chunk identifier", e))?;
+        }
+
+        Ok(())
+    }
+}
+
+impl ReadChunk for Header {
+    fn read_chunk<R>(read: &mut R) -> Result<Self>
+    where
+        R: Read + Seek,
+    {
+        // Magic and version ceremony.
+        let mut magic = [0u8; 4];
+        read.read_exact(&mut magic)
+            .map_err(|e| ErrorKind::io_error("Cannot read magic", e))?;
+
+        if magic != MAGIC {
+            return Err(ErrorKind::Format(format!(
+                "Expected 'FiFu' as magic, got: {}",
+                String::from_utf8_lossy(&magic).into_owned()
+            ))
+            .into());
+        }
+
+        let version = read
+            .read_u32::<LittleEndian>()
+            .map_err(|e| ErrorKind::io_error("Cannot read model version", e))?;
+        if version != MODEL_VERSION {
+            return Err(
+                ErrorKind::Format(format!("Unknown finalfusion version: {}", version)).into(),
+            );
+        }
+
+        // Read chunk identifiers.
+        let chunk_identifiers_len = read
+            .read_u32::<LittleEndian>()
+            .map_err(|e| ErrorKind::io_error("Cannot read chunk identifiers length", e))?
+            as usize;
+        let mut chunk_identifiers = Vec::with_capacity(chunk_identifiers_len);
+        for _ in 0..chunk_identifiers_len {
+            let identifier = read
+                .read_u32::<LittleEndian>()
+                .map_err(|e| ErrorKind::io_error("Cannot read chunk identifier", e))?;
+            let chunk_identifier = ChunkIdentifier::try_from(identifier)
+                .ok_or_else(|| {
+                    ErrorKind::Format(format!("Unknown chunk identifier: {}", identifier))
+                })
+                .map_err(Error::from)?;
+            chunk_identifiers.push(chunk_identifier);
+        }
+
+        Ok(Header { chunk_identifiers })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::{Cursor, Seek, SeekFrom};
+
+    use super::{ChunkIdentifier, Header, ReadChunk, WriteChunk};
+
+    #[test]
+    fn header_write_read_roundtrip() {
+        let check_header =
+            Header::new(vec![ChunkIdentifier::SimpleVocab, ChunkIdentifier::NdArray]);
+        let mut cursor = Cursor::new(Vec::new());
+        check_header.write_chunk(&mut cursor).unwrap();
+        cursor.seek(SeekFrom::Start(0)).unwrap();
+        let header = Header::read_chunk(&mut cursor).unwrap();
+        assert_eq!(header, check_header);
+    }
+}

src/metadata.rs renamed to src/chunks/metadata.rs

@@ -5,10 +5,8 @@ use std::io::{Read, Seek, Write};
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use toml::Value;
 
-use crate::io::{
-    private::{ChunkIdentifier, Header, ReadChunk, WriteChunk},
-    Error, ErrorKind, ReadMetadata, Result,
-};
+use super::io::{ChunkIdentifier, Header, ReadChunk, WriteChunk};
+use crate::io::{Error, ErrorKind, ReadMetadata, Result};
 
 /// Embeddings metadata.
 ///
@@ -101,7 +99,7 @@ mod tests {
     use toml::toml;
 
     use super::Metadata;
-    use crate::io::private::{ReadChunk, WriteChunk};
+    use crate::chunks::io::{ReadChunk, WriteChunk};
 
     fn read_chunk_size(read: &mut impl Read) -> u64 {
         // Skip identifier.

src/chunks/mod.rs

@@ -0,0 +1,11 @@
+//! finalfusion chunks
+
+pub(crate) mod io;
+
+pub mod metadata;
+
+pub mod norms;
+
+pub mod storage;
+
+pub mod vocab;

src/norms.rs renamed to src/chunks/norms.rs

@@ -6,7 +6,7 @@ use std::mem::size_of;
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use ndarray::Array1;
 
-use crate::io::private::{ChunkIdentifier, ReadChunk, TypeId, WriteChunk};
+use super::io::{ChunkIdentifier, ReadChunk, TypeId, WriteChunk};
 use crate::io::{ErrorKind, Result};
 use crate::util::padding;
 
@@ -117,8 +117,8 @@ mod tests {
     use byteorder::{LittleEndian, ReadBytesExt};
     use ndarray::Array1;
 
-    use crate::io::private::{ReadChunk, WriteChunk};
-    use crate::norms::NdNorms;
+    use super::NdNorms;
+    use crate::chunks::io::{ReadChunk, WriteChunk};
 
     const LEN: usize = 100;
 

src/storage.rs renamed to src/chunks/storage.rs

@@ -11,7 +11,7 @@ use rand::{FromEntropy, Rng};
 use rand_xorshift::XorShiftRng;
 use reductive::pq::{QuantizeVector, ReconstructVector, TrainPQ, PQ};
 
-use crate::io::private::{ChunkIdentifier, MmapChunk, ReadChunk, TypeId, WriteChunk};
+use super::io::{ChunkIdentifier, MmapChunk, ReadChunk, TypeId, WriteChunk};
 use crate::io::{Error, ErrorKind, Result};
 use crate::util::padding;
 
@@ -936,8 +936,8 @@ mod tests {
     use ndarray::Array2;
     use reductive::pq::PQ;
 
-    use crate::io::private::{ReadChunk, WriteChunk};
-    use crate::storage::{NdArray, Quantize, QuantizedArray, StorageView};
+    use crate::chunks::io::{ReadChunk, WriteChunk};
+    use crate::chunks::storage::{NdArray, Quantize, QuantizedArray, StorageView};
 
     const N_ROWS: usize = 100;
     const N_COLS: usize = 100;

src/vocab.rs renamed to src/chunks/vocab.rs

@@ -6,8 +6,8 @@ use std::mem::size_of;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
-use crate::fasttext::FastTextIndexer;
-use crate::io::private::{ChunkIdentifier, ReadChunk, WriteChunk};
+use super::io::{ChunkIdentifier, ReadChunk, WriteChunk};
+use crate::compat::fasttext::FastTextIndexer;
 use crate::io::{Error, ErrorKind, Result};
 use crate::subword::{BucketIndexer, FinalfusionHashIndexer, Indexer, SubwordIndices};
 
@@ -521,8 +521,8 @@ mod tests {
     use byteorder::{LittleEndian, ReadBytesExt};
 
     use super::{FastTextSubwordVocab, FinalfusionSubwordVocab, SimpleVocab, SubwordVocab};
-    use crate::fasttext::FastTextIndexer;
-    use crate::io::private::{ReadChunk, WriteChunk};
+    use crate::chunks::io::{ReadChunk, WriteChunk};
+    use crate::compat::fasttext::FastTextIndexer;
     use crate::subword::{BucketIndexer, FinalfusionHashIndexer};
 
     fn test_fasttext_subword_vocab() -> FastTextSubwordVocab {

src/fasttext/indexer.rs renamed to src/compat/fasttext/indexer.rs

@@ -5,14 +5,14 @@ use ndarray::{s, Array2, ErrorKind as ShapeErrorKind, ShapeError};
 use serde::Serialize;
 use toml::Value;
 
+use crate::chunks::metadata::Metadata;
+use crate::chunks::norms::NdNorms;
+use crate::chunks::storage::{NdArray, Storage, StorageViewMut};
+use crate::chunks::vocab::{FastTextSubwordVocab, Vocab};
 use crate::embeddings::Embeddings;
 use crate::io::{Error, ErrorKind, Result};
-use crate::metadata::Metadata;
-use crate::norms::NdNorms;
-use crate::storage::{NdArray, Storage, StorageViewMut};
 use crate::subword::BucketIndexer;
 use crate::util::{l2_normalize_array, read_string};
-use crate::vocab::{FastTextSubwordVocab, Vocab};
 
 use super::FastTextIndexer;
 

src/fasttext/io.rs renamed to src/compat/fasttext/io.rs

@@ -0,0 +1,7 @@
+//! Readers/writers for other embedding formats.
+
+pub mod fasttext;
+
+pub mod text;
+
+pub mod word2vec;