tempo_precompiles/storage/types/
set.rs

1//! OpenZeppelin's EnumerableSet implementation for EVM storage using Rust primitives.
2//! <https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/structs/EnumerableSet.sol>
3//!
4//! # Storage Layout
5//!
6//! EnumerableSet uses two storage structures:
7//! - **Values Vec**: A `Vec<T>` storing all set elements at `keccak256(base_slot)`
8//! - **Positions Mapping**: A `Mapping<T, u32>` at `base_slot + 1` storing 1-indexed positions
9//!   - Position 0 means the value is not in the set
10//!   - Position N means the value is at index N-1 in the values array
11//!
12//! # Design
13//!
14//! Two complementary types:
15//! - `Set<T>`: Read-only in-memory snapshot. `Vec<T>` wrapper. Ordered like storage.
16//! - `SetHandler<T>`: Storage operations.
17//!
18//! # Usage Patterns
19//!
20//! ## Single Operations (O(1) each)
21//! ```ignore
22//! handler.insert(addr)?;   // Direct storage write
23//! handler.remove(&addr)?;  // Direct storage write
24//! handler.contains(&addr)?; // Direct storage read
25//! ```
26//!
27//! ## Bulk Read
28//! ```ignore
29//! let set: Set<Address> = handler.read()?;
30//! for addr in &set {
31//!     // Iteration preserves storage order
32//!     // set[i] == handler.at(i)
33//! }
34//! ```
35//!
36//! ## Bulk Mutation
37//! ```ignore
38//! let mut vec: Vec<_> = handler.read()?.into();
39//! vec.push(new_addr);
40//! vec.retain(|a| a != &old_addr);
41//! handler.write(vec.into())?;  // `Set::from(vec)` deduplicates
42//! ```
43
44use alloy::primitives::{Address, U256};
45use std::{
46    collections::HashSet,
47    fmt,
48    hash::Hash,
49    ops::{Deref, Index},
50};
51
52use crate::{
53    error::{Result, TempoPrecompileError},
54    storage::{
55        Handler, Layout, LayoutCtx, Storable, StorableType, StorageKey, StorageOps,
56        types::{Mapping, Slot, vec::VecHandler},
57    },
58};
59
60/// Read-only snapshot of a set stored via [`SetHandler`].
61///
62/// Elements are ordered by their position in the underlying storage array.
63/// This order is **not** guaranteed to match insertion order: `SetHandler::remove`
64/// uses swap-and-pop semantics, so removing a non-tail element moves the last
65/// element into the vacated slot.
66///
67/// To mutate:
68/// 1. Convert to `Vec<T>` with `.into()`
69/// 2. Modify the Vec
70/// 3. Convert back with `Set::from(vec)` (deduplicates, preserves first-occurrence order)
71/// 4. Write with `handler.write(set)`
72///
73/// For single-element mutations, use `SetHandler` methods directly.
74///
75/// Implements `Deref<Target = [T]>`, so all slice methods are available:
76/// `len()`, `is_empty()`, `iter()`, `get()`, `contains()`, indexing, etc.
77#[derive(Debug, Clone, PartialEq, Eq, Default)]
78pub struct Set<T>(Vec<T>);
79
80impl<T> Set<T> {
81    /// Creates a new empty set.
82    #[inline]
83    pub fn new() -> Self {
84        Self(Vec::new())
85    }
86
87    /// Creates a set from a vector that is already known to contain no duplicates.
88    ///
89    /// # IMPORTANT
90    ///
91    /// The caller **must** guarantee that `vec` contains no duplicate elements.
92    /// Violating this breaks the position-mapping invariant in storage: two equal values would
93    /// share a single position slot, causing silent data corruption on subsequent `remove()` calls.
94    #[inline]
95    pub fn new_unchecked(vec: Vec<T>) -> Self {
96        Self(vec)
97    }
98}
99
100impl<T> Deref for Set<T> {
101    type Target = [T];
102
103    #[inline]
104    fn deref(&self) -> &[T] {
105        &self.0
106    }
107}
108
109impl<T> From<Set<T>> for Vec<T> {
110    #[inline]
111    fn from(set: Set<T>) -> Self {
112        set.0
113    }
114}
115
116impl<T: Eq + Hash + Clone> From<Vec<T>> for Set<T> {
117    /// Creates a set from a vector, removing duplicates.
118    ///
119    /// Preserves the order of first occurrences.
120    fn from(vec: Vec<T>) -> Self {
121        let (mut seen, mut deduped) = (HashSet::new(), Vec::new());
122        for item in vec {
123            if seen.insert(item.clone()) {
124                deduped.push(item);
125            }
126        }
127        Self(deduped)
128    }
129}
130
131impl<T: Eq + Hash + Clone> FromIterator<T> for Set<T> {
132    /// Creates a set from an iterator, removing duplicates.
133    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
134        let vec: Vec<T> = iter.into_iter().collect();
135        Self::from(vec)
136    }
137}
138
139impl<T> IntoIterator for Set<T> {
140    type Item = T;
141    type IntoIter = std::vec::IntoIter<T>;
142
143    #[inline]
144    fn into_iter(self) -> Self::IntoIter {
145        self.0.into_iter()
146    }
147}
148
149impl<'a, T> IntoIterator for &'a Set<T> {
150    type Item = &'a T;
151    type IntoIter = std::slice::Iter<'a, T>;
152
153    #[inline]
154    fn into_iter(self) -> Self::IntoIter {
155        self.0.iter()
156    }
157}
158
159/// Type-safe handler for accessing `Set<T>` in storage.
160///
161/// Provides the OZ storage operations but following the naming convention of `HashSet`:
162///
163/// | Method         | OZ equivalent    |
164/// |----------------|------------------|
165/// | `insert()`     | `add()`          |
166/// | `remove()`     | `remove()`       |
167/// | `contains()`   | `contains()`     |
168/// | `len()`        | `length()`       |
169/// | `at()`         | `at()`           |
170/// | `read()`       | `values()`       |
171/// | `read_range()` | `values_range()` |
172///
173/// Also implements `Handler<Set<T>>` for bulk operations:
174/// - `read`: Load all elements as `Set<T>`
175/// - `write`: Replace entire set
176/// - `delete`: Remove all elements
177pub struct SetHandler<T>
178where
179    T: Storable + StorageKey + Hash + Eq + Clone,
180{
181    /// Handler for the values vector (stores actual elements).
182    values: VecHandler<T>,
183    /// Handler for the positions mapping (value -> 1-indexed position).
184    positions: Mapping<T, u32>,
185    /// The base slot for the set.
186    base_slot: U256,
187    /// Contract address.
188    address: Address,
189}
190
191/// Set occupies 2 slots:
192///
193/// - Slot 0: `Vec` length slot, with data at `keccak256(slot)`
194/// - Slot 1: `Mapping` base slot for positions
195impl<T> StorableType for Set<T>
196where
197    T: Storable + StorageKey + Hash + Eq + Clone,
198{
199    const LAYOUT: Layout = Layout::Slots(2);
200    const IS_DYNAMIC: bool = true;
201    type Handler = SetHandler<T>;
202
203    fn handle(slot: U256, _ctx: LayoutCtx, address: Address) -> Self::Handler {
204        SetHandler::new(slot, address)
205    }
206}
207
208/// Storable implementation for `Set<T>`.
209impl<T> Storable for Set<T>
210where
211    T: Storable + StorageKey + Hash + Eq + Clone,
212    T::Handler: Handler<T>,
213{
214    fn load<S: StorageOps>(storage: &S, slot: U256, _ctx: LayoutCtx) -> Result<Self> {
215        let values: Vec<T> = Vec::load(storage, slot, LayoutCtx::FULL)?;
216        Ok(Self(values))
217    }
218
219    fn store<S: StorageOps>(&self, _storage: &mut S, _slot: U256, _ctx: LayoutCtx) -> Result<()> {
220        Err(TempoPrecompileError::Fatal(
221            "Set must be stored via SetHandler::write() to maintain position invariants".into(),
222        ))
223    }
224
225    fn delete<S: StorageOps>(storage: &mut S, slot: U256, ctx: LayoutCtx) -> Result<()> {
226        let values: Vec<T> = Vec::load(storage, slot, LayoutCtx::FULL)?;
227
228        for value in values {
229            let pos_slot = value.mapping_slot(slot + U256::ONE);
230            <U256 as Storable>::delete(storage, pos_slot, LayoutCtx::FULL)?;
231        }
232
233        <Vec<T> as Storable>::delete(storage, slot, ctx)
234    }
235}
236
237/// Converts a 0-based index to a 1-based position for storage.
238///
239/// Returns an error if the result would overflow `u32`, which would corrupt the
240/// sentinel value (`0` means "not present") used by `contains()` and `remove()`.
241#[inline]
242fn checked_position(index: usize) -> Result<u32> {
243    u32::try_from(index)
244        .ok()
245        .and_then(|i| i.checked_add(1))
246        .ok_or_else(TempoPrecompileError::under_overflow)
247}
248
249impl<T> SetHandler<T>
250where
251    T: Storable + StorageKey + Hash + Eq + Clone,
252{
253    /// Creates a new handler for the set at the given base slot.
254    ///
255    /// - `base_slot`: Used as the Vec's length slot
256    /// - `base_slot + 1`: Used as the Mapping's base slot
257    pub fn new(base_slot: U256, address: Address) -> Self {
258        Self {
259            values: VecHandler::new(base_slot, address),
260            positions: Mapping::new(base_slot + U256::ONE, address),
261            base_slot,
262            address,
263        }
264    }
265
266    /// Returns the base storage slot for this set.
267    #[inline]
268    pub fn base_slot(&self) -> U256 {
269        self.base_slot
270    }
271
272    /// Returns the number of elements in the set.
273    #[inline]
274    pub fn len(&self) -> Result<usize> {
275        self.values.len()
276    }
277
278    /// Returns whether the set is empty.
279    #[inline]
280    pub fn is_empty(&self) -> Result<bool> {
281        self.values.is_empty()
282    }
283
284    /// Returns true if the value is in the set.
285    pub fn contains(&self, value: &T) -> Result<bool>
286    where
287        T: StorageKey + Hash + Eq + Clone,
288    {
289        self.positions.at(value).read().map(|pos| pos != 0)
290    }
291
292    /// Inserts a value into the set.
293    ///
294    /// Returns `true` if the value was inserted (not already present).
295    /// Returns `false` if the value was already in the set.
296    #[inline]
297    pub fn insert(&mut self, value: T) -> Result<bool>
298    where
299        T: StorageKey + Hash + Eq + Clone,
300        T::Handler: Handler<T>,
301    {
302        // Check if already present
303        if self.contains(&value)? {
304            return Ok(false);
305        }
306
307        // Store position (1-indexed: position N means index N-1)
308        let length = self.values.len()?;
309        self.positions
310            .at_mut(&value)
311            .write(checked_position(length)?)?;
312
313        // Push value to the array
314        self.values.push(value)?;
315
316        Ok(true)
317    }
318
319    /// Removes a value from the set.
320    ///
321    /// Returns `true` if the value was removed. Otherwise, returns `false`.
322    #[inline]
323    pub fn remove(&mut self, value: &T) -> Result<bool>
324    where
325        T: StorageKey + Hash + Eq + Clone,
326        T::Handler: Handler<T>,
327    {
328        // Get position (1-indexed, 0 means not present)
329        let position = self.positions.at(value).read()?;
330        if position == 0 {
331            return Ok(false);
332        }
333
334        let len = self.values.len()?;
335        // Validate invariants
336        debug_assert!(
337            len != 0 && (position as usize) <= len,
338            "Set invariant violation: position exceeds length"
339        );
340
341        // Convert to 0-indexed
342        let last_index = len - 1;
343        let index = (position - 1) as usize;
344
345        // Swap with last element if not already last
346        if index != last_index {
347            let last_value = self.values[last_index].read()?;
348            self.positions.at_mut(&last_value).write(position)?;
349            self.values[index].write(last_value)?;
350        }
351
352        // Delete the last element and decrement its length.
353        // Equivalent to `self.values.pop()`, but without the OOB checks.
354        self.values[last_index].delete()?;
355        Slot::<U256>::new(self.values.len_slot(), self.address).write(U256::from(last_index))?;
356
357        // Clear removed value's position
358        self.positions.at_mut(value).delete()?;
359
360        Ok(true)
361    }
362
363    /// Returns the value at the given index with bounds checking.
364    ///
365    /// # Returns
366    /// - If the SLOAD to read the length fails, returns an error.
367    /// - If the index is OOB, returns `Ok(None)`.
368    /// - Otherwise, returns `Ok(Some(T))`.
369    pub fn at(&self, index: usize) -> Result<Option<T>>
370    where
371        T::Handler: Handler<T>,
372    {
373        if index >= self.len()? {
374            return Ok(None);
375        }
376        Ok(Some(self.values[index].read()?))
377    }
378
379    /// Reads a range of values from the set.
380    ///
381    /// This is a partial version of `read()` for when you only need a subset.
382    pub fn read_range(&self, start: usize, end: usize) -> Result<Vec<T>>
383    where
384        T::Handler: Handler<T>,
385    {
386        let len = self.len()?;
387        let end = end.min(len);
388        let start = start.min(end);
389
390        let mut result = Vec::new();
391        for i in start..end {
392            result.push(self.values[i].read()?);
393        }
394        Ok(result)
395    }
396}
397
398impl<T> Handler<Set<T>> for SetHandler<T>
399where
400    T: Storable + StorageKey + Hash + Eq + Clone,
401    T::Handler: Handler<T>,
402{
403    /// Reads all elements from storage as a `Set<T>`.
404    ///
405    /// The returned `Set` preserves storage order: `set[i] == handler.at(i)`.
406    fn read(&self) -> Result<Set<T>> {
407        let len = self.len()?;
408        let mut vec = Vec::new();
409
410        for i in 0..len {
411            vec.push(self.values[i].read()?);
412        }
413
414        Ok(Set(vec))
415    }
416
417    /// Replaces the entire set with new contents.
418    ///
419    /// The input Set is deduplicated by the `From<Vec<T>>` conversion.
420    fn write(&mut self, value: Set<T>) -> Result<()> {
421        let old_len = self.values.len()?;
422        let new_len = value.0.len();
423
424        // Clear old positions
425        for i in 0..old_len {
426            let old_value = self.values[i].read()?;
427            self.positions.at_mut(&old_value).delete()?;
428        }
429
430        // Write new values and positions (1-indexed)
431        for (index, new_value) in value.0.into_iter().enumerate() {
432            self.positions
433                .at_mut(&new_value)
434                .write(checked_position(index)?)?;
435            self.values[index].write(new_value)?;
436        }
437
438        // Update length
439        Slot::<U256>::new(self.values.len_slot(), self.address).write(U256::from(new_len))?;
440
441        // Clear leftover value slots if shrinking
442        for i in new_len..old_len {
443            self.values[i].delete()?;
444        }
445
446        Ok(())
447    }
448
449    /// Deletes all elements from the set.
450    ///
451    /// Clears both the values array and all position entries.
452    fn delete(&mut self) -> Result<()> {
453        let len = self.len()?;
454
455        // Clear all position entries
456        for i in 0..len {
457            let value = self.values[i].read()?;
458            self.positions.at_mut(&value).delete()?;
459        }
460
461        // Delete the underlying vector (clears length and data slots)
462        self.values.delete()
463    }
464
465    fn t_read(&self) -> Result<Set<T>> {
466        Err(TempoPrecompileError::Fatal(
467            "Set types don't support transient storage".into(),
468        ))
469    }
470
471    fn t_write(&mut self, _value: Set<T>) -> Result<()> {
472        Err(TempoPrecompileError::Fatal(
473            "Set types don't support transient storage".into(),
474        ))
475    }
476
477    fn t_delete(&mut self) -> Result<()> {
478        Err(TempoPrecompileError::Fatal(
479            "Set types don't support transient storage".into(),
480        ))
481    }
482}
483
484impl<T> fmt::Debug for SetHandler<T>
485where
486    T: Storable + StorageKey + Hash + Eq + Clone,
487{
488    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
489        f.debug_struct("SetHandler")
490            .field("base_slot", &self.base_slot)
491            .field("address", &self.address)
492            .finish()
493    }
494}
495
496impl<T> Clone for SetHandler<T>
497where
498    T: Storable + StorageKey + Hash + Eq + Clone,
499{
500    fn clone(&self) -> Self {
501        Self::new(self.base_slot, self.address)
502    }
503}
504
505impl<T> Index<usize> for SetHandler<T>
506where
507    T: Storable + StorageKey + Hash + Eq + Clone,
508{
509    type Output = T::Handler;
510
511    /// Returns a reference to the cached handler for the given index (unchecked).
512    ///
513    /// **WARNING:** Does not check bounds. Caller must ensure that the index is valid.
514    /// For checked access use `.at(index)` instead.
515    fn index(&self, index: usize) -> &Self::Output {
516        &self.values[index]
517    }
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523    use crate::{storage::StorageCtx, test_util::setup_storage};
524    use alloy::primitives::Address;
525    use proptest::prelude::*;
526
527    // -- SET TYPE TESTS -------------------------------------------------------
528
529    #[test]
530    fn test_set_from_vec_deduplicates() {
531        let vec = vec![1, 2, 3, 2, 1, 4];
532        let set = Set::from(vec);
533
534        assert_eq!(set.len(), 4);
535        assert_eq!(&set[..], &[1, 2, 3, 4]); // Order preserved
536    }
537
538    #[test]
539    fn test_set_from_iter_deduplicates() {
540        let set: Set<i32> = [1, 2, 3, 2, 1, 4].into_iter().collect();
541
542        assert_eq!(set.len(), 4);
543        assert!(set.contains(&1));
544        assert!(set.contains(&4));
545    }
546
547    #[test]
548    fn test_set_preserves_first_occurrence_order() {
549        let vec = vec!['a', 'b', 'c', 'b', 'a', 'd'];
550        let set = Set::from(vec);
551
552        assert_eq!(&set[..], &['a', 'b', 'c', 'd']);
553    }
554
555    #[test]
556    fn test_set_into_vec() {
557        let set = Set::from(vec![1, 2, 3]);
558        let vec: Vec<i32> = set.into();
559
560        assert_eq!(vec, vec![1, 2, 3]);
561    }
562
563    #[test]
564    fn test_set_iteration() {
565        let set = Set::from(vec![10, 20, 30]);
566
567        let collected: Vec<_> = set.iter().copied().collect();
568        assert_eq!(collected, vec![10, 20, 30]);
569
570        let collected2: Vec<_> = (&set).into_iter().copied().collect();
571        assert_eq!(collected2, vec![10, 20, 30]);
572    }
573
574    #[test]
575    fn test_set_get() {
576        let set = Set::from(vec!['a', 'b', 'c']);
577
578        assert_eq!(set.first(), Some(&'a'));
579        assert_eq!(set.get(1), Some(&'b'));
580        assert_eq!(set.get(2), Some(&'c'));
581        assert_eq!(set.get(3), None);
582    }
583
584    #[test]
585    fn test_set_deref_to_slice() {
586        let set = Set::from(vec![1, 2, 3]);
587
588        assert_eq!(set[0], 1);
589        assert_eq!(set[1], 2);
590        assert_eq!(set.len(), 3);
591    }
592
593    // -- HANDLER TESTS --------------------------------------------------------
594
595    /// Tests the read -> Vec -> mutate -> write pattern documented in the module.
596    #[test]
597    fn test_set_write_via_vec_mutation() -> eyre::Result<()> {
598        let (mut storage, address) = setup_storage();
599
600        StorageCtx::enter(&mut storage, || {
601            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
602
603            handler.insert(U256::ONE)?;
604            handler.insert(U256::from(2))?;
605            handler.insert(U256::from(3))?;
606
607            // Read, convert to Vec, mutate, convert back, write
608            let mut vec: Vec<U256> = handler.read()?.into();
609            vec.push(U256::from(4));
610            vec.push(U256::from(5));
611            vec.retain(|&x| x != U256::from(2));
612
613            handler.write(vec.into())?;
614
615            assert_eq!(handler.len()?, 4);
616            assert!(handler.contains(&U256::ONE)?);
617            assert!(!handler.contains(&U256::from(2))?);
618            assert!(handler.contains(&U256::from(3))?);
619            assert!(handler.contains(&U256::from(4))?);
620            assert!(handler.contains(&U256::from(5))?);
621
622            Ok(())
623        })
624    }
625
626    #[test]
627    fn test_set_constructors_and_edge_cases() {
628        assert!(Set::<i32>::new().is_empty());
629        assert!(Set::<i32>::default().is_empty());
630        assert!(Set::from(Vec::<i32>::new()).is_empty());
631
632        let set = Set::from(vec![5, 5, 5, 5]);
633        assert_eq!(set.len(), 1);
634        assert_eq!(&set[..], &[5]);
635
636        let collected: Vec<i32> = Set::from(vec![1, 2, 3]).into_iter().collect();
637        assert_eq!(collected, vec![1, 2, 3]);
638
639        assert_eq!(Set::from(vec![1, 2, 3]), Set::from(vec![1, 2, 3]));
640        assert_ne!(Set::from(vec![1, 2, 3]), Set::from(vec![3, 2, 1]));
641    }
642
643    // -- HANDLER TESTS --------------------------------------------------------
644
645    #[test]
646    fn test_handler_empty_state() -> eyre::Result<()> {
647        let (mut storage, address) = setup_storage();
648
649        StorageCtx::enter(&mut storage, || {
650            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
651
652            assert!(handler.is_empty()?);
653            assert_eq!(handler.len()?, 0);
654            assert!(!handler.contains(&U256::ONE)?);
655            assert!(!handler.remove(&U256::ONE)?);
656            assert_eq!(handler.at(0)?, None);
657            assert_eq!(handler.at(100)?, None);
658            assert!(handler.read()?.is_empty());
659            assert!(handler.read_range(0, 10)?.is_empty());
660
661            Ok(())
662        })
663    }
664
665    #[test]
666    fn test_handler_insert_remove_basics() -> eyre::Result<()> {
667        let (mut storage, address) = setup_storage();
668
669        StorageCtx::enter(&mut storage, || {
670            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
671
672            assert!(handler.insert(U256::ONE)?);
673            assert!(!handler.insert(U256::ONE)?);
674            assert_eq!(handler.len()?, 1);
675
676            assert!(handler.remove(&U256::ONE)?);
677            assert!(handler.is_empty()?);
678            assert!(!handler.contains(&U256::ONE)?);
679
680            handler.insert(U256::from(1))?;
681            handler.insert(U256::from(2))?;
682            handler.remove(&U256::from(1))?;
683            handler.insert(U256::from(3))?;
684            assert_eq!(handler.len()?, 2);
685            assert!(handler.contains(&U256::from(2))?);
686            assert!(handler.contains(&U256::from(3))?);
687            assert!(!handler.contains(&U256::from(1))?);
688
689            Ok(())
690        })
691    }
692
693    #[test]
694    fn test_handler_remove_swap_semantics() -> eyre::Result<()> {
695        let (mut storage, address) = setup_storage();
696
697        StorageCtx::enter(&mut storage, || {
698            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
699
700            handler.insert(U256::from(10))?;
701            handler.insert(U256::from(20))?;
702            handler.insert(U256::from(30))?;
703
704            // Remove last: no swap needed
705            assert!(handler.remove(&U256::from(30))?);
706            assert_eq!(&handler.read()?[..], &[U256::from(10), U256::from(20)]);
707
708            // Re-add and remove first: last swaps into position 0
709            handler.insert(U256::from(30))?;
710            assert!(handler.remove(&U256::from(10))?);
711            assert_eq!(&handler.read()?[..], &[U256::from(30), U256::from(20)]);
712
713            // Remove first of two
714            assert!(handler.remove(&U256::from(30))?);
715            assert_eq!(handler.len()?, 1);
716            assert!(handler.contains(&U256::from(20))?);
717
718            Ok(())
719        })
720    }
721
722    #[test]
723    fn test_handler_at_and_index() -> eyre::Result<()> {
724        let (mut storage, address) = setup_storage();
725
726        StorageCtx::enter(&mut storage, || {
727            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
728
729            handler.insert(U256::from(10))?;
730            handler.insert(U256::from(20))?;
731            handler.insert(U256::from(30))?;
732
733            assert_eq!(handler.at(0)?, Some(U256::from(10)));
734            assert_eq!(handler.at(1)?, Some(U256::from(20)));
735            assert_eq!(handler.at(2)?, Some(U256::from(30)));
736            assert_eq!(handler.at(3)?, None);
737
738            assert_eq!(handler[0].read()?, U256::from(10));
739            assert_eq!(handler[1].read()?, U256::from(20));
740
741            Ok(())
742        })
743    }
744
745    #[test]
746    fn test_handler_read_range() -> eyre::Result<()> {
747        let (mut storage, address) = setup_storage();
748
749        StorageCtx::enter(&mut storage, || {
750            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
751
752            for i in 0..5 {
753                handler.insert(U256::from(i))?;
754            }
755
756            assert_eq!(
757                handler.read_range(1, 4)?,
758                vec![U256::from(1), U256::from(2), U256::from(3)]
759            );
760            // end > len clamps
761            assert_eq!(handler.read_range(0, 100)?.len(), 5);
762            // start > end returns empty
763            assert!(handler.read_range(5, 3)?.is_empty());
764
765            Ok(())
766        })
767    }
768
769    #[test]
770    fn test_handler_write() -> eyre::Result<()> {
771        let (mut storage, address) = setup_storage();
772
773        StorageCtx::enter(&mut storage, || {
774            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
775
776            // Write to grow (1 → 3)
777            handler.insert(U256::from(1))?;
778            handler.write(Set::from(vec![
779                U256::from(10),
780                U256::from(20),
781                U256::from(30),
782            ]))?;
783            assert_eq!(handler.len()?, 3);
784            assert!(!handler.contains(&U256::from(1))?);
785            assert!(handler.contains(&U256::from(10))?);
786
787            // Write to shrink (3 → 2)
788            handler.write(Set::from(vec![U256::from(40), U256::from(50)]))?;
789            assert_eq!(handler.len()?, 2);
790            assert!(!handler.contains(&U256::from(10))?);
791
792            // Write empty
793            handler.write(Set::new())?;
794            assert!(handler.is_empty()?);
795
796            Ok(())
797        })
798    }
799
800    #[test]
801    fn test_handler_delete() -> eyre::Result<()> {
802        let (mut storage, address) = setup_storage();
803
804        StorageCtx::enter(&mut storage, || {
805            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
806
807            for i in 1..=3 {
808                handler.insert(U256::from(i))?;
809            }
810
811            handler.delete()?;
812            assert!(handler.is_empty()?);
813            for i in 1..=3 {
814                assert!(!handler.contains(&U256::from(i))?);
815            }
816
817            // Re-insert after delete: positions were properly cleared
818            handler.insert(U256::from(2))?;
819            assert_eq!(handler.at(0)?, Some(U256::from(2)));
820            assert_eq!(handler.len()?, 1);
821
822            Ok(())
823        })
824    }
825
826    #[test]
827    fn test_handler_transient_storage_errors() -> eyre::Result<()> {
828        let (mut storage, address) = setup_storage();
829
830        StorageCtx::enter(&mut storage, || {
831            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
832            assert!(handler.t_read().is_err());
833            assert!(handler.t_write(Set::new()).is_err());
834            assert!(handler.t_delete().is_err());
835            Ok(())
836        })
837    }
838
839    #[test]
840    fn test_checked_position() {
841        // Happy path: index 0 → position 1, last valid index → u32::MAX
842        assert_eq!(checked_position(0).unwrap(), 1);
843        assert_eq!(checked_position(1).unwrap(), 2);
844        assert_eq!(checked_position(u32::MAX as usize - 1).unwrap(), u32::MAX);
845
846        // Overflow: u32::MAX would produce position u32::MAX + 1, which wraps to 0
847        assert!(checked_position(u32::MAX as usize).is_err());
848        // usize values beyond u32::MAX also overflow
849        assert!(checked_position(u32::MAX as usize + 1).is_err());
850    }
851
852    #[test]
853    fn test_handler_insert_overflow() -> eyre::Result<()> {
854        let (mut storage, address) = setup_storage();
855
856        StorageCtx::enter(&mut storage, || {
857            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
858
859            // Simulate a full set by writing u32::MAX directly to the length slot.
860            // insert() must propagate an overflow error rather than wrapping position to 0.
861            Slot::<U256>::new(handler.base_slot(), address).write(U256::from(u32::MAX))?;
862            assert!(handler.insert(U256::ONE).is_err());
863
864            Ok(())
865        })
866    }
867
868    #[test]
869    fn test_handler_metadata() {
870        let address = Address::ZERO;
871        let handler = SetHandler::<U256>::new(U256::from(42), address);
872        assert_eq!(handler.base_slot(), U256::from(42));
873
874        let debug_str = format!("{handler:?}");
875        assert!(debug_str.contains("SetHandler"));
876
877        let cloned = handler.clone();
878        assert_eq!(cloned.base_slot(), handler.base_slot());
879    }
880
881    #[test]
882    fn test_handler_address_set() -> eyre::Result<()> {
883        let (mut storage, address) = setup_storage();
884
885        StorageCtx::enter(&mut storage, || {
886            let mut handler = SetHandler::<Address>::new(U256::ZERO, address);
887
888            let [a1, a2, a3] = [[1u8; 20], [2u8; 20], [3u8; 20]].map(Address::from);
889
890            for a in [a1, a2, a3] {
891                handler.insert(a)?;
892            }
893            assert_eq!(handler.len()?, 3);
894
895            handler.remove(&a2)?;
896            assert_eq!(handler.len()?, 2);
897            assert!(!handler.contains(&a2)?);
898            assert_eq!(handler.at(0)?, Some(a1));
899            assert_eq!(handler.at(1)?, Some(a3));
900
901            Ok(())
902        })
903    }
904
905    #[test]
906    fn test_handler_multiple_remove_insert_cycles() -> eyre::Result<()> {
907        let (mut storage, address) = setup_storage();
908
909        StorageCtx::enter(&mut storage, || {
910            let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
911
912            for i in 0..5 {
913                handler.insert(U256::from(i))?;
914            }
915            for i in 0..5 {
916                assert!(handler.remove(&U256::from(i))?);
917            }
918            assert!(handler.is_empty()?);
919
920            for i in 10..15 {
921                handler.insert(U256::from(i))?;
922            }
923            assert_eq!(handler.len()?, 5);
924            for i in 10..15 {
925                assert!(handler.contains(&U256::from(i))?);
926            }
927
928            Ok(())
929        })
930    }
931
932    // -- PROPERTY TESTS -------------------------------------------------------
933
934    fn arb_address() -> impl Strategy<Value = Address> {
935        any::<[u8; 20]>().prop_map(Address::from)
936    }
937
938    proptest! {
939        #![proptest_config(ProptestConfig::with_cases(100))]
940
941        #[test]
942        fn proptest_set_order_alignment(addresses in prop::collection::vec(arb_address(), 1..20)) {
943            let (mut storage, address) = setup_storage();
944
945            StorageCtx::enter(&mut storage, || -> std::result::Result<(), TestCaseError> {
946                let mut handler = SetHandler::<Address>::new(U256::ZERO, address);
947
948                for addr in &addresses {
949                    handler.insert(*addr)?;
950                }
951
952                let set = handler.read()?;
953
954                for i in 0..set.len() {
955                    prop_assert_eq!(set.get(i).cloned(), handler.at(i)?, "Order mismatch at index {}", i);
956                }
957
958                Ok(())
959            }).unwrap();
960        }
961
962        #[test]
963        fn proptest_insert_remove_contains(
964            ops in prop::collection::vec(
965                (any::<u64>(), any::<bool>()),
966                1..50
967            )
968        ) {
969            let (mut storage, address) = setup_storage();
970
971            StorageCtx::enter(&mut storage, || -> std::result::Result<(), TestCaseError> {
972                let mut handler = SetHandler::<U256>::new(U256::ZERO, address);
973                let mut reference: Vec<U256> = Vec::new();
974
975                for (val, insert) in ops {
976                    let value = U256::from(val % 20); // keep key space small for collisions
977                    if insert {
978                        let was_new = !reference.contains(&value);
979                        let result = handler.insert(value)?;
980                        prop_assert_eq!(result, was_new);
981                        if was_new {
982                            reference.push(value);
983                        }
984                    } else {
985                        let existed = reference.contains(&value);
986                        let result = handler.remove(&value)?;
987                        prop_assert_eq!(result, existed);
988                        if existed {
989                            reference.retain(|v| v != &value);
990                        }
991                    }
992                }
993
994                prop_assert_eq!(handler.len()?, reference.len());
995                for v in &reference {
996                    prop_assert!(handler.contains(v)?);
997                }
998
999                Ok(())
1000            }).unwrap();
1001        }
1002    }
1003}
tempo_precompiles/storage/types/set.rs

tempo_precompiles/storage/types/
set.rs
