Allow Echidna to collect inputs even after a revert

lib/Echidna/Campaign.hs

@@ -177,7 +177,7 @@ execTxOptC t = do
   memo <- use $ hasLens . bcMemo
   res  <- execTxWith vmExcept (execTxWithCov memo cov) t
   let vmr = getResult $ fst res
-  -- Update the coverage map with the proper binary according to the vm result
+  -- Update the coverage map with the proper result according to the vm result
   cov %= mapWithKey (\_ s -> DS.map (set _4 vmr) s)
   -- Update the global coverage map with the union of the result just obtained
   cov %= unionWith DS.union og

lib/Echidna/Exec.hs

@@ -14,15 +14,16 @@ import Data.Has (Has(..))
 import Data.Maybe (fromMaybe)
 import EVM
 import EVM.Exec (exec, vmForEthrunCreation)
-import EVM.Types (Buffer(..))
+import EVM.Types (Buffer(..), maybeLitWord)
+import EVM.Op (Op(..))
 import EVM.Symbolic (litWord)
 
 import qualified Data.Map as M
 import qualified Data.Set as S
 
 import Echidna.Transaction
 import Echidna.Types.Buffer (viewBuffer)
-import Echidna.Types.Coverage (CoverageMap)
+import Echidna.Types.Coverage (CoverageMap, RevertsSkipped)
 import Echidna.Types.Tx (TxCall(..), Tx, TxResult(..), call, dst, initialTimestamp, initialBlockNumber)
 
 import Echidna.Types.Signature (BytecodeMemo, lookupBytecodeMetadata)
@@ -148,34 +149,47 @@ handleErrorsAndConstruction onErr vmResult' vmBeforeTx tx' = case (vmResult', tx
 execTx :: (MonadState x m, Has VM x, MonadThrow m) => Tx -> m (VMResult, Int)
 execTx = execTxWith vmExcept $ liftSH exec
 
+skipRevert :: VM -> VM
+skipRevert jvm = case stk of
+  (x:y:xs) -> case maybeLitWord y of
+                    Just y' -> chstk jvm (x : litWord (1 - y') : xs)
+                    _       -> error "symbolic value?"
+  _        -> error "invalid stack?"
+  where stk = jvm ^. state . stack
+        chstk vm xs = set (state . stack) xs vm
+
+willJumpi :: VM -> Bool
+willJumpi vm = vmOp vm == Just OpJumpi
+
 -- | Execute a transaction, logging coverage at every step.
 execTxWithCov :: (MonadState x m, Has VM x) => BytecodeMemo -> Lens' x CoverageMap -> m VMResult
 execTxWithCov memo l = do
   vm :: VM          <- use hasLens
   cm :: CoverageMap <- use l
-  let (r, vm', cm') = loop vm cm
+  let (r, vm', cm') = loop vm vm 0 cm
   hasLens .= vm'
   l       .= cm'
   return r
   where
     -- | Repeatedly exec a step and add coverage until we have an end result
-    loop :: VM -> CoverageMap -> (VMResult, VM, CoverageMap)
-    loop vm cm = case _result vm of
-      Nothing  -> loop (stepVM vm) (addCoverage vm cm)
-      Just r   -> (r, vm, cm)
+    loop :: VM -> VM -> RevertsSkipped -> CoverageMap -> (VMResult, VM, CoverageMap)
+    loop jvm vm rs cm = case _result vm of
+      Nothing                                       -> loop (if willJumpi vm then vm else jvm) (stepVM vm) rs (addCoverage vm False rs cm)
+      Just f@(VMFailure (Revert _)) | willJumpi jvm -> let (_, x, y) = loop vm (stepVM $ skipRevert jvm) (rs+1) (addCoverage jvm False (rs+1) cm) in (f, x, y)
+      Just r                                        -> (r, vm, addCoverage vm True rs cm)
 
     -- | Execute one instruction on the EVM
     stepVM :: VM -> VM
     stepVM = execState exec1
 
     -- | Add current location to the CoverageMap
-    addCoverage :: VM -> CoverageMap -> CoverageMap
-    addCoverage vm = M.alter
-                       (Just . maybe mempty (S.insert $ currentCovLoc vm))
+    addCoverage :: VM -> Bool -> RevertsSkipped -> CoverageMap -> CoverageMap
+    addCoverage vm ends rs = M.alter
+                       (Just . maybe mempty (S.insert $ currentCovLoc vm ends rs))
                        (currentMeta vm)
 
     -- | Get the VM's current execution location
-    currentCovLoc vm = (vm ^. state . pc, fromMaybe 0 $ vmOpIx vm, length $ vm ^. frames, Stop)
+    currentCovLoc vm ends rs = (vm ^. state . pc, fromMaybe 0 $ vmOpIx vm, length $ vm ^. frames, Stop, ends, rs)
 
     -- | Get the current contract's bytecode metadata
     currentMeta vm = fromMaybe (error "no contract information on coverage") $ do

lib/Echidna/Output/Source.hs

@@ -23,7 +23,7 @@ import qualified Data.Map as M
 import qualified Data.Set as S
 import qualified Data.Text as T
 
-import Echidna.Types.Coverage (CoverageMap, CoverageInfo)
+import Echidna.Types.Coverage (CoverageMap, CoverageInfo, TxEnded)
 import Echidna.Types.Tx (TxResult(..))
 import Echidna.Types.Signature (getBytecodeMetadata)
 
@@ -52,24 +52,29 @@ ppCoveredCode sc cs s | s == mempty = "Coverage map is empty"
   in T.intercalate "\n" $ map ppFile allFiles
 
 -- | Mark one particular line, from a list of lines, keeping the order of them
-markLines :: V.Vector Text -> M.Map Int [TxResult] -> V.Vector Text
+markLines :: V.Vector Text -> M.Map Int [(TxResult, TxEnded)] -> V.Vector Text
 markLines codeLines resultMap = V.map markLine (V.indexed codeLines)
   where
     markLine (i, codeLine) =
       let results = fromMaybe [] (M.lookup (i+1) resultMap)
-      in pack $ printf "%-4s| %s" (sort $ nub $ getMarker <$> results) (unpack codeLine)
+      in pack $ printf "%-4s| %s" (sort $ simplifyMarkers $ nub $ getMarker <$> results) (unpack codeLine)
 
 -- | Select the proper marker, according to the result of the transaction
-getMarker :: TxResult -> Char
-getMarker ReturnTrue    = '*'
-getMarker ReturnFalse   = '*'
-getMarker Stop          = '*'
-getMarker ErrorRevert   = 'r' 
-getMarker ErrorOutOfGas = 'o'
-getMarker _             = 'e'
+getMarker :: (TxResult, TxEnded) -> Char
+getMarker (ReturnTrue,  _)   = '*'
+getMarker (ReturnFalse, _)   = '*'
+getMarker (Stop, _)          = '*'
+getMarker (ErrorRevert, b)   = if b then 'R' else 'r'
+getMarker (ErrorOutOfGas, _) = 'o'
+getMarker _                  = 'e'
+
+simplifyMarkers :: String -> String
+simplifyMarkers cs = if 'R' `elem` cs && 'r' `elem` cs
+  then filter (/= 'r') cs
+  else cs
 
 -- | 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 :: SourceCache -> CoverageMap -> [SolcContract] -> M.Map FilePathText (M.Map Int [(TxResult, TxEnded)])
 srcMapCov sc s contracts =
   M.map (M.fromListWith (++)) .
   M.fromListWith (++) .
@@ -84,13 +89,13 @@ srcMapCov sc s contracts =
       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 :: SourceCache -> (SrcMap, (TxResult, TxEnded)) -> Maybe (Text, Int, (TxResult, TxEnded))
 srcMapCodePosResult sc (n, r) = case srcMapCodePos sc n of
   Just (t,n') -> Just (t,n',r)
   _           -> Nothing
 
 -- | Given a contract, and tuple as coverage, return the corresponding mapped line (if any)
-srcMapForOpLocation :: SolcContract -> CoverageInfo -> Maybe (SrcMap, TxResult)
-srcMapForOpLocation c (_,n,_,r) = case preview (ix n) (c ^. runtimeSrcmap <> c ^. creationSrcmap) of
-  Just sm -> Just (sm,r)
+srcMapForOpLocation :: SolcContract -> CoverageInfo -> Maybe (SrcMap, (TxResult, TxEnded))
+srcMapForOpLocation c (_,n,_,r,b,_) = case preview (ix n) (c ^. runtimeSrcmap <> c ^. creationSrcmap) of
+  Just sm -> Just (sm,(r,b))
   _       -> Nothing

lib/Echidna/Types/Coverage.hs

@@ -13,8 +13,12 @@ type PC = Int
 type OpIx = Int
 -- Stack size from the EVM
 type FrameCount = Int
+-- Has the transaction ended
+type TxEnded = Bool
+-- How many reverts we skipped
+type RevertsSkipped = Int
 -- Basic coverage information
-type CoverageInfo = (PC, OpIx, FrameCount, TxResult)
+type CoverageInfo = (PC, OpIx, FrameCount, TxResult, TxEnded, RevertsSkipped)
 -- Map with the coverage information needed for fuzzing and source code printing 
 type CoverageMap = Map ByteString (Set CoverageInfo)
 

tests/solidity/basic/revert.sol

@@ -2,6 +2,7 @@ contract C {
 
   int private state = 0;
   function f(int x, address y, address z) public {
+    require(x > 0 || x <= 0);
     require(z != address(0x0));
     state = x;
   }