Skip to main content

evm_fork_cache/
mapping_probe.rs

1//! Trace-based discovery of hash-derived storage slots.
2//!
3//! Solidity, Vyper, and hand-written assembly contracts all place `mapping`
4//! entries (and dynamic arrays) at storage slots derived with `keccak256`. The
5//! canonical Solidity layout for `mapping(K => V)` entry `m[k]` is
6//! `keccak256(k ‖ slot)`; Vyper hashes the same two words in the opposite order
7//! (`keccak256(slot ‖ k)`); Solady-style assembly packs a key and a seed into a
8//! single 32-byte word. That diversity makes a *static* base-slot guess
9//! unreliable across tokens.
10//!
11//! This module derives the layout **dynamically** from a single simulated call.
12//! [`HashStorageProbe`] is a [`revm::Inspector`] that records, over one
13//! execution:
14//!   * every `KECCAK256` preimage, keyed by its hash output, and
15//!   * every `SLOAD` — its slot and the value it loaded.
16//!
17//! [`HashStorageProbe::accesses`] then factors each *hashed* `SLOAD` back into a
18//! [`HashSlotAccess`] — the mapping key chain, the declared base slot, the
19//! detected [`SlotLayout`], and the exact storage slot that was read — **without
20//! assuming a preimage byte order**. See [`resolve`](HashStorageProbe::accesses)
21//! for the disambiguation rules.
22//!
23//! A discovered single-level mapping can be captured as a [`TrackedMapping`],
24//! a small reusable descriptor whose [`TrackedMapping::slot_for`] recomputes the
25//! exact storage slot for *any* key using the discovered layout. That is the
26//! building block for "derive a token's balance slot once, then track these
27//! addresses" workflows: discover with the probe, then fan out cheaply.
28//!
29//! # Coupling
30//!
31//! This module depends only on `revm` and `alloy` primitives — it is decoupled
32//! from balances, allowances, or any ERC-20 semantics. [`EvmCache`] builds
33//! ergonomic wrappers on top (balance-slot discovery, layout-aware writes); the
34//! reactive freshness/prefetch layers consume [`TrackedMapping::slot_for`] to
35//! register derived slots.
36//!
37//! [`EvmCache`]: crate::cache::EvmCache
38
39use std::collections::HashMap;
40use std::fmt;
41
42use alloy_primitives::{Address, B256, U256, keccak256};
43use revm::Inspector;
44use revm::interpreter::interpreter_types::{Jumps, MemoryTr, StackTr};
45use revm::interpreter::{Interpreter, InterpreterTypes};
46
47/// `KECCAK256` (a.k.a. `SHA3`) opcode.
48const OP_KECCAK256: u8 = 0x20;
49/// `SLOAD` opcode.
50const OP_SLOAD: u8 = 0x54;
51/// Upper bound on a hashed preimage we will record. Mapping/array preimages are
52/// 32 or 64 bytes; this guards against hashing large memory regions.
53const MAX_PREIMAGE_LEN: usize = 4096;
54
55// ===========================================================================
56// Public result types
57// ===========================================================================
58
59/// The storage layout a [`HashSlotAccess`] was factored into.
60#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub enum SlotLayout {
62    /// `keccak256(key ‖ slot)` — Solidity's `mapping` layout.
63    SolidityMapping,
64    /// `keccak256(slot ‖ key)` — Vyper's `HashMap` layout (words swapped).
65    VyperMapping,
66    /// A nested mapping (2+ chained hashes), e.g. an ERC-20 allowance.
67    Nested,
68    /// `keccak256(addr ‖ seed)` packed into a single 32-byte word
69    /// (Solady/assembly), where `seed` is a per-mapping constant.
70    PackedSeed {
71        /// The packed low-word seed (the mapping's identifier).
72        seed: U256,
73    },
74    /// `keccak256(slot)` — a dynamic-array or base-pointer slot.
75    ArrayPointer,
76    /// Recognized as hash-derived but not factorable into a known shape.
77    Opaque,
78}
79
80impl fmt::Display for SlotLayout {
81    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82        match self {
83            SlotLayout::SolidityMapping => f.write_str("mapping(key‖slot) [Solidity]"),
84            SlotLayout::VyperMapping => f.write_str("mapping(slot‖key) [Vyper]"),
85            SlotLayout::Nested => f.write_str("nested mapping"),
86            SlotLayout::PackedSeed { seed } => write!(f, "packed(addr‖seed={seed:#x}) [Solady]"),
87            SlotLayout::ArrayPointer => f.write_str("array/base pointer"),
88            SlotLayout::Opaque => f.write_str("opaque"),
89        }
90    }
91}
92
93/// Confidence in a [`HashSlotAccess`] factoring, ordered weakest → strongest so
94/// `min` degrades correctly as evidence weakens.
95#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
96pub enum Confidence {
97    /// Ambiguous (e.g. both preimage halves look like keys).
98    Low,
99    /// A 32-byte packed/array fallback with no known-key anchor.
100    Heuristic,
101    /// Resolved by significant-byte magnitude (a small slot vs a wide key).
102    Medium,
103    /// A key half matched a caller-supplied known key.
104    High,
105}
106
107impl fmt::Display for Confidence {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        write!(f, "{self:?}")
110    }
111}
112
113/// One hash-derived `SLOAD` observed during a simulation, factored into its
114/// mapping structure.
115#[derive(Clone, Debug, PartialEq, Eq)]
116pub struct HashSlotAccess {
117    /// The exact storage slot that was read — the direct override/track target.
118    pub slot: B256,
119    /// The value the `SLOAD` returned.
120    pub value: U256,
121    /// Mapping keys, outer → inner. For `m[a][b]` this is `[b, a]`; empty for a
122    /// plain array/base pointer.
123    pub keys: Vec<B256>,
124    /// The inferred declared base slot (or packed seed) of the outer mapping.
125    pub base_slot: U256,
126    /// The detected storage layout.
127    pub layout: SlotLayout,
128    /// Nesting depth (number of hashes in the derivation chain).
129    pub depth: usize,
130    /// Confidence in the factoring.
131    pub confidence: Confidence,
132}
133
134impl HashSlotAccess {
135    /// True if any key in the chain equals `k` (matched as a full 32-byte word
136    /// or as an address in the word's low 20 bytes).
137    pub fn keyed_by(&self, k: B256) -> bool {
138        self.keys.iter().any(|key| word_matches(*key, k))
139    }
140
141    /// Capture a **single-level** mapping access as a reusable [`TrackedMapping`]
142    /// on `contract`. Returns `None` for nested, array, or opaque layouts (their
143    /// slot cannot be recomputed from a base slot and one key alone).
144    pub fn as_tracked(&self, contract: Address) -> Option<TrackedMapping> {
145        if self.keys.len() != 1 {
146            return None;
147        }
148        match self.layout {
149            SlotLayout::SolidityMapping
150            | SlotLayout::VyperMapping
151            | SlotLayout::PackedSeed { .. } => Some(TrackedMapping {
152                contract,
153                base_slot: self.base_slot,
154                layout: self.layout,
155            }),
156            _ => None,
157        }
158    }
159}
160
161/// A reusable descriptor for a single-level hash-derived mapping on one
162/// contract, from which the storage slot of any key can be recomputed.
163///
164/// Obtain one from [`HashSlotAccess::as_tracked`] (or
165/// [`EvmCache::discover_erc20_balance_slot`](crate::cache::EvmCache::discover_erc20_balance_slot)),
166/// then call [`slot_for`](Self::slot_for) for each key you want to track — no
167/// re-simulation required.
168#[derive(Clone, Copy, Debug, PartialEq, Eq)]
169pub struct TrackedMapping {
170    /// The contract whose storage holds the mapping.
171    pub contract: Address,
172    /// The mapping's declared base slot (or packed seed).
173    pub base_slot: U256,
174    /// The layout used to derive entry slots.
175    pub layout: SlotLayout,
176}
177
178impl TrackedMapping {
179    /// Build a descriptor explicitly (when the layout is already known).
180    pub fn new(contract: Address, base_slot: U256, layout: SlotLayout) -> Self {
181        Self {
182            contract,
183            base_slot,
184            layout,
185        }
186    }
187
188    /// The storage slot of the entry keyed by `key`, using the tracked layout.
189    ///
190    /// `key` is a full 32-byte word; for an address key pass
191    /// [`Address::into_word`]. Returns `None` if the layout is not a
192    /// single-level mapping shape.
193    pub fn slot_for(&self, key: B256) -> Option<B256> {
194        match self.layout {
195            SlotLayout::SolidityMapping => {
196                let mut pre = [0u8; 64];
197                pre[0..32].copy_from_slice(key.as_slice());
198                pre[32..64].copy_from_slice(self.base_slot_word().as_slice());
199                Some(keccak256(pre))
200            }
201            SlotLayout::VyperMapping => {
202                let mut pre = [0u8; 64];
203                pre[0..32].copy_from_slice(self.base_slot_word().as_slice());
204                pre[32..64].copy_from_slice(key.as_slice());
205                Some(keccak256(pre))
206            }
207            SlotLayout::PackedSeed { seed } => {
208                // Solady packs the address in the high 20 bytes and the seed in
209                // the low 12 bytes of a single word.
210                let mut pre = [0u8; 32];
211                pre[0..20].copy_from_slice(&Address::from_word(key).into_array());
212                let seed_be = seed.to_be_bytes::<32>();
213                pre[20..32].copy_from_slice(&seed_be[20..32]);
214                Some(keccak256(pre))
215            }
216            _ => None,
217        }
218    }
219
220    /// Compute `(key, slot)` for each key, skipping any the layout can't derive.
221    pub fn slots_for(&self, keys: impl IntoIterator<Item = B256>) -> Vec<(B256, B256)> {
222        keys.into_iter()
223            .filter_map(|k| self.slot_for(k).map(|s| (k, s)))
224            .collect()
225    }
226
227    fn base_slot_word(&self) -> B256 {
228        B256::from(self.base_slot.to_be_bytes::<32>())
229    }
230}
231
232/// A discovered balance mapping paired with each tracked holder's
233/// `(address, storage slot)` — the return of
234/// [`EvmCache::track_erc20_balances`](crate::cache::EvmCache::track_erc20_balances).
235pub type TrackedBalances = (TrackedMapping, Vec<(Address, B256)>);
236
237// ===========================================================================
238// The inspector
239// ===========================================================================
240
241#[derive(Clone, Debug)]
242struct SloadRecord {
243    slot: B256,
244    value: U256,
245}
246
247/// A [`revm::Inspector`] that records `KECCAK256` preimages and **every** `SLOAD`
248/// (slot + loaded value), so [`accesses`](Self::accesses) can reconstruct the
249/// mapping layout of hashed reads and [`slots_returning`](Self::slots_returning)
250/// can find the slot that drove a getter's return value (hashed *or* plain).
251///
252/// Attach it through
253/// [`EvmCache::call_raw_with_inspector`](crate::cache::EvmCache::call_raw_with_inspector)
254/// or [`EvmOverlay::call_raw_with_inspector`](crate::cache::EvmOverlay::call_raw_with_inspector),
255/// and it composes with other inspectors via
256/// [`InspectorStack`](crate::InspectorStack) so discovery can piggyback on a
257/// simulation you are already running.
258///
259/// ```
260/// use evm_fork_cache::mapping_probe::HashStorageProbe;
261/// let probe = HashStorageProbe::new();
262/// assert!(probe.accesses(&[]).is_empty()); // nothing executed yet
263/// ```
264#[derive(Clone, Debug, Default)]
265pub struct HashStorageProbe {
266    /// `keccak256(preimage) -> preimage`, every hash computed during the call.
267    preimages: HashMap<B256, Vec<u8>>,
268    /// Every `SLOAD` (slot + loaded value), in observation order.
269    reads: Vec<SloadRecord>,
270    /// Slot set on a `SLOAD` `step`, resolved to a value in the next `step_end`.
271    pending: Option<B256>,
272}
273
274impl HashStorageProbe {
275    /// Create an empty probe.
276    pub fn new() -> Self {
277        Self::default()
278    }
279
280    /// Number of distinct `KECCAK256` preimages observed.
281    pub fn preimage_count(&self) -> usize {
282        self.preimages.len()
283    }
284
285    /// Number of hash-derived `SLOAD`s observed (a subset of all reads).
286    pub fn hashed_read_count(&self) -> usize {
287        self.reads
288            .iter()
289            .filter(|r| self.preimages.contains_key(&r.slot))
290            .count()
291    }
292
293    /// Storage slots whose loaded value equalled `value`, deduplicated in
294    /// observation order.
295    ///
296    /// This is the building block for mocking a getter's return: the slot a view
297    /// call read that equals what it returned is (almost always) the slot that
298    /// *drives* the return. Includes plain (non-hashed) slots, so it covers
299    /// `totalSupply`-style variables as well as mapping entries.
300    pub fn slots_returning(&self, value: U256) -> Vec<B256> {
301        let mut seen = std::collections::HashSet::new();
302        self.reads
303            .iter()
304            .filter(|r| r.value == value)
305            .map(|r| r.slot)
306            .filter(|slot| seen.insert(*slot))
307            .collect()
308    }
309
310    /// Resolve every hash-derived `SLOAD` into a [`HashSlotAccess`].
311    ///
312    /// `known` are words (e.g. addresses via [`Address::into_word`]) the caller
313    /// wants matched as mapping keys; pass `&[]` to rely purely on the
314    /// significant-byte magnitude heuristic. Results are in observation order.
315    ///
316    /// # Disambiguation
317    ///
318    /// * **64-byte preimage `X ‖ Y`** — a mapping entry. If one half is itself a
319    ///   recorded hash it is the parent location (nested mapping) and the other
320    ///   is the key; recurse. Otherwise the key is the half matching a `known`
321    ///   word, else the half with more significant bytes (a 20-byte address
322    ///   dwarfs a small slot index); the remaining half is the base slot.
323    /// * **32-byte preimage** — Solady-style `addr ‖ seed` (detected when the
324    ///   high 20 bytes match a `known` address) or a `keccak(slot)` array
325    ///   pointer.
326    pub fn accesses(&self, known: &[B256]) -> Vec<HashSlotAccess> {
327        self.reads
328            .iter()
329            .filter(|r| self.preimages.contains_key(&r.slot))
330            .map(|r| resolve(r.slot, r.value, &self.preimages, known))
331            .collect()
332    }
333}
334
335impl<CTX, INTR: InterpreterTypes> Inspector<CTX, INTR> for HashStorageProbe {
336    fn step(&mut self, interp: &mut Interpreter<INTR>, _ctx: &mut CTX) {
337        match interp.bytecode.opcode() {
338            OP_KECCAK256 => {
339                // Stack (top first): offset, size.
340                let (offset, size) = {
341                    let s = interp.stack.data();
342                    let n = s.len();
343                    if n < 2 {
344                        return;
345                    }
346                    (s[n - 1], s[n - 2])
347                };
348                let (Some(offset), Some(size)) = (to_usize(offset), to_usize(size)) else {
349                    return;
350                };
351                if size == 0 || size > MAX_PREIMAGE_LEN {
352                    return;
353                }
354                let preimage = read_mem(interp, offset, size);
355                self.preimages.insert(keccak256(&preimage), preimage);
356            }
357            OP_SLOAD => {
358                // Record every SLOAD; `accesses` later filters to hashed slots,
359                // while `slots_returning` uses the full set for value-matching.
360                if let Some(k) = interp.stack.data().last() {
361                    self.pending = Some(word_from(*k));
362                }
363            }
364            _ => {}
365        }
366    }
367
368    fn step_end(&mut self, interp: &mut Interpreter<INTR>, _ctx: &mut CTX) {
369        if let Some(slot) = self.pending.take() {
370            // The SLOAD has executed; its result is now on top of the stack.
371            if let Some(value) = interp.stack.data().last().copied() {
372                self.reads.push(SloadRecord { slot, value });
373            }
374        }
375    }
376}
377
378// ===========================================================================
379// Resolver internals
380// ===========================================================================
381
382/// Factor a hashed `SLOAD` slot into its key chain, base slot, and layout.
383fn resolve(
384    slot: B256,
385    value: U256,
386    pre: &HashMap<B256, Vec<u8>>,
387    known: &[B256],
388) -> HashSlotAccess {
389    let mut keys: Vec<B256> = Vec::new();
390    let mut confidence = Confidence::High;
391    let mut cur = slot;
392
393    let (base_slot, layout) = loop {
394        let Some(p) = pre.get(&cur) else {
395            // Not a recorded preimage: `cur` is a literal base slot (or a hash
396            // we didn't observe).
397            let l = if keys.is_empty() {
398                SlotLayout::Opaque
399            } else {
400                SlotLayout::Nested
401            };
402            break (U256::from_be_slice(cur.as_slice()), l);
403        };
404
405        match p.len() {
406            64 => {
407                let aw = B256::from_slice(&p[0..32]);
408                let bw = B256::from_slice(&p[32..64]);
409                let a_parent = pre.contains_key(&aw);
410                let b_parent = pre.contains_key(&bw);
411                if a_parent ^ b_parent {
412                    // The known-hash half is the parent location; recurse.
413                    let (parent, key) = if a_parent { (aw, bw) } else { (bw, aw) };
414                    keys.push(key);
415                    cur = parent;
416                    continue;
417                }
418                let (slot_word, key_word, key_first, conf) =
419                    split_key_slot(&p[0..32], &p[32..64], known);
420                confidence = confidence.min(conf);
421                keys.push(key_word);
422                let layout = if keys.len() > 1 {
423                    SlotLayout::Nested
424                } else if key_first {
425                    SlotLayout::SolidityMapping
426                } else {
427                    SlotLayout::VyperMapping
428                };
429                break (U256::from_be_slice(slot_word.as_slice()), layout);
430            }
431            32 => {
432                // Solady-style packed: address in high 20 bytes + a low seed.
433                let hi = Address::from_slice(&p[0..20]);
434                if hi != Address::ZERO && known.iter().any(|k| Address::from_word(*k) == hi) {
435                    keys.push(hi.into_word());
436                    confidence = confidence.min(Confidence::Medium);
437                    let seed = U256::from_be_slice(&p[20..32]);
438                    break (seed, SlotLayout::PackedSeed { seed });
439                }
440                // Otherwise a keccak(slot) array/base pointer.
441                confidence = confidence.min(Confidence::Heuristic);
442                let l = if keys.is_empty() {
443                    SlotLayout::ArrayPointer
444                } else {
445                    SlotLayout::Nested
446                };
447                break (U256::from_be_slice(&p[0..32]), l);
448            }
449            _ => {
450                confidence = Confidence::Low;
451                break (U256::from_be_slice(cur.as_slice()), SlotLayout::Opaque);
452            }
453        }
454    };
455
456    let depth = keys.len().max(1);
457    HashSlotAccess {
458        slot,
459        value,
460        keys,
461        base_slot,
462        layout,
463        depth,
464        confidence,
465    }
466}
467
468/// Decide which half of a base-level 64-byte preimage is the key vs the slot.
469/// Returns `(slot_word, key_word, key_is_first, confidence)`.
470fn split_key_slot(a: &[u8], b: &[u8], known: &[B256]) -> (B256, B256, bool, Confidence) {
471    let aw = B256::from_slice(a);
472    let bw = B256::from_slice(b);
473    let a_known = known.iter().any(|k| word_matches(*k, aw));
474    let b_known = known.iter().any(|k| word_matches(*k, bw));
475    match (a_known, b_known) {
476        (true, false) => (bw, aw, true, Confidence::High), // a is key (first)
477        (false, true) => (aw, bw, false, Confidence::High), // b is key (second)
478        _ => {
479            // No unambiguous known-key anchor: the base slot is the numerically
480            // smaller word (a small slot index vs a 20-byte address/key).
481            let (sa, sb) = (sig(a), sig(b));
482            if sa < sb {
483                (aw, bw, false, Confidence::Medium)
484            } else if sb < sa {
485                (bw, aw, true, Confidence::Medium)
486            } else {
487                (aw, bw, false, Confidence::Low)
488            }
489        }
490    }
491}
492
493/// Read `len` bytes of local EVM memory at `offset`, zero-padding past the
494/// current memory size (EVM read-as-zero semantics) so we never panic.
495fn read_mem<INTR: InterpreterTypes>(
496    interp: &Interpreter<INTR>,
497    offset: usize,
498    len: usize,
499) -> Vec<u8> {
500    let mut out = vec![0u8; len];
501    let size = interp.memory.size();
502    if offset >= size || len == 0 {
503        return out;
504    }
505    let avail = (size - offset).min(len);
506    let chunk = interp.memory.slice_len(offset, avail);
507    out[..avail].copy_from_slice(&chunk[..]);
508    out
509}
510
511fn to_usize(x: U256) -> Option<usize> {
512    let limbs = x.as_limbs();
513    if limbs[1] | limbs[2] | limbs[3] != 0 {
514        return None;
515    }
516    usize::try_from(limbs[0]).ok()
517}
518
519fn word_from(x: U256) -> B256 {
520    B256::from(x.to_be_bytes::<32>())
521}
522
523/// Two words match if equal, or if they denote the same address in their low 20
524/// bytes (mapping keys are addresses stored left-padded).
525fn word_matches(a: B256, b: B256) -> bool {
526    a == b || Address::from_word(a) == Address::from_word(b)
527}
528
529/// Significant byte count = 32 minus leading zero bytes.
530fn sig(bytes: &[u8]) -> usize {
531    match bytes.iter().position(|&x| x != 0) {
532        Some(p) => bytes.len() - p,
533        None => 0,
534    }
535}
536
537// ===========================================================================
538// Tests
539// ===========================================================================
540
541#[cfg(test)]
542mod tests {
543    use super::*;
544    use alloy_primitives::address;
545
546    fn sol_slot(key: Address, base: u64) -> B256 {
547        let mut pre = [0u8; 64];
548        pre[0..32].copy_from_slice(key.into_word().as_slice());
549        pre[63] = base as u8;
550        keccak256(pre)
551    }
552
553    fn vyper_slot(base: u64, key: Address) -> B256 {
554        let mut pre = [0u8; 64];
555        pre[31] = base as u8;
556        pre[32..64].copy_from_slice(key.into_word().as_slice());
557        keccak256(pre)
558    }
559
560    fn solady_slot(key: Address, seed: u32) -> B256 {
561        let mut pre = [0u8; 32];
562        pre[0..20].copy_from_slice(&key.into_array());
563        pre[28..32].copy_from_slice(&seed.to_be_bytes());
564        keccak256(pre)
565    }
566
567    /// Build a probe with preimages seeded as if the given preimages were hashed.
568    fn probe_with(preimages: Vec<Vec<u8>>, reads: Vec<(B256, U256)>) -> HashStorageProbe {
569        let mut p = HashStorageProbe::new();
570        for pre in preimages {
571            p.preimages.insert(keccak256(&pre), pre);
572        }
573        for (slot, value) in reads {
574            p.reads.push(SloadRecord { slot, value });
575        }
576        p
577    }
578
579    #[test]
580    fn slots_returning_matches_value_and_dedups() {
581        // Two hashed reads + one plain read; slot_plain and one mapping slot
582        // share a value.
583        let key = address!("00000000000000000000000000000000000000A1");
584        let mut pre = vec![0u8; 64];
585        pre[0..32].copy_from_slice(key.into_word().as_slice());
586        pre[63] = 3;
587        let hashed_slot = keccak256(&pre);
588        let plain_slot = B256::from(U256::from(2u64).to_be_bytes::<32>()); // e.g. totalSupply
589
590        let mut probe = probe_with(vec![pre], vec![]);
591        // hashed slot read twice with value 100; plain slot read once with 100.
592        probe.reads.push(SloadRecord {
593            slot: hashed_slot,
594            value: U256::from(100u64),
595        });
596        probe.reads.push(SloadRecord {
597            slot: hashed_slot,
598            value: U256::from(100u64),
599        });
600        probe.reads.push(SloadRecord {
601            slot: plain_slot,
602            value: U256::from(100u64),
603        });
604        probe.reads.push(SloadRecord {
605            slot: plain_slot,
606            value: U256::from(7u64),
607        }); // different value
608
609        // Value 100 → the hashed slot then the plain slot (deduped, in order).
610        assert_eq!(
611            probe.slots_returning(U256::from(100u64)),
612            vec![hashed_slot, plain_slot]
613        );
614        // Plain slots participate even though they are not in `preimages`.
615        assert_eq!(probe.hashed_read_count(), 2); // both hashed reads counted; plain excluded
616        assert_eq!(probe.accesses(&[key.into_word()]).len(), 2); // only hashed reads resolve
617    }
618
619    #[test]
620    fn resolves_solidity_mapping() {
621        let key = address!("00000000000000000000000000000000000000A1");
622        let mut pre = vec![0u8; 64];
623        pre[0..32].copy_from_slice(key.into_word().as_slice());
624        pre[63] = 3;
625        let slot = keccak256(&pre);
626        let probe = probe_with(vec![pre], vec![(slot, U256::from(42u64))]);
627
628        let a = &probe.accesses(&[key.into_word()])[0];
629        assert_eq!(a.layout, SlotLayout::SolidityMapping);
630        assert_eq!(a.base_slot, U256::from(3u64));
631        assert_eq!(a.keys, vec![key.into_word()]);
632        assert_eq!(a.confidence, Confidence::High);
633    }
634
635    #[test]
636    fn resolves_vyper_order_without_known_key() {
637        // A realistic (full-width) address so the magnitude heuristic can
638        // separate the 20-byte key from the small slot index.
639        let key = address!("28C6c06298d514Db089934071355E5743bf21d60");
640        let mut pre = vec![0u8; 64];
641        pre[31] = 2;
642        pre[32..64].copy_from_slice(key.into_word().as_slice());
643        let slot = keccak256(&pre);
644        let probe = probe_with(vec![pre], vec![(slot, U256::from(1u64))]);
645
646        // No known keys: the magnitude heuristic must still find slot-first.
647        let a = &probe.accesses(&[])[0];
648        assert_eq!(a.layout, SlotLayout::VyperMapping);
649        assert_eq!(a.base_slot, U256::from(2u64));
650        assert_eq!(a.confidence, Confidence::Medium);
651    }
652
653    #[test]
654    fn resolves_solady_packed() {
655        let key = address!("00000000000000000000000000000000000000A1");
656        let seed = 0x87a2_11a2u32;
657        let mut pre = vec![0u8; 32];
658        pre[0..20].copy_from_slice(&key.into_array());
659        pre[28..32].copy_from_slice(&seed.to_be_bytes());
660        let slot = keccak256(&pre);
661        let probe = probe_with(vec![pre], vec![(slot, U256::from(7u64))]);
662
663        let a = &probe.accesses(&[key.into_word()])[0];
664        assert_eq!(
665            a.layout,
666            SlotLayout::PackedSeed {
667                seed: U256::from(seed)
668            }
669        );
670        assert!(a.keyed_by(key.into_word()));
671    }
672
673    #[test]
674    fn resolves_nested_mapping() {
675        let owner = address!("00000000000000000000000000000000000000A1");
676        let spender = address!("00000000000000000000000000000000000000B2");
677        // inner = keccak(owner ‖ 4); outer = keccak(spender ‖ inner)
678        let mut inner_pre = vec![0u8; 64];
679        inner_pre[0..32].copy_from_slice(owner.into_word().as_slice());
680        inner_pre[63] = 4;
681        let inner = keccak256(&inner_pre);
682        let mut outer_pre = vec![0u8; 64];
683        outer_pre[0..32].copy_from_slice(spender.into_word().as_slice());
684        outer_pre[32..64].copy_from_slice(inner.as_slice());
685        let outer = keccak256(&outer_pre);
686
687        let probe = probe_with(vec![inner_pre, outer_pre], vec![(outer, U256::from(9u64))]);
688        let a = &probe.accesses(&[owner.into_word(), spender.into_word()])[0];
689        assert_eq!(a.layout, SlotLayout::Nested);
690        assert_eq!(a.base_slot, U256::from(4u64));
691        assert_eq!(a.depth, 2);
692        assert_eq!(a.keys, vec![spender.into_word(), owner.into_word()]);
693    }
694
695    #[test]
696    fn tracked_mapping_round_trips_each_layout() {
697        let key = address!("00000000000000000000000000000000000000A1");
698
699        let t = TrackedMapping::new(Address::ZERO, U256::from(3u64), SlotLayout::SolidityMapping);
700        assert_eq!(t.slot_for(key.into_word()).unwrap(), sol_slot(key, 3));
701
702        let t = TrackedMapping::new(Address::ZERO, U256::from(2u64), SlotLayout::VyperMapping);
703        assert_eq!(t.slot_for(key.into_word()).unwrap(), vyper_slot(2, key));
704
705        let seed = 0x87a2_11a2u32;
706        let t = TrackedMapping::new(
707            Address::ZERO,
708            U256::from(seed),
709            SlotLayout::PackedSeed {
710                seed: U256::from(seed),
711            },
712        );
713        assert_eq!(t.slot_for(key.into_word()).unwrap(), solady_slot(key, seed));
714    }
715
716    #[test]
717    fn as_tracked_rejects_nested_and_arrays() {
718        let access = HashSlotAccess {
719            slot: B256::ZERO,
720            value: U256::ZERO,
721            keys: vec![B256::ZERO, B256::ZERO],
722            base_slot: U256::from(4u64),
723            layout: SlotLayout::Nested,
724            depth: 2,
725            confidence: Confidence::High,
726        };
727        assert!(access.as_tracked(Address::ZERO).is_none());
728    }
729}