content: {-# Language NamedFieldPuns #-}
  {-# Language DataKinds #-}
  {-# Language OverloadedStrings #-}
  {-# Language TypeApplications #-}
  {-# Language ScopedTypeVariables #-}
  
  module EVM.Symbolic where
  
  import Prelude hiding  (Word, LT, GT)
  import qualified Data.ByteString as BS
  import Data.ByteString (ByteString)
  import Control.Lens hiding (op, (:<), (|>), (.>))
  import Data.Maybe                   (fromMaybe, fromJust)
  
  import EVM.Types
  import qualified EVM.Concrete as Concrete
  import qualified Data.ByteArray       as BA
  import Data.SBV hiding (runSMT, newArray_, addAxiom, Word)
  import Data.SBV.Tools.Overflow
  import Crypto.Hash (Digest, SHA256)
  import qualified Crypto.Hash as Crypto
  
  litWord :: Word -> SymWord
  litWord (C whiff a) = S whiff (literal $ toSizzle a)
  
  litAddr :: Addr -> SAddr
  litAddr = SAddr . literal . toSizzle
  
  maybeLitAddr :: SAddr -> Maybe Addr
  maybeLitAddr (SAddr a) = fmap fromSizzle (unliteral a)
  
  maybeLitBytes :: [SWord 8] -> Maybe ByteString
  maybeLitBytes xs = fmap (\x -> BS.pack (fmap fromSized x)) (mapM unliteral xs)
  
  -- | Note: the (force*) functions are crude and in general,
  -- the continuation passing style `forceConcrete`
  -- alternatives should be prefered for better error
  -- handling when used during EVM execution
  forceLit :: SymWord -> Word
  forceLit (S whiff a) = case unliteral a of
    Just c -> C whiff (fromSizzle c)
    Nothing -> error "unexpected symbolic argument"
  
  forceLitBytes :: [SWord 8] -> ByteString
  forceLitBytes = BS.pack . fmap (fromSized . fromJust . unliteral)
  
  forceBuffer :: Buffer -> ByteString
  forceBuffer (ConcreteBuffer b) = b
  forceBuffer (SymbolicBuffer b) = forceLitBytes b
  
  sdiv :: SymWord -> SymWord -> SymWord
  sdiv (S a x) (S b y) = let sx, sy :: SInt 256
                             sx = sFromIntegral x
                             sy = sFromIntegral y
                         in S (Div a b) (sFromIntegral (sx `sQuot` sy))
  
  smod :: SymWord -> SymWord -> SymWord
  smod (S a x) (S b y) = let sx, sy :: SInt 256
                             sx = sFromIntegral x
                             sy = sFromIntegral y
                         in S (Mod a b) $ ite (y .== 0) 0 (sFromIntegral (sx `sRem` sy))
  
  addmod :: SymWord -> SymWord -> SymWord -> SymWord
  addmod (S a x) (S b y) (S c z) = let to512 :: SWord 256 -> SWord 512
                                       to512 = sFromIntegral
                                   in  S (Todo "addmod" [a, b, c]) $ ite (z .== 0) 0 $ sFromIntegral $ ((to512 x) + (to512 y)) `sMod` (to512 z)
  
  mulmod :: SymWord -> SymWord -> SymWord -> SymWord
  mulmod (S a x) (S b y) (S c z) = let to512 :: SWord 256 -> SWord 512
                                       to512 = sFromIntegral
                                   in S (Todo "mulmod" [a, b, c]) $ ite (z .== 0) 0 $ sFromIntegral $ ((to512 x) * (to512 y)) `sMod` (to512 z)
  
  -- | Signed less than
  slt :: SymWord -> SymWord -> SymWord
  slt (S xw x) (S yw y) =
    iteWhiff (SLT xw yw) (sFromIntegral x .< (sFromIntegral y :: (SInt 256))) 1 0
  
  -- | Signed greater than
  sgt :: SymWord -> SymWord -> SymWord
  sgt (S xw x) (S yw y) =
    iteWhiff (SGT xw yw) (sFromIntegral x .> (sFromIntegral y :: (SInt 256))) 1 0
  
  -- * Operations over symbolic memory (list of symbolic bytes)
  swordAt :: Int -> [SWord 8] -> SymWord
  swordAt i bs = let bs' = truncpad 32 $ drop i bs
                 in S (FromBytes (SymbolicBuffer bs')) (fromBytes bs')
  
  readByteOrZero' :: Int -> [SWord 8] -> SWord 8
  readByteOrZero' i bs = fromMaybe 0 (bs ^? ix i)
  
  sliceWithZero' :: Int -> Int -> [SWord 8] -> [SWord 8]
  sliceWithZero' o s m = truncpad s $ drop o m
  
  writeMemory' :: [SWord 8] -> Word -> Word -> Word -> [SWord 8] -> [SWord 8]
  writeMemory' bs1 (C _ n) (C _ src) (C _ dst) bs0 =
    let
      (a, b) = splitAt (num dst) bs0
      a'     = replicate (num dst - length a) 0
      c      = if src > num (length bs1)
               then replicate (num n) 0
               else sliceWithZero' (num src) (num n) bs1
      b'     = drop (num (n)) b
    in
      a <> a' <> c <> b'
  
  readMemoryWord' :: Word -> [SWord 8] -> SymWord
  readMemoryWord' (C _ i) m =
    let bs = truncpad 32 (drop (num i) m)
    in S (FromBytes (SymbolicBuffer bs)) (fromBytes bs)
  
  readMemoryWord32' :: Word -> [SWord 8] -> SWord 32
  readMemoryWord32' (C _ i) m = fromBytes $ truncpad 4 (drop (num i) m)
  
  setMemoryWord' :: Word -> SymWord -> [SWord 8] -> [SWord 8]
  setMemoryWord' (C _ i) (S _ x) =
    writeMemory' (toBytes x) 32 0 (num i)
  
  setMemoryByte' :: Word -> SWord 8 -> [SWord 8] -> [SWord 8]
  setMemoryByte' (C _ i) x =
    writeMemory' [x] 1 0 (num i)
  
  readSWord' :: Word -> [SWord 8] -> SymWord
  readSWord' (C _ i) x =
    if i > num (length x)
    then 0
    else swordAt (num i) x
  
  
  select' :: (Ord b, Num b, SymVal b, Mergeable a) => [a] -> a -> SBV b -> a
  select' xs err ind = walk xs ind err
      where walk []     _ acc = acc
            walk (e:es) i acc = walk es (i-1) (ite (i .== 0) e acc)
  
  -- | Read 32 bytes from index from a bounded list of bytes.
  readSWordWithBound :: SymWord -> Buffer -> SymWord -> SymWord
  readSWordWithBound sind@(S _ ind) (SymbolicBuffer xs) (S _ bound) = case (num <$> maybeLitWord sind, num <$> fromSizzle <$> unliteral bound) of
    (Just i, Just b) ->
      let bs = truncpad 32 $ drop i (take b xs)
      in S (FromBytes (SymbolicBuffer bs)) (fromBytes bs)
    _ ->
      -- Generates a ridiculously large set of constraints (roughly 25k) when
      -- the index is symbolic, but it still seems (kind of) manageable
      -- for the solvers.
  
      -- The proper solution here is to use smt arrays instead.
  
      let boundedList = [ite (i .<= bound) x' 0 | (x', i) <- zip xs [1..]]
          res = [select' boundedList 0 (ind + j) | j <- [0..31]]
      in S (FromBytes $ SymbolicBuffer res) $ fromBytes res
  
  readSWordWithBound sind (ConcreteBuffer xs) bound =
    case maybeLitWord sind of
      Nothing -> readSWordWithBound sind (SymbolicBuffer (litBytes xs)) bound
      Just x' ->
         -- INVARIANT: bound should always be length xs for concrete bytes
         -- so we should be able to safely ignore it here
           litWord $ Concrete.readMemoryWord x' xs
  
  -- a whole foldable instance seems overkill, but length is always good to have!
  len :: Buffer -> Int
  len (SymbolicBuffer bs) = length bs
  len (ConcreteBuffer bs) = BS.length bs
  
  readByteOrZero :: Int -> Buffer -> SWord 8
  readByteOrZero i (SymbolicBuffer bs) = readByteOrZero' i bs
  readByteOrZero i (ConcreteBuffer bs) = num $ Concrete.readByteOrZero i bs
  
  sliceWithZero :: Int -> Int -> Buffer -> Buffer
  sliceWithZero o s (SymbolicBuffer m) = SymbolicBuffer (sliceWithZero' o s m)
  sliceWithZero o s (ConcreteBuffer m) = ConcreteBuffer (Concrete.byteStringSliceWithDefaultZeroes o s m)
  
  writeMemory :: Buffer -> Word -> Word -> Word -> Buffer -> Buffer
  writeMemory (ConcreteBuffer bs1) n src dst (ConcreteBuffer bs0) =
    ConcreteBuffer (Concrete.writeMemory bs1 n src dst bs0)
  writeMemory (ConcreteBuffer bs1) n src dst (SymbolicBuffer bs0) =
    SymbolicBuffer (writeMemory' (litBytes bs1) n src dst bs0)
  writeMemory (SymbolicBuffer bs1) n src dst (ConcreteBuffer bs0) =
    SymbolicBuffer (writeMemory' bs1 n src dst (litBytes bs0))
  writeMemory (SymbolicBuffer bs1) n src dst (SymbolicBuffer bs0) =
    SymbolicBuffer (writeMemory' bs1 n src dst bs0)
  
  readMemoryWord :: Word -> Buffer -> SymWord
  readMemoryWord i (SymbolicBuffer m) = readMemoryWord' i m
  readMemoryWord i (ConcreteBuffer m) = litWord $ Concrete.readMemoryWord i m
  
  readMemoryWord32 :: Word -> Buffer -> SWord 32
  readMemoryWord32 i (SymbolicBuffer m) = readMemoryWord32' i m
  readMemoryWord32 i (ConcreteBuffer m) = num $ Concrete.readMemoryWord32 i m
  
  setMemoryWord :: Word -> SymWord -> Buffer -> Buffer
  setMemoryWord i x (SymbolicBuffer z) = SymbolicBuffer $ setMemoryWord' i x z
  setMemoryWord i x (ConcreteBuffer z) = case maybeLitWord x of
    Just x' -> ConcreteBuffer $ Concrete.setMemoryWord i x' z
    Nothing -> SymbolicBuffer $ setMemoryWord' i x (litBytes z)
  
  setMemoryByte :: Word -> SWord 8 -> Buffer -> Buffer
  setMemoryByte i x (SymbolicBuffer m) = SymbolicBuffer $ setMemoryByte' i x m
  setMemoryByte i x (ConcreteBuffer m) = case fromSized <$> unliteral x of
    Nothing -> SymbolicBuffer $ setMemoryByte' i x (litBytes m)
    Just x' -> ConcreteBuffer $ Concrete.setMemoryByte i x' m
  
  readSWord :: Word -> Buffer -> SymWord
  readSWord i (SymbolicBuffer x) = readSWord' i x
  readSWord i (ConcreteBuffer x) = num $ Concrete.readMemoryWord i x
  
  index :: Int -> Buffer -> SWord8
  index x (ConcreteBuffer b) = literal $ BS.index b x
  index x (SymbolicBuffer b) = fromSized $ b !! x
  
  -- * Uninterpreted functions
  
  symSHA256N :: SInteger -> SInteger -> SWord 256
  symSHA256N = uninterpret "sha256"
  
  symkeccakN :: SInteger -> SInteger -> SWord 256
  symkeccakN = uninterpret "keccak"
  
  toSInt :: [SWord 8] -> SInteger
  toSInt bs = sum $ zipWith (\a (i :: Integer) -> sFromIntegral a * 256 ^ i) bs [0..]
  
  
  -- | Although we'd like to define this directly as an uninterpreted function,
  -- we cannot because [a] is not a symbolic type. We must convert the list into a suitable
  -- symbolic type first. The only important property of this conversion is that it is injective.
  -- We embedd the bytestring as a pair of symbolic integers, this is a fairly easy solution.
  symkeccak' :: [SWord 8] -> SWord 256
  symkeccak' bytes = case length bytes of
    0 -> literal $ toSizzle $ keccak ""
    n -> symkeccakN (num n) (toSInt bytes)
  
  symSHA256 :: [SWord 8] -> [SWord 8]
  symSHA256 bytes = case length bytes of
    0 -> litBytes $ BS.pack $ BA.unpack $ (Crypto.hash BS.empty :: Digest SHA256)
    n -> toBytes $ symSHA256N (num n) (toSInt bytes)
  
  rawVal :: SymWord -> SWord 256
  rawVal (S _ v) = v
  
  -- | Reconstruct the smt/sbv value from a whiff
  -- Should satisfy (rawVal x .== whiffValue x)
  whiffValue :: Whiff -> SWord 256
  whiffValue w = case w of
    w'@(Todo _ _) -> error $ "unable to get value of " ++ show w'
    And x y       -> whiffValue x .&. whiffValue y
    Or x y        -> whiffValue x .|. whiffValue y
    Eq x y        -> ite (whiffValue x .== whiffValue y) 1 0
    LT x y        -> ite (whiffValue x .< whiffValue y) 1 0
    GT x y        -> ite (whiffValue x .> whiffValue y) 1 0
    ITE b x y     -> ite (whiffValue b .== 1) (whiffValue x) (whiffValue y)
    SLT x y       -> rawVal $ slt (S x (whiffValue x)) (S y (whiffValue y))
    SGT x y       -> rawVal $ sgt (S x (whiffValue x)) (S y (whiffValue y))
    IsZero x      -> ite (whiffValue x .== 0) 1 0
    SHL x y       -> sShiftLeft  (whiffValue x) (whiffValue y)
    SHR x y       -> sShiftRight (whiffValue x) (whiffValue y)
    SAR x y       -> sSignedShiftArithRight (whiffValue x) (whiffValue y)
    Add x y       -> whiffValue x + whiffValue y
    Sub x y       -> whiffValue x - whiffValue y
    Mul x y       -> whiffValue x * whiffValue y
    Div x y       -> whiffValue x `sDiv` whiffValue y
    Mod x y       -> whiffValue x `sMod` whiffValue y
    Exp x y       -> whiffValue x .^ whiffValue y
    Neg x         -> complement $ whiffValue x
    Var _ v       -> v
    FromKeccak (ConcreteBuffer bstr) -> literal $ num $ keccak bstr
    FromKeccak (SymbolicBuffer buf)  -> symkeccak' buf
    Literal x -> literal $ num $ x
    FromBytes buf -> rawVal $ readMemoryWord 0 buf
    FromStorage ind arr -> readArray arr (whiffValue ind)
  
  -- | Special cases that have proven useful in practice
  simplifyCondition :: SBool -> Whiff -> SBool
  simplifyCondition _ (IsZero (IsZero (IsZero a))) = whiffValue a .== 0
  
  
  
  -- | Overflow safe math can be difficult for smt solvers to deal with,
  -- especially for 256-bit words. When we recognize terms arising from
  -- overflow checks, we translate our queries into a more bespoke form,
  -- outlined in:
  -- Modular Bug-finding for Integer Overflows in the Large:
  -- Sound, Efficient, Bit-precise Static Analysis
  -- www.microsoft.com/en-us/research/wp-content/uploads/2016/02/z3prefix.pdf
  --
  -- Addition overflow.
  -- Written as
  --    require (x <= (x + y))
  -- or require (y <= (x + y))
  -- or require (!(y < (x + y)))
  simplifyCondition b (IsZero (IsZero (LT (Add x y) z))) =
    let x' = whiffValue x
        y' = whiffValue y
        z' = whiffValue z
        (_, overflow) = bvAddO x' y'
    in
      ite (x' .== z' .||
           y' .== z')
      overflow
      b
  
  -- Multiplication overflow.
  -- Written as
  --    require (y == 0 || x * y / y == x)
  -- or require (y == 0 || x == x * y / y)
  
  -- proveWith cvc4 $ \x y z -> ite (y .== (z :: SWord 8)) (((x * y) `sDiv` z ./= x) .<=> (snd (bvMulO x y) .|| (z .== 0 .&& x .> 0))) (sTrue)
  -- Q.E.D.
  simplifyCondition b (IsZero (Eq x (Div (Mul y z) w))) =
    simplifyCondition b (IsZero (Eq (Div (Mul y z) w) x))
  simplifyCondition b (IsZero (Eq (Div (Mul y z) w) x)) =
    let x' = whiffValue x
        y' = whiffValue y
        z' = whiffValue z
        w' = whiffValue w
        (_, overflow) = bvMulO y' z'
    in
      ite
      ((y' .== x' .&& z' .== w') .||
        (z' .== x' .&& y' .== w'))
      (overflow .|| (w' .== 0 .&& x' ./= 0))
      b
  simplifyCondition b _ = b
  