From 6af18e4f2cb65fc9eb5a9f25499bacbad4254aad Mon Sep 17 00:00:00 2001
From: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
Date: Wed, 28 Feb 2024 15:37:43 +0200
Subject: [PATCH] fix(cast): correctly compute mapping indexes (#7261)

* fix(cast): correctly compute mapping indexes

* chore: better error msg
---
 crates/cast/src/lib.rs | 57 +++++++++++++++++++++++++++++++++++-------
 1 file changed, 48 insertions(+), 9 deletions(-)

diff --git a/crates/cast/src/lib.rs b/crates/cast/src/lib.rs
index ca6eb83bec30..d8c2e995c123 100644
--- a/crates/cast/src/lib.rs
+++ b/crates/cast/src/lib.rs
@@ -2,7 +2,7 @@ use alloy_dyn_abi::{DynSolType, DynSolValue, FunctionExt};
 use alloy_json_abi::ContractObject;
 use alloy_primitives::{
     utils::{keccak256, ParseUnits, Unit},
-    Address, Bytes, B256, I256, U256,
+    Address, Bytes, Keccak256, B256, I256, U256,
 };
 use alloy_rlp::Decodable;
 use base::{Base, NumberWithBase, ToBase};
@@ -1656,16 +1656,20 @@ impl SimpleCast {
             .collect::<Result<Vec<InterfaceSource>>>()
     }
 
-    /// Prints the slot number for the specified mapping type and input data
-    /// Uses abi_encode to pad the data to 32 bytes.
-    /// For value types v, slot number of v is keccak256(concat(h(v) , p)) where h is the padding
-    /// function and p is slot number of the mapping.
+    /// Prints the slot number for the specified mapping type and input data.
+    ///
+    /// For value types `v`, slot number of `v` is `keccak256(concat(h(v), p))` where `h` is the
+    /// padding function for `v`'s type, and `p` is slot number of the mapping.
+    ///
+    /// See [the Solidity documentation](https://docs.soliditylang.org/en/latest/internals/layout_in_storage.html#mappings-and-dynamic-arrays)
+    /// for more details.
     ///
     /// # Example
     ///
     /// ```
     /// # use cast::SimpleCast as Cast;
     ///
+    /// // Value types.
     /// assert_eq!(
     ///     Cast::index("address", "0xD0074F4E6490ae3f888d1d4f7E3E43326bD3f0f5", "2").unwrap().as_str(),
     ///     "0x9525a448a9000053a4d151336329d6563b7e80b24f8e628e95527f218e8ab5fb"
@@ -1674,13 +1678,48 @@ impl SimpleCast {
     ///     Cast::index("uint256", "42", "6").unwrap().as_str(),
     ///     "0xfc808b0f31a1e6b9cf25ff6289feae9b51017b392cc8e25620a94a38dcdafcc1"
     /// );
+    ///
+    /// // Strings and byte arrays.
+    /// assert_eq!(
+    ///     Cast::index("string", "hello", "1").unwrap().as_str(),
+    ///     "0x8404bb4d805e9ca2bd5dd5c43a107e935c8ec393caa7851b353b3192cd5379ae"
+    /// );
     /// # Ok::<_, eyre::Report>(())
     /// ```
     pub fn index(from_type: &str, from_value: &str, slot_number: &str) -> Result<String> {
-        let sig = format!("x({from_type},uint256)");
-        let encoded = Self::abi_encode(&sig, &[from_value, slot_number])?;
-        let location: String = Self::keccak(&encoded)?;
-        Ok(location)
+        let mut hasher = Keccak256::new();
+
+        let v_ty = DynSolType::parse(from_type).wrap_err("Could not parse type")?;
+        let v = v_ty.coerce_str(from_value).wrap_err("Could not parse value")?;
+        match v_ty {
+            // For value types, `h` pads the value to 32 bytes in the same way as when storing the
+            // value in memory.
+            DynSolType::Bool |
+            DynSolType::Int(_) |
+            DynSolType::Uint(_) |
+            DynSolType::FixedBytes(_) |
+            DynSolType::Address |
+            DynSolType::Function => hasher.update(v.as_word().unwrap()),
+
+            // For strings and byte arrays, `h(k)` is just the unpadded data.
+            DynSolType::String | DynSolType::Bytes => hasher.update(v.as_packed_seq().unwrap()),
+
+            DynSolType::Array(..) |
+            DynSolType::FixedArray(..) |
+            DynSolType::Tuple(..) |
+            DynSolType::CustomStruct { .. } => {
+                eyre::bail!("Type `{v_ty}` is not supported as a mapping key")
+            }
+        }
+
+        let p = DynSolType::Uint(256)
+            .coerce_str(slot_number)
+            .wrap_err("Could not parse slot number")?;
+        let p = p.as_word().unwrap();
+        hasher.update(p);
+
+        let location = hasher.finalize();
+        Ok(location.to_string())
     }
 
     /// Converts ENS names to their namehash representation