Mutable coverage (#1003)

lib/Echidna/Campaign.hs

@@ -40,10 +40,9 @@ import Echidna.Types.Buffer (forceBuf)
 import Echidna.Types.Corpus (Corpus)
 import Echidna.Types.Campaign
 import Echidna.Types.Config
-import Echidna.Types.Coverage (coveragePoints)
 import Echidna.Types.Signature (makeBytecodeCache, FunctionName)
 import Echidna.Types.Test
-import Echidna.Types.Tx (TxCall(..), Tx(..), getResult, call)
+import Echidna.Types.Tx (TxCall(..), Tx(..), call)
 import Echidna.Types.World (World)
 
 instance MonadThrow m => MonadThrow (RandT g m) where
@@ -53,12 +52,12 @@ instance MonadCatch m => MonadCatch (RandT g m) where
 
 -- | Given a 'Campaign', checks if we can attempt any solves or shrinks without exceeding
 -- the limits defined in our 'CampaignConf'.
-isDone :: MonadReader EConfig m => Campaign -> m Bool
+isDone :: MonadReader Env m => GenericCampaign a -> m Bool
 isDone c | null c.tests = do
-  conf <- asks (.campaignConf)
+  conf <- asks (.cfg.campaignConf)
   pure $ c.ncallseqs * conf.seqLen >= conf.testLimit
 isDone c = do
-  conf <- asks (.campaignConf)
+  conf <- asks (.cfg.campaignConf)
   let
     result = \case
       Open i   -> if i >= conf.testLimit then Just True else Nothing
@@ -288,15 +287,9 @@ execTxOptC
   -> m (VMResult, Gas)
 execTxOptC tx = do
   (vm, camp@Campaign{coverage = oldCov}) <- get
-  ((res, txCov), vm') <- runStateT (execTxWithCov tx) vm
-  let
-    vmr = getResult $ fst res
-    -- Update the tx coverage map with the proper binary according to the vm result
-    txCov' = Map.mapWithKey (\_ s -> Set.map (set _4 vmr) s) txCov
-    -- Update the global coverage map with the one from this tx run
-    newCov = Map.unionWith Set.union oldCov txCov'
-  put (vm', camp { coverage = newCov })
-  when (coveragePoints oldCov < coveragePoints newCov) $ do
+  ((res, (cov', grew)), vm') <- runStateT (execTxWithCov tx oldCov) vm
+  put (vm', camp { coverage = cov' })
+  when grew $ do
     let dict' = case tx.call of
           SolCall c -> gaddCalls (Set.singleton c) camp.genDict
           _ -> camp.genDict

lib/Echidna/Exec.hs

@@ -5,20 +5,21 @@
 module Echidna.Exec where
 
 import Control.Lens
-import Control.Monad (when)
+import Control.Monad (forM_, when)
 import Control.Monad.Catch (MonadThrow(..))
 import Control.Monad.State.Strict (MonadState(get, put), execState, runStateT, MonadIO(liftIO))
 import Control.Monad.Reader (MonadReader, asks)
+import Data.Bits
+import Data.ByteString qualified as BS
 import Data.IORef (readIORef, atomicWriteIORef)
-import Data.Map qualified as M
 import Data.Map qualified as Map
-import Data.Maybe (fromMaybe)
-import Data.Set qualified as S
+import Data.Maybe (fromMaybe, fromJust)
 import Data.Text qualified as T
 import Data.Vector qualified as V
+import Data.Vector.Unboxed.Mutable qualified as V
 import System.Process (readProcessWithExitCode)
 
-import EVM hiding (Env, cache, contract, tx, value)
+import EVM hiding (pc, Env, cache, contract, tx, value)
 import EVM.ABI
 import EVM.Exec (exec, vmForEthrunCreation)
 import EVM.Fetch qualified
@@ -31,7 +32,7 @@ import Echidna.Types (ExecException(..), Gas, fromEVM, emptyAccount)
 import Echidna.Types.Buffer (forceBuf)
 import Echidna.Types.Coverage (CoverageMap)
 import Echidna.Types.Signature (MetadataCache, getBytecodeMetadata, lookupBytecodeMetadata)
-import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber)
+import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber, getResult)
 import Echidna.Types.Config (Env(..), EConfig(..), UIConf(..), OperationMode(..), OutputFormat(Text))
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Utility (timePrefix)
@@ -225,44 +226,85 @@ execTx
   -> m (VMResult, Gas)
 execTx = execTxWith id vmExcept $ fromEVM exec
 
+-- | A type alias for the context we carry while executing instructions
+type CoverageContext = (CoverageMap, Bool, Maybe (BS.ByteString, Int))
+
 -- | Execute a transaction, logging coverage at every step.
 execTxWithCov
   :: (MonadIO m, MonadState VM m, MonadReader Env m, MonadThrow m)
-  => Tx
-  -> m ((VMResult, Gas), CoverageMap)
-execTxWithCov tx = do
+  => Tx -> CoverageMap
+  -> m ((VMResult, Gas), (CoverageMap, Bool))
+execTxWithCov tx cov = do
   vm <- get
   metaCacheRef <- asks (.metadataCache)
   cache <- liftIO $ readIORef metaCacheRef
-  (r, (vm', cm)) <- runStateT (execTxWith _1 vmExcept (execCov cache) tx) (vm, mempty)
+  (r, (vm', (cm, grew, lastLoc))) <-
+    runStateT (execTxWith _1 vmExcept (execCov cache) tx) (vm, (cov, False, Nothing))
   put vm'
-  pure (r, cm)
+
+  -- Update the last valid location with the transaction result
+  grew' <- liftIO $ case lastLoc of
+    Just (meta, pc) ->
+      case Map.lookup meta cm of
+        Nothing -> pure False -- shouldn't happen
+        Just vec -> do
+          let txResultBit = fromEnum $ getResult $ fst r
+          V.read vec pc >>= \case
+            (opIx, depths, txResults) | not (txResults `testBit` txResultBit) -> do
+              V.write vec pc (opIx, depths, txResults `setBit` txResultBit)
+              pure True -- we count this as new coverage
+            _ -> pure False
+    _ -> pure False
+
+  pure (r, (cm, grew || grew'))
   where
     -- the same as EVM.exec but collects coverage, will stop on a query
     execCov cache = do
      (vm, cm) <- get
-     let (r, vm', cm') = loop cache vm cm
+     (r, vm', cm') <- liftIO $ loop cache vm cm
      put (vm', cm')
      pure r
 
     -- | Repeatedly exec a step and add coverage until we have an end result
-    loop :: MetadataCache -> VM -> CoverageMap -> (VMResult, VM, CoverageMap)
-    loop cache vm cm = case vm._result of
-      Nothing  -> loop cache (stepVM vm) (addCoverage cache vm cm)
-      Just r   -> (r, vm, cm)
+    loop :: MetadataCache -> VM -> CoverageContext -> IO (VMResult, VM, CoverageContext)
+    loop cache !vm !cc = case vm._result of
+      Nothing -> addCoverage cache vm cc >>= loop cache (stepVM vm)
+      Just r -> pure (r, vm, cc)
 
     -- | Execute one instruction on the EVM
     stepVM :: VM -> VM
     stepVM = execState exec1
 
     -- | Add current location to the CoverageMap
-    addCoverage :: MetadataCache -> VM -> CoverageMap -> CoverageMap
-    addCoverage cache vm = M.alter
-                       (Just . maybe mempty (S.insert $ currentCovLoc vm))
-                       (currentMeta cache vm)
+    addCoverage :: MetadataCache -> VM -> CoverageContext -> IO CoverageContext
+    addCoverage cache !vm (!cm, new, lastLoc) = do
+      let (pc, opIx, depth) = currentCovLoc vm
+          meta = currentMeta cache vm
+      case Map.lookup meta cm of
+        Nothing -> do
+          let size = BS.length . forceBuf . view bytecode . fromJust $
+                Map.lookup vm._state._contract vm._env._contracts
+          if size > 0 then do
+            vec <- V.new size
+            -- We use -1 for opIx to indicate that the location was not covered
+            forM_ [0..size-1] $ \i -> V.write vec i (-1, 0, 0)
+            V.write vec pc (opIx, fromIntegral depth, 0 `setBit` fromEnum Stop)
+            pure (Map.insert meta vec cm, True, Just (meta, 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 (cm, new, lastLoc)
+        Just vec -> do
+          V.read vec pc >>= \case
+            (_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
+              V.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
+              pure (cm, True, Just (meta, pc))
+            _ ->
+              pure (cm, new, Just (meta, pc))
 
     -- | Get the VM's current execution location
-    currentCovLoc vm = (vm._state._pc, fromMaybe 0 $ vmOpIx vm, length vm._frames, Stop)
+    currentCovLoc vm = (vm._state._pc, fromMaybe 0 $ vmOpIx vm, length vm._frames)
 
     -- | Get the current contract's bytecode metadata
     currentMeta cache vm = fromMaybe (error "no contract information on coverage") $ do

lib/Echidna/Output/JSON.hs

@@ -7,9 +7,9 @@ import Data.ByteString.Base16 qualified as BS16
 import Data.ByteString.Lazy (ByteString)
 import Data.Map (Map)
 import Data.Map qualified as Map
-import Data.Set qualified as Set
 import Data.Text
 import Data.Text.Encoding (decodeUtf8)
+import Data.Vector.Unboxed qualified as VU
 import Numeric (showHex)
 
 import EVM.Types (keccak')
@@ -93,14 +93,15 @@ instance ToJSON Transaction where
     , "gasprice" .= gasprice
     ]
 
-encodeCampaign :: C.Campaign -> ByteString
-encodeCampaign C.Campaign{..} = encode
-  Campaign
+encodeCampaign :: C.Campaign -> IO ByteString
+encodeCampaign C.Campaign{..} = do
+  frozenCov <- mapM VU.freeze coverage
+  pure $ encode Campaign
     { _success = True
     , _error = Nothing
     , _tests = mapTest <$> tests
     , seed = genDict.defSeed
-    , coverage = Map.mapKeys (("0x" ++) . (`showHex` "") . keccak') $ Set.toList <$> coverage
+    , coverage = Map.mapKeys (("0x" ++) . (`showHex` "") . keccak') $ VU.toList <$> frozenCov
     , gasInfo = Map.toList gasInfo
     }
 

lib/Echidna/Output/Source.hs

@@ -3,16 +3,18 @@ module Echidna.Output.Source where
 import Prelude hiding (writeFile)
 
 import Data.Foldable
-import Data.Maybe (fromMaybe, mapMaybe)
 import Data.List (nub, sort)
-import Data.Map qualified as M
+import Data.Maybe (fromMaybe, mapMaybe)
+import Data.Map (Map)
+import Data.Map qualified as Map
 import Data.Sequence qualified as Seq
 import Data.Set qualified as S
 import Data.Text (Text, pack)
 import Data.Text qualified as T
 import Data.Text.Encoding (decodeUtf8)
 import Data.Text.IO (writeFile)
 import Data.Vector qualified as V
+import Data.Vector.Unboxed qualified as VU
 import HTMLEntities.Text qualified as HTML
 import System.Directory (createDirectoryIfMissing)
 import System.FilePath ((</>))
@@ -21,24 +23,25 @@ import Text.Printf (printf)
 import EVM.Debug (srcMapCodePos)
 import EVM.Solidity (SourceCache(..), SrcMap, SolcContract(..))
 
-import Echidna.Types.Coverage (CoverageMap, CoverageInfo)
+import Echidna.Types.Coverage (CoverageMap, FrozenCoverageMap, OpIx, unpackTxResults)
 import Echidna.Types.Tx (TxResult(..))
 import Echidna.Types.Signature (getBytecodeMetadata)
 
 type FilePathText = Text
 
 saveCoverage :: Bool -> Int -> FilePath -> SourceCache -> [SolcContract] -> CoverageMap -> IO ()
-saveCoverage isHtml seed d sc cs s = let extension = if isHtml then ".html" else ".txt"
-                                         fn = d </> "covered." <> show seed <> extension
-                                         cc = ppCoveredCode isHtml sc cs s
-                                     in do
-                                       createDirectoryIfMissing True d
-                                       writeFile fn cc
+saveCoverage isHtml seed d sc cs covMap = do
+  frozenCovMap <- mapM VU.freeze covMap
+  let extension = if isHtml then ".html" else ".txt"
+      fn = d </> "covered." <> show seed <> extension
+      cc = ppCoveredCode isHtml sc cs frozenCovMap
+  createDirectoryIfMissing True d
+  writeFile fn cc
 
 -- | Pretty-print the covered code
-ppCoveredCode :: Bool -> SourceCache -> [SolcContract] -> CoverageMap -> Text
-ppCoveredCode isHtml sc cs s | s == mempty = "Coverage map is empty"
-                             | otherwise   =
+ppCoveredCode :: Bool -> SourceCache -> [SolcContract] -> FrozenCoverageMap -> Text
+ppCoveredCode isHtml sc cs s | null s    = "Coverage map is empty"
+                             | otherwise =
   let allFiles = zipWith (\(srcPath, _rawSource) srcLines -> (srcPath, V.map decodeUtf8 srcLines))
                    sc.files
                    sc.lines
@@ -48,8 +51,8 @@ ppCoveredCode isHtml sc cs s | s == mempty = "Coverage map is empty"
       runtimeLinesMap = buildRuntimeLinesMap sc cs
       -- ^ Excludes lines such as comments or blanks
       ppFile (srcPath, srcLines) =
-        let runtimeLines = fromMaybe mempty $ M.lookup srcPath runtimeLinesMap
-            marked = markLines isHtml srcLines runtimeLines (fromMaybe M.empty (M.lookup srcPath covLines))
+        let runtimeLines = fromMaybe mempty $ Map.lookup srcPath runtimeLinesMap
+            marked = markLines isHtml srcLines runtimeLines (fromMaybe mempty (Map.lookup srcPath covLines))
         in T.unlines (changeFileName srcPath : changeFileLines (V.toList marked))
       -- ^ Pretty print individual file coverage
       topHeader
@@ -72,13 +75,13 @@ ppCoveredCode isHtml sc cs s | s == mempty = "Coverage map is empty"
   in topHeader <> T.unlines (map ppFile allFiles)
 
 -- | Mark one particular line, from a list of lines, keeping the order of them
-markLines :: Bool -> V.Vector Text -> S.Set Int -> M.Map Int [TxResult] -> V.Vector Text
+markLines :: Bool -> V.Vector Text -> S.Set Int -> Map Int [TxResult] -> V.Vector Text
 markLines isHtml codeLines runtimeLines resultMap =
   V.map markLine (V.indexed codeLines)
   where
   markLine (i, codeLine) =
     let n = i + 1
-        results  = fromMaybe [] (M.lookup n resultMap)
+        results  = fromMaybe [] (Map.lookup n resultMap)
         markers = sort $ nub $ getMarker <$> results
         wrapLine :: Text -> Text
         wrapLine line
@@ -109,38 +112,39 @@ getMarker ErrorOutOfGas = 'o'
 getMarker _             = 'e'
 
 -- | Given a source cache, a coverage map, a contract returns a list of covered lines
-srcMapCov :: SourceCache -> CoverageMap -> [SolcContract] -> M.Map FilePathText (M.Map Int [TxResult])
-srcMapCov sc s contracts =
-  M.map (M.fromListWith (++)) .
-  M.fromListWith (++) .
-  map (\(srcPath, line, txResult) -> (srcPath, [(line, [txResult])])) .
-  nub .                                              -- Deduplicate results
-  mapMaybe (srcMapCodePosResult sc) $                -- Get the filename, number of line and tx result
-  concatMap mapContract contracts
+srcMapCov :: SourceCache -> FrozenCoverageMap -> [SolcContract] -> Map FilePathText (Map Int [TxResult])
+srcMapCov sc covMap contracts =
+  Map.unionsWith Map.union $ linesCovered <$> contracts
   where
-    mapContract c =
-      mapMaybe (srcMapForOpLocation c) .             -- Get the mapped line and tx result
-      S.toList . fromMaybe S.empty $                 -- Convert from Set to list
-      M.lookup (getBytecodeMetadata c.runtimeCode) s -- Get the coverage information of the current contract
-
--- | Given a source cache, a mapped line, return a tuple with the filename, number of line and tx result
-srcMapCodePosResult :: SourceCache -> (SrcMap, TxResult) -> Maybe (Text, Int, TxResult)
-srcMapCodePosResult sc (n, r) = case srcMapCodePos sc n of
-  Just (t,n') -> Just (t,n',r)
-  _           -> Nothing
+  linesCovered :: SolcContract -> Map Text (Map Int [TxResult])
+  linesCovered c =
+    case Map.lookup (getBytecodeMetadata c.runtimeCode) covMap of
+      Just vec -> VU.foldl' (\acc covInfo -> case covInfo of
+        (-1, _, _) -> acc -- not covered
+        (opIx, _stackDepths, txResults) ->
+          case srcMapForOpLocation c opIx of
+            Just srcMap ->
+              case srcMapCodePos sc srcMap of
+                Just (file, line) ->
+                  Map.alter
+                    (Just . Map.insert line (unpackTxResults txResults) . fromMaybe mempty)
+                    file
+                    acc
+                Nothing -> acc
+            Nothing -> acc
+        ) mempty vec
+      Nothing -> mempty
 
 -- | Given a contract, and tuple as coverage, return the corresponding mapped line (if any)
-srcMapForOpLocation :: SolcContract -> CoverageInfo -> Maybe (SrcMap, TxResult)
-srcMapForOpLocation contract (_,n,_,r) =
-  case Seq.lookup n (contract.runtimeSrcmap <> contract.creationSrcmap) of
-    Just sm -> Just (sm,r)
-    _       -> Nothing
+srcMapForOpLocation :: SolcContract -> OpIx -> Maybe SrcMap
+srcMapForOpLocation contract opIx =
+  Seq.lookup opIx (contract.runtimeSrcmap <> contract.creationSrcmap)
 
 -- | Builds a Map from file paths to lines that can be executed, this excludes
 -- for example lines with comments
-buildRuntimeLinesMap :: SourceCache -> [SolcContract] -> M.Map Text (S.Set Int)
+buildRuntimeLinesMap :: SourceCache -> [SolcContract] -> Map Text (S.Set Int)
 buildRuntimeLinesMap sc contracts =
-  M.fromListWith (<>)
+  Map.fromListWith (<>)
     [(k, S.singleton v) | (k, v) <- mapMaybe (srcMapCodePos sc) srcMaps]
   where
   srcMaps = concatMap

lib/Echidna/Transaction.hs

@@ -8,6 +8,7 @@ import Control.Monad (join)
 import Control.Monad.Random.Strict (MonadRandom, getRandomR, uniform)
 import Control.Monad.State.Strict (MonadState, gets, modify')
 import Data.Map (Map, toList)
+import Data.Map qualified as Map
 import Data.Maybe (mapMaybe)
 import Data.Set (Set)
 import Data.Set qualified as Set
@@ -24,8 +25,7 @@ import Echidna.Types.Buffer (forceBuf, forceLit)
 import Echidna.Types.Signature (SignatureMap, SolCall, ContractA, FunctionHash, MetadataCache, lookupBytecodeMetadata)
 import Echidna.Types.Tx
 import Echidna.Types.World (World(..))
-import Echidna.Types.Campaign (Campaign(..))
-import qualified Data.Map as Map
+import Echidna.Types.Campaign
 
 hasSelfdestructed :: VM -> Addr -> Bool
 hasSelfdestructed vm addr = addr `elem` vm._tx._substate._selfdestructs