[WIP-RFC] Inspect EIP-7201 bucket Schema

crates/forge/src/cmd/inspect.rs

@@ -1,5 +1,5 @@
 use alloy_json_abi::{EventParam, InternalType, JsonAbi, Param};
-use alloy_primitives::{hex, keccak256};
+use alloy_primitives::{U256, hex, keccak256};
 use clap::Parser;
 use comfy_table::{Cell, Table, modifiers::UTF8_ROUND_CORNERS, presets::ASCII_MARKDOWN};
 use eyre::{Result, eyre};
@@ -108,7 +108,9 @@ impl InspectArgs {
                 print_json(&artifact.gas_estimates)?;
             }
             ContractArtifactField::StorageLayout => {
-                print_storage_layout(artifact.storage_layout.as_ref(), wrap)?;
+                let bucket_rows =
+                    parse_storage_locations(artifact.raw_metadata.as_ref()).unwrap_or_default();
+                print_storage_layout(artifact.storage_layout.as_ref(), bucket_rows, wrap)?;
             }
             ContractArtifactField::DevDoc => {
                 print_json(&artifact.devdoc)?;
@@ -302,6 +304,7 @@ fn internal_ty(ty: &InternalType) -> String {
 
 pub fn print_storage_layout(
     storage_layout: Option<&StorageLayout>,
+    bucket_rows: Vec<(String, String)>,
     should_wrap: bool,
 ) -> Result<()> {
     let Some(storage_layout) = storage_layout else {
@@ -335,6 +338,16 @@ pub fn print_storage_layout(
                     &slot.contract,
                 ]);
             }
+            for (type_str, slot_dec) in &bucket_rows {
+                table.add_row([
+                    "storage-location",
+                    type_str.as_str(),
+                    slot_dec.as_str(),
+                    "0",
+                    "32",
+                    type_str,
+                ]);
+            }
         },
         should_wrap,
     )
@@ -639,6 +652,120 @@ fn missing_error(field: &str) -> eyre::Error {
     )
 }
 
+static STORAGE_LOC_HEAD_RE: LazyLock<Regex> =
+    LazyLock::new(|| Regex::new(r"(?i)erc[0-9]+\s*:").unwrap());
+
+static STORAGE_LOC_PAIR_RE: LazyLock<Regex> = LazyLock::new(|| {
+    Regex::new(
+        r"(?ix) ^
+        (?P<formula>erc[0-9]+)   # erc ID
+        \s*:\s*
+        (?P<ns>[A-Za-z0-9_.\-]+) # namespace (no colon)
+        $",
+    )
+    .unwrap()
+});
+
+fn split_erc_formulas(s: &str) -> Vec<(String, String)> {
+    let mut starts: Vec<usize> = STORAGE_LOC_HEAD_RE.find_iter(s).map(|m| m.start()).collect();
+
+    if starts.is_empty() {
+        return Vec::new();
+    }
+    starts.push(s.len());
+    let mut out = Vec::new();
+    for w in starts.windows(2) {
+        let (beg, end) = (w[0], w[1]);
+        let slice = s[beg..end].trim();
+        if let Some(caps) = STORAGE_LOC_PAIR_RE.captures(slice) {
+            let formula = caps.name("formula").unwrap().as_str().to_string();
+            let ns = caps.name("ns").unwrap().as_str().to_string();
+            out.push((formula, ns));
+        }
+    }
+    out
+}
+
+#[inline]
+fn compute_erc7201_slot_hex(ns: &str) -> String {
+    // Step 1: keccak256(bytes(id))
+    let ns_hash = keccak256(ns.as_bytes()); // 32 bytes
+
+    // Step 2: (uint256(keccak256(id)) - 1) as 32-byte big-endian
+    let mut u = U256::from_be_slice(ns_hash.as_slice());
+    u = u.wrapping_sub(U256::from(1u8));
+    let enc = u.to_be_bytes::<32>();
+
+    // Step 3: keccak256(abi.encode(uint256(...)))
+    let slot_hash = keccak256(enc);
+
+    // Step 4: & ~0xff (zero out the lowest byte)
+    let mut slot_u = U256::from_be_slice(slot_hash.as_slice());
+    slot_u &= !U256::from(0xffu8);
+
+    // 0x-prefixed 32-byte hex, optionally shorten with your helper
+    let full = hex::encode_prefixed(slot_u.to_be_bytes::<32>());
+    short_hex(&full)
+}
+
+// Simple “formula registry” so future EIPs can be added without touching the parser.
+fn derive_slot_hex(formula: &str, ns: &str) -> Option<String> {
+    match formula.to_ascii_lowercase().as_str() {
+        "erc7201" => Some(compute_erc7201_slot_hex(ns)),
+        // For future EIPs: add "erc1234" => Some(compute_erc1234_slot_hex(ns))
+        _ => None,
+    }
+}
+
+fn strings_from_json(val: &serde_json::Value) -> Vec<String> {
+    match val {
+        serde_json::Value::String(s) => vec![s.clone()],
+        serde_json::Value::Array(arr) => {
+            arr.iter().filter_map(|v| v.as_str().map(str::to_owned)).collect()
+        }
+        _ => vec![],
+    }
+}
+
+fn get_custom_tag_lines(devdoc: &serde_json::Value, key: &str) -> Vec<String> {
+    if let Some(v) = devdoc.get(key) {
+        let xs = strings_from_json(v);
+        if !xs.is_empty() {
+            return xs;
+        }
+    }
+    devdoc
+        .get("methods")
+        .and_then(|m| m.get("constructor"))
+        .and_then(|c| c.as_object())
+        .and_then(|obj| obj.get(key))
+        .map(strings_from_json)
+        .unwrap_or_default()
+}
+
+pub fn parse_storage_locations(raw_metadata: Option<&String>) -> Option<Vec<(String, String)>> {
+    let raw = raw_metadata?;
+    let v: serde_json::Value = serde_json::from_str(raw).ok()?;
+    let devdoc = v.get("output")?.get("devdoc")?;
+
+    let loc_lines = get_custom_tag_lines(devdoc, "custom:storage-location");
+    let out: Vec<(String, String)> = loc_lines
+        .iter()
+        .flat_map(|s| split_erc_formulas(s))
+        .filter_map(|(formula, ns)| derive_slot_hex(&formula, &ns).map(|slot_hex| (ns, slot_hex)))
+        .collect();
+
+    if !out.is_empty() {
+        return Some(out);
+    }
+    if !out.is_empty() { Some(out) } else { None }
+}
+
+fn short_hex(h: &str) -> String {
+    let s = h.strip_prefix("0x").unwrap_or(h);
+    if s.len() > 12 { format!("0x{}…{}", &s[..6], &s[s.len() - 4..]) } else { format!("0x{s}") }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -675,4 +802,58 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn parses_eip7201_storage_buckets_from_metadata() {
+        let raw_wrapped = r#"
+        {
+            "metadata": {
+                "compiler": { "version": "0.8.30+commit.73712a01" },
+                "language": "Solidity",
+                "output": {
+                    "abi": [],
+                    "devdoc": {
+                        "kind": "dev",
+                        "methods": {
+                            "constructor": {
+                                "custom:storage-location": "erc7201:openzeppelin.storage.ERC20erc7201:openzeppelin.storage.AccessControlDefaultAdminRules"
+                            }
+                        },
+                        "version": 1
+                    },
+                    "userdoc": { "kind": "user", "methods": {}, "version": 1 }
+                },
+                "settings": { "optimizer": { "enabled": false, "runs": 200 } },
+                "sources": {},
+                "version": 1
+            }
+        }"#;
+
+        let v: serde_json::Value = serde_json::from_str(raw_wrapped).unwrap();
+        let inner_meta_str = v.get("metadata").unwrap().to_string();
+
+        let rows = parse_storage_locations(Some(&inner_meta_str)).expect("parser returned None");
+        assert_eq!(rows.len(), 2, "expected two EIP-7201 buckets");
+
+        assert_eq!(rows[0].0, "openzeppelin.storage.ERC20");
+        assert_eq!(rows[1].0, "openzeppelin.storage.AccessControlDefaultAdminRules");
+
+        let expect_short = |h: &str| {
+            let hex_str = h.trim_start_matches("0x");
+            let slot = U256::from_str_radix(hex_str, 16).unwrap();
+            let full = alloy_primitives::hex::encode_prefixed(slot.to_be_bytes::<32>());
+            short_hex(&full)
+        };
+
+        let eip712_slot_hex =
+            expect_short("0x52c63247e1f47db19d5ce0460030c497f067ca4cebf71ba98eeadabe20bace00");
+        let nonces_slot_hex =
+            expect_short("0xeef3dac4538c82c8ace4063ab0acd2d15cdb5883aa1dff7c2673abb3d8698400");
+
+        assert_eq!(rows[0].1, eip712_slot_hex);
+        assert_eq!(rows[1].1, nonces_slot_hex);
+
+        assert!(rows[0].1.starts_with("0x") && rows[0].1.contains('…'));
+        assert!(rows[1].1.starts_with("0x") && rows[1].1.contains('…'));
+    }
 }