implemented push and pop

starkware-libs · Jan 24, 2025 · 9d188d4 · 9d188d4
1 parent e90be20
commit 9d188d4
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 0 deletions.
diff --git a/corelib/src/prelude/v2024_07.cairo b/corelib/src/prelude/v2024_07.cairo
@@ -18,5 +18,6 @@ pub use crate::traits::{
 };
 pub use crate::iter::{FromIterator, IntoIterator, Iterator};
 pub use crate::{assert, bool, felt252, starknet, usize};
+use crate::option::Option::{None, Some};
 
 pub use crate::zeroable::NonZero;
diff --git a/corelib/src/starknet/storage/vec.cairo b/corelib/src/starknet/storage/vec.cairo
@@ -321,6 +321,55 @@ pub trait MutableVecTrait<T> {
     /// }
     /// ```
     fn append(self: T) -> StoragePath<Mutable<Self::ElementType>>;
+
+    /// 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, StoragePointerWriteAccess};
+    ///
+    /// #[storage]
+    /// struct Storage {
+    ///     numbers: Vec<u256>,
+    /// }
+    ///
+    /// fn push_number(ref self: ContractState, number: u256) {
+    ///     self.numbers.push(number);
+    /// }
+    /// ```
+    fn push<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        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, StoragePointerWriteAccess};
+    ///
+    /// #[storage]
+    /// struct Storage {
+    ///     numbers: Vec<u256>,
+    /// }
+    ///
+    /// fn pop_number(ref self: ContractState) -> Option<u256> {
+    ///     self.numbers.pop()
+    /// }
+    /// ```
+    fn pop<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        self: T,
+    ) -> Option<Self::ElementType>;
 }
 
 /// Implement `MutableVecTrait` for `StoragePath<Mutable<Vec<T>>`.
@@ -350,6 +399,25 @@ impl MutableVecImpl<T> of MutableVecTrait<StoragePath<Mutable<Vec<T>>>> {
         self.as_ptr().write(vec_len + 1);
         self.update(vec_len)
     }
+
+    fn push<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        self: StoragePath<Mutable<Vec<T>>>, value: Self::ElementType,
+    ) {
+        let storage_path = self.append();
+        storage_path.write(value);
+    }
+
+    fn pop<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        self: StoragePath<Mutable<Vec<T>>>,
+    ) -> Option<Self::ElementType> {
+        let vec_len: u64 = self.len();
+        if vec_len == 0 {
+            return Option::None;
+        }
+        let last_element = self.at(vec_len - 1).read();
+        self.as_ptr().write(vec_len - 1);
+        Option::Some(last_element)
+    }
 }
 
 /// Implement `MutableVecTrait` for any type that implements StorageAsPath into a storage
@@ -377,6 +445,18 @@ impl PathableMutableVecImpl<
     fn append(self: T) -> StoragePath<Mutable<VecTraitImpl::ElementType>> {
         self.as_path().append()
     }
+
+    fn push<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        self: T, value: Self::ElementType,
+    ) {
+        self.as_path().push(value)
+    }
+
+    fn pop<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
+        self: T,
+    ) -> Option<Self::ElementType> {
+        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
@@ -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 vec_contract_state = contract_with_vec::contract_state_for_testing();
+    vec_contract_state.simple.push(10);
+    vec_contract_state.simple.push(20);
+    vec_contract_state.simple.push(30);
+    assert_eq!(vec_contract_state.simple.len(), 3);
+    assert_eq!(vec_contract_state.simple.at(0).read(), 10);
+    assert_eq!(vec_contract_state.simple.at(1).read(), 20);
+    assert_eq!(vec_contract_state.simple.at(2).read(), 30);
+}
+
+#[test]
+fn test_simple_member_pop_from_vec() {
+    let mut vec_contract_state = contract_with_vec::contract_state_for_testing();
+    vec_contract_state.simple.append().write(10);
+    vec_contract_state.simple.append().write(20);
+    vec_contract_state.simple.append().write(30);
+    assert_eq!(vec_contract_state.simple.pop(), Some(30));
+    assert_eq!(vec_contract_state.simple.pop(), Some(20));
+    assert_eq!(vec_contract_state.simple.pop(), Some(10));
+    assert_eq!(vec_contract_state.simple.len(), 0);
+    assert_eq!(vec_contract_state.simple.pop(), None);
+}