Repository: cberner/redb
Pull request #268 (2022-05-01 19:51:13)

Return old value when overwriting with insert()


Patch:
diff --git a/benches/common.rs b/benches/common.rs
index 85fdaa8..3c785da 100644
--- a/benches/common.rs
+++ b/benches/common.rs
@@ -117,7 +117,7 @@ pub struct RedbBenchInserter<'a> {
 
 impl BenchInserter for RedbBenchInserter<'_> {
     fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<(), ()> {
-        self.table.insert(key, value).map_err(|_| ())
+        self.table.insert(key, value).map(|_| ()).map_err(|_| ())
     }
 
     fn remove(&mut self, key: &[u8]) -> Result<(), ()> {
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
index 7bcfedf..9417521 100644
--- a/src/multimap_table.rs
+++ b/src/multimap_table.rs
@@ -66,7 +66,7 @@ impl<'a, K: RedbKey + ?Sized + 'a, V: RedbKey + ?Sized + 'a> RedbValue for Multi
         &self.data
     }
 
-    fn redb_type_name() -> &'static str {
+    fn redb_type_name() -> String {
         unreachable!()
     }
 }
@@ -323,13 +323,14 @@ impl<'db, 'txn, K: RedbKey + ?Sized, V: RedbKey + ?Sized> MultimapTable<'db, 'tx
     }
 
     /// Add the given value to the mapping of the key
-    // TODO: return bool indicating if the value already existed
-    pub fn insert(&mut self, key: &K, value: &V) -> Result {
+    ///
+    /// Returns `true` if the key-value pair was present
+    pub fn insert(&mut self, key: &K, value: &V) -> Result<bool> {
         let kv = MultimapKVPair::new_pair(key, value);
         // Safety: No other references to this table can exist.
         // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
         // and we borrow &mut self.
-        unsafe { self.tree.insert(&kv, b"") }
+        unsafe { self.tree.insert(&kv, b"").map(|x| x.is_some()) }
     }
 
     /// Removes the given key-value pair
diff --git a/src/table.rs b/src/table.rs
index 2a6dee6..2459895 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -37,8 +37,9 @@ impl<'db, 'txn, K: RedbKey + ?Sized, V: RedbValue + ?Sized> Table<'db, 'txn, K,
     }
 
     /// Insert mapping of the given key to the given value
-    // TODO: return the old value, if it exists
-    pub fn insert(&mut self, key: &K, value: &V) -> Result {
+    ///
+    /// Returns the old value, if the key was present in the table
+    pub fn insert(&mut self, key: &K, value: &V) -> Result<Option<AccessGuard<V>>> {
         // Safety: No other references to this table can exist.
         // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
         // and we borrow &mut self.
@@ -239,8 +240,8 @@ mod test {
                 &self.0
             }
 
-            fn redb_type_name() -> &'static str {
-                "ReverseKey"
+            fn redb_type_name() -> String {
+                "ReverseKey".to_string()
             }
         }
 
diff --git a/src/tree_store/btree.rs b/src/tree_store/btree.rs
index 11e412d..78df175 100644
--- a/src/tree_store/btree.rs
+++ b/src/tree_store/btree.rs
@@ -48,7 +48,7 @@ impl<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized> BtreeMut<'a, K, V> {
     }
 
     // Safety: caller must ensure that no uncommitted data is accessed within this tree, from other references
-    pub(crate) unsafe fn insert(&mut self, key: &K, value: &V) -> Result {
+    pub(crate) unsafe fn insert(&mut self, key: &K, value: &V) -> Result<Option<AccessGuard<V>>> {
         let mut freed_pages = self.freed_pages.borrow_mut();
         let mut operation = MutateHelper::new(
             &mut self.root,
@@ -56,8 +56,8 @@ impl<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized> BtreeMut<'a, K, V> {
             self.mem,
             freed_pages.as_mut(),
         );
-        operation.insert(key, value)?;
-        Ok(())
+        let (old_value, _) = operation.insert(key, value)?;
+        Ok(old_value)
     }
 
     /// Reserve space to insert a key-value pair
@@ -76,7 +76,7 @@ impl<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized> BtreeMut<'a, K, V> {
             self.mem,
             freed_pages.as_mut(),
         );
-        let guard = operation.insert(key, value.as_slice())?;
+        let (_, guard) = operation.insert(key, value.as_slice())?;
         Ok(guard)
     }
 
diff --git a/src/tree_store/btree_base.rs b/src/tree_store/btree_base.rs
index 94e46a8..b24ac5b 100644
--- a/src/tree_store/btree_base.rs
+++ b/src/tree_store/btree_base.rs
@@ -38,6 +38,13 @@ impl FreePolicy {
 
         Ok(())
     }
+
+    pub(crate) fn free_on_drop(&self, page: PageNumber, mem: &TransactionalMemory) -> bool {
+        match self {
+            FreePolicy::Never => false,
+            FreePolicy::Uncommitted => mem.uncommitted(page),
+        }
+    }
 }
 
 enum OnDrop {
@@ -49,6 +56,7 @@ enum OnDrop {
 enum EitherPage<'a> {
     Immutable(PageImpl<'a>),
     Mutable(PageMut<'a>),
+    OwnedMemory(Vec<u8>),
 }
 
 impl<'a> EitherPage<'a> {
@@ -56,6 +64,7 @@ impl<'a> EitherPage<'a> {
         match self {
             EitherPage::Immutable(page) => page.memory(),
             EitherPage::Mutable(page) => page.memory(),
+            EitherPage::OwnedMemory(mem) => mem.as_slice(),
         }
     }
 }
@@ -94,6 +103,18 @@ impl<'a, V: RedbValue + ?Sized> AccessGuard<'a, V> {
         }
     }
 
+    pub(super) fn with_owned_value(value: Vec<u8>, mem: &'a TransactionalMemory) -> Self {
+        let len = value.len();
+        Self {
+            page: EitherPage::OwnedMemory(value),
+            offset: 0,
+            len,
+            on_drop: OnDrop::None,
+            mem,
+            _value_type: Default::default(),
+        }
+    }
+
     // Safety: if free_on_drop is true, caller must guarantee that no other references to page exist,
     // and that no references will be created until this AccessGuard is dropped
     pub(super) unsafe fn remove_on_drop(
diff --git a/src/tree_store/btree_mutator.rs b/src/tree_store/btree_mutator.rs
index d5dd221..a146bee 100644
--- a/src/tree_store/btree_mutator.rs
+++ b/src/tree_store/btree_mutator.rs
@@ -26,6 +26,17 @@ enum DeletionResult {
     DeletedBranch(PageNumber),
 }
 
+struct InsertionResult<'a, V: RedbValue + ?Sized> {
+    // The new root page
+    new_root: PageNumber,
+    // Following sibling, if the root had to be split
+    additional_sibling: Option<(Vec<u8>, PageNumber)>,
+    // The inserted value for .insert_reserve() to use
+    inserted_value: AccessGuardMut<'a>,
+    // The previous value, if any
+    old_value: Option<AccessGuard<'a, V>>,
+}
+
 pub(crate) struct MutateHelper<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> {
     root: &'b mut Option<PageNumber>,
     free_policy: FreePolicy,
@@ -84,24 +95,28 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
     }
 
     // Safety: caller must ensure that no references to uncommitted pages in this tree exist
-    pub(crate) unsafe fn insert(&mut self, key: &K, value: &V) -> Result<AccessGuardMut<'a>> {
-        let (new_root, guard) = if let Some(p) = *self.root {
-            let (page1, more, guard) = self.insert_helper(
+    pub(crate) unsafe fn insert(
+        &mut self,
+        key: &K,
+        value: &V,
+    ) -> Result<(Option<AccessGuard<'a, V>>, AccessGuardMut<'a>)> {
+        let (new_root, old_value, guard) = if let Some(p) = *self.root {
+            let result = self.insert_helper(
                 self.mem.get_page(p),
                 key.as_bytes().as_ref(),
                 value.as_bytes().as_ref(),
             )?;
 
-            let new_root = if let Some((key, page2)) = more {
+            let new_root = if let Some((key, page2)) = result.additional_sibling {
                 let mut builder = BranchBuilder::new(self.mem, 2);
-                builder.push_child(page1);
+                builder.push_child(result.new_root);
                 builder.push_key(&key);
                 builder.push_child(page2);
                 builder.build()?.get_page_number()
             } else {
-                page1
+                result.new_root
             };
-            (new_root, guard)
+            (new_root, result.old_value, result.inserted_value)
         } else {
             let key_bytes = key.as_bytes();
             let value_bytes = value.as_bytes();
@@ -116,24 +131,19 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
             let page_num = page.get_page_number();
             let guard = AccessGuardMut::new(page, offset, value_bytes.len());
 
-            (page_num, guard)
+            (page_num, None, guard)
         };
         *self.root = Some(new_root);
-        Ok(guard)
+        Ok((old_value, guard))
     }
 
-    #[allow(clippy::type_complexity)]
     // Safety: caller must ensure that no references to uncommitted pages in this table exist
     unsafe fn insert_helper(
         &mut self,
         page: PageImpl<'a>,
         key: &[u8],
         value: &[u8],
-    ) -> Result<(
-        PageNumber,
-        Option<(Vec<u8>, PageNumber)>,
-        AccessGuardMut<'a>,
-    )> {
+    ) -> Result<InsertionResult<'a, V>> {
         let node_mem = page.memory();
         Ok(match node_mem[0] {
             LEAF => {
@@ -153,18 +163,20 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                     drop(new_page_accessor);
                     let guard = AccessGuardMut::new(new_page, offset, value.len());
                     return if position == 0 {
-                        Ok((
-                            new_page_number,
-                            Some((key.to_vec(), page.get_page_number())),
-                            guard,
-                        ))
+                        Ok(InsertionResult {
+                            new_root: new_page_number,
+                            additional_sibling: Some((key.to_vec(), page.get_page_number())),
+                            inserted_value: guard,
+                            old_value: None,
+                        })
                     } else {
                         let split_key = accessor.last_entry().key().to_vec();
-                        Ok((
-                            page.get_page_number(),
-                            Some((split_key, new_page_number)),
-                            guard,
-                        ))
+                        Ok(InsertionResult {
+                            new_root: page.get_page_number(),
+                            additional_sibling: Some((split_key, new_page_number)),
+                            inserted_value: guard,
+                            old_value: None,
+                        })
                     };
                 }
 
@@ -175,6 +187,12 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                     )
                 {
                     let page_number = page.get_page_number();
+                    let existing_value = if found {
+                        let copied_value = accessor.entry(position).unwrap().value().to_vec();
+                        Some(AccessGuard::with_owned_value(copied_value, self.mem))
+                    } else {
+                        None
+                    };
                     drop(page);
                     let mut page_mut = self.mem.get_page_mut(page_number);
                     let mut mutator = LeafMutator::new(&mut page_mut);
@@ -183,7 +201,12 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                     let offset = new_page_accessor.offset_of_value(position).unwrap();
                     drop(new_page_accessor);
                     let guard = AccessGuardMut::new(page_mut, offset, value.len());
-                    return Ok((page_number, None, guard));
+                    return Ok(InsertionResult {
+                        new_root: page_number,
+                        additional_sibling: None,
+                        inserted_value: guard,
+                        old_value: existing_value,
+                    });
                 }
 
                 let mut builder = LeafBuilder::new(self.mem, accessor.num_pairs() + 1);
@@ -203,23 +226,60 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                     let new_page = builder.build()?;
 
                     let page_number = page.get_page_number();
-                    drop(page);
-                    self.free_policy
-                        .conditional_free(page_number, self.freed, self.mem)?;
+                    let existing_value = if found {
+                        let (start, end) = accessor.value_range(position).unwrap();
+                        let free_on_drop = self.free_policy.free_on_drop(page_number, self.mem);
+                        if !free_on_drop {
+                            self.freed.push(page_number);
+                        }
+                        Some(AccessGuard::new(
+                            page,
+                            start,
+                            end - start,
+                            free_on_drop,
+                            self.mem,
+                        ))
+                    } else {
+                        drop(page);
+                        self.free_policy
+                            .conditional_free(page_number, self.freed, self.mem)?;
+                        None
+                    };
 
                     let new_page_number = new_page.get_page_number();
                     let accessor = LeafAccessor::new(&new_page);
                     let offset = accessor.offset_of_value(position).unwrap();
                     let guard = AccessGuardMut::new(new_page, offset, value.len());
 
-                    (new_page_number, None, guard)
+                    InsertionResult {
+                        new_root: new_page_number,
+                        additional_sibling: None,
+                        inserted_value: guard,
+                        old_value: existing_value,
+                    }
                 } else {
                     let (new_page1, split_key, new_page2) = builder.build_split()?;
-                    let page_number = page.get_page_number();
                     let split_key = split_key.to_vec();
-                    drop(page);
-                    self.free_policy
-                        .conditional_free(page_number, self.freed, self.mem)?;
+                    let page_number = page.get_page_number();
+                    let existing_value = if found {
+                        let (start, end) = accessor.value_range(position).unwrap();
+                        let free_on_drop = self.free_policy.free_on_drop(page_number, self.mem);
+                        if !free_on_drop {
+                            self.freed.push(page_number);
+                        }
+                        Some(AccessGuard::new(
+                            page,
+                            start,
+                            end - start,
+                            free_on_drop,
+                            self.mem,
+                        ))
+                    } else {
+                        drop(page);
+                        self.free_policy
+                            .conditional_free(page_number, self.freed, self.mem)?;
+                        None
+                    };
 
                     let new_page_number = new_page1.get_page_number();
                     let new_page_number2 = new_page2.get_page_number();
@@ -235,20 +295,29 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                         AccessGuardMut::new(new_page2, offset, value.len())
                     };
 
-                    (new_page_number, Some((split_key, new_page_number2)), guard)
+                    InsertionResult {
+                        new_root: new_page_number,
+                        additional_sibling: Some((split_key, new_page_number2)),
+                        inserted_value: guard,
+                        old_value: existing_value,
+                    }
                 }
             }
             BRANCH => {
                 let accessor = BranchAccessor::new(&page);
                 let (child_index, child_page) = accessor.child_for_key::<K>(key);
-                let (page1, more, guard) =
-                    self.insert_helper(self.mem.get_page(child_page), key, value)?;
+                let sub_result = self.insert_helper(self.mem.get_page(child_page), key, value)?;
 
-                if more.is_none() {
+                if sub_result.additional_sibling.is_none() {
                     // Check fast-path if no children were added
-                    if page1 == child_page {
+                    if sub_result.new_root == child_page {
                         // NO-OP. One of our descendants is uncommitted, so there was no change
-                        return Ok((page.get_page_number(), None, guard));
+                        return Ok(InsertionResult {
+                            new_root: page.get_page_number(),
+                            additional_sibling: None,
+                            inserted_value: sub_result.inserted_value,
+                            old_value: sub_result.old_value,
+                        });
                     } else if self.mem.uncommitted(page.get_page_number()) {
                         let page_number = page.get_page_number();
                         drop(page);
@@ -256,16 +325,21 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                         // and we just dropped our reference to it, on the line above
                         let mut mutpage = self.mem.get_page_mut(page_number);
                         let mut mutator = BranchMutator::new(&mut mutpage);
-                        mutator.write_child_page(child_index, page1);
-                        return Ok((mutpage.get_page_number(), None, guard));
+                        mutator.write_child_page(child_index, sub_result.new_root);
+                        return Ok(InsertionResult {
+                            new_root: mutpage.get_page_number(),
+                            additional_sibling: None,
+                            inserted_value: sub_result.inserted_value,
+                            old_value: sub_result.old_value,
+                        });
                     }
                 }
 
                 // A child was added, or we couldn't use the fast-path above
                 let mut builder = BranchBuilder::new(self.mem, accessor.count_children() + 1);
                 if child_index == 0 {
-                    builder.push_child(page1);
-                    if let Some((ref index_key2, page2)) = more {
+                    builder.push_child(sub_result.new_root);
+                    if let Some((ref index_key2, page2)) = sub_result.additional_sibling {
                         builder.push_key(index_key2);
                         builder.push_child(page2);
                     }
@@ -276,8 +350,8 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
                     if let Some(key) = accessor.key(i - 1) {
                         builder.push_key(key);
                         if i == child_index {
-                            builder.push_child(page1);
-                            if let Some((ref index_key2, page2)) = more {
+                            builder.push_child(sub_result.new_root);
+                            if let Some((ref index_key2, page2)) = sub_result.additional_sibling {
                                 builder.push_key(index_key2);
                                 builder.push_child(page2);
                             }
@@ -291,14 +365,20 @@ impl<'a, 'b, K: RedbKey + ?Sized, V: RedbValue + ?Sized> MutateHelper<'a, 'b, K,
 
                 let result = if builder.should_split() {
                     let (new_page1, split_key, new_page2) = builder.build_split()?;
-                    (
-                        new_page1.get_page_number(),
-                        Some((split_key.to_vec(), new_page2.get_page_number())),
-                        guard,
-                    )
+                    InsertionResult {
+                        new_root: new_page1.get_page_number(),
+                        additional_sibling: Some((split_key.to_vec(), new_page2.get_page_number())),
+                        inserted_value: sub_result.inserted_value,
+                        old_value: sub_result.old_value,
+                    }
                 } else {
                     let new_page = builder.build()?;
-                    (new_page.get_page_number(), None, guard)
+                    InsertionResult {
+                        new_root: new_page.get_page_number(),
+                        additional_sibling: None,
+                        inserted_value: sub_result.inserted_value,
+                        old_value: sub_result.old_value,
+                    }
                 };
                 // Free the original page, since we've replaced it
                 let page_number = page.get_page_number();
diff --git a/src/tree_store/table_tree.rs b/src/tree_store/table_tree.rs
index 5bfecdb..54bc986 100644
--- a/src/tree_store/table_tree.rs
+++ b/src/tree_store/table_tree.rs
@@ -36,8 +36,8 @@ impl RedbValue for FreedTableKey {
         result
     }
 
-    fn redb_type_name() -> &'static str {
-        "FreedTableKey"
+    fn redb_type_name() -> String {
+        "FreedTableKey".to_string()
     }
 }
 
@@ -156,8 +156,8 @@ impl RedbValue for InternalTableDefinition {
         result
     }
 
-    fn redb_type_name() -> &'static str {
-        "InternalTableDefinition"
+    fn redb_type_name() -> String {
+        "InternalTableDefinition".to_string()
     }
 }
 
@@ -301,8 +301,8 @@ impl<'txn> TableTree<'txn> {
         let table = InternalTableDefinition {
             table_root: None,
             table_type,
-            key_type: K::redb_type_name().to_string(),
-            value_type: V::redb_type_name().to_string(),
+            key_type: K::redb_type_name(),
+            value_type: V::redb_type_name(),
         };
         // Safety: References into the master table are never returned to the user
         unsafe { self.tree.insert(name, &table)? };
diff --git a/src/types.rs b/src/types.rs
index 8cb889f..47dd875 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -51,7 +51,7 @@ pub trait RedbValue {
     fn as_bytes(&self) -> <Self::ToBytes as AsBytesWithLifetime>::Out;
 
     /// Globally unique identifier for this type
-    fn redb_type_name() -> &'static str;
+    fn redb_type_name() -> String;
 }
 
 pub trait RedbKey: RedbValue {
@@ -71,8 +71,8 @@ impl RedbValue for [u8] {
         self
     }
 
-    fn redb_type_name() -> &'static str {
-        "[u8]"
+    fn redb_type_name() -> String {
+        "[u8]".to_string()
     }
 }
 
@@ -94,8 +94,8 @@ impl RedbValue for str {
         self
     }
 
-    fn redb_type_name() -> &'static str {
-        "str"
+    fn redb_type_name() -> String {
+        "str".to_string()
     }
 }
 
@@ -121,8 +121,8 @@ macro_rules! be_value {
                 self.to_le_bytes()
             }
 
-            fn redb_type_name() -> &'static str {
-                stringify!($t)
+            fn redb_type_name() -> String {
+                stringify!($t).to_string()
             }
         }
     };


Test patch:
diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
index 4535f96..9654f4b 100644
--- a/tests/basic_tests.rs
+++ b/tests/basic_tests.rs
@@ -171,7 +171,7 @@ fn insert_overwrite() {
     let write_txn = db.begin_write().unwrap();
     {
         let mut table = write_txn.open_table(SLICE_TABLE).unwrap();
-        table.insert(b"hello", b"world").unwrap();
+        assert!(table.insert(b"hello", b"world").unwrap().is_none());
     }
     write_txn.commit().unwrap();
 
@@ -182,7 +182,8 @@ fn insert_overwrite() {
     let write_txn = db.begin_write().unwrap();
     {
         let mut table = write_txn.open_table(SLICE_TABLE).unwrap();
-        table.insert(b"hello", b"replaced").unwrap();
+        let old_value = table.insert(b"hello", b"replaced").unwrap();
+        assert_eq!(old_value.unwrap().to_value(), b"world");
     }
     write_txn.commit().unwrap();
 
diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs
index 372c889..e8c0d75 100644
--- a/tests/multimap_tests.rs
+++ b/tests/multimap_tests.rs
@@ -59,8 +59,9 @@ fn insert() {
     let write_txn = db.begin_write().unwrap();
     {
         let mut table = write_txn.open_multimap_table(SLICE_TABLE).unwrap();
-        table.insert(b"hello", b"world").unwrap();
-        table.insert(b"hello", b"world2").unwrap();
+        assert!(!table.insert(b"hello", b"world").unwrap());
+        assert!(!table.insert(b"hello", b"world2").unwrap());
+        assert!(table.insert(b"hello", b"world2").unwrap());
     }
     write_txn.commit().unwrap();