Collect coverage during deploy

lib/Echidna/Deploy.hs

@@ -3,7 +3,7 @@ module Echidna.Deploy where
 import Control.Monad (foldM)
 import Control.Monad.Catch (MonadThrow(..), throwM)
 import Control.Monad.Reader (MonadReader, asks)
-import Control.Monad.State.Strict (MonadIO)
+import Control.Monad.State.Strict (MonadIO, runStateT)
 import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
 import Data.ByteString.Base16 qualified as BS16 (decode)
@@ -14,7 +14,7 @@ import Data.Text.Encoding (encodeUtf8)
 import EVM.Solidity
 import EVM.Types hiding (Env)
 
-import Echidna.Exec (execTx)
+import Echidna.Exec (execTxWithCov)
 import Echidna.Events (extractEvents)
 import Echidna.Types.Config (Env(..))
 import Echidna.Types.Solidity (SolException(..))
@@ -51,7 +51,7 @@ deployBytecodes' cs src initialVM = foldM deployOne initialVM cs
   where
   deployOne vm (dst, bytecode) = do
     (_, vm') <-
-      execTx vm $ createTx (bytecode <> zeros) src dst unlimitedGasPerBlock (0, 0)
+      runStateT (execTxWithCov $ createTx (bytecode <> zeros) src dst unlimitedGasPerBlock (0, 0)) vm
     case vm'.result of
       Just (VMSuccess _) -> pure vm'
       _ -> do

lib/Echidna/Exec.hs

@@ -13,27 +13,27 @@ import Control.Monad.State.Strict (MonadState(get, put), execState, runStateT, M
 import Control.Monad.Reader (MonadReader, ask, asks)
 import Control.Monad.ST (ST, stToIO, RealWorld)
 import Data.Bits
-import Data.ByteString qualified as BS
 import Data.IORef (readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
 import Data.Map qualified as Map
-import Data.Maybe (fromMaybe, fromJust)
+import Data.Maybe (fromMaybe)
 import Data.Text qualified as T
 import Data.Vector qualified as V
 import Data.Vector.Unboxed.Mutable qualified as VMut
 import System.Process (readProcessWithExitCode)
 
-import EVM (bytecode, replaceCodeOfSelf, loadContract, exec1, vmOpIx, clearTStorages)
+import EVM (replaceCodeOfSelf, loadContract, exec1, vmOpIx, clearTStorages)
 import EVM.ABI
-import EVM.Dapp (DappInfo)
+import EVM.Dapp (DappInfo(..))
 import EVM.Exec (exec, vmForEthrunCreation)
 import EVM.Fetch qualified
 import EVM.Format (hexText, showTraceTree)
+import EVM.Solidity (SolcContract(..))
 import EVM.Types hiding (Env, Gas)
 
 import Echidna.Events (emptyEvents)
 import Echidna.Onchain (safeFetchContractFrom, safeFetchSlotFrom)
-import Echidna.SourceMapping (lookupUsingCodehashOrInsert)
-import Echidna.Symbolic (forceBuf)
+import Echidna.SourceMapping (lookupUsingCodehashOrInsert, lookupCodehash)
+import Echidna.Symbolic (forceWord)
 import Echidna.Transaction
 import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
@@ -285,11 +285,17 @@ execTxWithCov tx = do
       -- | Add current location to the CoverageMap
       addCoverage :: VM Concrete RealWorld -> IO ()
       addCoverage !vm = do
-        let (pc, opIx, depth) = currentCovLoc vm
+        let
           contract = currentContract vm
+          getCodehash = lookupCodehash env.codehashMap (forceWord contract.codehash) contract env.dapp
+          getContractLengths = contractLengthsFromEnv <$> getCodehash
+          (pc, opIx_, depth) = currentCovLoc vm
+
+        opIx <- if isDeploy vm then (opIx_ +) . fst <$> getContractLengths else pure opIx_
 
         maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
+          (sizeA, sizeB) <- getContractLengths
+          let size = sizeA + sizeB
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
             vec <- VMut.new size
@@ -316,6 +322,13 @@ execTxWithCov tx = do
       -- | Get the VM's current execution location
       currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
 
+      isDeploy vm = case (.code) (currentContract vm) of
+        InitCode _ _ -> True
+        _ -> False
+
+      contractLengthsFromEnv codehash = maybe (0,0) (contractLengths . snd) $ Map.lookup codehash env.dapp.solcByHash
+      contractLengths contract = (length contract.runtimeSrcmap, length contract.creationSrcmap)
+
       -- | Get the current contract being executed
       currentContract vm = fromMaybe (error "no contract information on coverage") $
         vm ^? #env % #contracts % at vm.state.codeContract % _Just

lib/Echidna/Solidity.hs

@@ -4,6 +4,7 @@ import Optics.Core hiding (filtered)
 
 import Control.Monad (when, unless, forM_)
 import Control.Monad.Catch (MonadThrow(..))
+import Control.Monad.State (runStateT)
 import Control.Monad.Extra (whenM)
 import Control.Monad.Reader (ReaderT(runReaderT))
 import Control.Monad.ST (stToIO, RealWorld)
@@ -39,7 +40,7 @@ import Echidna.ABI
 import Echidna.Deploy (deployContracts, deployBytecodes)
 import Echidna.Etheno (loadEthenoBatch)
 import Echidna.Events (extractEvents)
-import Echidna.Exec (execTx, initialVM)
+import Echidna.Exec (execTx, initialVM, execTxWithCov)
 import Echidna.SourceAnalysis.Slither
 import Echidna.Test (createTests, isAssertionMode, isPropertyMode, isDapptestMode)
 import Echidna.Types.Config (EConfig(..), Env(..))
@@ -199,13 +200,13 @@ loadSpecified env mainContract cs = do
     vm2 <- deployBytecodes solConf.deployBytecodes solConf.deployer vm1
 
     -- main contract deployment
-    let deployment = execTx vm2 $ createTxWithValue
+    let deployment = runStateT (execTxWithCov (createTxWithValue
                                     mainContract.creationCode
                                     solConf.deployer
                                     solConf.contractAddr
                                     unlimitedGasPerBlock
                                     (fromIntegral solConf.balanceContract)
-                                    (0, 0)
+                                    (0, 0))) vm2
     (_, vm3) <- deployment
     when (isNothing $ currentContract vm3) $
       throwM $ DeploymentFailed solConf.contractAddr $ T.unlines $ extractEvents True env.dapp vm3