Cleanup RPC module

lib/Echidna/RPC.hs

@@ -9,15 +9,16 @@ import Control.Lens
 import Control.Monad (void)
 import Control.Monad.Catch (MonadThrow, throwM)
 import Control.Monad.Fail qualified as M (MonadFail(..))
-import Control.Monad.State.Strict (MonadState, get, put, execStateT)
+import Control.Monad.State.Strict (MonadState, get, put, execStateT, gets)
 import Data.Aeson (FromJSON(..), (.:), withObject, eitherDecodeFileStrict)
 import Data.ByteString.Base16 qualified as BS16 (decode)
 import Data.ByteString.Char8 (ByteString)
 import Data.ByteString.Char8 qualified as BS
 import Data.ByteString.Lazy qualified as LBS
+import Data.Map (member)
+import Data.Set (Set)
 import Data.Text qualified as T (drop)
 import Data.Text.Encoding (encodeUtf8)
-import Data.Map (member)
 import Data.Vector qualified as V (fromList, toList)
 import Text.Read (readMaybe)
 
@@ -26,20 +27,23 @@ import EVM.ABI (AbiType(..), AbiValue(..), decodeAbiValue, selector)
 import EVM.Exec (exec)
 import EVM.Types (Addr, W256, Expr(ConcreteBuf))
 
+import Echidna.ABI (encodeSig)
 import Echidna.Exec
 import Echidna.Transaction
 import Echidna.Types.Signature (SolSignature)
-import Echidna.ABI (encodeSig)
 import Echidna.Types (fromEVM)
 import Echidna.Types.Tx (TxCall(..), Tx(..), makeSingleTx, createTxWithValue, unlimitedGasPerBlock)
-import Data.Set (Set)
 
 -- | During initialization we can either call a function or create an account or contract
-data Etheno = AccountCreated Addr                                       -- ^ Registers an address with the echidna runtime
-            | ContractCreated Addr Addr Integer Integer ByteString W256 -- ^ A contract was constructed on the blockchain
-            | FunctionCall Addr Addr Integer Integer ByteString W256    -- ^ A contract function was executed
-            | BlockMined Integer Integer                                -- ^ A new block was mined contract
-
+data Etheno
+  -- | Registers an address with the echidna runtime
+  = AccountCreated Addr
+  -- | A contract was constructed on the blockchain
+  | ContractCreated Addr Addr Integer Integer ByteString W256
+  -- | A contract function was executed
+  | FunctionCall Addr Addr Integer Integer ByteString W256
+  -- | A new block was mined contract
+  | BlockMined Integer Integer
   deriving (Eq, Show)
 
 instance FromJSON Etheno where
@@ -100,8 +104,8 @@ extractFromEtheno ess ss = case ess of
 matchSignatureAndCreateTx :: SolSignature -> Etheno -> [Tx]
 matchSignatureAndCreateTx ("", []) _ = [] -- Not sure if we should match this.
 matchSignatureAndCreateTx (s,ts) (FunctionCall a d _ _ bs v) =
-  if BS.take 4 bs == selector (encodeSig (s,ts))
-  then makeSingleTx a d v $ SolCall (s, fromTuple $ decodeAbiValue t (LBS.fromStrict $ BS.drop 4 bs))
+  if BS.take 4 bs == selector (encodeSig (s,ts)) then
+    makeSingleTx a d v $ SolCall (s, fromTuple $ decodeAbiValue t (LBS.fromStrict $ BS.drop 4 bs))
   else []
   where t = AbiTupleType (V.fromList ts)
         fromTuple (AbiTuple xs) = V.toList xs
@@ -122,10 +126,12 @@ loadEthenoBatch ffi fp = do
 
 initAddress :: MonadState VM m => Addr -> m ()
 initAddress addr = do
-  cs <- use (env . EVM.contracts)
+  cs <- gets (._env._contracts)
   if addr `member` cs then pure ()
   else env . EVM.contracts . at addr .= Just account
- where account = initialContract (RuntimeCode (ConcreteRuntimeCode mempty)) & set nonce 0 & set balance 100000000000000000000 -- default balance for EOAs in etheno
+ where
+   account = (initialContract (RuntimeCode (ConcreteRuntimeCode mempty)))
+     { _nonce = 0, _balance = 100000000000000000000 } -- default balance for EOAs in etheno
 
 crashWithQueryError :: (MonadState VM m, MonadFail m, MonadThrow m) => Query -> Etheno -> m ()
 crashWithQueryError q et =
@@ -148,20 +154,28 @@ execEthenoTxs et = do
   vm <- get
   res <- fromEVM exec
   case (res, et) of
-       (_        , AccountCreated _)  -> return ()
-       (Reversion,   _)               -> void $ put vm
-       (VMFailure (Query q), _)       -> crashWithQueryError q et
-       (VMFailure x, _)               -> vmExcept x >> M.fail "impossible"
-       (VMSuccess (ConcreteBuf bc),
-        ContractCreated _ ca _ _ _ _) -> do
+    (_, AccountCreated _) ->
+      pure ()
+    (Reversion, _) ->
+      void $ put vm
+    (VMFailure (Query q), _) ->
+      crashWithQueryError q et
+    (VMFailure x, _) ->
+      vmExcept x >> M.fail "impossible"
+    (VMSuccess (ConcreteBuf bc), ContractCreated _ ca _ _ _ _) -> do
       env . contracts . at ca . _Just . contractcode .= InitCode mempty mempty
-          fromEVM (replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bc)) >> loadContract ca)
-          return ()
-       _                              -> return ()
+      fromEVM $ do
+        replaceCodeOfSelf (RuntimeCode (ConcreteRuntimeCode bc))
+        loadContract ca
+    _ -> pure ()
 
 -- | For an etheno txn, set up VM to execute txn
 setupEthenoTx :: MonadState VM m => Etheno -> m ()
-setupEthenoTx (AccountCreated f) = initAddress f -- TODO: improve etheno to include initial balance
-setupEthenoTx (ContractCreated f c _ _ d v) = setupTx $ createTxWithValue d f c unlimitedGasPerBlock v (1, 1)
-setupEthenoTx (FunctionCall f t _ _ d v) = setupTx $ Tx (SolCalldata d) f t unlimitedGasPerBlock 0 v (1, 1)
-setupEthenoTx (BlockMined n t) = setupTx $ Tx NoCall 0 0 0 0 0 (fromInteger t, fromInteger n)
+setupEthenoTx (AccountCreated f) =
+  initAddress f -- TODO: improve etheno to include initial balance
+setupEthenoTx (ContractCreated f c _ _ d v) =
+  setupTx $ createTxWithValue d f c unlimitedGasPerBlock v (1, 1)
+setupEthenoTx (FunctionCall f t _ _ d v) =
+  setupTx $ Tx (SolCalldata d) f t unlimitedGasPerBlock 0 v (1, 1)
+setupEthenoTx (BlockMined n t) =
+  setupTx $ Tx NoCall 0 0 0 0 0 (fromInteger t, fromInteger n)