@@ -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,264 @@ 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 inl = evalState (eval [] initEnv [topn] defs exp)
39+ (length args)
40+ -- do some case floating, which might arise as a result
41+ -- then, eta contract
42+ res = eta $ caseFloats 10 inl in
43+ case res of
44+ LLam args' body -> LFun opts topn (map snd initNames ++ args') body
45+ _ -> LFun opts topn (map snd initNames) res
46+ where
47+ caseFloats 0 tm = tm
48+ caseFloats n tm
49+ = let res = caseFloat tm in
50+ if res == tm
51+ then res
52+ else caseFloats (n-1) res
53+
54+ initNames = zipWith (\n i -> (n, newn n i)) args [0..]
55+ initEnv = map (\(n, n') -> (n, LV n')) initNames
56+ newn (UN n) i = MN i n
57+ newn _ i = sMN i "arg"
58+
59+unload :: [LExp] -> LExp -> LExp
60+unload [] e = e
61+unload stk (LApp tc e args) = LApp tc e (args ++ stk)
62+unload stk e = LApp False e stk
63+
64+takeStk :: [(Name, LExp)] -> [Name] -> [LExp] ->
65+ ([(Name, LExp)], [Name], [LExp])
66+takeStk env (a : args) (v : stk) = takeStk ((a, v) : env) args stk
67+takeStk env args stk = (env, args, stk)
68+
69+eval :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp -> State Int LExp
70+eval stk env rec defs (LLazyApp n es)
71+ = unload stk <$> LLazyApp n <$> (mapM (eval [] env rec defs) es)
72+eval stk env rec defs (LForce e)
73+ = do e' <- eval [] env rec defs e
74+ case e' of
75+ LLazyExp forced -> return $ unload stk forced
76+ LLazyApp n es -> return $ unload stk (LApp False (LV n) es)
77+ _ -> return (unload stk (LForce e'))
78+eval stk env rec defs (LLazyExp e)
79+ = unload stk <$> LLazyExp <$> eval [] env rec defs e
80+-- Special case for io_bind, because it needs to keep executing the first
81+-- action, and is worth inlining to avoid the thunk
82+eval [] env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
83+ | n == sUN "io_bind"
84+ = do w <- nextN
85+ let env' = (w, LV w) : env
86+ act' <- eval [] env' rec defs (LApp False act [LV w])
87+ argn <- nextN
88+ k' <- eval [] ((arg, LV argn) : env') rec defs (LApp False k [LV w])
89+ return $ LLam [w] (LLet argn act' k')
90+eval (world : stk) env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
91+ | n == sUN "io_bind"
92+ = do act' <- eval [] env rec defs (LApp False act [world])
93+ argn <- nextN
94+ k' <- eval stk ((arg, LV argn) : env) rec defs (LApp False k [world])
95+ return $ LLet argn act' k'
96+
97+eval stk env rec defs (LApp t f es)
98+ = do es' <- mapM (eval [] env rec defs) es
99+ eval (es' ++ stk) env rec defs f
100+eval stk env rec defs (LLet n val sc)
101+ = do n' <- nextN
102+ LLet n' <$> eval [] env rec defs val
103+ <*> eval stk ((n, LV n') : env) rec defs sc
104+eval stk env rec defs (LProj exp i)
105+ = unload stk <$> (LProj <$> eval [] env rec defs exp <*> return i)
106+eval stk env rec defs (LCon loc i n es)
107+ = unload stk <$> (LCon loc i n <$> mapM (eval [] env rec defs) es)
108+eval stk env rec defs (LCase ty e alts)
109+ = do alts' <- mapM (evalAlt stk env rec defs) alts
110+ e' <- eval [] env rec defs e
111+ -- If they're all lambdas, bind the lambda at the top
112+ let prefix = getLams (map getRHS alts')
113+ case prefix of
114+ [] -> return $ conOpt $ LCase ty e' (replaceInAlts e' alts')
115+ args -> do alts_red <- mapM (dropArgs args) alts'
116+ return $ LLam args
117+ (conOpt (LCase ty e' (replaceInAlts e' alts_red)))
118+eval stk env rec defs (LOp f es)
119+ = unload stk <$> LOp f <$> mapM (eval [] env rec defs) es
120+eval stk env rec defs (LForeign t s args)
121+ = unload stk <$>
122+ LForeign t s <$> mapM (\(t, e) -> do e' <- eval [] env rec defs e
123+ return (t, e')) args
124+-- save the interesting cases for the end:
125+-- lambdas, and names to reduce
126+eval stk env rec defs (LLam args sc)
127+ | (env', args', stk') <- takeStk env args stk
128+ = case args' of
129+ [] -> eval stk' env' rec defs sc
130+ as -> do ns' <- mapM (\n -> do n' <- nextN
131+ return (n, n')) args'
132+ unload stk' <$> LLam (map snd ns') <$>
133+ eval [] (map (\ (n, n') -> (n, LV n')) ns' ++ env')
134+ rec defs sc
135+eval stk env rec defs var@(LV n)
136+ = case lookup n env of
137+ Just t
138+ | t /= LV n && n `notElem` rec ->
139+ eval stk env (n : rec) defs t
140+ | otherwise -> return (unload stk t)
141+ Nothing
142+ | n `notElem` rec,
143+ Just (LFun opts _ args body) <- lookupCtxtExact n defs,
144+ Inline `elem` opts ->
145+ apply stk env (n : rec) defs var args body
146+ | Just (LConstructor n t a) <- lookupCtxtExact n defs ->
147+ return (LCon Nothing t n stk)
148+ | otherwise -> return (unload stk var)
149+eval stk env rec defs t = return (unload stk t)
150+
151+evalAlt stk env rec defs (LConCase i n es rhs)
152+ = do ns' <- mapM (\n -> do n' <- nextN
153+ return (n, n')) es
154+ LConCase i n (map snd ns') <$>
155+ eval stk (map (\ (n, n') -> (n, LV n')) ns' ++ env) rec defs rhs
156+evalAlt stk env rec defs (LConstCase c e)
157+ = LConstCase c <$> eval stk env rec defs e
158+evalAlt stk env rec defs (LDefaultCase e)
159+ = LDefaultCase <$> eval stk env rec defs e
160+
161+apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->
162+ [Name] -> LExp -> State Int LExp
163+apply stk env rec defs var args body
164+ = eval stk env rec defs (LLam args body)
165+
166+dropArgs :: [Name] -> LAlt -> State Int LAlt
167+dropArgs as (LConCase i n es (LLam args rhs))
168+ = do let old = take (length as) args
169+ rhs' <- eval [] (zipWith (\ o n -> (o, LV n)) old as) [] emptyContext rhs
170+ return (LConCase i n es rhs')
171+dropArgs as (LConstCase c (LLam args rhs))
172+ = do let old = take (length as) args
173+ rhs' <- eval [] (zipWith (\ o n -> (o, LV n)) old as) [] emptyContext rhs
174+ return (LConstCase c rhs')
175+dropArgs as (LDefaultCase (LLam args rhs))
176+ = do let old = take (length as) args
177+ rhs' <- eval [] (zipWith (\ o n -> (o, LV n)) old as) [] emptyContext rhs
178+ return (LDefaultCase rhs')
179+
180+caseFloat :: LExp -> LExp
181+caseFloat (LApp tc e es) = LApp tc (caseFloat e) (map caseFloat es)
182+caseFloat (LLazyExp e) = LLazyExp (caseFloat e)
183+caseFloat (LForce e) = LForce (caseFloat e)
184+caseFloat (LCon up i n es) = LCon up i n (map caseFloat es)
185+caseFloat (LOp f es) = LOp f (map caseFloat es)
186+caseFloat (LLam ns sc) = LLam ns (caseFloat sc)
187+caseFloat (LLet v val sc) = LLet v (caseFloat val) (caseFloat sc)
188+caseFloat (LCase _ (LCase ct exp alts) alts')
189+ | all conRHS alts || length alts == 1
190+ = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map (updateWith alts') alts))
191+ where
192+ conRHS (LConCase _ _ _ (LCon _ _ _ _)) = True
193+ conRHS (LConstCase _ (LCon _ _ _ _)) = True
194+ conRHS (LDefaultCase (LCon _ _ _ _)) = True
195+ conRHS _ = False
196+
197+ updateWith alts (LConCase i n es rhs) =
198+ LConCase i n es (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
199+ updateWith alts (LConstCase c rhs) =
200+ LConstCase c (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
201+ updateWith alts (LDefaultCase rhs) =
202+ LDefaultCase (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
203+
204+caseFloat (LCase ct exp alts')
205+ = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map cfAlt alts'))
206+ where
207+ cfAlt (LConCase i n es rhs) = LConCase i n es (caseFloat rhs)
208+ cfAlt (LConstCase c rhs) = LConstCase c (caseFloat rhs)
209+ cfAlt (LDefaultCase rhs) = LDefaultCase (caseFloat rhs)
210+caseFloat exp = exp
211+
212+-- Case of constructor
213+conOpt :: LExp -> LExp
214+conOpt (LCase ct (LCon _ t n args) alts)
215+ = pickAlt n args alts
216+ where
217+ pickAlt n args (LConCase i n' es rhs : as) | n == n'
218+ = substAll (zip es args) rhs
219+ pickAlt _ _ (LDefaultCase rhs : as) = rhs
220+ pickAlt n args (_ : as) = pickAlt n args as
221+ pickAlt n args [] = error "Can't happen pickAlt - impossible case found"
222+
223+ substAll [] rhs = rhs
224+ substAll ((n, tm) : ss) rhs = lsubst n tm (substAll ss rhs)
225+conOpt tm = tm
226+
227+replaceInCase :: LExp -> LExp
228+replaceInCase (LCase ty e alts)
229+ = LCase ty e (replaceInAlts e alts)
230+replaceInCase exp = exp
231+
232+replaceInAlts :: LExp -> [LAlt] -> [LAlt]
233+replaceInAlts exp alts = dropDups $ concatMap (replaceInAlt exp) alts
234+
235+-- Drop overlapping case (arising from case merging of overlapping
236+-- patterns)
237+dropDups (alt@(LConCase _ i n ns) : alts)
238+ = alt : dropDups (filter (notTag i) alts)
36239 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
76−
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
240+ notTag i (LConCase _ j n ns) = i /= j
241+ notTag _ _ = True
242+dropDups (c : alts) = c : dropDups alts
243+dropDups [] = []
244+
245+
246+replaceInAlt :: LExp -> LAlt -> [LAlt]
247+-- In an alternative, if the case appears on the right hand side, replace
248+-- it with the given expression, to preserve sharing
249+replaceInAlt exp@(LV _) (LConCase i con args rhs)
250+ = [LConCase i con args $
251+ replaceExp (LCon Nothing i con (map LV args)) exp rhs]
252+-- if a default case inspects the same variable as the case it's in,
253+-- remove the inspection and replace with the alternatives
254+-- (i.e. merge the inner case block)
255+replaceInAlt exp@(LV var) (LDefaultCase (LCase ty (LV var') alts))
256+ | var == var' = alts
257+replaceInAlt exp a = [a]
258+
259+replaceExp :: LExp -> LExp -> LExp -> LExp
260+replaceExp (LCon _ t n args) new (LCon _ t' n' args')
261+ | n == n' && args == args' = new
262+replaceExp (LCon _ t n args) new (LApp _ (LV n') args')
263+ | n == n' && args == args' = new
264+replaceExp old new tm = tm
265+
266+-- dropArgs as (LConstCase c rhs) = LConstCase c (dropRHS as rhs)
267+-- dropArgs as (LDefaultCase rhs) = LDefaultCase (dropRHS as rhs)
268+
269+getRHS (LConCase i n es rhs) = rhs
270+getRHS (LConstCase _ rhs) = rhs
271+getRHS (LDefaultCase rhs) = rhs
272+
273+getLams [] = []
274+getLams (LLam args tm : cs) = getLamPrefix args cs
275+getLams _ = []
276+
277+getLamPrefix as [] = as
278+getLamPrefix as (LLam args tm : cs)
279+ | length args < length as = getLamPrefix args cs
280+ | otherwise = getLamPrefix as cs
281+getLamPrefix as (_ : cs) = []
282+
283+-- eta contract ('\x -> f x' can just be compiled as 'f' when f is local)
284+eta :: LExp -> LExp
285+eta (LApp tc a es) = LApp tc (eta a) (map eta es)
286+eta (LLazyApp n es) = LLazyApp n (map eta es)
287+eta (LLazyExp e) = LLazyExp (eta e)
288+eta (LForce e) = LForce (eta e)
289+eta (LLet n val sc) = LLet n (eta val) (eta sc)
290+eta (LLam args (LApp tc f args'))
291+ | args' == map LV args = eta f
292+eta (LLam args e) = LLam args (eta e)
293+eta (LProj e i) = LProj (eta e) i
294+eta (LCon a t n es) = LCon a t n (map eta es)
295+eta (LCase ct e alts) = LCase ct (eta e) (map (fmap eta) alts)
296+eta (LOp f es) = LOp f (map eta es)
297+eta tm = tm
298+