From 915b5c55718977d7b2ead98d1347ec901365b1fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= <tomas@heliax.dev>
Date: Mon, 13 Mar 2023 07:45:13 +0000
Subject: [PATCH] core/lazy_vec: test and fix for address values

---
 .../storage_api/collections/lazy_vec.rs       | 81 +++++++++++++++++--
 1 file changed, 73 insertions(+), 8 deletions(-)

diff --git a/core/src/ledger/storage_api/collections/lazy_vec.rs b/core/src/ledger/storage_api/collections/lazy_vec.rs
index 47b5c95c754..67b1730c905 100644
--- a/core/src/ledger/storage_api/collections/lazy_vec.rs
+++ b/core/src/ledger/storage_api/collections/lazy_vec.rs
@@ -12,7 +12,7 @@ use super::LazyCollection;
 use crate::ledger::storage_api::validation::{self, Data};
 use crate::ledger::storage_api::{self, ResultExt, StorageRead, StorageWrite};
 use crate::ledger::vp_env::VpEnv;
-use crate::types::storage::{self, DbKeySeg};
+use crate::types::storage::{self, DbKeySeg, KeySeg};
 
 /// Subkey pointing to the length of the LazyVec
 pub const LEN_SUBKEY: &str = "len";
@@ -35,7 +35,7 @@ pub struct LazyVec<T> {
 }
 
 /// Possible sub-keys of a [`LazyVec`]
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum SubKey {
     /// Length sub-key
     Len,
@@ -144,6 +144,16 @@ where
             Some(Some(suffix)) => suffix,
         };
 
+        // A helper to validate the 2nd key segment
+        let validate_sub_key = |raw_sub_key| {
+            if let Ok(index) = storage::KeySeg::parse(raw_sub_key) {
+                Ok(Some(SubKey::Data(index)))
+            } else {
+                Err(ValidationError::InvalidSubKey(key.clone()))
+                    .into_storage_result()
+            }
+        };
+
         // Match the suffix against expected sub-keys
         match &suffix.segments[..] {
             [DbKeySeg::StringSeg(sub)] if sub == LEN_SUBKEY => {
@@ -152,12 +162,12 @@ where
             [DbKeySeg::StringSeg(sub_a), DbKeySeg::StringSeg(sub_b)]
                 if sub_a == DATA_SUBKEY =>
             {
-                if let Ok(index) = storage::KeySeg::parse(sub_b.clone()) {
-                    Ok(Some(SubKey::Data(index)))
-                } 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(),
@@ -477,6 +487,7 @@ where
 mod test {
     use super::*;
     use crate::ledger::storage::testing::TestWlStorage;
+    use crate::types::address::{self, Address};
 
     #[test]
     fn test_lazy_vec_basics() -> storage_api::Result<()> {
@@ -511,6 +522,60 @@ mod test {
         assert!(lazy_vec.get(&storage, 0)?.is_none());
         assert!(lazy_vec.get(&storage, 1)?.is_none());
 
+        let storage_key = lazy_vec.get_data_key(0);
+        assert_eq!(
+            lazy_vec.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(0))
+        );
+
+        let storage_key2 = lazy_vec.get_data_key(1);
+        assert_eq!(
+            lazy_vec.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(1))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_lazy_vec_with_addr() -> storage_api::Result<()> {
+        let mut storage = TestWlStorage::default();
+
+        let key = storage::Key::parse("test").unwrap();
+        let lazy_vec = LazyVec::<Address>::open(key);
+
+        // Push a new value and check that it's added
+        let val = address::testing::established_address_1();
+        lazy_vec.push(&mut storage, val.clone())?;
+        assert!(!lazy_vec.is_empty(&storage)?);
+        assert!(lazy_vec.len(&storage)? == 1);
+        assert_eq!(lazy_vec.iter(&storage)?.next().unwrap()?, val);
+        assert_eq!(lazy_vec.get(&storage, 0)?.unwrap(), val);
+        assert!(lazy_vec.get(&storage, 1)?.is_none());
+
+        let val2 = address::testing::established_address_2();
+        lazy_vec.push(&mut storage, val2.clone())?;
+
+        assert_eq!(lazy_vec.len(&storage)?, 2);
+        let mut iter = lazy_vec.iter(&storage)?;
+        // The iterator order follows the indices
+        assert_eq!(iter.next().unwrap()?, val);
+        assert_eq!(iter.next().unwrap()?, val2);
+        assert!(iter.next().is_none());
+        drop(iter);
+
+        let storage_key = lazy_vec.get_data_key(0);
+        assert_eq!(
+            lazy_vec.is_valid_sub_key(&storage_key).unwrap(),
+            Some(SubKey::Data(0))
+        );
+
+        let storage_key2 = lazy_vec.get_data_key(1);
+        assert_eq!(
+            lazy_vec.is_valid_sub_key(&storage_key2).unwrap(),
+            Some(SubKey::Data(1))
+        );
+
         Ok(())
     }
 }