module Grin.SSimplify(simplify) where

import qualified Data.IntSet as IS
import qualified Data.IntMap as IM
import qualified Data.Map as Map
import qualified Data.Set as Set

import StringTable.Atom
import Grin.Grin
import Grin.Noodle
import Util.Gen
import Util.RWS
import Support.CanType
import Support.FreeVars
import qualified Stats
import Stats(mtick)

-- This goes through and puts grin into a normal form, in addition, it carries out some straightforward
-- simplifications.
--
-- normalized form has the following properties
--
-- :>>= only appears in trailing position
-- Return [v0 .. vn] for n > 1 only appears in trailing position
--
-- all variables and function names are unique in their scope.


data SEnv = SEnv {
    envSubst :: IM.IntMap Val,   -- renaming substitution
    envCSE   :: Map.Map Exp (Atom,Exp),
    envPapp  :: IM.IntMap (Atom,[Val]),
    envPush  :: IM.IntMap Exp
    }
    {-! derive: Monoid !-}

newtype SState = SState { usedVars :: IS.IntSet }

data SCol = SCol {
    colStats :: Stats.Stat,
    colFreeVars :: Set.Set Var
    }
    {-! derive: Monoid !-}

newtype S a = S (RWS SEnv SCol SState a)
    deriving(Monad,Functor,MonadWriter SCol, MonadReader SEnv,MonadState SState)

instance Stats.MonadStats S where
    mtickStat s = S (tell mempty { colStats = s })
    mticks' n a = S (tell mempty { colStats = Stats.singleStat n a })


tellFV v = tell mempty { colFreeVars = freeVars v }


simplify :: Grin -> IO Grin
simplify grin = do
    let (fs,_,SCol { colStats = stats}) = runRWS fun mempty SState { usedVars = mempty }
        S fun = simpFuncs (grinFunctions grin)
    return grin { grinFunctions = fs, grinStats = grinStats grin `mappend` stats }


simpFuncs :: [FuncDef] -> S [FuncDef]
simpFuncs fd = do
    let f fd@FuncDef { funcDefBody = body } = do
            body' <- simpLam body
            return $ updateFuncDefProps fd { funcDefBody = body' }
    mapM f fd

simpLam :: Lam -> S Lam
simpLam (ps :-> e) = do
    (ps,env') <- renamePattern ps
    let f col = col { colFreeVars = colFreeVars col Set.\\ freeVars ps }
    (e,col) <- censor f $ listen $ local (env' `mappend`) $ simpExp e
    ps <- mapM (zeroVars (`Set.member` colFreeVars col)) ps
    return (ps :-> e)

simpDone :: Exp -> S Exp
simpDone e = do
    pmap <- asks envPapp
    case e of
        (App fap (Var (V vn) _:fs) ty) | fap == funcApply, Just (tl,gs) <- IM.lookup vn pmap -> do
            (cl,fn) <- tagUnfunction tl
            let ne = if cl == 1 then App fn (gs ++ fs) ty else Return [NodeC (partialTag fn (cl - 1)) (gs ++ fs)]
            mtick $ if cl == 1 then "Simplify.Apply.Papp.{" ++ show tl  else ("Simplify.Apply.App.{" ++ show fn)
            return ne
        _ -> do
            cmap <- asks envCSE
            case Map.lookup e cmap of
                Just (n,e') -> do mtick n; return e'
                Nothing -> return e

simpBind :: [Val] -> Exp -> S Exp -> S Exp
simpBind p e cont = f p e where
    cse name xs = do
        (z,col) <- listen $ local (\s -> s { envCSE = Map.fromList [ (x,(toAtom name,y)) | (x,y) <- xs] `Map.union` envCSE s }) cont
        e <- simpDone e
        if isOmittable e && Set.null (freeVars p `Set.intersection` colFreeVars col) then do
            mtick "Simplify.Omit.Bind"
            return z
         else return $ e :>>= (p :-> z)
    cse' name xs = cse name ((e,Return p):xs)
    f p app@(App a [v] _) | a == funcEval =  cse' "Simplify.CSE.eval" [(Fetch v,Return p)]
    f p (Fetch v@Var {}) =  cse' "Simplify.CSE.fetch" [(gEval v,Return p)]
    f [p@(Var (V vn) _)] (Return [v@(NodeC t vs)]) | not (isHoly v) = case tagUnfunction t of
        Nothing -> cse "Simplify.CSE.return-node" [(Return [p],Return [v]),(Store p,Store v)]
        Just (n,fn) -> local (\s -> s { envPapp = IM.insert vn (t,vs) (envPapp s) }) $ cse' "Simplify.CSE.return-node" [(Return [p],Return [v]),(Store p,Store v)]
    f [p] (Store v@Var {})  =  cse' "Simplify.CSE.demote" [(Fetch p,Return [v]),(gEval p,Return [v])]
    f [p@(Var (V vn) _)] (Store v@(NodeC t vs)) | not (isHoly v) = case tagIsWHNF t of
        True -> local (\s -> s { envPush = IM.insert vn (Store v) (envPush s) }) $ cse "Simplify.CSE.store-whnf" [(Fetch p,Return [v]),(gEval p,Return [v])]
        False -> cse' "Simplify.CSE.store" []
    f _ _ = cse "Simplify.CSE.NOT" []

extEnv :: Var -> Val -> SEnv -> SEnv
extEnv (V vn) v s = s { envSubst = IM.insert vn v (envSubst s) }



simpExp :: Exp -> S Exp
simpExp e = f e [] where
    f  (a :>>= (v :-> b)) xs = do
        env <- ask
        f a ((env,v,b):xs)

    -- simple transforms
    f (Fetch (Const x)) rs = do
        mtick "Grin.Simplify.fetch-const"
        f (Return [x]) rs
    f (Store x) rs | valIsNF x = do
        mtick "Grin.Simplify.store-normalform"
        f (Return [Const x]) rs
    f (App a [Const n] _) rs | a == funcEval = do
        mtick "Grin.Simplify.eval-const"
        f (Return [n]) rs
    f (Error s t) rs@(_:_) = do
        mtick "Grin.Simplify.error-discard"
        let (_,_,b) = last rs
        f (Error s (getType b)) []

    f (Return [v@Const {}]) ((senv,[Var vn _],b):rs) = do
        mtick "Grin.Simplify.Subst.const"
        fbind vn v senv b rs
    f (Return [v@Var {}]) ((senv,[Var vn _],b):rs) = do
        mtick "Grin.Simplify.Subst.var"
        fbind vn v senv b rs
    f a@(Return [NodeC t xs]) ((senv,[NodeC t' ys],b):rs) | t == t' = do
        mtick "Grin.Simplify.Assign.node-node"
        dtup xs ys senv b rs
    f (Return []) ((senv,[],b):rs) = do
        mtick "Grin.Simplify.Assign.unit-unit"
        dtup [] [] senv b rs
    f a@(Return (xs@(_:_:_))) ((senv,ys,b):rs) = do
        mtick "Grin.Simplify.Assign.tuple-tuple"
        dtup xs ys senv b rs
    f (Case v [l]) rs = do
        f (Return [v] :>>= l) rs
    f a ((senv,p,b):xs) = do
        a <- g a
        (p,env') <- renamePattern p
        let env'' = env' `mappend` senv
        local (const env'') $ simpBind p a (f b xs)
    f x [] = do
        e <- g x
        simpDone e
    fbind vn v senv b rs = do
        v' <- applySubst v
        local (\_ -> extEnv vn v' senv) $ f b rs

    dtup xs ys senv b rs | sameLength xs ys = do
        xs <- mapM applySubst xs
        (ys,env') <- renamePattern ys
        let env'' = env' `mappend` senv
        z <- local (const env'') $ f b rs
        ts <- mapM (return . Just) [([y],Return [x]) | x <- xs | y <- ys ]
        let h [] = z
            h ((p,v):rs) = v :>>= p :-> h rs
        return $ h [ (p,v) |  Just (p,v) <- ts]
    dtup _ _ _ _ _ = error "dtup: attempt to bind unequal lists"
    g (Case v as) = do
        v <- applySubst v
        as <- mapM simpLam as
        return $ Case v as

    g  lt@Let { expDefs = defs, expBody = body } = do
        body <- f body []
        defs <- simpFuncs defs
        return $ updateLetProps lt { expBody = body, expDefs = defs }
    g x = applySubstE x



applySubstE :: Exp -> S Exp
applySubstE x = mapExpVal applySubst x

applySubst x = f x where
    f var@(Var (V v) _) = do
        env <- asks envSubst
        case IM.lookup v env of
            Just n -> tellFV n >> return n
            Nothing -> tellFV var >> return var
    f x = mapValVal f x

zeroVars fn x = f x where
    f (Var v ty) | fn v || v == v0 = return (Var v ty)
                 | otherwise = do mtick $ "Simplify.ZeroVar.{" ++ show (Var v ty); return (Var v0 ty)
    f x = mapValVal f x

renamePattern :: [Val] ->  S ([Val],SEnv)
renamePattern x = runWriterT (mapM f x) where
    f :: Val -> WriterT SEnv S Val
    f (Var v@(V vn) t) = do
        v' <- lift $ newVarName v
        let nv = Var v' t
        tell (mempty { envSubst = IM.singleton vn nv })
        return nv
    f x = mapValVal f x

newVarName :: Var -> S Var
newVarName (V 0) = return (V 0)
newVarName (V sv) = do
    s <- gets usedVars
    let nv = v sv
        v n | n `IS.member` s = v (1 + n + IS.size s)
            | otherwise = n
    modify (\e -> e { usedVars = IS.insert nv s })
    return (V nv)




isHoly (NodeC _ as) | any isValUnknown as = True
isHoly n = False