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