feat: storage layout + read/write slots + mut keyword

contract-derive/src/lib.rs

@@ -1,5 +1,6 @@
 extern crate proc_macro;
-use alloy_core::primitives::keccak256;
+use alloy_core::primitives::{keccak256, U256};
+use alloy_sol_types::SolValue;
 use proc_macro::TokenStream;
 use quote::{format_ident, quote};
 use syn::{
@@ -257,11 +258,11 @@ pub fn contract(_attr: TokenStream, item: TokenStream) -> TokenStream {
             #emit_helper
 
             impl Contract for #struct_name {
-                fn call(&self) {
+                fn call(&mut self) {
                     self.call_with_data(&msg_data());
                 }
 
-                fn call_with_data(&self, calldata: &[u8]) {
+                fn call_with_data(&mut self, calldata: &[u8]) {
                     let selector = u32::from_be_bytes([calldata[0], calldata[1], calldata[2], calldata[3]]);
                     let calldata = &calldata[4..];
 
@@ -276,7 +277,7 @@ pub fn contract(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
             #[eth_riscv_runtime::entry]
             fn main() -> ! {
-                let contract = #struct_name::default();
+                let mut contract = #struct_name::default();
                 contract.call();
                 eth_riscv_runtime::return_riscv(0, 0)
             }
@@ -342,3 +343,52 @@ pub fn interface(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
     TokenStream::from(output)
 }
+
+#[proc_macro_attribute]
+pub fn storage(_attr: TokenStream, input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    let name = &input.ident;
+    let vis = &input.vis;
+
+    let fields = match &input.data {
+        Data::Struct(data) => match &data.fields {
+            Fields::Named(fields) => &fields.named,
+            _ => {
+                let output = quote! {
+                    #vis struct #name;
+                    impl #name { pub fn new() -> Self { Self {} } }
+                };
+                return TokenStream::from(output);
+            }
+        },
+        _ => panic!("Storage derive only works on structs"),
+    };
+
+    // Generate the struct definition with the same fields
+    let struct_fields = fields.iter().map(|f| {
+        let name = &f.ident;
+        let ty = &f.ty;
+        quote! { pub #name: #ty }
+    });
+
+    // Generate initialization code for each field
+    // TODO: PoC uses a naive strategy. Enhance to support complex types like tuples or custom structs.
+    let init_fields = fields.iter().enumerate().map(|(i, f)| {
+        let name = &f.ident;
+        let slot = U256::from(i);
+        let [limb0, limb1, limb2, limb3] = slot.as_limbs();
+        quote! { #name: StorageLayout::allocate(#limb0, #limb1, #limb2, #limb3) }
+    });
+
+    let expanded = quote! {
+        #vis struct #name { #(#struct_fields,)* }
+
+        impl #name {
+            pub fn default() -> Self {
+                Self { #(#init_fields,)* }
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}

erc20/src/lib.rs

@@ -3,22 +3,22 @@
 
 use core::default::Default;
 
-use contract_derive::{contract, payable, Event};
-use eth_riscv_runtime::types::Mapping;
+use contract_derive::{contract, payable, storage, Event};
+use eth_riscv_runtime::types::{Mapping, Slot, StorageLayout};
 
 use alloy_core::primitives::{address, Address, U256};
 
 extern crate alloc;
 use alloc::string::String;
 
-#[derive(Default)]
+#[storage]
 pub struct ERC20 {
-    balances: Mapping<Address, u64>,
-    allowances: Mapping<Address, Mapping<Address, u64>>,
-    total_supply: U256,
-    name: String,
-    symbol: String,
-    decimals: u8,
+    total_supply: Slot<U256>,
+    balances: Mapping<Address, U256>,
+    allowances: Mapping<Address, Mapping<Address, U256>>,
+    // name: String,
+    // symbol: String,
+    // decimals: u8,
 }
 
 #[derive(Event)]
@@ -27,7 +27,7 @@ pub struct Transfer {
     pub from: Address,
     #[indexed]
     pub to: Address,
-    pub value: u64,
+    pub value: U256,
 }
 
 #[derive(Event)]
@@ -36,16 +36,39 @@ pub struct Mint {
     pub from: Address,
     #[indexed]
     pub to: Address,
-    pub value: u64,
+    pub value: U256,
 }
 
 #[contract]
 impl ERC20 {
-    pub fn balance_of(&self, owner: Address) -> u64 {
-        self.balances.read(owner)
+    // -- STATE MODIFYING FUNCTIONS -------------------------------------------
+    #[payable]
+    pub fn mint(&mut self, to: Address, value: U256) -> bool {
+        // TODO: implement constructors and store contract owner
+        let _owner = msg_sender();
+
+        // increase user balance
+        let to_balance = self.balances.read(to);
+        self.balances.write(to, to_balance + value);
+        log::emit(Transfer::new(
+            address!("0000000000000000000000000000000000000000"),
+            to,
+            value,
+        ));
+
+        // increase total supply
+        self.total_supply += value;
+
+        true
     }
 
-    pub fn transfer(&self, to: Address, value: u64) -> bool {
+    pub fn approve(&mut self, spender: Address, value: U256) -> bool {
+        let mut spender_allowances = self.allowances.read(msg_sender());
+        spender_allowances.write(spender, value);
+        true
+    }
+
+    pub fn transfer(&mut self, to: Address, value: U256) -> bool {
         let from = msg_sender();
         let from_balance = self.balances.read(from);
         let to_balance = self.balances.read(to);
@@ -61,13 +84,7 @@ impl ERC20 {
         true
     }
 
-    pub fn approve(&self, spender: Address, value: u64) -> bool {
-        let spender_allowances = self.allowances.read(msg_sender());
-        spender_allowances.write(spender, value);
-        true
-    }
-
-    pub fn transfer_from(&self, sender: Address, recipient: Address, amount: u64) -> bool {
+    pub fn transfer_from(&mut self, sender: Address, recipient: Address, amount: U256) -> bool {
         let allowance = self.allowances.read(sender).read(msg_sender());
         let sender_balance = self.balances.read(sender);
         let recipient_balance = self.balances.read(recipient);
@@ -81,25 +98,17 @@ impl ERC20 {
         true
     }
 
+    // -- GETTER FUNCTIONS ----------------------------------------------------
+
     pub fn total_supply(&self) -> U256 {
-        self.total_supply
+        self.total_supply.read()
     }
 
-    pub fn allowance(&self, owner: Address, spender: Address) -> u64 {
-        self.allowances.read(owner).read(spender)
+    pub fn balance_of(&self, owner: Address) -> U256 {
+        self.balances.read(owner)
     }
 
-    #[payable]
-    pub fn mint(&self, to: Address, value: u64) -> bool {
-        let owner = msg_sender();
-
-        let to_balance = self.balances.read(to);
-        self.balances.write(to, to_balance + value);
-        log::emit(Transfer::new(
-            address!("0000000000000000000000000000000000000000"),
-            to,
-            value,
-        ));
-        true
+    pub fn allowance(&self, owner: Address, spender: Address) -> U256 {
+        self.allowances.read(owner).read(spender)
     }
 }

erc20x/src/lib.rs

@@ -15,7 +15,7 @@ pub struct ERC20x;
 
 #[contract]
 impl ERC20x {
-    pub fn x_balance_of(&self, owner: Address, target: Address) -> u64 {
+    pub fn x_balance_of(&self, owner: Address, target: Address) -> U256 {
         let token = IERC20::new(target);
         match token.balance_of(owner) {
             Some(balance) => balance,

erc20x_standalone/src/lib.rs

@@ -14,12 +14,12 @@ pub struct ERC20x;
 
 #[interface]
 trait IERC20 {
-    fn balance_of(&self, owner: Address) -> u64;
+    fn balance_of(&self, owner: Address) -> U256;
 }
 
 #[contract]
 impl ERC20x {
-    pub fn x_balance_of(&self, owner: Address, target: Address) -> u64 {
+    pub fn x_balance_of(&self, owner: Address, target: Address) -> U256 {
         let token = IERC20::new(target);
         match token.balance_of(owner) {
             Some(balance) => balance,

eth-riscv-runtime/src/lib.rs

@@ -10,8 +10,8 @@ pub use riscv_rt::entry;
 
 mod alloc;
 pub mod block;
-pub mod tx;
 pub mod types;
+pub mod tx;
 
 pub mod log;
 pub use log::{emit_log, Event};
@@ -22,8 +22,8 @@ pub use call::call_contract;
 const CALLDATA_ADDRESS: usize = 0x8000_0000;
 
 pub trait Contract {
-    fn call(&self);
-    fn call_with_data(&self, calldata: &[u8]);
+    fn call(&mut self);
+    fn call_with_data(&mut self, calldata: &[u8]);
 }
 
 pub unsafe fn slice_from_raw_parts(address: usize, length: usize) -> &'static [u8] {
@@ -56,21 +56,30 @@ pub fn return_riscv(addr: u64, offset: u64) -> ! {
     unreachable!()
 }
 
-pub fn sload(key: u64) -> U256 {
-    let first: u64;
-    let second: u64;
-    let third: u64;
-    let fourth: u64;
+pub fn sload(key: U256) -> U256 {
+    let key = key.as_limbs();
+    let (val0, val1, val2, val3): (u64, u64, u64, u64);
     unsafe {
-        asm!("ecall", lateout("a0") first, lateout("a1") second, lateout("a2") third, lateout("a3") fourth, in("a0") key, in("t0") u8::from(Syscall::SLoad));
+        asm!(
+            "ecall",
+            lateout("a0") val0, lateout("a1") val1, lateout("a2") val2, lateout("a3") val3,
+            in("a0") key[0], in("a1") key[1], in("a2") key[2], in("a3") key[3],
+            in("t0") u8::from(Syscall::SLoad));
     }
-    U256::from_limbs([first, second, third, fourth])
+    U256::from_limbs([val0, val1, val2, val3])
 }
 
-pub fn sstore(key: u64, value: U256) {
-    let limbs = value.as_limbs();
+pub fn sstore(key: U256, value: U256) {
+    let key = key.as_limbs();
+    let value = value.as_limbs();
+
     unsafe {
-        asm!("ecall", in("a0") key, in("a1") limbs[0], in("a2") limbs[1], in("a3") limbs[2], in("a4") limbs[3], in("t0") u8::from(Syscall::SStore));
+        asm!(
+            "ecall",
+            in("a0") key[0], in("a1") key[1], in("a2") key[2], in("a3") key[3],
+            in("a4") value[0], in("a5") value[1], in("a6") value[2], in("a7") value[3],
+            in("t0") u8::from(Syscall::SStore)
+        );
     }
 }
 
@@ -107,12 +116,8 @@ pub fn revert() -> ! {
     unreachable!()
 }
 
-pub fn keccak256(offset: u64, size: u64) -> B256 {
-    let first: u64;
-    let second: u64;
-    let third: u64;
-    let fourth: u64;
-
+pub fn keccak256(offset: u64, size: u64) -> U256 {
+    let (first, second, third, fourth): (u64, u64, u64, u64);
     unsafe {
         asm!(
             "ecall",
@@ -125,21 +130,11 @@ pub fn keccak256(offset: u64, size: u64) -> B256 {
             in("t0") u8::from(Syscall::Keccak256)
         );
     }
-
-    let mut bytes = [0u8; 32];
-
-    bytes[0..8].copy_from_slice(&first.to_be_bytes());
-    bytes[8..16].copy_from_slice(&second.to_be_bytes());
-    bytes[16..24].copy_from_slice(&third.to_be_bytes());
-    bytes[24..32].copy_from_slice(&fourth.to_be_bytes());
-
-    B256::from_slice(&bytes)
+    U256::from_limbs([first, second, third, fourth])
 }
 
 pub fn msg_sender() -> Address {
-    let first: u64;
-    let second: u64;
-    let third: u64;
+    let (first, second, third): (u64, u64, u64);
     unsafe {
         asm!("ecall", lateout("a0") first, lateout("a1") second, lateout("a2") third, in("t0") u8::from(Syscall::Caller));
     }
@@ -151,10 +146,7 @@ pub fn msg_sender() -> Address {
 }
 
 pub fn msg_value() -> U256 {
-    let first: u64;
-    let second: u64;
-    let third: u64;
-    let fourth: u64;
+    let (first, second, third, fourth): (u64, u64, u64, u64);
     unsafe {
         asm!("ecall", lateout("a0") first, lateout("a1") second, lateout("a2") third, lateout("a3") fourth, in("t0") u8::from(Syscall::CallValue));
     }