Merge pull request AleoNet#2219 from ljedrz/feat/map_len

Add length methods to MapRead & NestedMapRead

ledger/store/src/helpers/memory/internal/map.rs

@@ -242,6 +242,13 @@ impl<
         core::iter::Map<indexmap::map::IntoIter<K, Option<V>>, fn((K, Option<V>)) -> (Cow<'a, K>, Option<Cow<'a, V>>)>;
     type Values = core::iter::Map<btree_map::IntoValues<Vec<u8>, V>, fn(V) -> Cow<'a, V>>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_confirmed(&self) -> usize {
+        self.map.read().len()
+    }
+
     ///
     /// Returns `true` if the given key exists in the map.
     ///

ledger/store/src/helpers/memory/internal/nested_map.rs

@@ -263,6 +263,16 @@ impl<
     >;
     type Values = core::iter::Map<btree_map::IntoValues<Vec<u8>, V>, fn(V) -> Cow<'a, V>>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_map_confirmed(&self, map: &M) -> Result<usize> {
+        // Serialize 'm'.
+        let m = bincode::serialize(map)?;
+        // Retrieve the keys for the serialized map.
+        Ok(self.map.read().get(&m).map(|keys| keys.len()).unwrap_or_default())
+    }
+
     ///
     /// Returns `true` if the given key exists in the map.
     ///

ledger/store/src/helpers/rocksdb/internal/map.rs

@@ -243,6 +243,33 @@ impl<
         core::iter::Map<indexmap::map::IntoIter<K, Option<V>>, fn((K, Option<V>)) -> (Cow<'a, K>, Option<Cow<'a, V>>)>;
     type Values = Values<'a, V>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_confirmed(&self) -> usize {
+        // A raw iterator doesn't allocate.
+        let mut iter = self.database.raw_iterator();
+        // Find the first key with the map prefix.
+        iter.seek(&self.context);
+
+        // Count the number of keys belonging to the map.
+        let mut len = 0usize;
+        while let Some(key) = iter.key() {
+            // Only compare the map ID - the network ID is guaranteed to
+            // remain the same as long as there is more than a single map.
+            if key[2..][..2] != self.context[2..][..2] {
+                // If the map ID is different, it's the end of iteration.
+                break;
+            }
+
+            // Increment the length and go to the next record.
+            len += 1;
+            iter.next();
+        }
+
+        len
+    }
+
     ///
     /// Returns `true` if the given key exists in the map.
     ///
@@ -1387,9 +1414,13 @@ mod tests {
 
         // Ensure that all the items are present.
         assert_eq!(test_storage.own_map.iter_confirmed().count(), 1);
+        assert_eq!(test_storage.own_map.len_confirmed(), 1);
         assert_eq!(test_storage.extra_maps.own_map1.iter_confirmed().count(), 1);
+        assert_eq!(test_storage.extra_maps.own_map1.len_confirmed(), 1);
         assert_eq!(test_storage.extra_maps.own_map2.iter_confirmed().count(), 1);
+        assert_eq!(test_storage.extra_maps.own_map2.len_confirmed(), 1);
         assert_eq!(test_storage.extra_maps.extra_maps.own_map.iter_confirmed().count(), 1);
+        assert_eq!(test_storage.extra_maps.extra_maps.own_map.len_confirmed(), 1);
 
         // The atomic_write_batch macro uses ?, so the test returns a Result for simplicity.
         Ok(())

ledger/store/src/helpers/rocksdb/internal/nested_map.rs

@@ -349,6 +349,38 @@ impl<
     >;
     type Values = NestedValues<'a, V>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_map_confirmed(&self, map: &M) -> Result<usize> {
+        // Obtain the nested map prefix and its final part.
+        let prefix = self.create_prefixed_map(map)?;
+        let serialized_map = &prefix[PREFIX_LEN + 4..];
+
+        // A raw iterator doesn't allocate.
+        let mut iter = self.database.raw_iterator();
+        // Find the first key with the nested map prefix.
+        iter.seek(&prefix);
+
+        // Count the number of keys belonging to the nested map.
+        let mut len = 0usize;
+        while let Some(key) = iter.key() {
+            // Only compare the nested map - the network ID and the outer map
+            // ID are guaranteed to remain the same as long as there is more
+            // than a single map in the database.
+            if !key[PREFIX_LEN + 4..].starts_with(serialized_map) {
+                // If the nested map ID is different, it's the end of iteration.
+                break;
+            }
+
+            // Increment the length and go to the next record.
+            len += 1;
+            iter.next();
+        }
+
+        Ok(len)
+    }
+
     ///
     /// Returns `true` if the given map and key exists.
     ///

ledger/store/src/helpers/test_helpers/map/check_insert_and_get_speculative.rs

@@ -30,6 +30,9 @@ pub fn check_insert_and_get_speculative(map: impl for<'a> Map<'a, usize, String>
 
     // Check that the item is not yet in the map.
     assert!(map.get_confirmed(&0).unwrap().is_none());
+    // Check that the map has no entries.
+    assert_eq!(map.len_confirmed(), 0);
+    assert!(map.is_empty_confirmed());
     // Check that the item is in the batch.
     assert_eq!(map.get_pending(&0), Some(Some("0".to_string())));
     // Check that the item can be speculatively retrieved.
@@ -50,12 +53,16 @@ pub fn check_insert_and_get_speculative(map: impl for<'a> Map<'a, usize, String>
 
     // The map should still contain no items.
     assert!(map.iter_confirmed().next().is_none());
+    assert_eq!(map.len_confirmed(), 0);
+    assert!(map.is_empty_confirmed());
 
     // Finish the current atomic write batch.
     map.finish_atomic().unwrap();
 
     // Check that the item is present in the map now.
     assert_eq!(map.get_confirmed(&0).unwrap(), Some(Cow::Owned("9".to_string())));
+    assert_eq!(map.len_confirmed(), 1);
+    assert!(!map.is_empty_confirmed());
     // Check that the item is not in the batch.
     assert_eq!(map.get_pending(&0), None);
     // Check that the item can be speculatively retrieved.

ledger/store/src/helpers/test_helpers/nested_map/check_insert_and_get_value_speculative.rs

@@ -30,6 +30,8 @@ pub fn check_insert_and_get_value_speculative(map: impl for<'a> NestedMap<'a, us
 
     // Check that the item is not yet in the map.
     assert!(map.get_value_confirmed(&0, &0).unwrap().is_none());
+    assert_eq!(map.len_map_confirmed(&0).unwrap(), 0);
+    assert!(map.is_empty_map_confirmed(&0).unwrap());
     // Check that the item is in the batch.
     assert_eq!(map.get_value_pending(&0, &0), Some(Some("0".to_string())));
     // Check that the item can be speculatively retrieved.
@@ -50,12 +52,16 @@ pub fn check_insert_and_get_value_speculative(map: impl for<'a> NestedMap<'a, us
 
     // The map should still contain no items.
     assert!(map.iter_confirmed().next().is_none());
+    assert_eq!(map.len_map_confirmed(&0).unwrap(), 0);
+    assert!(map.is_empty_map_confirmed(&0).unwrap());
 
     // Finish the current atomic write batch.
     map.finish_atomic().unwrap();
 
     // Check that the item is present in the map now.
     assert_eq!(map.get_value_confirmed(&0, &0).unwrap(), Some(Cow::Owned("9".to_string())));
+    assert_eq!(map.len_map_confirmed(&0).unwrap(), 1);
+    assert!(!map.is_empty_map_confirmed(&0).unwrap());
     // Check that the item is not in the batch.
     assert_eq!(map.get_value_pending(&0, &0), None);
     // Check that the item can be speculatively retrieved.

ledger/store/src/helpers/traits/map.rs

@@ -87,6 +87,18 @@ pub trait MapRead<
     type Keys: Iterator<Item = Cow<'a, K>>;
     type Values: Iterator<Item = Cow<'a, V>>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_confirmed(&self) -> usize;
+
+    ///
+    /// Checks whether there are any confirmed entries in the map.
+    ///
+    fn is_empty_confirmed(&self) -> bool {
+        self.len_confirmed() == 0
+    }
+
     ///
     /// Returns `true` if the given key exists in the map.
     ///

ledger/store/src/helpers/traits/nested_map.rs

@@ -94,6 +94,18 @@ pub trait NestedMapRead<
     type Keys: Iterator<Item = (Cow<'a, M>, Cow<'a, K>)>;
     type Values: Iterator<Item = Cow<'a, V>>;
 
+    ///
+    /// Returns the number of confirmed entries in the map.
+    ///
+    fn len_map_confirmed(&self, map: &M) -> Result<usize>;
+
+    ///
+    /// Checks whether there are any confirmed entries in the map.
+    ///
+    fn is_empty_map_confirmed(&self, map: &M) -> Result<bool> {
+        Ok(self.len_map_confirmed(map)? == 0)
+    }
+
     ///
     /// Returns `true` if the given key exists in the map.
     ///