1use crate::db::TransactionGuard;
2use crate::sealed::Sealed;
3use crate::tree_store::{
4 AccessGuardMut, Btree, BtreeExtractIf, BtreeHeader, BtreeMut, BtreeRangeIter, MAX_PAIR_LENGTH,
5 MAX_VALUE_LENGTH, PageHint, PageNumber, PageTrackerPolicy, RawBtree, TransactionalMemory,
6};
7use crate::types::{Key, MutInPlaceValue, Value};
8use crate::{AccessGuard, StorageError, WriteTransaction};
9use crate::{Result, TableHandle};
10use std::borrow::Borrow;
11use std::fmt::{Debug, Formatter};
12use std::marker::PhantomData;
13use std::ops::RangeBounds;
14use std::sync::{Arc, Mutex};
15
16#[derive(Debug)]
18pub struct TableStats {
19 pub(crate) tree_height: u32,
20 pub(crate) leaf_pages: u64,
21 pub(crate) branch_pages: u64,
22 pub(crate) stored_leaf_bytes: u64,
23 pub(crate) metadata_bytes: u64,
24 pub(crate) fragmented_bytes: u64,
25}
26
27impl TableStats {
28 pub fn tree_height(&self) -> u32 {
30 self.tree_height
31 }
32
33 pub fn leaf_pages(&self) -> u64 {
35 self.leaf_pages
36 }
37
38 pub fn branch_pages(&self) -> u64 {
40 self.branch_pages
41 }
42
43 pub fn stored_bytes(&self) -> u64 {
46 self.stored_leaf_bytes
47 }
48
49 pub fn metadata_bytes(&self) -> u64 {
51 self.metadata_bytes
52 }
53
54 pub fn fragmented_bytes(&self) -> u64 {
56 self.fragmented_bytes
57 }
58}
59
60pub struct Table<'txn, K: Key + 'static, V: Value + 'static> {
62 name: String,
63 transaction: &'txn WriteTransaction,
64 tree: BtreeMut<'txn, K, V>,
65}
66
67impl<K: Key + 'static, V: Value + 'static> TableHandle for Table<'_, K, V> {
68 fn name(&self) -> &str {
69 &self.name
70 }
71}
72
73impl<'txn, K: Key + 'static, V: Value + 'static> Table<'txn, K, V> {
74 pub(crate) fn new(
75 name: &str,
76 table_root: Option<BtreeHeader>,
77 freed_pages: Arc<Mutex<Vec<PageNumber>>>,
78 allocated_pages: Arc<Mutex<PageTrackerPolicy>>,
79 mem: Arc<TransactionalMemory>,
80 transaction: &'txn WriteTransaction,
81 ) -> Table<'txn, K, V> {
82 Table {
83 name: name.to_string(),
84 transaction,
85 tree: BtreeMut::new(
86 table_root,
87 transaction.transaction_guard(),
88 mem,
89 freed_pages,
90 allocated_pages,
91 ),
92 }
93 }
94
95 #[allow(dead_code)]
96 pub(crate) fn print_debug(&self, include_values: bool) -> Result {
97 self.tree.print_debug(include_values)
98 }
99
100 pub fn pop_first(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
102 let first = self
104 .iter()?
105 .next()
106 .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
107 if let Some(owned_key) = first {
108 let owned_key = owned_key?;
109 let key = K::from_bytes(&owned_key);
110 let value = self.remove(&key)?.unwrap();
111 drop(key);
112 Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
113 } else {
114 Ok(None)
115 }
116 }
117
118 pub fn pop_last(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
120 let last = self
122 .iter()?
123 .next_back()
124 .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
125 if let Some(owned_key) = last {
126 let owned_key = owned_key?;
127 let key = K::from_bytes(&owned_key);
128 let value = self.remove(&key)?.unwrap();
129 drop(key);
130 Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
131 } else {
132 Ok(None)
133 }
134 }
135
136 pub fn extract_if<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
141 &mut self,
142 predicate: F,
143 ) -> Result<ExtractIf<K, V, F>> {
144 self.extract_from_if::<K::SelfType<'_>, F>(.., predicate)
145 }
146
147 pub fn extract_from_if<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
152 &mut self,
153 range: impl RangeBounds<KR> + 'a,
154 predicate: F,
155 ) -> Result<ExtractIf<K, V, F>>
156 where
157 KR: Borrow<K::SelfType<'a>> + 'a,
158 {
159 self.tree
160 .extract_from_if(&range, predicate)
161 .map(ExtractIf::new)
162 }
163
164 pub fn retain<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
168 &mut self,
169 predicate: F,
170 ) -> Result {
171 self.tree.retain_in::<K::SelfType<'_>, F>(predicate, ..)
172 }
173
174 pub fn retain_in<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
178 &mut self,
179 range: impl RangeBounds<KR> + 'a,
180 predicate: F,
181 ) -> Result
182 where
183 KR: Borrow<K::SelfType<'a>> + 'a,
184 {
185 self.tree.retain_in(predicate, range)
186 }
187
188 pub fn insert<'k, 'v>(
194 &mut self,
195 key: impl Borrow<K::SelfType<'k>>,
196 value: impl Borrow<V::SelfType<'v>>,
197 ) -> Result<Option<AccessGuard<V>>> {
198 let value_len = V::as_bytes(value.borrow()).as_ref().len();
199 if value_len > MAX_VALUE_LENGTH {
200 return Err(StorageError::ValueTooLarge(value_len));
201 }
202 let key_len = K::as_bytes(key.borrow()).as_ref().len();
203 if key_len > MAX_VALUE_LENGTH {
204 return Err(StorageError::ValueTooLarge(key_len));
205 }
206 if value_len + key_len > MAX_PAIR_LENGTH {
207 return Err(StorageError::ValueTooLarge(value_len + key_len));
208 }
209 self.tree.insert(key.borrow(), value.borrow())
210 }
211
212 pub fn remove<'a>(
216 &mut self,
217 key: impl Borrow<K::SelfType<'a>>,
218 ) -> Result<Option<AccessGuard<V>>> {
219 self.tree.remove(key.borrow())
220 }
221}
222
223impl<K: Key + 'static, V: MutInPlaceValue + 'static> Table<'_, K, V> {
224 pub fn insert_reserve<'a>(
230 &mut self,
231 key: impl Borrow<K::SelfType<'a>>,
232 value_length: u32,
233 ) -> Result<AccessGuardMut<V>> {
234 if value_length as usize > MAX_VALUE_LENGTH {
235 return Err(StorageError::ValueTooLarge(value_length as usize));
236 }
237 let key_len = K::as_bytes(key.borrow()).as_ref().len();
238 if key_len > MAX_VALUE_LENGTH {
239 return Err(StorageError::ValueTooLarge(key_len));
240 }
241 if value_length as usize + key_len > MAX_PAIR_LENGTH {
242 return Err(StorageError::ValueTooLarge(value_length as usize + key_len));
243 }
244 self.tree.insert_reserve(key.borrow(), value_length)
245 }
246}
247
248impl<K: Key + 'static, V: Value + 'static> ReadableTableMetadata for Table<'_, K, V> {
249 fn stats(&self) -> Result<TableStats> {
250 let tree_stats = self.tree.stats()?;
251
252 Ok(TableStats {
253 tree_height: tree_stats.tree_height,
254 leaf_pages: tree_stats.leaf_pages,
255 branch_pages: tree_stats.branch_pages,
256 stored_leaf_bytes: tree_stats.stored_leaf_bytes,
257 metadata_bytes: tree_stats.metadata_bytes,
258 fragmented_bytes: tree_stats.fragmented_bytes,
259 })
260 }
261
262 fn len(&self) -> Result<u64> {
263 self.tree.len()
264 }
265}
266
267impl<K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for Table<'_, K, V> {
268 fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
269 self.tree.get(key.borrow())
270 }
271
272 fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
273 where
274 KR: Borrow<K::SelfType<'a>> + 'a,
275 {
276 self.tree
277 .range(&range)
278 .map(|x| Range::new(x, self.transaction.transaction_guard()))
279 }
280
281 fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
282 self.tree.first()
283 }
284
285 fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
286 self.tree.last()
287 }
288}
289
290impl<K: Key, V: Value> Sealed for Table<'_, K, V> {}
291
292impl<K: Key + 'static, V: Value + 'static> Drop for Table<'_, K, V> {
293 fn drop(&mut self) {
294 self.transaction.close_table(
295 &self.name,
296 &self.tree,
297 self.tree.get_root().map(|x| x.length).unwrap_or_default(),
298 );
299 }
300}
301
302fn debug_helper<K: Key + 'static, V: Value + 'static>(
303 f: &mut Formatter<'_>,
304 name: &str,
305 len: Result<u64>,
306 first: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
307 last: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
308) -> std::fmt::Result {
309 write!(f, "Table [ name: \"{name}\", ")?;
310 if let Ok(len) = len {
311 if len == 0 {
312 write!(f, "No entries")?;
313 } else if len == 1 {
314 if let Ok(first) = first {
315 let (key, value) = first.as_ref().unwrap();
316 write!(f, "One key-value: {:?} = {:?}", key.value(), value.value())?;
317 } else {
318 write!(f, "I/O Error accessing table!")?;
319 }
320 } else {
321 if let Ok(first) = first {
322 let (key, value) = first.as_ref().unwrap();
323 write!(f, "first: {:?} = {:?}, ", key.value(), value.value())?;
324 } else {
325 write!(f, "I/O Error accessing table!")?;
326 }
327 if len > 2 {
328 write!(f, "...{} more entries..., ", len - 2)?;
329 }
330 if let Ok(last) = last {
331 let (key, value) = last.as_ref().unwrap();
332 write!(f, "last: {:?} = {:?}", key.value(), value.value())?;
333 } else {
334 write!(f, "I/O Error accessing table!")?;
335 }
336 }
337 } else {
338 write!(f, "I/O Error accessing table!")?;
339 }
340 write!(f, " ]")?;
341
342 Ok(())
343}
344
345impl<K: Key + 'static, V: Value + 'static> Debug for Table<'_, K, V> {
346 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
347 debug_helper(f, &self.name, self.len(), self.first(), self.last())
348 }
349}
350
351pub trait ReadableTableMetadata {
352 fn stats(&self) -> Result<TableStats>;
354
355 fn len(&self) -> Result<u64>;
357
358 fn is_empty(&self) -> Result<bool> {
360 Ok(self.len()? == 0)
361 }
362}
363
364pub trait ReadableTable<K: Key + 'static, V: Value + 'static>: ReadableTableMetadata {
365 fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>>;
367
368 fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
401 where
402 KR: Borrow<K::SelfType<'a>> + 'a;
403
404 fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
406
407 fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
409
410 fn iter(&self) -> Result<Range<K, V>> {
412 self.range::<K::SelfType<'_>>(..)
413 }
414}
415
416pub struct ReadOnlyUntypedTable {
418 tree: RawBtree,
419}
420
421impl Sealed for ReadOnlyUntypedTable {}
422
423impl ReadableTableMetadata for ReadOnlyUntypedTable {
424 fn stats(&self) -> Result<TableStats> {
426 let tree_stats = self.tree.stats()?;
427
428 Ok(TableStats {
429 tree_height: tree_stats.tree_height,
430 leaf_pages: tree_stats.leaf_pages,
431 branch_pages: tree_stats.branch_pages,
432 stored_leaf_bytes: tree_stats.stored_leaf_bytes,
433 metadata_bytes: tree_stats.metadata_bytes,
434 fragmented_bytes: tree_stats.fragmented_bytes,
435 })
436 }
437
438 fn len(&self) -> Result<u64> {
439 self.tree.len()
440 }
441}
442
443impl ReadOnlyUntypedTable {
444 pub(crate) fn new(
445 root_page: Option<BtreeHeader>,
446 fixed_key_size: Option<usize>,
447 fixed_value_size: Option<usize>,
448 mem: Arc<TransactionalMemory>,
449 ) -> Self {
450 Self {
451 tree: RawBtree::new(root_page, fixed_key_size, fixed_value_size, mem),
452 }
453 }
454}
455
456pub struct ReadOnlyTable<K: Key + 'static, V: Value + 'static> {
458 name: String,
459 tree: Btree<K, V>,
460 transaction_guard: Arc<TransactionGuard>,
461}
462
463impl<K: Key + 'static, V: Value + 'static> TableHandle for ReadOnlyTable<K, V> {
464 fn name(&self) -> &str {
465 &self.name
466 }
467}
468
469impl<K: Key + 'static, V: Value + 'static> ReadOnlyTable<K, V> {
470 pub(crate) fn new(
471 name: String,
472 root_page: Option<BtreeHeader>,
473 hint: PageHint,
474 guard: Arc<TransactionGuard>,
475 mem: Arc<TransactionalMemory>,
476 ) -> Result<ReadOnlyTable<K, V>> {
477 Ok(ReadOnlyTable {
478 name,
479 tree: Btree::new(root_page, hint, guard.clone(), mem)?,
480 transaction_guard: guard,
481 })
482 }
483
484 pub fn get<'a>(
485 &self,
486 key: impl Borrow<K::SelfType<'a>>,
487 ) -> Result<Option<AccessGuard<'static, V>>> {
488 self.tree.get(key.borrow())
489 }
490
491 pub fn range<'a, KR>(&self, range: impl RangeBounds<KR>) -> Result<Range<'static, K, V>>
494 where
495 KR: Borrow<K::SelfType<'a>>,
496 {
497 self.tree
498 .range(&range)
499 .map(|x| Range::new(x, self.transaction_guard.clone()))
500 }
501}
502
503impl<K: Key + 'static, V: Value + 'static> ReadableTableMetadata for ReadOnlyTable<K, V> {
504 fn stats(&self) -> Result<TableStats> {
505 let tree_stats = self.tree.stats()?;
506
507 Ok(TableStats {
508 tree_height: tree_stats.tree_height,
509 leaf_pages: tree_stats.leaf_pages,
510 branch_pages: tree_stats.branch_pages,
511 stored_leaf_bytes: tree_stats.stored_leaf_bytes,
512 metadata_bytes: tree_stats.metadata_bytes,
513 fragmented_bytes: tree_stats.fragmented_bytes,
514 })
515 }
516
517 fn len(&self) -> Result<u64> {
518 self.tree.len()
519 }
520}
521
522impl<K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for ReadOnlyTable<K, V> {
523 fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
524 self.tree.get(key.borrow())
525 }
526
527 fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
528 where
529 KR: Borrow<K::SelfType<'a>> + 'a,
530 {
531 self.tree
532 .range(&range)
533 .map(|x| Range::new(x, self.transaction_guard.clone()))
534 }
535
536 fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
537 self.tree.first()
538 }
539
540 fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
541 self.tree.last()
542 }
543}
544
545impl<K: Key, V: Value> Sealed for ReadOnlyTable<K, V> {}
546
547impl<K: Key + 'static, V: Value + 'static> Debug for ReadOnlyTable<K, V> {
548 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
549 debug_helper(f, &self.name, self.len(), self.first(), self.last())
550 }
551}
552
553pub struct ExtractIf<
554 'a,
555 K: Key + 'static,
556 V: Value + 'static,
557 F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
558> {
559 inner: BtreeExtractIf<'a, K, V, F>,
560}
561
562impl<
563 'a,
564 K: Key + 'static,
565 V: Value + 'static,
566 F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
567> ExtractIf<'a, K, V, F>
568{
569 pub(crate) fn new(inner: BtreeExtractIf<'a, K, V, F>) -> Self {
570 Self { inner }
571 }
572}
573
574impl<
575 'a,
576 K: Key + 'static,
577 V: Value + 'static,
578 F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
579> Iterator for ExtractIf<'a, K, V, F>
580{
581 type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
582
583 fn next(&mut self) -> Option<Self::Item> {
584 let entry = self.inner.next()?;
585 Some(entry.map(|entry| {
586 let (page, key_range, value_range) = entry.into_raw();
587 let key = AccessGuard::with_page(page.clone(), key_range);
588 let value = AccessGuard::with_page(page, value_range);
589 (key, value)
590 }))
591 }
592}
593
594impl<
595 K: Key + 'static,
596 V: Value + 'static,
597 F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
598> DoubleEndedIterator for ExtractIf<'_, K, V, F>
599{
600 fn next_back(&mut self) -> Option<Self::Item> {
601 let entry = self.inner.next_back()?;
602 Some(entry.map(|entry| {
603 let (page, key_range, value_range) = entry.into_raw();
604 let key = AccessGuard::with_page(page.clone(), key_range);
605 let value = AccessGuard::with_page(page, value_range);
606 (key, value)
607 }))
608 }
609}
610
611#[derive(Clone)]
612pub struct Range<'a, K: Key + 'static, V: Value + 'static> {
613 inner: BtreeRangeIter<K, V>,
614 _transaction_guard: Arc<TransactionGuard>,
615 _lifetime: PhantomData<&'a ()>,
617}
618
619impl<K: Key + 'static, V: Value + 'static> Range<'_, K, V> {
620 pub(super) fn new(inner: BtreeRangeIter<K, V>, guard: Arc<TransactionGuard>) -> Self {
621 Self {
622 inner,
623 _transaction_guard: guard,
624 _lifetime: Default::default(),
625 }
626 }
627}
628
629impl<'a, K: Key + 'static, V: Value + 'static> Iterator for Range<'a, K, V> {
630 type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
631
632 fn next(&mut self) -> Option<Self::Item> {
633 self.inner.next().map(|x| {
634 x.map(|entry| {
635 let (page, key_range, value_range) = entry.into_raw();
636 let key = AccessGuard::with_page(page.clone(), key_range);
637 let value = AccessGuard::with_page(page, value_range);
638 (key, value)
639 })
640 })
641 }
642}
643
644impl<K: Key + 'static, V: Value + 'static> DoubleEndedIterator for Range<'_, K, V> {
645 fn next_back(&mut self) -> Option<Self::Item> {
646 self.inner.next_back().map(|x| {
647 x.map(|entry| {
648 let (page, key_range, value_range) = entry.into_raw();
649 let key = AccessGuard::with_page(page.clone(), key_range);
650 let value = AccessGuard::with_page(page, value_range);
651 (key, value)
652 })
653 })
654 }
655}