bal_layout/
ext.rs

1//! Beyond one solc layout:
2//! - **namespaced storage** (ERC-7201, Diamond): a struct laid out by another
3//!   compile mounted at a computed or explicit base slot;
4//! - **dynamic `bytes`/`string`**: short values live in their slot, long ones
5//!   at `keccak(slot) + i` — the extra slots to read and how to assemble them;
6//! - **mapping keys**: a slot can be named when a candidate key is at hand
7//!   (`keccak` is one-way, but `keccak(key ‖ base)` is one keccak per guess).
8
9use crate::{Encoding, Layout, LayoutError, Location, Result, StorageEntry, TypeInfo, Value};
10use alloy_primitives::{keccak256, Address, B256, U256};
11use std::path::Path;
12
13/// A layout manifest: a base layout plus namespaces mounted into it.
14///
15/// ```json
16/// {
17///   "base": "out/Vault.sol/Vault.json",
18///   "namespaces": [
19///     { "prefix": "erc20", "layout": "out/ERC20Storage.sol/ERC20Storage.json",
20///       "erc7201": "openzeppelin.storage.ERC20" },
21///     { "prefix": "diamond", "layout": "out/AppStorage.sol/AppStorage.json",
22///       "slot": "0x…" }
23///   ]
24/// }
25/// ```
26///
27/// Paths are relative to the manifest. The namespace layout is any solc
28/// `storageLayout` whose top-level variables are the struct's members — the
29/// usual trick is a one-line contract that declares the struct as its only
30/// state variable, or the storage-layout output of the library itself.
31#[derive(Debug, serde::Deserialize)]
32#[serde(deny_unknown_fields)]
33struct Manifest {
34    base: Option<String>,
35    #[serde(default)]
36    namespaces: Vec<Namespace>,
37}
38
39#[derive(Debug, serde::Deserialize)]
40#[serde(deny_unknown_fields)]
41struct Namespace {
42    prefix: String,
43    layout: String,
44    erc7201: Option<String>,
45    slot: Option<String>,
46}
47
48impl Layout {
49    /// ERC-7201 base slot: `keccak256(abi.encode(uint256(keccak256(id)) - 1)) & ~0xff`.
50    pub fn erc7201_slot(id: &str) -> B256 {
51        let h = U256::from_be_bytes(keccak256(id.as_bytes()).0).wrapping_sub(U256::from(1));
52        let mut out = keccak256(h.to_be_bytes::<32>()).0;
53        out[31] = 0;
54        B256::from(out)
55    }
56
57    /// Mount `other`'s top-level variables as members of a struct named
58    /// `prefix` that starts at `base`: `prefix.var`, `prefix.map[key]`, …
59    /// Types are merged by id (solc ids are canonical, so equal types agree).
60    pub fn mount(&mut self, prefix: &str, other: &Layout, base: B256) -> Result<()> {
61        if prefix.is_empty() || prefix.contains(['.', '[', ']']) {
62            return Err(LayoutError::Syntax(prefix.into()));
63        }
64        if self.storage.iter().any(|e| e.label == prefix) {
65            return Err(LayoutError::Syntax(format!("{prefix}: name already used")));
66        }
67        for (id, t) in &other.types {
68            self.types.entry(id.clone()).or_insert_with(|| t.clone());
69        }
70        let words = other
71            .storage
72            .iter()
73            .map(|e| {
74                let size = other
75                    .types
76                    .get(&e.type_id)
77                    .map(|t| t.number_of_bytes)
78                    .unwrap_or(32);
79                e.slot + U256::from(size.div_ceil(32))
80            })
81            .max()
82            .unwrap_or(U256::ZERO);
83        let type_id = format!("t_struct(namespace {prefix})_storage");
84        self.types.insert(
85            type_id.clone(),
86            TypeInfo {
87                encoding: Encoding::Inplace,
88                label: format!("struct {prefix}"),
89                number_of_bytes: words.saturating_to::<usize>().saturating_mul(32),
90                key: None,
91                value: None,
92                base: None,
93                members: Some(other.storage.clone()),
94            },
95        );
96        self.storage.push(StorageEntry {
97            label: prefix.to_string(),
98            slot: U256::from_be_bytes(base.0),
99            offset: 0,
100            type_id,
101        });
102        Ok(())
103    }
104
105    /// Build from a manifest file (see [`Manifest`]).
106    pub fn from_manifest(path: impl AsRef<Path>) -> Result<Self> {
107        let path = path.as_ref();
108        let text = std::fs::read_to_string(path)?;
109        let m: Manifest =
110            serde_json::from_str(&text).map_err(|e| LayoutError::Json(e.to_string()))?;
111        let dir = path.parent().unwrap_or(Path::new("."));
112        let mut layout = match &m.base {
113            Some(b) => Layout::from_artifact(dir.join(b))?,
114            None => Layout::from_json(r#"{"storage":[],"types":{}}"#)?,
115        };
116        for ns in &m.namespaces {
117            let other = Layout::from_artifact(dir.join(&ns.layout))?;
118            let base = match (&ns.erc7201, &ns.slot) {
119                (Some(id), None) => Layout::erc7201_slot(id),
120                (None, Some(s)) => s
121                    .parse::<B256>()
122                    .map_err(|_| LayoutError::Syntax(format!("{}: slot {s}", ns.prefix)))?,
123                _ => {
124                    return Err(LayoutError::Syntax(format!(
125                        "{}: exactly one of erc7201 / slot",
126                        ns.prefix
127                    )))
128                }
129            };
130            layout.mount(&ns.prefix, &other, base)?;
131        }
132        Ok(layout)
133    }
134
135    /// Is this location a dynamic `bytes` / `string`?
136    pub fn is_dynamic_bytes(&self, loc: &Location) -> bool {
137        self.types
138            .get(&loc.type_id)
139            .is_some_and(|t| t.encoding == Encoding::Bytes)
140    }
141
142    /// For a dynamic `bytes`/`string` whose slot holds `word`: the extra
143    /// slots holding the data, in order. Empty for a short value (≤ 31
144    /// bytes, stored in place) or for anything that is not dynamic bytes.
145    pub fn bytes_data_slots(&self, loc: &Location, word: B256) -> Vec<B256> {
146        if !self.is_dynamic_bytes(loc) || word.0[31] & 1 == 0 {
147            return Vec::new();
148        }
149        let len = (U256::from_be_bytes(word.0) - U256::from(1)) / U256::from(2);
150        // A crafted word could claim a gigantic length; cap what we ask for.
151        let words = len
152            .div_ceil(U256::from(32))
153            .min(U256::from(4096))
154            .to::<u64>();
155        let data = U256::from_be_bytes(keccak256(loc.slot.as_slice()).0);
156        (0..words)
157            .map(|i| B256::from(data.wrapping_add(U256::from(i)).to_be_bytes::<32>()))
158            .collect()
159    }
160
161    /// Assemble a dynamic `bytes`/`string` from its slot word and the data
162    /// slots from [`Layout::bytes_data_slots`] (in order). Strings that are
163    /// not valid UTF-8 come back as bytes.
164    pub fn decode_bytes(&self, loc: &Location, word: B256, chunks: &[B256]) -> Value {
165        if !self.is_dynamic_bytes(loc) {
166            return self.decode(loc, word);
167        }
168        let raw: Vec<u8> = if word.0[31] & 1 == 0 {
169            let len = (word.0[31] / 2) as usize;
170            word.0[..len.min(31)].to_vec()
171        } else {
172            let len = ((U256::from_be_bytes(word.0) - U256::from(1)) / U256::from(2))
173                .min(U256::from(chunks.len() * 32))
174                .to::<usize>();
175            let mut out = Vec::with_capacity(len);
176            for c in chunks {
177                out.extend_from_slice(c.as_slice());
178            }
179            out.truncate(len);
180            out
181        };
182        let is_string = self
183            .types
184            .get(&loc.type_id)
185            .is_some_and(|t| t.label == "string" || t.label == "string storage ref");
186        if is_string {
187            match String::from_utf8(raw) {
188                Ok(s) => Value::Str(s),
189                Err(e) => Value::Bytes(e.into_bytes()),
190            }
191        } else {
192            Value::Bytes(raw)
193        }
194    }
195
196    /// [`Layout::describe_slot`] plus mapping entries whose key is one of
197    /// `keys` (each a 32-byte word: an address left-padded, a number, …).
198    /// Two mapping levels are tried with the same candidates, and a struct
199    /// or array under a mapping is walked as usual.
200    pub fn describe_slot_with_keys(
201        &self,
202        slot: B256,
203        array_probe: u64,
204        keys: &[B256],
205    ) -> Vec<(String, Location)> {
206        let mut out = self.describe_slot(slot, array_probe);
207        if keys.is_empty() {
208            return out;
209        }
210        let target = U256::from_be_bytes(slot.0);
211        for e in &self.storage {
212            self.describe_mappings_in(
213                &e.label,
214                e.slot,
215                &e.type_id,
216                target,
217                array_probe,
218                keys,
219                0,
220                &mut out,
221            );
222        }
223        out
224    }
225
226    #[allow(clippy::too_many_arguments)]
227    fn describe_mappings_in(
228        &self,
229        name: &str,
230        base: U256,
231        type_id: &str,
232        target: U256,
233        probe: u64,
234        keys: &[B256],
235        depth: usize,
236        out: &mut Vec<(String, Location)>,
237    ) {
238        if depth > 8 {
239            return;
240        }
241        let Ok(t) = self.ty(type_id) else { return };
242        match (t.encoding, t.members.as_deref()) {
243            (Encoding::Inplace, Some(members)) => {
244                for m in members {
245                    self.describe_mappings_in(
246                        &format!("{name}.{}", m.label),
247                        base + m.slot,
248                        &m.type_id,
249                        target,
250                        probe,
251                        keys,
252                        depth + 1,
253                        out,
254                    );
255                }
256            }
257            (Encoding::Mapping, _) => {
258                let (Some(kt), Some(vt)) = (t.key.as_deref(), t.value.as_deref()) else {
259                    return;
260                };
261                let key_label = self.ty(kt).map(|k| k.label.clone()).unwrap_or_default();
262                for key in keys {
263                    let mut buf = [0u8; 64];
264                    buf[..32].copy_from_slice(key.as_slice());
265                    buf[32..].copy_from_slice(&base.to_be_bytes::<32>());
266                    let s1 = U256::from_be_bytes(keccak256(buf).0);
267                    let entry = format!("{name}[{}]", show_key(key, &key_label));
268                    let Ok(v) = self.ty(vt) else { continue };
269                    match (v.encoding, v.members.is_some(), v.base.is_some()) {
270                        (Encoding::Mapping, _, _) => {
271                            self.describe_mappings_in(
272                                &entry,
273                                s1,
274                                vt,
275                                target,
276                                probe,
277                                keys,
278                                depth + 1,
279                                out,
280                            );
281                        }
282                        (Encoding::Inplace, false, false) | (Encoding::Bytes, _, _) => {
283                            if s1 == target {
284                                out.push((
285                                    entry,
286                                    Location {
287                                        slot: slot_b(target),
288                                        offset: 0,
289                                        size: v.number_of_bytes,
290                                        type_id: vt.to_string(),
291                                    },
292                                ));
293                            }
294                        }
295                        _ => {
296                            // struct / array under the mapping: walk it
297                            self.describe_in(&entry, s1, 0, vt, target, probe, out);
298                            self.describe_mappings_in(
299                                &entry,
300                                s1,
301                                vt,
302                                target,
303                                probe,
304                                keys,
305                                depth + 1,
306                                out,
307                            );
308                        }
309                    }
310                }
311            }
312            _ => {}
313        }
314    }
315}
316
317fn slot_b(u: U256) -> B256 {
318    B256::from(u.to_be_bytes::<32>())
319}
320
321/// How a candidate key reads in a path: addresses checksummed, everything
322/// else decimal.
323fn show_key(key: &B256, key_label: &str) -> String {
324    if key_label.starts_with("address") || key_label.starts_with("contract ") {
325        Address::from_slice(&key.0[12..]).to_string()
326    } else {
327        U256::from_be_bytes(key.0).to_string()
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    #![allow(clippy::unwrap_used)]
334    use super::*;
335
336    const FIXTURE: &str = include_str!("../tests/fixtures/Playground.layout.json");
337
338    #[test]
339    fn erc7201_matches_openzeppelin() {
340        // OpenZeppelin ERC20Upgradeable: ERC20StorageLocation.
341        let s = Layout::erc7201_slot("openzeppelin.storage.ERC20");
342        assert_eq!(
343            format!("{s}"),
344            "0x52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00"
345        );
346        assert_eq!(s.0[31], 0, "low byte cleared");
347    }
348
349    #[test]
350    fn mount_namespace_and_read_through_it() {
351        let mut base = Layout::from_json(r#"{"storage":[],"types":{}}"#).unwrap();
352        let ns = Layout::from_json(FIXTURE).unwrap();
353        let at = Layout::erc7201_slot("test.playground");
354        base.mount("pg", &ns, at).unwrap();
355        let counter = base.locate("pg.counter").unwrap();
356        assert_eq!(counter.slot, at);
357        let bal = base
358            .locate("pg.balances[0x35825972e2ca90851b14576C531F13dA0B5d53ce]")
359            .unwrap();
360        assert_ne!(bal.slot, at);
361        // The same slot names back through the namespace.
362        let names = base.describe_slot(counter.slot, 16);
363        assert!(names.iter().any(|(n, _)| n == "pg.counter"), "{names:?}");
364        assert!(
365            base.mount("pg", &ns, at).is_err(),
366            "duplicate prefix refused"
367        );
368        assert!(base.mount("a.b", &ns, at).is_err(), "dots refused");
369        assert!(base.typescript("V").contains("readonly pg: {"));
370    }
371
372    #[test]
373    fn mapping_slots_are_named_from_candidate_keys() {
374        let l = Layout::from_json(FIXTURE).unwrap();
375        let user: Address = "0x35825972e2ca90851b14576C531F13dA0B5d53ce"
376            .parse()
377            .unwrap();
378        let key = B256::left_padding_from(user.as_slice());
379        let loc = l.locate(&format!("balances[{user}]")).unwrap();
380        assert!(
381            l.describe_slot(loc.slot, 16).is_empty(),
382            "one-way without a key"
383        );
384        let named = l.describe_slot_with_keys(loc.slot, 16, &[key]);
385        assert_eq!(named.len(), 1);
386        assert_eq!(named[0].0, format!("balances[{user}]"));
387        // A wrong candidate names nothing.
388        let other = B256::left_padding_from(Address::repeat_byte(9).as_slice());
389        assert!(l.describe_slot_with_keys(loc.slot, 16, &[other]).is_empty());
390    }
391
392    #[test]
393    fn dynamic_bytes_short_and_long() {
394        // A minimal layout with one `string`.
395        let l = Layout::from_json(
396            r#"{"storage":[{"label":"name","slot":"0","offset":0,"type":"t_string_storage"}],
397                "types":{"t_string_storage":{"encoding":"bytes","label":"string","numberOfBytes":"32"}}}"#,
398        )
399        .unwrap();
400        let loc = l.locate("name").unwrap();
401        assert!(l.is_dynamic_bytes(&loc));
402        // Short: "hi" → data left-aligned, last byte = len*2.
403        let mut w = [0u8; 32];
404        w[..2].copy_from_slice(b"hi");
405        w[31] = 4;
406        let word = B256::from(w);
407        assert!(l.bytes_data_slots(&loc, word).is_empty());
408        assert_eq!(l.decode_bytes(&loc, word, &[]), Value::Str("hi".into()));
409        // Long: 40 bytes → word = len*2+1, data in keccak(slot)+0..1.
410        let text = "0123456789012345678901234567890123456789";
411        let lenw = B256::from(U256::from(text.len() * 2 + 1).to_be_bytes::<32>());
412        let slots = l.bytes_data_slots(&loc, lenw);
413        assert_eq!(slots.len(), 2);
414        let data = U256::from_be_bytes(keccak256(loc.slot.as_slice()).0);
415        assert_eq!(slots[0], B256::from(data.to_be_bytes::<32>()));
416        let mut c0 = [0u8; 32];
417        c0.copy_from_slice(&text.as_bytes()[..32]);
418        let mut c1 = [0u8; 32];
419        c1[..8].copy_from_slice(&text.as_bytes()[32..]);
420        assert_eq!(
421            l.decode_bytes(&loc, lenw, &[B256::from(c0), B256::from(c1)]),
422            Value::Str(text.into())
423        );
424    }
425}
bal_layout/ext.rs

bal_layout/
ext.rs
