From 141bad2e6be732bef28f121c54e8038c0697c59e Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Fri, 5 Jan 2024 16:46:02 -0500
Subject: [PATCH 01/13] WIP implementation using codehash map

---
 lib/Echidna/Exec.hs              | 62 +++++++++++++-------------------
 lib/Echidna/Output/Source.hs     |  6 ++--
 lib/Echidna/Types/CodehashMap.hs | 39 ++++++++++++++++++++
 lib/Echidna/Types/Config.hs      |  2 ++
 lib/Echidna/Types/Coverage.hs    |  3 +-
 5 files changed, 71 insertions(+), 41 deletions(-)
 create mode 100644 lib/Echidna/Types/CodehashMap.hs

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index f85301f1f..4f9483949 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -14,7 +14,7 @@ import Control.Monad.Reader (MonadReader, asks)
 import Control.Monad.ST (ST, stToIO, RealWorld)
 import Data.Bits
 import Data.ByteString qualified as BS
-import Data.IORef (readIORef, atomicWriteIORef, atomicModifyIORef', newIORef, writeIORef, modifyIORef')
+import Data.IORef (readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe, fromJust)
 import Data.Text qualified as T
@@ -36,7 +36,8 @@ import Echidna.Transaction
 import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
 import Echidna.Types.Coverage (CoverageInfo)
-import Echidna.Types.Signature (getBytecodeMetadata, lookupBytecodeMetadata)
+import Echidna.Types.CodehashMap (lookupUsingCodehash)
+import Echidna.Types.Signature (getBytecodeMetadata)
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber, getResult)
 import Echidna.Utility (getTimestamp, timePrefix)
@@ -247,12 +248,12 @@ execTxWithCov
   -> m ((VMResult RealWorld, Gas), Bool)
 execTxWithCov tx = do
   covRef <- asks (.coverageRef)
-  metaCacheRef <- asks (.metadataCache)
-  cache <- liftIO $ readIORef metaCacheRef
+  codehashMap <- asks (.codehashMap)
+  dapp <- asks (.dapp)
 
   covContextRef <- liftIO $ newIORef (False, Nothing)
 
-  r <- execTxWith (execCov covRef covContextRef cache) tx
+  r <- execTxWith (execCov covRef codehashMap dapp covContextRef) tx
 
   (grew, lastLoc) <- liftIO $ readIORef covContextRef
 
@@ -270,7 +271,7 @@ execTxWithCov tx = do
   pure (r, grew || grew')
   where
     -- the same as EVM.exec but collects coverage, will stop on a query
-    execCov covRef covContextRef cache = do
+    execCov covRef codehashMap dapp covContextRef = do
       vm <- get
       (r, vm') <- liftIO $ loop vm
       put vm'
@@ -292,32 +293,21 @@ execTxWithCov tx = do
       addCoverage :: VM RealWorld -> IO ()
       addCoverage !vm = do
         let (pc, opIx, depth) = currentCovLoc vm
-            meta = currentMeta vm
-        cov <- readIORef covRef
-        case Map.lookup meta cov of
-          Nothing -> do
-            let size = BS.length . forceBuf . fromJust . view bytecode . fromJust $
-                  Map.lookup vm.state.contract vm.env.contracts
-            if size > 0 then do
-              vec <- VMut.new size
-              -- We use -1 for opIx to indicate that the location was not covered
-              forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
-
-              vec' <- atomicModifyIORef' covRef $ \cm ->
-                -- this should reduce races
-                case Map.lookup meta cm of
-                  Nothing -> (Map.insert meta vec cm, vec)
-                  Just vec' -> (cm, vec')
-
-              VMut.write vec' pc (opIx, fromIntegral depth, 0 `setBit` fromEnum Stop)
-
-              writeIORef covContextRef (True, Just (vec', pc))
-            else do
-              -- TODO: should we collect the coverage here? Even if there is no
-              -- bytecode for external contract, we could have a "virtual" location
-              -- that PC landed at and record that.
-              pure ()
-          Just vec ->
+            contr = currentContract vm
+
+        maybeMetaVec <- lookupUsingCodehash codehashMap contr dapp covRef $ do
+          let size = BS.length . forceBuf . fromJust . view bytecode . fromJust $
+                Map.lookup vm.state.contract vm.env.contracts
+          if size == 0 then pure Nothing else do
+            -- IO for making a new vec
+            vec <- VMut.new size
+            -- We use -1 for opIx to indicate that the location was not covered
+            forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
+            pure $ Just vec
+
+        case maybeMetaVec of
+          Nothing -> pure ()
+          Just (meta, vec) -> do
             -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
             -- we observed this in some real-world scenarios. This is likely a
             -- bug in another place, investigate.
@@ -332,11 +322,9 @@ execTxWithCov tx = do
       -- | Get the VM's current execution location
       currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
 
-      -- | Get the current contract's bytecode metadata
-      currentMeta vm = fromMaybe (error "no contract information on coverage") $ do
-        buffer <- vm ^? #env % #contracts % at vm.state.codeContract % _Just % bytecode
-        let bc = forceBuf $ fromJust buffer
-        pure $ lookupBytecodeMetadata cache bc
+      -- | Get the current contract
+      currentContract vm = fromMaybe (error "no contract information on coverage") $
+        vm ^? #env % #contracts % at vm.state.codeContract % _Just
 
 initialVM :: Bool -> ST s (VM s)
 initialVM ffi = do
diff --git a/lib/Echidna/Output/Source.hs b/lib/Echidna/Output/Source.hs
index 6be9f0096..ca1044acd 100644
--- a/lib/Echidna/Output/Source.hs
+++ b/lib/Echidna/Output/Source.hs
@@ -162,8 +162,8 @@ srcMapCov sc covMap contracts = do
   Map.unionsWith Map.union <$> mapM linesCovered contracts
   where
   linesCovered :: SolcContract -> IO (Map FilePath (Map Int [TxResult]))
-  linesCovered c =
-    case Map.lookup (getBytecodeMetadata c.runtimeCode) covMap of
+  linesCovered c = undefined -- TODO
+    {- case Map.lookup (getBytecodeMetadata c.runtimeCode) covMap of
       Just vec -> VU.foldl' (\acc covInfo -> case covInfo of
         (-1, _, _) -> acc -- not covered
         (opIx, _stackDepths, txResults) ->
@@ -183,7 +183,7 @@ srcMapCov sc covMap contracts = do
                 Nothing -> acc
             Nothing -> acc
         ) mempty vec
-      Nothing -> mempty
+      Nothing -> mempty -}
 
 -- | Given a contract, and tuple as coverage, return the corresponding mapped line (if any)
 srcMapForOpLocation :: SolcContract -> OpIx -> Maybe SrcMap
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
new file mode 100644
index 000000000..05977207f
--- /dev/null
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -0,0 +1,39 @@
+module Echidna.Types.CodehashMap where
+
+import Data.IORef (IORef, readIORef, atomicModifyIORef')
+import Data.Map.Strict qualified as Map
+import Data.Map.Strict (Map)
+import Data.Maybe (fromMaybe)
+import EVM.Dapp (DappInfo, findSrc)
+import EVM.Solidity (SolcContract(..))
+import EVM.Types (Contract(..), W256, maybeLitWord)
+
+type CodehashMap = IORef (Map W256 W256)
+
+lookupCodehash :: CodehashMap -> W256 -> Contract -> DappInfo -> IO W256
+lookupCodehash chmap codehash contr dapp = Map.lookup codehash <$> readIORef chmap >>= \case
+  Just val -> pure val
+  Nothing -> do
+    let originalCodehash = fromMaybe codehash $ (.creationCodehash) <$> findSrc contr dapp
+    atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
+    pure originalCodehash
+
+lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe (W256, a))
+lookupUsingCodehash chmap contr dapp mapRef make = do
+  mapVal <- readIORef mapRef
+  ifNotFound codehash mapVal $ do
+    codehash' <- lookupCodehash chmap codehash contr dapp
+    ifNotFound codehash' mapVal $ do
+      retVal <- make
+      retVal' <- atomicModifyIORef' mapRef $ modifyFn codehash' retVal
+      pure $ (codehash', ) <$> retVal'
+  where
+    codehash = fromMaybe (error "TODO make error msg") $ maybeLitWord contr.codehash
+    ifNotFound key mapVal notFoundCase = case (Map.lookup key mapVal) of
+      Nothing -> notFoundCase
+      Just val -> pure $ Just (key, val)
+
+    modifyFn _ Nothing oldMap = (oldMap, Nothing)
+    modifyFn key (Just val) oldMap = case (Map.lookup key oldMap) of
+      Just val' -> (oldMap, Just val')
+      Nothing -> (Map.insert key val oldMap, Just val)
diff --git a/lib/Echidna/Types/Config.hs b/lib/Echidna/Types/Config.hs
index 0098bbce7..4dcf000da 100644
--- a/lib/Echidna/Types/Config.hs
+++ b/lib/Echidna/Types/Config.hs
@@ -13,6 +13,7 @@ import EVM.Dapp (DappInfo)
 import EVM.Types (Addr, Contract, W256)
 
 import Echidna.Types.Campaign (CampaignConf, CampaignEvent)
+import Echidna.Types.CodehashMap (CodehashMap)
 import Echidna.Types.Corpus (Corpus)
 import Echidna.Types.Coverage (CoverageMap)
 import Echidna.Types.Signature (MetadataCache)
@@ -69,6 +70,7 @@ data Env = Env
 
   , testsRef :: IORef [EchidnaTest]
   , coverageRef :: IORef CoverageMap
+  , codehashMap :: CodehashMap
   , corpusRef :: IORef Corpus
 
   , metadataCache :: IORef MetadataCache
diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index f793abf5a..6508e8829 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -8,11 +8,12 @@ import Data.Map.Strict (Map)
 import Data.Vector.Unboxed.Mutable (IOVector)
 import Data.Vector.Unboxed.Mutable qualified as V
 import Data.Word (Word64)
+import EVM.Types (W256)
 
 import Echidna.Types.Tx (TxResult)
 
 -- | Map with the coverage information needed for fuzzing and source code printing
-type CoverageMap = Map ByteString (IOVector CoverageInfo)
+type CoverageMap = Map W256 (IOVector CoverageInfo)
 
 -- | Basic coverage information
 type CoverageInfo = (OpIx, StackDepths, TxResults)

From 61ec23c3cba9b60b1c7fbeb401038e58feea1833 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Mon, 8 Jan 2024 11:35:40 -0500
Subject: [PATCH 02/13] compiles and runs; significantly faster than master

---
 lib/Echidna/Exec.hs              |  2 +-
 lib/Echidna/Output/JSON.hs       |  3 +--
 lib/Echidna/Output/Source.hs     |  7 +++----
 lib/Echidna/Types/CodehashMap.hs | 22 ++++++++++++----------
 lib/Echidna/Types/Coverage.hs    |  1 -
 src/Main.hs                      |  2 ++
 src/test/Common.hs               |  4 ++++
 src/test/Tests/Compile.hs        |  2 ++
 8 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 4f9483949..9e20bafdc 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -307,7 +307,7 @@ execTxWithCov tx = do
 
         case maybeMetaVec of
           Nothing -> pure ()
-          Just (meta, vec) -> do
+          Just vec -> do
             -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
             -- we observed this in some real-world scenarios. This is likely a
             -- bug in another place, investigate.
diff --git a/lib/Echidna/Output/JSON.hs b/lib/Echidna/Output/JSON.hs
index a1abe5a7c..ebacc8183 100644
--- a/lib/Echidna/Output/JSON.hs
+++ b/lib/Echidna/Output/JSON.hs
@@ -14,7 +14,6 @@ import Data.Vector.Unboxed qualified as VU
 import Numeric (showHex)
 
 import EVM.Dapp (DappInfo)
-import EVM.Types (keccak')
 
 import Echidna.ABI (ppAbiValue, GenDict(..))
 import Echidna.Events (Events, extractEvents)
@@ -110,7 +109,7 @@ encodeCampaign env workerStates = do
     , _error = Nothing
     , _tests = mapTest env.dapp <$> tests
     , seed = worker0.genDict.defSeed
-    , coverage = Map.mapKeys (("0x" ++) . (`showHex` "") . keccak') $ VU.toList <$> frozenCov
+    , coverage = Map.mapKeys (("0x" ++) . (`showHex` "")) $ VU.toList <$> frozenCov
     , gasInfo = Map.toList $ Map.unionsWith max ((.gasInfo) <$> workerStates)
     }
 
diff --git a/lib/Echidna/Output/Source.hs b/lib/Echidna/Output/Source.hs
index ca1044acd..6e8075fe0 100644
--- a/lib/Echidna/Output/Source.hs
+++ b/lib/Echidna/Output/Source.hs
@@ -29,7 +29,6 @@ import EVM.Solidity (SourceCache(..), SrcMap, SolcContract(..))
 
 import Echidna.Types.Coverage (OpIx, unpackTxResults, CoverageMap)
 import Echidna.Types.Tx (TxResult(..))
-import Echidna.Types.Signature (getBytecodeMetadata)
 
 saveCoverages
   :: [CoverageFileType]
@@ -162,8 +161,8 @@ srcMapCov sc covMap contracts = do
   Map.unionsWith Map.union <$> mapM linesCovered contracts
   where
   linesCovered :: SolcContract -> IO (Map FilePath (Map Int [TxResult]))
-  linesCovered c = undefined -- TODO
-    {- case Map.lookup (getBytecodeMetadata c.runtimeCode) covMap of
+  linesCovered c =
+    case Map.lookup c.runtimeCodehash covMap of
       Just vec -> VU.foldl' (\acc covInfo -> case covInfo of
         (-1, _, _) -> acc -- not covered
         (opIx, _stackDepths, txResults) ->
@@ -183,7 +182,7 @@ srcMapCov sc covMap contracts = do
                 Nothing -> acc
             Nothing -> acc
         ) mempty vec
-      Nothing -> mempty -}
+      Nothing -> mempty
 
 -- | Given a contract, and tuple as coverage, return the corresponding mapped line (if any)
 srcMapForOpLocation :: SolcContract -> OpIx -> Maybe SrcMap
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index 05977207f..4d3dd4707 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -4,9 +4,10 @@ import Data.IORef (IORef, readIORef, atomicModifyIORef')
 import Data.Map.Strict qualified as Map
 import Data.Map.Strict (Map)
 import Data.Maybe (fromMaybe)
+import Echidna.Symbolic (forceWord)
 import EVM.Dapp (DappInfo, findSrc)
 import EVM.Solidity (SolcContract(..))
-import EVM.Types (Contract(..), W256, maybeLitWord)
+import EVM.Types (Contract(..), W256)
 
 type CodehashMap = IORef (Map W256 W256)
 
@@ -14,26 +15,27 @@ lookupCodehash :: CodehashMap -> W256 -> Contract -> DappInfo -> IO W256
 lookupCodehash chmap codehash contr dapp = Map.lookup codehash <$> readIORef chmap >>= \case
   Just val -> pure val
   Nothing -> do
-    let originalCodehash = fromMaybe codehash $ (.creationCodehash) <$> findSrc contr dapp
+    let originalCodehash = fromMaybe codehash $ (.runtimeCodehash) <$> findSrc contr dapp
     atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
     pure originalCodehash
 
-lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe (W256, a))
+lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe a)
 lookupUsingCodehash chmap contr dapp mapRef make = do
   mapVal <- readIORef mapRef
   ifNotFound codehash mapVal $ do
     codehash' <- lookupCodehash chmap codehash contr dapp
     ifNotFound codehash' mapVal $ do
-      retVal <- make
-      retVal' <- atomicModifyIORef' mapRef $ modifyFn codehash' retVal
-      pure $ (codehash', ) <$> retVal'
+      toInsert <- make
+      applyModification codehash' toInsert
   where
-    codehash = fromMaybe (error "TODO make error msg") $ maybeLitWord contr.codehash
+    codehash = forceWord contr.codehash
     ifNotFound key mapVal notFoundCase = case (Map.lookup key mapVal) of
       Nothing -> notFoundCase
-      Just val -> pure $ Just (key, val)
+      Just val -> pure (Just val)
 
-    modifyFn _ Nothing oldMap = (oldMap, Nothing)
-    modifyFn key (Just val) oldMap = case (Map.lookup key oldMap) of
+    applyModification _ Nothing = pure Nothing
+    applyModification key (Just val) = atomicModifyIORef' mapRef $ modifyFn key val
+
+    modifyFn key val oldMap = case (Map.lookup key oldMap) of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)
diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index 6508e8829..0a2c82503 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -1,7 +1,6 @@
 module Echidna.Types.Coverage where
 
 import Data.Bits (testBit)
-import Data.ByteString (ByteString)
 import Data.List (foldl')
 import Data.Map qualified as Map
 import Data.Map.Strict (Map)
diff --git a/src/Main.hs b/src/Main.hs
index 7645ec055..171b76885 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -90,6 +90,7 @@ main = withUtf8 $ withCP65001 $ do
   cacheContractsRef <- newIORef $ fromMaybe mempty loadedContractsCache
   cacheSlotsRef <- newIORef $ fromMaybe mempty loadedSlotsCache
   cacheMetaRef <- newIORef mempty
+  codehashMap <- newIORef mempty
   chainId <- RPC.fetchChainId cfg.rpcUrl
   eventQueue <- newChan
   coverageRef <- newIORef mempty
@@ -103,6 +104,7 @@ main = withUtf8 $ withCP65001 $ do
                 -- TODO put in real path
               , dapp = dappInfo "/" buildOutput
               , metadataCache = cacheMetaRef
+              , codehashMap = codehashMap
               , fetchContractCache = cacheContractsRef
               , fetchSlotCache = cacheSlotsRef
               , chainId = chainId
diff --git a/src/test/Common.hs b/src/test/Common.hs
index 6b3a022f8..3cc8f9624 100644
--- a/src/test/Common.hs
+++ b/src/test/Common.hs
@@ -98,6 +98,7 @@ runContract f selectedContract cfg = do
     contracts = Map.elems . Map.unions $ (\(BuildOutput (Contracts c) _) -> c) <$> buildOutputs
 
   metadataCache <- newIORef mempty
+  codehashMap <- newIORef mempty
   fetchContractCache <- newIORef mempty
   fetchSlotCache <- newIORef mempty
   coverageRef <- newIORef mempty
@@ -107,6 +108,7 @@ runContract f selectedContract cfg = do
   let env = Env { cfg = cfg
                 , dapp = dappInfo "/" buildOutput
                 , metadataCache
+                , codehashMap
                 , fetchContractCache
                 , fetchSlotCache
                 , coverageRef
@@ -162,6 +164,7 @@ testContract' fp n v configPath s expectations = testCase fp $ withSolcVersion v
 checkConstructorConditions :: FilePath -> String -> TestTree
 checkConstructorConditions fp as = testCase fp $ do
   cacheMeta <- newIORef mempty
+  codehashMap <- newIORef mempty
   cacheContracts <- newIORef mempty
   cacheSlots <- newIORef mempty
   coverageRef <- newIORef mempty
@@ -171,6 +174,7 @@ checkConstructorConditions fp as = testCase fp $ do
   let env = Env { cfg = testConfig
                 , dapp = emptyDapp
                 , metadataCache = cacheMeta
+                , codehashMap
                 , fetchContractCache = cacheContracts
                 , fetchSlotCache = cacheSlots
                 , coverageRef
diff --git a/src/test/Tests/Compile.hs b/src/test/Tests/Compile.hs
index 54a9afa19..980ec13be 100644
--- a/src/test/Tests/Compile.hs
+++ b/src/test/Tests/Compile.hs
@@ -43,6 +43,7 @@ loadFails :: FilePath -> Maybe Text -> String -> (SolException -> Bool) -> TestT
 loadFails fp c e p = testCase fp . catch tryLoad $ assertBool e . p where
   tryLoad = do
     cacheMeta <- newIORef mempty
+    codehashMap <- newIORef mempty
     cacheContracts <- newIORef mempty
     cacheSlots <- newIORef mempty
     eventQueue <- newChan
@@ -52,6 +53,7 @@ loadFails fp c e p = testCase fp . catch tryLoad $ assertBool e . p where
     let env = Env { cfg = testConfig
                   , dapp = emptyDapp
                   , metadataCache = cacheMeta
+                  , codehashMap
                   , fetchContractCache = cacheContracts
                   , fetchSlotCache = cacheSlots
                   , chainId = Nothing

From 7e70e18d56480d5ea6ae85e4e802c9f275657d0f Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Mon, 8 Jan 2024 14:54:31 -0500
Subject: [PATCH 03/13] use codehash for signaturemap

---
 lib/Echidna/Solidity.hs          |  6 +++---
 lib/Echidna/Transaction.hs       | 32 +++++++++++++++++---------------
 lib/Echidna/Types/CodehashMap.hs | 11 +++++++++++
 lib/Echidna/Types/Signature.hs   |  4 ++--
 4 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/lib/Echidna/Solidity.hs b/lib/Echidna/Solidity.hs
index 1300730e3..7f2547333 100644
--- a/lib/Echidna/Solidity.hs
+++ b/lib/Echidna/Solidity.hs
@@ -44,7 +44,7 @@ import Echidna.Symbolic (forceAddr)
 import Echidna.Test (createTests, isAssertionMode, isPropertyMode, isDapptestMode)
 import Echidna.Types.Config (EConfig(..), Env(..))
 import Echidna.Types.Signature
-  (ContractName, SolSignature, SignatureMap, getBytecodeMetadata, FunctionName)
+  (ContractName, SolSignature, SignatureMap, FunctionName)
 import Echidna.Types.Solidity
 import Echidna.Types.Test (EchidnaTest(..))
 import Echidna.Types.Tx
@@ -219,11 +219,11 @@ loadSpecified env name cs = do
             let filtered = filterMethods contract.contractName
                                          solConf.methodFilter
                                          (abiOf solConf.prefix contract)
-            in (getBytecodeMetadata contract.runtimeCode,) <$> NE.nonEmpty filtered)
+            in (contract.runtimeCodehash,) <$> NE.nonEmpty filtered)
           cs
       else
         case NE.nonEmpty fabiOfc of
-          Just ne -> Map.singleton (getBytecodeMetadata mainContract.runtimeCode) ne
+          Just ne -> Map.singleton mainContract.runtimeCodehash ne
           Nothing -> mempty
 
   -- Set up initial VM, either with chosen contract or Etheno initialization file
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index f78b43da2..946afea74 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -7,27 +7,30 @@ import Optics.Core
 import Optics.State.Operators
 
 import Control.Monad (join)
-import Control.Monad.Random.Strict (MonadRandom, getRandomR, uniform, MonadIO)
+import Control.Monad.IO.Class (MonadIO, liftIO)
+import Control.Monad.Random.Strict (MonadRandom, getRandomR, uniform)
+import Control.Monad.Reader (MonadReader, ask)
 import Control.Monad.State.Strict (MonadState, gets, modify', execState)
 import Control.Monad.ST (RealWorld)
 import Data.Map (Map, toList)
-import Data.Map qualified as Map
-import Data.Maybe (mapMaybe, fromJust)
+import Data.Maybe (mapMaybe)
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Vector qualified as V
 
-import EVM (initialContract, loadContract, bytecode, resetState)
+import EVM (initialContract, loadContract, resetState)
 import EVM.ABI (abiValueType)
-import EVM.Types hiding (VMOpts(timestamp, gasprice))
+import EVM.Types hiding (Env, VMOpts(timestamp, gasprice))
 
 import Echidna.ABI
 import Echidna.Orphans.JSON ()
-import Echidna.Symbolic (forceBuf, forceWord, forceAddr)
+import Echidna.Symbolic (forceWord, forceAddr)
 import Echidna.Types (fromEVM)
+import Echidna.Types.CodehashMap (lookupUsingCodehashNoInsert)
+import Echidna.Types.Config (Env(..))
 import Echidna.Types.Random
 import Echidna.Types.Signature
-  (SignatureMap, SolCall, ContractA, MetadataCache, lookupBytecodeMetadata)
+  (SignatureMap, SolCall, ContractA, MetadataCache)
 import Echidna.Types.Tx
 import Echidna.Types.World (World(..))
 import Echidna.Types.Campaign
@@ -53,18 +56,19 @@ getSignatures hmm (Just lmm) =
 
 -- | Generate a random 'Transaction' with either synthesis or mutation of dictionary entries.
 genTx
-  :: (MonadRandom m, MonadState WorkerState m)
+  :: (MonadIO m, MonadRandom m, MonadState WorkerState m, MonadReader Env m)
   => MetadataCache
   -> World
   -> TxConf
   -> Map (Expr EAddr) Contract
   -> m Tx
 genTx memo world txConf deployedContracts = do
+  env <- ask
   genDict <- gets (.genDict)
   sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
   sender <- rElem' world.senders
-  (dstAddr, dstAbis) <- rElem' $ Set.fromList $
-    mapMaybe (toContractA sigMap) (toList deployedContracts)
+  mappedList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
+  (dstAddr, dstAbis) <- rElem' $ Set.fromList $ mapMaybe id mappedList
   solCall <- genInteractionsM genDict dstAbis
   value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall
   ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues
@@ -78,11 +82,9 @@ genTx memo world txConf deployedContracts = do
             , delay = level ts
             }
   where
-    toContractA :: SignatureMap -> (Expr EAddr, Contract) -> Maybe ContractA
-    toContractA sigMap (addr, c) =
-      let bc = forceBuf $ fromJust $ view bytecode c
-          metadata = lookupBytecodeMetadata memo bc
-      in (forceAddr addr,) <$> Map.lookup metadata sigMap
+    toContractA :: Env -> SignatureMap -> (Expr EAddr, Contract) -> IO (Maybe ContractA)
+    toContractA env sigMap (addr, c) =
+      fmap (forceAddr addr,) <$> lookupUsingCodehashNoInsert env.codehashMap c env.dapp sigMap
 
 genDelay :: MonadRandom m => W256 -> Set W256 -> m W256
 genDelay mv ds = do
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index 4d3dd4707..7029dbdc1 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -39,3 +39,14 @@ lookupUsingCodehash chmap contr dapp mapRef make = do
     modifyFn key val oldMap = case (Map.lookup key oldMap) of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)
+
+lookupUsingCodehashNoInsert :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (Maybe a)
+lookupUsingCodehashNoInsert chmap contr dapp mapVal = do
+  ifNotFound codehash $ do
+    codehash' <- lookupCodehash chmap codehash contr dapp
+    ifNotFound codehash' $ pure Nothing
+  where
+    codehash = forceWord contr.codehash
+    ifNotFound key notFoundCase = case (Map.lookup key mapVal) of
+      Nothing -> notFoundCase
+      Just val -> pure (Just val)
diff --git a/lib/Echidna/Types/Signature.hs b/lib/Echidna/Types/Signature.hs
index 6a420b678..dc71a0cca 100644
--- a/lib/Echidna/Types/Signature.hs
+++ b/lib/Echidna/Types/Signature.hs
@@ -11,7 +11,7 @@ import Data.Maybe (fromMaybe)
 import Data.Text (Text)
 
 import EVM.ABI (AbiType, AbiValue)
-import EVM.Types (Addr)
+import EVM.Types (Addr, W256)
 import Data.Map (Map)
 
 -- | Name of the contract
@@ -34,7 +34,7 @@ type ContractA = (Addr, NonEmpty SolSignature)
 -- | Used to memoize results of getBytecodeMetadata
 type MetadataCache = Map ByteString ByteString
 
-type SignatureMap = Map ByteString (NonEmpty SolSignature)
+type SignatureMap = Map W256 (NonEmpty SolSignature)
 
 getBytecodeMetadata :: ByteString -> ByteString
 getBytecodeMetadata bs =

From 964cbafd323961fa4f33980b8ca7ee3bf6b4c4f6 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Mon, 8 Jan 2024 15:25:04 -0500
Subject: [PATCH 04/13] remove metadata cache

---
 lib/Echidna/Campaign.hs        | 23 ++++++-----------------
 lib/Echidna/Exec.hs            |  6 ------
 lib/Echidna/Transaction.hs     |  7 +++----
 lib/Echidna/Types/Config.hs    |  2 --
 lib/Echidna/Types/Signature.hs | 20 --------------------
 src/Main.hs                    |  2 --
 src/test/Common.hs             |  4 ----
 src/test/Tests/Compile.hs      |  2 --
 8 files changed, 9 insertions(+), 57 deletions(-)

diff --git a/lib/Echidna/Campaign.hs b/lib/Echidna/Campaign.hs
index 31768d81b..0b739e0e1 100644
--- a/lib/Echidna/Campaign.hs
+++ b/lib/Echidna/Campaign.hs
@@ -3,8 +3,6 @@
 
 module Echidna.Campaign where
 
-import Optics.Core hiding ((|>))
-
 import Control.Concurrent (writeChan)
 import Control.DeepSeq (force)
 import Control.Monad (replicateM, when, void, forM_)
@@ -17,16 +15,16 @@ import Control.Monad.ST (RealWorld)
 import Control.Monad.Trans (lift)
 import Data.Binary.Get (runGetOrFail)
 import Data.ByteString.Lazy qualified as LBS
-import Data.IORef (readIORef, writeIORef, atomicModifyIORef')
+import Data.IORef (readIORef, atomicModifyIORef')
 import Data.Map qualified as Map
 import Data.Map (Map, (\\))
-import Data.Maybe (isJust, mapMaybe, fromMaybe, fromJust)
+import Data.Maybe (isJust, mapMaybe, fromMaybe)
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Text (Text)
 import System.Random (mkStdGen)
 
-import EVM (bytecode, cheatCode)
+import EVM (cheatCode)
 import EVM.ABI (getAbi, AbiType(AbiAddressType), AbiValue(AbiAddress))
 import EVM.Types hiding (Env, Frame(state))
 
@@ -34,7 +32,7 @@ import Echidna.ABI
 import Echidna.Exec
 import Echidna.Mutator.Corpus
 import Echidna.Shrink (shrinkTest)
-import Echidna.Symbolic (forceBuf, forceAddr)
+import Echidna.Symbolic (forceAddr)
 import Echidna.Test
 import Echidna.Transaction
 import Echidna.Types (Gas)
@@ -42,7 +40,7 @@ import Echidna.Types.Campaign
 import Echidna.Types.Corpus (Corpus, corpusSize)
 import Echidna.Types.Coverage (scoveragePoints)
 import Echidna.Types.Config
-import Echidna.Types.Signature (makeBytecodeCache, FunctionName)
+import Echidna.Types.Signature (FunctionName)
 import Echidna.Types.Test
 import Echidna.Types.Test qualified as Test
 import Echidna.Types.Tx (TxCall(..), Tx(..), call)
@@ -86,12 +84,6 @@ runWorker
   -> Int     -- ^ Test limit for this worker
   -> m (WorkerStopReason, WorkerState)
 runWorker callback vm world dict workerId initialCorpus testLimit = do
-  metaCacheRef <- asks (.metadataCache)
-  fetchContractCacheRef <- asks (.fetchContractCache)
-  external <- liftIO $ Map.mapMaybe id <$> readIORef fetchContractCacheRef
-  let concretizeKeys = Map.foldrWithKey (Map.insert . forceAddr) mempty
-  liftIO $ writeIORef metaCacheRef (mkMemo (concretizeKeys vm.env.contracts <> external))
-
   let
     effectiveSeed = dict.defSeed + workerId
     effectiveGenDict = dict { defSeed = effectiveSeed }
@@ -152,8 +144,6 @@ runWorker callback vm world dict workerId initialCorpus testLimit = do
 
   continue = runUpdate (shrinkTest vm) >> lift callback >> run
 
-  mkMemo = makeBytecodeCache . map (forceBuf . fromJust . (^. bytecode)) . Map.elems
-
 -- | Generate a new sequences of transactions, either using the corpus or with
 -- randomly created transactions
 randseq
@@ -163,7 +153,6 @@ randseq
   -> m [Tx]
 randseq deployedContracts world = do
   env <- ask
-  memo <- liftIO $ readIORef env.metadataCache
 
   let
     mutConsts = env.cfg.campaignConf.mutConsts
@@ -174,7 +163,7 @@ randseq deployedContracts world = do
   --let rs = filter (not . null) $ map (.testReproducer) $ ca._tests
 
   -- Generate new random transactions
-  randTxs <- replicateM seqLen (genTx memo world txConf deployedContracts)
+  randTxs <- replicateM seqLen (genTx world txConf deployedContracts)
   -- Generate a random mutator
   cmut <- if seqLen == 1 then seqMutatorsStateless (fromConsts mutConsts)
                          else seqMutatorsStateful (fromConsts mutConsts)
diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 9e20bafdc..37ef8b7ca 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -37,7 +37,6 @@ import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
 import Echidna.Types.Coverage (CoverageInfo)
 import Echidna.Types.CodehashMap (lookupUsingCodehash)
-import Echidna.Types.Signature (getBytecodeMetadata)
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber, getResult)
 import Echidna.Utility (getTimestamp, timePrefix)
@@ -122,11 +121,6 @@ execTxWith executeTx tx = do
                 case ret of
                   -- TODO: fix hevm to not return an empty contract in case of an error
                   Just contract | contract.code /= RuntimeCode (ConcreteRuntimeCode "") -> do
-                    metaCacheRef <- asks (.metadataCache)
-                    metaCache <- liftIO $ readIORef metaCacheRef
-                    let bc = forceBuf $ fromJust (contract ^. bytecode)
-                    liftIO $ atomicWriteIORef metaCacheRef $ Map.insert bc (getBytecodeMetadata bc) metaCache
-
                     fromEVM (continuation contract)
                     liftIO $ atomicWriteIORef cacheRef $ Map.insert addr (Just contract) cache
                   _ -> do
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index 946afea74..e3492f8aa 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -30,7 +30,7 @@ import Echidna.Types.CodehashMap (lookupUsingCodehashNoInsert)
 import Echidna.Types.Config (Env(..))
 import Echidna.Types.Random
 import Echidna.Types.Signature
-  (SignatureMap, SolCall, ContractA, MetadataCache)
+  (SignatureMap, SolCall, ContractA)
 import Echidna.Types.Tx
 import Echidna.Types.World (World(..))
 import Echidna.Types.Campaign
@@ -57,12 +57,11 @@ getSignatures hmm (Just lmm) =
 -- | Generate a random 'Transaction' with either synthesis or mutation of dictionary entries.
 genTx
   :: (MonadIO m, MonadRandom m, MonadState WorkerState m, MonadReader Env m)
-  => MetadataCache
-  -> World
+  => World
   -> TxConf
   -> Map (Expr EAddr) Contract
   -> m Tx
-genTx memo world txConf deployedContracts = do
+genTx world txConf deployedContracts = do
   env <- ask
   genDict <- gets (.genDict)
   sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
diff --git a/lib/Echidna/Types/Config.hs b/lib/Echidna/Types/Config.hs
index 4dcf000da..0aac2a5a5 100644
--- a/lib/Echidna/Types/Config.hs
+++ b/lib/Echidna/Types/Config.hs
@@ -16,7 +16,6 @@ import Echidna.Types.Campaign (CampaignConf, CampaignEvent)
 import Echidna.Types.CodehashMap (CodehashMap)
 import Echidna.Types.Corpus (Corpus)
 import Echidna.Types.Coverage (CoverageMap)
-import Echidna.Types.Signature (MetadataCache)
 import Echidna.Types.Solidity (SolConf)
 import Echidna.Types.Test (TestConf, EchidnaTest)
 import Echidna.Types.Tx (TxConf)
@@ -73,7 +72,6 @@ data Env = Env
   , codehashMap :: CodehashMap
   , corpusRef :: IORef Corpus
 
-  , metadataCache :: IORef MetadataCache
   , fetchContractCache :: IORef (Map Addr (Maybe Contract))
   , fetchSlotCache :: IORef (Map Addr (Map W256 (Maybe W256)))
   , chainId :: Maybe W256
diff --git a/lib/Echidna/Types/Signature.hs b/lib/Echidna/Types/Signature.hs
index dc71a0cca..c7f38bbec 100644
--- a/lib/Echidna/Types/Signature.hs
+++ b/lib/Echidna/Types/Signature.hs
@@ -4,10 +4,7 @@ module Echidna.Types.Signature where
 
 import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
-import Data.Foldable (find)
 import Data.List.NonEmpty (NonEmpty)
-import Data.Map.Strict qualified as M
-import Data.Maybe (fromMaybe)
 import Data.Text (Text)
 
 import EVM.ABI (AbiType, AbiValue)
@@ -31,25 +28,8 @@ type SolCall = (FunctionName, [AbiValue])
 -- | A contract is just an address with an ABI (for our purposes).
 type ContractA = (Addr, NonEmpty SolSignature)
 
--- | Used to memoize results of getBytecodeMetadata
-type MetadataCache = Map ByteString ByteString
-
 type SignatureMap = Map W256 (NonEmpty SolSignature)
 
-getBytecodeMetadata :: ByteString -> ByteString
-getBytecodeMetadata bs =
-  let stripCandidates = flip BS.breakSubstring bs <$> knownBzzrPrefixes in
-    case find ((/= mempty) . snd) stripCandidates of
-      Nothing     -> bs -- if no metadata is found, return the complete bytecode
-      Just (_, m) -> m
-
-lookupBytecodeMetadata :: MetadataCache -> ByteString -> ByteString
-lookupBytecodeMetadata memo bs = fromMaybe (getBytecodeMetadata bs) (memo M.!? bs)
-
--- | Precalculate getBytecodeMetadata for all contracts in a list
-makeBytecodeCache :: [ByteString] -> MetadataCache
-makeBytecodeCache bss = M.fromList $ bss `zip` (getBytecodeMetadata <$> bss)
-
 knownBzzrPrefixes :: [ByteString]
 knownBzzrPrefixes =
   -- a1 65 "bzzr0" 0x58 0x20 (solc <= 0.5.8)
diff --git a/src/Main.hs b/src/Main.hs
index 171b76885..f79d28aa7 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -89,7 +89,6 @@ main = withUtf8 $ withCP65001 $ do
   buildOutputs <- compileContracts cfg.solConf cliFilePath
   cacheContractsRef <- newIORef $ fromMaybe mempty loadedContractsCache
   cacheSlotsRef <- newIORef $ fromMaybe mempty loadedSlotsCache
-  cacheMetaRef <- newIORef mempty
   codehashMap <- newIORef mempty
   chainId <- RPC.fetchChainId cfg.rpcUrl
   eventQueue <- newChan
@@ -103,7 +102,6 @@ main = withUtf8 $ withCP65001 $ do
     env = Env { cfg
                 -- TODO put in real path
               , dapp = dappInfo "/" buildOutput
-              , metadataCache = cacheMetaRef
               , codehashMap = codehashMap
               , fetchContractCache = cacheContractsRef
               , fetchSlotCache = cacheSlotsRef
diff --git a/src/test/Common.hs b/src/test/Common.hs
index 3cc8f9624..17a1918f6 100644
--- a/src/test/Common.hs
+++ b/src/test/Common.hs
@@ -97,7 +97,6 @@ runContract f selectedContract cfg = do
     buildOutput = selectBuildOutput selectedContract buildOutputs
     contracts = Map.elems . Map.unions $ (\(BuildOutput (Contracts c) _) -> c) <$> buildOutputs
 
-  metadataCache <- newIORef mempty
   codehashMap <- newIORef mempty
   fetchContractCache <- newIORef mempty
   fetchSlotCache <- newIORef mempty
@@ -107,7 +106,6 @@ runContract f selectedContract cfg = do
   testsRef <- newIORef mempty
   let env = Env { cfg = cfg
                 , dapp = dappInfo "/" buildOutput
-                , metadataCache
                 , codehashMap
                 , fetchContractCache
                 , fetchSlotCache
@@ -163,7 +161,6 @@ testContract' fp n v configPath s expectations = testCase fp $ withSolcVersion v
 
 checkConstructorConditions :: FilePath -> String -> TestTree
 checkConstructorConditions fp as = testCase fp $ do
-  cacheMeta <- newIORef mempty
   codehashMap <- newIORef mempty
   cacheContracts <- newIORef mempty
   cacheSlots <- newIORef mempty
@@ -173,7 +170,6 @@ checkConstructorConditions fp as = testCase fp $ do
   eventQueue <- newChan
   let env = Env { cfg = testConfig
                 , dapp = emptyDapp
-                , metadataCache = cacheMeta
                 , codehashMap
                 , fetchContractCache = cacheContracts
                 , fetchSlotCache = cacheSlots
diff --git a/src/test/Tests/Compile.hs b/src/test/Tests/Compile.hs
index 980ec13be..54415f265 100644
--- a/src/test/Tests/Compile.hs
+++ b/src/test/Tests/Compile.hs
@@ -42,7 +42,6 @@ compilationTests = testGroup "Compilation and loading tests"
 loadFails :: FilePath -> Maybe Text -> String -> (SolException -> Bool) -> TestTree
 loadFails fp c e p = testCase fp . catch tryLoad $ assertBool e . p where
   tryLoad = do
-    cacheMeta <- newIORef mempty
     codehashMap <- newIORef mempty
     cacheContracts <- newIORef mempty
     cacheSlots <- newIORef mempty
@@ -52,7 +51,6 @@ loadFails fp c e p = testCase fp . catch tryLoad $ assertBool e . p where
     testsRef <- newIORef mempty
     let env = Env { cfg = testConfig
                   , dapp = emptyDapp
-                  , metadataCache = cacheMeta
                   , codehashMap
                   , fetchContractCache = cacheContracts
                   , fetchSlotCache = cacheSlots

From dc797eb0b3bf4661b6d104ef24a287bcbc6ece0e Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Tue, 9 Jan 2024 10:31:15 -0500
Subject: [PATCH 05/13] use codeContract rather than contract; use env rather
 than individual pieces of env

---
 lib/Echidna/Exec.hs | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 37ef8b7ca..99ae64a2d 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -10,7 +10,7 @@ import Optics.State.Operators
 import Control.Monad (when, forM_)
 import Control.Monad.Catch (MonadThrow(..))
 import Control.Monad.State.Strict (MonadState(get, put), execState, runStateT, MonadIO(liftIO), gets, modify', execStateT)
-import Control.Monad.Reader (MonadReader, asks)
+import Control.Monad.Reader (MonadReader, ask, asks)
 import Control.Monad.ST (ST, stToIO, RealWorld)
 import Data.Bits
 import Data.ByteString qualified as BS
@@ -241,13 +241,11 @@ execTxWithCov
   => Tx
   -> m ((VMResult RealWorld, Gas), Bool)
 execTxWithCov tx = do
-  covRef <- asks (.coverageRef)
-  codehashMap <- asks (.codehashMap)
-  dapp <- asks (.dapp)
+  env <- ask
 
   covContextRef <- liftIO $ newIORef (False, Nothing)
 
-  r <- execTxWith (execCov covRef codehashMap dapp covContextRef) tx
+  r <- execTxWith (execCov env covContextRef) tx
 
   (grew, lastLoc) <- liftIO $ readIORef covContextRef
 
@@ -265,7 +263,7 @@ execTxWithCov tx = do
   pure (r, grew || grew')
   where
     -- the same as EVM.exec but collects coverage, will stop on a query
-    execCov covRef codehashMap dapp covContextRef = do
+    execCov env covContextRef = do
       vm <- get
       (r, vm') <- liftIO $ loop vm
       put vm'
@@ -287,11 +285,10 @@ execTxWithCov tx = do
       addCoverage :: VM RealWorld -> IO ()
       addCoverage !vm = do
         let (pc, opIx, depth) = currentCovLoc vm
-            contr = currentContract vm
+            contract = currentContract vm
 
-        maybeMetaVec <- lookupUsingCodehash codehashMap contr dapp covRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode . fromJust $
-                Map.lookup vm.state.contract vm.env.contracts
+        maybeMetaVec <- lookupUsingCodehash env.codehashMap contract env.dapp env.coverageRef $ do
+          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
             vec <- VMut.new size

From 292168a8aa089f58d4f2821ba54b1f99341966b1 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Tue, 9 Jan 2024 10:50:07 -0500
Subject: [PATCH 06/13] refactor codehash helper functions

---
 lib/Echidna/Exec.hs              |  4 ++--
 lib/Echidna/Transaction.hs       |  4 ++--
 lib/Echidna/Types/CodehashMap.hs | 34 +++++++++++++-------------------
 3 files changed, 18 insertions(+), 24 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 99ae64a2d..b0b2d49f4 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -36,7 +36,7 @@ import Echidna.Transaction
 import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
 import Echidna.Types.Coverage (CoverageInfo)
-import Echidna.Types.CodehashMap (lookupUsingCodehash)
+import Echidna.Types.CodehashMap (lookupUsingCodehashOrInsert)
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber, getResult)
 import Echidna.Utility (getTimestamp, timePrefix)
@@ -287,7 +287,7 @@ execTxWithCov tx = do
         let (pc, opIx, depth) = currentCovLoc vm
             contract = currentContract vm
 
-        maybeMetaVec <- lookupUsingCodehash env.codehashMap contract env.dapp env.coverageRef $ do
+        maybeMetaVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
           let size = BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index e3492f8aa..70547fee3 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -26,7 +26,7 @@ import Echidna.ABI
 import Echidna.Orphans.JSON ()
 import Echidna.Symbolic (forceWord, forceAddr)
 import Echidna.Types (fromEVM)
-import Echidna.Types.CodehashMap (lookupUsingCodehashNoInsert)
+import Echidna.Types.CodehashMap (lookupUsingCodehash)
 import Echidna.Types.Config (Env(..))
 import Echidna.Types.Random
 import Echidna.Types.Signature
@@ -83,7 +83,7 @@ genTx world txConf deployedContracts = do
   where
     toContractA :: Env -> SignatureMap -> (Expr EAddr, Contract) -> IO (Maybe ContractA)
     toContractA env sigMap (addr, c) =
-      fmap (forceAddr addr,) <$> lookupUsingCodehashNoInsert env.codehashMap c env.dapp sigMap
+      fmap (forceAddr addr,) . snd <$> lookupUsingCodehash env.codehashMap c env.dapp sigMap
 
 genDelay :: MonadRandom m => W256 -> Set W256 -> m W256
 genDelay mv ds = do
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index 7029dbdc1..e591da80c 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -19,34 +19,28 @@ lookupCodehash chmap codehash contr dapp = Map.lookup codehash <$> readIORef chm
     atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
     pure originalCodehash
 
-lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe a)
-lookupUsingCodehash chmap contr dapp mapRef make = do
-  mapVal <- readIORef mapRef
-  ifNotFound codehash mapVal $ do
+lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (W256, Maybe a)
+lookupUsingCodehash chmap contr dapp mapVal = do
+  ifNotFound codehash $ do
     codehash' <- lookupCodehash chmap codehash contr dapp
-    ifNotFound codehash' mapVal $ do
-      toInsert <- make
-      applyModification codehash' toInsert
+    ifNotFound codehash' $ pure (codehash', Nothing)
   where
     codehash = forceWord contr.codehash
-    ifNotFound key mapVal notFoundCase = case (Map.lookup key mapVal) of
+    ifNotFound key notFoundCase = case (Map.lookup key mapVal) of
       Nothing -> notFoundCase
-      Just val -> pure (Just val)
+      Just val -> pure (key, Just val)
 
+lookupUsingCodehashOrInsert :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe a)
+lookupUsingCodehashOrInsert chmap contr dapp mapRef make = do
+  mapVal <- readIORef mapRef
+  (key, valFound) <- lookupUsingCodehash chmap contr dapp mapVal
+  case valFound of
+    Just val -> pure (Just val)
+    Nothing -> applyModification key =<< make
+  where
     applyModification _ Nothing = pure Nothing
     applyModification key (Just val) = atomicModifyIORef' mapRef $ modifyFn key val
 
     modifyFn key val oldMap = case (Map.lookup key oldMap) of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)
-
-lookupUsingCodehashNoInsert :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (Maybe a)
-lookupUsingCodehashNoInsert chmap contr dapp mapVal = do
-  ifNotFound codehash $ do
-    codehash' <- lookupCodehash chmap codehash contr dapp
-    ifNotFound codehash' $ pure Nothing
-  where
-    codehash = forceWord contr.codehash
-    ifNotFound key notFoundCase = case (Map.lookup key mapVal) of
-      Nothing -> notFoundCase
-      Just val -> pure (Just val)

From 7c8fdbc793d024c9c9194f72513159476d262772 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Tue, 9 Jan 2024 12:04:33 -0500
Subject: [PATCH 07/13] hlint

---
 lib/Echidna/Transaction.hs       |  4 ++--
 lib/Echidna/Types/CodehashMap.hs | 19 ++++++++++---------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index 70547fee3..eba92a707 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -13,7 +13,7 @@ import Control.Monad.Reader (MonadReader, ask)
 import Control.Monad.State.Strict (MonadState, gets, modify', execState)
 import Control.Monad.ST (RealWorld)
 import Data.Map (Map, toList)
-import Data.Maybe (mapMaybe)
+import Data.Maybe (catMaybes)
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Vector qualified as V
@@ -67,7 +67,7 @@ genTx world txConf deployedContracts = do
   sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
   sender <- rElem' world.senders
   mappedList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
-  (dstAddr, dstAbis) <- rElem' $ Set.fromList $ mapMaybe id mappedList
+  (dstAddr, dstAbis) <- rElem' $ Set.fromList $ catMaybes mappedList
   solCall <- genInteractionsM genDict dstAbis
   value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall
   ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index e591da80c..cbef9a47a 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -3,7 +3,6 @@ module Echidna.Types.CodehashMap where
 import Data.IORef (IORef, readIORef, atomicModifyIORef')
 import Data.Map.Strict qualified as Map
 import Data.Map.Strict (Map)
-import Data.Maybe (fromMaybe)
 import Echidna.Symbolic (forceWord)
 import EVM.Dapp (DappInfo, findSrc)
 import EVM.Solidity (SolcContract(..))
@@ -12,12 +11,14 @@ import EVM.Types (Contract(..), W256)
 type CodehashMap = IORef (Map W256 W256)
 
 lookupCodehash :: CodehashMap -> W256 -> Contract -> DappInfo -> IO W256
-lookupCodehash chmap codehash contr dapp = Map.lookup codehash <$> readIORef chmap >>= \case
-  Just val -> pure val
-  Nothing -> do
-    let originalCodehash = fromMaybe codehash $ (.runtimeCodehash) <$> findSrc contr dapp
-    atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
-    pure originalCodehash
+lookupCodehash chmap codehash contr dapp = do
+  chmapVal <- readIORef chmap
+  case Map.lookup codehash chmapVal of
+    Just val -> pure val
+    Nothing -> do
+      let originalCodehash = maybe codehash (.runtimeCodehash) (findSrc contr dapp)
+      atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
+      pure originalCodehash
 
 lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (W256, Maybe a)
 lookupUsingCodehash chmap contr dapp mapVal = do
@@ -26,7 +27,7 @@ lookupUsingCodehash chmap contr dapp mapVal = do
     ifNotFound codehash' $ pure (codehash', Nothing)
   where
     codehash = forceWord contr.codehash
-    ifNotFound key notFoundCase = case (Map.lookup key mapVal) of
+    ifNotFound key notFoundCase = case Map.lookup key mapVal of
       Nothing -> notFoundCase
       Just val -> pure (key, Just val)
 
@@ -41,6 +42,6 @@ lookupUsingCodehashOrInsert chmap contr dapp mapRef make = do
     applyModification _ Nothing = pure Nothing
     applyModification key (Just val) = atomicModifyIORef' mapRef $ modifyFn key val
 
-    modifyFn key val oldMap = case (Map.lookup key oldMap) of
+    modifyFn key val oldMap = case Map.lookup key oldMap of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)

From b78e6a7af416e0fe36901257f40b4b50e1d7fab6 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Tue, 9 Jan 2024 13:46:00 -0500
Subject: [PATCH 08/13] cleanup

---
 lib/Echidna/Exec.hs              |  4 +++-
 lib/Echidna/Transaction.hs       |  4 ++--
 lib/Echidna/Types/CodehashMap.hs | 21 +++++++++++++++++++--
 lib/Echidna/Types/Config.hs      |  2 +-
 lib/Echidna/Types/Coverage.hs    |  3 ++-
 lib/Echidna/Types/Signature.hs   |  1 +
 6 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index b0b2d49f4..d490daad7 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -302,6 +302,8 @@ execTxWithCov tx = do
             -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
             -- we observed this in some real-world scenarios. This is likely a
             -- bug in another place, investigate.
+            -- ... this should be fixed now, since we use `codeContract` instead
+            -- of `contract` for everything; it may be safe to remove this check.
             when (pc < VMut.length vec) $
               VMut.read vec pc >>= \case
                 (_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
@@ -313,7 +315,7 @@ execTxWithCov tx = do
       -- | Get the VM's current execution location
       currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
 
-      -- | Get the current contract
+      -- | Get the current contract being executed
       currentContract vm = fromMaybe (error "no contract information on coverage") $
         vm ^? #env % #contracts % at vm.state.codeContract % _Just
 
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index eba92a707..0b5b6f6f0 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -66,8 +66,8 @@ genTx world txConf deployedContracts = do
   genDict <- gets (.genDict)
   sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
   sender <- rElem' world.senders
-  mappedList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
-  (dstAddr, dstAbis) <- rElem' $ Set.fromList $ catMaybes mappedList
+  contractAList <- liftIO $ mapM (toContractA env sigMap) (toList deployedContracts)
+  (dstAddr, dstAbis) <- rElem' $ Set.fromList $ catMaybes contractAList
   solCall <- genInteractionsM genDict dstAbis
   value <- genValue txConf.maxValue genDict.dictValues world.payableSigs solCall
   ts <- (,) <$> genDelay txConf.maxTimeDelay genDict.dictValues
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index cbef9a47a..632e28433 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -8,8 +8,15 @@ import EVM.Dapp (DappInfo, findSrc)
 import EVM.Solidity (SolcContract(..))
 import EVM.Types (Contract(..), W256)
 
+-- | Map from contracts' codehashes to their "real" (compile-time) codehash.
+-- This is relevant when the immutables solidity feature is used;
+-- when this feature is not used, the map will just end up being an identity map.
+-- `CodehashMap` is used in signature map and coverage map lookups.
 type CodehashMap = IORef (Map W256 W256)
 
+-- | Lookup a codehash in the `CodehashMap`.
+-- In the case that it's not found, find the "real" (compile-time) codehash and add it to the map.
+-- This is done using hevm's `findSrc` function.
 lookupCodehash :: CodehashMap -> W256 -> Contract -> DappInfo -> IO W256
 lookupCodehash chmap codehash contr dapp = do
   chmapVal <- readIORef chmap
@@ -20,17 +27,26 @@ lookupCodehash chmap codehash contr dapp = do
       atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
       pure originalCodehash
 
+-- | Given a map from codehash to some values of type `a`, lookup a contract in the map using its codehash.
+-- In current use, the `Map W256 a` will be either a `SignatureMap` or a `CoverageMap`.
+-- Returns the "real" codehash, and the map entry if it is found.
 lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (W256, Maybe a)
-lookupUsingCodehash chmap contr dapp mapVal = do
+lookupUsingCodehash chmap contr dapp mapVal =
   ifNotFound codehash $ do
     codehash' <- lookupCodehash chmap codehash contr dapp
-    ifNotFound codehash' $ pure (codehash', Nothing)
+    ifNotFound codehash' $
+      pure (codehash', Nothing)
   where
     codehash = forceWord contr.codehash
     ifNotFound key notFoundCase = case Map.lookup key mapVal of
       Nothing -> notFoundCase
       Just val -> pure (key, Just val)
 
+-- | Same as `lookupUsingCodehash`, except we add to the map if we don't find anything.
+-- The `make` argument is the IO to generate a new element;
+-- it is only run if nothing is found in the map.
+-- In the case that `make` returns `Nothing`, the map will be unchanged.
+-- Returns the map entry, if it is found or generated.
 lookupUsingCodehashOrInsert :: CodehashMap -> Contract -> DappInfo -> IORef (Map W256 a) -> IO (Maybe a) -> IO (Maybe a)
 lookupUsingCodehashOrInsert chmap contr dapp mapRef make = do
   mapVal <- readIORef mapRef
@@ -42,6 +58,7 @@ lookupUsingCodehashOrInsert chmap contr dapp mapRef make = do
     applyModification _ Nothing = pure Nothing
     applyModification key (Just val) = atomicModifyIORef' mapRef $ modifyFn key val
 
+    -- Take care of multithreaded edge case
     modifyFn key val oldMap = case Map.lookup key oldMap of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)
diff --git a/lib/Echidna/Types/Config.hs b/lib/Echidna/Types/Config.hs
index 0aac2a5a5..b85d4f4a0 100644
--- a/lib/Echidna/Types/Config.hs
+++ b/lib/Echidna/Types/Config.hs
@@ -69,9 +69,9 @@ data Env = Env
 
   , testsRef :: IORef [EchidnaTest]
   , coverageRef :: IORef CoverageMap
-  , codehashMap :: CodehashMap
   , corpusRef :: IORef Corpus
 
+  , codehashMap :: CodehashMap
   , fetchContractCache :: IORef (Map Addr (Maybe Contract))
   , fetchSlotCache :: IORef (Map Addr (Map W256 (Maybe W256)))
   , chainId :: Maybe W256
diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index 0a2c82503..25832b2d9 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -11,7 +11,8 @@ import EVM.Types (W256)
 
 import Echidna.Types.Tx (TxResult)
 
--- | Map with the coverage information needed for fuzzing and source code printing
+-- | Map with the coverage information needed for fuzzing and source code printing.
+-- Indexed by contracts' "real" codehash; see `CodehashMap`.
 type CoverageMap = Map W256 (IOVector CoverageInfo)
 
 -- | Basic coverage information
diff --git a/lib/Echidna/Types/Signature.hs b/lib/Echidna/Types/Signature.hs
index c7f38bbec..ad366e608 100644
--- a/lib/Echidna/Types/Signature.hs
+++ b/lib/Echidna/Types/Signature.hs
@@ -28,6 +28,7 @@ type SolCall = (FunctionName, [AbiValue])
 -- | A contract is just an address with an ABI (for our purposes).
 type ContractA = (Addr, NonEmpty SolSignature)
 
+-- | Indexed by contracts' "real" codehash; see `CodehashMap`.
 type SignatureMap = Map W256 (NonEmpty SolSignature)
 
 knownBzzrPrefixes :: [ByteString]

From 514220265d295939b1e5393c62c3d8d5e9ba7195 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Tue, 9 Jan 2024 15:00:28 -0500
Subject: [PATCH 09/13] added test (don't know if it works yet)

---
 src/test/Tests/Integration.hs        |  2 ++
 tests/solidity/basic/immutable-2.sol | 15 +++++++++++++++
 tests/solidity/basic/immutable-3.sol | 12 ++++++++++++
 3 files changed, 29 insertions(+)
 create mode 100644 tests/solidity/basic/immutable-2.sol
 create mode 100644 tests/solidity/basic/immutable-3.sol

diff --git a/src/test/Tests/Integration.hs b/src/test/Tests/Integration.hs
index bb0303609..287c17170 100644
--- a/src/test/Tests/Integration.hs
+++ b/src/test/Tests/Integration.hs
@@ -70,6 +70,8 @@ integrationTests = testGroup "Solidity Integration Testing"
       , ("echidna_timestamp passed",               solved    "echidna_timestamp") ]
   , testContractV "basic/immutable.sol"    (Just (>= solcV (0,6,0))) Nothing
       [ ("echidna_test passed",                    solved      "echidna_test") ]
+  , testContractV "basic/immutable-2.sol"    (Just (>= solcV (0,6,0))) Nothing
+      [ ("echidna_test passed",                    solved      "echidna_test") ]
   , testContract "basic/construct.sol"    Nothing
       [ ("echidna_construct passed",               solved      "echidna_construct") ]
   , testContract "basic/gasprice.sol"     (Just "basic/gasprice.yaml")
diff --git a/tests/solidity/basic/immutable-2.sol b/tests/solidity/basic/immutable-2.sol
new file mode 100644
index 000000000..d7c1d1ee7
--- /dev/null
+++ b/tests/solidity/basic/immutable-2.sol
@@ -0,0 +1,15 @@
+import "./immutable-3.sol";
+
+contract C {
+  D d;
+  constructor() {
+    d = new D(0);
+  }
+  function set(uint256 n, uint256 m) external {
+    d = new D(n);
+    d.set(m);
+  }
+  function echidna_test() public returns (bool) {
+    return d.state();
+  }
+}
diff --git a/tests/solidity/basic/immutable-3.sol b/tests/solidity/basic/immutable-3.sol
new file mode 100644
index 000000000..696dd066b
--- /dev/null
+++ b/tests/solidity/basic/immutable-3.sol
@@ -0,0 +1,12 @@
+contract D {
+  uint256 public immutable n;
+  bool public state = true;
+  constructor(uint256 _n) {
+    n = _n;
+  }
+  function set(uint256 m) external {
+    if (n+1 != 101) revert();
+    if (m+1 != 104) revert();
+    state = false;
+  }
+}

From 6a0e32f7deccb25ea350eb4371f8bc4aa41bdf10 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Wed, 10 Jan 2024 16:02:37 -0500
Subject: [PATCH 10/13] Add @arcz 's suggestions

---
 lib/Echidna/Campaign.hs          | 3 +--
 lib/Echidna/Exec.hs              | 4 ++--
 lib/Echidna/Transaction.hs       | 6 +++---
 lib/Echidna/Types/CodehashMap.hs | 6 +++---
 lib/Echidna/Types/Coverage.hs    | 2 +-
 lib/Echidna/Types/Signature.hs   | 2 +-
 6 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/lib/Echidna/Campaign.hs b/lib/Echidna/Campaign.hs
index 0b739e0e1..9c7c9e144 100644
--- a/lib/Echidna/Campaign.hs
+++ b/lib/Echidna/Campaign.hs
@@ -156,14 +156,13 @@ randseq deployedContracts world = do
 
   let
     mutConsts = env.cfg.campaignConf.mutConsts
-    txConf = env.cfg.txConf
     seqLen = env.cfg.campaignConf.seqLen
 
   -- TODO: include reproducer when optimizing
   --let rs = filter (not . null) $ map (.testReproducer) $ ca._tests
 
   -- Generate new random transactions
-  randTxs <- replicateM seqLen (genTx world txConf deployedContracts)
+  randTxs <- replicateM seqLen (genTx world deployedContracts)
   -- Generate a random mutator
   cmut <- if seqLen == 1 then seqMutatorsStateless (fromConsts mutConsts)
                          else seqMutatorsStateful (fromConsts mutConsts)
diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index d490daad7..7c9bdb240 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -287,7 +287,7 @@ execTxWithCov tx = do
         let (pc, opIx, depth) = currentCovLoc vm
             contract = currentContract vm
 
-        maybeMetaVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
+        maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
           let size = BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
@@ -296,7 +296,7 @@ execTxWithCov tx = do
             forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
             pure $ Just vec
 
-        case maybeMetaVec of
+        case maybeCovVec of
           Nothing -> pure ()
           Just vec -> do
             -- TODO: no-op when pc is out-of-bounds. This shouldn't happen but
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index 0b5b6f6f0..7b8d0b347 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -27,7 +27,7 @@ import Echidna.Orphans.JSON ()
 import Echidna.Symbolic (forceWord, forceAddr)
 import Echidna.Types (fromEVM)
 import Echidna.Types.CodehashMap (lookupUsingCodehash)
-import Echidna.Types.Config (Env(..))
+import Echidna.Types.Config (Env(..), EConfig(..))
 import Echidna.Types.Random
 import Echidna.Types.Signature
   (SignatureMap, SolCall, ContractA)
@@ -58,11 +58,11 @@ getSignatures hmm (Just lmm) =
 genTx
   :: (MonadIO m, MonadRandom m, MonadState WorkerState m, MonadReader Env m)
   => World
-  -> TxConf
   -> Map (Expr EAddr) Contract
   -> m Tx
-genTx world txConf deployedContracts = do
+genTx world deployedContracts = do
   env <- ask
+  let txConf = env.cfg.txConf
   genDict <- gets (.genDict)
   sigMap <- getSignatures world.highSignatureMap world.lowSignatureMap
   sender <- rElem' world.senders
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index 632e28433..aab66fdb8 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -8,14 +8,14 @@ import EVM.Dapp (DappInfo, findSrc)
 import EVM.Solidity (SolcContract(..))
 import EVM.Types (Contract(..), W256)
 
--- | Map from contracts' codehashes to their "real" (compile-time) codehash.
+-- | Map from contracts' codehashes to their compile-time codehash.
 -- This is relevant when the immutables solidity feature is used;
 -- when this feature is not used, the map will just end up being an identity map.
 -- `CodehashMap` is used in signature map and coverage map lookups.
 type CodehashMap = IORef (Map W256 W256)
 
 -- | Lookup a codehash in the `CodehashMap`.
--- In the case that it's not found, find the "real" (compile-time) codehash and add it to the map.
+-- In the case that it's not found, find the compile-time codehash and add it to the map.
 -- This is done using hevm's `findSrc` function.
 lookupCodehash :: CodehashMap -> W256 -> Contract -> DappInfo -> IO W256
 lookupCodehash chmap codehash contr dapp = do
@@ -29,7 +29,7 @@ lookupCodehash chmap codehash contr dapp = do
 
 -- | Given a map from codehash to some values of type `a`, lookup a contract in the map using its codehash.
 -- In current use, the `Map W256 a` will be either a `SignatureMap` or a `CoverageMap`.
--- Returns the "real" codehash, and the map entry if it is found.
+-- Returns the compile-time codehash, and the map entry if it is found.
 lookupUsingCodehash :: CodehashMap -> Contract -> DappInfo -> Map W256 a -> IO (W256, Maybe a)
 lookupUsingCodehash chmap contr dapp mapVal =
   ifNotFound codehash $ do
diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index 25832b2d9..36075b7bd 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -12,7 +12,7 @@ import EVM.Types (W256)
 import Echidna.Types.Tx (TxResult)
 
 -- | Map with the coverage information needed for fuzzing and source code printing.
--- Indexed by contracts' "real" codehash; see `CodehashMap`.
+-- Indexed by contracts' compile-time codehash; see `CodehashMap`.
 type CoverageMap = Map W256 (IOVector CoverageInfo)
 
 -- | Basic coverage information
diff --git a/lib/Echidna/Types/Signature.hs b/lib/Echidna/Types/Signature.hs
index ad366e608..b2638f478 100644
--- a/lib/Echidna/Types/Signature.hs
+++ b/lib/Echidna/Types/Signature.hs
@@ -28,7 +28,7 @@ type SolCall = (FunctionName, [AbiValue])
 -- | A contract is just an address with an ABI (for our purposes).
 type ContractA = (Addr, NonEmpty SolSignature)
 
--- | Indexed by contracts' "real" codehash; see `CodehashMap`.
+-- | Indexed by contracts' compile-time codehash; see `CodehashMap`.
 type SignatureMap = Map W256 (NonEmpty SolSignature)
 
 knownBzzrPrefixes :: [ByteString]

From bad2d741ac6e97dde6960228b49711dba0882529 Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Thu, 11 Jan 2024 12:17:46 -0500
Subject: [PATCH 11/13] fallback on bytecode metadata if findSrc doesn't work

---
 lib/Echidna/Types/CodehashMap.hs     | 44 ++++++++++++++++++++++++++--
 tests/solidity/basic/immutable-2.sol |  4 +--
 tests/solidity/basic/immutable-3.sol |  5 ++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/Types/CodehashMap.hs
index aab66fdb8..50c02a314 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/Types/CodehashMap.hs
@@ -1,12 +1,18 @@
 module Echidna.Types.CodehashMap where
 
+import Control.Applicative ((<|>))
+import Data.ByteString (ByteString)
+import Data.ByteString qualified as BS
 import Data.IORef (IORef, readIORef, atomicModifyIORef')
+import Data.List (find)
 import Data.Map.Strict qualified as Map
 import Data.Map.Strict (Map)
+import Data.Maybe (mapMaybe)
+import Data.Vector qualified as V
 import Echidna.Symbolic (forceWord)
-import EVM.Dapp (DappInfo, findSrc)
+import EVM.Dapp (DappInfo(..), findSrc)
 import EVM.Solidity (SolcContract(..))
-import EVM.Types (Contract(..), W256)
+import EVM.Types (Contract(..), ContractCode(..), RuntimeCode(..), W256, maybeLitByte)
 
 -- | Map from contracts' codehashes to their compile-time codehash.
 -- This is relevant when the immutables solidity feature is used;
@@ -23,7 +29,9 @@ lookupCodehash chmap codehash contr dapp = do
   case Map.lookup codehash chmapVal of
     Just val -> pure val
     Nothing -> do
-      let originalCodehash = maybe codehash (.runtimeCodehash) (findSrc contr dapp)
+      -- hevm's `findSrc` doesn't always work, since `SolcContract.immutableReferences` isn't always populated
+      let solcContract = findSrc contr dapp <|> findSrcByMetadata contr dapp
+          originalCodehash = maybe codehash (.runtimeCodehash) solcContract
       atomicModifyIORef' chmap $ (, ()) . Map.insert codehash originalCodehash
       pure originalCodehash
 
@@ -62,3 +70,33 @@ lookupUsingCodehashOrInsert chmap contr dapp mapRef make = do
     modifyFn key val oldMap = case Map.lookup key oldMap of
       Just val' -> (oldMap, Just val')
       Nothing -> (Map.insert key val oldMap, Just val)
+
+-- | Try to find a SolcContract with a matching bytecode metadata
+findSrcByMetadata :: Contract -> DappInfo -> Maybe SolcContract
+findSrcByMetadata contr dapp = find compareMetadata (snd <$> Map.elems dapp.solcByHash) where
+  compareMetadata solc = contrMeta == Just (getBytecodeMetadata solc.runtimeCode)
+  contrMeta = getBytecodeMetadata <$> contrCode
+  contrCode = case contr.code of
+    (UnknownCode _) -> Nothing
+    (InitCode c _) -> Just c
+    (RuntimeCode (ConcreteRuntimeCode c)) -> Just c
+    (RuntimeCode (SymbolicRuntimeCode c)) -> Just $ BS.pack $ mapMaybe maybeLitByte $ V.toList c
+
+getBytecodeMetadata :: ByteString -> ByteString
+getBytecodeMetadata bs =
+  let stripCandidates = flip BS.breakSubstring bs <$> knownBzzrPrefixes in
+    case find ((/= mempty) . snd) stripCandidates of
+      Nothing     -> bs -- if no metadata is found, return the complete bytecode
+      Just (_, m) -> m
+
+knownBzzrPrefixes :: [ByteString]
+knownBzzrPrefixes =
+  -- a1 65 "bzzr0" 0x58 0x20 (solc <= 0.5.8)
+  [ BS.pack [0xa1, 0x65, 98, 122, 122, 114, 48, 0x58, 0x20]
+  -- a2 65 "bzzr0" 0x58 0x20 (solc >= 0.5.9)
+  , BS.pack [0xa2, 0x65, 98, 122, 122, 114, 48, 0x58, 0x20]
+  -- a2 65 "bzzr1" 0x58 0x20 (solc >= 0.5.11)
+  , BS.pack [0xa2, 0x65, 98, 122, 122, 114, 49, 0x58, 0x20]
+  -- a2 64 "ipfs" 0x58 0x22 (solc >= 0.6.0)
+  , BS.pack [0xa2, 0x64, 0x69, 0x70, 0x66, 0x73, 0x58, 0x22]
+  ]
diff --git a/tests/solidity/basic/immutable-2.sol b/tests/solidity/basic/immutable-2.sol
index d7c1d1ee7..7275027ce 100644
--- a/tests/solidity/basic/immutable-2.sol
+++ b/tests/solidity/basic/immutable-2.sol
@@ -5,9 +5,9 @@ contract C {
   constructor() {
     d = new D(0);
   }
-  function set(uint256 n, uint256 m) external {
+  function set(uint256 n) external {
     d = new D(n);
-    d.set(m);
+    d.set();
   }
   function echidna_test() public returns (bool) {
     return d.state();
diff --git a/tests/solidity/basic/immutable-3.sol b/tests/solidity/basic/immutable-3.sol
index 696dd066b..84b7691a8 100644
--- a/tests/solidity/basic/immutable-3.sol
+++ b/tests/solidity/basic/immutable-3.sol
@@ -4,9 +4,8 @@ contract D {
   constructor(uint256 _n) {
     n = _n;
   }
-  function set(uint256 m) external {
-    if (n+1 != 101) revert();
-    if (m+1 != 104) revert();
+  function set() external {
+    if (n != 1) revert();
     state = false;
   }
 }

From dcbc558ad7859cd33b1d7a404eefe14485b10a5c Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Thu, 11 Jan 2024 12:33:32 -0500
Subject: [PATCH 12/13] fix tests

---
 tests/solidity/basic/immutable-2.sol | 4 ++--
 tests/solidity/basic/immutable-3.sol | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/solidity/basic/immutable-2.sol b/tests/solidity/basic/immutable-2.sol
index 7275027ce..88f16fe2e 100644
--- a/tests/solidity/basic/immutable-2.sol
+++ b/tests/solidity/basic/immutable-2.sol
@@ -2,14 +2,14 @@ import "./immutable-3.sol";
 
 contract C {
   D d;
-  constructor() {
+  constructor() public {
     d = new D(0);
   }
   function set(uint256 n) external {
     d = new D(n);
     d.set();
   }
-  function echidna_test() public returns (bool) {
+  function echidna_test() external returns (bool) {
     return d.state();
   }
 }
diff --git a/tests/solidity/basic/immutable-3.sol b/tests/solidity/basic/immutable-3.sol
index 84b7691a8..5a61c26ff 100644
--- a/tests/solidity/basic/immutable-3.sol
+++ b/tests/solidity/basic/immutable-3.sol
@@ -1,7 +1,7 @@
 contract D {
   uint256 public immutable n;
   bool public state = true;
-  constructor(uint256 _n) {
+  constructor(uint256 _n) public {
     n = _n;
   }
   function set() external {

From 87c635b59e036ad6640d127223138fcd526781bb Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Thu, 11 Jan 2024 13:48:49 -0500
Subject: [PATCH 13/13] rename Echidna.Types.CodehashMap to
 Echidna.SignatureMapping

---
 lib/Echidna/Exec.hs                                    | 2 +-
 lib/Echidna/{Types/CodehashMap.hs => SourceMapping.hs} | 2 +-
 lib/Echidna/Transaction.hs                             | 2 +-
 lib/Echidna/Types/Config.hs                            | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)
 rename lib/Echidna/{Types/CodehashMap.hs => SourceMapping.hs} (99%)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 7c9bdb240..c7495f904 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -31,12 +31,12 @@ import EVM.Types hiding (Env)
 
 import Echidna.Events (emptyEvents)
 import Echidna.RPC (safeFetchContractFrom, safeFetchSlotFrom)
+import Echidna.SourceMapping (lookupUsingCodehashOrInsert)
 import Echidna.Symbolic (forceBuf)
 import Echidna.Transaction
 import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
 import Echidna.Types.Coverage (CoverageInfo)
-import Echidna.Types.CodehashMap (lookupUsingCodehashOrInsert)
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber, getResult)
 import Echidna.Utility (getTimestamp, timePrefix)
diff --git a/lib/Echidna/Types/CodehashMap.hs b/lib/Echidna/SourceMapping.hs
similarity index 99%
rename from lib/Echidna/Types/CodehashMap.hs
rename to lib/Echidna/SourceMapping.hs
index 50c02a314..e0c476d85 100644
--- a/lib/Echidna/Types/CodehashMap.hs
+++ b/lib/Echidna/SourceMapping.hs
@@ -1,4 +1,4 @@
-module Echidna.Types.CodehashMap where
+module Echidna.SourceMapping where
 
 import Control.Applicative ((<|>))
 import Data.ByteString (ByteString)
diff --git a/lib/Echidna/Transaction.hs b/lib/Echidna/Transaction.hs
index 7b8d0b347..3a5dc528e 100644
--- a/lib/Echidna/Transaction.hs
+++ b/lib/Echidna/Transaction.hs
@@ -24,9 +24,9 @@ import EVM.Types hiding (Env, VMOpts(timestamp, gasprice))
 
 import Echidna.ABI
 import Echidna.Orphans.JSON ()
+import Echidna.SourceMapping (lookupUsingCodehash)
 import Echidna.Symbolic (forceWord, forceAddr)
 import Echidna.Types (fromEVM)
-import Echidna.Types.CodehashMap (lookupUsingCodehash)
 import Echidna.Types.Config (Env(..), EConfig(..))
 import Echidna.Types.Random
 import Echidna.Types.Signature
diff --git a/lib/Echidna/Types/Config.hs b/lib/Echidna/Types/Config.hs
index b85d4f4a0..0f2dc39d5 100644
--- a/lib/Echidna/Types/Config.hs
+++ b/lib/Echidna/Types/Config.hs
@@ -12,8 +12,8 @@ import Data.Word (Word64)
 import EVM.Dapp (DappInfo)
 import EVM.Types (Addr, Contract, W256)
 
+import Echidna.SourceMapping (CodehashMap)
 import Echidna.Types.Campaign (CampaignConf, CampaignEvent)
-import Echidna.Types.CodehashMap (CodehashMap)
 import Echidna.Types.Corpus (Corpus)
 import Echidna.Types.Coverage (CoverageMap)
 import Echidna.Types.Solidity (SolConf)