{-# Language DataKinds #-}

{- |
    Module: EVM.CSE
    Description: Common subexpression elimination for Expr ast
-}

module EVM.CSE (BufEnv, StoreEnv, eliminateExpr, eliminateProps) where

import Prelude hiding (Word, LT, GT)

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

import EVM.Types
import EVM.Traversals

-- maps expressions to variable names
data BuilderState = BuilderState
  { BuilderState -> Map (Expr 'Buf) Int
bufs :: Map (Expr Buf) Int
  , BuilderState -> Map (Expr 'Storage) Int
stores :: Map (Expr Storage) Int
  , BuilderState -> Int
count :: Int
  }
  deriving (Int -> BuilderState -> ShowS
[BuilderState] -> ShowS
BuilderState -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BuilderState] -> ShowS
$cshowList :: [BuilderState] -> ShowS
show :: BuilderState -> String
$cshow :: BuilderState -> String
showsPrec :: Int -> BuilderState -> ShowS
$cshowsPrec :: Int -> BuilderState -> ShowS
Show)

type BufEnv = Map Int (Expr Buf)
type StoreEnv = Map Int (Expr Storage)

initState :: BuilderState
initState :: BuilderState
initState = BuilderState
  { bufs :: Map (Expr 'Buf) Int
bufs = forall a. Monoid a => a
mempty
  , stores :: Map (Expr 'Storage) Int
stores = forall a. Monoid a => a
mempty
  , count :: Int
count = Int
0
  }


go :: Expr a -> State BuilderState (Expr a)
go :: forall (a :: EType). Expr a -> State BuilderState (Expr a)
go = \case
  -- buffers
  e :: Expr a
e@(WriteWord {}) -> do
    BuilderState
s <- forall s (m :: * -> *). MonadState s m => m s
get
    case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr a
e (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s) of
      Just Int
v -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
v)
      Maybe Int
Nothing -> do
        let
          next :: Int
next = BuilderState -> Int
count BuilderState
s
          bs' :: Map (Expr a) Int
bs' = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Expr a
e Int
next (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s)
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{bufs :: Map (Expr 'Buf) Int
bufs=Map (Expr a) Int
bs', count :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1}
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
next)
  e :: Expr a
e@(WriteByte {}) -> do
    BuilderState
s <- forall s (m :: * -> *). MonadState s m => m s
get
    case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr a
e (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s) of
      Just Int
v -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
v)
      Maybe Int
Nothing -> do
        let
          next :: Int
next = BuilderState -> Int
count BuilderState
s
          bs' :: Map (Expr a) Int
bs' = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Expr a
e Int
next (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s)
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{bufs :: Map (Expr 'Buf) Int
bufs=Map (Expr a) Int
bs', count :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1}
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
next)
  e :: Expr a
e@(CopySlice {}) -> do
    BuilderState
s <- forall s (m :: * -> *). MonadState s m => m s
get
    case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr a
e (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s) of
      Just Int
v -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
v)
      Maybe Int
Nothing -> do
        let
          next :: Int
next = BuilderState -> Int
count BuilderState
s
          bs' :: Map (Expr a) Int
bs' = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Expr a
e Int
next (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
s)
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{count :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1, bufs :: Map (Expr 'Buf) Int
bufs=Map (Expr a) Int
bs'}
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Buf
BufVar Int
next)
  -- storage
  e :: Expr a
e@(SStore {}) -> do
    BuilderState
s <- forall s (m :: * -> *). MonadState s m => m s
get
    case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr a
e (BuilderState -> Map (Expr 'Storage) Int
stores BuilderState
s) of
      Just Int
v -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Storage
StoreVar Int
v)
      Maybe Int
Nothing -> do
        let
          next :: Int
next = BuilderState -> Int
count BuilderState
s
          ss' :: Map (Expr a) Int
ss' = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Expr a
e Int
next (BuilderState -> Map (Expr 'Storage) Int
stores BuilderState
s)
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{count :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1, stores :: Map (Expr 'Storage) Int
stores=Map (Expr a) Int
ss'}
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (a :: EType). GVar a -> Expr a
GVar (Int -> GVar 'Storage
StoreVar Int
next)
  Expr a
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr a
e

invertKeyVal :: forall a. Map a Int -> Map Int a
invertKeyVal :: forall a. Map a Int -> Map Int a
invertKeyVal =  forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\(a
x, Int
y) -> (Int
y, a
x)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList

-- | Common subexpression elimination pass for Expr
eliminateExpr' :: Expr a -> State BuilderState (Expr a)
eliminateExpr' :: forall (a :: EType). Expr a -> State BuilderState (Expr a)
eliminateExpr' Expr a
e = forall (m :: * -> *) (b :: EType).
Monad m =>
(forall (a :: EType). Expr a -> m (Expr a)) -> Expr b -> m (Expr b)
mapExprM forall (a :: EType). Expr a -> State BuilderState (Expr a)
go Expr a
e

eliminateExpr :: Expr a -> (Expr a, BufEnv, StoreEnv)
eliminateExpr :: forall (a :: EType). Expr a -> (Expr a, BufEnv, StoreEnv)
eliminateExpr Expr a
e =
  let (Expr a
e', BuilderState
st) = forall s a. State s a -> s -> (a, s)
runState (forall (a :: EType). Expr a -> State BuilderState (Expr a)
eliminateExpr' Expr a
e) BuilderState
initState in
  (Expr a
e', forall a. Map a Int -> Map Int a
invertKeyVal (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
st), forall a. Map a Int -> Map Int a
invertKeyVal (BuilderState -> Map (Expr 'Storage) Int
stores BuilderState
st))

-- | Common subexpression elimination pass for Prop
eliminateProp' :: Prop -> State BuilderState Prop
eliminateProp' :: Prop -> State BuilderState Prop
eliminateProp' Prop
prop = forall (m :: * -> *).
Monad m =>
(forall (a :: EType). Expr a -> m (Expr a)) -> Prop -> m Prop
mapPropM forall (a :: EType). Expr a -> State BuilderState (Expr a)
go Prop
prop

-- | Common subexpression elimination pass for list of Prop
eliminateProps' :: [Prop] -> State BuilderState [Prop]
eliminateProps' :: [Prop] -> State BuilderState [Prop]
eliminateProps' [Prop]
props = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Prop -> State BuilderState Prop
eliminateProp' [Prop]
props


-- | Common subexpression elimination pass for list of Prop
eliminateProps :: [Prop] -> ([Prop], BufEnv, StoreEnv)
eliminateProps :: [Prop] -> ([Prop], BufEnv, StoreEnv)
eliminateProps [Prop]
props =
  let ([Prop]
props', BuilderState
st) = forall s a. State s a -> s -> (a, s)
runState ([Prop] -> State BuilderState [Prop]
eliminateProps' [Prop]
props) BuilderState
initState in
  ([Prop]
props',  forall a. Map a Int -> Map Int a
invertKeyVal (BuilderState -> Map (Expr 'Buf) Int
bufs BuilderState
st),  forall a. Map a Int -> Map Int a
invertKeyVal (BuilderState -> Map (Expr 'Storage) Int
stores BuilderState
st))