chore: allow creating custom snapshot (#60)

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "vart"
 publish = true
-version = "0.6.1"
+version = "0.6.2"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"

src/art.rs

@@ -1353,6 +1353,31 @@ impl<P: KeyTrait, V: Clone> Tree<P, V> {
         Snapshot::new(root, version)
     }
 
+    /// Creates a new snapshot of the Trie at the specified version.
+    ///
+    /// This function creates a snapshot of the Trie at the given version. The snapshot
+    /// captures the state of the Trie as it was at the specified version. This can be useful
+    /// for versioned data access, allowing you to interact with the Trie as it existed at
+    /// a particular point in time.
+    ///
+    /// # Arguments
+    ///
+    /// * `version`: The version number at which to create the snapshot.
+    ///
+    /// # Returns
+    ///
+    /// Returns a `Snapshot` that can be used to interact with the Trie at the specified version.
+    pub fn create_snapshot_at_version(&self, version: u64) -> Result<Snapshot<P, V>, TrieError> {
+        if let Some(root) = self.root.as_ref() {
+            if version < root.version() {
+                return Err(TrieError::SnapshotOlderThanRoot);
+            }
+        }
+
+        let root = self.root.as_ref().cloned();
+        Ok(Snapshot::new(root, version))
+    }
+
     /// Creates an iterator over the Trie's key-value pairs.
     ///
     /// This function creates and returns an iterator that can be used to traverse the key-value pairs

src/lib.rs

@@ -331,6 +331,7 @@ pub enum TrieError {
     FixedSizeKeyLengthExceeded,
     VersionIsOld,
     RootIsNotUniquelyOwned,
+    SnapshotOlderThanRoot,
 }
 
 impl Error for TrieError {}
@@ -344,6 +345,7 @@ impl fmt::Display for TrieError {
                 write!(f, "Given version is older than root's current version")
             }
             TrieError::RootIsNotUniquelyOwned => write!(f, "Root arc is not uniquely owned"),
+            TrieError::SnapshotOlderThanRoot => write!(f, "Snapshot is older than root"),
         }
     }
 }

src/snapshot.rs

@@ -10,22 +10,22 @@ use crate::KeyTrait;
 #[derive(Clone)]
 /// Represents a snapshot of the data within the Trie.
 pub struct Snapshot<P: KeyTrait, V: Clone> {
-    pub(crate) ts: u64,
+    pub(crate) version: u64,
     pub(crate) root: Option<Arc<Node<P, V>>>,
 }
 
 impl<P: KeyTrait, V: Clone> Snapshot<P, V> {
     /// Creates a new Snapshot instance with the provided snapshot_id and root node.
-    pub(crate) fn new(root: Option<Arc<Node<P, V>>>, ts: u64) -> Self {
-        Snapshot { ts, root }
+    pub(crate) fn new(root: Option<Arc<Node<P, V>>>, version: u64) -> Self {
+        Snapshot { version, root }
     }
 
     /// Inserts a key-value pair into the snapshot.
     pub fn insert(&mut self, key: &P, value: V, ts: u64) {
         // Insert the key-value pair into the root node using a recursive function
         match &self.root {
             Some(root) => {
-                let new_node = Node::insert_recurse(root, key, value, self.ts, ts, 0, false);
+                let new_node = Node::insert_recurse(root, key, value, self.version, ts, 0, false);
                 // Update the root node with the new node after insertion
                 self.root = Some(new_node);
             }
@@ -34,7 +34,7 @@ impl<P: KeyTrait, V: Clone> Snapshot<P, V> {
                     key.as_slice().into(),
                     key.as_slice().into(),
                     value,
-                    self.ts,
+                    self.version,
                     ts,
                 )))
             }
@@ -76,7 +76,7 @@ impl<P: KeyTrait, V: Clone> Snapshot<P, V> {
     }
 
     pub fn ts(&self) -> u64 {
-        self.ts
+        self.version
     }
 
     pub fn iter(&self) -> Iter<P, V> {
@@ -828,4 +828,89 @@ mod tests {
         let keys = snap.keys_at_ts(RangeFull {}, 0);
         assert_eq!(keys.len(), 1);
     }
+
+    #[test]
+    fn snapshot_creation_at_version() {
+        let mut tree: Tree<VariableSizeKey, i32> = Tree::<VariableSizeKey, i32>::new();
+        let keys = ["key_1", "key_2", "key_3"];
+
+        for key in keys.iter() {
+            assert!(tree
+                .insert(&VariableSizeKey::from_str(key).unwrap(), 1, 0, 0)
+                .is_ok());
+        }
+
+        // Create a snapshot at the current version of the tree
+        let current_version = tree.version();
+        assert_eq!(current_version, keys.len() as u64);
+
+        let mut snap1 = tree.create_snapshot_at_version(current_version).unwrap();
+
+        let key_to_insert = "key_1";
+        snap1.insert(&VariableSizeKey::from_str(key_to_insert).unwrap(), 1, 0);
+
+        let expected_snap_ts = current_version;
+        assert_eq!(snap1.version(), expected_snap_ts);
+
+        let expected_tree_ts = current_version;
+        assert_eq!(tree.version(), expected_tree_ts);
+    }
+
+    #[test]
+    fn snapshot_isolation_at_version() {
+        let mut tree: Tree<VariableSizeKey, i32> = Tree::<VariableSizeKey, i32>::new();
+        let key_1 = VariableSizeKey::from_str("key_1").unwrap();
+        let key_2 = VariableSizeKey::from_str("key_2").unwrap();
+        let key_3_snap1 = VariableSizeKey::from_str("key_3_snap1").unwrap();
+        let key_3_snap2 = VariableSizeKey::from_str("key_3_snap2").unwrap();
+
+        assert!(tree.insert(&key_1, 1, 0, 0).is_ok());
+
+        // Create snapshots at the current version of the tree
+        let current_version = tree.version() + 1;
+        let mut snap1 = tree.create_snapshot_at_version(current_version).unwrap();
+        assert_eq!(snap1.get(&key_1).unwrap(), (1, 1, 0));
+
+        let mut snap2 = tree.create_snapshot_at_version(current_version).unwrap();
+        assert_eq!(snap2.get(&key_1).unwrap(), (1, 1, 0));
+
+        // Keys inserted after snapshot creation should not be visible to other snapshots
+        assert!(tree.insert(&key_2, 1, 0, 0).is_ok());
+        assert!(snap1.get(&key_2).is_none());
+        assert!(snap2.get(&key_2).is_none());
+
+        // Keys inserted after snapshot creation should be visible to the snapshot that inserted them
+        snap1.insert(&key_3_snap1, 2, 0);
+        assert_eq!(snap1.get(&key_3_snap1).unwrap(), (2, current_version, 0));
+
+        snap2.insert(&key_3_snap2, 3, 0);
+        assert_eq!(snap2.get(&key_3_snap2).unwrap(), (3, current_version, 0));
+
+        // Keys inserted after snapshot creation should not be visible to other snapshots
+        assert!(snap1.get(&key_3_snap2).is_none());
+        assert!(snap2.get(&key_3_snap1).is_none());
+    }
+
+    #[test]
+    fn snapshot_creation_at_invalid_version() {
+        let mut tree: Tree<VariableSizeKey, i32> = Tree::<VariableSizeKey, i32>::new();
+        let keys = ["key_1", "key_2", "key_3"];
+
+        for key in keys.iter() {
+            assert!(tree
+                .insert(&VariableSizeKey::from_str(key).unwrap(), 1, 0, 0)
+                .is_ok());
+        }
+
+        // Create a snapshot at the current version of the tree
+        let current_version = tree.version();
+        assert_eq!(current_version, keys.len() as u64);
+
+        // Attempt to create a snapshot at a version less than the current version
+        let invalid_version = current_version - 1;
+        let result = tree.create_snapshot_at_version(invalid_version);
+
+        // Ensure that an error is returned
+        assert!(result.is_err());
+    }
 }