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
- contract = currentContract 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