-- GRIN-like backend for Yhc Core.
-- Inline function calls and
-- flatten inline blocks where possible.

module Yhc.Core.GRIN.FlatInline where

import Yhc.Core.GRIN.Type
import Yhc.Core.GRIN.FindType
import Yhc.Core.GRIN.HeapPointsTo
import Yhc.Core.GRIN.SubstVars
import Yhc.Core.GRIN.ElimDead
import Yhc.Core.GRIN.OptCases

import Data.List
import qualified Data.Set as S
import qualified Data.Map as M


-- Flatten inline blocks. The logic is:
--  
--  a <- do ...
--       b <- foo
--       return b
--  
--  is same as
--  ...
--  a <- foo
--
--  b is used only within the block being flattened.

flattenInline :: GRIN -> GRIN

flattenInline gr = gr {gFuncs = map f (gFuncs gr)} where
  f gf = gf {gFuncBody = filter (/= GNop) (ff (gFuncBody gf))}
  ff [] = []
  ff (GBind (GInline [GSimple (GUnit (GVal (GVar newvar)))]) bv : bs) =
    let oldvar = unGVar bv
        submap = M.singleton oldvar newvar
    in  ff (map (substGX submap) bs)
  ff (GBind (GInline ib) bv : bs) = flatten (ff ib) bv ++ ff bs
  ff (GCase cv cbrs : bs) = 
    let pats = map fst cbrs
        blks = map snd cbrs
        blks' = map ff blks
        b' = GCase cv (zip pats blks')
    in  b' : ff bs
  ff (b:bs) = b : ff bs

flatten ib bv =
  case reverse ib of
    GSimple (GUnit iret) : GBind irsx iret' : ibrest | iret' == iret ->
      reverse ibrest ++ [GBind irsx bv]
    GSimple (GUnit t@(GTagged _ _)) : ibrest ->
      reverse ibrest ++ [GBind (GStore t) bv]
    _ -> [GBind (GInline ib) bv]


-- Determine if a function is inlineable and get the list of its
-- variables to substitute. Function should be neither a CAF nor
-- recursive. Recursive means explicitly calling itself; P- or F-tag
-- for itself is OK. If the function does not qualify, return an empty
-- list. Non-linear functions (containing case statements) are also excluded
-- because their blocks cannot be flattened (NB may be able to inline
-- tail calls)


isInlineable :: GFunc -> S.Set GName

isInlineable gf | null (gFuncArgs gf) = S.empty  -- CAF

isInlineable gf = S.filter isNode $ 
                  notrec $ 
                  useVarBlk (S.fromList $ gFuncArgs gf) (gFuncBody gf) where
  notrec l = if gFuncName gf `S.notMember` l then l else S.empty


-- Determine if a function is linear, that is, it does not contain
-- cases, inline blocks, and function calls (CAFs only). 
-- This would in most cases apply to dictionary building functions
-- and lift functions.

isLinear :: GFunc -> S.Set GName

isLinear (GFunc {gFuncBody = [GBind (GCall "GRIN;eval" [GVal (GVar ev)]) v1, GSimple (GUnit v2)]
                ,gFuncArgs = [v0]}) | ev == v0 && v1 == v2 = S.fromList [v0, unGVar v1, unGVar v2]

isLinear gf = 
  let bl = gFuncBody gf
      storec (GBind (GStore (GTagged ('C':'@':_) _)) _) = True
      storec _ = False
      nstorec = length (filter storec bl)
      narg = length $ gFuncArgs gf
      pred (GCase _ _) = True
      pred (GBind (GInline ib) _) = any pred ib
      pred (GBind (GCall _ _) _) = narg == 0
      pred _ = narg == 0 && nstorec > 1 -- this inlines only CAF dictionaries
      notrec l = if gFuncName gf `S.notMember` l then l else S.empty
  in  if any pred bl then S.empty
      else S.filter isNode $ notrec $ useVarBlk (S.fromList $ gFuncArgs gf) bl

-- Inline qualifying functions. An explicit function call is replaced
-- with an inline block bound to the variable which was initially bound
-- to the call. All variables within the block will be renamed by prepending
-- the bound-to variable name, except for the function parameters
-- which will be substituted from the call. To be conservative, unaltered
-- versions of functions will be inlined (no infinite inlining).


inlineCalls :: (GFunc -> S.Set GName) -> GRIN -> GRIN

inlineCalls iqf g =
  let oldfuncs = M.fromList $ zip (map gFuncName (gFuncs g)) (gFuncs g)
      inl g@(GFunc {gFuncName = n}) | n `elem` ["GRIN;eval", "GRIN;apply"] = g
      inl gf = gf {gFuncBody = ff (gFuncBody gf)} where
        ff [] = []
        ff (b:bs) = fff b : ff bs
        fff (GCase cv cbrs) =
          let pats = map fst cbrs
              newbrs = map (ff . snd) cbrs
          in  GCase cv (zip pats newbrs)
        fff (GBind (GInline bl) gv) = GBind (GInline (ff bl)) gv
        fff b@(GBind (GCall fn ps) gv) = case M.lookup fn oldfuncs of
          Nothing -> b
          Just fdf -> 
            let q = S.toList (iqf fdf)
                args = gFuncArgs fdf
                (vargs, nodes) = partition (`elem` args) q
                subarg = zip vargs (map unGVar ps)
                subnod = zip nodes (map (inname gv) nodes)
                submap = M.fromList (subarg ++ subnod)
                inname gv n = if isShrNode n 
                  then 's' : (unGVar gv ++ "@" ++ tail n)
                  else if isNode n then unGVar gv ++ "@" ++ n
                                   else n
                inbody = map (substGX submap) (gFuncBody fdf)
            in  case null q of
                  True -> b
                  False -> GBind (GInline inbody) gv
        fff z = z
  in  g {gFuncs = map inl (gFuncs g)}