git.iko.soy/Idris-dev
on 385dcbf help sign in
Idris-dev/src/IRTS/Defunctionalise.hs
377 lines 16 kB history raw
1{-|
2Module : IRTS.Defunctionalise
3Description : Defunctionalise Idris' IR.
4Copyright :
5License : BSD3
6Maintainer : The Idris Community.
8To defunctionalise:
101. Create a data constructor for each function
112. Create a data constructor for each underapplication of a function
123. Convert underapplications to their corresponding constructors
134. Create an EVAL function which calls the appropriate function for data constructors
14 created as part of step 1
155. Create an APPLY function which adds an argument to each underapplication (or calls
16 APPLY again for an exact application)
176. Wrap overapplications in chains of APPLY
187. Wrap unknown applications (i.e. applications of local variables) in chains of APPLY
198. Add explicit EVAL to case, primitives, and foreign calls
21-}
22{-# LANGUAGE PatternGuards #-}
23module IRTS.Defunctionalise(module IRTS.Defunctionalise
24 , module IRTS.Lang
25 ) where
27import Idris.Core.CaseTree
28import Idris.Core.TT
29import IRTS.Lang
31import Control.Monad
32import Control.Monad.State
33import Data.List
34import Data.Maybe
35import Debug.Trace
37data DExp = DV LVar
38 | DApp Bool Name [DExp] -- True = tail call
39 | DLet Name DExp DExp -- name just for pretty printing
40 | DUpdate Name DExp -- eval expression, then update var with it
41 | DProj DExp Int
42 | DC (Maybe LVar) Int Name [DExp]
43 | DCase CaseType DExp [DAlt]
44 | DChkCase DExp [DAlt] -- a case where the type is unknown (for EVAL/APPLY)
45 | DConst Const
46 | DForeign FDesc FDesc [(FDesc, DExp)]
47 | DOp PrimFn [DExp]
48 | DNothing -- erased value, can be compiled to anything since it'll never
49 -- be inspected
50 | DError String
51 deriving Eq
53data DAlt = DConCase Int Name [Name] DExp
54 | DConstCase Const DExp
55 | DDefaultCase DExp
56 deriving (Show, Eq)
58data DDecl = DFun Name [Name] DExp -- name, arg names, definition
59 | DConstructor Name Int Int -- constructor name, tag, arity
60 deriving (Show, Eq)
62type DDefs = Ctxt DDecl
64defunctionalise :: Int -> LDefs -> DDefs
65defunctionalise nexttag defs
66 = let all = toAlist defs
67 -- sort newcons so that EVAL and APPLY cons get sequential tags
68 (allD, (enames, anames)) = runState (mapM (addApps defs) all) ([], [])
69 anames' = sort (nub anames)
70 enames' = nub enames
71 newecons = sortBy conord $ concatMap (toCons enames') (getFn all)
72 newacons = sortBy conord $ concatMap (toConsA anames') (getFn all)
73 eval = mkEval newecons
74 app = mkApply newacons
75 app2 = mkApply2 newacons
76 condecls = declare nexttag (newecons ++ newacons) in
77 addAlist (eval : app : app2 : condecls ++ allD) emptyContext
78 where conord (n, _, _) (n', _, _) = compare n n'
80getFn :: [(Name, LDecl)] -> [(Name, Int)]
81getFn xs = mapMaybe fnData xs
82 where fnData (n, LFun _ _ args _) = Just (n, length args)
83 fnData _ = Nothing
85addApps :: LDefs -> (Name, LDecl) -> State ([Name], [(Name, Int)]) (Name, DDecl)
86addApps defs o@(n, LConstructor _ t a)
87 = return (n, DConstructor n t a)
88addApps defs (n, LFun _ _ args e)
89 = do e' <- aa args e
90 return (n, DFun n args e')
91 where
92 aa :: [Name] -> LExp -> State ([Name], [(Name, Int)]) DExp
93 aa env (LV (Glob n)) | n `elem` env = return $ DV (Glob n)
94 | otherwise = aa env (LApp False (LV (Glob n)) [])
95 aa env (LApp tc (LV (Glob n)) args)
96 = do args' <- mapM (aa env) args
97 case lookupCtxtExact n defs of
98 Just (LConstructor _ i ar) -> return $ DApp tc n args'
99 Just (LFun _ _ as _) -> let arity = length as in
100 fixApply tc n args' arity
101 Nothing -> return $ chainAPPLY (DV (Glob n)) args'
102 aa env (LLazyApp n args)
103 = do args' <- mapM (aa env) args
104 case lookupCtxtExact n defs of
105 Just (LConstructor _ i ar) -> return $ DApp False n args'
106 Just (LFun _ _ as _) -> let arity = length as in
107 fixLazyApply n args' arity
108 Nothing -> return $ chainAPPLY (DV (Glob n)) args'
109 aa env (LForce (LLazyApp n args)) = aa env (LApp False (LV (Glob n)) args)
110 aa env (LForce e) = liftM eEVAL (aa env e)
111 aa env (LLet n v sc) = liftM2 (DLet n) (aa env v) (aa (n : env) sc)
112 aa env (LCon loc i n args) = liftM (DC loc i n) (mapM (aa env) args)
113 aa env (LProj t@(LV (Glob n)) i)
114 | n `elem` env = do t' <- aa env t
115 return $ DProj (DUpdate n t') i
116 aa env (LProj t i) = do t' <- aa env t
117 return $ DProj t' i
118 aa env (LCase up e alts) = do e' <- aa env e
119 alts' <- mapM (aaAlt env) alts
120 return $ DCase up e' alts'
121 aa env (LConst c) = return $ DConst c
122 aa env (LForeign t n args)
123 = do args' <- mapM (aaF env) args
124 return $ DForeign t n args'
125 aa env (LOp LFork args) = liftM (DOp LFork) (mapM (aa env) args)
126 aa env (LOp f args) = do args' <- mapM (aa env) args
127 return $ DOp f args'
128 aa env LNothing = return DNothing
129 aa env (LError e) = return $ DError e
131 aaF env (t, e) = do e' <- aa env e
132 return (t, e')
134 aaAlt env (LConCase i n args e)
135 = liftM (DConCase i n args) (aa (args ++ env) e)
136 aaAlt env (LConstCase c e) = liftM (DConstCase c) (aa env e)
137 aaAlt env (LDefaultCase e) = liftM DDefaultCase (aa env e)
139 fixApply tc n args ar
140 | length args == ar
141 = return $ DApp tc n args
142 | length args < ar
143 = do (ens, ans) <- get
144 let alln = map (\x -> (n, x)) [length args .. ar]
145 put (ens, alln ++ ans)
146 return $ DApp tc (mkUnderCon n (ar - length args)) args
147 | length args > ar
148 = return $ chainAPPLY (DApp tc n (take ar args)) (drop ar args)
150 fixLazyApply n args ar
151 | length args == ar
152 = do (ens, ans) <- get
153 put (n : ens, ans)
154 return $ DApp False (mkFnCon n) args
155 | length args < ar
156 = do (ens, ans) <- get
157 let alln = map (\x -> (n, x)) [length args .. ar]
158 put (ens, alln ++ ans)
159 return $ DApp False (mkUnderCon n (ar - length args)) args
160 | length args > ar
161 = return $ chainAPPLY (DApp False n (take ar args)) (drop ar args)
163 chainAPPLY f [] = f
164-- chainAPPLY f (a : b : as)
165-- = chainAPPLY (DApp False (sMN 0 "APPLY2") [f, a, b]) as
166 chainAPPLY f (a : as) = chainAPPLY (DApp False (sMN 0 "APPLY") [f, a]) as
168 -- if anything in the DExp is projected from, we'll need to evaluate it,
169 -- but we only want to do it once, rather than every time we project.
171 preEval [] t = t
172 preEval (x : xs) t
173 | needsEval x t = DLet x (DV (Glob x)) (preEval xs t)
174 | otherwise = preEval xs t
176 needsEval x (DApp _ _ args) = any (needsEval x) args
177 needsEval x (DC _ _ _ args) = any (needsEval x) args
178 needsEval x (DCase up e alts) = needsEval x e || any nec alts
179 where nec (DConCase _ _ _ e) = needsEval x e
180 nec (DConstCase _ e) = needsEval x e
181 nec (DDefaultCase e) = needsEval x e
182 needsEval x (DChkCase e alts) = needsEval x e || any nec alts
183 where nec (DConCase _ _ _ e) = needsEval x e
184 nec (DConstCase _ e) = needsEval x e
185 nec (DDefaultCase e) = needsEval x e
186 needsEval x (DLet n v e)
187 | x == n = needsEval x v
188 | otherwise = needsEval x v || needsEval x e
189 needsEval x (DForeign _ _ args) = any (needsEval x) (map snd args)
190 needsEval x (DOp op args) = any (needsEval x) args
191 needsEval x (DProj (DV (Glob x')) _) = x == x'
192 needsEval x _ = False
194eEVAL x = DApp False (sMN 0 "EVAL") [x]
196data EvalApply a = EvalCase (Name -> a)
197 | ApplyCase a
198 | Apply2Case a
200-- For a function name, generate a list of
201-- data constuctors, and whether to handle them in EVAL or APPLY
203toCons :: [Name] -> (Name, Int) -> [(Name, Int, EvalApply DAlt)]
204toCons ns (n, i)
205 | n `elem` ns
206 = (mkFnCon n, i,
207 EvalCase (\tlarg ->
208 (DConCase (-1) (mkFnCon n) (take i (genArgs 0))
209 (dupdate tlarg
210 (DApp False n (map (DV . Glob) (take i (genArgs 0))))))))
211 : [] -- mkApplyCase n 0 i
212 | otherwise = []
213 where dupdate tlarg x = DUpdate tlarg x
215toConsA :: [(Name, Int)] -> (Name, Int) -> [(Name, Int, EvalApply DAlt)]
216toConsA ns (n, i)
217 | Just ar <- lookup n ns
218-- = (mkFnCon n, i,
219-- EvalCase (\tlarg ->
220-- (DConCase (-1) (mkFnCon n) (take i (genArgs 0))
221-- (dupdate tlarg
222-- (DApp False n (map (DV . Glob) (take i (genArgs 0))))))))
223 = mkApplyCase n ar i
224 | otherwise = []
225 where dupdate tlarg x = x
227mkApplyCase fname n ar | n == ar = []
228mkApplyCase fname n ar
229 = let nm = mkUnderCon fname (ar - n) in
230 (nm, n, ApplyCase (DConCase (-1) nm (take n (genArgs 0))
231 (DApp False (mkUnderCon fname (ar - (n + 1)))
232 (map (DV . Glob) (take n (genArgs 0) ++
233 [sMN 0 "arg"])))))
234 :
235 if (ar - (n + 2) >=0 )
236 then (nm, n, Apply2Case (DConCase (-1) nm (take n (genArgs 0))
237 (DApp False (mkUnderCon fname (ar - (n + 2)))
238 (map (DV . Glob) (take n (genArgs 0) ++
239 [sMN 0 "arg0", sMN 0 "arg1"])))))
240 :
241 mkApplyCase fname (n + 1) ar
242 else mkApplyCase fname (n + 1) ar
244mkEval :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
245mkEval xs = (sMN 0 "EVAL", DFun (sMN 0 "EVAL") [sMN 0 "arg"]
246 (mkBigCase (sMN 0 "EVAL") 256 (DV (Glob (sMN 0 "arg")))
247 (mapMaybe evalCase xs ++
248 [DDefaultCase (DV (Glob (sMN 0 "arg")))])))
249 where
250 evalCase (n, t, EvalCase x) = Just (x (sMN 0 "arg"))
251 evalCase _ = Nothing
253mkApply :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
254mkApply xs = (sMN 0 "APPLY", DFun (sMN 0 "APPLY") [sMN 0 "fn", sMN 0 "arg"]
255 (case mapMaybe applyCase xs of
256 [] -> DNothing
257 cases ->
258 mkBigCase (sMN 0 "APPLY") 256
259 (DV (Glob (sMN 0 "fn")))
260 (cases ++
261 [DDefaultCase DNothing])))
262 where
263 applyCase (n, t, ApplyCase x) = Just x
264 applyCase _ = Nothing
266mkApply2 :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
267mkApply2 xs = (sMN 0 "APPLY2", DFun (sMN 0 "APPLY2") [sMN 0 "fn", sMN 0 "arg0", sMN 0 "arg1"]
268 (case mapMaybe applyCase xs of
269 [] -> DNothing
270 cases ->
271 mkBigCase (sMN 0 "APPLY") 256
272 (DV (Glob (sMN 0 "fn")))
273 (cases ++
274 [DDefaultCase
275 (DApp False (sMN 0 "APPLY")
276 [DApp False (sMN 0 "APPLY")
277 [DV (Glob (sMN 0 "fn")),
278 DV (Glob (sMN 0 "arg0"))],
279 DV (Glob (sMN 0 "arg1"))])
280 ])))
281 where
282 applyCase (n, t, Apply2Case x) = Just x
283 applyCase _ = Nothing
286declare :: Int -> [(Name, Int, EvalApply DAlt)] -> [(Name, DDecl)]
287declare t xs = dec' t xs [] where
288 dec' t [] acc = reverse acc
289 dec' t ((n, ar, _) : xs) acc = dec' (t + 1) xs ((n, DConstructor n t ar) : acc)
292genArgs i = sMN i "P_c" : genArgs (i + 1)
294mkFnCon n = sMN 0 ("P_" ++ show n)
295mkUnderCon n 0 = n
296mkUnderCon n missing = sMN missing ("U_" ++ show n)
298instance Show DExp where
299 show e = show' [] e where
300 show' env (DV (Loc i)) = "var " ++ env!!i
301 show' env (DV (Glob n)) = "GLOB " ++ show n
302 show' env (DApp _ e args) = show e ++ "(" ++
303 showSep ", " (map (show' env) args) ++")"
304 show' env (DLet n v e) = "let " ++ show n ++ " = " ++ show' env v ++ " in " ++
305 show' (env ++ [show n]) e
306 show' env (DUpdate n e) = "!update " ++ show n ++ "(" ++ show' env e ++ ")"
307 show' env (DC loc i n args) = atloc loc ++ "CON " ++ show n ++ "(" ++ showSep ", " (map (show' env) args) ++ ")"
308 where atloc Nothing = ""
309 atloc (Just l) = "@" ++ show (LV l) ++ ":"
310 show' env (DProj t i) = show t ++ "!" ++ show i
311 show' env (DCase up e alts) = "case" ++ update ++ show' env e ++ " of {\n\t" ++
312 showSep "\n\t| " (map (showAlt env) alts)
313 where update = case up of
314 Shared -> " "
315 Updatable -> "! "
316 show' env (DChkCase e alts) = "case' " ++ show' env e ++ " of {\n\t" ++
317 showSep "\n\t| " (map (showAlt env) alts)
318 show' env (DConst c) = show c
319 show' env (DForeign ty n args)
320 = "foreign " ++ show n ++ "(" ++ showSep ", " (map (show' env) (map snd args)) ++ ")"
321 show' env (DOp f args) = show f ++ "(" ++ showSep ", " (map (show' env) args) ++ ")"
322 show' env (DError str) = "error " ++ show str
323 show' env DNothing = "____"
325 showAlt env (DConCase _ n args e)
326 = show n ++ "(" ++ showSep ", " (map show args) ++ ") => "
327 ++ show' env e
328 showAlt env (DConstCase c e) = show c ++ " => " ++ show' env e
329 showAlt env (DDefaultCase e) = "_ => " ++ show' env e
331-- | Divide up a large case expression so that each has a maximum of
332-- 'max' branches
333mkBigCase cn max arg branches
334 | length branches <= max = DChkCase arg branches
335 | otherwise = -- DChkCase arg branches -- until I think of something...
336 -- divide the branches into groups of at most max (by tag),
337 -- generate a new case and shrink, recursively
338 let bs = sortBy tagOrd branches
339 (all, def) = case (last bs) of
340 DDefaultCase t -> (init all, Just (DDefaultCase t))
341 _ -> (all, Nothing)
342 bss = groupsOf max all
343 cs = map mkCase bss in
344 DChkCase arg branches
346 where mkCase bs = DChkCase arg bs
348 tagOrd (DConCase t _ _ _) (DConCase t' _ _ _) = compare t t'
349 tagOrd (DConstCase c _) (DConstCase c' _) = compare c c'
350 tagOrd (DDefaultCase _) (DDefaultCase _) = EQ
352 tagOrd (DConCase _ _ _ _) (DDefaultCase _) = LT
353 tagOrd (DConCase _ _ _ _) (DConstCase _ _) = LT
354 tagOrd (DConstCase _ _) (DDefaultCase _) = LT
356 tagOrd (DDefaultCase _) (DConCase _ _ _ _) = GT
357 tagOrd (DConstCase _ _) (DConCase _ _ _ _) = GT
358 tagOrd (DDefaultCase _) (DConstCase _ _) = GT
360groupsOf :: Int -> [DAlt] -> [[DAlt]]
361groupsOf x [] = []
362groupsOf x xs = let (batch, rest) = span (tagLT (x + tagHead xs)) xs in
363 batch : groupsOf x rest
364 where tagHead (DConstCase (I i) _ : _) = i
365 tagHead (DConCase t _ _ _ : _) = t
366 tagHead (DDefaultCase _ : _) = -1 -- must be the end
368 tagLT i (DConstCase (I j) _) = i < j
369 tagLT i (DConCase j _ _ _) = i < j
370 tagLT i (DDefaultCase _) = False
372dumpDefuns :: DDefs -> String
373dumpDefuns ds = showSep "\n" $ map showDef (toAlist ds)
374 where showDef (x, DFun fn args exp)
375 = show fn ++ "(" ++ showSep ", " (map show args) ++ ") = \n\t" ++
376 show exp ++ "\n"
377 showDef (x, DConstructor n t a) = "Constructor " ++ show n ++ " " ++ show t