{-# LANGUAGE DataKinds #-}

{- |
    Module: EVM.Keccak
    Description: Expr passes to determine Keccak assumptions
-}
module EVM.Keccak (keccakAssumptions, keccakCompute) where

import Control.Monad.State
import Data.Set (Set)
import Data.Set qualified as Set

import EVM.Traversals
import EVM.Types
import EVM.Expr


newtype KeccakStore = KeccakStore
  { KeccakStore -> Set (Expr 'EWord)
keccakEqs :: Set (Expr EWord) }
  deriving (Int -> KeccakStore -> ShowS
[KeccakStore] -> ShowS
KeccakStore -> String
(Int -> KeccakStore -> ShowS)
-> (KeccakStore -> String)
-> ([KeccakStore] -> ShowS)
-> Show KeccakStore
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KeccakStore -> ShowS
showsPrec :: Int -> KeccakStore -> ShowS
$cshow :: KeccakStore -> String
show :: KeccakStore -> String
$cshowList :: [KeccakStore] -> ShowS
showList :: [KeccakStore] -> ShowS
Show)

initState :: KeccakStore
initState :: KeccakStore
initState = KeccakStore { $sel:keccakEqs:KeccakStore :: Set (Expr 'EWord)
keccakEqs = Set (Expr 'EWord)
forall a. Set a
Set.empty }

keccakFinder :: forall a. Expr a -> State KeccakStore (Expr a)
keccakFinder :: forall (a :: EType). Expr a -> State KeccakStore (Expr a)
keccakFinder = \case
  e :: Expr a
e@(Keccak Expr 'Buf
_) -> do
    KeccakStore
s <- StateT KeccakStore Identity KeccakStore
forall s (m :: * -> *). MonadState s m => m s
get
    KeccakStore -> StateT KeccakStore Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (KeccakStore -> StateT KeccakStore Identity ())
-> KeccakStore -> StateT KeccakStore Identity ()
forall a b. (a -> b) -> a -> b
$ KeccakStore
s{$sel:keccakEqs:KeccakStore :: Set (Expr 'EWord)
keccakEqs=Expr 'EWord -> Set (Expr 'EWord) -> Set (Expr 'EWord)
forall a. Ord a => a -> Set a -> Set a
Set.insert Expr a
Expr 'EWord
e KeccakStore
s.keccakEqs}
    Expr a -> State KeccakStore (Expr a)
forall a. a -> StateT KeccakStore Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr a
e
  Expr a
e -> Expr a -> State KeccakStore (Expr a)
forall a. a -> StateT KeccakStore Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr a
e

findKeccakExpr :: forall a. Expr a -> State KeccakStore (Expr a)
findKeccakExpr :: forall (a :: EType). Expr a -> State KeccakStore (Expr a)
findKeccakExpr Expr a
e = (forall (a :: EType). Expr a -> State KeccakStore (Expr a))
-> Expr a -> StateT KeccakStore Identity (Expr a)
forall (m :: * -> *) (b :: EType).
Monad m =>
(forall (a :: EType). Expr a -> m (Expr a)) -> Expr b -> m (Expr b)
mapExprM Expr a -> State KeccakStore (Expr a)
forall (a :: EType). Expr a -> State KeccakStore (Expr a)
keccakFinder Expr a
e

findKeccakProp :: Prop -> State KeccakStore Prop
findKeccakProp :: Prop -> State KeccakStore Prop
findKeccakProp Prop
p = (forall (a :: EType). Expr a -> State KeccakStore (Expr a))
-> Prop -> State KeccakStore Prop
forall (m :: * -> *).
Monad m =>
(forall (a :: EType). Expr a -> m (Expr a)) -> Prop -> m Prop
mapPropM Expr a -> State KeccakStore (Expr a)
forall (a :: EType). Expr a -> State KeccakStore (Expr a)
keccakFinder Prop
p

findKeccakPropsExprs :: [Prop] -> [Expr Buf]  -> [Expr Storage]-> State KeccakStore ()
findKeccakPropsExprs :: [Prop]
-> [Expr 'Buf] -> [Expr 'Storage] -> StateT KeccakStore Identity ()
findKeccakPropsExprs [Prop]
ps [Expr 'Buf]
bufs [Expr 'Storage]
stores = do
  (Prop -> State KeccakStore Prop)
-> [Prop] -> StateT KeccakStore Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Prop -> State KeccakStore Prop
findKeccakProp [Prop]
ps
  (Expr 'Buf -> StateT KeccakStore Identity (Expr 'Buf))
-> [Expr 'Buf] -> StateT KeccakStore Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Expr 'Buf -> StateT KeccakStore Identity (Expr 'Buf)
forall (a :: EType). Expr a -> State KeccakStore (Expr a)
findKeccakExpr [Expr 'Buf]
bufs
  (Expr 'Storage -> StateT KeccakStore Identity (Expr 'Storage))
-> [Expr 'Storage] -> StateT KeccakStore Identity ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Expr 'Storage -> StateT KeccakStore Identity (Expr 'Storage)
forall (a :: EType). Expr a -> State KeccakStore (Expr a)
findKeccakExpr [Expr 'Storage]
stores


combine :: [a] -> [(a,a)]
combine :: forall a. [a] -> [(a, a)]
combine [a]
lst = [a] -> [[(a, a)]] -> [(a, a)]
forall {b}. [b] -> [[(b, b)]] -> [(b, b)]
combine' [a]
lst []
  where
    combine' :: [b] -> [[(b, b)]] -> [(b, b)]
combine' [] [[(b, b)]]
acc = [[(b, b)]] -> [(b, b)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(b, b)]]
acc
    combine' (b
x:[b]
xs) [[(b, b)]]
acc =
      let xcomb :: [(b, b)]
xcomb = [ (b
x, b
y) | b
y <- [b]
xs] in
      [b] -> [[(b, b)]] -> [(b, b)]
combine' [b]
xs ([(b, b)]
xcomb[(b, b)] -> [[(b, b)]] -> [[(b, b)]]
forall a. a -> [a] -> [a]
:[[(b, b)]]
acc)

minProp :: Expr EWord -> Prop
minProp :: Expr 'EWord -> Prop
minProp k :: Expr 'EWord
k@(Keccak Expr 'Buf
_) = Expr 'EWord -> Expr 'EWord -> Prop
PGT Expr 'EWord
k (W256 -> Expr 'EWord
Lit W256
256)
minProp Expr 'EWord
_ = String -> Prop
forall a. HasCallStack => String -> a
internalError String
"expected keccak expression"

concVal :: Expr EWord -> Prop
concVal :: Expr 'EWord -> Prop
concVal k :: Expr 'EWord
k@(Keccak (ConcreteBuf ByteString
bs)) = Expr 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
PEq (W256 -> Expr 'EWord
Lit (ByteString -> W256
keccak' ByteString
bs)) Expr 'EWord
k
concVal Expr 'EWord
_ = Bool -> Prop
PBool Bool
True

injProp :: (Expr EWord, Expr EWord) -> Prop
injProp :: (Expr 'EWord, Expr 'EWord) -> Prop
injProp (k1 :: Expr 'EWord
k1@(Keccak Expr 'Buf
b1), k2 :: Expr 'EWord
k2@(Keccak Expr 'Buf
b2)) =
  Prop -> Prop -> Prop
POr ((Expr 'Buf
b1 Expr 'Buf -> Expr 'Buf -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
.== Expr 'Buf
b2) Prop -> Prop -> Prop
.&& (Expr 'Buf -> Expr 'EWord
bufLength Expr 'Buf
b1 Expr 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
.== Expr 'Buf -> Expr 'EWord
bufLength Expr 'Buf
b2)) (Prop -> Prop
PNeg (Expr 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
PEq Expr 'EWord
k1 Expr 'EWord
k2))
injProp (Expr 'EWord, Expr 'EWord)
_ = String -> Prop
forall a. HasCallStack => String -> a
internalError String
"expected keccak expression"


-- Takes a list of props, find all keccak occurrences and generates two kinds of assumptions:
--   1. Minimum output value: That the output of the invocation is greater than
--      256 (needed to avoid spurious counterexamples due to storage collisions
--      with solidity mappings & value type storage slots)
--   2. Injectivity: That keccak is an injective function (we avoid quantifiers
--      here by making this claim for each unique pair of keccak invocations
--      discovered in the input expressions)
keccakAssumptions :: [Prop] -> [Expr Buf] -> [Expr Storage] -> [Prop]
keccakAssumptions :: [Prop] -> [Expr 'Buf] -> [Expr 'Storage] -> [Prop]
keccakAssumptions [Prop]
ps [Expr 'Buf]
bufs [Expr 'Storage]
stores = [Prop]
injectivity [Prop] -> [Prop] -> [Prop]
forall a. Semigroup a => a -> a -> a
<> [Prop]
minValue [Prop] -> [Prop] -> [Prop]
forall a. Semigroup a => a -> a -> a
<> [Prop]
minDiffOfPairs [Prop] -> [Prop] -> [Prop]
forall a. Semigroup a => a -> a -> a
<> [Prop]
concValues
  where
    (()
_, KeccakStore
st) = StateT KeccakStore Identity () -> KeccakStore -> ((), KeccakStore)
forall s a. State s a -> s -> (a, s)
runState ([Prop]
-> [Expr 'Buf] -> [Expr 'Storage] -> StateT KeccakStore Identity ()
findKeccakPropsExprs [Prop]
ps [Expr 'Buf]
bufs [Expr 'Storage]
stores) KeccakStore
initState

    injectivity :: [Prop]
injectivity = ((Expr 'EWord, Expr 'EWord) -> Prop)
-> [(Expr 'EWord, Expr 'EWord)] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Expr 'EWord, Expr 'EWord) -> Prop
injProp ([(Expr 'EWord, Expr 'EWord)] -> [Prop])
-> [(Expr 'EWord, Expr 'EWord)] -> [Prop]
forall a b. (a -> b) -> a -> b
$ [Expr 'EWord] -> [(Expr 'EWord, Expr 'EWord)]
forall a. [a] -> [(a, a)]
combine (Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.toList KeccakStore
st.keccakEqs)
    concValues :: [Prop]
concValues = (Expr 'EWord -> Prop) -> [Expr 'EWord] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr 'EWord -> Prop
concVal (Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.toList KeccakStore
st.keccakEqs)
    minValue :: [Prop]
minValue = (Expr 'EWord -> Prop) -> [Expr 'EWord] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr 'EWord -> Prop
minProp (Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.toList KeccakStore
st.keccakEqs)
    minDiffOfPairs :: [Prop]
minDiffOfPairs = ((Expr 'EWord, Expr 'EWord) -> Prop)
-> [(Expr 'EWord, Expr 'EWord)] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
map (Expr 'EWord, Expr 'EWord) -> Prop
minDistance ([(Expr 'EWord, Expr 'EWord)] -> [Prop])
-> [(Expr 'EWord, Expr 'EWord)] -> [Prop]
forall a b. (a -> b) -> a -> b
$ ((Expr 'EWord, Expr 'EWord) -> Bool)
-> [(Expr 'EWord, Expr 'EWord)] -> [(Expr 'EWord, Expr 'EWord)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Expr 'EWord -> Expr 'EWord -> Bool)
-> (Expr 'EWord, Expr 'EWord) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Expr 'EWord -> Expr 'EWord -> Bool
forall a. Eq a => a -> a -> Bool
(/=)) [(Expr 'EWord
a,Expr 'EWord
b) | Expr 'EWord
a<-(Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.elems KeccakStore
st.keccakEqs), Expr 'EWord
b<-(Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.elems KeccakStore
st.keccakEqs)]
     where
      minDistance :: (Expr EWord, Expr EWord) -> Prop
      minDistance :: (Expr 'EWord, Expr 'EWord) -> Prop
minDistance (ka :: Expr 'EWord
ka@(Keccak Expr 'Buf
a), kb :: Expr 'EWord
kb@(Keccak Expr 'Buf
b)) = Prop -> Prop -> Prop
PImpl (Expr 'Buf
a Expr 'Buf -> Expr 'Buf -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
./= Expr 'Buf
b) (Prop -> Prop -> Prop
PAnd Prop
req1 Prop
req2)
        where
          req1 :: Prop
req1 = (Expr 'EWord -> Expr 'EWord -> Prop
PGEq (Expr 'EWord -> Expr 'EWord -> Expr 'EWord
Sub Expr 'EWord
ka Expr 'EWord
kb) (W256 -> Expr 'EWord
Lit W256
256))
          req2 :: Prop
req2 = (Expr 'EWord -> Expr 'EWord -> Prop
PGEq (Expr 'EWord -> Expr 'EWord -> Expr 'EWord
Sub Expr 'EWord
kb Expr 'EWord
ka) (W256 -> Expr 'EWord
Lit W256
256))
      minDistance (Expr 'EWord, Expr 'EWord)
_ = String -> Prop
forall a. HasCallStack => String -> a
internalError String
"expected Keccak expression"

compute :: forall a. Expr a -> [Prop]
compute :: forall (a :: EType). Expr a -> [Prop]
compute = \case
  e :: Expr a
e@(Keccak Expr 'Buf
buf) -> do
    let b :: Expr 'Buf
b = Expr 'Buf -> Expr 'Buf
forall (a :: EType). Expr a -> Expr a
simplify Expr 'Buf
buf
    case Expr 'Buf -> Expr 'EWord
keccak Expr 'Buf
b of
      lit :: Expr 'EWord
lit@(Lit W256
_) -> [Expr a -> Expr a -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
PEq Expr a
e Expr a
Expr 'EWord
lit]
      Expr 'EWord
_ -> []
  Expr a
_ -> []

computeKeccakExpr :: forall a. Expr a -> [Prop]
computeKeccakExpr :: forall (a :: EType). Expr a -> [Prop]
computeKeccakExpr Expr a
e = (forall (a :: EType). Expr a -> [Prop])
-> [Prop] -> Expr a -> [Prop]
forall b (c :: EType).
Monoid b =>
(forall (a :: EType). Expr a -> b) -> b -> Expr c -> b
foldExpr Expr a -> [Prop]
forall (a :: EType). Expr a -> [Prop]
compute [] Expr a
e

computeKeccakProp :: Prop -> [Prop]
computeKeccakProp :: Prop -> [Prop]
computeKeccakProp Prop
p = (forall (a :: EType). Expr a -> [Prop]) -> [Prop] -> Prop -> [Prop]
forall b.
Monoid b =>
(forall (a :: EType). Expr a -> b) -> b -> Prop -> b
foldProp Expr a -> [Prop]
forall (a :: EType). Expr a -> [Prop]
compute [] Prop
p

keccakCompute :: [Prop] -> [Expr Buf] -> [Expr Storage] -> [Prop]
keccakCompute :: [Prop] -> [Expr 'Buf] -> [Expr 'Storage] -> [Prop]
keccakCompute [Prop]
ps [Expr 'Buf]
buf [Expr 'Storage]
stores =
  (Prop -> [Prop]) -> [Prop] -> [Prop]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Prop -> [Prop]
computeKeccakProp [Prop]
ps [Prop] -> [Prop] -> [Prop]
forall a. Semigroup a => a -> a -> a
<>
  (Expr 'Buf -> [Prop]) -> [Expr 'Buf] -> [Prop]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr 'Buf -> [Prop]
forall (a :: EType). Expr a -> [Prop]
computeKeccakExpr [Expr 'Buf]
buf [Prop] -> [Prop] -> [Prop]
forall a. Semigroup a => a -> a -> a
<>
  (Expr 'Storage -> [Prop]) -> [Expr 'Storage] -> [Prop]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr 'Storage -> [Prop]
forall (a :: EType). Expr a -> [Prop]
computeKeccakExpr [Expr 'Storage]
stores