module Grin.Optimize(grinPush,grinSpeculate) where

import Control.Monad.State
import List
import qualified Data.Set as Set

import StringTable.Atom
import C.Prims
import Grin.Grin
import Grin.Noodle
import Stats hiding(null)
import Support.CanType
import Support.FreeVars
import Util.Graph
import Util.SetLike


data PExp = PExp {
    pexpUniq :: Int,
    pexpBind :: [Val],
    pexpExp  :: Exp,
    pexpProvides :: [Var],
    pexpDeps :: [Int]
    } deriving(Show)

instance Eq PExp where
    a == b = pexpUniq a == pexpUniq b

makeDeps :: [PExp] -> PExp -> PExp
makeDeps cs pexp = pexp { pexpProvides = freeVars (pexpBind pexp), pexpDeps = deps } where
    deps = [ pexpUniq c | c <- cs, not $ null $ fvs `intersect` pexpProvides c ]
    fvs = freeVars (pexpExp pexp)

justDeps :: [PExp] -> [Var] -> [Int]
justDeps cs fs = deps where
    deps = [ pexpUniq c | c <- cs, not $ null $ fs `intersect` pexpProvides c ]

-- | grinPush pushes the definitions of variables as far inward as they can go so
-- peephole optimizations have a better chance of firing. when the order of definitons
-- doesn't matter, it uses heuristics to decide which one to push to allow the most
-- peephole optimizations.

grinPush :: Stats -> Lam -> IO Lam
grinPush stats (l :-> e) = ans where
    ans = do
--        putStrLn "@@@ grinPush"
        e' <- evalStateT (f e) (1,[])
        return (l :-> e')
    f (exp :>>= v :-> e2) | isOmittable exp = do
        (nn,cv) <- get
        let npexp = makeDeps cv PExp { pexpUniq = nn, pexpBind = v, pexpExp = exp, pexpDeps = undefined, pexpProvides = undefined }
        put (nn+1,npexp:cv)
        f e2
    f (exp :>>= v :-> e2) = do
        exp <- fixupLet exp
        (v',exp') <- dropAny (Just v) exp
        e2' <- f e2
        return $ exp' :>>= v' :-> e2'
    f exp = do
        exp <- fixupLet exp
        (_,exp') <- dropAny Nothing exp
        return exp'

    fixupLet lt@Let { expDefs = defs, expBody = b } = do
        let def = (Set.fromList $ map funcDefName defs)
            f (e :>>= l :-> r) | Set.null (freeVars e `Set.intersection` def) = do
                exp <- f r
                return (e :>>= l :-> exp)
            f r = return $ updateLetProps lt {  expBody = r }
        f b
    fixupLet exp = return exp
    dropAny mv (exp::Exp) = do
        (nn,xs) <- get
        let graph = newGraph xs pexpUniq pexpDeps
            deps = justDeps xs (freeVars exp)
            reached = reachable graph deps
            --dropped = case prefered reached exp of
            --    Just (x:_) | [] <- [ r | r <- reached, pexpUniq x `elem` pexpDeps r ] -> (reverse $ topSort $ newGraph (filter (/= x) reached) pexpUniq pexpDeps) ++ [x]
            --    _ -> reverse $ topSort $ newGraph reached pexpUniq pexpDeps
            dropped =  reverse $ topSort $ newGraph reached pexpUniq pexpDeps
            ff pexp exp = pexpExp pexp :>>= pexpBind pexp :-> exp
            ebinds = [ Var v t | (v,t) <- Set.toList $ freeVars (map pexpBind dropped) ]
            (exp',mv') | Just vv <- mv = let mv' = vv ++ ebinds in (exp :>>= vv :-> Return mv',mv')
                       | otherwise = (exp,[])
        put (nn,[ x | x <- xs, pexpUniq x `notElem` (map pexpUniq reached) ])
--        when (not $ null dropped) $ lift $ do
--            putStrLn "@@@ dropped"
--            mapM_ Prelude.print dropped
        return (mv',foldr ff exp' dropped :: Exp)
    -- | preferentially pull definitons of the variable this returns right next to it as it admits a peephole optimization
--    prefer (Store v@Var {}) = return v
--    prefer (App fn [v@Var {}] _)  | fn == funcEval = return v
--    prefer (App fn [v@Var {},_] _)| fn == funcApply = return v
--    prefer (App fn [v@Var {}] _)  | fn == funcApply = return v
--    prefer (Update _ v@Var {}) = return v
--    prefer (Update v@Var {} _) = return v
--    prefer _ = fail "no preference"
--    _prefered pexps exp = do
--        v <- prefer exp
--        return [ p | p <- pexps, v == pexpBind p]

--grinPush :: Stats -> Lam -> IO Lam
--grinPush stats lam = ans where
--    ans = do
--        putStrLn "@@@ grinPush"
--        (ans,_) <- evalStateT (whiz subBlock doexp finalExp whizState lam) (1,[])
--        return ans
--    subBlock _ action = do
--        (nn,x) <- get
--        put (nn,mempty)
--        r <- action
--        (nn,_) <- get
--        put (nn,x)
--        return r
--    doexp (v, exp) | isOmittable exp = do
--        (nn,cv) <- get
--        let npexp = makeDeps cv PExp { pexpUniq = nn, pexpBind = v, pexpExp = exp, pexpDeps = undefined, pexpProvides = undefined }
--        put (nn+1,npexp:cv)
--        return Nothing
--    doexp (v, exp) = do
--        exp <- fixupLet exp
--        (v',exp') <- dropAny (Just v) exp
--        return $ Just (v',exp')
--    finalExp (exp::Exp) = do
--        exp <- fixupLet exp
--        (_,exp') <- dropAny Nothing exp
--        return (exp'::Exp)
--    fixupLet lt@Let { expDefs = defs, expBody = b } = do
--        let def = (Set.fromList $ map funcDefName defs)
--            f (e :>>= l :-> r) | Set.null (freeVars e `Set.intersection` def) = do
--                exp <- f r
--                return (e :>>= l :-> exp)
--            f r = return $ updateLetProps lt {  expBody = r }
--        f b
--    fixupLet exp = return exp
--    dropAny mv (exp::Exp) = do
--        (nn,xs) <- get
--        let graph = newGraph xs pexpUniq pexpDeps
--            deps = justDeps xs (freeVars exp)
--            reached = reachable graph deps
--            --dropped = case prefered reached exp of
--            --    Just (x:_) | [] <- [ r | r <- reached, pexpUniq x `elem` pexpDeps r ] -> (reverse $ topSort $ newGraph (filter (/= x) reached) pexpUniq pexpDeps) ++ [x]
--            --    _ -> reverse $ topSort $ newGraph reached pexpUniq pexpDeps
--            dropped =  reverse $ topSort $ newGraph reached pexpUniq pexpDeps
--            ff pexp exp = pexpExp pexp :>>= pexpBind pexp :-> exp
--            ebinds = [ Var v t | (v,t) <- Set.toList $ freeVars (map pexpBind dropped) ]
--            (exp',mv') | Just vv <- mv = let mv' = tuple $ fromTuple vv ++ ebinds in (exp :>>= vv :-> Return mv',mv')
--                       | otherwise = (exp,unit)
--        put (nn,[ x | x <- xs, pexpUniq x `notElem` (map pexpUniq reached) ])
--        when (not $ null dropped) $ lift $ do
--            putStrLn "@@@ dropped"
--            mapM_ Prelude.print dropped
--        return (mv',foldr ff exp' dropped :: Exp)
--    -- | preferentially pull definitons of the variable this returns right next to it as it admits a peephole optimization
--    prefer (Store v@Var {}) = return v
--    prefer (App fn [v@Var {}] _)  | fn == funcEval = return v
--    prefer (App fn [v@Var {},_] _)| fn == funcApply = return v
--    prefer (App fn [v@Var {}] _)  | fn == funcApply = return v
--    prefer (Update _ v@Var {}) = return v
--    prefer (Update v@Var {} _) = return v
--    prefer _ = fail "no preference"
--    prefered pexps exp = do
--        v <- prefer exp
--        return [ p | p <- pexps, v == pexpBind p]



grinSpeculate :: Grin -> IO Grin
grinSpeculate grin = do
    let ss = findSpeculatable grin
    putStrLn "Speculatable:"
    mapM_ Prelude.print ss
    let (grin',stats) = runStatM (performSpeculate ss grin)
    Stats.printStat "Speculate" stats
    return grin'


performSpeculate specs grin = do
    let sset = Set.fromList (map tagFlipFunction specs)
    let f (a,l) = mapBodyM h l  >>= \l' -> return (a,l')
        h (Store (NodeC t xs)) | t `member` sset = do
            let t' = tagFlipFunction t
            mtick $ "Optimize.speculate.store.{" ++ show t'
            return (App t' xs [TyNode] :>>= [n1] :-> Store n1)
--        h (Update v (NodeC t xs)) | not (isMutableNodeTag t), t `member` sset = do
--            let t' = tagFlipFunction t
--            mtick $ "Optimize.speculate.update.{" ++ show t'
--            return (App t' xs [TyNode] :>>= [n1] :-> Update v n1)
        h e = mapExpExp h e
    fs <- mapM f (grinFuncs grin)
    return $ setGrinFunctions fs grin

findSpeculatable :: Grin -> [Atom]
findSpeculatable grin = ans where
    ans = [ x | Left (x,_) <- scc graph ]
    graph = newGraph [ (a,concatMap f (freeVars l)) | (a,_ :-> l) <- grinFuncs grin, isSpeculatable l, getType l == [TyNode] ] fst snd
    f t | tagIsSuspFunction t = [tagFlipFunction t]
        | tagIsFunction t = [t]
        | otherwise = []
    isSpeculatable Return {} = True
    isSpeculatable Store {} = True
    isSpeculatable (x :>>= _ :-> y) = isSpeculatable x && isSpeculatable y
    isSpeculatable (Case e as) = all isSpeculatable [ e | _ :-> e <- as]
    isSpeculatable Prim { expPrimitive = APrim p _ } = primIsConstant p
    isSpeculatable _ = False