1use super::*;
2use crate::storage::overflow::OverflowChain;
3
4impl<P: PageStore> BTree<P> {
5 fn free_chain_if_pointer(&self, stored: &[u8]) -> BTreeResult<()> {
17 if let Some(head) = value_layout::pointer_head(stored) {
18 OverflowChain::new(self.pager.as_ref())
19 .free(head)
20 .map_err(value_layout::ValueLayoutError::from)?;
21 }
22 Ok(())
23 }
24 pub fn new(pager: Arc<P>) -> Self {
26 Self {
27 retirement_pin: Arc::new(()),
28 pager,
29 root_page_id: RwLock::new(0),
30 rightmost_leaf: RwLock::new(None),
31 writer: parking_lot::ReentrantMutex::new(()),
32 }
33 }
34
35 pub fn with_root(pager: Arc<P>, root_page_id: u32) -> Self {
37 Self {
38 retirement_pin: Arc::new(()),
39 pager,
40 root_page_id: RwLock::new(root_page_id),
41 rightmost_leaf: RwLock::new(None),
42 writer: parking_lot::ReentrantMutex::new(()),
43 }
44 }
45
46 pub fn root_page_id(&self) -> u32 {
48 *self
49 .root_page_id
50 .read()
51 .unwrap_or_else(|poisoned| poisoned.into_inner())
52 }
53
54 pub fn is_empty(&self) -> bool {
56 self.root_page_id() == 0
57 }
58
59 pub fn get(&self, key: &[u8]) -> BTreeResult<Option<Vec<u8>>> {
61 self.get_with_budget(key, usize::MAX)
62 }
63
64 pub fn get_with_budget(&self, key: &[u8], limit_bytes: usize) -> BTreeResult<Option<Vec<u8>>> {
66 let root_id = self.root_page_id();
67 if root_id == 0 {
68 return Ok(None);
69 }
70
71 let (leaf_id, _) = self.find_leaf(root_id, key)?;
73
74 const MAX_RIGHTLINK_HOPS: usize = 32;
83 let mut current_id = leaf_id;
84 for _ in 0..MAX_RIGHTLINK_HOPS {
85 let page = self.pager.read_page(current_id)?;
86
87 match search_leaf(&page, key)? {
88 SearchResult::Found(pos) => {
89 let (_, stored) = read_leaf_cell(&page, pos)?;
90 let value = value_layout::decode_with_budget(
91 self.pager.as_ref(),
92 &stored,
93 limit_bytes,
94 )?;
95 return Ok(Some(value));
96 }
97 SearchResult::NotFound(_) => {
98 let right = leaf_right_sibling(&page);
100 if right != 0 {
101 if let Some(high_key) = leaf_high_key(&page)? {
102 if key > high_key.as_slice() {
103 current_id = right;
105 continue;
106 }
107 }
108 }
109 return Ok(None);
111 }
112 }
113 }
114
115 let (leaf_id, _) = self.find_leaf(self.root_page_id(), key)?;
117 let page = self.pager.read_page(leaf_id)?;
118 match search_leaf(&page, key)? {
119 SearchResult::Found(pos) => {
120 let (_, stored) = read_leaf_cell(&page, pos)?;
121 let value =
122 value_layout::decode_with_budget(self.pager.as_ref(), &stored, limit_bytes)?;
123 Ok(Some(value))
124 }
125 SearchResult::NotFound(_) => Ok(None),
126 }
127 }
128
129 pub fn insert(&self, key: &[u8], value: &[u8]) -> BTreeResult<()> {
138 let _writer = self.writer.lock();
139 if key.len() > MAX_KEY_SIZE {
140 return Err(BTreeError::KeyTooLarge(key.len()));
141 }
142 let stored = value_layout::encode(self.pager.as_ref(), value)?;
143 self.insert_encoded(key, &stored)
144 }
145
146 pub(crate) fn insert_encoded(&self, key: &[u8], value: &[u8]) -> BTreeResult<()> {
154 let _writer = self.writer.lock();
155 let root_id = self.root_page_id();
156
157 if root_id == 0 {
159 let mut page = self.pager.allocate_page(PageType::BTreeLeaf)?;
160 clear_leaf_cells(&mut page);
161 insert_into_leaf(&mut page, key, value)?;
162 init_leaf_links(&mut page, 0, 0);
163 page.update_checksum();
164 let new_root = page.page_id();
165 self.pager.write_page(new_root, page)?;
166 *self.root_page_id.write().map_err(|e| {
167 BTreeError::LockPoisoned(format!("insert: root_page_id write lock: {e}"))
168 })? = new_root;
169 return Ok(());
170 }
171
172 let cached = {
177 self.rightmost_leaf
178 .read()
179 .map_err(|e| BTreeError::LockPoisoned(format!("rightmost_leaf read: {e}")))?
180 .clone()
181 };
182 if let Some((cached_page_id, ref high_key)) = cached {
183 if key > high_key.as_slice() {
184 let mut page = self.pager.read_page(cached_page_id)?;
185 let right_sibling = leaf_right_sibling(&page);
189 if right_sibling == 0 {
190 if can_insert_leaf(&page, key, value) {
192 insert_into_leaf(&mut page, key, value)?;
193 let page_id = page.page_id();
197 *self.rightmost_leaf.write().map_err(|e| {
199 BTreeError::LockPoisoned(format!("rightmost_leaf write: {e}"))
200 })? = Some((page_id, key.to_vec()));
201 self.pager.write_page(page_id, page)?;
202 return Ok(());
203 }
204 let (new_leaf, separator_key) = self.split_leaf(&mut page, key, value)?;
207 let new_leaf_id = new_leaf.page_id();
208 page.update_checksum();
209 let page_id = page.page_id();
210 self.pager.write_page(page_id, page.clone())?;
211 self.cache_rightmost_leaf(&new_leaf)?;
213 let (_, path) = self.find_leaf(root_id, &separator_key)?;
216 self.insert_into_parent(path, cached_page_id, &separator_key, new_leaf_id)?;
217 return Ok(());
218 }
219 *self.rightmost_leaf.write().map_err(|e| {
223 BTreeError::LockPoisoned(format!("rightmost_leaf write: {e}"))
224 })? = None;
225 }
226 }
227
228 let (leaf_id, path) = self.find_leaf(root_id, key)?;
230 let mut page = self.pager.read_page(leaf_id)?;
231
232 if let SearchResult::Found(_) = search_leaf(&page, key)? {
234 return Err(BTreeError::DuplicateKey);
235 }
236
237 if can_insert_leaf(&page, key, value) {
239 insert_into_leaf(&mut page, key, value)?;
240 let page_id = page.page_id();
242 if leaf_right_sibling(&page) == 0 {
244 self.cache_rightmost_leaf(&page)?;
245 }
246 self.pager.write_page(page_id, page)?;
247 return Ok(());
248 }
249
250 let (new_leaf, separator_key) = self.split_leaf(&mut page, key, value)?;
252 let new_leaf_id = new_leaf.page_id();
253 page.update_checksum();
254 let page_id = page.page_id();
255 self.pager.write_page(page_id, page.clone())?;
256 self.cache_rightmost_leaf(&new_leaf)?;
258
259 self.insert_into_parent(path, page_id, &separator_key, new_leaf_id)?;
261
262 Ok(())
263 }
264
265 pub fn upsert(&self, key: &[u8], value: &[u8]) -> BTreeResult<()> {
277 let _writer = self.writer.lock();
278 let root_id = self.root_page_id();
279 if root_id == 0 {
280 return self.insert(key, value);
281 }
282
283 let new_stored = value_layout::encode(self.pager.as_ref(), value)?;
288
289 let (leaf_id, _path) = self.find_leaf(root_id, key)?;
290 let mut page = self.pager.read_page(leaf_id)?;
291
292 if let SearchResult::Found(pos) = search_leaf(&page, key)? {
293 let (_, old_value) = read_leaf_cell(&page, pos)?;
294 if new_stored.len() <= old_value.len() {
295 self.free_chain_if_pointer(&old_value)?;
300 overwrite_leaf_value(&mut page, pos, &new_stored)?;
301 let page_id = page.page_id();
302 self.pager.write_page(page_id, page)?;
303 return Ok(());
304 }
305 let _ = self.delete(key);
308 return self.insert_encoded(key, &new_stored);
309 }
310
311 self.insert_encoded(key, &new_stored)
314 }
315
316 pub fn try_upsert_inline_without_splits(
321 &self,
322 items: &[(Vec<u8>, Vec<u8>)],
323 ) -> BTreeResult<bool> {
324 let _writer = self.writer.lock();
325 let root_id = self.root_page_id();
326 if root_id == 0 && !items.is_empty() {
327 return Ok(false);
328 }
329 let mut pages = std::collections::BTreeMap::<u32, (Page, Vec<(Vec<u8>, Vec<u8>)>)>::new();
330 for (key, value) in items {
331 if key.len() > MAX_KEY_SIZE {
332 return Err(BTreeError::KeyTooLarge(key.len()));
333 }
334 let Some(new_value) = value_layout::encode_inline_without_pages(value)? else {
335 return Ok(false);
336 };
337 let (leaf_id, _) = self.find_leaf(root_id, key)?;
338 if let std::collections::btree_map::Entry::Vacant(entry) = pages.entry(leaf_id) {
339 let page = self.pager.read_page(leaf_id)?;
340 let entries = read_leaf_entries(&page)?;
341 entry.insert((page, entries));
342 }
343 let (_, entries) = pages.get_mut(&leaf_id).expect("leaf was loaded");
344 match entries.binary_search_by(|(existing, _)| existing.as_slice().cmp(key)) {
345 Ok(position) => {
346 if value_layout::pointer_head(&entries[position].1).is_some() {
347 return Ok(false);
348 }
349 entries[position].1 = new_value;
350 }
351 Err(position) => entries.insert(position, (key.clone(), new_value)),
352 }
353 }
354 for (page, entries) in pages.values_mut() {
357 if leaf_entries_size(entries) > PAGE_SIZE - LEAF_SLOT_ARRAY_OFFSET {
358 return Ok(false);
359 }
360 write_leaf_entries(page, entries)?;
361 }
362 for (page_id, (page, _)) in pages {
363 self.pager.write_page(page_id, page)?;
364 }
365 Ok(true)
366 }
367
368 pub fn upsert_batch_sorted(&self, items: &[(Vec<u8>, Vec<u8>)]) -> BTreeResult<()> {
379 let _writer = self.writer.lock();
380 if items.is_empty() {
381 return Ok(());
382 }
383 let root_id = self.root_page_id();
384 if root_id == 0 {
385 for (k, v) in items {
386 self.upsert(k, v)?;
387 }
388 return Ok(());
389 }
390
391 let mut current_leaf: Option<(u32, Page, bool)> = None; for (key, value) in items {
393 let (leaf_id, _path) = self.find_leaf(root_id, key)?;
396 if current_leaf.as_ref().map(|(id, _, _)| *id) != Some(leaf_id) {
397 if let Some((id, page, dirty)) = current_leaf.take() {
398 if dirty {
399 self.pager.write_page(id, page)?;
400 }
401 }
402 let page = self.pager.read_page(leaf_id)?;
403 current_leaf = Some((leaf_id, page, false));
404 }
405
406 let applied_in_place = {
412 let (_, page, dirty) = current_leaf.as_mut().expect("set above");
413 if let SearchResult::Found(pos) = search_leaf(page, key)? {
414 let (_, old_value) = read_leaf_cell(page, pos)?;
415 match value_layout::encode(self.pager.as_ref(), value) {
416 Ok(new_stored) if new_stored.len() <= old_value.len() => {
417 self.free_chain_if_pointer(&old_value)?;
421 overwrite_leaf_value(page, pos, &new_stored)?;
422 *dirty = true;
423 true
424 }
425 _ => false,
426 }
427 } else {
428 false
429 }
430 };
431 if !applied_in_place {
432 if let Some((id, page, dirty)) = current_leaf.take() {
433 if dirty {
434 self.pager.write_page(id, page)?;
435 }
436 }
437 self.upsert(key, value)?;
438 }
439 }
440
441 if let Some((id, page, dirty)) = current_leaf.take() {
442 if dirty {
443 self.pager.write_page(id, page)?;
444 }
445 }
446 Ok(())
447 }
448
449 pub fn bulk_insert_sorted(&self, items: &[(Vec<u8>, Vec<u8>)]) -> BTreeResult<()> {
458 let _writer = self.writer.lock();
459 if items.is_empty() {
460 return Ok(());
461 }
462 for (key, _) in items {
468 if key.len() > MAX_KEY_SIZE {
469 return Err(BTreeError::KeyTooLarge(key.len()));
470 }
471 }
472
473 let mut encoded: Vec<(&[u8], Vec<u8>)> = Vec::with_capacity(items.len());
480 let mut first_oversize: Option<usize> = None;
481 for (key, value) in items {
482 match value_layout::encode(self.pager.as_ref(), value) {
483 Ok(stored) => encoded.push((key.as_slice(), stored)),
484 Err(value_layout::ValueLayoutError::ValueTooLarge(n)) => {
485 if first_oversize.is_none() {
486 first_oversize = Some(n);
487 }
488 }
489 Err(other) => return Err(other.into()),
490 }
491 }
492
493 let items: Vec<(&[u8], &[u8])> = encoded.iter().map(|(k, v)| (*k, v.as_slice())).collect();
494
495 let mut last_filled_leaf: Option<u32> = None;
503 let mut i = 0;
504 while i < items.len() {
505 let root_id = self.root_page_id();
506 if root_id == 0 {
507 let (k, v) = items[i];
511 self.insert_encoded(k, v)?;
512 i += 1;
513 last_filled_leaf = None;
514 continue;
515 }
516
517 let hint_leaf_id: Option<u32> = if let Some(prev) = last_filled_leaf {
523 let prev_page = self.pager.read_page(prev)?;
524 let sibling = leaf_right_sibling(&prev_page);
525 if sibling != 0 {
526 let sib_page = self.pager.read_page(sibling)?;
527 match leaf_high_key(&sib_page)? {
528 None => Some(sibling), Some(hk)
530 if leaf_right_sibling(&sib_page) == 0 && items[i].0 > hk.as_slice() =>
531 {
532 Some(sibling)
533 }
534 Some(_) => None,
535 }
536 } else {
537 None
538 }
539 } else {
540 None
541 };
542
543 let (leaf_id, _path) = match hint_leaf_id {
544 Some(id) => (id, Vec::new()),
545 None => self.find_leaf(root_id, items[i].0)?,
546 };
547 let mut page = self.pager.read_page(leaf_id)?;
548
549 let mut last_key_in_leaf: Option<Vec<u8>> = {
554 let cell_count = page.cell_count() as usize;
555 if cell_count == 0 {
556 None
557 } else {
558 Some(read_leaf_cell(&page, cell_count - 1)?.0)
559 }
560 };
561
562 let mut inserted = 0usize;
568 while i + inserted < items.len() {
569 let (key, value) = items[i + inserted];
570
571 if let Some(lk) = &last_key_in_leaf {
574 match key.cmp(lk.as_slice()) {
575 Ordering::Greater => {}
576 Ordering::Equal => return Err(BTreeError::DuplicateKey),
577 Ordering::Less => break,
578 }
579 }
580
581 let cell_size = 4 + key.len() + value.len();
586 let slot_end_after =
587 LEAF_SLOT_ARRAY_OFFSET + (page.cell_count() as usize + inserted + 1) * 2;
588 let cells_start = leaf_cells_start(&page);
589 if slot_end_after + cell_size > cells_start {
590 break;
591 }
592
593 let cell_offset = cells_start - cell_size;
598 {
599 let data = page.as_bytes_mut();
600 data[cell_offset..cell_offset + 2]
601 .copy_from_slice(&(key.len() as u16).to_le_bytes());
602 data[cell_offset + 2..cell_offset + 4]
603 .copy_from_slice(&(value.len() as u16).to_le_bytes());
604 data[cell_offset + 4..cell_offset + 4 + key.len()].copy_from_slice(key);
605 data[cell_offset + 4 + key.len()..cell_offset + cell_size]
606 .copy_from_slice(value);
607 }
608 page.set_free_end(cell_offset as u16);
609
610 let new_slot_index = page.cell_count() as usize + inserted;
611 let slot_pos = leaf_slot_offset_for(new_slot_index);
612 {
613 let data = page.as_bytes_mut();
614 data[slot_pos..slot_pos + 2]
615 .copy_from_slice(&(cell_offset as u16).to_le_bytes());
616 }
617
618 last_key_in_leaf = Some(key.to_vec());
619 inserted += 1;
620 }
621
622 if inserted > 0 {
623 let new_count = page.cell_count() as usize + inserted;
624 page.set_cell_count(new_count as u16);
625 page.set_free_start((LEAF_SLOT_ARRAY_OFFSET + new_count * 2) as u16);
626 page.update_checksum();
627 self.cache_rightmost_leaf(&page)?;
628 self.pager.write_page(leaf_id, page)?;
629 i += inserted;
630 last_filled_leaf = Some(leaf_id);
633 } else {
634 let (k, v) = items[i];
641 self.insert_encoded(k, v)?;
642 i += 1;
643 last_filled_leaf = None;
644 }
645 }
646
647 if let Some(n) = first_oversize {
652 return Err(BTreeError::ValueTooLarge(n));
653 }
654
655 Ok(())
656 }
657
658 pub fn delete(&self, key: &[u8]) -> BTreeResult<bool> {
660 let _writer = self.writer.lock();
661 let root_id = self.root_page_id();
662 if root_id == 0 {
663 return Ok(false);
664 }
665
666 let (leaf_id, path) = self.find_leaf(root_id, key)?;
670 let mut current_id = leaf_id;
671 const MAX_RIGHTLINK_HOPS: usize = 32;
672 for _ in 0..MAX_RIGHTLINK_HOPS {
673 let mut page = self.pager.read_page(current_id)?;
674 match search_leaf(&page, key)? {
675 SearchResult::Found(pos) => {
676 let (_, stored) = read_leaf_cell(&page, pos)?;
681 self.free_chain_if_pointer(&stored)?;
682 delete_from_leaf(&mut page, pos)?;
683 page.update_checksum();
684 let page_id = page.page_id();
685 self.pager.write_page(page_id, page.clone())?;
686
687 *self
688 .rightmost_leaf
689 .write()
690 .unwrap_or_else(|e| e.into_inner()) = None;
691
692 if page.cell_count() == 0 && page_id == root_id {
693 self.pager.free_page(root_id)?;
694 *self.root_page_id.write().map_err(|e| {
695 BTreeError::LockPoisoned(format!(
696 "delete: root_page_id write lock: {e}"
697 ))
698 })? = 0;
699 } else {
700 self.rebalance_leaf(current_id, path)?;
701 }
702 return Ok(true);
703 }
704 SearchResult::NotFound(_) => {
705 let right = leaf_right_sibling(&page);
706 if right != 0 {
707 if let Some(high_key) = leaf_high_key(&page)? {
708 if key > high_key.as_slice() {
709 current_id = right;
710 continue;
711 }
712 }
713 }
714 return Ok(false);
715 }
716 }
717 }
718 Ok(false)
719 }
720
721 pub fn cursor_first(&self) -> BTreeResult<BTreeCursor<P>> {
723 let root_id = self.root_page_id();
724 if root_id == 0 {
725 return Ok(BTreeCursor {
726 _retirement_pin: Arc::clone(&self.retirement_pin),
727 leaf_page_id: 0,
728 position: 0,
729 pager: self.pager.clone(),
730 prefetched_next: false,
731 immutable: false,
732 immutable_leaf: None,
733 });
734 }
735
736 let first_leaf = self.find_first_leaf(root_id)?;
738
739 Ok(BTreeCursor {
740 _retirement_pin: Arc::clone(&self.retirement_pin),
741 leaf_page_id: first_leaf,
742 position: 0,
743 pager: self.pager.clone(),
744 prefetched_next: false,
745 immutable: false,
746 immutable_leaf: None,
747 })
748 }
749
750 pub fn cursor_seek(&self, key: &[u8]) -> BTreeResult<BTreeCursor<P>> {
752 let root_id = self.root_page_id();
753 if root_id == 0 {
754 return Ok(BTreeCursor {
755 _retirement_pin: Arc::clone(&self.retirement_pin),
756 leaf_page_id: 0,
757 position: 0,
758 pager: self.pager.clone(),
759 prefetched_next: false,
760 immutable: false,
761 immutable_leaf: None,
762 });
763 }
764
765 let (leaf_id, _) = self.find_leaf(root_id, key)?;
766 let page = self.pager.read_page(leaf_id)?;
767
768 let position = match search_leaf(&page, key)? {
769 SearchResult::Found(pos) => pos,
770 SearchResult::NotFound(pos) => pos,
771 };
772
773 Ok(BTreeCursor {
774 _retirement_pin: Arc::clone(&self.retirement_pin),
775 leaf_page_id: leaf_id,
776 position,
777 pager: self.pager.clone(),
778 prefetched_next: false,
779 immutable: false,
780 immutable_leaf: None,
781 })
782 }
783
784 pub fn range(&self, start_key: &[u8], end_key: &[u8]) -> BTreeResult<Vec<(Vec<u8>, Vec<u8>)>> {
786 let mut results = Vec::new();
787 let mut cursor = self.cursor_seek(start_key)?;
788
789 while let Some((key, value)) = cursor.next()? {
790 if key.as_slice() > end_key {
791 break;
792 }
793 results.push((key, value));
794 }
795
796 Ok(results)
797 }
798
799 pub fn count(&self) -> BTreeResult<usize> {
801 let root_id = self.root_page_id();
802 if root_id == 0 {
803 return Ok(0);
804 }
805
806 let mut count = 0;
807 let mut cursor = self.cursor_first()?;
808 while cursor.next()?.is_some() {
809 count += 1;
810 }
811
812 Ok(count)
813 }
814
815 fn cache_rightmost_leaf(&self, page: &Page) -> BTreeResult<()> {
818 let cached = if leaf_right_sibling(page) == 0 {
822 leaf_high_key(page)?.map(|high| (page.page_id(), high))
823 } else {
824 None
825 };
826 *self.rightmost_leaf.write().map_err(|error| {
827 BTreeError::LockPoisoned(format!("rightmost_leaf write: {error}"))
828 })? = cached;
829 Ok(())
830 }
831
832 fn find_leaf(&self, page_id: u32, key: &[u8]) -> BTreeResult<(u32, Vec<u32>)> {
834 let mut current_id = page_id;
835 let mut path = Vec::new();
836
837 for _ in 0..64 {
840 let page = self.pager.read_page(current_id)?;
841
842 match page.page_type()? {
843 PageType::BTreeLeaf => {
844 return Ok((current_id, path));
845 }
846 PageType::BTreeInterior => {
847 path.push(current_id);
848 current_id = find_child(&page, key)?;
849 }
850 other => {
851 return Err(BTreeError::Corrupted(format!(
852 "Unexpected page type in B-tree: {:?}",
853 other
854 )));
855 }
856 }
857 }
858 Err(BTreeError::Corrupted(
859 "B-tree descent exceeds maximum depth".into(),
860 ))
861 }
862
863 pub(crate) fn find_first_leaf(&self, page_id: u32) -> BTreeResult<u32> {
865 let mut current_id = page_id;
866
867 for _ in 0..64 {
868 let page = self.pager.read_page(current_id)?;
869
870 match page.page_type()? {
871 PageType::BTreeLeaf => return Ok(current_id),
872 PageType::BTreeInterior => {
873 current_id = find_first_child(&page)?;
875 }
876 _ => {
877 return Err(BTreeError::Corrupted("Invalid page type".into()));
878 }
879 }
880 }
881 Err(BTreeError::Corrupted(
882 "leftmost-leaf descent exceeds maximum depth".into(),
883 ))
884 }
885
886 fn split_leaf(
888 &self,
889 page: &mut Page,
890 new_key: &[u8],
891 new_value: &[u8],
892 ) -> BTreeResult<(Page, Vec<u8>)> {
893 let mut entries: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
895 let cell_count = page.cell_count() as usize;
896
897 for i in 0..cell_count {
898 entries.push(read_leaf_cell(page, i)?);
899 }
900
901 let insert_pos = entries.partition_point(|(k, _)| k.as_slice() < new_key);
903 entries.insert(insert_pos, (new_key.to_vec(), new_value.to_vec()));
904
905 let capacity = PAGE_SIZE - LEAF_SLOT_ARRAY_OFFSET;
908 let total = leaf_entries_size(&entries);
909 let mut left_bytes = 0;
910 let mut split = None;
911 let mut best_balance = usize::MAX;
912 for index in 1..entries.len() {
913 left_bytes += leaf_entry_size(&entries[index - 1]);
914 let right_bytes = total - left_bytes;
915 if left_bytes <= capacity && right_bytes <= capacity {
916 let balance = left_bytes.abs_diff(right_bytes);
917 if balance < best_balance {
918 split = Some(index);
919 best_balance = balance;
920 }
921 }
922 }
923 let mid = split
924 .ok_or_else(|| BTreeError::Corrupted("leaf entries cannot fit in two pages".into()))?;
925
926 let mut new_page = self.pager.allocate_page(PageType::BTreeLeaf)?;
928
929 let old_next = read_next_leaf(page);
931 init_leaf_links(&mut new_page, page.page_id(), old_next);
932 set_next_leaf(page, new_page.page_id());
933
934 write_leaf_entries(page, &entries[..mid])?;
937 write_leaf_entries(&mut new_page, &entries[mid..])?;
938
939 let separator = entries[mid].0.clone();
941
942 new_page.update_checksum();
943 let new_page_id = new_page.page_id();
944 self.pager.write_page(new_page_id, new_page.clone())?;
945 if old_next != 0 {
946 let mut successor = self.pager.read_page(old_next)?;
949 if successor.page_type()? != PageType::BTreeLeaf {
950 return Err(BTreeError::Corrupted("invalid split successor type".into()));
951 }
952 set_prev_leaf(&mut successor, new_page_id);
953 successor.update_checksum();
954 self.pager.write_page(old_next, successor)?;
955 }
956
957 Ok((new_page, separator))
958 }
959
960 fn insert_into_parent(
962 &self,
963 mut path: Vec<u32>,
964 left_child: u32,
965 key: &[u8],
966 right_child: u32,
967 ) -> BTreeResult<()> {
968 if path.is_empty() {
970 let mut new_root = self.pager.allocate_page(PageType::BTreeInterior)?;
971 write_interior_entries(&mut new_root, &[key.to_vec()], &[left_child, right_child])?;
974
975 new_root.update_checksum();
976 let new_root_id = new_root.page_id();
977 self.pager.write_page(new_root_id, new_root)?;
978 *self.root_page_id.write().map_err(|e| {
979 BTreeError::LockPoisoned(format!(
980 "insert_into_parent: root_page_id write lock: {e}"
981 ))
982 })? = new_root_id;
983 return Ok(());
984 }
985
986 let parent_id = path.pop().ok_or_else(|| {
988 BTreeError::Corrupted("insert_into_parent: path unexpectedly empty".into())
989 })?;
990 let mut parent = self.pager.read_page(parent_id)?;
991
992 if can_insert_interior(&parent, key) {
994 insert_into_interior(&mut parent, key, left_child, right_child)?;
995 parent.update_checksum();
996 self.pager.write_page(parent_id, parent)?;
997 return Ok(());
998 }
999
1000 let (new_interior, separator) =
1002 self.split_interior(&mut parent, key, left_child, right_child)?;
1003 parent.update_checksum();
1004 self.pager.write_page(parent_id, parent.clone())?;
1005
1006 self.insert_into_parent(path, parent.page_id(), &separator, new_interior.page_id())
1008 }
1009
1010 fn split_interior(
1012 &self,
1013 page: &mut Page,
1014 new_key: &[u8],
1015 left_child: u32,
1016 right_child: u32,
1017 ) -> BTreeResult<(Page, Vec<u8>)> {
1018 let (mut keys, mut children) = read_interior_keys_children(page)?;
1019 let key_insert_pos = keys.partition_point(|key| key.as_slice() < new_key);
1020 let child_insert_pos = children
1021 .iter()
1022 .position(|child| *child == left_child)
1023 .unwrap_or(key_insert_pos);
1024
1025 keys.insert(key_insert_pos, new_key.to_vec());
1026 children.insert(child_insert_pos + 1, right_child);
1027
1028 let mid = keys.len() / 2;
1031 let separator = keys[mid].clone();
1032
1033 let mut new_page = self.pager.allocate_page(PageType::BTreeInterior)?;
1035
1036 write_interior_entries(page, &keys[..mid], &children[..mid + 1])?;
1037 write_interior_entries(&mut new_page, &keys[mid + 1..], &children[mid + 1..])?;
1038
1039 new_page.update_checksum();
1040 let new_page_id = new_page.page_id();
1041 self.pager.write_page(new_page_id, new_page.clone())?;
1042
1043 Ok((new_page, separator))
1044 }
1045
1046 fn rebalance_leaf(&self, leaf_id: u32, mut path: Vec<u32>) -> BTreeResult<()> {
1047 if let Some(parent_id) = self.rebalance_leaf_step(leaf_id, &path)? {
1050 path.pop();
1051 self.rebalance_interior(parent_id, path)?;
1052 }
1053 Ok(())
1054 }
1055
1056 fn rebalance_leaf_step(&self, leaf_id: u32, path: &[u32]) -> BTreeResult<Option<u32>> {
1057 if path.is_empty() {
1058 return Ok(None);
1059 }
1060
1061 let root_id = self.root_page_id();
1062 if leaf_id == root_id {
1063 return Ok(None);
1064 }
1065
1066 let mut leaf = self.pager.read_page(leaf_id)?;
1067 let mut leaf_entries = read_leaf_entries(&leaf)?;
1068 let min_bytes = leaf_min_bytes();
1069
1070 let parent_id = *path.last().ok_or_else(|| {
1071 BTreeError::Corrupted("rebalance_leaf: path unexpectedly empty".into())
1072 })?;
1073 let mut parent = self.pager.read_page(parent_id)?;
1074 let (mut parent_keys, mut parent_children) = read_interior_keys_children(&parent)?;
1075
1076 let child_index = parent_children
1077 .iter()
1078 .position(|&id| id == leaf_id)
1079 .ok_or_else(|| BTreeError::Corrupted("Leaf missing from parent".into()))?;
1080
1081 if child_index > 0 {
1082 if let Some((first_key, _)) = leaf_entries.first() {
1083 if parent_keys.get(child_index - 1).map(|k| k.as_slice())
1084 != Some(first_key.as_slice())
1085 {
1086 parent_keys[child_index - 1] = first_key.clone();
1087 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1088 parent.update_checksum();
1089 self.pager.write_page(parent_id, parent.clone())?;
1090 }
1091 }
1092 }
1093
1094 if leaf_entries_size(&leaf_entries) >= min_bytes {
1095 return Ok(None);
1096 }
1097
1098 if child_index > 0 {
1099 let left_id = parent_children[child_index - 1];
1100 let mut left = self.pager.read_page(left_id)?;
1101 let mut left_entries = read_leaf_entries(&left)?;
1102 let mut borrowed = false;
1103
1104 while leaf_entries_size(&leaf_entries) < min_bytes {
1105 let Some(entry) = left_entries.pop() else {
1106 break;
1107 };
1108 if leaf_entries_size(&left_entries) < min_bytes {
1109 left_entries.push(entry);
1110 break;
1111 }
1112 leaf_entries.insert(0, entry);
1113 borrowed = true;
1114 }
1115
1116 if borrowed {
1117 write_leaf_entries(&mut left, &left_entries)?;
1118 left.update_checksum();
1119 self.pager.write_page(left_id, left)?;
1120
1121 write_leaf_entries(&mut leaf, &leaf_entries)?;
1122 leaf.update_checksum();
1123 self.pager.write_page(leaf_id, leaf)?;
1124
1125 if let Some((first_key, _)) = leaf_entries.first() {
1126 parent_keys[child_index - 1] = first_key.clone();
1127 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1128 parent.update_checksum();
1129 self.pager.write_page(parent_id, parent)?;
1130 }
1131
1132 return Ok(None);
1133 }
1134 }
1135
1136 if child_index + 1 < parent_children.len() {
1137 let right_id = parent_children[child_index + 1];
1138 let mut right = self.pager.read_page(right_id)?;
1139 let mut right_entries = read_leaf_entries(&right)?;
1140 let mut borrowed = false;
1141
1142 while leaf_entries_size(&leaf_entries) < min_bytes {
1143 if right_entries.is_empty() {
1144 break;
1145 }
1146 let entry = right_entries.remove(0);
1147 if leaf_entries_size(&right_entries) < min_bytes {
1148 right_entries.insert(0, entry);
1149 break;
1150 }
1151 leaf_entries.push(entry);
1152 borrowed = true;
1153 }
1154
1155 if borrowed {
1156 write_leaf_entries(&mut right, &right_entries)?;
1157 right.update_checksum();
1158 self.pager.write_page(right_id, right)?;
1159
1160 write_leaf_entries(&mut leaf, &leaf_entries)?;
1161 leaf.update_checksum();
1162 self.pager.write_page(leaf_id, leaf)?;
1163
1164 if let Some((first_key, _)) = right_entries.first() {
1165 parent_keys[child_index] = first_key.clone();
1166 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1167 parent.update_checksum();
1168 self.pager.write_page(parent_id, parent)?;
1169 }
1170
1171 return Ok(None);
1172 }
1173 }
1174
1175 if child_index > 0 {
1176 let left_id = parent_children[child_index - 1];
1177 let mut left = self.pager.read_page(left_id)?;
1178 let mut left_entries = read_leaf_entries(&left)?;
1179
1180 left_entries.extend(leaf_entries);
1181 write_leaf_entries(&mut left, &left_entries)?;
1182
1183 let next_leaf = read_next_leaf(&leaf);
1184 set_next_leaf(&mut left, next_leaf);
1185 if next_leaf != 0 {
1186 let mut next = self.pager.read_page(next_leaf)?;
1187 set_prev_leaf(&mut next, left_id);
1188 next.update_checksum();
1189 self.pager.write_page(next_leaf, next)?;
1190 }
1191
1192 left.update_checksum();
1193 self.pager.write_page(left_id, left)?;
1194 self.pager.free_page(leaf_id)?;
1195
1196 parent_keys.remove(child_index - 1);
1197 parent_children.remove(child_index);
1198 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1199 parent.update_checksum();
1200 self.pager.write_page(parent_id, parent)?;
1201
1202 return Ok(Some(parent_id));
1203 }
1204
1205 if child_index + 1 < parent_children.len() {
1206 let right_id = parent_children[child_index + 1];
1207 let right = self.pager.read_page(right_id)?;
1208 let right_entries = read_leaf_entries(&right)?;
1209
1210 leaf_entries.extend(right_entries);
1211 write_leaf_entries(&mut leaf, &leaf_entries)?;
1212
1213 let next_leaf = read_next_leaf(&right);
1214 set_next_leaf(&mut leaf, next_leaf);
1215 if next_leaf != 0 {
1216 let mut next = self.pager.read_page(next_leaf)?;
1217 set_prev_leaf(&mut next, leaf_id);
1218 next.update_checksum();
1219 self.pager.write_page(next_leaf, next)?;
1220 }
1221
1222 leaf.update_checksum();
1223 self.pager.write_page(leaf_id, leaf)?;
1224 self.pager.free_page(right_id)?;
1225
1226 parent_keys.remove(child_index);
1227 parent_children.remove(child_index + 1);
1228 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1229 parent.update_checksum();
1230 self.pager.write_page(parent_id, parent)?;
1231
1232 return Ok(Some(parent_id));
1233 }
1234
1235 Ok(None)
1236 }
1237
1238 fn rebalance_interior(&self, mut node_id: u32, mut path: Vec<u32>) -> BTreeResult<()> {
1239 while let Some(parent_id) = self.rebalance_interior_step(node_id, &mut path)? {
1241 node_id = parent_id;
1242 }
1243 Ok(())
1244 }
1245
1246 fn rebalance_interior_step(
1247 &self,
1248 node_id: u32,
1249 path: &mut Vec<u32>,
1250 ) -> BTreeResult<Option<u32>> {
1251 let root_id = self.root_page_id();
1252 let mut node = self.pager.read_page(node_id)?;
1253 let (mut node_keys, mut node_children) = read_interior_keys_children(&node)?;
1254 let min_bytes = interior_min_bytes();
1255
1256 if node_id == root_id {
1257 if node_keys.is_empty() {
1258 let next_root = node_children.first().copied().unwrap_or(0);
1259 self.pager.free_page(node_id)?;
1260 *self.root_page_id.write().map_err(|e| {
1261 BTreeError::LockPoisoned(format!(
1262 "rebalance_interior: root_page_id write lock: {e}"
1263 ))
1264 })? = next_root;
1265 }
1266 return Ok(None);
1267 }
1268
1269 if interior_entries_size(&node_keys) >= min_bytes {
1270 return Ok(None);
1271 }
1272
1273 let parent_id = match path.pop() {
1274 Some(id) => id,
1275 None => return Ok(None),
1276 };
1277 let mut parent = self.pager.read_page(parent_id)?;
1278 let (mut parent_keys, mut parent_children) = read_interior_keys_children(&parent)?;
1279
1280 let child_index = parent_children
1281 .iter()
1282 .position(|&id| id == node_id)
1283 .ok_or_else(|| BTreeError::Corrupted("Interior missing from parent".into()))?;
1284
1285 if child_index > 0 {
1286 let left_id = parent_children[child_index - 1];
1287 let mut left = self.pager.read_page(left_id)?;
1288 let (mut left_keys, mut left_children) = read_interior_keys_children(&left)?;
1289
1290 if let Some(borrow_key) = left_keys.last().cloned() {
1291 let borrow_size = interior_key_size(&borrow_key);
1292 if interior_entries_size(&left_keys).saturating_sub(borrow_size) >= min_bytes {
1293 let parent_key = parent_keys[child_index - 1].clone();
1294 let borrowed_key = left_keys.pop().ok_or_else(|| {
1295 BTreeError::Corrupted(
1296 "rebalance_interior: left_keys empty after check".into(),
1297 )
1298 })?;
1299 let borrowed_child = left_children.pop().ok_or_else(|| {
1300 BTreeError::Corrupted("rebalance_interior: left_children empty".into())
1301 })?;
1302
1303 node_keys.insert(0, parent_key);
1304 node_children.insert(0, borrowed_child);
1305 parent_keys[child_index - 1] = borrowed_key;
1306
1307 write_interior_entries(&mut left, &left_keys, &left_children)?;
1308 left.update_checksum();
1309 self.pager.write_page(left_id, left)?;
1310
1311 write_interior_entries(&mut node, &node_keys, &node_children)?;
1312 node.update_checksum();
1313 self.pager.write_page(node_id, node)?;
1314
1315 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1316 parent.update_checksum();
1317 self.pager.write_page(parent_id, parent)?;
1318
1319 return Ok(None);
1320 }
1321 }
1322 }
1323
1324 if child_index + 1 < parent_children.len() {
1325 let right_id = parent_children[child_index + 1];
1326 let mut right = self.pager.read_page(right_id)?;
1327 let (mut right_keys, mut right_children) = read_interior_keys_children(&right)?;
1328
1329 if let Some(borrow_key) = right_keys.first().cloned() {
1330 let borrow_size = interior_key_size(&borrow_key);
1331 if interior_entries_size(&right_keys).saturating_sub(borrow_size) >= min_bytes {
1332 let parent_key = parent_keys[child_index].clone();
1333 let new_parent_key = right_keys.remove(0);
1334 let borrowed_child = right_children.remove(0);
1335
1336 node_keys.push(parent_key);
1337 node_children.push(borrowed_child);
1338 parent_keys[child_index] = new_parent_key;
1339
1340 write_interior_entries(&mut right, &right_keys, &right_children)?;
1341 right.update_checksum();
1342 self.pager.write_page(right_id, right)?;
1343
1344 write_interior_entries(&mut node, &node_keys, &node_children)?;
1345 node.update_checksum();
1346 self.pager.write_page(node_id, node)?;
1347
1348 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1349 parent.update_checksum();
1350 self.pager.write_page(parent_id, parent)?;
1351
1352 return Ok(None);
1353 }
1354 }
1355 }
1356
1357 if child_index > 0 {
1358 let left_id = parent_children[child_index - 1];
1359 let mut left = self.pager.read_page(left_id)?;
1360 let (mut left_keys, mut left_children) = read_interior_keys_children(&left)?;
1361 let parent_key = parent_keys.remove(child_index - 1);
1362 parent_children.remove(child_index);
1363
1364 left_keys.push(parent_key);
1365 left_keys.extend(node_keys);
1366 left_children.extend(node_children);
1367
1368 write_interior_entries(&mut left, &left_keys, &left_children)?;
1369 left.update_checksum();
1370 self.pager.write_page(left_id, left)?;
1371 self.pager.free_page(node_id)?;
1372
1373 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1374 parent.update_checksum();
1375 self.pager.write_page(parent_id, parent)?;
1376
1377 return Ok(Some(parent_id));
1378 }
1379
1380 if child_index + 1 < parent_children.len() {
1381 let right_id = parent_children[child_index + 1];
1382 let right = self.pager.read_page(right_id)?;
1383 let (right_keys, right_children) = read_interior_keys_children(&right)?;
1384 let parent_key = parent_keys.remove(child_index);
1385 parent_children.remove(child_index + 1);
1386
1387 node_keys.push(parent_key);
1388 node_keys.extend(right_keys);
1389 node_children.extend(right_children);
1390
1391 write_interior_entries(&mut node, &node_keys, &node_children)?;
1392 node.update_checksum();
1393 self.pager.write_page(node_id, node)?;
1394 self.pager.free_page(right_id)?;
1395
1396 write_interior_entries(&mut parent, &parent_keys, &parent_children)?;
1397 parent.update_checksum();
1398 self.pager.write_page(parent_id, parent)?;
1399
1400 return Ok(Some(parent_id));
1401 }
1402
1403 Ok(None)
1404 }
1405}