bao_tree/io/
mod.rs

1//! Implementation of bao streaming for std io and tokio io
2use std::pin::Pin;
3
4use crate::{blake3, BlockSize, ChunkNum, ChunkRanges, TreeNode};
5use bytes::Bytes;
6
7mod error;
8pub use error::*;
9use range_collections::{range_set::RangeSetRange, RangeSetRef};
10use std::future::Future;
11
12#[cfg(feature = "tokio_fsm")]
13pub mod fsm;
14pub mod outboard;
15pub mod sync;
16
17/// A parent hash pair.
18#[derive(Debug)]
19pub struct Parent {
20    /// The node in the tree for which the hashes are.
21    pub node: TreeNode,
22    /// The pair of hashes for the node.
23    pub pair: (blake3::Hash, blake3::Hash),
24}
25
26/// A leaf node.
27#[derive(Debug)]
28pub struct Leaf {
29    /// The byte offset of the leaf in the file.
30    pub offset: u64,
31    /// The data of the leaf.
32    pub data: Bytes,
33}
34
35/// A content item for the bao streaming protocol.
36///
37/// After reading the initial header, the only possible items are `Parent` and
38/// `Leaf`.
39#[derive(Debug)]
40pub enum BaoContentItem {
41    /// a parent node, to update the outboard
42    Parent(Parent),
43    /// a leaf node, to write to the file
44    Leaf(Leaf),
45}
46
47impl From<Parent> for BaoContentItem {
48    fn from(p: Parent) -> Self {
49        Self::Parent(p)
50    }
51}
52
53impl From<Leaf> for BaoContentItem {
54    fn from(l: Leaf) -> Self {
55        Self::Leaf(l)
56    }
57}
58
59/// Given a range set of byte ranges, round it up to full chunks.
60///
61/// E.g. a byte range from 1..3 will be converted into the chunk range 0..1 (0..1024 bytes).
62pub fn round_up_to_chunks(ranges: &RangeSetRef<u64>) -> ChunkRanges {
63    let mut res = ChunkRanges::empty();
64    // we don't know if the ranges are overlapping, so we just compute the union
65    for item in ranges.iter() {
66        // full_chunks() rounds down, chunks() rounds up
67        match item {
68            RangeSetRange::RangeFrom(range) => {
69                res |= ChunkRanges::from(ChunkNum::full_chunks(*range.start)..)
70            }
71            RangeSetRange::Range(range) => {
72                res |= ChunkRanges::from(
73                    ChunkNum::full_chunks(*range.start)..ChunkNum::chunks(*range.end),
74                )
75            }
76        }
77    }
78    res
79}
80
81/// Given a range set of byte ranges, round it up to chunk groups.
82///
83/// If we store outboard data at a level of granularity of `block_size`, we can only
84/// share full chunk groups because we don't have proofs for anything below a chunk group.
85pub fn full_chunk_groups(ranges: &ChunkRanges, block_size: BlockSize) -> ChunkRanges {
86    fn floor(value: u64, shift: u8) -> u64 {
87        value >> shift << shift
88    }
89
90    fn ceil(value: u64, shift: u8) -> u64 {
91        (value + (1 << shift) - 1) >> shift << shift
92    }
93    let mut res = ChunkRanges::empty();
94    for item in ranges.iter() {
95        match item {
96            RangeSetRange::RangeFrom(range) => {
97                let start = ceil(range.start.0, block_size.0);
98                res |= ChunkRanges::from(ChunkNum(start)..)
99            }
100            RangeSetRange::Range(range) => {
101                let start = ceil(range.start.0, block_size.0);
102                let end = floor(range.end.0, block_size.0);
103                if start < end {
104                    res |= ChunkRanges::from(ChunkNum(start)..ChunkNum(end))
105                }
106            }
107        }
108    }
109    res
110}
111
112pub(crate) fn combine_hash_pair(l: &blake3::Hash, r: &blake3::Hash) -> [u8; 64] {
113    let mut res = [0u8; 64];
114    let lb: &mut [u8; 32] = (&mut res[0..32]).try_into().unwrap();
115    *lb = *l.as_bytes();
116    let rb: &mut [u8; 32] = (&mut res[32..]).try_into().unwrap();
117    *rb = *r.as_bytes();
118    res
119}
120
121pub(crate) type LocalBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;