-- |Handles the inlining of single function applications and inlining
-- residuals through a whole program
module Optimus.Inline (
			maybeInline,
			progInline
		) where
	import Control.Monad
	import Data.List
	import qualified Data.Map as Map
	import qualified Data.Set as Set
	import Data.Maybe

	import Flite.Fresh
	import Flite.Syntax
	import Flite.Traversals
	
	import Optimus.Uniplate
	import Optimus.Util
	
	-- |Build an inlined form of a saturated function application given
	-- the correct function definition. Assumes that equations have
	-- been desugared. Simplify afterwards for efficiency.
	maybeInline :: [Exp] -> Decl -> Maybe (Fresh Exp)
	maybeInline ys (Func _ args rhs)
		| length_args > length_ys = fail "Unsaturated application."
		| otherwise = Just $ do
				vs <- mapM (const fresh) args
				let (bound, remaining) = splitAt (length args) ys
				let rhs' = substMany rhs (zip (map Var vs) (concatMap patVars args))
				return $ mkApp (mkLet (zip vs bound) rhs') remaining
		where
			length_ys = length ys
			length_args = length args
			
	-- |Make an Application only if there are arguments.
	mkApp :: Exp -> [Exp] -> Exp
	mkApp x [] = x
	mkApp x ys = App x ys
	
	-- |Make a Let only if there are bindings.
	mkLet :: [Binding] -> Exp -> Exp
	mkLet [] y = y
	mkLet bs y = Let bs y
			
	-- |Inline all non-recursive function applications.
	progInline :: Prog -> Fresh Prog
	progInline p = progInline' (byFuncName p) p
		where
			callCounts = map (\xs@(x:_) -> (x, length xs)) . group . sort . concatMap (calls . funcRhs) $ p
			highlyCalled = Set.fromList [ f | (f, n) <- callCounts, n > 10 ]
			progInline' :: Map.Map Id Decl -> Prog -> Fresh Prog
			progInline' m [] = return $ Map.elems m
			progInline' m (d:ds) = do
				d'@(Func f _ _) <- declInline (`Map.lookup` m) highlyCalled d
				progInline' (Map.insert f d' m) ds
	
	-- |Given a mapping of function names to declarations, inline
	-- all non-recursive function applications
	declInline :: (Id -> Maybe Decl) -> Set.Set Id -> Decl -> Fresh Decl
	declInline progmap base (Func f a r) = liftM (Func f a) (expInline (f `Set.insert` base) r)
		where
			-- |Given a trace of functions visited, inline all non-
			-- recursive function applications.
			expInline :: Set.Set Id -> Exp -> Fresh Exp
			expInline fs (Fun f) | f `Set.notMember` fs && isJust inlined
				= fromJust inlined >>= descendM (expInline $ f `Set.insert` fs)
				where
					inlined = progmap f >>= notRecursive >>= maybeInline []
			expInline fs (App (Fun f) xs) | f `Set.notMember` fs && isJust inlined
				= fromJust inlined >>= descendM (expInline $ f `Set.insert` fs)
				where
					inlined = progmap f >>= notRecursive >>= maybeInline xs
			expInline fs e = descendM (expInline fs) e
			
	-- |Only return functions that are not immediately recursive.
	notRecursive :: Decl -> Maybe Decl
	notRecursive d@(Func f _ r) | f `elem` calls r 	= Nothing
								| otherwise			= Just d