From da67116b4d9469e87642ded8f289ccfe55900393 Mon Sep 17 00:00:00 2001
From: Bjerg <onbjerg@users.noreply.github.com>
Date: Thu, 17 Mar 2022 11:29:14 +0100
Subject: [PATCH] feat: support hardhat artifacts in `vm.getCode` (#956)

Ports #903
---
 .../src/executor/inspector/cheatcodes/ext.rs  | 33 +++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/forge/src/executor/inspector/cheatcodes/ext.rs b/forge/src/executor/inspector/cheatcodes/ext.rs
index e139ae90fbf7..30d4d874a0bc 100644
--- a/forge/src/executor/inspector/cheatcodes/ext.rs
+++ b/forge/src/executor/inspector/cheatcodes/ext.rs
@@ -4,6 +4,7 @@ use ethers::{
     abi::{self, AbiEncode, Token},
     prelude::{artifacts::CompactContractBytecode, ProjectPathsConfig},
 };
+use serde::Deserialize;
 use std::{fs::File, io::Read, path::Path, process::Command};
 
 fn ffi(args: &[String]) -> Result<Bytes, Bytes> {
@@ -18,6 +19,34 @@ fn ffi(args: &[String]) -> Result<Bytes, Bytes> {
     Ok(abi::encode(&[Token::Bytes(decoded.to_vec())]).into())
 }
 
+/// An enum which unifies the deserialization of Hardhat-style artifacts with Forge-style artifacts
+/// to get their bytecode.
+#[derive(Deserialize)]
+#[serde(untagged)]
+#[allow(clippy::large_enum_variant)]
+enum ArtifactBytecode {
+    Hardhat(HardhatArtifact),
+    Forge(CompactContractBytecode),
+}
+
+impl ArtifactBytecode {
+    fn into_inner(self) -> Option<ethers::types::Bytes> {
+        match self {
+            ArtifactBytecode::Hardhat(inner) => Some(inner.bytecode),
+            ArtifactBytecode::Forge(inner) => {
+                inner.bytecode.and_then(|bytecode| bytecode.object.into_bytes())
+            }
+        }
+    }
+}
+
+/// A thin wrapper around a Hardhat-style artifact that only extracts the bytecode.
+#[derive(Deserialize)]
+struct HardhatArtifact {
+    #[serde(deserialize_with = "ethers::solc::artifacts::deserialize_bytes")]
+    bytecode: ethers::types::Bytes,
+}
+
 fn get_code(path: &str) -> Result<Bytes, Bytes> {
     let path = if path.ends_with(".json") {
         Path::new(&path).to_path_buf()
@@ -37,10 +66,10 @@ fn get_code(path: &str) -> Result<Bytes, Bytes> {
         .read_to_string(&mut buffer)
         .map_err(|err| err.to_string().encode())?;
 
-    let artifact = serde_json::from_str::<CompactContractBytecode>(&buffer)
+    let bytecode = serde_json::from_str::<ArtifactBytecode>(&buffer)
         .map_err(|err| err.to_string().encode())?;
 
-    if let Some(bin) = artifact.bytecode.and_then(|bytecode| bytecode.object.into_bytes()) {
+    if let Some(bin) = bytecode.into_inner() {
         Ok(abi::encode(&[Token::Bytes(bin.to_vec())]).into())
     } else {
         Err("No bytecode for contract. Is it abstract or unlinked?".to_string().encode().into())