1use crate::format::bytes::{read_le_addr as read_addr, read_le_uint as read_uint};
32use crate::format::{FormatError, FormatResult};
33
34pub const BTREE_V1_SIGNATURE: [u8; 4] = *b"TREE";
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct BTreeV1Config {
47 pub sym_leaf_k: u16,
50 pub snode_internal_k: u16,
53 pub chunk_internal_k: u16,
56}
57
58impl Default for BTreeV1Config {
59 fn default() -> Self {
60 Self {
61 sym_leaf_k: 4,
62 snode_internal_k: 16,
63 chunk_internal_k: 32,
64 }
65 }
66}
67
68use crate::format::superblock::symbol_table_entry_size;
69
70impl BTreeV1Config {
71 pub fn sym_leaf_max_entries(&self) -> u16 {
73 self.sym_leaf_k.saturating_mul(2)
74 }
75
76 pub fn snode_max_entries(&self) -> u16 {
78 self.snode_internal_k.saturating_mul(2)
79 }
80
81 pub fn chunk_max_entries(&self) -> u16 {
83 self.chunk_internal_k.saturating_mul(2)
84 }
85
86 pub fn snode_btree_node_size(&self, sizeof_addr: usize, sizeof_size: usize) -> usize {
88 btree_node_size(self.snode_max_entries(), sizeof_addr, sizeof_size)
89 }
90
91 pub fn chunk_btree_node_size(&self, sizeof_addr: usize, rank: usize) -> usize {
94 btree_node_size(
95 self.chunk_max_entries(),
96 sizeof_addr,
97 4 + 4 + (rank + 1) * 8,
98 )
99 }
100
101 pub fn symbol_table_node_size(&self, sizeof_addr: usize, sizeof_size: usize) -> usize {
104 8 + (self.sym_leaf_k as usize) * 2 * symbol_table_entry_size(sizeof_addr, sizeof_size)
105 }
106}
107
108fn btree_node_size(two_k: u16, sizeof_addr: usize, key_size: usize) -> usize {
110 let two_k = two_k as usize;
111 8 + 2 * sizeof_addr + two_k * sizeof_addr + (two_k + 1) * key_size
112}
113
114#[derive(Debug, Clone)]
116pub struct BTreeV1Node {
117 pub node_type: u8,
119 pub level: u8,
121 pub entries_used: u16,
123 pub left_sibling: u64,
125 pub right_sibling: u64,
127 pub keys: Vec<u64>,
129 pub children: Vec<u64>,
131}
132
133impl BTreeV1Node {
134 pub fn encode(
143 &self,
144 node_size: usize,
145 sizeof_addr: usize,
146 sizeof_size: usize,
147 ) -> FormatResult<Vec<u8>> {
148 if self.node_type != 0 {
149 return Err(FormatError::UnsupportedFeature(format!(
150 "B-tree v1 type {} is not encoded by BTreeV1Node (only type 0, \
151 symbol-table nodes)",
152 self.node_type
153 )));
154 }
155 if self.keys.len() != self.children.len() + 1 {
156 return Err(FormatError::InvalidData(format!(
157 "B-tree v1 node has {} keys for {} children; a v1 node stores both \
158 bounds of every child, so it needs exactly one more key than children",
159 self.keys.len(),
160 self.children.len()
161 )));
162 }
163 if self.entries_used as usize != self.children.len() {
164 return Err(FormatError::InvalidData(format!(
165 "B-tree v1 node declares {} entries but carries {} children",
166 self.entries_used,
167 self.children.len()
168 )));
169 }
170 let needed =
171 8 + 2 * sizeof_addr + self.children.len() * (sizeof_size + sizeof_addr) + sizeof_size;
172 if needed > node_size {
173 return Err(FormatError::InvalidData(format!(
174 "B-tree v1 node needs {needed} bytes for {} children, more than the \
175 {node_size}-byte record the file's 'K' value allows",
176 self.children.len()
177 )));
178 }
179
180 let mut buf = Vec::with_capacity(node_size);
181 buf.extend_from_slice(&BTREE_V1_SIGNATURE);
182 buf.push(self.node_type);
183 buf.push(self.level);
184 buf.extend_from_slice(&self.entries_used.to_le_bytes());
185 buf.extend_from_slice(&self.left_sibling.to_le_bytes()[..sizeof_addr]);
186 buf.extend_from_slice(&self.right_sibling.to_le_bytes()[..sizeof_addr]);
187 for (i, &child) in self.children.iter().enumerate() {
188 buf.extend_from_slice(&self.keys[i].to_le_bytes()[..sizeof_size]);
189 buf.extend_from_slice(&child.to_le_bytes()[..sizeof_addr]);
190 }
191 buf.extend_from_slice(&self.keys[self.children.len()].to_le_bytes()[..sizeof_size]);
192 buf.resize(node_size, 0);
193 Ok(buf)
194 }
195
196 pub fn decode(
202 buf: &[u8],
203 sizeof_addr: usize,
204 sizeof_size: usize,
205 max_entries: u16,
206 ) -> FormatResult<Self> {
207 let header_size = 4 + 1 + 1 + 2 + sizeof_addr * 2;
208 if buf.len() < header_size {
209 return Err(FormatError::BufferTooShort {
210 needed: header_size,
211 available: buf.len(),
212 });
213 }
214
215 if buf[0..4] != BTREE_V1_SIGNATURE {
216 return Err(FormatError::InvalidSignature);
217 }
218
219 let node_type = buf[4];
220 let level = buf[5];
221 let entries_used = u16::from_le_bytes([buf[6], buf[7]]);
222 if entries_used > max_entries {
223 return Err(FormatError::InvalidData(format!(
224 "B-tree v1 node declares {entries_used} entries, more than the \
225 {max_entries} its 'K' value allows"
226 )));
227 }
228
229 let mut pos = 8;
230 let left_sibling = read_addr(&buf[pos..], sizeof_addr);
231 pos += sizeof_addr;
232 let right_sibling = read_addr(&buf[pos..], sizeof_addr);
233 pos += sizeof_addr;
234
235 let n = entries_used as usize;
239
240 if node_type == 0 {
241 let key_size = sizeof_size;
243 let child_size = sizeof_addr;
244 let data_size = (n + 1) * key_size + n * child_size;
246 let needed = pos + data_size;
247 if buf.len() < needed {
248 return Err(FormatError::BufferTooShort {
249 needed,
250 available: buf.len(),
251 });
252 }
253
254 let mut keys = Vec::with_capacity(n + 1);
255 let mut children = Vec::with_capacity(n);
256
257 for _i in 0..n {
258 keys.push(read_uint(&buf[pos..], key_size));
260 pos += key_size;
261 children.push(read_uint(&buf[pos..], child_size));
263 pos += child_size;
264 }
265 keys.push(read_uint(&buf[pos..], key_size));
267
268 Ok(BTreeV1Node {
269 node_type,
270 level,
271 entries_used,
272 left_sibling,
273 right_sibling,
274 keys,
275 children,
276 })
277 } else {
278 Err(FormatError::UnsupportedFeature(format!(
282 "B-tree v1 type {} not supported by BTreeV1Node (use ChunkBTreeV1Node)",
283 node_type
284 )))
285 }
286 }
287}
288
289#[derive(Debug, Clone, PartialEq, Eq)]
291pub struct ChunkKey {
292 pub chunk_size: u32,
294 pub filter_mask: u32,
296 pub offsets: Vec<u64>,
300}
301
302impl ChunkKey {
303 pub fn for_chunk(scaled: &[u64], dims: &[u64], chunk_size: u32, filter_mask: u32) -> Self {
312 let mut offsets: Vec<u64> = scaled
313 .iter()
314 .zip(dims)
315 .map(|(&s, &d)| s.saturating_mul(d))
316 .collect();
317 offsets.push(0);
318 Self {
319 chunk_size,
320 filter_mask,
321 offsets,
322 }
323 }
324
325 pub fn right_bound(scaled: &[u64], dims: &[u64]) -> Self {
337 let mut offsets: Vec<u64> = scaled
338 .iter()
339 .zip(dims)
340 .map(|(&s, &d)| s.saturating_mul(d))
341 .collect();
342 offsets.push(*dims.last().unwrap_or(&0));
343 Self {
344 chunk_size: 0,
345 filter_mask: 0,
346 offsets,
347 }
348 }
349
350 fn encoded_size(rank: usize) -> usize {
353 4 + 4 + (rank + 1) * 8
354 }
355
356 fn encode_into(&self, buf: &mut Vec<u8>) {
357 buf.extend_from_slice(&self.chunk_size.to_le_bytes());
358 buf.extend_from_slice(&self.filter_mask.to_le_bytes());
359 for &o in &self.offsets {
360 buf.extend_from_slice(&o.to_le_bytes());
361 }
362 }
363}
364
365#[derive(Debug, Clone)]
367pub struct ChunkBTreeV1Node {
368 pub level: u8,
371 pub entries_used: u16,
373 pub left_sibling: u64,
375 pub right_sibling: u64,
377 pub keys: Vec<ChunkKey>,
380 pub children: Vec<u64>,
382}
383
384impl ChunkBTreeV1Node {
385 pub fn encode(&self, node_size: usize, sizeof_addr: usize) -> FormatResult<Vec<u8>> {
394 if self.keys.len() != self.children.len() + 1 {
395 return Err(FormatError::InvalidData(format!(
396 "chunk B-tree v1 node has {} keys for {} children; a v1 node stores \
397 both bounds of every child, so it needs exactly one more key than \
398 children",
399 self.keys.len(),
400 self.children.len()
401 )));
402 }
403 if self.entries_used as usize != self.children.len() {
404 return Err(FormatError::InvalidData(format!(
405 "chunk B-tree v1 node declares {} entries but carries {} children",
406 self.entries_used,
407 self.children.len()
408 )));
409 }
410 let key_size = self.keys[0].offsets.len() * 8 + 8;
413 if let Some(k) = self
414 .keys
415 .iter()
416 .find(|k| k.offsets.len() * 8 + 8 != key_size)
417 {
418 return Err(FormatError::InvalidData(format!(
419 "chunk B-tree v1 node mixes keys of {} and {} offsets; every key in \
420 one tree describes the same dataset",
421 self.keys[0].offsets.len(),
422 k.offsets.len()
423 )));
424 }
425 let needed =
426 8 + 2 * sizeof_addr + self.children.len() * sizeof_addr + self.keys.len() * key_size;
427 if needed > node_size {
428 return Err(FormatError::InvalidData(format!(
429 "chunk B-tree v1 node needs {needed} bytes for {} children, more than \
430 the {node_size}-byte record the file's 'K' value allows",
431 self.children.len()
432 )));
433 }
434
435 let mut buf = Vec::with_capacity(node_size);
436 buf.extend_from_slice(&BTREE_V1_SIGNATURE);
437 buf.push(1);
438 buf.push(self.level);
439 buf.extend_from_slice(&self.entries_used.to_le_bytes());
440 buf.extend_from_slice(&self.left_sibling.to_le_bytes()[..sizeof_addr]);
441 buf.extend_from_slice(&self.right_sibling.to_le_bytes()[..sizeof_addr]);
442 for (i, &child) in self.children.iter().enumerate() {
443 self.keys[i].encode_into(&mut buf);
444 buf.extend_from_slice(&child.to_le_bytes()[..sizeof_addr]);
445 }
446 self.keys[self.children.len()].encode_into(&mut buf);
447 buf.resize(node_size, 0);
448 Ok(buf)
449 }
450
451 pub fn decode(
459 buf: &[u8],
460 sizeof_addr: usize,
461 rank: usize,
462 max_entries: u16,
463 ) -> FormatResult<Self> {
464 let header_size = 4 + 1 + 1 + 2 + sizeof_addr * 2;
465 if buf.len() < header_size {
466 return Err(FormatError::BufferTooShort {
467 needed: header_size,
468 available: buf.len(),
469 });
470 }
471
472 if buf[0..4] != BTREE_V1_SIGNATURE {
473 return Err(FormatError::InvalidSignature);
474 }
475
476 let node_type = buf[4];
477 if node_type != 1 {
478 return Err(FormatError::UnsupportedFeature(format!(
479 "expected B-tree v1 chunk node (type 1), found type {node_type}"
480 )));
481 }
482 let level = buf[5];
483 let entries_used = u16::from_le_bytes([buf[6], buf[7]]);
484 if entries_used > max_entries {
485 return Err(FormatError::InvalidData(format!(
486 "chunk B-tree v1 node declares {entries_used} entries, more than \
487 the {max_entries} its 'K' value allows"
488 )));
489 }
490
491 let mut pos = 8;
492 let left_sibling = read_addr(&buf[pos..], sizeof_addr);
493 pos += sizeof_addr;
494 let right_sibling = read_addr(&buf[pos..], sizeof_addr);
495 pos += sizeof_addr;
496
497 let n = entries_used as usize;
498 let key_size = ChunkKey::encoded_size(rank);
500 let data_size = (n + 1) * key_size + n * sizeof_addr;
502 let needed = pos + data_size;
503 if buf.len() < needed {
504 return Err(FormatError::BufferTooShort {
505 needed,
506 available: buf.len(),
507 });
508 }
509
510 let decode_key = |slice: &[u8]| -> ChunkKey {
511 let chunk_size = u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]);
512 let filter_mask = u32::from_le_bytes([slice[4], slice[5], slice[6], slice[7]]);
513 let mut offsets = Vec::with_capacity(rank + 1);
514 let mut o = 8;
515 for _ in 0..(rank + 1) {
516 offsets.push(read_uint(&slice[o..], 8));
517 o += 8;
518 }
519 ChunkKey {
520 chunk_size,
521 filter_mask,
522 offsets,
523 }
524 };
525
526 let mut keys = Vec::with_capacity(n + 1);
527 let mut children = Vec::with_capacity(n);
528 for _ in 0..n {
529 keys.push(decode_key(&buf[pos..pos + key_size]));
530 pos += key_size;
531 children.push(read_addr(&buf[pos..], sizeof_addr));
532 pos += sizeof_addr;
533 }
534 keys.push(decode_key(&buf[pos..pos + key_size]));
536
537 Ok(ChunkBTreeV1Node {
538 level,
539 entries_used,
540 left_sibling,
541 right_sibling,
542 keys,
543 children,
544 })
545 }
546}
547
548pub struct ChunkBTreeV1Tree {
562 nodes: Vec<TreeNode>,
564 node_size: usize,
566 sizeof_addr: usize,
567}
568
569struct TreeNode {
571 level: u8,
572 keys: Vec<ChunkKey>,
575 children: TreeChildren,
576 left: Option<usize>,
578 right: Option<usize>,
579}
580
581enum TreeChildren {
585 Chunks(Vec<u64>),
587 Nodes(Vec<usize>),
589}
590
591impl TreeChildren {
592 fn len(&self) -> usize {
593 match self {
594 Self::Chunks(v) => v.len(),
595 Self::Nodes(v) => v.len(),
596 }
597 }
598}
599
600impl ChunkBTreeV1Tree {
601 pub fn build(
615 entries: &[(ChunkKey, u64)],
616 end_key: ChunkKey,
617 config: &BTreeV1Config,
618 sizeof_addr: usize,
619 ) -> Self {
620 let rank = end_key.offsets.len().saturating_sub(1);
621 let node_size = config.chunk_btree_node_size(sizeof_addr, rank);
622 let cap = (config.chunk_max_entries() as usize).max(1);
623 let mut nodes: Vec<TreeNode> = Vec::new();
624
625 if !entries.is_empty() {
626 let mut level_range = spread(entries.len(), cap)
628 .into_iter()
629 .scan(0usize, |start, m| {
630 let range = *start..*start + m;
631 *start += m;
632 Some(range)
633 })
634 .map(|r| {
635 let mut keys: Vec<ChunkKey> =
636 entries[r.clone()].iter().map(|(k, _)| k.clone()).collect();
637 keys.push(match entries.get(r.end) {
638 Some((k, _)) => k.clone(),
639 None => end_key.clone(),
640 });
641 TreeNode {
642 level: 0,
643 keys,
644 children: TreeChildren::Chunks(
645 entries[r].iter().map(|&(_, a)| a).collect(),
646 ),
647 left: None,
648 right: None,
649 }
650 })
651 .collect::<Vec<_>>();
652 let mut level: u8 = 0;
653 loop {
654 let base = nodes.len();
655 let count = level_range.len();
656 for (i, mut node) in level_range.into_iter().enumerate() {
657 node.left = (i > 0).then(|| base + i - 1);
658 node.right = (i + 1 < count).then(|| base + i + 1);
659 nodes.push(node);
660 }
661 if count == 1 {
662 break;
663 }
664 let children: Vec<usize> = (base..base + count).collect();
667 level += 1;
668 let mut start = 0usize;
669 level_range = spread(count, cap)
670 .into_iter()
671 .map(|m| {
672 let run = &children[start..start + m];
673 start += m;
674 let mut keys: Vec<ChunkKey> =
675 run.iter().map(|&c| nodes[c].keys[0].clone()).collect();
676 keys.push(match children.get(start) {
677 Some(&next) => nodes[next].keys[0].clone(),
678 None => end_key.clone(),
679 });
680 TreeNode {
681 level,
682 keys,
683 children: TreeChildren::Nodes(run.to_vec()),
684 left: None,
685 right: None,
686 }
687 })
688 .collect();
689 }
690 }
691
692 Self {
693 nodes,
694 node_size,
695 sizeof_addr,
696 }
697 }
698
699 pub fn node_count(&self) -> usize {
701 self.nodes.len()
702 }
703
704 pub fn node_size(&self) -> usize {
706 self.node_size
707 }
708
709 pub fn root_address(&self, addrs: &[u64]) -> u64 {
712 match self.nodes.len() {
713 0 => crate::format::UNDEF_ADDR,
714 n => addrs[n - 1],
715 }
716 }
717
718 pub fn encode(&self, addrs: &[u64]) -> FormatResult<Vec<Vec<u8>>> {
723 let sibling = |i: Option<usize>| i.map_or(crate::format::UNDEF_ADDR, |j| addrs[j]);
724 self.nodes
725 .iter()
726 .map(|n| {
727 let children = match &n.children {
728 TreeChildren::Chunks(v) => v.clone(),
729 TreeChildren::Nodes(v) => v.iter().map(|&j| addrs[j]).collect(),
730 };
731 ChunkBTreeV1Node {
732 level: n.level,
733 entries_used: n.children.len() as u16,
734 left_sibling: sibling(n.left),
735 right_sibling: sibling(n.right),
736 keys: n.keys.clone(),
737 children,
738 }
739 .encode(self.node_size, self.sizeof_addr)
740 })
741 .collect()
742 }
743}
744
745fn spread(n: usize, cap: usize) -> Vec<usize> {
749 let k = n.div_ceil(cap);
750 if k == 0 {
751 return Vec::new();
752 }
753 let (base, extra) = (n / k, n % k);
754 (0..k).map(|i| base + usize::from(i < extra)).collect()
755}
756
757#[cfg(test)]
760mod tests {
761 use super::*;
762 use crate::format::UNDEF_ADDR;
763
764 fn build_group_btree(
766 level: u8,
767 keys: &[u64],
768 children: &[u64],
769 sizeof_addr: usize,
770 sizeof_size: usize,
771 ) -> Vec<u8> {
772 assert_eq!(keys.len(), children.len() + 1);
773 let entries_used = children.len() as u16;
774
775 let mut buf = Vec::new();
776 buf.extend_from_slice(&BTREE_V1_SIGNATURE);
777 buf.push(0); buf.push(level);
779 buf.extend_from_slice(&entries_used.to_le_bytes());
780 buf.extend_from_slice(&UNDEF_ADDR.to_le_bytes()[..sizeof_addr]);
782 buf.extend_from_slice(&UNDEF_ADDR.to_le_bytes()[..sizeof_addr]);
784
785 for i in 0..children.len() {
787 buf.extend_from_slice(&keys[i].to_le_bytes()[..sizeof_size]);
788 buf.extend_from_slice(&children[i].to_le_bytes()[..sizeof_addr]);
789 }
790 buf.extend_from_slice(&keys[children.len()].to_le_bytes()[..sizeof_size]);
792
793 buf
794 }
795
796 #[test]
797 fn decode_leaf_node() {
798 let buf = build_group_btree(
799 0, &[0, 8, 16], &[0x100, 0x200], 8,
803 8,
804 );
805 let node = BTreeV1Node::decode(&buf, 8, 8, 32).unwrap();
806 assert_eq!(node.node_type, 0);
807 assert_eq!(node.level, 0);
808 assert_eq!(node.entries_used, 2);
809 assert_eq!(node.keys, vec![0, 8, 16]);
810 assert_eq!(node.children, vec![0x100, 0x200]);
811 assert_eq!(node.left_sibling, UNDEF_ADDR);
812 assert_eq!(node.right_sibling, UNDEF_ADDR);
813 }
814
815 #[test]
816 fn decode_internal_node() {
817 let buf = build_group_btree(
818 1, &[0, 100], &[0x500], 8,
822 8,
823 );
824 let node = BTreeV1Node::decode(&buf, 8, 8, 32).unwrap();
825 assert_eq!(node.level, 1);
826 assert_eq!(node.entries_used, 1);
827 assert_eq!(node.children, vec![0x500]);
828 }
829
830 #[test]
831 fn decode_single_entry() {
832 let buf = build_group_btree(0, &[0, 8], &[0x100], 8, 8);
833 let node = BTreeV1Node::decode(&buf, 8, 8, 32).unwrap();
834 assert_eq!(node.entries_used, 1);
835 assert_eq!(node.children.len(), 1);
836 }
837
838 #[test]
839 fn decode_4byte() {
840 let buf = build_group_btree(0, &[0, 4], &[0x80], 4, 4);
841 let node = BTreeV1Node::decode(&buf, 4, 4, 32).unwrap();
842 assert_eq!(node.entries_used, 1);
843 assert_eq!(node.children, vec![0x80]);
844 }
845
846 #[test]
847 fn decode_bad_sig() {
848 let mut buf = build_group_btree(0, &[0, 8], &[0x100], 8, 8);
849 buf[0] = b'X';
850 assert!(matches!(
851 BTreeV1Node::decode(&buf, 8, 8, 32).unwrap_err(),
852 FormatError::InvalidSignature
853 ));
854 }
855
856 #[test]
857 fn decode_too_short() {
858 assert!(matches!(
859 BTreeV1Node::decode(&[0u8; 4], 8, 8, 32).unwrap_err(),
860 FormatError::BufferTooShort { .. }
861 ));
862 }
863
864 #[test]
865 fn decode_unsupported_type() {
866 let mut buf = build_group_btree(0, &[0, 8], &[0x100], 8, 8);
867 buf[4] = 1; assert!(matches!(
869 BTreeV1Node::decode(&buf, 8, 8, 32).unwrap_err(),
870 FormatError::UnsupportedFeature(_)
871 ));
872 }
873
874 fn build_chunk_btree(
877 level: u8,
878 keys: &[ChunkKey],
879 children: &[u64],
880 sizeof_addr: usize,
881 ) -> Vec<u8> {
882 assert_eq!(keys.len(), children.len() + 1);
883 let entries_used = children.len() as u16;
884
885 let mut buf = Vec::new();
886 buf.extend_from_slice(&BTREE_V1_SIGNATURE);
887 buf.push(1); buf.push(level);
889 buf.extend_from_slice(&entries_used.to_le_bytes());
890 buf.extend_from_slice(&UNDEF_ADDR.to_le_bytes()[..sizeof_addr]); buf.extend_from_slice(&UNDEF_ADDR.to_le_bytes()[..sizeof_addr]); let encode_key = |buf: &mut Vec<u8>, k: &ChunkKey| {
894 buf.extend_from_slice(&k.chunk_size.to_le_bytes());
895 buf.extend_from_slice(&k.filter_mask.to_le_bytes());
896 for &o in &k.offsets {
897 buf.extend_from_slice(&o.to_le_bytes());
898 }
899 };
900
901 for i in 0..children.len() {
902 encode_key(&mut buf, &keys[i]);
903 buf.extend_from_slice(&children[i].to_le_bytes()[..sizeof_addr]);
904 }
905 encode_key(&mut buf, &keys[children.len()]);
906 buf
907 }
908
909 fn chunk_key(size: u32, mask: u32, offsets: &[u64]) -> ChunkKey {
910 ChunkKey {
911 chunk_size: size,
912 filter_mask: mask,
913 offsets: offsets.to_vec(),
914 }
915 }
916
917 #[test]
918 fn decode_chunk_leaf_1d() {
919 let keys = [
921 chunk_key(32, 0, &[0, 0]),
922 chunk_key(32, 0, &[8, 0]),
923 chunk_key(0, 0, &[16, 0]),
924 ];
925 let buf = build_chunk_btree(0, &keys, &[0x400, 0x800], 8);
926 let node = ChunkBTreeV1Node::decode(&buf, 8, 1, 64).unwrap();
927 assert_eq!(node.level, 0);
928 assert_eq!(node.entries_used, 2);
929 assert_eq!(node.children, vec![0x400, 0x800]);
930 assert_eq!(node.keys.len(), 3);
931 assert_eq!(node.keys[0].chunk_size, 32);
932 assert_eq!(node.keys[1].offsets, vec![8, 0]);
933 }
934
935 #[test]
936 fn decode_chunk_internal_2d() {
937 let keys = [chunk_key(64, 0, &[0, 0, 0]), chunk_key(64, 0, &[4, 4, 0])];
939 let buf = build_chunk_btree(1, &keys, &[0x1000], 8);
940 let node = ChunkBTreeV1Node::decode(&buf, 8, 2, 64).unwrap();
941 assert_eq!(node.level, 1);
942 assert_eq!(node.entries_used, 1);
943 assert_eq!(node.children, vec![0x1000]);
944 assert_eq!(node.keys[0].offsets, vec![0, 0, 0]);
945 }
946
947 #[test]
948 fn decode_chunk_filtered_key() {
949 let keys = [chunk_key(17, 0x1, &[0, 0]), chunk_key(0, 0, &[8, 0])];
950 let buf = build_chunk_btree(0, &keys, &[0x200], 8);
951 let node = ChunkBTreeV1Node::decode(&buf, 8, 1, 64).unwrap();
952 assert_eq!(node.keys[0].chunk_size, 17);
953 assert_eq!(node.keys[0].filter_mask, 0x1);
954 }
955
956 #[test]
957 fn decode_chunk_rejects_group_node() {
958 let buf = build_group_btree(0, &[0, 8], &[0x100], 8, 8);
959 assert!(matches!(
960 ChunkBTreeV1Node::decode(&buf, 8, 1, 64).unwrap_err(),
961 FormatError::UnsupportedFeature(_)
962 ));
963 }
964
965 #[test]
966 fn decode_chunk_too_short() {
967 assert!(matches!(
968 ChunkBTreeV1Node::decode(&[0u8; 4], 8, 1, 64).unwrap_err(),
969 FormatError::BufferTooShort { .. }
970 ));
971 }
972
973 #[test]
974 fn decode_chunk_bad_sig() {
975 let keys = [chunk_key(8, 0, &[0, 0]), chunk_key(0, 0, &[8, 0])];
976 let mut buf = build_chunk_btree(0, &keys, &[0x100], 8);
977 buf[0] = b'X';
978 assert!(matches!(
979 ChunkBTreeV1Node::decode(&buf, 8, 1, 64).unwrap_err(),
980 FormatError::InvalidSignature
981 ));
982 }
983
984 #[test]
985 fn node_sizes_match_upstream_formula() {
986 let cfg = BTreeV1Config::default();
987 assert_eq!(cfg.snode_btree_node_size(8, 8), 8 + 16 + 32 * 8 + 33 * 8);
989 assert_eq!(cfg.chunk_btree_node_size(8, 1), 8 + 16 + 64 * 8 + 65 * 24);
991 assert_eq!(cfg.symbol_table_node_size(8, 8), 8 + 8 * 40);
993 }
994
995 #[test]
996 fn non_default_k_scales_every_node_size() {
997 let cfg = BTreeV1Config {
998 sym_leaf_k: 128,
999 snode_internal_k: 512,
1000 chunk_internal_k: 256,
1001 };
1002 assert_eq!(cfg.snode_max_entries(), 1024);
1003 assert_eq!(cfg.chunk_max_entries(), 512);
1004 assert!(cfg.snode_btree_node_size(8, 8) > 8192);
1007 assert!(cfg.chunk_btree_node_size(8, 1) > 8192);
1008 assert!(cfg.symbol_table_node_size(8, 8) > 8192);
1009 }
1010
1011 #[test]
1012 fn decode_rejects_entries_beyond_two_k() {
1013 let buf = build_group_btree(0, &[0, 8, 16], &[0x100, 0x200], 8, 8);
1014 assert!(matches!(
1016 BTreeV1Node::decode(&buf, 8, 8, 0).unwrap_err(),
1017 FormatError::InvalidData(_)
1018 ));
1019 assert!(BTreeV1Node::decode(&buf, 8, 8, 2).is_ok());
1021 }
1022
1023 #[test]
1024 fn decode_chunk_rejects_entries_beyond_two_k() {
1025 let keys = [
1026 chunk_key(32, 0, &[0, 0]),
1027 chunk_key(32, 0, &[8, 0]),
1028 chunk_key(0, 0, &[16, 0]),
1029 ];
1030 let buf = build_chunk_btree(0, &keys, &[0x400, 0x800], 8);
1031 assert!(matches!(
1032 ChunkBTreeV1Node::decode(&buf, 8, 1, 1).unwrap_err(),
1033 FormatError::InvalidData(_)
1034 ));
1035 assert!(ChunkBTreeV1Node::decode(&buf, 8, 1, 2).is_ok());
1036 }
1037
1038 #[test]
1039 fn decode_chunk_4byte_addr() {
1040 let keys = [chunk_key(16, 0, &[0, 0]), chunk_key(0, 0, &[4, 0])];
1041 let buf = build_chunk_btree(0, &keys, &[0x80], 4);
1042 let node = ChunkBTreeV1Node::decode(&buf, 4, 1, 64).unwrap();
1043 assert_eq!(node.children, vec![0x80]);
1044 }
1045
1046 #[test]
1050 fn an_encoded_group_btree_node_matches_the_bytes_libhdf5_wrote() {
1051 let node = BTreeV1Node {
1052 node_type: 0,
1053 level: 0,
1054 entries_used: 1,
1055 left_sibling: UNDEF_ADDR,
1056 right_sibling: UNDEF_ADDR,
1057 keys: vec![0, 24],
1058 children: vec![0x430],
1059 };
1060 let node_size = BTreeV1Config::default().snode_btree_node_size(8, 8);
1061 assert_eq!(node_size, 544);
1062 let encoded = node.encode(node_size, 8, 8).unwrap();
1063 let mut expected = Vec::new();
1064 expected.extend_from_slice(b"TREE");
1065 expected.extend_from_slice(&[0, 0, 1, 0]); expected.extend_from_slice(&[0xff; 16]); expected.extend_from_slice(&0u64.to_le_bytes()); expected.extend_from_slice(&0x430u64.to_le_bytes()); expected.extend_from_slice(&24u64.to_le_bytes()); assert_eq!(&encoded[..expected.len()], &expected[..]);
1071 assert!(encoded[expected.len()..].iter().all(|&b| b == 0));
1072 assert_eq!(encoded.len(), node_size);
1073 }
1074
1075 #[test]
1079 fn an_encoded_group_btree_node_round_trips_an_interior_level() {
1080 let cfg = BTreeV1Config::default();
1081 let node_size = cfg.snode_btree_node_size(8, 8);
1082 let node = BTreeV1Node {
1083 node_type: 0,
1084 level: 1,
1085 entries_used: 2,
1086 left_sibling: UNDEF_ADDR,
1087 right_sibling: UNDEF_ADDR,
1088 keys: vec![0, 896, 1600],
1089 children: vec![0x1a2a8, 0x1a088],
1090 };
1091 let encoded = node.encode(node_size, 8, 8).unwrap();
1092 let decoded = BTreeV1Node::decode(&encoded, 8, 8, cfg.snode_max_entries()).unwrap();
1093 assert_eq!(decoded.level, 1);
1094 assert_eq!(decoded.entries_used, 2);
1095 assert_eq!(decoded.keys, node.keys);
1096 assert_eq!(decoded.children, node.children);
1097 assert_eq!(decoded.left_sibling, UNDEF_ADDR);
1098 }
1099
1100 #[test]
1103 fn an_encoded_group_btree_node_round_trips_an_empty_root() {
1104 let cfg = BTreeV1Config::default();
1105 let node = BTreeV1Node {
1106 node_type: 0,
1107 level: 0,
1108 entries_used: 0,
1109 left_sibling: UNDEF_ADDR,
1110 right_sibling: UNDEF_ADDR,
1111 keys: vec![0],
1112 children: vec![],
1113 };
1114 let encoded = node.encode(cfg.snode_btree_node_size(8, 8), 8, 8).unwrap();
1115 let decoded = BTreeV1Node::decode(&encoded, 8, 8, cfg.snode_max_entries()).unwrap();
1116 assert_eq!(decoded.entries_used, 0);
1117 assert!(decoded.children.is_empty());
1118 assert_eq!(decoded.keys, vec![0]);
1119 }
1120
1121 #[test]
1125 fn a_group_btree_node_refuses_a_key_count_that_does_not_bound_its_children() {
1126 let cfg = BTreeV1Config::default();
1127 let node = BTreeV1Node {
1128 node_type: 0,
1129 level: 0,
1130 entries_used: 2,
1131 left_sibling: UNDEF_ADDR,
1132 right_sibling: UNDEF_ADDR,
1133 keys: vec![0, 8],
1134 children: vec![0x400, 0x800],
1135 };
1136 assert!(matches!(
1137 node.encode(cfg.snode_btree_node_size(8, 8), 8, 8)
1138 .unwrap_err(),
1139 FormatError::InvalidData(_)
1140 ));
1141 }
1142
1143 fn dense_1d_tree(nchunks: u64, chunk: u64, elem: u64, cfg: &BTreeV1Config) -> ChunkBTreeV1Tree {
1147 let dims = [chunk, elem];
1148 let nbytes = (chunk * elem) as u32;
1149 let entries: Vec<(ChunkKey, u64)> = (0..nchunks)
1150 .map(|i| {
1151 (
1152 ChunkKey::for_chunk(&[i], &dims, nbytes, 0),
1153 0x1000 + i * chunk * elem,
1154 )
1155 })
1156 .collect();
1157 let end = ChunkKey::right_bound(&[nchunks - 1], &dims);
1158 ChunkBTreeV1Tree::build(&entries, end, cfg, 8)
1159 }
1160
1161 #[test]
1166 fn a_bulk_loaded_chunk_tree_keys_the_way_libhdf5_does() {
1167 let cfg = BTreeV1Config::default();
1168 let tree = dense_1d_tree(2, 4, 4, &cfg);
1169 assert_eq!(tree.node_count(), 1);
1170 assert_eq!(tree.node_size(), cfg.chunk_btree_node_size(8, 1));
1171
1172 let addrs = [0x578u64];
1173 let images = tree.encode(&addrs).unwrap();
1174 let node = ChunkBTreeV1Node::decode(&images[0], 8, 1, cfg.chunk_max_entries()).unwrap();
1175 assert_eq!(images[0].len(), tree.node_size());
1176 assert_eq!(node.level, 0);
1177 assert_eq!(node.entries_used, 2);
1178 assert_eq!(node.children, vec![0x1000, 0x1010]);
1179 assert_eq!(node.left_sibling, UNDEF_ADDR);
1180 assert_eq!(node.right_sibling, UNDEF_ADDR);
1181 let offsets: Vec<&[u64]> = node.keys.iter().map(|k| k.offsets.as_slice()).collect();
1182 assert_eq!(offsets, vec![&[0, 0], &[4, 0], &[4, 4]]);
1183 assert_eq!(
1184 node.keys.iter().map(|k| k.chunk_size).collect::<Vec<_>>(),
1185 vec![16, 16, 0]
1186 );
1187 assert_eq!(tree.root_address(&addrs), 0x578);
1188 }
1189
1190 #[test]
1195 fn a_bulk_loaded_chunk_tree_grows_a_level_past_2k() {
1196 let cfg = BTreeV1Config::default();
1197 let two_k = cfg.chunk_max_entries() as u64;
1198 let nchunks = two_k * 3 + 1;
1199 let tree = dense_1d_tree(nchunks, 4, 4, &cfg);
1200 assert_eq!(tree.node_count(), 5);
1203
1204 let addrs: Vec<u64> = (0..tree.node_count() as u64)
1205 .map(|i| 0x1_0000 + i * 4096)
1206 .collect();
1207 let images = tree.encode(&addrs).unwrap();
1208 let nodes: Vec<ChunkBTreeV1Node> = images
1209 .iter()
1210 .map(|img| ChunkBTreeV1Node::decode(img, 8, 1, cfg.chunk_max_entries()).unwrap())
1211 .collect();
1212
1213 let (leaves, root) = nodes.split_at(4);
1214 let root = &root[0];
1215 assert!(leaves.iter().all(|n| n.level == 0));
1216 assert_eq!(root.level, 1);
1217 assert_eq!(root.entries_used, 4);
1218 assert_eq!(root.children, addrs[..4]);
1219 assert_eq!(root.left_sibling, UNDEF_ADDR);
1220 assert_eq!(root.right_sibling, UNDEF_ADDR);
1221
1222 assert_eq!(
1224 leaves.iter().map(|n| n.entries_used as u64).sum::<u64>(),
1225 nchunks
1226 );
1227 assert!(leaves.iter().all(|n| n.entries_used as u64 <= two_k));
1228
1229 for (i, leaf) in leaves.iter().enumerate() {
1232 assert_eq!(
1233 leaf.left_sibling,
1234 if i == 0 { UNDEF_ADDR } else { addrs[i - 1] }
1235 );
1236 assert_eq!(
1237 leaf.right_sibling,
1238 if i + 1 == leaves.len() {
1239 UNDEF_ADDR
1240 } else {
1241 addrs[i + 1]
1242 }
1243 );
1244 assert_eq!(root.keys[i], leaf.keys[0]);
1245 if let Some(next) = leaves.get(i + 1) {
1246 assert_eq!(leaf.keys[leaf.keys.len() - 1], next.keys[0]);
1247 }
1248 }
1249 let end = ChunkKey::right_bound(&[nchunks - 1], &[4, 4]);
1252 assert_eq!(*leaves[3].keys.last().unwrap(), end);
1253 assert_eq!(*root.keys.last().unwrap(), end);
1254 assert_eq!(end.offsets, vec![(nchunks - 1) * 4, 4]);
1255 }
1256
1257 #[test]
1261 fn an_empty_chunk_tree_has_no_root() {
1262 let cfg = BTreeV1Config::default();
1263 let tree = ChunkBTreeV1Tree::build(&[], ChunkKey::right_bound(&[0], &[4, 4]), &cfg, 8);
1264 assert_eq!(tree.node_count(), 0);
1265 assert!(tree.encode(&[]).unwrap().is_empty());
1266 assert_eq!(tree.root_address(&[]), UNDEF_ADDR);
1267 }
1268
1269 #[test]
1273 fn a_chunk_node_wider_than_its_record_is_refused() {
1274 let node = ChunkBTreeV1Node {
1275 level: 0,
1276 entries_used: 2,
1277 left_sibling: UNDEF_ADDR,
1278 right_sibling: UNDEF_ADDR,
1279 keys: (0..3)
1280 .map(|i| ChunkKey::for_chunk(&[i], &[4, 4], 16, 0))
1281 .collect(),
1282 children: vec![0x100, 0x200],
1283 };
1284 assert!(matches!(
1285 node.encode(64, 8).unwrap_err(),
1286 FormatError::InvalidData(_)
1287 ));
1288 let mut broken = node.clone();
1291 broken.keys.pop();
1292 assert!(matches!(
1293 broken.encode(4096, 8).unwrap_err(),
1294 FormatError::InvalidData(_)
1295 ));
1296 }
1297
1298 #[test]
1301 fn a_chunk_btree_node_is_not_encoded_by_the_group_encoder() {
1302 let node = BTreeV1Node {
1303 node_type: 1,
1304 level: 0,
1305 entries_used: 0,
1306 left_sibling: UNDEF_ADDR,
1307 right_sibling: UNDEF_ADDR,
1308 keys: vec![0],
1309 children: vec![],
1310 };
1311 assert!(matches!(
1312 node.encode(4096, 8, 8).unwrap_err(),
1313 FormatError::UnsupportedFeature(_)
1314 ));
1315 }
1316}