{-# 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
  { $sel:bufs:BuilderState :: Map (Expr 'Buf) Int
bufs = forall a. Monoid a => a
mempty
  , $sel:stores:BuilderState :: Map (Expr 'Storage) Int
stores = forall a. Monoid a => a
mempty
  , $sel:count:BuilderState :: 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
s.bufs 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
s.count
          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
s.bufs
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{$sel:bufs:BuilderState :: Map (Expr 'Buf) Int
bufs=Map (Expr a) Int
bs', $sel:count:BuilderState :: 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
s.bufs 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
s.count
          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
s.bufs
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{$sel:bufs:BuilderState :: Map (Expr 'Buf) Int
bufs=Map (Expr a) Int
bs', $sel:count:BuilderState :: 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
s.bufs 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
s.count
          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
s.bufs
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{$sel:count:BuilderState :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1, $sel:bufs:BuilderState :: 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
s.stores 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
s.count
          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
s.stores
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BuilderState
s{$sel:count:BuilderState :: Int
count=Int
nextforall a. Num a => a -> a -> a
+Int
1, $sel:stores:BuilderState :: 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
st.bufs, forall a. Map a Int -> Map Int a
invertKeyVal BuilderState
st.stores)

-- | 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
st.bufs,  forall a. Map a Int -> Map Int a
invertKeyVal BuilderState
st.stores)