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

bal_layout/
ext.rs
