Add helper fn to compute full chunk groups.

src/io/mod.rs

@@ -1,5 +1,5 @@
 //! Implementation of bao streaming for std io and tokio io
-use crate::{blake3, BaoTree, BlockSize, ByteNum, ChunkRanges, TreeNode};
+use crate::{blake3, BaoTree, BlockSize, ByteNum, ChunkNum, ChunkRanges, TreeNode};
 use bytes::Bytes;
 
 mod error;
@@ -78,3 +78,35 @@ pub fn round_up_to_chunks(ranges: &RangeSetRef<u64>) -> ChunkRanges {
     }
     res
 }
+
+/// Given a range set of byte ranges, round it up to chunk groups
+/// Given a range set of chunk ranges, return the full chunk groups.
+///
+/// If we store outboard data at a level of granularity of `block_size`, we can only
+/// share full chunk groups because we don't have proofs for anything below a chunk group.
+pub fn full_chunk_groups(ranges: &ChunkRanges, block_size: BlockSize) -> ChunkRanges {
+    fn floor(value: u64, shift: u8) -> u64 {
+        value >> shift << shift
+    }
+
+    fn ceil(value: u64, shift: u8) -> u64 {
+        (value + (1 << shift) - 1) >> shift << shift
+    }
+    let mut res = ChunkRanges::empty();
+    for item in ranges.iter() {
+        match item {
+            RangeSetRange::RangeFrom(range) => {
+                let start = ceil(range.start.0, block_size.0);
+                res |= ChunkRanges::from(ChunkNum(start)..)
+            }
+            RangeSetRange::Range(range) => {
+                let start = ceil(range.start.0, block_size.0);
+                let end = floor(range.end.0, block_size.0);
+                if start < end {
+                    res |= ChunkRanges::from(ChunkNum(start)..ChunkNum(end))
+                }
+            }
+        }
+    }
+    res
+}

src/tests.rs

@@ -11,6 +11,7 @@ use range_collections::RangeSet2;
 use crate::{
     assert_tuple_eq, blake3,
     io::{
+        full_chunk_groups,
         outboard::PreOrderMemOutboard,
         sync::{DecodeResponseItem, Outboard},
         Leaf,
@@ -913,6 +914,32 @@ fn encode_last_chunk_cases() {
     }
 }
 
+#[test]
+fn test_full_chunk_groups() {
+    let cases = vec![
+        (
+            ChunkRanges::from(ChunkNum(8)..),
+            ChunkRanges::from(ChunkNum(16)..),
+        ),
+        (
+            ChunkRanges::from(ChunkNum(8)..ChunkNum(16)),
+            ChunkRanges::empty(),
+        ),
+        (
+            ChunkRanges::from(ChunkNum(11)..ChunkNum(34)),
+            ChunkRanges::from(ChunkNum(16)..ChunkNum(32)),
+        ),
+        (
+            ChunkRanges::from(..ChunkNum(35)),
+            ChunkRanges::from(..ChunkNum(32)),
+        ),
+    ];
+    for (case, expected) in cases {
+        let res = full_chunk_groups(&case, BlockSize(4));
+        assert_eq!(res, expected);
+    }
+}
+
 #[test]
 fn sub_chunk_group_query() {
     let tree = BaoTree::new(ByteNum(1024 * 32), BlockSize(4));