Merge branch 'tomas/nested-lazy-vec-iter' (#1218)

* tomas/nested-lazy-vec-iter:
  changelog: add #1218
  core/storage_api/lazy_map: fix `iter` for map with nested collections
  core/storage_api/collections: add `is_data_sub_key` helper method
  core/storage_api: add predicate filtering version of prefix_iter
  test/lazy_vec: add test case to reproduce nested collection iter issue

.changelog/v0.15.0/bug-fixes/1218-nested-lazy-vec-iter.md

@@ -0,0 +1,3 @@
+- Fixed an issue with the iterator of LazyMap with a nested LazyVec collection
+  that would match non-data keys and fail to decode those with the data decoder.
+  ([#1218](https://github.com/anoma/namada/pull/1218))

core/src/ledger/storage_api/collections/lazy_map.rs

@@ -178,6 +178,21 @@ where
         }
     }
 
+    fn is_data_sub_key(&self, key: &storage::Key) -> bool {
+        let sub_key = self.is_valid_sub_key(key);
+        match sub_key {
+            Ok(Some(NestedSubKey::Data {
+                key: parsed_key,
+                nested_sub_key: _,
+            })) => {
+                let sub = self.at(&parsed_key);
+                // Check in the nested collection
+                sub.is_data_sub_key(key)
+            }
+            _ => false,
+        }
+    }
+
     fn read_sub_key_data<ENV>(
         env: &ENV,
         storage_key: &storage::Key,
@@ -303,6 +318,10 @@ where
         }
     }
 
+    fn is_data_sub_key(&self, key: &storage::Key) -> bool {
+        matches!(self.is_valid_sub_key(key), Ok(Some(_)))
+    }
+
     fn read_sub_key_data<ENV>(
         env: &ENV,
         storage_key: &storage::Key,
@@ -392,7 +411,11 @@ where
             )>,
         > + 'iter,
     > {
-        let iter = storage_api::iter_prefix(storage, &self.get_data_prefix())?;
+        let iter = storage_api::iter_prefix_with_filter(
+            storage,
+            &self.get_data_prefix(),
+            |key| self.is_data_sub_key(key),
+        )?;
         Ok(iter.map(|key_val_res| {
             let (key, val) = key_val_res?;
             let sub_key = LazyCollection::is_valid_sub_key(self, &key)?

core/src/ledger/storage_api/collections/lazy_set.rs

@@ -107,6 +107,10 @@ where
         }
     }
 
+    fn is_data_sub_key(&self, key: &storage::Key) -> bool {
+        matches!(self.is_valid_sub_key(key), Ok(Some(_)))
+    }
+
     fn read_sub_key_data<ENV>(
         env: &ENV,
         storage_key: &storage::Key,

core/src/ledger/storage_api/collections/lazy_vec.rs

@@ -174,6 +174,12 @@ where
         }
     }
 
+    fn is_data_sub_key(&self, key: &storage::Key) -> bool {
+        let sub_key = self.is_valid_sub_key(key);
+        // The `SubKey::Len` is not data sub-key
+        matches!(sub_key, Ok(Some(SubKey::Data(_))))
+    }
+
     fn read_sub_key_data<ENV>(
         env: &ENV,
         storage_key: &storage::Key,
@@ -487,6 +493,7 @@ where
 mod test {
     use super::*;
     use crate::ledger::storage::testing::TestWlStorage;
+    use crate::ledger::storage_api::collections::lazy_map::{self, NestedMap};
     use crate::types::address::{self, Address};
 
     #[test]
@@ -578,4 +585,67 @@ mod test {
 
         Ok(())
     }
+
+    /// Test iterator on a `LazyVec` nested inside a `LazyMap`
+    #[test]
+    fn test_nested_lazy_vec_iter() -> storage_api::Result<()> {
+        let mut storage = TestWlStorage::default();
+
+        let prefix = storage::Key::parse("test").unwrap();
+        let handle = NestedMap::<Address, LazyVec<u32>>::open(prefix);
+
+        let key = address::testing::established_address_1();
+
+        // Push first value and check iterator
+        handle.at(&key).push(&mut storage, 15)?;
+        let expected = (
+            lazy_map::NestedSubKey::Data {
+                key: key.clone(),                // LazyMap key
+                nested_sub_key: SubKey::Data(0), // LazyVec index
+            },
+            15, // the value
+        );
+
+        let mut iter = handle.iter(&storage)?;
+        assert_eq!(iter.next().unwrap()?, expected);
+        assert!(iter.next().is_none());
+        drop(iter);
+
+        // Push second value and check iterator again
+        handle.at(&key).push(&mut storage, 1)?;
+        let expected2 = (
+            lazy_map::NestedSubKey::Data {
+                key: key.clone(),                // LazyMap key
+                nested_sub_key: SubKey::Data(1), // LazyVec index
+            },
+            1, // the value
+        );
+
+        let mut iter = handle.iter(&storage)?;
+        assert_eq!(iter.next().unwrap()?, expected);
+        assert_eq!(iter.next().unwrap()?, expected2);
+        assert!(iter.next().is_none());
+        drop(iter);
+
+        let key2 = address::testing::established_address_2();
+        // Push third value on a different outer key and check iterator again
+        handle.at(&key2).push(&mut storage, 9)?;
+        let expected3 = (
+            lazy_map::NestedSubKey::Data {
+                key: key2.clone(),               // LazyMap key
+                nested_sub_key: SubKey::Data(0), // LazyVec index
+            },
+            9, // the value
+        );
+
+        let mut iter = handle.iter(&storage)?;
+        assert!(key < key2, "sanity check - this influences the iter order");
+        assert_eq!(iter.next().unwrap()?, expected);
+        assert_eq!(iter.next().unwrap()?, expected2);
+        assert_eq!(iter.next().unwrap()?, expected3);
+        assert!(iter.next().is_none());
+        drop(iter);
+
+        Ok(())
+    }
 }

core/src/ledger/storage_api/collections/mod.rs

@@ -75,6 +75,13 @@ pub trait LazyCollection {
         key: &storage::Key,
     ) -> storage_api::Result<Option<Self::SubKey>>;
 
+    /// Check if the given storage key is a valid data key.
+    ///
+    /// For most collections, this is the same as `is_valid_sub_key`, but for
+    /// example for `LazyVec`, which has an additional sub-key for length of the
+    /// vec, only the element data sub-keys would return `true`.
+    fn is_data_sub_key(&self, key: &storage::Key) -> bool;
+
     /// Try to read and decode the data for each change storage key in prior and
     /// posterior state. If there is no value in neither prior or posterior
     /// state (which is a possible state when transaction e.g. writes and then

core/src/ledger/storage_api/mod.rs

@@ -183,3 +183,63 @@ where
     });
     Ok(iter)
 }
+
+/// Iterate Borsh encoded items matching the given prefix and passing the given
+/// `filter` predicate, ordered by the storage keys.
+///
+/// The `filter` predicate is a function from a storage key to bool and only
+/// the items that return `true` will be returned from the iterator.
+///
+/// Note that this is preferable over the regular `iter_prefix` combined with
+/// the iterator's `filter` function as it avoids trying to decode values that
+/// don't pass the filter. For `iter_prefix_bytes`, `filter` works fine.
+pub fn iter_prefix_with_filter<'a, T, F>(
+    storage: &'a impl StorageRead,
+    prefix: &crate::types::storage::Key,
+    filter: F,
+) -> Result<impl Iterator<Item = Result<(storage::Key, T)>> + 'a>
+where
+    T: BorshDeserialize,
+    F: Fn(&storage::Key) -> bool + 'a,
+{
+    let iter = storage.iter_prefix(prefix)?;
+    let iter = itertools::unfold(iter, move |iter| {
+        // The loop is for applying filter - we `continue` when the current key
+        // doesn't pass the predicate.
+        loop {
+            match storage.iter_next(iter) {
+                Ok(Some((key, val))) => {
+                    let key =
+                        match storage::Key::parse(key).into_storage_result() {
+                            Ok(key) => key,
+                            Err(err) => {
+                                // Propagate key encoding errors into Iterator's
+                                // Item
+                                return Some(Err(err));
+                            }
+                        };
+                    // Check the predicate
+                    if !filter(&key) {
+                        continue;
+                    }
+                    let val =
+                        match T::try_from_slice(&val).into_storage_result() {
+                            Ok(val) => val,
+                            Err(err) => {
+                                // Propagate val encoding errors into Iterator's
+                                // Item
+                                return Some(Err(err));
+                            }
+                        };
+                    return Some(Ok((key, val)));
+                }
+                Ok(None) => return None,
+                Err(err) => {
+                    // Propagate `iter_next` errors into Iterator's Item
+                    return Some(Err(err));
+                }
+            }
+        }
+    });
+    Ok(iter)
+}