feat(zink): auto generate storage keys (#240)

* feat(zink): generate storage keys for non-wasm targets

* feat(zink): import keccak in non-wasm targets

* feat(zink): generate storage key for mappings

* feat(zink): auto generate storage keys for double key mappings

Cargo.toml

@@ -100,6 +100,9 @@ path = "zink/src/lib.rs"
 paste.workspace = true
 zink-codegen.workspace = true
 
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+tiny-keccak.workspace = true
+
 [target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
 anyhow.workspace = true
 filetests.workspace = true

abi/src/abi.rs

@@ -0,0 +1,77 @@
+//! Zink ABI implementation
+//!
+//! Currently just a wrapper of solidity ABI.
+
+use core::ops::{Deref, DerefMut};
+
+/// Function ABI.
+#[derive(Clone, Debug, Default)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct Abi(sol_abi::Abi);
+
+impl Deref for Abi {
+    type Target = sol_abi::Abi;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for Abi {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+#[cfg(feature = "bytes")]
+impl Abi {
+    /// Convert [`Abi`] to bytes.
+    pub fn to_bytes(&self) -> postcard::Result<Vec<u8>> {
+        postcard::to_stdvec(self).map_err(Into::into)
+    }
+
+    /// Convert bytes to [`Abi`].
+    pub fn from_bytes(bytes: impl AsRef<[u8]>) -> postcard::Result<Self> {
+        postcard::from_bytes(bytes.as_ref()).map_err(Into::into)
+    }
+}
+
+#[cfg(feature = "hex")]
+mod hex_impl {
+    use crate::{result::Result, Abi};
+    use core::fmt;
+
+    impl Abi {
+        /// Convert [`Abi`] to hex string.
+        pub fn to_hex(&self) -> Result<String> {
+            Ok("0x".to_string() + &hex::encode(self.to_bytes()?))
+        }
+
+        /// Convert hex string to [`Abi`].
+        pub fn from_hex(hex: impl AsRef<str>) -> Result<Self> {
+            Self::from_bytes(hex::decode(hex.as_ref().trim_start_matches("0x"))?)
+                .map_err(Into::into)
+        }
+    }
+
+    impl fmt::Display for Abi {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(f, "{}", self.to_hex().unwrap_or_default())
+        }
+    }
+
+    impl core::str::FromStr for Abi {
+        type Err = crate::result::Error;
+
+        fn from_str(hex: &str) -> Result<Self> {
+            Self::from_hex(hex)
+        }
+    }
+}
+
+#[cfg(feature = "syn")]
+impl From<&syn::Signature> for Abi {
+    fn from(sig: &syn::Signature) -> Self {
+        Self(sol_abi::Abi::from(sig))
+    }
+}

abi/src/lib.rs

@@ -2,79 +2,11 @@
 //!
 //! Currently just a wrapper of solidity ABI.
 
+mod abi;
 pub mod result;
 pub mod selector;
 
-use core::ops::{Deref, DerefMut};
+pub use abi::Abi;
 
-/// Function ABI.
-#[derive(Clone, Debug, Default)]
-#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub struct Abi(sol_abi::Abi);
-
-impl Deref for Abi {
-    type Target = sol_abi::Abi;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl DerefMut for Abi {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
-
-#[cfg(feature = "bytes")]
-impl Abi {
-    /// Convert [`Abi`] to bytes.
-    pub fn to_bytes(&self) -> postcard::Result<Vec<u8>> {
-        postcard::to_stdvec(self).map_err(Into::into)
-    }
-
-    /// Convert bytes to [`Abi`].
-    pub fn from_bytes(bytes: impl AsRef<[u8]>) -> postcard::Result<Self> {
-        postcard::from_bytes(bytes.as_ref()).map_err(Into::into)
-    }
-}
-
-#[cfg(feature = "hex")]
-mod hex_impl {
-    use crate::{result::Result, Abi};
-    use core::fmt;
-
-    impl Abi {
-        /// Convert [`Abi`] to hex string.
-        pub fn to_hex(&self) -> Result<String> {
-            Ok("0x".to_string() + &hex::encode(self.to_bytes()?))
-        }
-
-        /// Convert hex string to [`Abi`].
-        pub fn from_hex(hex: impl AsRef<str>) -> Result<Self> {
-            Self::from_bytes(hex::decode(hex.as_ref().trim_start_matches("0x"))?)
-                .map_err(Into::into)
-        }
-    }
-
-    impl fmt::Display for Abi {
-        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-            write!(f, "{}", self.to_hex().unwrap_or_default())
-        }
-    }
-
-    impl core::str::FromStr for Abi {
-        type Err = crate::result::Error;
-
-        fn from_str(hex: &str) -> Result<Self> {
-            Self::from_hex(hex)
-        }
-    }
-}
-
-#[cfg(feature = "syn")]
-impl From<&syn::Signature> for Abi {
-    fn from(sig: &syn::Signature) -> Self {
-        Self(sol_abi::Abi::from(sig))
-    }
-}
+#[cfg(feature = "selector")]
+pub use selector::keccak256;

examples/dkmapping.rs

@@ -41,7 +41,8 @@ fn storage_double_key_mapping() -> anyhow::Result<()> {
     assert!(info.ret.is_empty());
 
     // verify result with database
-    let storage_key = zint::keccak256(&[[0; 32], [0; 32], 1.to_bytes32()].concat());
+    let _sk = zint::keccak256(&[[0; 32], [0; 32], 1.to_bytes32()].concat());
+    let storage_key = DoubleKeyMapping::storage_key(key1, key2);
     assert_eq!(
         evm.storage(contract.address, storage_key)?,
         value.to_bytes32(),

examples/mapping.rs

@@ -39,7 +39,7 @@ fn storage_mapping() -> anyhow::Result<()> {
     assert!(info.ret.is_empty());
 
     // verify result with database
-    let storage_key = zint::keccak256(&[0; 0x40]);
+    let storage_key = Mapping::storage_key(key);
     assert_eq!(
         evm.storage(contract.address, storage_key)?,
         value.to_bytes32(),

examples/storage.rs

@@ -26,11 +26,13 @@ fn value() -> anyhow::Result<()> {
     let mut contract = Contract::search("storage")?.compile()?;
 
     {
-        let key = 0;
         let value: i32 = 42;
         let info = contract.execute(&[b"set(int32)".to_vec(), value.to_bytes32().to_vec()])?;
         assert!(info.ret.is_empty());
-        assert_eq!(info.storage.get(&U256::from(key)), Some(&U256::from(value)));
+        assert_eq!(
+            info.storage.get(&U256::from_le_bytes(Counter::STORAGE_KEY)),
+            Some(&U256::from(value))
+        );
     }
 
     {

zink/codegen/Cargo.toml

@@ -14,6 +14,7 @@ proc-macro = true
 
 [dependencies]
 heck.workspace = true
+hex.workspace = true
 proc-macro2.workspace = true
 quote.workspace = true
 syn.workspace = true

zink/codegen/src/storage.rs

@@ -2,7 +2,7 @@ extern crate proc_macro;
 
 use heck::AsSnakeCase;
 use proc_macro::TokenStream;
-use proc_macro2::{Span, TokenTree};
+use proc_macro2::{Literal, Span, TokenTree};
 use quote::quote;
 use std::{cell::RefCell, collections::HashSet};
 use syn::{
@@ -47,12 +47,19 @@ impl Storage {
         let is = &self.target;
         let name = self.target.ident.clone();
         let slot = storage_slot(name.to_string());
+        let mut key = [0; 32];
+        key[28..].copy_from_slice(&slot.to_le_bytes());
+
+        let keyl = Literal::byte_string(&key);
         let mut expanded = quote! {
             #is
 
             impl zink::storage::Storage for #name {
-                type Value = #value;
+                #[cfg(not(target_family = "wasm"))]
+                const STORAGE_KEY: [u8; 32] = *#keyl;
                 const STORAGE_SLOT: i32 = #slot;
+
+                type Value = #value;
             }
         };
 
@@ -75,6 +82,10 @@ impl Storage {
         let is = &self.target;
         let name = self.target.ident.clone();
         let slot = storage_slot(name.to_string());
+        let mut seed = [0; 64];
+        seed[29..=32].copy_from_slice(&slot.to_le_bytes());
+
+        let seedl = Literal::byte_string(&seed);
         let mut expanded = quote! {
             #is
 
@@ -83,6 +94,15 @@ impl Storage {
 
                 type Key = #key;
                 type Value = #value;
+
+                #[cfg(not(target_family = "wasm"))]
+                fn storage_key(key: Self::Key) -> [u8; 32] {
+                    use zink::Asm;
+
+                    let mut seed = *#seedl;
+                    seed[32..].copy_from_slice(&key.bytes32());
+                    zink::keccak256(&seed)
+                }
             }
         };
 
@@ -105,6 +125,10 @@ impl Storage {
         let is = &self.target;
         let name = self.target.ident.clone();
         let slot = storage_slot(name.to_string());
+        let mut seed = [0; 96];
+        seed[29..=32].copy_from_slice(&slot.to_le_bytes());
+
+        let seedl = Literal::byte_string(&seed);
         let mut expanded = quote! {
             #is
 
@@ -114,6 +138,16 @@ impl Storage {
                 type Key1 = #key1;
                 type Key2 = #key2;
                 type Value = #value;
+
+                #[cfg(not(target_family = "wasm"))]
+                fn storage_key(key1: Self::Key1, key2: Self::Key2) -> [u8; 32] {
+                    use zink::Asm;
+
+                    let mut seed = *#seedl;
+                    seed[33..=64].copy_from_slice(&key1.bytes32());
+                    seed[64..].copy_from_slice(&key2.bytes32());
+                    zink::keccak256(&seed)
+                }
             }
         };
 

zink/src/asm.rs

@@ -7,6 +7,9 @@ use paste::paste;
 pub trait Asm: Copy {
     /// Push self on the stack.
     fn push(self);
+
+    #[cfg(not(target_family = "wasm"))]
+    fn bytes32(&self) -> [u8; 32];
 }
 
 macro_rules! impl_asm {
@@ -17,6 +20,11 @@ macro_rules! impl_asm {
                     paste! { ffi::asm::[<push_ $ty>](self); }
                 }
             }
+
+            #[cfg(not(target_family = "wasm"))]
+            fn bytes32(&self) -> [u8; 32] {
+                crate::to_bytes32(&self.to_le_bytes())
+            }
         }
     };
     ($($ty:tt),+) => {

zink/src/lib.rs

@@ -2,6 +2,9 @@
 
 #![no_std]
 
+#[cfg(not(target_family = "wasm"))]
+extern crate alloc;
+
 mod asm;
 mod event;
 pub mod ffi;
@@ -12,6 +15,37 @@ pub use self::{asm::Asm, event::Event};
 pub use storage::{DoubleKeyMapping, Mapping, Storage};
 pub use zink_codegen::{external, storage, Event};
 
+/// Generate a keccak hash of the input (sha3)
+#[cfg(not(target_family = "wasm"))]
+pub fn keccak256(input: &[u8]) -> [u8; 32] {
+    use tiny_keccak::{Hasher, Keccak};
+    let mut hasher = Keccak::v256();
+    let mut output = [0; 32];
+    hasher.update(input);
+    hasher.finalize(&mut output);
+    output
+}
+
+/// Convert bytes to ls bytes
+#[cfg(not(target_family = "wasm"))]
+pub fn to_bytes32(src: &[u8]) -> [u8; 32] {
+    use alloc::vec::Vec;
+    let mut bytes = [0u8; 32];
+    let ls_bytes = {
+        src.iter()
+            .cloned()
+            .rev()
+            .skip_while(|b| *b == 0)
+            .collect::<Vec<_>>()
+            .into_iter()
+            .rev()
+            .collect::<Vec<_>>()
+    };
+
+    bytes[(32 - ls_bytes.len())..].copy_from_slice(&ls_bytes);
+    bytes
+}
+
 // Panic hook implementation
 #[cfg(target_arch = "wasm32")]
 #[panic_handler]