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