lazy collections with Address

.changelog/unreleased/bug-fixes/1212-lazy-collection-sub-key.md

@@ -0,0 +1,3 @@
+- Fixed an issue with lazy collections sub-key validation with the `Address`
+  type. This issue was also affecting the iterator of nested `LazyMap`.
+  ([#1212](https://github.com/anoma/namada/pull/1212))

.changelog/unreleased/bug-fixes/1256-fix-addr-storage-key-ord.md

@@ -0,0 +1,3 @@
+- Addresses are now being ordered by their string format (bech32m)
+  to ensure that their order is preserved inside raw storage keys.
+  ([#1256](https://github.com/anoma/namada/pull/1256))

core/src/ledger/storage/wl_storage.rs

@@ -183,10 +183,9 @@ where
                         what = Next::ReturnStorage;
                     }
                     (Some((storage_key, _, _)), Some((wl_key, _))) => {
-                        let wl_key = wl_key.to_string();
-                        if &wl_key <= storage_key {
+                        if wl_key <= storage_key {
                             what = Next::ReturnWl {
-                                advance_storage: &wl_key == storage_key,
+                                advance_storage: wl_key == storage_key,
                             };
                         } else {
                             what = Next::ReturnStorage;

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

@@ -40,7 +40,7 @@ pub struct LazyMap<K, V, SON = super::Simple> {
 pub type NestedMap<K, V> = LazyMap<K, V, super::Nested>;
 
 /// Possible sub-keys of a [`LazyMap`]
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum SubKey<K> {
     /// Data sub-key, further sub-keyed by its literal map key
     Data(K),
@@ -81,7 +81,7 @@ pub enum NestedAction<K, A> {
 }
 
 /// Possible sub-keys of a nested [`LazyMap`]
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum NestedSubKey<K, S> {
     /// Data sub-key
     Data {
@@ -141,27 +141,37 @@ where
             Some(Some(suffix)) => suffix,
         };
 
+        // A helper to validate the 2nd key segment
+        let validate_sub_key = |raw_sub_key| {
+            if let Ok(key_in_kv) = storage::KeySeg::parse(raw_sub_key) {
+                let nested = self.at(&key_in_kv).is_valid_sub_key(key)?;
+                match nested {
+                    Some(nested_sub_key) => Ok(Some(NestedSubKey::Data {
+                        key: key_in_kv,
+                        nested_sub_key,
+                    })),
+                    None => {
+                        Err(ValidationError::InvalidNestedSubKey(key.clone()))
+                            .into_storage_result()
+                    }
+                }
+            } else {
+                Err(ValidationError::InvalidSubKey(key.clone()))
+                    .into_storage_result()
+            }
+        };
+
         // Match the suffix against expected sub-keys
         match &suffix.segments[..2] {
             [DbKeySeg::StringSeg(sub_a), DbKeySeg::StringSeg(sub_b)]
                 if sub_a == DATA_SUBKEY =>
             {
-                if let Ok(key_in_kv) = storage::KeySeg::parse(sub_b.clone()) {
-                    let nested = self.at(&key_in_kv).is_valid_sub_key(key)?;
-                    match nested {
-                        Some(nested_sub_key) => Ok(Some(NestedSubKey::Data {
-                            key: key_in_kv,
-                            nested_sub_key,
-                        })),
-                        None => Err(ValidationError::InvalidNestedSubKey(
-                            key.clone(),
-                        ))
-                        .into_storage_result(),
-                    }
-                } else {
-                    Err(ValidationError::InvalidSubKey(key.clone()))
-                        .into_storage_result()
-                }
+                validate_sub_key(sub_b.clone())
+            }
+            [DbKeySeg::StringSeg(sub_a), DbKeySeg::AddressSeg(sub_b)]
+                if sub_a == DATA_SUBKEY =>
+            {
+                validate_sub_key(sub_b.raw())
             }
             _ => Err(ValidationError::InvalidSubKey(key.clone()))
                 .into_storage_result(),
@@ -266,17 +276,27 @@ where
             Some(Some(suffix)) => suffix,
         };
 
+        // A helper to validate the 2nd key segment
+        let validate_sub_key = |raw_sub_key| {
+            if let Ok(key_in_kv) = storage::KeySeg::parse(raw_sub_key) {
+                Ok(Some(SubKey::Data(key_in_kv)))
+            } else {
+                Err(ValidationError::InvalidSubKey(key.clone()))
+                    .into_storage_result()
+            }
+        };
+
         // Match the suffix against expected sub-keys
         match &suffix.segments[..] {
             [DbKeySeg::StringSeg(sub_a), DbKeySeg::StringSeg(sub_b)]
                 if sub_a == DATA_SUBKEY =>
             {
-                if let Ok(key_in_kv) = storage::KeySeg::parse(sub_b.clone()) {
-                    Ok(Some(SubKey::Data(key_in_kv)))
-                } else {
-                    Err(ValidationError::InvalidSubKey(key.clone()))
-                        .into_storage_result()
-                }
+                validate_sub_key(sub_b.clone())
+            }
+            [DbKeySeg::StringSeg(sub_a), DbKeySeg::AddressSeg(sub_b)]
+                if sub_a == DATA_SUBKEY =>
+            {
+                validate_sub_key(sub_b.raw())
             }
             _ => Err(ValidationError::InvalidSubKey(key.clone()))
                 .into_storage_result(),
@@ -523,6 +543,7 @@ where
 mod test {
     use super::*;
     use crate::ledger::storage::testing::TestWlStorage;
+    use crate::types::address::{self, Address};
 
     #[test]
     fn test_lazy_map_basics() -> storage_api::Result<()> {
@@ -533,7 +554,7 @@ mod test {
 
         // The map should be empty at first
         assert!(lazy_map.is_empty(&storage)?);
-        assert!(lazy_map.len(&storage)? == 0);
+        assert_eq!(lazy_map.len(&storage)?, 0);
         assert!(!lazy_map.contains(&storage, &0)?);
         assert!(!lazy_map.contains(&storage, &1)?);
         assert!(lazy_map.iter(&storage)?.next().is_none());
@@ -552,7 +573,7 @@ mod test {
         assert!(!lazy_map.contains(&storage, &0)?);
         assert!(lazy_map.contains(&storage, &key)?);
         assert!(!lazy_map.is_empty(&storage)?);
-        assert!(lazy_map.len(&storage)? == 2);
+        assert_eq!(lazy_map.len(&storage)?, 2);
         let mut map_it = lazy_map.iter(&storage)?;
         assert_eq!(map_it.next().unwrap()?, (key, val.clone()));
         assert_eq!(map_it.next().unwrap()?, (key2, val2.clone()));
@@ -566,7 +587,7 @@ mod test {
         let removed = lazy_map.remove(&mut storage, &key)?.unwrap();
         assert_eq!(removed, val);
         assert!(!lazy_map.is_empty(&storage)?);
-        assert!(lazy_map.len(&storage)? == 1);
+        assert_eq!(lazy_map.len(&storage)?, 1);
         assert!(!lazy_map.contains(&storage, &0)?);
         assert!(!lazy_map.contains(&storage, &1)?);
         assert!(!lazy_map.contains(&storage, &123)?);
@@ -579,7 +600,120 @@ mod test {
         let removed = lazy_map.remove(&mut storage, &key2)?.unwrap();
         assert_eq!(removed, val2);
         assert!(lazy_map.is_empty(&storage)?);
-        assert!(lazy_map.len(&storage)? == 0);
+        assert_eq!(lazy_map.len(&storage)?, 0);
+
+        let storage_key = lazy_map.get_data_key(&key);
+        assert_eq!(
+            lazy_map.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(key))
+        );
+
+        let storage_key2 = lazy_map.get_data_key(&key2);
+        assert_eq!(
+            lazy_map.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(key2))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_lazy_map_with_addr_key() -> storage_api::Result<()> {
+        let mut storage = TestWlStorage::default();
+
+        let key = storage::Key::parse("test").unwrap();
+        let lazy_map = LazyMap::<Address, String>::open(key);
+
+        // Insert a new value and check that it's added
+        let (key, val) = (
+            address::testing::established_address_1(),
+            "Test".to_string(),
+        );
+        lazy_map.insert(&mut storage, key.clone(), val.clone())?;
+
+        assert_eq!(lazy_map.len(&storage)?, 1);
+        let mut map_it = lazy_map.iter(&storage)?;
+        assert_eq!(map_it.next().unwrap()?, (key.clone(), val.clone()));
+        drop(map_it);
+
+        let (key2, val2) = (
+            address::testing::established_address_2(),
+            "Test2".to_string(),
+        );
+        lazy_map.insert(&mut storage, key2.clone(), val2.clone())?;
+
+        assert_eq!(lazy_map.len(&storage)?, 2);
+        let mut map_it = lazy_map.iter(&storage)?;
+        assert!(key < key2, "sanity check - this influences the iter order");
+        assert_eq!(map_it.next().unwrap()?, (key.clone(), val));
+        assert_eq!(map_it.next().unwrap()?, (key2.clone(), val2));
+        assert!(map_it.next().is_none());
+        drop(map_it);
+
+        let storage_key = lazy_map.get_data_key(&key);
+        assert_eq!(
+            lazy_map.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(key))
+        );
+
+        let storage_key2 = lazy_map.get_data_key(&key2);
+        assert_eq!(
+            lazy_map.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(key2))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_nested_lazy_map_with_addr_key() -> storage_api::Result<()> {
+        let mut storage = TestWlStorage::default();
+
+        let key = storage::Key::parse("test").unwrap();
+        let lazy_map = NestedMap::<Address, LazyMap<u64, String>>::open(key);
+
+        // Insert a new value and check that it's added
+        let (key, sub_key, val) = (
+            address::testing::established_address_1(),
+            1_u64,
+            "Test".to_string(),
+        );
+        lazy_map
+            .at(&key)
+            .insert(&mut storage, sub_key, val.clone())?;
+
+        assert_eq!(lazy_map.at(&key).len(&storage)?, 1);
+        let mut map_it = lazy_map.iter(&storage)?;
+        let expected_key = NestedSubKey::Data {
+            key: key.clone(),
+            nested_sub_key: SubKey::Data(sub_key),
+        };
+        assert_eq!(
+            map_it.next().unwrap()?,
+            (expected_key.clone(), val.clone())
+        );
+        drop(map_it);
+
+        let (key2, sub_key2, val2) = (
+            address::testing::established_address_2(),
+            2_u64,
+            "Test2".to_string(),
+        );
+        lazy_map
+            .at(&key2)
+            .insert(&mut storage, sub_key2, val2.clone())?;
+
+        assert_eq!(lazy_map.at(&key2).len(&storage)?, 1);
+        let mut map_it = lazy_map.iter(&storage)?;
+        assert!(key < key2, "sanity check - this influences the iter order");
+        let expected_key2 = NestedSubKey::Data {
+            key: key2,
+            nested_sub_key: SubKey::Data(sub_key2),
+        };
+        assert_eq!(map_it.next().unwrap()?, (expected_key, val));
+        assert_eq!(map_it.next().unwrap()?, (expected_key2, val2));
+        assert!(map_it.next().is_none());
+        drop(map_it);
 
         Ok(())
     }

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

@@ -28,7 +28,7 @@ pub struct LazySet<K> {
 }
 
 /// Possible sub-keys of a [`LazySet`]
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum SubKey<K> {
     /// Literal set key
     Data(K),
@@ -88,16 +88,20 @@ where
             Some(Some(suffix)) => suffix,
         };
 
+        // A helper to validate the 2nd key segment
+        let validate_sub_key = |raw_sub_key| {
+            if let Ok(key) = storage::KeySeg::parse(raw_sub_key) {
+                Ok(Some(SubKey::Data(key)))
+            } else {
+                Err(ValidationError::InvalidSubKey(key.clone()))
+                    .into_storage_result()
+            }
+        };
+
         // Match the suffix against expected sub-keys
         match &suffix.segments[..] {
-            [DbKeySeg::StringSeg(sub)] => {
-                if let Ok(key) = storage::KeySeg::parse(sub.clone()) {
-                    Ok(Some(SubKey::Data(key)))
-                } else {
-                    Err(ValidationError::InvalidSubKey(key.clone()))
-                        .into_storage_result()
-                }
-            }
+            [DbKeySeg::StringSeg(sub)] => validate_sub_key(sub.clone()),
+            [DbKeySeg::AddressSeg(sub)] => validate_sub_key(sub.raw()),
             _ => Err(ValidationError::InvalidSubKey(key.clone()))
                 .into_storage_result(),
         }
@@ -265,6 +269,7 @@ where
 mod test {
     use super::*;
     use crate::ledger::storage::testing::TestWlStorage;
+    use crate::types::address::{self, Address};
 
     #[test]
     fn test_lazy_set_basics() -> storage_api::Result<()> {
@@ -331,6 +336,60 @@ mod test {
         assert!(lazy_set.try_insert(&mut storage, key).is_ok());
         assert!(lazy_set.try_insert(&mut storage, key).is_err());
 
+        let storage_key = lazy_set.get_key(&key);
+        assert_eq!(
+            lazy_set.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(key))
+        );
+
+        let storage_key2 = lazy_set.get_key(&key2);
+        assert_eq!(
+            lazy_set.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(key2))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_lazy_set_with_addr_key() -> storage_api::Result<()> {
+        let mut storage = TestWlStorage::default();
+
+        let key = storage::Key::parse("test").unwrap();
+        let lazy_set = LazySet::<Address>::open(key);
+
+        // Insert a new value and check that it's added
+        let key = address::testing::established_address_1();
+        lazy_set.insert(&mut storage, key.clone())?;
+
+        assert_eq!(lazy_set.len(&storage)?, 1);
+        let mut map_it = lazy_set.iter(&storage)?;
+        assert_eq!(map_it.next().unwrap()?, key);
+        drop(map_it);
+
+        let key2 = address::testing::established_address_2();
+        lazy_set.insert(&mut storage, key2.clone())?;
+
+        assert_eq!(lazy_set.len(&storage)?, 2);
+        let mut iter = lazy_set.iter(&storage)?;
+        assert!(key < key2, "sanity check - this influences the iter order");
+        assert_eq!(iter.next().unwrap()?, key);
+        assert_eq!(iter.next().unwrap()?, key2);
+        assert!(iter.next().is_none());
+        drop(iter);
+
+        let storage_key = lazy_set.get_key(&key);
+        assert_eq!(
+            lazy_set.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(key))
+        );
+
+        let storage_key2 = lazy_set.get_key(&key2);
+        assert_eq!(
+            lazy_set.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(key2))
+        );
+
         Ok(())
     }
 }