content: {-# Language DataKinds #-}
  {-# Language OverloadedStrings #-}
  {-# Language TypeApplications #-}
  
  module EVM.SymExec where
  
  import Prelude hiding (Word)
  
  import Control.Lens hiding (pre)
  import EVM hiding (Query, push)
  import qualified EVM
  import EVM.Exec
  import qualified EVM.Fetch as Fetch
  import EVM.ABI
  import EVM.Stepper (Stepper)
  import qualified EVM.Stepper as Stepper
  import qualified Control.Monad.Operational as Operational
  import Control.Monad.State.Strict hiding (state)
  import Data.Maybe (catMaybes, fromMaybe)
  import EVM.Types
  import EVM.Concrete (createAddress)
  import qualified EVM.FeeSchedule as FeeSchedule
  import Data.SBV.Trans.Control
  import Data.SBV.Trans hiding (distinct, Word)
  import Data.SBV hiding (runSMT, newArray_, addAxiom, distinct, sWord8s, Word)
  import Data.Vector (toList, fromList)
  import Data.Tree
  import Data.DoubleWord (Word256)
  
  import Data.ByteString (ByteString, pack)
  import qualified Data.ByteString.Lazy as Lazy
  import qualified Data.ByteString as BS
  import Data.Text (Text, splitOn, unpack)
  import qualified Control.Monad.State.Class as State
  import Control.Applicative
  
  data ProofResult a b c = Qed a | Cex b | Timeout c
  type VerifyResult = ProofResult (Tree BranchInfo) (Tree BranchInfo) (Tree BranchInfo)
  type EquivalenceResult = ProofResult ([VM], [VM]) VM ()
  
  -- | Convenience functions for generating large symbolic byte strings
  sbytes32, sbytes128, sbytes256, sbytes512, sbytes1024 :: Query ([SWord 8])
  sbytes32 = toBytes <$> freshVar_ @ (WordN 256)
  sbytes128 = toBytes <$> freshVar_ @ (WordN 1024)
  sbytes256 = liftA2 (++) sbytes128 sbytes128
  sbytes512 = liftA2 (++) sbytes256 sbytes256
  sbytes1024 = liftA2 (++) sbytes512 sbytes512
  
  mkByte :: Query [SWord 8]
  mkByte = do x <- freshVar_
              return [x]
  
  -- | Abstract calldata argument generation
  symAbiArg :: AbiType -> Query ([SWord 8], W256)
  symAbiArg (AbiUIntType n) | n `mod` 8 == 0 && n <= 256 =
    do x <- concatMapM (const mkByte) [0..(n `div` 8) - 1]
       return (padLeft' 32 x, 32)
                            | otherwise = error "bad type"
  
  symAbiArg (AbiIntType n)  | n `mod` 8 == 0 && n <= 256 =
    do x <- concatMapM (const mkByte) [(0 :: Int) ..(n `div` 8) - 1]
       return (padLeft' 32 x, 32)
  
                            | otherwise = error "bad type"
  symAbiArg AbiBoolType =
    do x <- mkByte
       return (padLeft' 32 x, 32)
  
  symAbiArg AbiAddressType =
    do x <- concatMapM (const mkByte) [(0 :: Int)..19]
       return (padLeft' 32 x, 32)
  
  symAbiArg (AbiBytesType n) | n <= 32 =
    do x <- concatMapM (const mkByte) [0..n - 1]
       return (padLeft' 32 x, 32)
  
                             | otherwise = error "bad type"
  
  -- TODO: is this encoding correct?
  symAbiArg (AbiArrayType len typ) =
    do args <- mapM symAbiArg (replicate len typ)
       return (litBytes (encodeAbiValue (AbiUInt 256 (fromIntegral len))) <> (concat $ fst <$> args),
               32 + (sum $ snd <$> args))
  
  symAbiArg (AbiTupleType tuple) =
    do args <- mapM symAbiArg (toList tuple)
       return (concat $ fst <$> args, sum $ snd <$> args)
  symAbiArg n =
    error $ "Unsupported symbolic abiencoding for"
      <> show n
      <> ". Please file an issue at https://github.com/dapphub/dapptools if you really need this."
  
  -- | Generates calldata matching given type signature, optionally specialized
  -- with concrete arguments.
  -- Any argument given as "<symbolic>" or omitted at the tail of the list are
  -- kept symbolic.
  symCalldata :: Text -> [AbiType] -> [String] -> Query ([SWord 8], W256)
  symCalldata sig typesignature concreteArgs =
    let args = concreteArgs <> replicate (length typesignature - length concreteArgs)  "<symbolic>"
        mkArg typ "<symbolic>" = symAbiArg typ
        mkArg typ arg = let n = litBytes . encodeAbiValue $ makeAbiValue typ arg
                        in return (n, num (length n))
        sig' = litBytes $ selector sig
    in do calldatas <- zipWithM mkArg typesignature args
          return (sig' <> concat (fst <$> calldatas), 4 + (sum $ snd <$> calldatas))
  
  abstractVM :: Maybe (Text, [AbiType]) -> [String] -> ByteString -> StorageModel -> Query VM
  abstractVM typesignature concreteArgs x storagemodel = do
    (cd', cdlen, cdconstraint) <-
      case typesignature of
        Nothing -> do cd <- sbytes256
                      len <- freshVar_
                      return (cd, var "calldataLength" len, (len .<= 256, Todo "calldatalength < 256" []))
        Just (name, typs) -> do (cd, cdlen) <- symCalldata name typs concreteArgs
                                return (cd, S (Literal cdlen) (literal $ num cdlen), (sTrue, Todo "Trivial" []))
    symstore <- case storagemodel of
      SymbolicS -> Symbolic [] <$> freshArray_ Nothing
      InitialS -> Symbolic [] <$> freshArray_ (Just 0)
      ConcreteS -> return $ Concrete mempty
    c <- SAddr <$> freshVar_
    value' <- var "CALLVALUE" <$> freshVar_
    return $ loadSymVM (RuntimeCode (ConcreteBuffer x)) symstore storagemodel c value' (SymbolicBuffer cd', cdlen) & over constraints ((<>) [cdconstraint])
  
  loadSymVM :: ContractCode -> Storage -> StorageModel -> SAddr -> SymWord -> (Buffer, SymWord) -> VM
  loadSymVM x initStore model addr callvalue' calldata' =
    (makeVm $ VMOpts
      { vmoptContract = contractWithStore x initStore
      , vmoptCalldata = calldata'
      , vmoptValue = callvalue'
      , vmoptAddress = createAddress ethrunAddress 1
      , vmoptCaller = addr
      , vmoptOrigin = ethrunAddress --todo: generalize
      , vmoptCoinbase = 0
      , vmoptNumber = 0
      , vmoptTimestamp = 0
      , vmoptBlockGaslimit = 0
      , vmoptGasprice = 0
      , vmoptDifficulty = 0
      , vmoptGas = 0xffffffffffffffff
      , vmoptGaslimit = 0xffffffffffffffff
      , vmoptBaseFee = 0
      , vmoptPriorityFee = 0
      , vmoptMaxCodeSize = 0xffffffff
      , vmoptSchedule = FeeSchedule.berlin
      , vmoptChainId = 1
      , vmoptCreate = False
      , vmoptStorageModel = model
      , vmoptTxAccessList = mempty
      , vmoptAllowFFI = False
      }) & set (env . contracts . at (createAddress ethrunAddress 1))
               (Just (contractWithStore x initStore))
  
  data BranchInfo = BranchInfo
    { _vm                 :: VM,
      _branchCondition    :: Maybe Whiff
    }
  
  doInterpret :: Fetch.Fetcher -> Maybe Integer -> Maybe Integer -> VM -> Query (Tree BranchInfo)
  doInterpret fetcher maxIter askSmtIters vm = let
        f (vm', cs) = Node (BranchInfo (if null cs then vm' else vm) Nothing) cs
      in f <$> interpret' fetcher maxIter askSmtIters vm
  
  interpret' :: Fetch.Fetcher -> Maybe Integer -> Maybe Integer -> VM -> Query (VM, [Tree BranchInfo])
  interpret' fetcher maxIter askSmtIters vm = let
    cont s = interpret' fetcher maxIter askSmtIters $ execState s vm
    in case view EVM.result vm of
  
      Nothing -> cont exec1
  
      Just (VMFailure (EVM.Query q@(PleaseAskSMT _ _ continue))) -> let
        codelocation = getCodeLocation vm
        iteration = num $ fromMaybe 0 $ view (iterations . at codelocation) vm
        -- as an optimization, we skip consulting smt
        -- if we've been at the location less than 5 times
        in if iteration < (fromMaybe 5 askSmtIters)
           then cont $ continue EVM.Unknown
           else io (fetcher q) >>= cont
  
      Just (VMFailure (EVM.Query q)) -> io (fetcher q) >>= cont
  
      Just (VMFailure (Choose (EVM.PleaseChoosePath whiff continue)))
        -> case maxIterationsReached vm maxIter of
          Nothing -> let
            lvm = execState (continue True) vm
            rvm = execState (continue False) vm
            in do
              push 1
              (leftvm, left) <- interpret' fetcher maxIter askSmtIters lvm
              pop 1
              push 1
              (rightvm, right) <- interpret' fetcher maxIter askSmtIters rvm
              pop 1
              return (vm, [Node (BranchInfo leftvm (Just whiff)) left, Node (BranchInfo rightvm (Just whiff)) right])
          Just n -> cont $ continue (not n)
  
      Just _
        -> return (vm, [])
  
  -- | Interpreter which explores all paths at
  -- | branching points.
  -- | returns a list of possible final evm states
  interpret
    :: Fetch.Fetcher
    -> Maybe Integer -- max iterations
    -> Maybe Integer -- ask smt iterations
    -> Stepper a
    -> StateT VM Query [a]
  interpret fetcher maxIter askSmtIters =
    eval . Operational.view
  
    where
      eval
        :: Operational.ProgramView Stepper.Action a
        -> StateT VM Query [a]
  
      eval (Operational.Return x) =
        pure [x]
  
      eval (action Operational.:>>= k) =
        case action of
          Stepper.Exec ->
            exec >>= interpret fetcher maxIter askSmtIters . k
          Stepper.Run ->
            run >>= interpret fetcher maxIter askSmtIters . k
          Stepper.IOAct q ->
            mapStateT io q >>= interpret fetcher maxIter askSmtIters . k
          Stepper.Ask (EVM.PleaseChoosePath _ continue) -> do
            vm <- get
            case maxIterationsReached vm maxIter of
              Nothing -> do
                push 1
                a <- interpret fetcher maxIter askSmtIters (Stepper.evm (continue True) >>= k)
                put vm
                pop 1
                push 1
                b <- interpret fetcher maxIter askSmtIters (Stepper.evm (continue False) >>= k)
                pop 1
                return $ a <> b
              Just n ->
                interpret fetcher maxIter askSmtIters (Stepper.evm (continue (not n)) >>= k)
          Stepper.Wait q -> do
            let performQuery = do
                  m <- liftIO (fetcher q)
                  interpret fetcher maxIter askSmtIters (Stepper.evm m >>= k)
  
            case q of
              PleaseAskSMT _ _ continue -> do
                codelocation <- getCodeLocation <$> get
                iteration <- num . fromMaybe 0 <$> use (iterations . at codelocation)
  
                -- if this is the first time we are branching at this point,
                -- explore both branches without consulting SMT.
                -- Exploring too many branches is a lot cheaper than
                -- consulting our SMT solver.
                if iteration < (fromMaybe 5 askSmtIters)
                then interpret fetcher maxIter askSmtIters (Stepper.evm (continue EVM.Unknown) >>= k)
                else performQuery
  
              _ -> performQuery
  
          Stepper.EVM m ->
            State.state (runState m) >>= interpret fetcher maxIter askSmtIters . k
  
  maxIterationsReached :: VM -> Maybe Integer -> Maybe Bool
  maxIterationsReached _ Nothing = Nothing
  maxIterationsReached vm (Just maxIter) =
    let codelocation = getCodeLocation vm
        iters = view (iterations . at codelocation . non 0) vm
    in if num maxIter <= iters
       then view (cache . path . at (codelocation, iters - 1)) vm
       else Nothing
  
  type Precondition = VM -> SBool
  type Postcondition = (VM, VM) -> SBool
  
  checkAssert :: [Word256] -> ByteString -> Maybe (Text, [AbiType]) -> [String] -> Query (VerifyResult, VM)
  checkAssert errs c signature' concreteArgs = verifyContract c signature' concreteArgs SymbolicS (const sTrue) (Just $ checkAssertions errs)
  
  {- |Checks if an assertion violation has been encountered
  
    hevm recognises the following as an assertion violation:
  
    1. the invalid opcode (0xfe) (solc < 0.8)
    2. a revert with a reason of the form `abi.encodeWithSelector("Panic(uint256)", code)`, where code is one of the following (solc >= 0.8):
      - 0x00: Used for generic compiler inserted panics.
      - 0x01: If you call assert with an argument that evaluates to false.
      - 0x11: If an arithmetic operation results in underflow or overflow outside of an unchecked { ... } block.
      - 0x12; If you divide or modulo by zero (e.g. 5 / 0 or 23 % 0).
      - 0x21: If you convert a value that is too big or negative into an enum type.
      - 0x22: If you access a storage byte array that is incorrectly encoded.
      - 0x31: If you call .pop() on an empty array.
      - 0x32: If you access an array, bytesN or an array slice at an out-of-bounds or negative index (i.e. x[i] where i >= x.length or i < 0).
      - 0x41: If you allocate too much memory or create an array that is too large.
      - 0x51: If you call a zero-initialized variable of internal function type.
  
    see: https://docs.soliditylang.org/en/v0.8.6/control-structures.html?highlight=Panic#panic-via-assert-and-error-via-require
  -}
  checkAssertions :: [Word256] -> Postcondition
  checkAssertions errs (_, out) = case view result out of
    Just (EVM.VMFailure (EVM.UnrecognizedOpcode 254)) -> sFalse
    Just (EVM.VMFailure (EVM.Revert msg)) -> if msg `elem` (fmap panicMsg errs) then sFalse else sTrue
    _ -> sTrue
  
  -- |By default hevm checks for all assertions except those which result from arithmetic overflow
  defaultPanicCodes :: [Word256]
  defaultPanicCodes = [ 0x00, 0x01, 0x12, 0x21, 0x22, 0x31, 0x32, 0x41, 0x51 ]
  
  allPanicCodes :: [Word256]
  allPanicCodes = [ 0x00, 0x01, 0x11, 0x12, 0x21, 0x22, 0x31, 0x32, 0x41, 0x51 ]
  
  -- |Produces the revert message for solc >=0.8 assertion violations
  panicMsg :: Word256 -> ByteString
  panicMsg err = (selector "Panic(uint256)") <> (encodeAbiValue $ AbiUInt 256 err)
  
  verifyContract :: ByteString -> Maybe (Text, [AbiType]) -> [String] -> StorageModel -> Precondition -> Maybe Postcondition -> Query (VerifyResult, VM)
  verifyContract theCode signature' concreteArgs storagemodel pre maybepost = do
      preStateRaw <- abstractVM signature' concreteArgs theCode  storagemodel
      -- add the pre condition to the pathconditions to ensure that we are only exploring valid paths
      let preState = over constraints ((++) [(pre preStateRaw, Todo "assumptions" [])]) preStateRaw
      v <- verify preState Nothing Nothing Nothing maybepost
      return (v, preState)
  
  pruneDeadPaths :: [VM] -> [VM]
  pruneDeadPaths =
    filter $ \vm -> case view result vm of
      Just (VMFailure DeadPath) -> False
      _ -> True
  
  consistentPath :: VM -> Query (Maybe VM)
  consistentPath vm = do
    resetAssertions
    constrain $ sAnd $ fst <$> view constraints vm
    checkSat >>= \case
      Sat -> return $ Just vm
      Unk -> return $ Just vm -- the path may still be consistent
      Unsat -> return Nothing
      DSat _ -> error "unexpected DSAT"
  
  consistentTree :: Tree BranchInfo -> Query (Maybe (Tree BranchInfo))
  consistentTree (Node (BranchInfo vm w) []) = do
    consistentPath vm >>= \case
      Nothing  -> return Nothing
      Just vm' -> return $ Just $ Node (BranchInfo vm' w) []
  consistentTree (Node b xs) = do
    consistentChildren <- catMaybes <$> forM xs consistentTree
    if null consistentChildren then
      return Nothing
    else
      return $ Just (Node b consistentChildren)
  
  
  leaves :: Tree BranchInfo -> [VM]
  leaves (Node x []) = [_vm x]
  leaves (Node _ xs) = concatMap leaves xs
  
  -- | Symbolically execute the VM and check all endstates against the postcondition, if available.
  verify :: VM -> Maybe Integer -> Maybe Integer -> Maybe (Fetch.BlockNumber, Text) -> Maybe Postcondition -> Query VerifyResult
  verify preState maxIter askSmtIters rpcinfo maybepost = do
    smtState <- queryState
    tree <- doInterpret (Fetch.oracle (Just smtState) rpcinfo False) maxIter askSmtIters preState
    case maybepost of
      (Just post) -> do
        let livePaths = pruneDeadPaths $ leaves tree
            -- have we hit max iterations at any point in a given path
            maxReached :: VM -> Bool
            maxReached p = case maxIter of
              Just maxI -> any (>= (fromInteger maxI)) (view iterations p)
              Nothing -> False
            -- is there any path which can possibly violate the postcondition?
            -- can also do these queries individually (even concurrently!). Could save time and report multiple violations
            postC = sOr $ fmap (\postState -> (sAnd (fst <$> view constraints postState)) .&& sNot (post (preState, postState))) livePaths
        resetAssertions
        constrain postC
        io $ putStrLn "checking postcondition..."
        checkSat >>= \case
          Unk -> do io $ putStrLn "postcondition query timed out"
                    return $ Timeout tree
          Unsat -> do
            if any maxReached livePaths
              then io $ putStrLn "WARNING: max iterations reached, execution halted prematurely"
              else io $ putStrLn "Q.E.D."
            return $ Qed tree
          Sat -> return $ Cex tree
          DSat _ -> error "unexpected DSAT"
  
      Nothing -> do io $ putStrLn "Nothing to check"
                    return $ Qed tree
  
  -- | Compares two contract runtimes for trace equivalence by running two VMs and comparing the end states.
  equivalenceCheck :: ByteString -> ByteString -> Maybe Integer -> Maybe Integer -> Maybe (Text, [AbiType]) -> Query EquivalenceResult
  equivalenceCheck bytecodeA bytecodeB maxiter askSmtIters signature' = do
    let
      bytecodeA' = if BS.null bytecodeA then BS.pack [0] else bytecodeA
      bytecodeB' = if BS.null bytecodeB then BS.pack [0] else bytecodeB
    preStateA <- abstractVM signature' [] bytecodeA' SymbolicS
  
    let preself = preStateA ^. state . contract
        precaller = preStateA ^. state . caller
        callvalue' = preStateA ^. state . callvalue
        prestorage = preStateA ^?! env . contracts . ix preself . storage
        (calldata', cdlen) = view (state . calldata) preStateA
        pathconds = view constraints preStateA
        preStateB = loadSymVM (RuntimeCode (ConcreteBuffer bytecodeB')) prestorage SymbolicS precaller callvalue' (calldata', cdlen) & set constraints pathconds
  
    smtState <- queryState
    push 1
    aVMs <- doInterpret (Fetch.oracle (Just smtState) Nothing False) maxiter askSmtIters preStateA
    pop 1
    push 1
    bVMs <- doInterpret (Fetch.oracle (Just smtState) Nothing False) maxiter askSmtIters preStateB
    pop 1
    -- Check each pair of endstates for equality:
    let differingEndStates = uncurry distinct <$> [(a,b) | a <- pruneDeadPaths (leaves aVMs), b <- pruneDeadPaths (leaves bVMs)]
        distinct a b =
          let (aPath, bPath) = both' (view constraints) (a, b)
              (aSelf, bSelf) = both' (view (state . contract)) (a, b)
              (aEnv, bEnv) = both' (view (env . contracts)) (a, b)
              (aResult, bResult) = both' (view result) (a, b)
              (Symbolic _ aStorage, Symbolic _ bStorage) = (view storage (aEnv ^?! ix aSelf), view storage (bEnv ^?! ix bSelf))
              differingResults = case (aResult, bResult) of
  
                (Just (VMSuccess aOut), Just (VMSuccess bOut)) ->
                  aOut ./= bOut .|| aStorage ./= bStorage .|| fromBool (aSelf /= bSelf)
  
                (Just (VMFailure UnexpectedSymbolicArg), _) ->
                  error $ "Unexpected symbolic argument at opcode: " <> maybe "??" show (vmOp a) <> ". Not supported (yet!)"
  
                (_, Just (VMFailure UnexpectedSymbolicArg)) ->
                  error $ "Unexpected symbolic argument at opcode: " <> maybe "??" show (vmOp a) <> ". Not supported (yet!)"
  
                (Just (VMFailure _), Just (VMFailure _)) -> sFalse
  
                (Just _, Just _) -> sTrue
  
                errormsg -> error $ show errormsg
  
          in sAnd (fst <$> aPath) .&& sAnd (fst <$> bPath) .&& differingResults
    -- If there exists a pair of endstates where this is not the case,
    -- the following constraint is satisfiable
    constrain $ sOr differingEndStates
  
    checkSat >>= \case
       Unk -> return $ Timeout ()
       Sat -> return $ Cex preStateA
       Unsat -> return $ Qed (leaves aVMs, leaves bVMs)
       DSat _ -> error "unexpected DSAT"
  
  both' :: (a -> b) -> (a, a) -> (b, b)
  both' f (x, y) = (f x, f y)
  
  showCounterexample :: VM -> Maybe (Text, [AbiType]) -> Query ()
  showCounterexample vm maybesig = do
    let (calldata', S _ cdlen) = view (EVM.state . EVM.calldata) vm
        S _ cvalue = view (EVM.state . EVM.callvalue) vm
        SAddr caller' = view (EVM.state . EVM.caller) vm
    cdlen' <- num <$> getValue cdlen
    calldatainput <- case calldata' of
      SymbolicBuffer cd -> mapM (getValue.fromSized) (take cdlen' cd) >>= return . pack
      ConcreteBuffer cd -> return $ BS.take cdlen' cd
    callvalue' <- getValue cvalue
    caller'' <- num <$> getValue caller'
    io $ do
      putStrLn "Calldata:"
      print $ ByteStringS calldatainput
  
      -- pretty print calldata input if signature is available
      case maybesig of
        Just (name, types) -> putStrLn $ unpack (head (splitOn "(" name)) ++
          show (decodeAbiValue (AbiTupleType (fromList types)) $ Lazy.fromStrict (BS.drop 4 calldatainput))
        Nothing -> return ()
  
      putStrLn "Caller:"
      print (Addr caller'')
      putStrLn "Callvalue:"
      print callvalue'
  