Simplify Exec.hs

lib/Echidna/Exec.hs

@@ -14,7 +14,7 @@ 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.IORef (IORef, readIORef, atomicWriteIORef, newIORef, writeIORef, modifyIORef')
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe, fromJust)
 import Data.Text qualified as T
@@ -102,17 +102,51 @@ execTxWith executeTx tx = do
         pure (vmResult, gasLeftBeforeTx - gasLeftAfterTx)
   where
   runFully = do
+    vmResult <- executeTx
+    maybe (pure vmResult) (\q -> handleQuery q >> runFully) (getQuery vmResult)
+
+  -- | Handles reverts, failures and contract creations that might be the result
+  -- (`vmResult`) of executing transaction `tx`.
+  handleErrorsAndConstruction vmResult vmBeforeTx = case (vmResult, tx.call) of
+    (Reversion, _) -> do
+      tracesBeforeVMReset <- gets (.traces)
+      codeContractBeforeVMReset <- gets (.state.codeContract)
+      calldataBeforeVMReset <- gets (.state.calldata)
+      callvalueBeforeVMReset <- gets (.state.callvalue)
+      -- If a transaction reverts reset VM to state before the transaction.
+      put vmBeforeTx
+      -- Undo reset of some of the VM state.
+      -- Otherwise we'd loose all information about the reverted transaction like
+      -- contract address, calldata, result and traces.
+      #result ?= vmResult
+      #state % #calldata .= calldataBeforeVMReset
+      #state % #callvalue .= callvalueBeforeVMReset
+      #traces .= tracesBeforeVMReset
+      #state % #codeContract .= codeContractBeforeVMReset
+    (VMFailure x, _) -> do
+      dapp <- asks (.dapp)
+      vm <- get
+      vmExcept (Just (dapp, vm)) x
+    (VMSuccess (ConcreteBuf bytecode'), SolCreate _) -> do
+      -- Handle contract creation.
+      #env % #contracts % at (LitAddr tx.dst) % _Just % #code .= InitCode mempty mempty
+      fromEVM $ replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bytecode'))
+      modify' $ execState $ loadContract (LitAddr tx.dst)
+    _ -> pure ()
+
+  getRpcInfo = do
     config <- asks (.cfg)
     -- TODO: Is the latest block a good default? It makes fuzzing hard to reproduce. Rethink this.
     let rpcBlock = maybe EVM.Fetch.Latest (EVM.Fetch.BlockNumber . fromIntegral) config.rpcBlock
+    return (config.rpcUrl, rpcBlock)
+
 
-    vmResult <- executeTx
   -- For queries, we halt execution because the VM needs some additional
-    -- information from the outside. We provide this information and resume
-    -- the execution by recursively calling `runFully`.
-    case getQuery vmResult of
+  -- information from the outside. We provide this information, and then
+  -- the execution is resumed.
+
   -- A previously unknown contract is required
-      Just q@(PleaseFetchContract addr _ continuation) -> do
+  handleQuery q@(PleaseFetchContract addr _ continuation) = do
     cacheRef <- asks (.fetchContractCache)
     cache <- liftIO $ readIORef cacheRef
     case Map.lookup addr cache of
@@ -123,7 +157,8 @@ execTxWith executeTx tx = do
         put v'
       Nothing -> do
         logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
           Just rpcUrl -> do
             ret <- liftIO $ safeFetchContractFrom rpcBlock rpcUrl addr
             case ret of
@@ -145,10 +180,9 @@ execTxWith executeTx tx = do
             -- TODO: How should we fail here? RPC is not configured but VM
             -- wants to fetch
             fromEVM (continuation emptyAccount)
-        runFully -- resume execution
 
   -- A previously unknown slot is required
-      Just q@(PleaseFetchSlot addr slot continuation) -> do
+  handleQuery q@(PleaseFetchSlot addr slot continuation) = do
     cacheRef <- asks (.fetchSlotCache)
     cache <- liftIO $ readIORef cacheRef
     case Map.lookup addr cache >>= Map.lookup slot of
@@ -156,7 +190,8 @@ execTxWith executeTx tx = do
       Just Nothing -> fromEVM (continuation 0)
       Nothing -> do
         logMsg $ "INFO: Performing RPC: " <> show q
-            case config.rpcUrl of
+        (maybeRpcUrl, rpcBlock) <- getRpcInfo
+        case maybeRpcUrl of
           Just rpcUrl -> do
             ret <- liftIO $ safeFetchSlotFrom rpcBlock rpcUrl addr slot
             case ret of
@@ -175,47 +210,15 @@ execTxWith executeTx tx = do
             logMsg $ "ERROR: Requested RPC but it is not configured: " <> show q
             -- Use the zero slot
             fromEVM (continuation 0)
-        runFully -- resume execution
 
   -- Execute a FFI call
-      Just (PleaseDoFFI (cmd : args) continuation) -> do
+  handleQuery (PleaseDoFFI (cmd : args) continuation) = do
     (_, stdout, _) <- liftIO $ readProcessWithExitCode cmd args ""
     let encodedResponse = encodeAbiValue $
           AbiTuple (V.fromList [AbiBytesDynamic . hexText . T.pack $ stdout])
     fromEVM (continuation encodedResponse)
-        runFully
-
-      -- No queries to answer, the tx is fully executed and the result is final
-      _ -> pure vmResult
 
-  -- | Handles reverts, failures and contract creations that might be the result
-  -- (`vmResult`) of executing transaction `tx`.
-  handleErrorsAndConstruction vmResult vmBeforeTx = case (vmResult, tx.call) of
-    (Reversion, _) -> do
-      tracesBeforeVMReset <- gets (.traces)
-      codeContractBeforeVMReset <- gets (.state.codeContract)
-      calldataBeforeVMReset <- gets (.state.calldata)
-      callvalueBeforeVMReset <- gets (.state.callvalue)
-      -- If a transaction reverts reset VM to state before the transaction.
-      put vmBeforeTx
-      -- Undo reset of some of the VM state.
-      -- Otherwise we'd loose all information about the reverted transaction like
-      -- contract address, calldata, result and traces.
-      #result ?= vmResult
-      #state % #calldata .= calldataBeforeVMReset
-      #state % #callvalue .= callvalueBeforeVMReset
-      #traces .= tracesBeforeVMReset
-      #state % #codeContract .= codeContractBeforeVMReset
-    (VMFailure x, _) -> do
-      dapp <- asks (.dapp)
-      vm <- get
-      vmExcept (Just (dapp, vm)) x
-    (VMSuccess (ConcreteBuf bytecode'), SolCreate _) -> do
-      -- Handle contract creation.
-      #env % #contracts % at (LitAddr tx.dst) % _Just % #code .= InitCode mempty mempty
-      fromEVM $ replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bytecode'))
-      modify' $ execState $ loadContract (LitAddr tx.dst)
-    _ -> pure ()
+  handleQuery (PleaseDoFFI [] _) = error "Malformed FFI call"
 
 logMsg :: (MonadIO m, MonadReader Env m) => String -> m ()
 logMsg msg = do
@@ -262,8 +265,13 @@ execTxWithCov tx = do
     _ -> pure False
 
   pure (r, grew || grew')
-  where
-    -- the same as EVM.exec but collects coverage, will stop on a query
+
+-- | The same as EVM.exec but collects coverage, will stop on a query
+execCov
+  :: (MonadIO m, MonadState (VM Concrete RealWorld) m, MonadThrow m)
+  => Env
+  -> IORef (Bool, Maybe (VMut.IOVector CoverageInfo, Int))
+  -> m (VMResult Concrete RealWorld)
 execCov env covContextRef = do
   vm <- get
   (r, vm') <- liftIO $ loop vm
@@ -288,18 +296,11 @@ execTxWithCov tx = do
     let (pc, opIx, depth) = currentCovLoc vm
         contract = currentContract vm
 
-        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
-            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
+    maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ createCoverageVec contract
 
     case maybeCovVec of
       Nothing -> pure ()
-          Just vec -> do
+      Just vec ->
         -- 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.
@@ -313,6 +314,15 @@ execTxWithCov tx = do
             _ ->
               modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
 
+  createCoverageVec contract = 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
+      -- 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
+
   -- | Get the VM's current execution location
   currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)