module Yhc.Core.Inline(
coreInline, InlineMode(..),
coreInlineFunc, coreInlineFuncLambda
) where
import Yhc.Core.Type
import Yhc.Core.Uniplate
import Yhc.Core.FreeVar
import qualified Data.Map as Map
import Data.Maybe
import Data.List
data InlineMode = InlineNone
| InlineAlias
| InlineForward
| InlineCallOnce
| InlineFull
coreInline :: InlineMode -> Core -> Core
coreInline InlineNone core = core
coreInline InlineAlias core = coreInlineAlias core
coreInline InlineForward core = inlineNormal (analyseForward core) core
coreInline InlineCallOnce core = inlineNormal (analyseForward core `Map.union` analyseCallOnce core) core
coreInline InlineFull core = inlineNormal (analyseFull core) $ coreInlineAlias core
coreInlineAlias core = inlineAlias (analyseAlias core) core
inlineAlias :: Map.Map CoreFuncName CoreFuncName -> Core -> Core
inlineAlias rep core = transformExpr f core
where
f (CoreFun x) = CoreFun $ Map.findWithDefault x x rep
f x = x
inlineNormal :: Map.Map CoreFuncName CoreFunc -> Core -> Core
inlineNormal rep core = applyFuncCore f core
where
f (CoreFunc name args body) = CoreFunc name args $ transformExpr (g [name]) body
f x = x
g done x = fromMaybe x $ do
(CoreFun fn,args) <- return $ fromCoreApp x
func <- Map.lookup fn rep
True <- return $ fn `notElem` done
res <- coreInlineFunc func args
return $ transformExpr (g (fn:done)) res
analyseAlias :: Core -> Map.Map CoreFuncName CoreFuncName
analyseAlias core = transForward
where
basicForward :: Map.Map String String
basicForward = Map.fromList $ concatMap f (coreFuncs core)
where
f (CoreFunc name args (CoreApp (CoreFun x) xs)) | map CoreVar args == xs = [(name,x)]
f _ = []
transForward :: Map.Map String String
transForward = Map.mapWithKey (\k v -> f [k] v) basicForward
where
f done name =
case Map.lookup name basicForward of
Just x | name `notElem` done -> f (name:done) x
_ -> name
analyseForward :: Core -> Map.Map CoreFuncName CoreFunc
analyseForward core = Map.fromList
[(name, func) | func@(CoreFunc name _ bod) <- coreFuncs core, canInline bod]
where
canInline (CorePos _ x) = canInline x
canInline (CoreApp x xs) = isGoodFun x && all isGoodArg xs && disjoint [i | CoreVar i <- xs]
canInline x = isCoreCon x || isCoreFun x || isGoodArg x
isGoodFun x = isCoreFun x || isCoreCon x
isGoodArg x = isCoreVar x || isSmallConst x
isSmallConst x = isCoreLit x && not (isCoreStr $ fromCoreLit x)
analyseCallOnce :: Core -> Map.Map CoreFuncName CoreFunc
analyseCallOnce core = Map.fromList
[(name,func) | func@(CoreFunc name (_:_) _) <- coreFuncs core, Just True == Map.lookup name once]
where
once :: Map.Map CoreFuncName Bool
once = foldl f Map.empty [x | CoreFun x <- universeExpr core]
f mp x = Map.insertWith (\_ _ -> False) x True mp
analyseFull :: Core -> Map.Map CoreFuncName CoreFunc
analyseFull core = Map.fromList [(name,func) | func@(CoreFunc name (_:_) _) <- coreFuncs core]
disjoint x = length (nub x) == length x
coreInlineFunc :: CoreFunc -> [CoreExpr] -> Maybe CoreExpr
coreInlineFunc func@(CoreFunc name params2 body2) args
| nparams > nargs = Nothing
| otherwise = Just res
where
res = coreApp subst (drop nparams args)
(nargs, nparams) = (length args, length params2)
argvars = concatMap collectAllVars args
allvars = ['v':show i | i <- [1..]] \\ (params2 ++ argvars ++ collectAllVars body2)
(params,rest) = splitAt nparams allvars
body = uniqueBoundVarsWith rest $ replaceFreeVars (zip params2 (map CoreVar params)) body2
newvars = rest \\ collectAllVars body
(dupe,once) = partition (\(lhs,rhs) -> requiresLet rhs && countFreeVar lhs body > 1) (zip params args)
requiresLet x = not (isCoreVar x || isCoreFun x)
dupnew = zip newvars dupe
binds = [(new,a) | (new,(p,a)) <- dupnew]
reps = [(p,CoreVar new) | (new,(p,a)) <- dupnew] ++ once
subst = coreLet binds (replaceFreeVars reps body)
coreInlineFuncLambda :: CoreFunc -> [CoreExpr] -> ([String], CoreExpr)
coreInlineFuncLambda func@(CoreFunc name params body) args =
(extraArgs, fromJust $ coreInlineFunc func (args ++ map CoreVar extraArgs))
where
extraArgs = drop (length args) (coreFuncArgs func)