diff --git a/src/lib/rainbowutils.rs b/src/lib/rainbowutils.rs index fe3fae3..19fd6ad 100644 --- a/src/lib/rainbowutils.rs +++ b/src/lib/rainbowutils.rs @@ -4,21 +4,29 @@ use rayon::prelude::*; use std::collections::HashMap; use std::convert::{TryFrom, TryInto}; use std::fs::File; -use std::io::{BufReader, Read}; +use std::io::{BufReader, BufWriter, Read}; use std::io::{Error, ErrorKind}; +pub const LZMA: &str = "lzma"; + pub const BDF_HDR: &[u8; 11] = b"BDF\x01RAINBOW"; pub const NULL_BYTES: &[u8; 4] = &[0u8; 4]; pub const META_CHUNK_NAME: &str = "META"; pub const HTBL_CHUNK_NAME: &str = "HTBL"; pub const DTBL_CHUNK_NAME: &str = "DTBL"; -pub struct BinaryDictionaryFile { +pub struct BDFReader { reader: BufReader, metadata: Option, lookup_table: Option, } +pub struct BDFWriter { + writer: BufWriter, + metadata: MetaChunk, + lookup_table: HashLookupTable, +} + #[derive(Debug, Clone)] pub struct GenericChunk { length: u32, @@ -53,7 +61,17 @@ pub struct DataEntry { hashes: HashMap>, } -impl BinaryDictionaryFile { +impl BDFWriter { + pub fn new(writer: BufWriter, compress: bool) -> Self { + Self { + metadata: MetaChunk::new(0, 0, compress), + lookup_table: HashLookupTable::new(HashMap::new()), + writer, + } + } +} + +impl BDFReader { pub fn new(reader: BufReader) -> Self { Self { metadata: None, @@ -100,6 +118,7 @@ impl BinaryDictionaryFile { } } + /// Validates the header of the file fn validate_header(&mut self) -> bool { let mut header = [0u8; 11]; let _ = self.reader.read(&mut header); @@ -131,6 +150,7 @@ impl BinaryDictionaryFile { } impl GenericChunk { + /// Serializes the chunk to a vector of bytes pub fn serialize(&mut self) -> Vec { let mut serialized: Vec = Vec::new(); let mut length_raw = [0u8; 4]; @@ -238,6 +258,24 @@ impl From for GenericChunk { } impl MetaChunk { + /// Creates a new meta chunk + pub fn new(entry_count: u32, entries_per_chunk: u32, compress: bool) -> Self { + let compression_method = if compress { + Some(LZMA.to_string()) + } else { + None + }; + let chunk_count = (entry_count as f32 / entries_per_chunk as f32).ceil() as u32; + + Self { + chunk_count, + entry_count, + entries_per_chunk, + compression_method, + } + } + + /// Serializes the chunk into bytes pub fn serialize(&self) -> Vec { let mut serialized_data: Vec = Vec::new(); let mut chunk_count_raw = [0u8; 4]; @@ -299,6 +337,11 @@ impl TryFrom for MetaChunk { } impl HashLookupTable { + pub fn new(entries: HashMap) -> Self { + Self { entries } + } + + /// Returns an entry by the name of the hash function pub fn get_entry(&self, name: &String) -> Option<(&u32, &HashEntry)> { self.entries.iter().find(|(_, entry)| entry.name == *name) }