ic_stable_structures/
writer.rs

1use crate::{safe_write, GrowFailed, Memory};
2use std::io;
3
4#[cfg(test)]
5mod tests;
6
7/// A writer that writes sequentially to memory.
8///
9/// Warning: This will overwrite any existing data in stable memory as it writes, so ensure you set
10/// the `offset` value accordingly if you wish to preserve existing data.
11///
12/// The writer Will attempt to grow the memory as it writes.
13pub struct Writer<'a, M> {
14    /// The offset of the next write.
15    offset: u64,
16
17    /// The stable memory to write data to.
18    memory: &'a mut M,
19}
20
21impl<'a, M: Memory> Writer<'a, M> {
22    /// Creates a new `Writer` which writes to the selected memory starting from the offset.
23    pub fn new(memory: &'a mut M, offset: u64) -> Self {
24        Self { offset, memory }
25    }
26
27    /// Writes a byte slice to the underlying memory directly.
28    ///
29    /// Note: The writer will first attempt to grow the memory enough to write _all_ the data and
30    /// only then start writing. If the memory can not be grown sufficiently, the write is aborted
31    /// without writing any data.
32    pub fn write(&mut self, buf: &[u8]) -> Result<(), GrowFailed> {
33        safe_write(self.memory, self.offset, buf)?;
34        self.offset += buf.len() as u64;
35        Ok(())
36    }
37}
38
39impl<M: Memory> io::Write for Writer<'_, M> {
40    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
41        self.write(buf)
42            .map_err(|e| io::Error::new(io::ErrorKind::OutOfMemory, e))?;
43        Ok(buf.len())
44    }
45
46    fn flush(&mut self) -> Result<(), io::Error> {
47        // Noop.
48        Ok(())
49    }
50}
51
52/// A writer to the stable memory which first writes the bytes to an in memory buffer and flushes
53/// the buffer to stable memory each time it becomes full.
54///
55/// Warning: This will overwrite any existing data in stable memory as it writes, so ensure you set
56/// the `offset` value accordingly if you wish to preserve existing data.
57///
58/// Note: Each call to grow or write to stable memory is a relatively expensive operation, so pick a
59/// buffer size large enough to avoid excessive calls to stable memory.
60pub struct BufferedWriter<'a, M: Memory> {
61    inner: io::BufWriter<Writer<'a, M>>,
62}
63
64impl<M: Memory> BufferedWriter<'_, M> {
65    /// Creates a new `BufferedStableWriter` which writes to the selected memory
66    pub fn new(buffer_size: usize, writer: Writer<M>) -> BufferedWriter<M> {
67        BufferedWriter {
68            inner: io::BufWriter::with_capacity(buffer_size, writer),
69        }
70    }
71}
72
73impl<M: Memory> io::Write for BufferedWriter<'_, M> {
74    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
75        self.inner.write(buf)
76    }
77
78    fn flush(&mut self) -> io::Result<()> {
79        self.inner.flush()
80    }
81}