module Grin.Unboxing(unboxReturnValues) where

import Maybe
import Monad
import qualified Data.Map as Map
import qualified Data.Set as Set

import StringTable.Atom
import GenUtil
import Grin.Grin
import Support.CanType
import Support.Tuple
import Util.Graph


tailcalls :: Lam -> Set.Set Atom
tailcalls (_ :-> e) = f e where
    f (_ :>>= l) = tailcalls l
    f App { expFunction = fn } = Set.singleton fn
    f Case { expAlts = as } = Set.unions (map tailcalls as)
    f _ = Set.empty

unboxingCandidate :: Item -> Bool
unboxingCandidate item = isJust (unboxFunction undefined item)


isEnum (NV _ []) = True
isEnum _ = False


unboxFunction :: Monad m => Atom -> Item -> m (Exp -> Exp, Exp -> Exp, Ty, Item)
unboxFunction _ x | getType x == tyUnit = fail "unboxFunction: return type is already ()"
-- get rid of any fully constant values in return
unboxFunction fn item | any isLeft rvs = return (unboxReturn, unboxCall, returnType, nvs) where
    vs = fromTuple item
    rvs = [ case constantItem v of Just x -> Left x ; _ -> Right v | v <- vs ]
    nvs = tuple (rights rvs)
    returnType = getType nvs
    unboxReturn e = e :>>= tuple vars :-> Return (tuple vars')
    unboxCall (App a as _) | a == fn = App a as returnType :>>= tuple vars' :-> Return (tuple [ case x of Left x -> x ; Right _ -> v |  v <- vars | x <- rvs ])
    vars  = [Var v t | v <- [v1 ..] | t <- map getType vs ]
    vars' = concat [ perhapsM (isRight r) (Var v t)  | v <- [v1 ..] | t <- map getType vs | r <- rvs ]

-- unbox enumerated types
unboxFunction fn (NodeValue vs) | all isEnum (Set.toList vs) = return (unboxReturn, unboxCall, TyTag, itemTag) where
    unboxReturn (Return (NodeC t [])) = Return (Tag t)
    unboxReturn e = e :>>= nodev :-> Return var
    unboxCall (App a as ty) = App a as TyTag :>>= var :-> Return nodev
    var = Var v1 TyTag
    nodev = NodeV v1 []

-- returning a known node type
unboxFunction fn (NodeValue vs) | [NV t args] <- Set.toList vs  =  let
    returnType = tuple (map getType args)
    unboxReturn (Return (NodeC t' xs))
        | t == t' = Return (tuple xs)
        | otherwise = error "returning wrong node"
    unboxReturn e = e :>>= NodeC t vars :-> Return (tuple vars)
    unboxCall (App a as _) | a == fn = App a as returnType :>>= tuple vars :-> Return (NodeC t vars)
    vars  = [Var v t | v <- [v1 ..] | t <- map getType args ]
    in return (unboxReturn, unboxCall, returnType, tuple args)

unboxFunction _ item = fail "function not unboxable" -- (id,id,getType item)

constantItem (NodeValue vs) | [NV t xs] <- Set.toList vs  = do
    xs <- mapM constantItem xs
    return (NodeC t xs)
constantItem (TupledValue xs) = do
    xs <- mapM constantItem xs
    return (Tup xs)
constantItem (HeapValue vs) | [HV _ (Right val)] <- Set.toList vs  = do
    return (Const val)
constantItem _ = fail "not constant item"

{-# NOINLINE unboxReturnValues #-}
unboxReturnValues :: Grin -> IO Grin
unboxReturnValues grin = do
    let tcgraph = newGraph [ (n, Set.toList $ tailcalls body) | (n,body) <- grinFuncs  grin] fst snd
        ubc a | Just v <- Map.lookup a (grinReturnTags grin) = unboxingCandidate v
        ubc _ = False
        cfns = filter ubc (fsts $ grinFuncs grin)
        pf fn | Just item <- Map.lookup fn (grinReturnTags grin) =
            do x <- unboxFunction fn item ; return $ Map.singleton fn x
        fns = Map.unions $ concatMap pf cfns
        retTag fn _ | Just (_,_,_,ret) <- Map.lookup fn fns = ret
        retTag _ x = x
        retTe fn tyty | Just (_,_,ret,_) <- Map.lookup fn fns = tyty { tyReturn = ret }
        retTe _ x = x
        mtenv (TyEnv mp) = TyEnv $ Map.mapWithKey retTe mp
        doFunc (fn,lam) | Just (unboxReturn,_,_,_) <- Map.lookup fn fns = doFunc' (fn,convertReturns unboxReturn lam)
        doFunc (fn,lam) = doFunc' (fn,lam)
        doFunc' (fn,lam) = (fn, convertApps doApp lam)
        doApp ap@(App fn _ _) | Just (_,f,_,_) <- Map.lookup fn fns = f ap
        doApp e = e
    putStrLn "Unboxed return values"
    mapM_ putStrLn [ "  " ++ show fn ++ " - " ++  show nt | (fn,(_,_,nt,_)) <- Map.toList fns]

    let newgrin = setGrinFunctions (map doFunc (grinFuncs grin)) grin {
        grinReturnTags = Map.mapWithKey retTag (grinReturnTags grin),
        grinTypeEnv = mtenv (grinTypeEnv grin)
        }
    if Map.null fns then return newgrin else unboxReturnValues newgrin


convertReturns unboxReturn lam = g lam where
    g (l :-> e) = l :-> f e
    f (e :>>= l) = e :>>= g l
    f e@Case { expAlts = as } = e { expAlts = map g as }
    f e@Let { expBody = b } = e { expBody = f b }
    f e@MkCont { expCont = c , expLam = b } = e { expCont = g c, expLam = g b }
    f e = unboxReturn e

convertApps doApp lam = g lam where
    g (l :-> e) = l :-> f e
    f (e :>>= l) = f e :>>= g l
    f e@Case { expAlts = as } = e { expAlts = map g as }
    f e@Let { expDefs = defs, expBody = b } = e { expBody = f b, expDefs = [ createFuncDef True (funcDefName d) (g $ funcDefBody d) | d <- defs ] }
    f e@MkCont { expCont = c , expLam = b } = e { expCont = g c, expLam = g b }
    f e = doApp e