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}