Skip to main content

stylus_sdk/storage/
array.rs

1// Copyright 2023-2024, Offchain Labs, Inc.
2// For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/main/licenses/COPYRIGHT.md
3
4use core::marker::PhantomData;
5
6use alloy_primitives::U256;
7use stylus_core::HostAccess;
8
9use super::{Erase, StorageGuard, StorageGuardMut, StorageType};
10use crate::host::VM;
11
12/// Accessor for a storage-backed array.
13pub struct StorageArray<S: StorageType, const N: usize> {
14    slot: U256,
15    marker: PhantomData<S>,
16    __stylus_host: VM,
17}
18
19impl<S: StorageType, const N: usize> StorageType for StorageArray<S, N> {
20    type Wraps<'a>
21        = StorageGuard<'a, StorageArray<S, N>>
22    where
23        Self: 'a;
24    type WrapsMut<'a>
25        = StorageGuardMut<'a, StorageArray<S, N>>
26    where
27        Self: 'a;
28
29    const REQUIRED_SLOTS: usize = Self::required_slots();
30
31    unsafe fn new(slot: U256, offset: u8, host: VM) -> Self {
32        debug_assert!(offset == 0);
33        Self {
34            slot,
35            marker: PhantomData,
36            __stylus_host: host,
37        }
38    }
39
40    fn load<'s>(self) -> Self::Wraps<'s> {
41        StorageGuard::new(self)
42    }
43
44    fn load_mut<'s>(self) -> Self::WrapsMut<'s> {
45        StorageGuardMut::new(self)
46    }
47}
48
49impl<S: StorageType, const N: usize> HostAccess for StorageArray<S, N> {
50    type Host = VM;
51
52    #[inline]
53    fn vm(&self) -> &Self::Host {
54        &self.__stylus_host
55    }
56}
57
58#[cfg(feature = "stylus-test")]
59impl<const N: usize, S, T> From<&T> for StorageArray<S, N>
60where
61    T: stylus_core::Host + Clone + 'static,
62    S: StorageType,
63{
64    fn from(host: &T) -> Self {
65        unsafe {
66            Self::new(
67                U256::ZERO,
68                0,
69                crate::host::VM {
70                    host: alloc::boxed::Box::new(host.clone()),
71                },
72            )
73        }
74    }
75}
76
77impl<S: StorageType, const N: usize> StorageArray<S, N> {
78    /// Gets the number of elements stored.
79    ///
80    /// Although this type will always have the same length, this method is still provided for
81    /// consistency with [`StorageVec`].
82    #[allow(clippy::len_without_is_empty)]
83    pub const fn len(&self) -> usize {
84        N
85    }
86
87    /// Gets an accessor to the element at a given index, if it exists.
88    /// Note: the accessor is protected by a [`StorageGuard`], which restricts
89    /// its lifetime to that of `&self`.
90    pub fn getter(&self, index: impl TryInto<usize>) -> Option<StorageGuard<'_, S>> {
91        let store = unsafe { self.accessor(index)? };
92        Some(StorageGuard::new(store))
93    }
94
95    /// Gets a mutable accessor to the element at a given index, if it exists.
96    /// Note: the accessor is protected by a [`StorageGuardMut`], which restricts
97    /// its lifetime to that of `&mut self`.
98    pub fn setter(&mut self, index: impl TryInto<usize>) -> Option<StorageGuardMut<'_, S>> {
99        let store = unsafe { self.accessor(index)? };
100        Some(StorageGuardMut::new(store))
101    }
102
103    /// Gets the underlying accessor to the element at a given index, if it exists.
104    ///
105    /// # Safety
106    ///
107    /// Enables aliasing.
108    unsafe fn accessor(&self, index: impl TryInto<usize>) -> Option<S> {
109        let index = index.try_into().ok()?;
110        if index >= N {
111            return None;
112        }
113        let (slot, offset) = self.index_slot(index);
114        Some(S::new(slot, offset, self.__stylus_host.clone()))
115    }
116
117    /// Gets the underlying accessor to the element at a given index, even if out of bounds.
118    ///
119    /// # Safety
120    ///
121    /// Enables aliasing. UB if out of bounds.
122    unsafe fn accessor_unchecked(&self, index: usize) -> S {
123        let (slot, offset) = self.index_slot(index);
124        S::new(slot, offset, self.__stylus_host.clone())
125    }
126
127    /// Gets the element at the given index, if it exists.
128    pub fn get(&self, index: impl TryInto<usize>) -> Option<S::Wraps<'_>> {
129        let store = unsafe { self.accessor(index)? };
130        Some(store.load())
131    }
132
133    /// Gets a mutable accessor to the element at a given index, if it exists.
134    pub fn get_mut(&mut self, index: impl TryInto<usize>) -> Option<S::WrapsMut<'_>> {
135        let store = unsafe { self.accessor(index)? };
136        Some(store.load_mut())
137    }
138
139    /// Determines the slot and offset for the element at an index.
140    fn index_slot(&self, index: usize) -> (U256, u8) {
141        let width = S::SLOT_BYTES;
142        let words = S::REQUIRED_SLOTS.max(1);
143        let density = Self::density();
144
145        let slot = self.slot + U256::from(words * index / density);
146        let offset = 32 - (width * (1 + index % density)) as u8;
147        (slot, offset)
148    }
149
150    /// Number of elements per slot.
151    const fn density() -> usize {
152        32 / S::SLOT_BYTES
153    }
154
155    /// Required slots for the storage array.
156    const fn required_slots() -> usize {
157        let reserved = N * S::REQUIRED_SLOTS;
158        let density = Self::density();
159        let packed = N.div_ceil(density);
160        if reserved > packed {
161            return reserved;
162        }
163        packed
164    }
165}
166
167impl<S: Erase, const N: usize> Erase for StorageArray<S, N> {
168    fn erase(&mut self) {
169        for i in 0..N {
170            let mut store = unsafe { self.accessor_unchecked(i) };
171            store.erase()
172        }
173    }
174}