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