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