\ignore{
\begin{code}
module GCodeUtils ( localAllocator, tailCallify, qualify, fallsThrough ) where

import GCode
import GCodeSysDeps

import FastString (nullFS)
import Outputable
import OrdList
import UniqFM
import UniqSet
import Util

#include "HsVersions.h"
\end{code}
}

\subsection{The GCode local variable allocator}

\begin{code}
localAllocator :: Module -> GenModule Int
localAllocator m = m
    {m_closures=map (\c -> c
        {c_code=fst (runLaM (laStmnts (c_code c)) env)
        }) (m_closures m)
    ,m_init = fst (runLaM (laStmnts (m_init m)) top_env)}
    where
        -- The initial environment is sort of a hack. We never want to nuke
        -- or reuse slot 0 (thisNode) so we give it a size of 0 which 
        -- prevents it from being added to the freelist
        env = LaEnv (unitUFM thisUnique (0,0)) [] 1
        top_env = LaEnv emptyUFM [] 0


data LaEnv = LaEnv
    { la_binds :: UniqFM (Int,Int)
    , la_free :: [(Int,Bool)]
    , la_next :: Int }

newtype LaM a = LaM { runLaM :: LaEnv -> (a,LaEnv) }

instance Monad LaM where
    return x = LaM $ \s -> (x,s)
    m >>= k = LaM  $ \s -> let (a,s') = runLaM m s in runLaM (k a) s'

laGet :: LaM LaEnv
laGet = LaM $ \s -> (s,s)

laGets :: (LaEnv -> a) -> LaM a
laGets sel = laGet >>= return . sel

laPut :: LaEnv -> LaM ()
laPut s = LaM $ \_ -> ((),s)

laModify :: (LaEnv -> LaEnv) -> LaM ()
laModify f = LaM $ \s -> ((),f s)
    
laStmnts :: Stmnts -> LaM LStmnts
laStmnts ss = do
    sss <- mapM laStmnt (fromOL ss)
    return (concatOL sss)

laStmnt :: Stmnt -> LaM LStmnts
laStmnt (Move src dst) = do
    src' <- laLoc False src
    dst' <- laLoc True dst
    return (unitOL (Move src' dst'))
laStmnt (Call call) = do
    nukes <- laNuke
    call' <- laCall call
    return (nukes `snocOL` Call call')
laStmnt (TailCall call) = do
    -- No laNuke here, tailcalling will kill the locals
    call' <- laCall call 
    return (unitOL (TailCall call'))
laStmnt Return = return (unitOL Return)
laStmnt (Update loc) = laLoc False loc >>= \loc' -> return (unitOL (Update loc'))
laStmnt (Alloc qn loc) = laLoc True loc >>= \loc' -> return (unitOL (Alloc qn loc'))
laStmnt (Switch loc arms def) = do
    loc' <- laLoc False loc
    ((_,def'):arms') <- laFork laArm ((undefined,def):arms)
    return (unitOL (Switch loc' arms' def'))
laStmnt (PrimOp op tys args rets) = do
    args' <- mapM (laLoc False) args
    rets' <- mapM (laLoc True) rets
    return (unitOL (PrimOp op tys args' rets'))
laStmnt (Foreign fun sig args rets) = do
    args' <- mapM (laLoc False) args
    rets' <- mapM (laLoc True) rets
    return (unitOL (Foreign fun sig args' rets'))
laStmnt (Trim live) = do
    laKeep live
    binds <- laGets la_binds
    return (unitOL (Trim (mkUniqSet [i|Just (i,_) <- map (lookupUFM binds) (uniqSetToList live)])))

laFork :: (a -> LaM b) -> [a] -> LaM [b]
laFork f xs = do
    env <- laGet
    let (xs',envs) = mapAndUnzip (\x -> runLaM (f x) env) xs
    laPut (foldl1 merge envs)
    return xs'
    where
        merge (LaEnv binds1 free1 next1) (LaEnv binds2 free2 next2)
            = LaEnv
                (intersectUFM_C merge_bind binds1 binds2 `delFromUFM` mkPseudoUniqueC 1)
                (merge_free free1 free2)
                (max next1 next2)
        merge_bind x y
            | x == y = x
            | otherwise = pprPanic "merge_bind" (ppr x <+> ppr y)
        merge_free xs [] = xs
        merge_free [] ys = ys
        merge_free (x@(i,p):xs) (y@(j,q):ys) = case compare i j of
            LT -> x : merge_free xs (y:ys)
            EQ -> (i,p || q) : merge_free xs ys
            GT -> y : merge_free (x:xs) ys

laLoc ::Bool ->  Loc -> LaM LLoc
laLoc _ (Reg t i) = return (Reg t i)
laLoc w (Var t v) = do
    v' <- laVar w v (typeSize t)
    return (Var t v')
laLoc _ (Field l t i) = laLoc False l >>= \l' -> return (Field l' t i)
laLoc _ (Global t qn) = return (Global t qn)
laLoc _ (Literal literal) = return (Literal literal)

laCall :: Call -> LaM (GenCall Int)
laCall (FunCall x) = laLoc False x >>= \x' -> return (FunCall x')
laCall (RtsCall n) = return (RtsCall n)

laArm :: (Tag,Stmnts) -> LaM (Tag,LStmnts)
laArm (tag,ss) = laStmnts ss >>= \ss' -> return (tag,ss')

-- Based on CgStackery.allocPrimStack
laVar :: Bool -> Unique -> Int -> LaM Int
laVar write v size = do
    binds <- laGets la_binds
    case binds `lookupUFM` v of
        Just (i,_) -> ASSERT2(not write, ppr v) return i
        Nothing -> ASSERT2(write, ppr v) do
            free <- laGets la_free
            i <- case find_free free of
                Just i -> do
                    laModify (\s -> s{la_free=delete_free free i})
                    return i
                Nothing -> do
                    i <- laGets la_next
                    laModify (\s -> s{la_next=i+size})
                    return i
            laModify (\s -> s{la_binds=addToUFM binds v (i,size)})
            return i
    where
        find_free [] = Nothing
        find_free (x@(i,_):xs)
            | take size [j | (j,_) <- x:xs] == [i..i+size-1] = Just i
            | otherwise = find_free xs
        delete_free xs slot = [x | x@(i,_) <- xs, i >= slot+size || i < slot]

-- Based on CgBindery.nukeDeadBindings
laKeep :: UniqSet Unique -> LaM ()
laKeep live_vars = do
    binds <- laGets la_binds
    let (extra_free,binds') = dead_slots [] [] (ufmToList binds)
    free <- laGets la_free
    laModify (\s->s
        {la_binds=listToUFM binds'
        ,la_free=add_free_slots free [(i,True)|i<-sortLe (<=) extra_free]})
    where
        dead_slots fbs ds [] = (ds,reverse fbs)
        dead_slots fbs ds (b@(v,(i,size)):bs)
            | v `elementOfUniqSet` live_vars = dead_slots (b:fbs) ds bs
            | size > 0 = dead_slots fbs ([i..i+size-1]++ds) bs
            | otherwise = dead_slots fbs ds bs
        add_free_slots xs [] = xs
        add_free_slots [] ys = ys
        add_free_slots (x@(i,_):xs) (y@(j,_):ys)
            | i < j     = x : add_free_slots xs (y:ys)
            | otherwise = y : add_free_slots (x:xs) ys

laNuke :: LaM LStmnts
laNuke = do
    free <- laGets la_free
    laModify (\s->s{la_free=[(i,False)|(i,_)<-free]})
    return $ toOL $ 
        [Move (Literal (LNull LiftedTy)) (Var LiftedTy i)
        |(i,needsnuke) <- free, needsnuke]
\end{code}

\begin{code}
tailCallify :: Outputable a => GenModule a -> GenModule a
tailCallify m = m{m_closures=map (\c -> c{c_code=tcStmnts (c_code c)}) (m_closures m)}

tcStmnts :: Outputable a => GenStmnts a -> GenStmnts a
tcStmnts ss
    | not (isNilOL ss), (ss',Return) <- unsnocOL ss
    , not (isNilOL ss'), (body,final) <- unsnocOL ss'
    , Just final' <- tcStmnt final
    = body `snocOL` final'
    | otherwise = ss
    where

tcStmnt :: Outputable a => GenStmnt a -> Maybe (GenStmnt a)
tcStmnt (Switch loc arms def) = Just (Switch loc arms' def')
    where
        arms' = [(tag,tcStmnts (ss`snocOL`Return))|(tag,ss)<-arms] 
        def' = tcStmnts (def`snocOL`Return)
tcStmnt (Call call) = Just (TailCall call)
#ifdef DEBUG
tcStmnt s@(TailCall _) = pprPanic "tcStmnt" (ppr s)
tcStmnt s@Return = pprPanic "tcStmnt" (ppr s)
#endif
tcStmnt _ = Nothing
\end{code}

\subsection{Qualifier}

\begin{code}
qualify :: GenModule a -> GenModule a
qualify m = m
    {m_closures=map (\c -> c{c_code=qStmnts (m_name m) (c_code c)}) (m_closures m)
    ,m_init = qStmnts (m_name m) (m_init m)}

qStmnts :: Name -> GenStmnts a -> GenStmnts a
qStmnts m = mapOL (qStmnt m)

qStmnt :: Name -> GenStmnt a -> GenStmnt a
qStmnt m (Move src dst) = Move (qLoc m src) (qLoc m dst)
qStmnt m (Alloc qn loc) = Alloc (qQName m qn) (qLoc m loc)
qStmnt m (Switch loc arms def) = Switch (qLoc m loc) [(tag,qStmnts m ss)|(tag,ss)<-arms] (qStmnts m def)
qStmnt m (PrimOp op tys args rets) = PrimOp op (map (qType m) tys) (map (qLoc m) args) (map (qLoc m) rets)
qStmnt m (Foreign spec (atys,rtys) args rets) = Foreign spec (map (qType m) atys,map (qType m) rtys) (map (qLoc m) args) (map (qLoc m) rets)
qStmnt m (Call loc) = Call (qCall m loc)
qStmnt m (TailCall loc) = TailCall (qCall m loc)
qStmnt m (Update loc) = Update (qLoc m loc)
qStmnt _ s = s

qLoc :: Name -> GenLoc a -> GenLoc a
qLoc n (Var ty u) = Var (qType n ty) u
qLoc n (Field loc ty i) = Field (qLoc n loc) (qType n ty) i
qLoc n (Reg ty i) = Reg (qType n ty) i
qLoc n (Global ty qn) = Global ty (qQName n qn)
qLoc n (Literal (LNull ty)) = Literal (LNull (qType n ty))
qLoc _ loc = loc

qCall :: Name -> GenCall a -> GenCall a
qCall n (FunCall loc) = FunCall (qLoc n loc)
qCall _ call = call

qType :: Name -> Type -> Type
qType n (PrimTy tc args) = PrimTy tc (map (qType n) args)
qType _ LiftedTy = LiftedTy
qType n (ClosureTy qn) = ClosureTy (qQName n qn)

qQName :: Name -> QName -> QName
qQName m' qn@(m,n)
    | nullFS m = (m',n)
    | otherwise = qn 
\end{code}

\subsection{Misc}

\begin{code}
fallsThrough :: GenStmnts a -> Bool
fallsThrough ss 
    | isNilOL ss = True
    | otherwise  = ftStmnt (snd (unsnocOL ss))

ftStmnt :: GenStmnt a -> Bool
ftStmnt TailCall{} = False
ftStmnt Return{} = False
ftStmnt (Switch _ arms def) = any fallsThrough (def:[ss|(_,ss)<-arms])
ftStmnt _ = True
\end{code}