fix: store code when inserting account in LayeredState

crates/rethnet_evm/src/state/hybrid.rs

@@ -88,7 +88,7 @@ impl StateDebug for HybridState<RethnetLayer> {
         account_info: AccountInfo,
     ) -> Result<(), Self::Error> {
         self.trie.insert_account(address, account_info.clone())?;
-        self.changes.account_or_insert_mut(&address).info = account_info;
+        self.changes.insert_account(&address, account_info);
 
         Ok(())
     }

crates/rethnet_evm/src/state/layered.rs

@@ -94,7 +94,7 @@ impl StateDebug for LayeredState<RethnetLayer> {
         address: Address,
         account_info: AccountInfo,
     ) -> Result<(), Self::Error> {
-        self.changes.account_or_insert_mut(&address).info = account_info;
+        self.changes.insert_account(&address, account_info);
 
         Ok(())
     }
@@ -275,3 +275,24 @@ impl StateHistory for LayeredState<RethnetLayer> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use rethnet_eth::Bytes;
+
+    #[test]
+    fn code_by_hash_success() {
+        let mut state = LayeredState::<RethnetLayer>::default();
+        let inserted_bytecode = Bytecode::new_raw(Bytes::from("0x11"));
+        state
+            .insert_account(
+                Address::from_low_u64_ne(1234),
+                AccountInfo::new(U256::ZERO, 0, inserted_bytecode.clone()),
+            )
+            .unwrap();
+        let retrieved_bytecode = state.code_by_hash(inserted_bytecode.hash()).unwrap();
+        assert_eq!(retrieved_bytecode, inserted_bytecode);
+    }
+}

crates/rethnet_evm/src/state/layered/changes.rs

@@ -344,6 +344,14 @@ impl LayeredChanges<RethnetLayer> {
     pub fn remove_code(&mut self, code_hash: &B256) {
         self.last_layer_mut().contracts.remove(code_hash);
     }
+
+    pub fn insert_account(&mut self, address: &Address, mut account_info: AccountInfo) {
+        if let Some(code) = account_info.code.take() {
+            self.insert_code(code);
+        }
+
+        self.account_or_insert_mut(address).info = account_info;
+    }
 }
 
 impl From<&LayeredChanges<RethnetLayer>> for SharedMap<B256, Bytecode, true> {