diff --git a/packages/utils/src/storage.cairo b/packages/utils/src/storage.cairo index da6cb02..f206267 100644 --- a/packages/utils/src/storage.cairo +++ b/packages/utils/src/storage.cairo @@ -1,5 +1,6 @@ pub(crate) mod bit_mask; pub(crate) mod bit_set; +pub mod fast_iterable_map; pub mod iterable_map; #[cfg(test)] diff --git a/packages/utils/src/storage/fast_iterable_map.cairo b/packages/utils/src/storage/fast_iterable_map.cairo new file mode 100644 index 0000000..9d6b2a0 --- /dev/null +++ b/packages/utils/src/storage/fast_iterable_map.cairo @@ -0,0 +1,249 @@ +use core::hash::Hash; +use core::iter::{IntoIterator, Iterator}; +use core::pedersen::HashState; +use starknet::Store; +use starknet::storage::{ + Map, Mutable, StorageAsPath, StorageMapReadAccess, StorageMapWriteAccess, StoragePath, +}; + +#[derive(Copy, Drop, starknet::Store)] +struct Entry { + _value: V, + _key: K, +} + +#[generate_trait] +impl EntryImpl of EntryTrait { + fn key(self: Entry) -> K { + self._key + } + fn value(self: Entry) -> V { + self._value + } +} + +/// FastIterableMap is a map that is optimized for fast iteration. +#[starknet::storage_node] +#[allow(starknet::invalid_storage_member_types)] +pub struct FastIterableMap { + _inner_map: Map>, + _initialized_map: Map, + _head: K, + _tail: K, + _length: u64, +} + +/// Trait for the interface of a iterable map. +pub trait FastIterableMapTrait { + type Key; + fn len(self: T) -> u64; +} + +impl StoragePathFastIterableMapImpl< + K, V, +Drop, +Drop, +Store, +Hash, +> of FastIterableMapTrait>> { + type Key = K; + fn len(self: StoragePath>) -> u64 { + self._length.read() + } +} + +impl StoragePathMutableFastIterableMapImpl< + K, V, +Drop, +Drop, +Store, +Hash, +> of FastIterableMapTrait>>> { + type Key = K; + fn len(self: StoragePath>>) -> u64 { + self._length.read() + } +} + +pub impl FastIterableMapTraitImpl< + T, + +Drop, + impl StorageAsPathImpl: StorageAsPath, + impl StoragePathImpl: FastIterableMapTrait>, + +Drop, +> of FastIterableMapTrait { + type Key = StoragePathImpl::Key; + fn len(self: T) -> u64 { + self.as_path().len() + } +} + +/// Read and write access trait implementations: +impl StoragePathFastIterableMapReadAccessImpl< + K, V, +Drop, +Drop, +Store, +Store, +Hash, +> of StorageMapReadAccess>> { + type Key = K; + type Value = V; + fn read(self: StoragePath>, key: Self::Key) -> Option { + if self._initialized_map.entry(key).read() == true { + Option::Some(self._inner_map.entry(key).read().value()) + } else { + Option::None + } + } +} + +impl StoragePathMutableFastIterableMapReadAccessImpl< + K, V, +Drop, +Store, +Default, +Hash, +> of StorageMapReadAccess>>> { + type Key = K; + type Value = V; + fn read( + self: StoragePath>>, key: Self::Key, + ) -> Option { + if self._initialized_map.entry(key).read() == true { + Option::Some(self._inner_map.entry(key).read().value()) + } else { + Option::None + } + } +} + +pub impl FastIterableMapReadAccessImpl< + T, + +Drop, + impl StorageAsPathImpl: StorageAsPath, + impl StoragePathImpl: StorageMapReadAccess>, + +FastIterableMapTrait>, + +Drop, +> of StorageMapReadAccess { + type Key = StoragePathImpl::Key; + type Value = StoragePathImpl::Value; + fn read(self: T, key: Self::Key) -> Self::Value { + self.as_path().read(key) + } +} + +impl StoragePathFastIterableMapWriteAccessImpl< + K, V, +Drop, +Drop, +Store, +Store, +Hash, +Copy, +> of StorageMapWriteAccess>>> { + type Key = K; + type Value = V; + fn write( + self: StoragePath>>, key: Self::Key, value: Self::Value, + ) { + if self._initialized_map.entry(key).read() == true { + self._inner_map.entry(key).read().value().write(value); + } else { + let current_length = self._length.read(); + if current_length != 0 { + let current_tail_key = self._tail.read(); + self._inner_map.entry(current_tail_key).read().key().write(key); + } else { + self._head.write(key); + } + self._tail.write(key); + self._initialized_map.entry(key).write(true); + self._length.write(current_length + 1); + self._inner_map.entry(key).read().value().write(value); + } + } +} + +pub impl FastIterableMapWriteAccessImpl< + T, + +Drop, + impl StorageAsPathImpl: StorageAsPath, + impl StoragePathImpl: StorageMapWriteAccess>, + +FastIterableMapTrait>, + +Drop, + +Drop, +> of StorageMapWriteAccess { + type Key = StoragePathImpl::Key; + type Value = StoragePathImpl::Value; + fn write(self: T, key: Self::Key, value: Self::Value) { + self.as_path().write(key, value) + } +} + +/// Iterator and IntoItarator implementations: +#[derive(Copy, Drop)] +struct FastMapIterator { + _inner_map: Map>, + _current_key: K, + _current_index: u64, + _length: u64, +} + +#[derive(Copy, Drop)] +struct FastMapIteratorMut { + _inner_map: Map>, + _current_key: K, + _current_index: u64, + _length: u64, +} + +pub impl FastIterableMapIteratorImpl< + K, V, +Drop, +Store, +Hash, +Copy, +Store, +> of Iterator> { + type Item = (K, V); + fn next(ref self: FastMapIterator) -> Option { + if self._current_index >= self._length { + return Option::None; + } + let current_key = self._current_key; + let current_entry = self._inner_map.entry(current_key).read(); + self._current_key = current_entry.key(); + self._current_index += 1; + Option::Some((current_key, current_entry.value())) + } +} + +impl StoragePathFastIterableMapIntoIteratorImpl< + K, V, +Drop, +Store, +Hash, +Copy, +Store, +> of IntoIterator>> { + type IntoIter = FastMapIterator; + fn into_iter(self: StoragePath>) -> Self::IntoIter { + FastMapIterator { + _inner_map: self._inner_map, + _current_key: self._head.read(), + _current_index: 0, + _length: self._length.read(), + } + } +} + +pub impl FastIterableMapIteratorMutImpl< + K, V, +Drop, +Store, +Hash, +Copy, +Store, +> of Iterator> { + type Item = (K, V); + fn next(ref self: FastMapIteratorMut) -> Option { + if self._current_index >= self._length { + return Option::None; + } + let current_key = self._current_key; + let current_entry = self._inner_map.entry(current_key).read(); + self._current_key = current_entry.key(); + self._current_index += 1; + Option::Some((current_key, current_entry.value())) + } +} + +impl StoragePathMutableFastIterableMapIntoIteratorImpl< + K, V, +Drop, +Store, +Hash, +Copy, +Store, +> of IntoIterator>>> { + type IntoIter = FastMapIteratorMut; + fn into_iter(self: StoragePath>>) -> Self::IntoIter { + FastMapIteratorMut { + _inner_map: self._inner_map, + _current_key: self._head.read(), + _current_index: 0, + _length: self._length.read(), + } + } +} + +pub impl FastIterableMapIntoIterImpl< + T, + +Drop, + impl StorageAsPathImpl: StorageAsPath, + impl StoragePathImpl: IntoIterator>, + +FastIterableMapTrait>, +> of IntoIterator { + type IntoIter = StoragePathImpl::IntoIter; + fn into_iter(self: T) -> Self::IntoIter { + self.as_path().into_iter() + } +}