bao_tree/io/
outboard.rs

1//! Implementations of the outboard traits
2//!
3//! A number of implementations for the sync and async outboard traits are provided.
4//! Implementations for in-memory outboards, for outboards where the data resides on disk,
5//! and a special implementation [EmptyOutboard] that just ignores all writes.
6use std::io;
7
8use crate::{blake3, BaoTree, BlockSize, TreeNode};
9
10/// An empty outboard, that just returns 0 hashes for all nodes.
11///
12/// Also allows you to write and will immediately discard the data, a bit like /dev/null
13#[derive(Debug)]
14pub struct EmptyOutboard {
15    /// tree defining the geometry
16    pub tree: BaoTree,
17    /// root hash
18    pub root: blake3::Hash,
19}
20
21impl crate::io::sync::Outboard for EmptyOutboard {
22    fn root(&self) -> blake3::Hash {
23        self.root
24    }
25    fn tree(&self) -> BaoTree {
26        self.tree
27    }
28    fn load(&self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
29        Ok(if self.tree.is_relevant_for_outboard(node) {
30            // behave as if it was an outboard file filled with 0s
31            Some((blake3::Hash::from([0; 32]), blake3::Hash::from([0; 32])))
32        } else {
33            None
34        })
35    }
36}
37
38#[cfg(feature = "tokio_fsm")]
39impl crate::io::fsm::Outboard for EmptyOutboard {
40    fn root(&self) -> blake3::Hash {
41        self.root
42    }
43    fn tree(&self) -> BaoTree {
44        self.tree
45    }
46    async fn load(&mut self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
47        Ok(if self.tree.is_relevant_for_outboard(node) {
48            // behave as if it was an outboard file filled with 0s
49            Some((blake3::Hash::from([0; 32]), blake3::Hash::from([0; 32])))
50        } else {
51            None
52        })
53    }
54}
55
56impl crate::io::sync::OutboardMut for EmptyOutboard {
57    fn save(&mut self, node: TreeNode, _pair: &(blake3::Hash, blake3::Hash)) -> io::Result<()> {
58        if self.tree.is_relevant_for_outboard(node) {
59            Ok(())
60        } else {
61            Err(io::Error::new(
62                io::ErrorKind::InvalidInput,
63                "invalid node for this outboard",
64            ))
65        }
66    }
67
68    fn sync(&mut self) -> io::Result<()> {
69        Ok(())
70    }
71}
72
73#[cfg(feature = "tokio_fsm")]
74impl crate::io::fsm::OutboardMut for EmptyOutboard {
75    async fn save(
76        &mut self,
77        node: TreeNode,
78        _pair: &(blake3::Hash, blake3::Hash),
79    ) -> io::Result<()> {
80        if self.tree.is_relevant_for_outboard(node) {
81            Ok(())
82        } else {
83            Err(io::Error::new(
84                io::ErrorKind::InvalidInput,
85                "invalid node for this outboard",
86            ))
87        }
88    }
89
90    async fn sync(&mut self) -> io::Result<()> {
91        Ok(())
92    }
93}
94
95/// A generic outboard in pre order
96///
97/// All fields are public, so an outboard does not enforce any invariants.
98/// This is necessary since we want outboards to exist in an incomplete state
99/// where data does not match the root hash.
100///
101/// Caution: unlike the outboard implementation in the bao crate, this
102/// implementation does not assume an 8 byte size prefix.
103#[derive(Debug, Clone)]
104pub struct PreOrderOutboard<D = Vec<u8>> {
105    /// root hash
106    pub root: blake3::Hash,
107    /// tree defining the data
108    pub tree: BaoTree,
109    /// hashes with length prefix
110    pub data: D,
111}
112
113impl<R: Default> Default for PreOrderOutboard<R> {
114    fn default() -> Self {
115        Self {
116            root: blake3::hash(&[]),
117            tree: BaoTree::new(0, BlockSize::ZERO),
118            data: Default::default(),
119        }
120    }
121}
122
123/// A generic outboard in post order
124///
125/// All fields are public, so an outboard does not enforce any invariants.
126/// This is necessary since we want outboards to exist in an incomplete state
127/// where data does not match the root hash.
128#[derive(Debug, Clone)]
129pub struct PostOrderOutboard<D = Vec<u8>> {
130    /// root hash
131    pub root: blake3::Hash,
132    /// tree defining the data
133    pub tree: BaoTree,
134    /// hashes with length prefix
135    pub data: D,
136}
137
138impl<D: Default> Default for PostOrderOutboard<D> {
139    fn default() -> Self {
140        Self {
141            root: blake3::hash(&[]),
142            tree: BaoTree::new(0, BlockSize::ZERO),
143            data: Default::default(),
144        }
145    }
146}
147
148/// A post order outboard that is optimized for memory storage.
149///
150/// The [PostOrderOutboard] will work just fine using e.g. a [`Vec<u8>`] as storage.
151/// However, it will not work for types such as [bytes::Bytes] or [bytes::BytesMut]
152/// that do not implement the io traits.
153///
154/// The traits are implemented for fixed size slices or mutable slices, so
155/// unlike the [PostOrderOutboard], you must make sure that the data is already
156/// the right size.
157#[derive(Debug, Clone, PartialEq, Eq)]
158pub struct PostOrderMemOutboard<T = Vec<u8>> {
159    /// root hash
160    pub root: blake3::Hash,
161    /// tree defining the data
162    pub tree: BaoTree,
163    /// hashes without length suffix
164    pub data: T,
165}
166
167impl<T: Default> Default for PostOrderMemOutboard<T> {
168    fn default() -> Self {
169        Self {
170            root: blake3::hash(&[]),
171            tree: BaoTree::new(0, BlockSize::ZERO),
172            data: Default::default(),
173        }
174    }
175}
176
177impl PostOrderMemOutboard {
178    /// Create a new outboard from `data` and a `block_size`.
179    ///
180    /// This will hash the data and create an outboard.
181    ///
182    /// It is just a shortcut that calls [crate::io::sync::outboard_post_order].
183    pub fn create(data: impl AsRef<[u8]>, block_size: BlockSize) -> Self {
184        let data = data.as_ref();
185        let size = data.len() as u64;
186        let tree = BaoTree::new(size, block_size);
187        let mut outboard = Vec::with_capacity(tree.outboard_size().try_into().unwrap());
188        let root = crate::io::sync::outboard_post_order(data, tree, &mut outboard).unwrap();
189        Self {
190            root,
191            tree,
192            data: outboard,
193        }
194    }
195
196    /// returns the outboard data, with the length suffix.
197    pub fn into_inner_with_suffix(self) -> Vec<u8> {
198        let mut res = self.data;
199        res.extend_from_slice(self.tree.size.to_le_bytes().as_slice());
200        res
201    }
202}
203
204impl<T> PostOrderMemOutboard<T> {
205    /// Map the outboard data to a new type.
206    pub fn map_data<F, U>(self, f: F) -> PostOrderMemOutboard<U>
207    where
208        F: FnOnce(T) -> U,
209        U: AsRef<[u8]>,
210    {
211        PostOrderMemOutboard {
212            root: self.root,
213            tree: self.tree,
214            data: f(self.data),
215        }
216    }
217
218    /// Flip the outboard to pre order.
219    pub fn flip(&self) -> PreOrderMemOutboard
220    where
221        T: AsRef<[u8]>,
222    {
223        let mut target = PreOrderMemOutboard {
224            root: self.root,
225            tree: self.tree,
226            data: vec![0; self.tree.outboard_size().try_into().unwrap()],
227        };
228        crate::io::sync::copy(self, &mut target).unwrap();
229        target
230    }
231}
232
233impl<T: AsRef<[u8]>> crate::io::sync::Outboard for PostOrderMemOutboard<T> {
234    fn root(&self) -> blake3::Hash {
235        self.root
236    }
237    fn tree(&self) -> BaoTree {
238        self.tree
239    }
240    fn load(&self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
241        Ok(load_post(&self.tree, self.data.as_ref(), node))
242    }
243}
244
245#[cfg(feature = "tokio_fsm")]
246impl<T: AsRef<[u8]>> crate::io::fsm::Outboard for PostOrderMemOutboard<T> {
247    fn root(&self) -> blake3::Hash {
248        self.root
249    }
250    fn tree(&self) -> BaoTree {
251        self.tree
252    }
253    async fn load(&mut self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
254        Ok(load_post(&self.tree, self.data.as_ref(), node))
255    }
256}
257
258impl<T: AsMut<[u8]>> crate::io::sync::OutboardMut for PostOrderMemOutboard<T> {
259    fn save(&mut self, node: TreeNode, pair: &(blake3::Hash, blake3::Hash)) -> io::Result<()> {
260        match self.tree.post_order_offset(node) {
261            Some(offset) => {
262                let offset = usize::try_from(offset.value() * 64).unwrap();
263                let data = self.data.as_mut();
264                data[offset..offset + 32].copy_from_slice(pair.0.as_bytes());
265                data[offset + 32..offset + 64].copy_from_slice(pair.1.as_bytes());
266                Ok(())
267            }
268            None => Err(io::Error::new(
269                io::ErrorKind::InvalidInput,
270                "invalid node for this outboard",
271            )),
272        }
273    }
274
275    fn sync(&mut self) -> io::Result<()> {
276        Ok(())
277    }
278}
279
280#[cfg(feature = "tokio_fsm")]
281impl<T: AsMut<[u8]>> crate::io::fsm::OutboardMut for PostOrderMemOutboard<T> {
282    async fn save(
283        &mut self,
284        node: TreeNode,
285        pair: &(blake3::Hash, blake3::Hash),
286    ) -> io::Result<()> {
287        match self.tree.post_order_offset(node) {
288            Some(offset) => {
289                let offset = usize::try_from(offset.value() * 64).unwrap();
290                let data = self.data.as_mut();
291                data[offset..offset + 32].copy_from_slice(pair.0.as_bytes());
292                data[offset + 32..offset + 64].copy_from_slice(pair.1.as_bytes());
293                Ok(())
294            }
295            None => Err(io::Error::new(
296                io::ErrorKind::InvalidInput,
297                "invalid node for this outboard",
298            )),
299        }
300    }
301
302    async fn sync(&mut self) -> io::Result<()> {
303        Ok(())
304    }
305}
306
307fn load_raw_post_mem(tree: &BaoTree, data: &[u8], node: TreeNode) -> Option<[u8; 64]> {
308    let offset = tree.post_order_offset(node)?.value();
309    let offset = usize::try_from(offset * 64).unwrap();
310    let slice = &data[offset..offset + 64];
311    Some(slice.try_into().unwrap())
312}
313
314fn load_post(tree: &BaoTree, data: &[u8], node: TreeNode) -> Option<(blake3::Hash, blake3::Hash)> {
315    load_raw_post_mem(tree, data, node).map(parse_hash_pair)
316}
317
318/// A pre order outboard that is optimized for memory storage.
319///
320/// The traits are implemented for fixed size slices or mutable slices, so you
321/// must make sure that the data is already the right size.
322#[derive(Debug, Clone, PartialEq, Eq)]
323pub struct PreOrderMemOutboard<T = Vec<u8>> {
324    /// root hash
325    pub root: blake3::Hash,
326    /// tree defining the data
327    pub tree: BaoTree,
328    /// hashes with length prefix
329    pub data: T,
330}
331
332impl<T: Default> Default for PreOrderMemOutboard<T> {
333    fn default() -> Self {
334        Self {
335            root: blake3::hash(&[]),
336            tree: BaoTree::new(0, BlockSize::ZERO),
337            data: Default::default(),
338        }
339    }
340}
341
342impl PreOrderMemOutboard {
343    /// returns the outboard data, with the length prefix added.
344    pub fn into_inner_with_prefix(self) -> Vec<u8> {
345        let mut res = self.data;
346        res.splice(0..0, self.tree.size.to_le_bytes());
347        res
348    }
349
350    /// Create a new outboard from `data` and a `block_size`.
351    ///
352    /// This will hash the data and create an outboard
353    pub fn create(data: impl AsRef<[u8]>, block_size: BlockSize) -> Self {
354        let data = data.as_ref();
355        let size = data.len() as u64;
356        let tree = BaoTree::new(size, block_size);
357        // the outboard impl for PreOrderMemOutboard requires just AsMut<[u8]>,
358        // so the data must already be the right size.
359        let outboard = vec![0u8; tree.outboard_size().try_into().unwrap()];
360        let mut res = Self {
361            root: blake3::Hash::from([0; 32]),
362            tree,
363            data: outboard,
364        };
365        let root = crate::io::sync::outboard(data, tree, &mut res).unwrap();
366        res.root = root;
367        res
368    }
369}
370
371impl<T> PreOrderMemOutboard<T> {
372    /// Map the outboard data to a new type.
373    pub fn map_data<F, U>(self, f: F) -> PreOrderMemOutboard<U>
374    where
375        F: FnOnce(T) -> U,
376        U: AsRef<[u8]>,
377    {
378        PreOrderMemOutboard {
379            root: self.root,
380            tree: self.tree,
381            data: f(self.data),
382        }
383    }
384
385    /// Flip the outboard to a post order outboard.
386    pub fn flip(&self) -> PostOrderMemOutboard
387    where
388        T: AsRef<[u8]>,
389    {
390        let mut target = PostOrderMemOutboard {
391            root: self.root,
392            tree: self.tree,
393            data: vec![0; self.tree.outboard_size().try_into().unwrap()],
394        };
395        crate::io::sync::copy(self, &mut target).unwrap();
396        target
397    }
398}
399
400impl<T: AsRef<[u8]>> crate::io::sync::Outboard for PreOrderMemOutboard<T> {
401    fn root(&self) -> blake3::Hash {
402        self.root
403    }
404    fn tree(&self) -> BaoTree {
405        self.tree
406    }
407    fn load(&self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
408        Ok(load_pre(&self.tree, self.data.as_ref(), node))
409    }
410}
411
412impl<T: AsMut<[u8]>> crate::io::sync::OutboardMut for PreOrderMemOutboard<T> {
413    fn save(&mut self, node: TreeNode, pair: &(blake3::Hash, blake3::Hash)) -> io::Result<()> {
414        match self.tree.pre_order_offset(node) {
415            Some(offset) => {
416                let offset_u64 = offset * 64;
417                let offset = usize::try_from(offset_u64).unwrap();
418                let data = self.data.as_mut();
419                data[offset..offset + 32].copy_from_slice(pair.0.as_bytes());
420                data[offset + 32..offset + 64].copy_from_slice(pair.1.as_bytes());
421                Ok(())
422            }
423            None => Err(io::Error::new(
424                io::ErrorKind::InvalidInput,
425                "invalid node for this outboard",
426            )),
427        }
428    }
429
430    fn sync(&mut self) -> io::Result<()> {
431        Ok(())
432    }
433}
434
435#[cfg(feature = "tokio_fsm")]
436impl<T: AsRef<[u8]>> crate::io::fsm::Outboard for PreOrderMemOutboard<T> {
437    fn root(&self) -> blake3::Hash {
438        self.root
439    }
440    fn tree(&self) -> BaoTree {
441        self.tree
442    }
443    async fn load(&mut self, node: TreeNode) -> io::Result<Option<(blake3::Hash, blake3::Hash)>> {
444        Ok(load_raw_pre_mem(&self.tree, self.data.as_ref(), node).map(parse_hash_pair))
445    }
446}
447
448#[cfg(feature = "tokio_fsm")]
449impl<T: AsMut<[u8]>> crate::io::fsm::OutboardMut for PreOrderMemOutboard<T> {
450    async fn save(
451        &mut self,
452        node: TreeNode,
453        pair: &(blake3::Hash, blake3::Hash),
454    ) -> io::Result<()> {
455        match self.tree.pre_order_offset(node) {
456            Some(offset) => {
457                let offset_u64 = offset * 64;
458                let offset = usize::try_from(offset_u64).unwrap();
459                let data = self.data.as_mut();
460                data[offset..offset + 32].copy_from_slice(pair.0.as_bytes());
461                data[offset + 32..offset + 64].copy_from_slice(pair.1.as_bytes());
462                Ok(())
463            }
464            None => Err(io::Error::new(
465                io::ErrorKind::InvalidInput,
466                "invalid node for this outboard",
467            )),
468        }
469    }
470
471    async fn sync(&mut self) -> io::Result<()> {
472        Ok(())
473    }
474}
475
476fn load_raw_pre_mem(tree: &BaoTree, data: &[u8], node: TreeNode) -> Option<[u8; 64]> {
477    // this is a bit slow because pre_order_offset uses a loop.
478    // pretty sure there is a way to write it as a single expression if you spend the time.
479    // but profiling still has this in the nanosecond range, so this is unlikely to be a
480    // bottleneck.
481    let offset = tree.pre_order_offset(node)?;
482    let offset = usize::try_from(offset * 64).unwrap();
483    let slice = &data[offset..offset + 64];
484    Some(slice.try_into().unwrap())
485}
486
487fn load_pre(tree: &BaoTree, data: &[u8], node: TreeNode) -> Option<(blake3::Hash, blake3::Hash)> {
488    load_raw_pre_mem(tree, data, node).map(parse_hash_pair)
489}
490
491pub(crate) fn parse_hash_pair(buf: [u8; 64]) -> (blake3::Hash, blake3::Hash) {
492    let l_hash = blake3::Hash::from(<[u8; 32]>::try_from(&buf[..32]).unwrap());
493    let r_hash = blake3::Hash::from(<[u8; 32]>::try_from(&buf[32..]).unwrap());
494    (l_hash, r_hash)
495}