Add basic storage array. (#5974)

corelib/src/starknet/storage.cairo

@@ -5,6 +5,9 @@ use starknet::storage_access::StorageBaseAddress;
 use starknet::SyscallResult;
 use starknet::storage_access::storage_base_address_from_felt252;
 
+mod vec;
+pub use vec::{Vec, VecTrait, MutableVecTrait};
+
 
 /// A pointer to an address in storage, can be used to read and write values, if the generic type
 /// supports it (e.g. basic types like `felt252`).

corelib/src/starknet/storage/vec.cairo

@@ -0,0 +1,120 @@
+use super::{
+    StorageAsPath, StorageAsPointer, StoragePath, StoragePointer0Offset, Mutable, StoragePathTrait,
+    StoragePathUpdateTrait, StoragePointerReadAccess, StoragePointerWriteAccess
+};
+use core::Option;
+
+/// A type to represent a vec in storage. The length of the storage is stored in the storage
+/// base, while the elements are stored in hash(storage_base, index).
+#[phantom]
+pub struct Vec<T> {}
+
+impl VecDrop<T> of Drop<Vec<T>> {}
+impl VecCopy<T> of Copy<Vec<T>> {}
+
+/// Implement as_ptr for Vec.
+impl VecAsPointer<T> of StorageAsPointer<StoragePath<Vec<T>>> {
+    type Value = u64;
+    fn as_ptr(self: @StoragePath<Vec<T>>) -> StoragePointer0Offset<u64> {
+        StoragePointer0Offset { address: (*self).finalize() }
+    }
+}
+
+/// Implement as_ptr for Mutable<Vec>.
+impl MutableVecAsPointer<T> of StorageAsPointer<StoragePath<Mutable<Vec<T>>>> {
+    type Value = Mutable<u64>;
+    fn as_ptr(self: @StoragePath<Mutable<Vec<T>>>) -> StoragePointer0Offset<Mutable<u64>> {
+        StoragePointer0Offset { address: (*self).finalize() }
+    }
+}
+
+
+/// Trait for the interface of a storage vec.
+pub trait VecTrait<T> {
+    type ElementType;
+    fn get(self: T, index: u64) -> Option<StoragePath<Self::ElementType>>;
+    fn len(self: T) -> u64;
+}
+
+/// Implement `VecTrait` for `StoragePath<Vec<T>>`.
+impl VecImpl<T> of VecTrait<StoragePath<Vec<T>>> {
+    type ElementType = T;
+    fn get(self: StoragePath<Vec<T>>, index: u64) -> Option<StoragePath<T>> {
+        let vec_len = self.len();
+        if index < vec_len {
+            Option::Some(self.update(index))
+        } else {
+            Option::None
+        }
+    }
+    fn len(self: StoragePath<Vec<T>>) -> u64 {
+        self.as_ptr().read()
+    }
+}
+
+/// Implement `VecTrait` for any type that implements StorageAsPath into a storage path
+/// that implements VecTrait.
+impl PathableVecImpl<
+    T,
+    +Drop<T>,
+    impl PathImpl: StorageAsPath<T>,
+    impl VecTraitImpl: VecTrait<StoragePath<PathImpl::Value>>
+> of VecTrait<T> {
+    type ElementType = VecTraitImpl::ElementType;
+    fn get(self: T, index: u64) -> Option<StoragePath<VecTraitImpl::ElementType>> {
+        self.as_path().get(index)
+    }
+    fn len(self: T) -> u64 {
+        self.as_path().len()
+    }
+}
+
+/// Trait for the interface of a mutable storage vec.
+pub trait MutableVecTrait<T> {
+    type ElementType;
+    fn get(self: T, index: u64) -> Option<StoragePath<Mutable<Self::ElementType>>>;
+    fn len(self: T) -> u64;
+    fn append(self: T) -> StoragePath<Mutable<Self::ElementType>>;
+}
+
+/// Implement `MutableVecTrait` for `StoragePath<Mutable<Vec<T>>`.
+impl MutableVecImpl<T, +Drop<T>> of MutableVecTrait<StoragePath<Mutable<Vec<T>>>> {
+    type ElementType = T;
+    fn get(self: StoragePath<Mutable<Vec<T>>>, index: u64) -> Option<StoragePath<Mutable<T>>> {
+        let vec_len = self.len();
+        if index < vec_len {
+            Option::Some(self.update(index))
+        } else {
+            Option::None
+        }
+    }
+    fn len(self: StoragePath<Mutable<Vec<T>>>) -> u64 {
+        self.as_ptr().read()
+    }
+    fn append(self: StoragePath<Mutable<Vec<T>>>) -> StoragePath<Mutable<T>> {
+        let vec_len = self.len();
+        self.as_ptr().write(vec_len + 1);
+        self.update(vec_len)
+    }
+}
+
+/// Implement `MutableVecTrait` for any type that implements StorageAsPath into a storage
+/// path that implements MutableVecTrait.
+impl PathableMutableVecImpl<
+    T,
+    +Drop<T>,
+    impl PathImpl: StorageAsPath<T>,
+    impl VecTraitImpl: MutableVecTrait<StoragePath<PathImpl::Value>>
+> of MutableVecTrait<T> {
+    type ElementType = VecTraitImpl::ElementType;
+    fn get(self: T, index: u64) -> Option<StoragePath<Mutable<VecTraitImpl::ElementType>>> {
+        self.as_path().get(index)
+    }
+    fn len(self: T) -> u64 {
+        self.as_path().len()
+    }
+    fn append(self: T) -> StoragePath<Mutable<VecTraitImpl::ElementType>> {
+        self.as_path().append()
+    }
+}
+

crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo

@@ -0,0 +1,58 @@
+use core::starknet::storage::StoragePointerReadAccess;
+use core::starknet::storage::MutableVecTrait;
+use core::starknet::storage::StoragePathEntry;
+
+
+#[starknet::contract]
+mod contract_with_map {
+    use starknet::storage::Map;
+    #[storage]
+    struct Storage {
+        simple: Map<u64, u32>,
+        nested: Map<u64, Map<u64, u32>>,
+    }
+}
+
+#[starknet::contract]
+mod contract_with_vec {
+    use starknet::storage::Vec;
+    #[storage]
+    struct Storage {
+        simple: Vec<u32>,
+        nested: Vec<Vec<u32>>,
+    }
+}
+
+#[test]
+fn test_simple_member_write_to_map() {
+    let mut map_contract_state = contract_with_map::contract_state_for_testing();
+    let mut vec_contract_state = contract_with_vec::contract_state_for_testing();
+    let vec_entry = vec_contract_state.simple.append();
+    map_contract_state.simple.entry(0).write(1);
+    assert_eq!(vec_entry.read(), 1);
+}
+
+#[test]
+fn test_simple_member_write_to_vec() {
+    let mut map_contract_state = contract_with_map::contract_state_for_testing();
+    let mut vec_contract_state = contract_with_vec::contract_state_for_testing();
+    vec_contract_state.simple.append().write(1);
+    assert_eq!(map_contract_state.simple.entry(0).read(), 1);
+}
+
+#[test]
+fn test_nested_member_write_to_map() {
+    let mut map_contract_state = contract_with_map::contract_state_for_testing();
+    let mut vec_contract_state = contract_with_vec::contract_state_for_testing();
+    let vec_entry = vec_contract_state.nested.append().append();
+    map_contract_state.nested.entry(0).entry(0).write(1);
+    assert_eq!(vec_entry.read(), 1);
+}
+
+#[test]
+fn test_nested_member_write_to_vec() {
+    let mut map_contract_state = contract_with_map::contract_state_for_testing();
+    let mut vec_contract_state = contract_with_vec::contract_state_for_testing();
+    vec_contract_state.nested.append().append().write(1);
+    assert_eq!(map_contract_state.nested.entry(0).entry(0).read(), 1);
+}

crates/cairo-lang-starknet/cairo_level_tests/lib.cairo

@@ -25,3 +25,5 @@ mod storage_access;
 #[cfg(test)]
 mod contract_address_test;
 mod utils;
+#[cfg(test)]
+mod collections_test;

crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo

@@ -4,6 +4,7 @@ use super::utils::{deserialized, serialized};
 use core::integer::BoundedInt;
 use core::num::traits::Zero;
 use core::byte_array::ByteArrayTrait;
+use starknet::storage::Vec;
 
 impl StorageAddressPartialEq of PartialEq<StorageAddress> {
     fn eq(lhs: @StorageAddress, rhs: @StorageAddress) -> bool {
@@ -78,15 +79,25 @@ struct NonZeros {
     value_felt252: NonZero<felt252>,
 }
 
+#[starknet::storage_node]
+struct Vecs {
+    vec: Vec<u32>,
+    vec_of_vecs: Vec<Vec<u32>>,
+}
+
 #[starknet::contract]
 mod test_contract {
-    use super::{AbcEtc, ByteArrays, NonZeros};
+    use core::option::OptionTrait;
+    use core::starknet::storage::StoragePointerWriteAccess;
+    use super::{AbcEtc, ByteArrays, NonZeros, Vecs,};
+    use starknet::storage::{VecTrait, MutableVecTrait, StorageAsPath,};
 
     #[storage]
     struct Storage {
         data: AbcEtc,
         byte_arrays: ByteArrays,
         non_zeros: NonZeros,
+        vecs: Vecs,
     }
 
     #[external(v0)]
@@ -128,6 +139,46 @@ mod test_contract {
     pub fn get_non_zeros(self: @ContractState) -> NonZeros {
         self.non_zeros.read()
     }
+
+    #[external(v0)]
+    pub fn append_to_vec(ref self: ContractState, value: u32) {
+        self.vecs.vec.append().write(value);
+    }
+
+    #[external(v0)]
+    pub fn get_vec_length(self: @ContractState) -> u64 {
+        self.vecs.vec.len()
+    }
+
+    #[external(v0)]
+    pub fn get_vec_element(self: @ContractState, index: u64) -> u32 {
+        self.vecs.vec.get(index).unwrap().read()
+    }
+
+    #[external(v0)]
+    pub fn append_an_vec(ref self: ContractState) {
+        self.vecs.vec_of_vecs.append();
+    }
+
+    #[external(v0)]
+    pub fn append_to_nested_vec(ref self: ContractState, index: u64, value: u32) {
+        self.vecs.vec_of_vecs.get(index).unwrap().append().write(value);
+    }
+
+    #[external(v0)]
+    pub fn get_vec_of_vecs_length(self: @ContractState) -> u64 {
+        self.vecs.vec_of_vecs.len()
+    }
+
+    #[external(v0)]
+    pub fn get_nested_vec_length(self: @ContractState, index: u64) -> u64 {
+        self.vecs.vec_of_vecs.get(index).unwrap().len()
+    }
+
+    #[external(v0)]
+    pub fn get_nested_vec_element(self: @ContractState, index: u64, nested_index: u64) -> u32 {
+        self.vecs.vec_of_vecs.get(index).unwrap().get(nested_index).unwrap().read()
+    }
 }
 
 #[test]
@@ -224,3 +275,52 @@ fn test_read_write_non_zero() {
     assert!(test_contract::__external::set_non_zeros(serialized(x.clone())).is_empty());
     assert_eq!(deserialized(test_contract::__external::get_non_zeros(serialized(()))), x);
 }
+
+#[test]
+fn test_storage_array() {
+    assert!(test_contract::__external::append_to_vec(serialized(1_u32)).is_empty());
+    assert!(test_contract::__external::append_to_vec(serialized(2_u32)).is_empty());
+    assert!(test_contract::__external::append_to_vec(serialized(3_u32)).is_empty());
+    assert_eq!(deserialized(test_contract::__external::get_vec_length(serialized(()))), 3);
+    assert_eq!(deserialized(test_contract::__external::get_vec_element(serialized(0_u64))), 1);
+    assert_eq!(deserialized(test_contract::__external::get_vec_element(serialized(1_u64))), 2);
+    assert_eq!(deserialized(test_contract::__external::get_vec_element(serialized(2_u64))), 3);
+}
+
+#[test]
+fn test_storage_vec_of_vecs() {
+    assert!(test_contract::__external::append_an_vec(serialized(())).is_empty());
+    assert!(test_contract::__external::append_to_nested_vec(serialized((0_u64, 1_u32))).is_empty());
+    assert!(test_contract::__external::append_to_nested_vec(serialized((0_u64, 2_u32))).is_empty());
+    assert!(test_contract::__external::append_to_nested_vec(serialized((0_u64, 3_u32))).is_empty());
+    assert!(test_contract::__external::append_an_vec(serialized(())).is_empty());
+    assert!(test_contract::__external::append_to_nested_vec(serialized((1_u64, 4_u32))).is_empty());
+    assert!(test_contract::__external::append_to_nested_vec(serialized((1_u64, 5_u32))).is_empty());
+    assert_eq!(deserialized(test_contract::__external::get_vec_of_vecs_length(serialized(()))), 2);
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_length(serialized(0_u64))), 3
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_element(serialized((0_u64, 0_u64)))),
+        1
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_element(serialized((0_u64, 1_u64)))),
+        2
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_element(serialized((0_u64, 2_u64)))),
+        3
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_length(serialized(1_u64))), 2
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_element(serialized((1_u64, 0_u64)))),
+        4
+    );
+    assert_eq!(
+        deserialized(test_contract::__external::get_nested_vec_element(serialized((1_u64, 1_u64)))),
+        5
+    );
+}