{-# 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
  { bufs :: Map (Expr Buf) Int
  , stores :: Map (Expr Storage) Int
  , count :: Int
  }
  deriving (Show)

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

initState :: BuilderState
initState = BuilderState
  { bufs = mempty
  , stores = mempty
  , count = 0
  }


go :: Expr a -> State BuilderState (Expr a)
go = \case
  -- buffers
  e@(WriteWord {}) -> do
    s <- get
    case Map.lookup e s.bufs of
      Just v -> pure $ GVar (BufVar v)
      Nothing -> do
        let
          next = s.count
          bs' = Map.insert e next s.bufs
        put $ s{bufs=bs', count=next+1}
        pure $ GVar (BufVar next)
  e@(WriteByte {}) -> do
    s <- get
    case Map.lookup e s.bufs of
      Just v -> pure $ GVar (BufVar v)
      Nothing -> do
        let
          next = s.count
          bs' = Map.insert e next s.bufs
        put $ s{bufs=bs', count=next+1}
        pure $ GVar (BufVar next)
  e@(CopySlice {}) -> do
    s <- get
    case Map.lookup e s.bufs of
      Just v -> pure $ GVar (BufVar v)
      Nothing -> do
        let
          next = s.count
          bs' = Map.insert e next s.bufs
        put $ s{count=next+1, bufs=bs'}
        pure $ GVar (BufVar next)
  -- storage
  e@(SStore {}) -> do
    s <- get
    case Map.lookup e s.stores of
      Just v -> pure $ GVar (StoreVar v)
      Nothing -> do
        let
          next = s.count
          ss' = Map.insert e next s.stores
        put $ s{count=next+1, stores=ss'}
        pure $ GVar (StoreVar next)
  e -> pure e

invertKeyVal :: forall a. Map a Int -> Map Int a
invertKeyVal =  Map.fromList . map (\(x, y) -> (y, x)) . Map.toList

-- | Common subexpression elimination pass for Expr
eliminateExpr' :: Expr a -> State BuilderState (Expr a)
eliminateExpr' e = mapExprM go e

eliminateExpr :: Expr a -> (Expr a, BufEnv, StoreEnv)
eliminateExpr e =
  let (e', st) = runState (eliminateExpr' e) initState in
  (e', invertKeyVal st.bufs, invertKeyVal st.stores)

-- | Common subexpression elimination pass for Prop
eliminateProp' :: Prop -> State BuilderState Prop
eliminateProp' prop = mapPropM go prop

-- | Common subexpression elimination pass for list of Prop
eliminateProps' :: [Prop] -> State BuilderState [Prop]
eliminateProps' props = mapM eliminateProp' props


-- | Common subexpression elimination pass for list of Prop
eliminateProps :: [Prop] -> ([Prop], BufEnv, StoreEnv)
eliminateProps props =
  let (props', st) = runState (eliminateProps' props) initState in
  (props',  invertKeyVal st.bufs,  invertKeyVal st.stores)