@@ -7,12 +7,16 @@ Maintainer : The Idris Community.
77-}
88{-# LANGUAGE DeriveFunctor, PatternGuards #-}
99
10−module IRTS.LangOpts where
10+module IRTS.LangOpts(inlineAll) where
1111
12+import Idris.Core.CaseTree
1213import Idris.Core.TT
1314import IRTS.Lang
1415
1516import Control.Monad.State hiding (lift)
17+import Data.List
18+
19+import Debug.Trace
1620
1721inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
1822inlineAll lds = let defs = addAlist lds emptyContext in
@@ -31,53 +35,108 @@ nextN = do i <- get
3135doInline :: LDefs -> LDecl -> LDecl
3236doInline defs d@(LConstructor _ _ _) = d
3337doInline defs (LFun opts topn args exp)
34− = let res = evalState (inlineWith [topn] (map (\n -> (n, LV n)) args) exp) 0 in
35− LFun opts topn args res
38+ = let res = evalState (eval [] initEnv [topn] defs exp)
39+ (length args) in
40+-- = let res = evalState (inlineWith [topn] (map (\n -> (n, LV n)) args) exp) 0 in
41+ case res of
42+ LLam args' body -> LFun opts topn (map snd initNames ++ args') body
43+ _ -> LFun opts topn (map snd initNames) res
3644 where
37− inlineWith :: [Name] -> [(Name, LExp)] -> LExp -> State Int LExp
38− inlineWith done env var@(LV n)
39− = case lookup n env of
40− Just t -> return t
41− Nothing -> return var
42− inlineWith done env (LLazyApp n es) = LLazyApp n <$> (mapM (inlineWith done env) es)
43− inlineWith done env (LForce e) = LForce <$> inlineWith done env e
44− inlineWith done env (LLazyExp e) = LLazyExp <$> inlineWith done env e
45− -- Extend the environment for Let and Lam so that bound names aren't
46− -- expanded with any top level argument definitions they shadow
47− inlineWith done env (LLet n val sc)
48− = do n' <- nextN
49− LLet n' <$> inlineWith done env val <*>
50− inlineWith done ((n, LV n') : env) sc
51− inlineWith done env (LLam args sc)
52− = do ns' <- mapM (\n -> do n' <- nextN
53− return (n, n')) args
54− LLam (map snd ns') <$>
55− inlineWith done (map (\ (n,n') -> (n, LV n')) ns' ++ env) sc
56− inlineWith done env (LProj exp i) = LProj <$> inlineWith done env exp <*> return i
57− inlineWith done env (LCon loc i n es)
58− = LCon loc i n <$> mapM (inlineWith done env) es
59− inlineWith done env (LCase ty e alts)
60− = LCase ty <$> inlineWith done env e <*> mapM (inlineWithAlt done env) alts
61− inlineWith done env (LOp f es) = LOp f <$> mapM (inlineWith done env) es
62− -- the interesting case!
63− inlineWith done env (LApp t (LV n) es)
64− | n `notElem` done,
65− [LFun opts _ args body] <- lookupCtxt n defs,
66− Inline `elem` opts,
67− length es == length args
68− = do es' <- mapM (inlineWith done env) es
69− inlineWith (n : done) (zip args es' ++ env) body
70− inlineWith done env (LApp t f es)
71− = LApp t <$> inlineWith done env f <*> mapM (inlineWith done env) es
72− inlineWith done env (LForeign t s args)
73− = LForeign t s <$> mapM (\(t, e) -> do e' <- inlineWith done env e
74− return (t, e')) args
75− inlineWith done env t = return t
45+ initNames = zipWith (\n i -> (n, newn n i)) args [0..]
46+ initEnv = map (\(n, n') -> (n, LV n')) initNames
47+ newn (UN n) i = MN i n
48+ newn _ i = sMN i "arg"
49+
50+unload :: [LExp] -> LExp -> LExp
51+unload [] e = e
52+unload stk (LApp tc e args) = LApp tc e (args ++ stk)
53+unload stk e = LApp False e stk
54+
55+takeStk :: [(Name, LExp)] -> [Name] -> [LExp] ->
56+ ([(Name, LExp)], [Name], [LExp])
57+takeStk env (a : args) (v : stk) = takeStk ((a, v) : env) args stk
58+takeStk env args stk = (env, args, stk)
59+
60+eval :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp -> State Int LExp
61+eval stk env rec defs (LLazyApp n es)
62+ = unload stk <$> LLazyApp n <$> (mapM (eval [] env rec defs) es)
63+eval stk env rec defs (LForce e)
64+ = unload stk <$> LForce <$> eval [] env rec defs e
65+eval stk env rec defs (LLazyExp e)
66+ = unload stk <$> LLazyExp <$> eval [] env rec defs e
67+-- Special case for io_bind, because it needs to keep executing the first
68+-- action, and is worth inlining to avoid the thunk
69+eval [] env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
70+ | n == sUN "io_bind"
71+ = do w <- nextN
72+ let env' = (w, LV w) : env
73+ act' <- eval [] env' rec defs (LApp False act [LV w])
74+ argn <- nextN
75+ k' <- eval [] ((arg, LV argn) : env') rec defs (LApp False k [LV w])
76+ return $ LLam [w] (LLet argn act' k')
77+eval (world : stk) env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
78+ | n == sUN "io_bind"
79+ = do act' <- eval [] env rec defs (LApp False act [world])
80+ argn <- nextN
81+ k' <- eval stk ((arg, LV argn) : env) rec defs (LApp False k [world])
82+ return $ LLet argn act' k'
83+
84+eval stk env rec defs (LApp t f es)
85+ = do es' <- mapM (eval [] env rec defs) es
86+ eval (es' ++ stk) env rec defs f
87+eval stk env rec defs (LLet n val sc)
88+ = do n' <- nextN
89+ LLet n' <$> eval [] env rec defs val
90+ <*> eval stk ((n, LV n') : env) rec defs sc
91+eval stk env rec defs (LProj exp i)
92+ = unload stk <$> (LProj <$> eval [] env rec defs exp <*> return i)
93+eval stk env rec defs (LCon loc i n es)
94+ = unload stk <$> (LCon loc i n <$> mapM (eval [] env rec defs) es)
95+eval stk env rec defs (LCase ty e alts)
96+ = LCase ty <$> eval [] env rec defs e
97+ <*> mapM (evalAlt stk env rec defs) alts
98+eval stk env rec defs (LOp f es)
99+ = unload stk <$> LOp f <$> mapM (eval [] env rec defs) es
100+eval stk env rec defs (LForeign t s args)
101+ = unload stk <$>
102+ LForeign t s <$> mapM (\(t, e) -> do e' <- eval [] env rec defs e
103+ return (t, e')) args
104+-- save the interesting cases for the end:
105+-- lambdas, and names to reduce
106+eval stk env rec defs (LLam args sc)
107+ | (env', args', stk') <- takeStk env args stk
108+ = case args' of
109+ [] -> eval stk' env' rec defs sc
110+ as -> do ns' <- mapM (\n -> do n' <- nextN
111+ return (n, n')) args'
112+ LLam (map snd ns') <$>
113+ eval [] (map (\ (n, n') -> (n, LV n')) ns' ++ env')
114+ rec defs sc
115+eval stk env rec defs var@(LV n)
116+ = case lookup n env of
117+ Just t
118+ | t /= LV n && n `notElem` rec ->
119+ eval stk env (n : rec) defs t
120+ | otherwise -> return (unload stk t)
121+ Nothing
122+ | n `notElem` rec,
123+ Just (LFun opts _ args body) <- lookupCtxtExact n defs,
124+ Inline `elem` opts ->
125+ apply stk env (n : rec) defs var args body
126+ | otherwise -> return (unload stk var)
127+eval stk env rec defs t = return (unload stk t)
128+
129+evalAlt stk env rec defs (LConCase i n es rhs)
130+ = do ns' <- mapM (\n -> do n' <- nextN
131+ return (n, n')) es
132+ LConCase i n (map snd ns') <$>
133+ eval stk (map (\ (n, n') -> (n, LV n')) ns' ++ env) rec defs rhs
134+evalAlt stk env rec defs (LConstCase c e)
135+ = LConstCase c <$> eval stk env rec defs e
136+evalAlt stk env rec defs (LDefaultCase e)
137+ = LDefaultCase <$> eval stk env rec defs e
76138
77− inlineWithAlt done env (LConCase i n es rhs)
78− = do ns' <- mapM (\n -> do n' <- nextN
79− return (n, n')) es
80− LConCase i n (map snd ns') <$>
81− inlineWith done (map (\ (n,n') -> (n, LV n')) ns' ++ env) rhs
82− inlineWithAlt done env (LConstCase c e) = LConstCase c <$> inlineWith done env e
83− inlineWithAlt done env (LDefaultCase e) = LDefaultCase <$> inlineWith done env e
139+apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->
140+ [Name] -> LExp -> State Int LExp
141+apply stk env rec defs var args body
142+ = eval stk env rec defs (LLam args body)