git.iko.soy/Idris-dev
on master help sign in
Idris-dev/src/IRTS/LangOpts.hs
327 lines 13 kB history raw
1{-|
2Module : IRTS.LangOpts
3Description : Transformations to apply to Idris' IR.
5License : BSD3
6Maintainer : The Idris Community.
7-}
8{-# LANGUAGE DeriveFunctor, PatternGuards #-}
10module IRTS.LangOpts(inlineAll) where
12import Idris.Core.CaseTree
13import Idris.Core.TT
14import IRTS.Lang
16import Control.Monad.State hiding (lift)
17import Data.List
19import Debug.Trace
21inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
22inlineAll lds = let defs = addAlist lds emptyContext in
23 map (\ (n, def) -> (n, doInline defs def)) lds
25nextN :: State Int Name
26nextN = do i <- get
27 put (i + 1)
28 return $ sMN i "in"
30-- | Inline inside a declaration.
31--
32-- Variables are still Name at this stage. Need to preserve
33-- uniqueness of variable names in the resulting definition, so invent
34-- a new name for every variable we encounter
35doInline :: LDefs -> LDecl -> LDecl
36doInline = doInline' 1
38doInline' :: Int -> LDefs -> LDecl -> LDecl
39doInline' 0 defs d = d
40doInline' i defs d@(LConstructor _ _ _) = d
41doInline' i defs (LFun opts topn args exp)
42 = let inl = evalState (eval [] initEnv [topn] defs exp)
43 (length args)
44 -- do some case floating, which might arise as a result
45 -- then, eta contract
46 res = eta $ caseFloats 10 inl in
47 case res of
48 LLam args' body ->
49 doInline' (i - 1) defs $
50 LFun opts topn (map snd initNames ++ args') body
51 _ -> doInline' (i - 1) defs $
52 LFun opts topn (map snd initNames) res
53 where
54 caseFloats 0 tm = tm
55 caseFloats n tm
56 = let res = caseFloat tm in
57 if res == tm
58 then res
59 else caseFloats (n-1) res
61 initNames = zipWith (\n i -> (n, newn n i)) args [0..]
62 initEnv = map (\(n, n') -> (n, LV n')) initNames
63 newn (UN n) i = MN i n
64 newn _ i = sMN i "arg"
66unload :: [LExp] -> LExp -> LExp
67unload [] e = e
68unload stk (LApp tc e args) = LApp tc e (args ++ stk)
69unload stk e = LApp False e stk
71takeStk :: [(Name, LExp)] -> [Name] -> [LExp] ->
72 ([(Name, LExp)], [Name], [LExp])
73takeStk env (a : args) (v : stk) = takeStk ((a, v) : env) args stk
74takeStk env args stk = (env, args, stk)
76eval :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp -> State Int LExp
77eval stk env rec defs (LLazyApp n es)
78 = unload stk <$> LLazyApp n <$> (mapM (eval [] env rec defs) es)
79eval stk env rec defs (LForce e)
80 = do e' <- eval [] env rec defs e
81 case e' of
82 LLazyExp forced -> eval stk env rec defs forced
83 LLazyApp n es -> eval stk env rec defs (LApp False (LV n) es)
84 _ -> return (unload stk (LForce e'))
85eval stk env rec defs (LLazyExp e)
86 = unload stk <$> LLazyExp <$> eval [] env rec defs e
87-- Special case for io_bind, because it needs to keep executing the first
88-- action, and is worth inlining to avoid the thunk
89eval [] env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
90 | n == sUN "io_bind"
91 = do w <- nextN
92 let env' = (w, LV w) : env
93 act' <- eval [] env' rec defs (LApp False act [LV w])
94 argn <- nextN
95 k' <- eval [] ((arg, LV argn) : env') rec defs (LApp False k [LV w])
96 return $ LLam [w] (LLet argn act' k')
97eval (world : stk) env rec defs (LApp t (LV n) [_, _, _, act, (LLam [arg] k)])
98 | n == sUN "io_bind"
99 = do act' <- eval [] env rec defs (LApp False act [world])
100 argn <- nextN
101 k' <- eval stk ((arg, LV argn) : env) rec defs (LApp False k [world])
102 -- Needs to be a LLet to make sure the action gets evaluated
103 return $ LLet argn act' k'
105eval stk env rec defs (LApp t f es)
106 = do es' <- mapM (eval [] env rec defs) es
107 eval (es' ++ stk) env rec defs f
108eval stk env rec defs (LLet n val sc)
109 = do n' <- nextN
110 LLet n' <$> eval [] env rec defs val
111 <*> eval stk ((n, LV n') : env) rec defs sc
112eval stk env rec defs (LProj exp i)
113 = unload stk <$> (LProj <$> eval [] env rec defs exp <*> return i)
114eval stk env rec defs (LCon loc i n es)
115 = unload stk <$> (LCon loc i n <$> mapM (eval [] env rec defs) es)
116eval stk env rec defs (LCase ty e [])
117 = pure LNothing
118eval stk env rec defs (LCase ty e alts)
119 = do e' <- eval [] env rec defs e
120 case evalAlts e' alts of
121 Just (env', tm) -> eval stk env' rec defs tm
122 Nothing ->
123 do alts' <- mapM (evalAlt stk env rec defs) alts
124 -- If they're all lambdas, bind the lambda at the top
125 let prefix = getLams (map getRHS alts')
126 case prefix of
127 [] -> return $ LCase ty e' (replaceInAlts e' alts')
128 args -> do alts_red <- mapM (dropArgs args) alts'
129 return $ LLam args
130 (LCase ty e' (replaceInAlts e' alts_red))
131 where
132 evalAlts e' [] = Nothing
133 evalAlts (LCon _ t n args) (LConCase i n' es rhs : as)
134 | n == n' = Just (zip es args ++ env, rhs)
135 evalAlts (LConst c) (LConstCase c' rhs : as)
136 | c == c' = Just (env, rhs)
137 evalAlts (LCon _ _ _ _) (LDefaultCase rhs : as) = Just (env, rhs)
138 evalAlts (LConst _) (LDefaultCase rhs : as) = Just (env, rhs)
139 evalAlts tm (_ : as) = evalAlts tm as
140eval stk env rec defs (LOp f es)
141 = unload stk <$> LOp f <$> mapM (eval [] env rec defs) es
142eval stk env rec defs (LForeign t s args)
143 = unload stk <$>
144 LForeign t s <$> mapM (\(t, e) -> do e' <- eval [] env rec defs e
145 return (t, e')) args
146-- save the interesting cases for the end:
147-- lambdas, and names to reduce
148eval stk env rec defs (LLam args sc)
149 | (env', args', stk') <- takeStk env args stk
150 = case args' of
151 [] -> eval stk' env' rec defs sc
152 as -> do ns' <- mapM (\n -> do n' <- nextN
153 return (n, n')) args'
154 LLam (map snd ns') <$>
155 eval stk' (map (\ (n, n') -> (n, LV n')) ns' ++ env')
156 rec defs sc
157eval stk env rec defs var@(LV n)
158 = case lookup n env of
159 Just t
160 | t /= LV n && n `notElem` rec ->
161 eval stk env (n : rec) defs t
162 | otherwise -> return (unload stk t)
163 Nothing
164 | n `notElem` rec,
165 Just (LFun opts _ args body) <- lookupCtxtExact n defs,
166 Inline `elem` opts ->
167 apply stk env (n : rec) defs var args body
168 | Just (LConstructor n t a) <- lookupCtxtExact n defs ->
169 return (LCon Nothing t n stk)
170 | otherwise -> return (unload stk var)
171eval stk env rec defs t = return (unload stk t)
173evalAlt stk env rec defs (LConCase i n es rhs)
174 = do ns' <- mapM (\n -> do n' <- nextN
175 return (n, n')) es
176 LConCase i n (map snd ns') <$>
177 eval stk (map (\ (n, n') -> (n, LV n')) ns' ++ env) rec defs rhs
178evalAlt stk env rec defs (LConstCase c e)
179 = LConstCase c <$> eval stk env rec defs e
180evalAlt stk env rec defs (LDefaultCase e)
181 = LDefaultCase <$> eval stk env rec defs e
183apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->
184 [Name] -> LExp -> State Int LExp
185apply stk env rec defs var args body
186 = eval stk env rec defs (LLam args body)
188dropArgs :: [Name] -> LAlt -> State Int LAlt
189dropArgs as (LConCase i n es t)
190 = do rhs' <- dropArgsTm as t
191 return (LConCase i n es rhs')
192dropArgs as (LConstCase c t)
193 = do rhs' <- dropArgsTm as t
194 return (LConstCase c rhs')
195dropArgs as (LDefaultCase t)
196 = do rhs' <- dropArgsTm as t
197 return (LDefaultCase rhs')
199dropArgsTm as (LLam args rhs)
200 = do let old = take (length as) args
201 eval [] (zipWith (\ o n -> (o, LV n)) old as) [] emptyContext rhs
202dropArgsTm as (LLet n val rhs)
203 = do rhs' <- dropArgsTm as rhs
204 pure (LLet n val rhs')
205dropArgsTm as tm = return tm
207caseFloat :: LExp -> LExp
208caseFloat (LApp tc e es) = LApp tc (caseFloat e) (map caseFloat es)
209caseFloat (LLazyExp e) = LLazyExp (caseFloat e)
210caseFloat (LForce e) = LForce (caseFloat e)
211caseFloat (LCon up i n es) = LCon up i n (map caseFloat es)
212caseFloat (LOp f es) = LOp f (map caseFloat es)
213caseFloat (LLam ns sc) = LLam ns (caseFloat sc)
214caseFloat (LLet v val sc) = LLet v (caseFloat val) (caseFloat sc)
215caseFloat (LCase _ (LCase ct exp alts) alts')
216 | all conRHS alts || length alts == 1
217 = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map (updateWith alts') alts))
218 where
219 conRHS (LConCase _ _ _ (LCon _ _ _ _)) = True
220 conRHS (LConstCase _ (LCon _ _ _ _)) = True
221 conRHS (LDefaultCase (LCon _ _ _ _)) = True
222 conRHS _ = False
224 updateWith alts (LConCase i n es rhs) =
225 LConCase i n es (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
226 updateWith alts (LConstCase c rhs) =
227 LConstCase c (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
228 updateWith alts (LDefaultCase rhs) =
229 LDefaultCase (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))
231caseFloat (LCase ct exp alts')
232 = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map cfAlt alts'))
233 where
234 cfAlt (LConCase i n es rhs) = LConCase i n es (caseFloat rhs)
235 cfAlt (LConstCase c rhs) = LConstCase c (caseFloat rhs)
236 cfAlt (LDefaultCase rhs) = LDefaultCase (caseFloat rhs)
237caseFloat exp = exp
239-- Case of constructor
240conOpt :: LExp -> LExp
241conOpt (LCase ct (LCon _ t n args) alts)
242 = pickAlt n args alts
243 where
244 pickAlt n args (LConCase i n' es rhs : as) | n == n'
245 = substAll (zip es args) rhs
246 pickAlt _ _ (LDefaultCase rhs : as) = rhs
247 pickAlt n args (_ : as) = pickAlt n args as
248 pickAlt n args [] = error "Can't happen pickAlt - impossible case found"
250 substAll [] rhs = rhs
251 substAll ((n, tm) : ss) rhs = lsubst n tm (substAll ss rhs)
252conOpt tm = tm
254replaceInCase :: LExp -> LExp
255replaceInCase (LCase ty e alts)
256 = LCase ty e (replaceInAlts e alts)
257replaceInCase exp = exp
259replaceInAlts :: LExp -> [LAlt] -> [LAlt]
260replaceInAlts exp alts = dropDups $ concatMap (replaceInAlt exp) alts
262-- Drop overlapping case (arising from case merging of overlapping
263-- patterns)
264dropDups (alt@(LConCase _ i n ns) : alts)
265 = alt : dropDups (filter (notTag i) alts)
266 where
267 notTag i (LConCase _ j n ns) = i /= j
268 notTag _ _ = True
269dropDups (c : alts) = c : dropDups alts
270dropDups [] = []
273replaceInAlt :: LExp -> LAlt -> [LAlt]
274-- In an alternative, if the case appears on the right hand side, replace
275-- it with the given expression, to preserve sharing
276replaceInAlt exp@(LV _) (LConCase i con args rhs)
277 = [LConCase i con args $
278 replaceExp (LCon Nothing i con (map LV args)) exp rhs]
279-- if a default case inspects the same variable as the case it's in,
280-- remove the inspection and replace with the alternatives
281-- (i.e. merge the inner case block)
282replaceInAlt exp@(LV var) (LDefaultCase (LCase ty (LV var') alts))
283 | var == var' = alts
284replaceInAlt exp a = [a]
286replaceExp :: LExp -> LExp -> LExp -> LExp
287replaceExp (LCon _ t n args) new (LCon _ t' n' args')
288 | n == n' && args == args' = new
289replaceExp (LCon _ t n args) new (LApp _ (LV n') args')
290 | n == n' && args == args' = new
291replaceExp old new tm = tm
293-- dropArgs as (LConstCase c rhs) = LConstCase c (dropRHS as rhs)
294-- dropArgs as (LDefaultCase rhs) = LDefaultCase (dropRHS as rhs)
296getRHS (LConCase i n es rhs) = rhs
297getRHS (LConstCase _ rhs) = rhs
298getRHS (LDefaultCase rhs) = rhs
300getLams [] = []
301getLams (LLam args tm : cs) = getLamPrefix args cs
302getLams (LLet n val exp : cs) = getLams (exp : cs)
303getLams _ = []
305getLamPrefix as [] = as
306getLamPrefix as (LLam args tm : cs)
307 | length args < length as = getLamPrefix args cs
308 | otherwise = getLamPrefix as cs
309getLamPrefix as (LLet n val exp : cs) = getLamPrefix as (exp : cs)
310getLamPrefix as (_ : cs) = []
312-- eta contract ('\x -> f x' can just be compiled as 'f' when f is local)
313eta :: LExp -> LExp
314eta (LApp tc a es) = LApp tc (eta a) (map eta es)
315eta (LLazyApp n es) = LLazyApp n (map eta es)
316eta (LLazyExp e) = LLazyExp (eta e)
317eta (LForce e) = LForce (eta e)
318eta (LLet n val sc) = LLet n (eta val) (eta sc)
319eta (LLam args (LApp tc f args'))
320 | args' == map LV args = eta f
321eta (LLam args e) = LLam args (eta e)
322eta (LProj e i) = LProj (eta e) i
323eta (LCon a t n es) = LCon a t n (map eta es)
324eta (LCase ct e alts) = LCase ct (eta e) (map (fmap eta) alts)
325eta (LOp f es) = LOp f (map eta es)
326eta tm = tm