foundry_evm_fuzz/strategies/
uint.rs

1use alloy_dyn_abi::{DynSolType, DynSolValue};
2use alloy_primitives::U256;
3use proptest::{
4    prelude::Rng,
5    strategy::{NewTree, Strategy, ValueTree},
6    test_runner::TestRunner,
7};
8
9/// Value tree for unsigned ints (up to uint256).
10pub struct UintValueTree {
11    /// Lower base
12    lo: U256,
13    /// Current value
14    curr: U256,
15    /// Higher base
16    hi: U256,
17    /// If true cannot be simplified or complexified
18    fixed: bool,
19}
20
21impl UintValueTree {
22    /// Create a new tree
23    /// # Arguments
24    /// * `start` - Starting value for the tree
25    /// * `fixed` - If `true` the tree would only contain one element and won't be simplified.
26    fn new(start: U256, fixed: bool) -> Self {
27        Self { lo: U256::ZERO, curr: start, hi: start, fixed }
28    }
29
30    fn reposition(&mut self) -> bool {
31        let interval = self.hi - self.lo;
32        let new_mid = self.lo + interval / U256::from(2);
33
34        if new_mid == self.curr {
35            false
36        } else {
37            self.curr = new_mid;
38            true
39        }
40    }
41}
42
43impl ValueTree for UintValueTree {
44    type Value = U256;
45
46    fn current(&self) -> Self::Value {
47        self.curr
48    }
49
50    fn simplify(&mut self) -> bool {
51        if self.fixed || (self.hi <= self.lo) {
52            return false;
53        }
54        self.hi = self.curr;
55        self.reposition()
56    }
57
58    fn complicate(&mut self) -> bool {
59        if self.fixed || (self.hi <= self.lo) {
60            return false;
61        }
62
63        self.lo = self.curr + U256::from(1);
64        self.reposition()
65    }
66}
67
68/// Value tree for unsigned ints (up to uint256).
69/// The strategy combines 3 different strategies, each assigned a specific weight:
70/// 1. Generate purely random value in a range. This will first choose bit size uniformly (up `bits`
71///    param). Then generate a value for this bit size.
72/// 2. Generate a random value around the edges (+/- 3 around 0 and max possible value)
73/// 3. Generate a value from a predefined fixtures set
74///
75/// To define uint fixtures:
76/// - return an array of possible values for a parameter named `amount` declare a function `function
77///   fixture_amount() public returns (uint32[] memory)`.
78/// - use `amount` named parameter in fuzzed test in order to include fixtures in fuzzed values
79///   `function testFuzz_uint32(uint32 amount)`.
80///
81/// If fixture is not a valid uint type then error is raised and random value generated.
82#[derive(Debug)]
83pub struct UintStrategy {
84    /// Bit size of uint (e.g. 256)
85    bits: usize,
86    /// A set of fixtures to be generated
87    fixtures: Vec<DynSolValue>,
88    /// The weight for edge cases (+/- 3 around 0 and max possible value)
89    edge_weight: usize,
90    /// The weight for fixtures
91    fixtures_weight: usize,
92    /// The weight for purely random values
93    random_weight: usize,
94}
95
96impl UintStrategy {
97    /// Create a new strategy.
98    /// #Arguments
99    /// * `bits` - Size of uint in bits
100    /// * `fixtures` - A set of fixed values to be generated (according to fixtures weight)
101    pub fn new(bits: usize, fixtures: Option<&[DynSolValue]>) -> Self {
102        Self {
103            bits,
104            fixtures: Vec::from(fixtures.unwrap_or_default()),
105            edge_weight: 10usize,
106            fixtures_weight: 40usize,
107            random_weight: 50usize,
108        }
109    }
110
111    fn generate_edge_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
112        let rng = runner.rng();
113        // Choose if we want values around 0 or max
114        let is_min = rng.random::<bool>();
115        let offset = U256::from(rng.random_range(0..4));
116        let start = if is_min { offset } else { self.type_max().saturating_sub(offset) };
117        Ok(UintValueTree::new(start, false))
118    }
119
120    fn generate_fixtures_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
121        // generate random cases if there's no fixtures
122        if self.fixtures.is_empty() {
123            return self.generate_random_tree(runner);
124        }
125
126        // Generate value tree from fixture.
127        let fixture = &self.fixtures[runner.rng().random_range(0..self.fixtures.len())];
128        if let Some(uint_fixture) = fixture.as_uint()
129            && uint_fixture.1 == self.bits
130        {
131            return Ok(UintValueTree::new(uint_fixture.0, false));
132        }
133
134        // If fixture is not a valid type, raise error and generate random value.
135        error!("{:?} is not a valid {} fixture", fixture, DynSolType::Uint(self.bits));
136        self.generate_random_tree(runner)
137    }
138
139    fn generate_random_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
140        let rng = runner.rng();
141
142        // generate random number of bits uniformly
143        let bits = rng.random_range(0..=self.bits);
144
145        // init 2 128-bit randoms
146        let mut higher: u128 = rng.random_range(0..=u128::MAX);
147        let mut lower: u128 = rng.random_range(0..=u128::MAX);
148
149        // cut 2 randoms according to bits size
150        match bits {
151            x if x < 128 => {
152                lower &= (1u128 << x) - 1;
153                higher = 0;
154            }
155            x if (128..256).contains(&x) => higher &= (1u128 << (x - 128)) - 1,
156            _ => {}
157        };
158
159        // init U256 from 2 randoms
160        let mut inner: [u64; 4] = [0; 4];
161        let mask64 = (1 << 65) - 1;
162        inner[0] = (lower & mask64) as u64;
163        inner[1] = (lower >> 64) as u64;
164        inner[2] = (higher & mask64) as u64;
165        inner[3] = (higher >> 64) as u64;
166        let start: U256 = U256::from_limbs(inner);
167
168        Ok(UintValueTree::new(start, false))
169    }
170
171    fn type_max(&self) -> U256 {
172        if self.bits < 256 { (U256::from(1) << self.bits) - U256::from(1) } else { U256::MAX }
173    }
174}
175
176impl Strategy for UintStrategy {
177    type Tree = UintValueTree;
178    type Value = U256;
179    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
180        let total_weight = self.random_weight + self.fixtures_weight + self.edge_weight;
181        let bias = runner.rng().random_range(0..total_weight);
182        // randomly select one of 3 strategies
183        match bias {
184            x if x < self.edge_weight => self.generate_edge_tree(runner),
185            x if x < self.edge_weight + self.fixtures_weight => self.generate_fixtures_tree(runner),
186            _ => self.generate_random_tree(runner),
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use crate::strategies::uint::UintValueTree;
194    use alloy_primitives::U256;
195    use proptest::strategy::ValueTree;
196
197    #[test]
198    fn test_uint_tree_complicate_max() {
199        let mut uint_tree = UintValueTree::new(U256::MAX, false);
200        assert_eq!(uint_tree.hi, U256::MAX);
201        assert_eq!(uint_tree.curr, U256::MAX);
202        uint_tree.complicate();
203        assert_eq!(uint_tree.lo, U256::MIN);
204    }
205}