diff --git a/corelib/src/starknet/storage/vec.cairo b/corelib/src/starknet/storage/vec.cairo index a8fa4ca085d..0fdbf204e5b 100644 --- a/corelib/src/starknet/storage/vec.cairo +++ b/corelib/src/starknet/storage/vec.cairo @@ -321,6 +321,55 @@ pub trait MutableVecTrait { /// } /// ``` fn append(self: T) -> StoragePath>; + + /// Pushes a new value onto the vector. + /// + /// This operation: + /// 1. Increments the vector's length. + /// 2. Writes the provided value to the new storage location at the end of the vector. + /// + /// # Examples + /// + /// ``` + /// use core::starknet::storage::{Vec, MutableVecTrait}; + /// + /// #[storage] + /// struct Storage { + /// numbers: Vec, + /// } + /// + /// fn push_number(ref self: ContractState, number: u256) { + /// self.numbers.push(number); + /// } + /// ``` + fn push<+Drop, +starknet::Store>( + self: T, value: Self::ElementType, + ); + + /// Pops the last value off the vector. + /// + /// This operation: + /// 1. Retrieves the value stored at the last position in the vector. + /// 2. Decrements the vector's length. + /// 3. Returns the retrieved value or `None` if the vector is empty. + /// + /// # Examples + /// + /// ``` + /// use core::starknet::storage::{Vec, MutableVecTrait}; + /// + /// #[storage] + /// struct Storage { + /// numbers: Vec, + /// } + /// + /// fn pop_number(ref self: ContractState) -> Option { + /// self.numbers.pop() + /// } + /// ``` + fn pop<+Drop, +starknet::Store>( + self: T, + ) -> Option; } /// Implement `MutableVecTrait` for `StoragePath>`. @@ -350,8 +399,34 @@ impl MutableVecImpl of MutableVecTrait>>> { self.as_ptr().write(vec_len + 1); self.update(vec_len) } -} + fn push<+Drop, +starknet::Store>( + self: StoragePath>>, value: Self::ElementType, + ) { + self.append().write(value); + } + + fn pop<+Drop, +starknet::Store>( + self: StoragePath>>, + ) -> Option { + let len_ptr = self.as_ptr(); + let vec_len: u64 = len_ptr.read(); + if vec_len == 0 { + return None; + } + let entry: StoragePath> = self.update(vec_len - 1); + let last_element = entry.read(); + // Remove the element's data from the storage. + let entry_ptr = entry.as_ptr(); + starknet::SyscallResultTrait::unwrap_syscall( + starknet::Store::< + Self::ElementType, + >::scrub(0, entry_ptr.__storage_pointer_address__, 0), + ); + len_ptr.write(vec_len - 1); + Some(last_element) + } +} /// Implement `MutableVecTrait` for any type that implements StorageAsPath into a storage /// path that implements MutableVecTrait. impl PathableMutableVecImpl< @@ -377,6 +452,18 @@ impl PathableMutableVecImpl< fn append(self: T) -> StoragePath> { self.as_path().append() } + + fn push<+Drop, +starknet::Store>( + self: T, value: Self::ElementType, + ) { + self.as_path().push(value) + } + + fn pop<+Drop, +starknet::Store>( + self: T, + ) -> Option { + self.as_path().pop() + } } pub impl VecIndexView< diff --git a/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo b/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo index e5f492a7b49..eff1667e5e5 100644 --- a/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo +++ b/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo @@ -101,3 +101,28 @@ fn test_nested_member_write_to_vec() { vec_contract_state.nested.append().append().write(1); assert_eq!(map_contract_state.nested.entry(0).entry(0).read(), 1); } + +#[test] +fn test_simple_member_push_to_vec() { + let mut state = contract_with_vec::contract_state_for_testing(); + state.simple.push(10); + state.simple.push(20); + state.simple.push(30); + assert_eq!(state.simple.len(), 3); + assert_eq!(state.simple.at(0).read(), 10); + assert_eq!(state.simple.at(1).read(), 20); + assert_eq!(state.simple.at(2).read(), 30); +} + +#[test] +fn test_simple_member_pop_from_vec() { + let mut state = contract_with_vec::contract_state_for_testing(); + state.simple.append().write(10); + state.simple.append().write(20); + state.simple.append().write(30); + assert_eq!(state.simple.pop(), Some(30)); + assert_eq!(state.simple.pop(), Some(20)); + assert_eq!(state.simple.pop(), Some(10)); + assert_eq!(state.simple.len(), 0); + assert_eq!(state.simple.pop(), None); +} diff --git a/crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo b/crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo index 176aa3bbebc..59cdd27bba4 100644 --- a/crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo +++ b/crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo @@ -1,7 +1,7 @@ #[feature("deprecated-bounded-int-trait")] use core::integer::BoundedInt; use core::num::traits::Zero; -use starknet::storage::Vec; +use starknet::storage::{StoragePointerReadAccess, StoragePointerWriteAccess, Vec}; use starknet::{ClassHash, ContractAddress, EthAddress, StorageAddress}; use super::utils::{deserialized, serialized};