mmap_cache/
builder.rs

1use crate::Error;
2
3use std::fs;
4use std::io;
5use std::path::Path;
6
7/// Serializes a stream of `([u8], [u8])` key-value pairs.
8///
9/// Serialization happens by writing key-value pairs in sorted order. A value is always written before its corresponding key,
10/// because the index will map that key to the starting byte offset of the value that was written.
11///
12/// Many calls of `append_value_bytes` can be made before committing the key --> offset mapping:
13///
14/// ```
15/// # use mmap_cache::Error;
16/// # fn example() -> Result<(), Error> {
17/// use mmap_cache::{Builder, Cache};
18///
19/// let mut index_bytes = Vec::new();
20/// let mut value_bytes = Vec::new();
21///
22/// let mut builder = Builder::new(&mut index_bytes, &mut value_bytes)?;
23///
24/// // Write a value with multiple append calls.
25/// builder.append_value_bytes(&777u32.to_be_bytes())?;
26/// builder.append_value_bytes(&777u32.to_be_bytes())?;
27/// builder.commit_entry(b"hot_garbage")?;
28///
29/// // Or equivalently, use just one insert call.
30/// let mut buf = [0; 8];
31/// buf[0..4].copy_from_slice(&777u32.to_be_bytes());
32/// buf[4..8].copy_from_slice(&777u32.to_be_bytes());
33/// builder.insert(b"lots_of_garbage", &buf)?;
34///
35/// builder.finish()?;
36///
37/// let cache = Cache::new(&index_bytes, &value_bytes)?;
38/// assert_eq!(cache.get_value_offset(b"hot_garbage"), Some(0));
39/// assert_eq!(unsafe { cache.get_transmuted_value(b"hot_garbage") }, Some(&buf));
40/// assert_eq!(cache.get_value_offset(b"lots_of_garbage"), Some(8));
41/// assert_eq!(unsafe { cache.get_transmuted_value(b"lots_of_garbage") }, Some(&buf));
42/// # Ok(())
43/// # }
44/// # example().unwrap();
45/// ```
46pub struct Builder<WK, WV> {
47    map_builder: fst::MapBuilder<WK>,
48    value_writer: WV,
49    value_cursor: usize,
50    committed_value_cursor: usize,
51}
52
53impl<WK, WV> Builder<WK, WV>
54where
55    WK: io::Write,
56    WV: io::Write,
57{
58    /// Creates a new [`Builder`] for serializing a collection of key-value pairs.
59    ///
60    /// - `index_writer`: Writes the serialized [`fst::Map`] which stores the value offsets.
61    /// - `value_writer`: Writes the values pointed to by the byte offsets stored in the [`fst::Map`].
62    ///
63    /// ## Warning
64    ///
65    /// This crate has no control over the alignment guarantees provided by the given writers. Be careful to preserve alignment
66    /// when using [`memmap2`].
67    pub fn new(index_writer: WK, value_writer: WV) -> Result<Self, Error> {
68        Ok(Self {
69            map_builder: fst::MapBuilder::new(index_writer)?,
70            value_writer,
71            committed_value_cursor: 0,
72            value_cursor: 0,
73        })
74    }
75
76    /// Writes `value` into the value stream and commits the entry, storing the value's [`u64`] byte offset along with the `key`
77    /// in the [`fst::Map`].
78    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
79        self.append_value_bytes(value)?;
80        self.commit_entry(key)
81    }
82
83    /// Finishes writing the current value, associating the starting byte offset of the value with `key`.
84    pub fn commit_entry(&mut self, key: &[u8]) -> Result<(), Error> {
85        self.map_builder
86            .insert(key, u64::try_from(self.committed_value_cursor).unwrap())?;
87        self.committed_value_cursor = self.value_cursor;
88        Ok(())
89    }
90
91    /// Writes `value` into the value stream.
92    ///
93    /// The caller may continue appending more value bytes as needed before calling `commit_entry` to finish the current entry
94    /// and start a new one.
95    pub fn append_value_bytes(&mut self, value: &[u8]) -> Result<(), Error> {
96        self.value_writer.write_all(value)?;
97        self.value_cursor += value.len();
98        Ok(())
99    }
100
101    /// Completes the serialization and flushes any outstanding IO.
102    pub fn finish(mut self) -> Result<(), Error> {
103        self.value_writer.flush()?;
104        Ok(self.map_builder.finish()?)
105    }
106}
107
108pub type FileBuilder = Builder<io::BufWriter<fs::File>, io::BufWriter<fs::File>>;
109
110impl FileBuilder {
111    /// Creates a new [`Builder`], using the file at `index_path` for an index writer and the file at `value_path` as a value
112    /// writer.
113    ///
114    /// This always overwrites the given files.
115    ///
116    /// After calling `finish`, these same files can be used with `Cache::map_files`.
117    pub fn create_files(
118        index_path: impl AsRef<Path>,
119        value_path: impl AsRef<Path>,
120    ) -> Result<Self, Error> {
121        let index_writer = io::BufWriter::new(fs::File::create(index_path)?);
122        let value_writer = io::BufWriter::new(fs::File::create(value_path)?);
123        Builder::new(index_writer, value_writer)
124    }
125}