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

tempo_precompiles/storage/types/
set.rs
