Skip to content

Commit

Permalink
implemented push and pop
Browse files Browse the repository at this point in the history
  • Loading branch information
dean-starkware committed Jan 28, 2025
1 parent 00dcc3d commit 91d7a4e
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 1 deletion.
89 changes: 88 additions & 1 deletion corelib/src/starknet/storage/vec.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -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};
///
/// #[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};
///
/// #[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>>`.
Expand Down Expand Up @@ -350,8 +399,34 @@ 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,
) {
self.append().write(value);
}

fn pop<+Drop<Self::ElementType>, +starknet::Store<Self::ElementType>>(
self: StoragePath<Mutable<Vec<T>>>,
) -> Option<Self::ElementType> {
let len_ptr = self.as_ptr();
let vec_len: u64 = len_ptr.read();
if vec_len == 0 {
return None;
}
let entry: StoragePath<Mutable<T>> = 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<
Expand All @@ -377,6 +452,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<
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

0 comments on commit 91d7a4e

Please sign in to comment.