Skip to main content

bao_tree/
lib.rs

1//! # Efficient BLAKE3 based verified streaming
2//!
3//! This crate is similar to the [bao crate](https://crates.io/crates/bao), but
4//! takes a slightly different approach.
5//!
6//! The core struct is [BaoTree], which describes the geometry of the tree and
7//! various ways to traverse it. An individual tree node is identified by
8//! [TreeNode], which is just a newtype wrapper for an u64.
9//!
10//! [TreeNode] provides various helpers to e.g. get the offset of a node in
11//! different traversal orders.
12//!
13//! There are newtypes for the different kinds of integers used in the
14//! tree:
15//! [ChunkNum] is an u64 number of chunks,
16//! [TreeNode] is an u64 tree node identifier,
17//! and [BlockSize] is the log base 2 of the chunk group size.
18//!
19//! All this is then used in the [io] module to implement the actual io, both
20//! [synchronous](io::sync) and [asynchronous](io::tokio).
21//!
22//! # Basic usage
23//!
24//! The basic workflow is like this: you have some existing data, for which
25//! you want to enable verified streaming. This data can either be in memory,
26//! in a file, or even a remote resource such as an HTTP server.
27//!
28//! ## Outboard creation
29//!
30//! You create an outboard using the [CreateOutboard](io::sync::CreateOutboard)
31//! trait. It has functions to [create](io::sync::CreateOutboard::create) an
32//! outboard from scratch or to [initialize](io::sync::CreateOutboard::init_from)
33//! data and root hash from existing data.
34//!
35//! ## Serving requests
36//!
37//! You serve streaming requests by using the
38//! [encode_ranges](io::sync::encode_ranges) or
39//! [encode_ranges_validated](io::sync::encode_ranges_validated) functions
40//! in the sync or async io module. For this you need data and a matching
41//! outboard.
42//!
43//! The difference between the two functions is that the validated version
44//! will check the hash of each chunk against the bao tree encoded in the
45//! outboard, so you will detect data corruption before sending out data
46//! to the requester. When using the unvalidated version, you might send out
47//! corrupted data without ever noticing and earn a bad reputation.
48//!
49//! Due to the speed of the blake3 hash function, validation is not a
50//! significant performance overhead compared to network operations and
51//! encryption.
52//!
53//! The requester will send a set of chunk ranges they are interested in.
54//! To compute chunk ranges from byte ranges, there is a helper function
55//! [round_up_to_chunks](io::round_up_to_chunks) that takes a byte range and
56//! rounds up to chunk ranges.
57//!
58//! If you just want to stream the entire blob, you can use [ChunkRanges::all]
59//! as the range.
60//!
61//! ## Processing requests
62//!
63//! You process requests by using the [decode_ranges](io::sync::decode_ranges)
64//! function in the sync or async io module. This function requires prior
65//! knowledge of the tree geometry (total data size and block size). A common
66//! way to get this information is to have the block size as a common parameter
67//! of both sides, and send the total data size as a prefix of the encoded data.
68//! E.g. the original bao crate uses a little endian u64 as the prefix.
69//!
70//! This function will perform validation in any case, there is no variant
71//! that skips validation since that would defeat the purpose of verified
72//! streaming.
73//!
74//! ## Simple end to end example
75//!
76//! ```no_run
77//! use std::io;
78//!
79//! use bao_tree::{
80//!     io::{
81//!         outboard::PreOrderOutboard,
82//!         round_up_to_chunks,
83//!         sync::{decode_ranges, encode_ranges_validated, valid_ranges, CreateOutboard},
84//!     },
85//!     BlockSize, ByteRanges, ChunkRanges,
86//! };
87//!
88//! /// Use a block size of 16 KiB, a good default for most cases
89//! const BLOCK_SIZE: BlockSize = BlockSize::from_chunk_log(4);
90//!
91//! # fn main() -> io::Result<()> {
92//! // The file we want to serve
93//! let file = std::fs::File::open("video.mp4")?;
94//! // Create an outboard for the file, using the current size
95//! let ob = PreOrderOutboard::<Vec<u8>>::create(&file, BLOCK_SIZE)?;
96//! // Encode the first 100000 bytes of the file
97//! let ranges = ByteRanges::from(0..100000);
98//! let ranges = round_up_to_chunks(&ranges);
99//! // Stream of data to client. Needs to implement `io::Write`. We just use a vec here.
100//! let mut to_client = vec![];
101//! encode_ranges_validated(&file, &ob, &ranges, &mut to_client)?;
102//!
103//! // Stream of data from client. Needs to implement `io::Read`. We just wrap the vec in a cursor.
104//! let from_server = io::Cursor::new(to_client);
105//! let root = ob.root;
106//! let tree = ob.tree;
107//!
108//! // Decode the encoded data into a file
109//! let mut decoded = std::fs::File::create("copy.mp4")?;
110//! let mut ob = PreOrderOutboard {
111//!     tree,
112//!     root,
113//!     data: vec![],
114//! };
115//! decode_ranges(from_server, &ranges, &mut decoded, &mut ob)?;
116//!
117//! // the first 100000 bytes of the file should now be in `decoded`
118//! // in addition, the required part of the tree to validate that the data is
119//! // correct are in `ob.data`
120//!
121//! // Print the valid ranges of the file
122//! for range in valid_ranges(&ob, &decoded, &ChunkRanges::all()) {
123//!     println!("{:?}", range);
124//! }
125//! # Ok(())
126//! # }
127//! ```
128//!
129//! # Async end to end example
130//!
131//! The async version is very similar to the sync version, except that it needs
132//! an async context. All functions that do IO are async. The file has to be
133//! an [iroh_io::File], which is just a wrapper for [std::fs::File] that implements
134//! async random access via the [AsyncSliceReader](iroh_io::AsyncSliceReader) trait.
135//!
136//! We use [futures_lite] crate, but using the normal futures crate will also work.
137//!
138//! ```no_run
139//! use std::io;
140//!
141//! use bao_tree::{
142//!     io::{
143//!         fsm::{decode_ranges, encode_ranges_validated, valid_ranges, CreateOutboard},
144//!         outboard::PreOrderOutboard,
145//!         round_up_to_chunks,
146//!     },
147//!     BlockSize, ByteRanges, ChunkRanges,
148//! };
149//! use bytes::BytesMut;
150//! use futures_lite::StreamExt;
151//!
152//! /// Use a block size of 16 KiB, a good default for most cases
153//! const BLOCK_SIZE: BlockSize = BlockSize::from_chunk_log(4);
154//!
155//! # #[tokio::main]
156//! # async fn main() -> io::Result<()> {
157//! // The file we want to serve
158//! let mut file = iroh_io::File::open("video.mp4".into()).await?;
159//! // Create an outboard for the file, using the current size
160//! let mut ob = PreOrderOutboard::<BytesMut>::create(&mut file, BLOCK_SIZE).await?;
161//! // Encode the first 100000 bytes of the file
162//! let ranges = ByteRanges::from(0..100000);
163//! let ranges = round_up_to_chunks(&ranges);
164//! // Stream of data to client. Needs to implement `io::Write`. We just use a vec here.
165//! let mut to_client = Vec::new();
166//! encode_ranges_validated(file, &mut ob, &ranges, &mut to_client).await?;
167//!
168//! // Stream of data from client. Needs to implement `io::Read`. We just wrap the vec in a cursor.
169//! let from_server = io::Cursor::new(to_client.as_slice());
170//! let root = ob.root;
171//! let tree = ob.tree;
172//!
173//! // Decode the encoded data into a file
174//! let mut decoded = iroh_io::File::open("copy.mp4".into()).await?;
175//! let mut ob = PreOrderOutboard {
176//!     tree,
177//!     root,
178//!     data: BytesMut::new(),
179//! };
180//! decode_ranges(from_server, ranges, &mut decoded, &mut ob).await?;
181//!
182//! // the first 100000 bytes of the file should now be in `decoded`
183//! // in addition, the required part of the tree to validate that the data is
184//! // correct are in `ob.data`
185//!
186//! // Print the valid ranges of the file
187//! let ranges = ChunkRanges::all();
188//! let mut stream = valid_ranges(&mut ob, &mut decoded, &ranges);
189//! while let Some(range) = stream.next().await {
190//!     println!("{:?}", range);
191//! }
192//! # Ok(())
193//! # }
194//! ```
195//!
196//! # Keyed hashing
197//!
198//! For domain-separated trees, use the `keyed_*` functions. They mirror the
199//! standard API with an additional `key: &[u8; 32]` argument, like
200//! [`blake3::keyed_hash`] mirrors [`blake3::hash`]. The key is out-of-band
201//! metadata and is not included in the encoded stream.
202//!
203//! # Compatibility with the [bao crate](https://crates.io/crates/bao)
204//!
205//! This crate will be compatible with the bao crate, provided you do the
206//! following:
207//!
208//! - use a block size of 1024, so no chunk groups
209//! - use a little endian u64 as the prefix for the encoded data
210//! - use only a single range
211#![deny(missing_docs)]
212use std::{
213    fmt::{self, Debug},
214    ops::Range,
215};
216
217use range_collections::RangeSetRef;
218pub mod iter;
219mod rec;
220mod tree;
221use iter::*;
222pub use tree::{BlockSize, ChunkNum};
223pub mod io;
224pub use blake3;
225
226#[cfg(all(test, feature = "tokio_fsm"))]
227mod tests;
228#[cfg(all(test, feature = "tokio_fsm"))]
229mod tests2;
230
231/// A set of chunk ranges
232pub type ChunkRanges = range_collections::RangeSet2<ChunkNum>;
233
234/// A set of byte ranges
235pub type ByteRanges = range_collections::RangeSet2<u64>;
236
237/// A referenceable set of chunk ranges
238///
239/// [ChunkRanges] implements [`AsRef<ChunkRangesRef>`].
240pub type ChunkRangesRef = range_collections::RangeSetRef<ChunkNum>;
241
242/// Hashing mode for shared encode and decode paths, either standard or keyed BLAKE3.
243#[derive(Clone, Copy)]
244pub(crate) enum HashMode {
245    Standard,
246    Keyed([u8; 32]),
247}
248
249impl std::fmt::Debug for HashMode {
250    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251        match self {
252            HashMode::Standard => f.write_str("Standard"),
253            HashMode::Keyed(_) => f.debug_struct("Keyed").finish_non_exhaustive(),
254        }
255    }
256}
257
258impl HashMode {
259    pub(crate) fn hash_subtree(
260        &self,
261        start_chunk: u64,
262        data: &[u8],
263        is_root: bool,
264    ) -> blake3::Hash {
265        use blake3::hazmat::{ChainingValue, HasherExt};
266        if is_root {
267            debug_assert!(start_chunk == 0);
268            match self {
269                HashMode::Standard => blake3::hash(data),
270                HashMode::Keyed(key) => blake3::keyed_hash(key, data),
271            }
272        } else {
273            let mut hasher = match self {
274                HashMode::Standard => blake3::Hasher::new(),
275                HashMode::Keyed(key) => blake3::Hasher::new_keyed(key),
276            };
277            hasher.set_input_offset(start_chunk * 1024);
278            hasher.update(data);
279            let non_root_hash: ChainingValue = hasher.finalize_non_root();
280            blake3::Hash::from(non_root_hash)
281        }
282    }
283
284    pub(crate) fn parent_cv(
285        &self,
286        left_child: &blake3::Hash,
287        right_child: &blake3::Hash,
288        is_root: bool,
289    ) -> blake3::Hash {
290        use blake3::hazmat::{merge_subtrees_non_root, merge_subtrees_root, ChainingValue, Mode};
291        let left_child: ChainingValue = *left_child.as_bytes();
292        let right_child: ChainingValue = *right_child.as_bytes();
293        let mode = match self {
294            HashMode::Standard => Mode::Hash,
295            HashMode::Keyed(key) => Mode::KeyedHash(key),
296        };
297        if is_root {
298            merge_subtrees_root(&left_child, &right_child, mode)
299        } else {
300            blake3::Hash::from(merge_subtrees_non_root(&left_child, &right_child, mode))
301        }
302    }
303}
304
305/// Compute the hash of a subtree using BLAKE3 keyed mode.
306///
307/// See [keyed_parent_cv] for merging child hashes in keyed mode.
308#[inline]
309pub fn keyed_hash_subtree(
310    start_chunk: u64,
311    data: &[u8],
312    is_root: bool,
313    key: &[u8; 32],
314) -> blake3::Hash {
315    HashMode::Keyed(*key).hash_subtree(start_chunk, data, is_root)
316}
317
318/// Merge two child subtree hashes using BLAKE3 keyed mode.
319#[inline]
320pub fn keyed_parent_cv(
321    left_child: &blake3::Hash,
322    right_child: &blake3::Hash,
323    is_root: bool,
324    key: &[u8; 32],
325) -> blake3::Hash {
326    HashMode::Keyed(*key).parent_cv(left_child, right_child, is_root)
327}
328
329/// Defines a Bao tree.
330///
331/// This is just the specification of the tree, it does not contain any actual data.
332///
333/// Usually trees are self-contained. This means that the tree starts at chunk 0,
334/// and the hash of the root node is computed with the is_root flag set to true.
335///
336/// For some internal use, it is also possible to create trees that are just subtrees
337/// of a larger tree. In this case, the start_chunk is the chunk number of the first
338/// chunk in the tree, and the is_root flag can be false.
339#[derive(Debug, Clone, Copy, PartialEq, Eq)]
340pub struct BaoTree {
341    /// Total number of bytes in the file
342    size: u64,
343    /// Log base 2 of the chunk group size
344    block_size: BlockSize,
345}
346
347/// An offset of a node in a post-order outboard
348#[derive(Debug, Clone, Copy)]
349pub enum PostOrderOffset {
350    /// the node is stable and won't change when appending data
351    Stable(u64),
352    /// the node is unstable and will change when appending data
353    Unstable(u64),
354}
355
356impl PostOrderOffset {
357    /// Just get the offset value, ignoring whether it's stable or unstable
358    pub fn value(self) -> u64 {
359        match self {
360            Self::Stable(n) => n,
361            Self::Unstable(n) => n,
362        }
363    }
364}
365
366impl BaoTree {
367    /// Create a new self contained BaoTree
368    pub fn new(size: u64, block_size: BlockSize) -> Self {
369        Self { size, block_size }
370    }
371
372    /// The size of the blob from which this tree was constructed, in bytes
373    pub fn size(&self) -> u64 {
374        self.size
375    }
376
377    /// The block size of the tree
378    pub fn block_size(&self) -> BlockSize {
379        self.block_size
380    }
381
382    /// Given a tree of size `size` and block size `block_size`,
383    /// compute the root node and the number of nodes for a shifted tree.
384    pub(crate) fn shifted(&self) -> (TreeNode, TreeNode) {
385        let level = self.block_size.0;
386        let size = self.size;
387        let shift = 10 + level;
388        let mask = (1 << shift) - 1;
389        // number of full blocks of size 1024 << level
390        let full_blocks = size >> shift;
391        // 1 if the last block is non zero, 0 otherwise
392        let open_block = ((size & mask) != 0) as u64;
393        // total number of blocks, rounding up to 1 if there are no blocks
394        let blocks = (full_blocks + open_block).max(1);
395        let n = blocks.div_ceil(2);
396        // root node
397        let root = n.next_power_of_two() - 1;
398        // number of nodes in the tree
399        let filled_size = n + n.saturating_sub(1);
400        (TreeNode(root), TreeNode(filled_size))
401    }
402
403    fn byte_range(&self, node: TreeNode) -> Range<u64> {
404        let start = node.chunk_range().start.to_bytes();
405        let end = node.chunk_range().end.to_bytes();
406        start..end.min(self.size)
407    }
408
409    /// Compute the byte ranges for a leaf node
410    ///
411    /// Returns two ranges, the first is the left range, the second is the right range
412    /// If the leaf is partially contained in the tree, the right range will be empty
413    fn leaf_byte_ranges3(&self, leaf: TreeNode) -> (u64, u64, u64) {
414        let Range { start, end } = leaf.byte_range();
415        let mid = leaf.mid().to_bytes();
416        if !(start < self.size || (start == 0 && self.size == 0)) {
417            debug_assert!(start < self.size || (start == 0 && self.size == 0));
418        }
419        (start, mid.min(self.size), end.min(self.size))
420    }
421
422    /// Traverse the entire tree in post order as [BaoChunk]s
423    ///
424    /// This iterator is used by both the sync and async io code for computing
425    /// an outboard from existing data
426    pub fn post_order_chunks_iter(&self) -> PostOrderChunkIter {
427        PostOrderChunkIter::new(*self)
428    }
429
430    /// Traverse the part of the tree that is relevant for a ranges query
431    /// in pre order as [BaoChunk]s
432    ///
433    /// This iterator is used by both the sync and async io code for encoding
434    /// from an outboard and ranges as well as decoding an encoded stream.
435    pub fn ranges_pre_order_chunks_iter_ref<'a>(
436        &self,
437        ranges: &'a RangeSetRef<ChunkNum>,
438        min_level: u8,
439    ) -> PreOrderPartialChunkIterRef<'a> {
440        PreOrderPartialChunkIterRef::new(*self, ranges, min_level)
441    }
442
443    /// Traverse the entire tree in post order as [TreeNode]s,
444    /// down to the level given by the block size.
445    pub fn post_order_nodes_iter(&self) -> impl Iterator<Item = TreeNode> {
446        let (root, len) = self.shifted();
447        let shift = self.block_size.0;
448        PostOrderNodeIter::new(root, len).map(move |x| x.subtract_block_size(shift))
449    }
450
451    /// Traverse the entire tree in pre order as [TreeNode]s,
452    /// down to the level given by the block size.
453    pub fn pre_order_nodes_iter(&self) -> impl Iterator<Item = TreeNode> {
454        let (root, len) = self.shifted();
455        let shift = self.block_size.0;
456        PreOrderNodeIter::new(root, len).map(move |x| x.subtract_block_size(shift))
457    }
458
459    /// Traverse the part of the tree that is relevant for a ranges querys
460    /// in pre order as [NodeInfo]s
461    ///
462    /// This is mostly used internally.
463    ///
464    /// When `min_level` is set to a value greater than 0, the iterator will
465    /// skip all branch nodes that are at a level < min_level if they are fully
466    /// covered by the ranges.
467    #[cfg(test)]
468    pub fn ranges_pre_order_nodes_iter<'a>(
469        &self,
470        ranges: &'a RangeSetRef<ChunkNum>,
471        min_level: u8,
472    ) -> PreOrderPartialIterRef<'a> {
473        PreOrderPartialIterRef::new(*self, ranges, min_level)
474    }
475
476    /// Root of the tree
477    ///
478    /// Does not consider block size
479    pub fn root(&self) -> TreeNode {
480        let shift = 10;
481        let mask = (1 << shift) - 1;
482        let full_blocks = self.size >> shift;
483        let open_block = ((self.size & mask) != 0) as u64;
484        let blocks = (full_blocks + open_block).max(1);
485        let chunks = ChunkNum(blocks);
486        TreeNode::root(chunks)
487    }
488
489    /// Number of blocks in the tree
490    ///
491    /// At chunk group size 1, this is the same as the number of chunks
492    /// Even a tree with 0 bytes size has a single block
493    pub fn blocks(&self) -> u64 {
494        // handle the case of an empty tree having 1 block
495        blocks(self.size, self.block_size).max(1)
496    }
497
498    /// Number of chunks in the tree
499    pub fn chunks(&self) -> ChunkNum {
500        ChunkNum::chunks(self.size)
501    }
502
503    /// Number of hash pairs in the outboard
504    fn outboard_hash_pairs(&self) -> u64 {
505        self.blocks() - 1
506    }
507
508    /// The outboard size for this tree.
509    ///
510    /// This is the outboard size *without* the size prefix.
511    pub fn outboard_size(&self) -> u64 {
512        self.outboard_hash_pairs() * 64
513    }
514
515    #[allow(dead_code)]
516    fn filled_size(&self) -> TreeNode {
517        let blocks = self.chunks();
518        let n = blocks.0.div_ceil(2);
519        TreeNode(n + n.saturating_sub(1))
520    }
521
522    /// true if the node is a leaf for this tree
523    ///
524    /// If a tree has a non-zero block size, this is different than the node
525    /// being a leaf (level=0).
526    #[cfg(test)]
527    const fn is_leaf(&self, node: TreeNode) -> bool {
528        node.level() == self.block_size.to_u32()
529    }
530
531    /// true if the given node is persisted
532    ///
533    /// the only node that is not persisted is the last leaf node, if it is
534    /// less than half full
535    #[inline]
536    #[cfg(test)]
537    const fn is_persisted(&self, node: TreeNode) -> bool {
538        !self.is_leaf(node) || node.mid().to_bytes() < self.size
539    }
540
541    /// true if this is a node that is relevant for the outboard
542    #[inline]
543    const fn is_relevant_for_outboard(&self, node: TreeNode) -> bool {
544        let level = node.level();
545        if level < self.block_size.to_u32() {
546            // too small, this outboard does not track it
547            false
548        } else if level > self.block_size.to_u32() {
549            // a parent node, always relevant
550            true
551        } else {
552            node.mid().to_bytes() < self.size
553        }
554    }
555
556    /// The offset of the given node in the pre order traversal
557    pub fn pre_order_offset(&self, node: TreeNode) -> Option<u64> {
558        // if the node has a level less than block_size, this will return None
559        let shifted = node.add_block_size(self.block_size.0)?;
560        let is_half_leaf = shifted.is_leaf() && node.mid().to_bytes() >= self.size;
561        if !is_half_leaf {
562            let (_, tree_filled_size) = self.shifted();
563            Some(pre_order_offset_loop(shifted.0, tree_filled_size.0))
564        } else {
565            None
566        }
567    }
568
569    /// The offset of the given node in the post order traversal
570    pub fn post_order_offset(&self, node: TreeNode) -> Option<PostOrderOffset> {
571        // if the node has a level less than block_size, this will return None
572        let shifted = node.add_block_size(self.block_size.0)?;
573        if node.byte_range().end <= self.size {
574            // stable node, use post_order_offset
575            Some(PostOrderOffset::Stable(shifted.post_order_offset()))
576        } else {
577            // unstable node
578            if shifted.is_leaf() && node.mid().to_bytes() >= self.size {
579                // half full leaf node, not considered
580                None
581            } else {
582                // compute the offset based on the total size and the height of the node
583                self.outboard_hash_pairs()
584                    .checked_sub(u64::from(node.right_count()) + 1)
585                    .map(PostOrderOffset::Unstable)
586            }
587        }
588    }
589
590    const fn chunk_group_chunks(&self) -> ChunkNum {
591        ChunkNum(1 << self.block_size.0)
592    }
593
594    fn chunk_group_bytes(&self) -> usize {
595        self.chunk_group_chunks().to_bytes().try_into().unwrap()
596    }
597}
598
599/// number of blocks that this number of bytes covers,
600/// given a block size
601pub(crate) const fn blocks(size: u64, block_size: BlockSize) -> u64 {
602    let chunk_group_log = block_size.0;
603    let block_bits = chunk_group_log + 10;
604    let block_mask = (1 << block_bits) - 1;
605    let full_blocks = size >> block_bits;
606    let open_block = ((size & block_mask) != 0) as u64;
607    full_blocks + open_block
608}
609
610/// An u64 that defines a node in a bao tree.
611///
612/// You typically don't have to use this, but it can be useful for debugging
613/// and error handling. Hash validation errors contain a `TreeNode` that allows
614/// you to find the position where validation failed.
615#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
616#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
617pub struct TreeNode(u64);
618
619impl fmt::Display for TreeNode {
620    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
621        write!(f, "{}", self.0)
622    }
623}
624
625impl fmt::Debug for TreeNode {
626    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
627        if !f.alternate() {
628            write!(f, "TreeNode({})", self.0)
629        } else if self.is_leaf() {
630            write!(f, "TreeNode::Leaf({})", self.0)
631        } else {
632            write!(f, "TreeNode::Branch({}, level={})", self.0, self.level())
633        }
634    }
635}
636
637impl TreeNode {
638    /// Create a new tree node from a start chunk and a level
639    ///
640    /// The start chunk must be the start of a subtree with the given level.
641    /// So for level 0, the start chunk must even. For level 1, the start chunk
642    /// must be divisible by 4, etc.
643    ///
644    /// This is a bridge from the recursive reference implementation to the node
645    /// based implementations, and is therefore only used in tests.
646    #[cfg(all(test, feature = "tokio_fsm"))]
647    fn from_start_chunk_and_level(start_chunk: ChunkNum, level: BlockSize) -> Self {
648        let start_chunk = start_chunk.0;
649        let level = level.0;
650        // check that the start chunk a start of a subtree with level `level`
651        // this ensures that there is a 0 at bit `level`.
652        let check_mask = (1 << (level + 1)) - 1;
653        debug_assert_eq!(start_chunk & check_mask, 0);
654        let level_mask = (1 << level) - 1;
655        // set the trailing `level` bits to 1.
656        // The level is the number of trailing ones.
657        Self(start_chunk | level_mask)
658    }
659
660    /// Given a number of blocks, gives root node
661    fn root(chunks: ChunkNum) -> TreeNode {
662        Self(chunks.0.div_ceil(2).next_power_of_two() - 1)
663    }
664
665    /// the middle of the tree node, in blocks
666    pub const fn mid(&self) -> ChunkNum {
667        ChunkNum(self.0 + 1)
668    }
669
670    #[inline]
671    const fn half_span(&self) -> u64 {
672        1 << self.level()
673    }
674
675    /// The level of the node in the tree, 0 for leafs.
676    #[inline]
677    pub const fn level(&self) -> u32 {
678        self.0.trailing_ones()
679    }
680
681    /// True if this is a leaf node.
682    #[inline]
683    pub const fn is_leaf(&self) -> bool {
684        (self.0 & 1) == 0
685    }
686
687    /// Convert a node to a node in a tree with a smaller block size
688    ///
689    /// E.g. a leaf node in a tree with block size 4 will become a node
690    /// with level 4 in a tree with block size 0.
691    ///
692    /// This works by just adding n trailing 1 bits to the node by shifting
693    /// to the left.
694    #[inline]
695    pub const fn subtract_block_size(&self, n: u8) -> Self {
696        let shifted = !(!self.0 << n);
697        Self(shifted)
698    }
699
700    /// Convert a node to a node in a tree with a larger block size
701    ///
702    /// If the nodes has n trailing 1 bits, they are removed by shifting
703    /// the node to the right by n bits.
704    ///
705    /// If the node has less than n trailing 1 bits, the node is too small
706    /// to be represented in the target tree.
707    #[inline]
708    pub const fn add_block_size(&self, n: u8) -> Option<Self> {
709        let mask = (1 << n) - 1;
710        // check if the node has a high enough level
711        if self.0 & mask == mask {
712            Some(Self(self.0 >> n))
713        } else {
714            None
715        }
716    }
717
718    /// Range of blocks that this node covers, given a block size
719    ///
720    /// Note that this will give the untruncated range, which may be larger than
721    /// the actual tree. To get the exact byte range for a tree, use
722    /// [BaoTree::byte_range];
723    fn byte_range(&self) -> Range<u64> {
724        let range = self.chunk_range();
725        range.start.to_bytes()..range.end.to_bytes()
726    }
727
728    /// Number of nodes below this node, excluding this node.
729    #[inline]
730    pub const fn count_below(&self) -> u64 {
731        // go to representation where trailing zeros are the level
732        let x = self.0 + 1;
733        // isolate the lowest bit
734        let lowest_bit = x & (-(x as i64) as u64);
735        // number of nodes is n * 2 - 1, subtract 1 for the node itself
736        lowest_bit * 2 - 2
737    }
738
739    /// Get the next left ancestor of this node, or None if there is none.
740    pub fn next_left_ancestor(&self) -> Option<Self> {
741        self.next_left_ancestor0().map(Self)
742    }
743
744    /// Get the left child of this node, or None if it is a child node.
745    pub fn left_child(&self) -> Option<Self> {
746        let offset = 1 << self.level().checked_sub(1)?;
747        Some(Self(self.0 - offset))
748    }
749
750    /// Get the right child of this node, or None if it is a child node.
751    pub fn right_child(&self) -> Option<Self> {
752        let offset = 1 << self.level().checked_sub(1)?;
753        Some(Self(self.0 + offset))
754    }
755
756    /// Unrestricted parent, can only be None if we are at the top
757    pub fn parent(&self) -> Option<Self> {
758        let level = self.level();
759        if level == 63 {
760            return None;
761        }
762        let span = 1u64 << level;
763        let offset = self.0;
764        Some(Self(if (offset & (span * 2)) == 0 {
765            offset + span
766        } else {
767            offset - span
768        }))
769    }
770
771    /// Restricted parent, will be None if we call parent on the root
772    pub fn restricted_parent(&self, len: Self) -> Option<Self> {
773        let mut curr = *self;
774        while let Some(parent) = curr.parent() {
775            if parent.0 < len.0 {
776                return Some(parent);
777            }
778            curr = parent;
779        }
780        // we hit the top
781        None
782    }
783
784    /// Get a valid right descendant for an offset
785    pub(crate) fn right_descendant(&self, len: Self) -> Option<Self> {
786        let mut node = self.right_child()?;
787        while node >= len {
788            node = node.left_child()?;
789        }
790        Some(node)
791    }
792
793    /// Get the range of nodes this node covers
794    pub const fn node_range(&self) -> Range<Self> {
795        let half_span = self.half_span();
796        let nn = self.0;
797        let r = nn + half_span;
798        let l = nn + 1 - half_span;
799        Self(l)..Self(r)
800    }
801
802    /// Get the range of blocks this node covers
803    pub fn chunk_range(&self) -> Range<ChunkNum> {
804        let level = self.level();
805        let span = 1 << level;
806        let mid = self.0 + 1;
807        // at level 0 (leaf), range will be nn..nn+2
808        // at level >0 (branch), range will be centered on nn+1
809        ChunkNum(mid - span)..ChunkNum(mid + span)
810    }
811
812    /// the number of times you have to go right from the root to get to this node
813    ///
814    /// 0 for a root node
815    pub fn right_count(&self) -> u32 {
816        (self.0 + 1).count_ones() - 1
817    }
818
819    /// Get the post order offset of this node
820    #[inline]
821    pub const fn post_order_offset(&self) -> u64 {
822        // compute number of nodes below me
823        let below_me = self.count_below();
824        // compute next ancestor that is to the left
825        let next_left_ancestor = self.next_left_ancestor0();
826        // compute offset
827        match next_left_ancestor {
828            Some(nla) => below_me + nla + 1 - ((nla + 1).count_ones() as u64),
829            None => below_me,
830        }
831    }
832
833    /// Get the range of post order offsets this node covers
834    pub const fn post_order_range(&self) -> Range<u64> {
835        let offset = self.post_order_offset();
836        let end = offset + 1;
837        let start = offset - self.count_below();
838        start..end
839    }
840
841    /// Get the next left ancestor, or None if we don't have one
842    ///
843    /// this is a separate fn so it can be const.
844    #[inline]
845    const fn next_left_ancestor0(&self) -> Option<u64> {
846        // add 1 to go to the representation where trailing zeroes = level
847        let x = self.0 + 1;
848        // clear the lowest bit
849        let without_lowest_bit = x & (x - 1);
850        // go back to the normal representation,
851        // producing None if without_lowest_bit is 0, which means that there is no next left ancestor
852        without_lowest_bit.checked_sub(1)
853    }
854}
855
856/// Iterative way to find the offset of a node in a pre-order traversal.
857///
858/// I am sure there is a way that does not require a loop, but this will do for now.
859/// It is slower than the direct formula, but it is still in the nanosecond range,
860/// so at a block size of 16 KiB it should not be the limiting factor for anything.
861fn pre_order_offset_loop(node: u64, len: u64) -> u64 {
862    // node level, 0 for leaf nodes
863    let level = (!node).trailing_zeros();
864    // span of the node, 1 for leaf nodes
865    let span = 1u64 << level;
866    // nodes to the left of the tree of this node
867    let left = node + 1 - span;
868    // count the parents with a loop
869    let mut parent_count = 0;
870    let mut offset = node;
871    let mut span = span;
872    // loop until we reach the root, adding valid parents
873    loop {
874        let pspan = span * 2;
875        // find parent
876        offset = if (offset & pspan) == 0 {
877            offset + span
878        } else {
879            offset - span
880        };
881        // if parent is inside the tree, increase parent count
882        if offset < len {
883            parent_count += 1;
884        }
885        if pspan >= len {
886            // we are at the root
887            break;
888        }
889        span = pspan;
890    }
891    left - (left.count_ones() as u64) + parent_count
892}
893
894/// Split a range set into range sets for the left and right half of a node
895///
896/// Requires that the range set is minimal, it should not contain any redundant
897/// boundaries outside of the range of the node. The values outside of the node
898/// range don't matter, so any change outside the range must be omitted.
899///
900/// Produces two range sets that are also minimal. A range set for left or right
901/// that covers the entire range of the node will be replaced with the set of
902/// all chunks, so an is_all() check can be used to check if further recursion
903/// is necessary.
904pub(crate) fn split(
905    ranges: &RangeSetRef<ChunkNum>,
906    node: TreeNode,
907) -> (&RangeSetRef<ChunkNum>, &RangeSetRef<ChunkNum>) {
908    let mid = node.mid();
909    let start = node.chunk_range().start;
910    split_inner(ranges, start, mid)
911}
912
913/// The actual implementation of split. This is used from split and from the
914/// recursive reference implementation.
915pub(crate) fn split_inner(
916    ranges: &RangeSetRef<ChunkNum>,
917    start: ChunkNum,
918    mid: ChunkNum,
919) -> (&RangeSetRef<ChunkNum>, &RangeSetRef<ChunkNum>) {
920    let (mut a, mut b) = ranges.split(mid);
921    // check that a does not contain a redundant boundary at or after mid
922    debug_assert!(a.boundaries().last() < Some(&mid));
923    // Replace a with the canonicalized version if it is a single interval that
924    // starts at or before start. This is necessary to be able to check it with
925    // RangeSetRef::is_all()
926    if a.boundaries().len() == 1 && a.boundaries()[0] <= start {
927        a = RangeSetRef::new(&[ChunkNum(0)]).unwrap();
928    }
929    // Replace b with the canonicalized version if it is a single interval that
930    // starts at or before mid. This is necessary to be able to check it with
931    // RangeSetRef::is_all()
932    if b.boundaries().len() == 1 && b.boundaries()[0] <= mid {
933        b = RangeSetRef::new(&[ChunkNum(0)]).unwrap();
934    }
935    (a, b)
936}
937
938// Module that handles io::Error serialization/deserialization
939#[cfg(feature = "serde")]
940mod io_error_serde {
941    use std::{fmt, io};
942
943    use serde::{
944        de::{self, Visitor},
945        Deserializer, Serializer,
946    };
947
948    pub fn serialize<S>(error: &io::Error, serializer: S) -> Result<S::Ok, S::Error>
949    where
950        S: Serializer,
951    {
952        // Serialize the error kind and message
953        serializer.serialize_str(&format!("{:?}:{}", error.kind(), error))
954    }
955
956    pub fn deserialize<'de, D>(deserializer: D) -> Result<io::Error, D::Error>
957    where
958        D: Deserializer<'de>,
959    {
960        struct IoErrorVisitor;
961
962        impl Visitor<'_> for IoErrorVisitor {
963            type Value = io::Error;
964
965            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
966                formatter.write_str("an io::Error string representation")
967            }
968
969            fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
970            where
971                E: de::Error,
972            {
973                // For simplicity, create a generic error
974                // In a real app, you might want to parse the kind from the string
975                Ok(io::Error::other(value))
976            }
977        }
978
979        deserializer.deserialize_str(IoErrorVisitor)
980    }
981}