From 7ac1201c03a5791eb2f8446651c103ec3ec88b99 Mon Sep 17 00:00:00 2001
From: Wodann <Wodann@users.noreply.github.com>
Date: Tue, 10 Oct 2023 19:43:28 -0400
Subject: [PATCH] feat: account storage override

---
 .../rethnet_evm_napi/src/state/overrides.rs   | 58 +++++++++----------
 .../hardhat-network/provider/vm/rethnet.ts    | 32 +++++++++-
 2 files changed, 57 insertions(+), 33 deletions(-)

diff --git a/crates/rethnet_evm_napi/src/state/overrides.rs b/crates/rethnet_evm_napi/src/state/overrides.rs
index fa1dda20338..4922b4ebd33 100644
--- a/crates/rethnet_evm_napi/src/state/overrides.rs
+++ b/crates/rethnet_evm_napi/src/state/overrides.rs
@@ -1,39 +1,10 @@
-use napi::{
-    bindgen_prelude::{BigInt, Buffer},
-    Either,
-};
+use napi::bindgen_prelude::{BigInt, Buffer};
 use napi_derive::napi;
 use rethnet_eth::{Address, Bytes};
 use rethnet_evm::HashMap;
 
 use crate::cast::TryCast;
 
-/// Type representing either a diff or full set of overrides for storage information.
-pub type StorageOverride = Either<Vec<StorageSlotChange>, Vec<StorageSlotChange>>;
-
-impl TryCast<rethnet_evm::state::StorageOverride> for StorageOverride {
-    type Error = napi::Error;
-
-    fn try_cast(self) -> napi::Result<rethnet_evm::state::StorageOverride> {
-        match self {
-            Either::A(diff) => Ok(rethnet_evm::state::StorageOverride::Diff(
-                diff.into_iter()
-                    .map(|StorageSlotChange { index, value }| {
-                        Ok((BigInt::try_cast(index)?, BigInt::try_cast(value)?))
-                    })
-                    .collect::<napi::Result<_>>()?,
-            )),
-            Either::B(full) => Ok(rethnet_evm::state::StorageOverride::Full(
-                full.into_iter()
-                    .map(|StorageSlotChange { index, value }| {
-                        Ok((BigInt::try_cast(index)?, BigInt::try_cast(value)?))
-                    })
-                    .collect::<napi::Result<_>>()?,
-            )),
-        }
-    }
-}
-
 #[napi(object)]
 pub struct StorageSlotChange {
     pub index: BigInt,
@@ -46,7 +17,8 @@ pub struct AccountOverride {
     pub balance: Option<BigInt>,
     pub nonce: Option<BigInt>,
     pub code: Option<Buffer>,
-    pub storage: Option<StorageOverride>,
+    pub storage: Option<Vec<StorageSlotChange>>,
+    pub storage_diff: Option<Vec<StorageSlotChange>>,
 }
 
 impl TryCast<rethnet_evm::state::AccountOverride> for AccountOverride {
@@ -58,7 +30,29 @@ impl TryCast<rethnet_evm::state::AccountOverride> for AccountOverride {
         let code = self
             .code
             .map(|code| rethnet_evm::Bytecode::new_raw(Bytes::copy_from_slice(code.as_ref())));
-        let storage = self.storage.map(StorageOverride::try_cast).transpose()?;
+
+        let storage =
+            match (self.storage, self.storage_diff) {
+                (None, None) => None,
+                (None, Some(diff)) => Some(rethnet_evm::state::StorageOverride::Diff(
+                    diff.into_iter()
+                        .map(|StorageSlotChange { index, value }| {
+                            Ok((BigInt::try_cast(index)?, BigInt::try_cast(value)?))
+                        })
+                        .collect::<napi::Result<_>>()?,
+                )),
+                (Some(full), None) => Some(rethnet_evm::state::StorageOverride::Full(
+                    full.into_iter()
+                        .map(|StorageSlotChange { index, value }| {
+                            Ok((BigInt::try_cast(index)?, BigInt::try_cast(value)?))
+                        })
+                        .collect::<napi::Result<_>>()?,
+                )),
+                (Some(_), Some(_)) => return Err(napi::Error::new(
+                    napi::Status::InvalidArg,
+                    "Both `storage` and `storageDiff` were specified. Only one can be provided.",
+                )),
+            };
 
         Ok(rethnet_evm::state::AccountOverride {
             balance,
diff --git a/packages/hardhat-core/src/internal/hardhat-network/provider/vm/rethnet.ts b/packages/hardhat-core/src/internal/hardhat-network/provider/vm/rethnet.ts
index 1c9a06a7670..c2eb91ba06f 100644
--- a/packages/hardhat-core/src/internal/hardhat-network/provider/vm/rethnet.ts
+++ b/packages/hardhat-core/src/internal/hardhat-network/provider/vm/rethnet.ts
@@ -3,7 +3,9 @@ import { Common } from "@nomicfoundation/ethereumjs-common";
 import {
   Account,
   Address,
+  bufferToBigInt,
   KECCAK256_NULL,
+  toBuffer,
 } from "@nomicfoundation/ethereumjs-util";
 import { TypedTransaction } from "@nomicfoundation/ethereumjs-tx";
 import {
@@ -141,13 +143,41 @@ export class RethnetAdapter implements VMAdapter {
       this._state.asInner(),
       new StateOverrides(
         Object.entries(stateOverrideSet).map(([address, account]) => {
+          const storage =
+            account.state !== undefined
+              ? Object.entries(account.state).map(([key, value]) => {
+                  const index = bufferToBigInt(toBuffer(key));
+                  const number = bufferToBigInt(toBuffer(value));
+
+                  return {
+                    index,
+                    value: number,
+                  };
+                })
+              : undefined;
+
+          const storageDiff =
+            account.stateDiff !== undefined
+              ? Object.entries(account.stateDiff).map(([key, value]) => {
+                  const index = bufferToBigInt(toBuffer(key));
+                  const number = bufferToBigInt(toBuffer(value));
+
+                  return {
+                    index,
+                    value: number,
+                  };
+                })
+              : undefined;
+
           const accountOverride: AccountOverride = {
             balance: account.balance,
             nonce: account.nonce,
             code: account.code,
+            storage,
+            storageDiff,
           };
 
-          return [Address.fromString(address).buf, accountOverride];
+          return [toBuffer(address), accountOverride];
         })
       ),
       config,