Fix #[state(codec_builder = ...)] test (Sovereign-Labs#723)

Signed-off-by: Filippo Costa <[email protected]>

module-system/module-implementations/examples/sov-vec-setter/tests/tests.rs

@@ -1,7 +1,6 @@
 use sov_modules_api::default_context::DefaultContext;
 use sov_modules_api::{Address, Context, Module};
-use sov_state::codec::StateValueCodec;
-use sov_state::{ProverStorage, StateVec, Storage, WorkingSet};
+use sov_state::{ProverStorage, WorkingSet};
 use sov_vec_setter::{CallMessage, VecSetter, VecSetterConfig};
 
 // rustfmt doesn't like long lines, but it's easier to read in this case.
@@ -47,23 +46,10 @@ fn test_vec_setter_calls() {
         let call_result = vec_setter.call(call, &context, &mut working_set);
 
         if call_result.is_ok() {
-            let vec_contents = state_vec_get_all(&vec_setter.vector, &mut working_set);
+            let vec_contents = vec_setter.vector.iter(&mut working_set).collect::<Vec<_>>();
             assert_eq!(Some(vec_contents), expected_contents);
         } else {
             assert_eq!(expected_contents, None);
         }
     }
 }
-
-fn state_vec_get_all<T, VC, C>(sv: &StateVec<T, VC>, ws: &mut WorkingSet<C>) -> Vec<T>
-where
-    VC: StateValueCodec<T> + StateValueCodec<usize>,
-    C: Storage,
-{
-    let mut result = Vec::new();
-    let len = sv.len(ws);
-    for i in 0..len {
-        result.push(sv.get(i, ws).unwrap());
-    }
-    result
-}

module-system/sov-modules-macros/tests/all_tests.rs

@@ -6,6 +6,7 @@ fn module_info_tests() {
     t.pass("tests/module_info/use_address_trait.rs");
     t.pass("tests/module_info/not_supported_attribute.rs");
     t.pass("tests/module_info/custom_codec_builder.rs");
+    t.pass("tests/custom_codec_must_be_used.rs");
     t.compile_fail("tests/module_info/derive_on_enum_not_supported.rs");
     t.compile_fail("tests/module_info/field_missing_attribute.rs");
     t.compile_fail("tests/module_info/missing_address.rs");

module-system/sov-modules-macros/tests/custom_codec_must_be_used.rs

@@ -1,42 +1,47 @@
 use std::panic::catch_unwind;
 
-use sov_modules_api::default_context::ZkDefaultContext;
+use sov_modules_api::default_context::DefaultContext;
 use sov_modules_api::{Context, ModuleInfo};
-use sov_state::{DefaultStorageSpec, ProverStorage, StateValue, StateValueCodec, WorkingSet};
+use sov_state::codec::StateValueCodec;
+use sov_state::{DefaultStorageSpec, ProverStorage, StateValue, WorkingSet};
 
 #[derive(ModuleInfo)]
 struct TestModule<C>
 where
     C: Context,
 {
     #[address]
-    pub address: C::Address,
+    address: C::Address,
 
-    #[state(codec_builder = "Default::default")]
-    pub state_value: StateValue<u32, CustomCodec>,
+    #[state(codec_builder = "crate::CustomCodec::new")]
+    state_value: StateValue<u32, CustomCodec>,
 }
 
 #[derive(Default)]
-pub struct CustomCodec;
+struct CustomCodec;
+
+impl CustomCodec {
+    fn new() -> Self {
+        Self
+    }
+}
 
 impl<V> StateValueCodec<V> for CustomCodec {
-    type ValueError = String;
+    type Error = String;
 
     fn encode_value(&self, _value: &V) -> Vec<u8> {
         unimplemented!()
     }
 
-    fn try_decode_value(&self, _bytes: &[u8]) -> Result<V, Self::ValueError> {
+    fn try_decode_value(&self, _bytes: &[u8]) -> Result<V, Self::Error> {
         unimplemented!()
     }
 }
 
-#[test]
 fn main() {
-    let module: TestModule<ZkDefaultContext> = Default::default();
-
     let tempdir = tempfile::tempdir().unwrap();
     let storage: ProverStorage<DefaultStorageSpec> = ProverStorage::with_path(&tempdir).unwrap();
+    let module: TestModule<DefaultContext> = TestModule::default();
 
     catch_unwind(|| {
         let mut working_set = WorkingSet::new(storage);

module-system/sov-state/src/vec.rs

@@ -125,17 +125,19 @@ where
     /// Pops a value from the end of the [`StateVec`] and returns it.
     pub fn pop<S: Storage>(&self, working_set: &mut WorkingSet<S>) -> Option<V> {
         let len = self.len(working_set);
-        let new_len = len.checked_sub(1)?;
+        let last_i = len.checked_sub(1)?;
+        let elem = self.elems.remove(&last_i, working_set)?;
 
-        let elem = self.get(new_len, working_set)?;
+        let new_len = last_i;
         self.set_len(new_len, working_set);
+
         Some(elem)
     }
 
     pub fn clear<S: Storage>(&self, working_set: &mut WorkingSet<S>) {
         let len = self.len_value.remove(working_set).unwrap_or_default();
 
-        for i in 0..len + 1 {
+        for i in 0..len {
             self.elems.delete(&i, working_set);
         }
     }
@@ -177,49 +179,47 @@ where
 /// An [`Iterator`] over a [`StateVec`]
 ///
 /// See [`StateVec::iter`] for more details.
-pub struct StateVecIter<'a, 'ws, V, C, S>
+pub struct StateVecIter<'a, 'ws, V, VC, S>
 where
-    C: StateValueCodec<V>,
+    VC: StateValueCodec<V>,
     S: Storage,
 {
-    state_vec: &'a StateVec<V, C>,
+    state_vec: &'a StateVec<V, VC>,
     ws: &'ws mut WorkingSet<S>,
     len: usize,
     next_i: usize,
 }
 
-impl<'a, 'ws, V, C, S> Iterator for StateVecIter<'a, 'ws, V, C, S>
+impl<'a, 'ws, V, VC, S> Iterator for StateVecIter<'a, 'ws, V, VC, S>
 where
-    C: StateValueCodec<V>,
+    VC: StateValueCodec<V>,
     S: Storage,
 {
     type Item = V;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.next_i == self.state_vec.len(self.ws) {
-            return None;
+        let elem = self.state_vec.get(self.next_i, self.ws);
+        if elem.is_some() {
+            self.next_i += 1;
         }
 
-        let elem = self.state_vec.get(self.next_i, self.ws);
-        debug_assert!(elem.is_some());
-        self.next_i += 1;
         elem
     }
 }
 
-impl<'a, 'ws, V, C, S> ExactSizeIterator for StateVecIter<'a, 'ws, V, C, S>
+impl<'a, 'ws, V, VC, S> ExactSizeIterator for StateVecIter<'a, 'ws, V, VC, S>
 where
-    C: StateValueCodec<V>,
+    VC: StateValueCodec<V>,
     S: Storage,
 {
     fn len(&self) -> usize {
         self.len - self.next_i
     }
 }
 
-impl<'a, 'ws, V, C, S> FusedIterator for StateVecIter<'a, 'ws, V, C, S>
+impl<'a, 'ws, V, VC, S> FusedIterator for StateVecIter<'a, 'ws, V, VC, S>
 where
-    C: StateValueCodec<V>,
+    VC: StateValueCodec<V>,
     S: Storage,
 {
 }