{- |
    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 Data.List (tails)

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


newtype KeccakStore = KeccakStore
  { KeccakStore -> Set (Expr 'EWord)
keccakExprs :: 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:keccakExprs:KeccakStore :: Set (Expr 'EWord)
keccakExprs = Set (Expr 'EWord)
forall a. Set a
Set.empty }

keccakFinder :: forall a. Expr a -> State KeccakStore ()
keccakFinder :: forall (a :: EType). Expr a -> State KeccakStore ()
keccakFinder = \case
  e :: Expr a
e@(Keccak Expr 'Buf
_) -> (KeccakStore -> KeccakStore) -> State KeccakStore ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\KeccakStore
s -> KeccakStore
s{keccakExprs=Set.insert e s.keccakExprs})
  Expr a
_ -> () -> State KeccakStore ()
forall a. a -> StateT KeccakStore Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

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

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

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


uniquePairs :: [a] -> [(a,a)]
uniquePairs :: forall a. [a] -> [(a, a)]
uniquePairs [a]
xs = [(a
x,a
y) | (a
x:[a]
ys) <- [a] -> [[a]]
forall a. [a] -> [[a]]
Data.List.tails [a]
xs, a
y <- [a]
ys]

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"

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
PImpl (Expr 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
PEq Expr 'EWord
k1 Expr 'EWord
k2) ((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))
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
  where
    (()
_, KeccakStore
st) = State KeccakStore () -> KeccakStore -> ((), KeccakStore)
forall s a. State s a -> s -> (a, s)
runState ([Prop] -> [Expr 'Buf] -> [Expr 'Storage] -> State KeccakStore ()
findKeccakPropsExprs [Prop]
ps [Expr 'Buf]
bufs [Expr 'Storage]
stores) KeccakStore
initState

    keccakPairs :: [(Expr 'EWord, Expr 'EWord)]
keccakPairs = [Expr 'EWord] -> [(Expr 'EWord, Expr 'EWord)]
forall a. [a] -> [(a, a)]
uniquePairs (Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.toList KeccakStore
st.keccakExprs)
    injectivity :: [Prop]
injectivity = ((Expr 'EWord, Expr 'EWord) -> Prop)
-> [(Expr 'EWord, Expr 'EWord)] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
map (Expr 'EWord, Expr 'EWord) -> Prop
injProp [(Expr 'EWord, Expr 'EWord)]
keccakPairs
    minValue :: [Prop]
minValue = (Expr 'EWord -> Prop) -> [Expr 'EWord] -> [Prop]
forall a b. (a -> b) -> [a] -> [b]
map Expr 'EWord -> Prop
minProp (Set (Expr 'EWord) -> [Expr 'EWord]
forall a. Set a -> [a]
Set.toList KeccakStore
st.keccakExprs)
    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)]
keccakPairs
     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
.|| (Expr 'Buf -> Expr 'EWord
bufLength Expr 'Buf
a Expr 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
./= Expr 'Buf -> Expr 'EWord
bufLength 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 'EWord -> Expr 'EWord -> Prop
forall (a :: EType). Typeable a => Expr a -> Expr a -> Prop
PEq Expr 'EWord
lit Expr a
Expr 'EWord
e]
      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