on aa7262e help sign in
Idris-dev/src/IRTS/Defunctionalise.hs
248 lines 10 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
11data DExp = DV LVar
12 | DApp Bool Name [DExp] -- True = tail call
13 | DLet Name DExp DExp -- name just for pretty printing
14 | DProj DExp Int
15 | DC Int Name [DExp]
16 | DCase DExp [DAlt]
17 | DConst Const
18 | DForeign FLang FType String [(FType, DExp)]
19 | DOp PrimFn [DExp]
20 | DNothing -- erased value, can be compiled to anything since it'll never
21 -- be inspected
22 | DError String
23 deriving Eq
25data DAlt = DConCase Int Name [Name] DExp
26 | DConstCase Const DExp
27 | DDefaultCase DExp
28 deriving (Show, Eq)
30data DDecl = DFun Name [Name] DExp -- name, arg names, definition
31 | DConstructor Name Int Int -- constructor name, tag, arity
32 deriving (Show, Eq)
34type DDefs = Ctxt DDecl
36defunctionalise :: Int -> LDefs -> DDefs
37defunctionalise nexttag defs
38 = let all = toAlist defs
39 -- sort newcons so that EVAL and APPLY cons get sequential tags
40 newcons = sortBy conord $ concatMap toCons (getFn all)
41 eval = mkEval newcons
42 app = mkApply newcons
43 condecls = declare nexttag newcons in
44 addAlist (eval : app : condecls ++ (map (addApps defs) all)) emptyContext
45 where conord (n, _, _) (n', _, _) = compare n n'
47getFn :: [(Name, LDecl)] -> [(Name, Int)]
48getFn xs = mapMaybe fnData xs
49 where fnData (n, LFun _ args _) = Just (n, length args)
50 fnData _ = Nothing
52-- To defunctionalise:
53--
54-- 1 Create a data constructor for each function
55-- 2 Create a data constructor for each underapplication of a function
56-- 3 Convert underapplications to their corresponding constructors
57-- 4 Create an EVAL function which calls the appropriate function for data constructors
58-- created as part of step 1
59-- 5 Create an APPLY function which adds an argument to each underapplication (or calls
60-- APPLY again for an exact application)
61-- 6 Wrap overapplications in chains of APPLY
62-- 7 Wrap unknown applications (i.e. applications of local variables) in chains of APPLY
63-- 8 Add explicit EVAL to case, primitives, and foreign calls
65addApps :: LDefs -> (Name, LDecl) -> (Name, DDecl)
66addApps defs o@(n, LConstructor _ t a) = (n, DConstructor n t a)
67addApps defs (n, LFun _ args e) = (n, DFun n args (aa args e))
68 where
69 aa :: [Name] -> LExp -> DExp
70 aa env (LV (Glob n)) | n `elem` env = DV (Glob n)
71 | otherwise = aa env (LApp False (LV (Glob n)) [])
72-- aa env e@(LApp tc (MN 0 "EVAL") [a]) = e
73 aa env (LApp tc (LV (Glob n)) args)
74 = let args' = map (aa env) args in
75 case lookupCtxt Nothing n defs of
76 [LConstructor _ i ar] -> DApp tc n args'
77 [LFun _ as _] -> let arity = length as in
78 fixApply tc n args' arity
79 [] -> chainAPPLY (DV (Glob n)) args'
80 aa env (LLazyApp n args)
81 = let args' = map (aa env) args in
82 case lookupCtxt Nothing n defs of
83 [LConstructor _ i ar] -> DApp False n args'
84 [LFun _ as _] -> let arity = length as in
85 fixLazyApply n args' arity
86 [] -> chainAPPLY (DV (Glob n)) args'
87 aa env (LForce e) = eEVAL (aa env e)
88 aa env (LLet n v sc) = DLet n (aa env v) (aa (n : env) sc)
89 aa env (LCon i n args) = DC i n (map (aa env) args)
90 aa env (LProj t i) = DProj (eEVAL (aa env t)) i
91 aa env (LCase e alts) = DCase (eEVAL (aa env e)) (map (aaAlt env) alts)
92 aa env (LConst c) = DConst c
93 aa env (LForeign l t n args) = DForeign l t n (map (aaF env) args)
94 aa env (LOp LFork args) = DOp LFork (map (aa env) args)
95 aa env (LOp f args) = DOp f (map (eEVAL . (aa env)) args)
96 aa env LNothing = DNothing
97 aa env (LError e) = DError e
99 aaF env (t, e) = (t, eEVAL (aa env e))
101 aaAlt env (LConCase i n args e) = DConCase i n args (aa (args ++ env) e)
102 aaAlt env (LConstCase c e) = DConstCase c (aa env e)
103 aaAlt env (LDefaultCase e) = DDefaultCase (aa env e)
105 fixApply tc n args ar
106 | length args == ar = DApp tc n args
107 | length args < ar = DApp tc (mkUnderCon n (ar - length args)) args
108 | length args > ar = chainAPPLY (DApp tc n (take ar args)) (drop ar args)
110 fixLazyApply n args ar
111 | length args == ar = DApp False (mkFnCon n) args
112 | length args < ar = DApp False (mkUnderCon n (ar - length args)) args
113 | length args > ar = chainAPPLY (DApp False n (take ar args)) (drop ar args)
115 chainAPPLY f [] = f
116 chainAPPLY f (a : as) = chainAPPLY (DApp False (MN 0 "APPLY") [f, a]) as
118eEVAL x = DApp False (MN 0 "EVAL") [x]
120data EvalApply a = EvalCase a
121 | ApplyCase a
122 deriving Show
124-- For a function name, generate a list of
125-- data constuctors, and whether to handle them in EVAL or APPLY
127toCons :: (Name, Int) -> [(Name, Int, EvalApply DAlt)]
128toCons (n, i)
129 = (mkFnCon n, i,
130 EvalCase (DConCase (-1) (mkFnCon n) (take i (genArgs 0))
131 (eEVAL (DApp False n (map (DV . Glob) (take i (genArgs 0)))))))
132 : mkApplyCase n 0 i
134mkApplyCase fname n ar | n == ar = []
135mkApplyCase fname n ar
136 = let nm = mkUnderCon fname (ar - n) in
137 (nm, n, ApplyCase (DConCase (-1) nm (take n (genArgs 0))
138 (DApp False (mkUnderCon fname (ar - (n + 1)))
139 (map (DV . Glob) (take n (genArgs 0) ++
140 [MN 0 "arg"])))))
141 : mkApplyCase fname (n + 1) ar
143mkEval :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
144mkEval xs = (MN 0 "EVAL", DFun (MN 0 "EVAL") [MN 0 "arg"]
145 (mkBigCase (MN 0 "EVAL") 256 (DV (Glob (MN 0 "arg")))
146 (mapMaybe evalCase xs ++
147 [DDefaultCase (DV (Glob (MN 0 "arg")))])))
148 where
149 evalCase (n, t, EvalCase x) = Just x
150 evalCase _ = Nothing
152mkApply :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
153mkApply xs = (MN 0 "APPLY", DFun (MN 0 "APPLY") [MN 0 "fn", MN 0 "arg"]
154 (mkBigCase (MN 0 "APPLY")
155 256 (DApp False (MN 0 "EVAL")
156 [DV (Glob (MN 0 "fn"))])
157 (mapMaybe applyCase xs)))
158 where
159 applyCase (n, t, ApplyCase x) = Just x
160 applyCase _ = Nothing
163declare :: Int -> [(Name, Int, EvalApply DAlt)] -> [(Name, DDecl)]
164declare t xs = dec' t xs [] where
165 dec' t [] acc = reverse acc
166 dec' t ((n, ar, _) : xs) acc = dec' (t + 1) xs ((n, DConstructor n t ar) : acc)
169genArgs i = MN i "P_c" : genArgs (i + 1)
171mkFnCon n = MN 0 ("P_" ++ show n)
172mkUnderCon n 0 = n
173mkUnderCon n missing = MN missing ("U_" ++ show n)
175instance Show DExp where
176 show e = show' [] e where
177 show' env (DV (Loc i)) = "var " ++ env!!i
178 show' env (DV (Glob n)) = show n
179 show' env (DApp _ e args) = show e ++ "(" ++
180 showSep ", " (map (show' env) args) ++")"
181 show' env (DLet n v e) = "let " ++ show n ++ " = " ++ show' env v ++ " in " ++
182 show' (env ++ [show n]) e
183 show' env (DC i n args) = show n ++ ")" ++ showSep ", " (map (show' env) args) ++ ")"
184 show' env (DProj t i) = show t ++ "!" ++ show i
185 show' env (DCase e alts) = "case " ++ show' env e ++ " of {\n\t" ++
186 showSep "\n\t| " (map (showAlt env) alts)
187 show' env (DConst c) = show c
188 show' env (DForeign lang ty n args)
189 = "foreign " ++ n ++ "(" ++ showSep ", " (map (show' env) (map snd args)) ++ ")"
190 show' env (DOp f args) = show f ++ "(" ++ showSep ", " (map (show' env) args) ++ ")"
191 show' env (DError str) = "error " ++ show str
192 show' env DNothing = "____"
194 showAlt env (DConCase _ n args e)
195 = show n ++ "(" ++ showSep ", " (map show args) ++ ") => "
196 ++ show' env e
197 showAlt env (DConstCase c e) = show c ++ " => " ++ show' env e
198 showAlt env (DDefaultCase e) = "_ => " ++ show' env e
200-- Divide up a large case expression so that each has a maximum of
201-- 'max' branches
203mkBigCase cn max arg branches
204 | length branches <= max = DCase arg branches
205 | otherwise = -- DCase arg branches -- until I think of something...
206 -- divide the branches into groups of at most max (by tag),
207 -- generate a new case and shrink, recursively
208 let bs = sortBy tagOrd branches
209 (all, def) = case (last bs) of
210 DDefaultCase t -> (init all, Just (DDefaultCase t))
211 _ -> (all, Nothing)
212 bss = groupsOf max all
213 cs = map mkCase bss in
214 DCase arg branches
216 where mkCase bs = DCase arg bs
218 tagOrd (DConCase t _ _ _) (DConCase t' _ _ _) = compare t t'
219 tagOrd (DConstCase c _) (DConstCase c' _) = compare c c'
220 tagOrd (DDefaultCase _) (DDefaultCase _) = EQ
221 tagOrd (DConCase _ _ _ _) (DDefaultCase _) = LT
222 tagOrd (DConCase _ _ _ _) (DConstCase _ _) = LT
223 tagOrd (DConstCase _ _) (DDefaultCase _) = LT
224 tagOrd (DDefaultCase _) (DConCase _ _ _ _) = GT
225 tagOrd (DConstCase _ _) (DConCase _ _ _ _) = GT
226 tagOrd (DDefaultCase _) (DConstCase _ _) = GT
229groupsOf :: Int -> [DAlt] -> [[DAlt]]
230groupsOf x [] = []
231groupsOf x xs = let (batch, rest) = span (tagLT (x + tagHead xs)) xs in
232 batch : groupsOf x rest
233 where tagHead (DConstCase (I i) _ : _) = i
234 tagHead (DConCase t _ _ _ : _) = t
235 tagHead (DDefaultCase _ : _) = -1 -- must be the end
237 tagLT i (DConstCase (I j) _) = i < j
238 tagLT i (DConCase j _ _ _) = i < j
239 tagLT i (DDefaultCase _) = False
241dumpDefuns :: DDefs -> String
242dumpDefuns ds = showSep "\n" $ map showDef (toAlist ds)
243 where showDef (x, DFun fn args exp)
244 = show fn ++ "(" ++ showSep ", " (map show args) ++ ") = \n\t" ++
245 show exp ++ "\n"
246 showDef (x, DConstructor n t a) = "Constructor " ++ show n ++ " " ++ show t