use crate::file::{index::Entry, Index}; mod write_chunk { use std::collections::VecDeque; use crate::file::index; /// A [`Write`][std::io::Write] implementation that validates chunk sizes while allowing the user to know /// which chunk is to be written next. pub struct Chunk { chunks_to_write: VecDeque, inner: W, next_chunk: Option, written_bytes: usize, } impl Chunk where W: std::io::Write, { pub(crate) fn new(out: W, chunks: VecDeque) -> Chunk where W: std::io::Write, { Chunk { chunks_to_write: chunks, inner: out, next_chunk: None, written_bytes: 0, } } } impl std::io::Write for Chunk where W: std::io::Write, { fn write(&mut self, buf: &[u8]) -> std::io::Result { let written = self.inner.write(buf)?; self.written_bytes += written; Ok(written) } fn flush(&mut self) -> std::io::Result<()> { self.inner.flush() } } impl Chunk { /// Return the inner writer - should only be called once there is no more chunk to write. pub fn into_inner(self) -> W { self.inner } /// Return the next chunk-id to write, if there is one. pub fn next_chunk(&mut self) -> Option { if let Some(entry) = self.next_chunk.take() { assert_eq!( entry.offset.end, self.written_bytes as u64, "BUG: expected to write {} bytes, but only wrote {} for chunk {:?}", entry.offset.end, self.written_bytes, std::str::from_utf8(&entry.kind) ) } self.written_bytes = 0; self.next_chunk = self.chunks_to_write.pop_front(); self.next_chunk.as_ref().map(|e| e.kind) } } } pub use write_chunk::Chunk; /// Writing impl Index { /// Create a new index whose sole purpose is to be receiving chunks using [`plan_chunk()`][Index::plan_chunk()] and to be written to /// an output using [`into_write()`][Index::into_write()] pub fn for_writing() -> Self { Index { will_write: true, chunks: Vec::new(), } } /// Plan to write a new chunk as part of the index when [`into_write()`][Index::into_write()] is called. pub fn plan_chunk(&mut self, chunk: crate::Id, exact_size_on_disk: u64) { assert!(self.will_write, "BUG: create the index with `for_writing()`"); assert!( !self.chunks.iter().any(|e| e.kind == chunk), "BUG: must not add chunk of same kind twice: {:?}", std::str::from_utf8(&chunk) ); self.chunks.push(Entry { kind: chunk, offset: 0..exact_size_on_disk, }) } /// Return the total size of all planned chunks thus far. pub fn planned_storage_size(&self) -> u64 { assert!(self.will_write, "BUG: create the index with `for_writing()`"); self.chunks.iter().map(|e| e.offset.end).sum() } /// Return the amount of chunks we currently know. pub fn num_chunks(&self) -> usize { self.chunks.len() } /// After [planning all chunks][Index::plan_chunk()] call this method with the destination to write the chunks to. /// Use the [Chunk] writer to write each chunk in order. /// `current_offset` is the byte position at which `out` will continue writing. pub fn into_write(self, mut out: W, current_offset: usize) -> std::io::Result> where W: std::io::Write, { assert!( self.will_write, "BUG: create the index with `for_writing()`, cannot write decoded indices" ); // First chunk starts past the table of contents let mut current_offset = (current_offset + Self::size_for_entries(self.num_chunks())) as u64; for entry in &self.chunks { out.write_all(&entry.kind)?; out.write_all(¤t_offset.to_be_bytes())?; current_offset += entry.offset.end; } // sentinel to mark end of chunks out.write_all(&0u32.to_be_bytes())?; out.write_all(¤t_offset.to_be_bytes())?; Ok(Chunk::new(out, self.chunks.into())) } }