| 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<W> { |
| chunks_to_write: VecDeque<index::Entry>, |
| inner: W, |
| next_chunk: Option<index::Entry>, |
| written_bytes: usize, |
| } |
| |
| impl<W> Chunk<W> |
| where |
| W: std::io::Write, |
| { |
| pub(crate) fn new(out: W, chunks: VecDeque<index::Entry>) -> Chunk<W> |
| where |
| W: std::io::Write, |
| { |
| Chunk { |
| chunks_to_write: chunks, |
| inner: out, |
| next_chunk: None, |
| written_bytes: 0, |
| } |
| } |
| } |
| |
| impl<W> std::io::Write for Chunk<W> |
| where |
| W: std::io::Write, |
| { |
| fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { |
| let written = self.inner.write(buf)?; |
| self.written_bytes += written; |
| Ok(written) |
| } |
| |
| fn flush(&mut self) -> std::io::Result<()> { |
| self.inner.flush() |
| } |
| } |
| |
| impl<W> Chunk<W> { |
| /// 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<crate::Id> { |
| 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<W>(self, mut out: W, current_offset: usize) -> std::io::Result<Chunk<W>> |
| 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())) |
| } |
| } |