git.iko.soy/Idris-dev
on d90165e help sign in
Idris-dev/src/Core/CaseTree.hs
422 lines 16 kB history raw
1{-# LANGUAGE PatternGuards #-}
3module Core.CaseTree(CaseDef(..), SC(..), CaseAlt(..), Phase(..), CaseTree,
4 simpleCase, small, namesUsed, findCalls, findUsedArgs) where
6import Core.TT
8import Control.Monad.State
9import Data.Maybe
10import Data.List hiding (partition)
11import Debug.Trace
13data CaseDef = CaseDef [Name] SC [Term]
14 deriving Show
16data SC = Case Name [CaseAlt] -- invariant: lowest tags first
17 | ProjCase Term [CaseAlt] -- special case for projections
18 | STerm Term
19 | UnmatchedCase String -- error message
20 | ImpossibleCase -- already checked to be impossible
21 deriving (Eq, Ord)
22{-!
23deriving instance Binary SC
24!-}
26data CaseAlt = ConCase Name Int [Name] SC
27 | ConstCase Const SC
28 | DefaultCase SC
29 deriving (Show, Eq, Ord)
30{-!
31deriving instance Binary CaseAlt
32!-}
34instance Show SC where
35 show sc = show' 1 sc
36 where
37 show' i (Case n alts) = "case " ++ show n ++ " of\n" ++ indent i ++
38 showSep ("\n" ++ indent i) (map (showA i) alts)
39 show' i (ProjCase tm alts) = "case " ++ show tm ++ " of " ++
40 showSep ("\n" ++ indent i) (map (showA i) alts)
41 show' i (STerm tm) = show tm
42 show' i (UnmatchedCase str) = "error " ++ show str
43 show' i ImpossibleCase = "impossible"
45 indent i = concat $ take i (repeat " ")
47 showA i (ConCase n t args sc)
48 = show n ++ "(" ++ showSep (", ") (map show args) ++ ") => "
49 ++ show' (i+1) sc
50 showA i (ConstCase t sc)
51 = show t ++ " => " ++ show' (i+1) sc
52 showA i (DefaultCase sc)
53 = "_ => " ++ show' (i+1) sc
56type CaseTree = SC
57type Clause = ([Pat], (Term, Term))
58type CS = ([Term], Int)
60instance TermSize SC where
61 termsize n (Case n' as) = termsize n as
62 termsize n (STerm t) = termsize n t
63 termsize n _ = 1
65instance TermSize CaseAlt where
66 termsize n (ConCase _ _ _ s) = termsize n s
67 termsize n (ConstCase _ s) = termsize n s
68 termsize n (DefaultCase s) = termsize n s
70-- simple terms can be inlined trivially - good for primitives in particular
71small :: Name -> SC -> Bool
72small n t = termsize n t < 50
74namesUsed :: SC -> [Name]
75namesUsed sc = nub $ nu' [] sc where
76 nu' ps (Case n alts) = nub (concatMap (nua ps) alts) \\ [n]
77 nu' ps (STerm t) = nub $ nut ps t
78 nu' ps _ = []
80 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc) \\ args
81 nua ps (ConstCase _ sc) = nu' ps sc
82 nua ps (DefaultCase sc) = nu' ps sc
84 nut ps (P _ n _) | n `elem` ps = []
85 | otherwise = [n]
86 nut ps (App f a) = nut ps f ++ nut ps a
87 nut ps (Bind n (Let t v) sc) = nut ps v ++ nut (n:ps) sc
88 nut ps (Bind n b sc) = nut (n:ps) sc
89 nut ps _ = []
91-- Return all called functions, and which arguments are used in each argument position
92-- for the call, in order to help reduce compilation time, and trace all unused
93-- arguments
95findCalls :: SC -> [Name] -> [(Name, [[Name]])]
96findCalls sc topargs = nub $ nu' topargs sc where
97 nu' ps (Case n alts) = nub (concatMap (nua (n : ps)) alts)
98 nu' ps (STerm t) = nub $ nut ps t
99 nu' ps _ = []
101 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc)
102 nua ps (ConstCase _ sc) = nu' ps sc
103 nua ps (DefaultCase sc) = nu' ps sc
105 nut ps (P Ref n _) | n `elem` ps = []
106 | otherwise = [(n, [])] -- tmp
107 nut ps fn@(App f a)
108 | (P Ref n _, args) <- unApply fn
109 = if n `elem` ps then nut ps f ++ nut ps a
110 else [(n, map argNames args)] ++ concatMap (nut ps) args
111 | otherwise = nut ps f ++ nut ps a
112 nut ps (Bind n (Let t v) sc) = nut ps v ++ nut (n:ps) sc
113 nut ps (Bind n b sc) = nut (n:ps) sc
114 nut ps _ = []
116 argNames tm = let ns = directUse tm in
117 filter (\x -> x `elem` ns) topargs
119-- Find names which are used directly (i.e. not in a function call) in a term
121directUse :: Eq n => TT n -> [n]
122directUse (P _ n _) = [n]
123directUse (Bind n (Let t v) sc) = nub $ directUse v ++ (directUse sc \\ [n])
124 ++ directUse t
125directUse (Bind n b sc) = nub $ directUse (binderTy b) ++ (directUse sc \\ [n])
126directUse fn@(App f a)
127 | (P Ref n _, args) <- unApply fn = [] -- need to know what n does with them
128 | otherwise = nub $ directUse f ++ directUse a
129directUse (Proj x i) = nub $ directUse x
130directUse _ = []
132-- Find all directly used arguments (i.e. used but not in function calls)
134findUsedArgs :: SC -> [Name] -> [Name]
135findUsedArgs sc topargs = filter (\x -> x `elem` topargs) (nub $ nu' sc) where
136 nu' (Case n alts) = n : concatMap nua alts
137 nu' (STerm t) = directUse t
138 nu' _ = []
140 nua (ConCase n i args sc) = nu' sc
141 nua (ConstCase _ sc) = nu' sc
142 nua (DefaultCase sc) = nu' sc
144data Phase = CompileTime | RunTime
145 deriving (Show, Eq)
147-- Generate a simple case tree
148-- Work Left to Right at Compile Time
150simpleCase :: Bool -> Bool -> Phase -> FC -> [([Name], Term, Term)] -> TC CaseDef
151simpleCase tc cover phase fc []
152 = return $ CaseDef [] (UnmatchedCase "No pattern clauses") []
153simpleCase tc cover phase fc cs
154 = let proj = phase == RunTime
155 pats = map (\ (avs, l, r) ->
156 (avs, rev phase (toPats tc l), (l, r))) cs
157 chkPats = mapM chkAccessible pats in
158 case chkPats of
159 OK pats ->
160 let numargs = length (fst (head pats))
161 ns = take numargs args
162 (tree, st) = runState
163 (match (rev phase ns) pats (defaultCase cover)) ([], numargs)
164 t = CaseDef ns (prune proj (depatt ns tree)) (fst st) in
165 if proj then return (stripLambdas t) else return t
166 Error err -> Error (At fc err)
167 where args = map (\i -> MN i "e") [0..]
168 defaultCase True = STerm Erased
169 defaultCase False = UnmatchedCase "Error"
171 chkAccessible (avs, l, c)
172 | phase == RunTime = return (l, c)
173 | otherwise = do mapM_ (acc l) avs
174 return (l, c)
176 acc [] n = Error (Inaccessible n)
177 acc (PV x : xs) n | x == n = OK ()
178 acc (PCon _ _ ps : xs) n = acc (ps ++ xs) n
179 acc (_ : xs) n = acc xs n
181rev CompileTime = id
182rev _ = reverse
184data Pat = PCon Name Int [Pat]
185 | PConst Const
186 | PV Name
187 | PAny
188 deriving Show
190-- If there are repeated variables, take the *last* one (could be name shadowing
191-- in a where clause, so take the most recent).
193toPats :: Bool -> Term -> [Pat]
194toPats tc f = reverse (toPat tc (getArgs f)) where
195 getArgs (App f a) = a : getArgs f
196 getArgs _ = []
198toPat :: Bool -> [Term] -> [Pat]
199toPat tc tms = evalState (mapM (\x -> toPat' x []) tms) []
200 where
201 toPat' (P (DCon t a) n _) args = do args' <- mapM (\x -> toPat' x []) args
202 return $ PCon n t args'
203 -- Typecase
204 toPat' (P (TCon t a) n _) args | tc
205 = do args' <- mapM (\x -> toPat' x []) args
206 return $ PCon n t args'
207 toPat' (Constant IType) [] | tc = return $ PCon (UN "Int") 1 []
208 toPat' (Constant FlType) [] | tc = return $ PCon (UN "Float") 2 []
209 toPat' (Constant ChType) [] | tc = return $ PCon (UN "Char") 3 []
210 toPat' (Constant StrType) [] | tc = return $ PCon (UN "String") 4 []
211 toPat' (Constant PtrType) [] | tc = return $ PCon (UN "Ptr") 5 []
212 toPat' (Constant BIType) [] | tc = return $ PCon (UN "Integer") 6 []
214 toPat' (P Bound n _) [] = do ns <- get
215 if n `elem` ns
216 then return PAny
217 else do put (n : ns)
218 return (PV n)
219 toPat' (App f a) args = toPat' f (a : args)
220 toPat' (Constant x) [] = return $ PConst x
221 toPat' _ _ = return PAny
224data Partition = Cons [Clause]
225 | Vars [Clause]
226 deriving Show
228isVarPat (PV _ : ps , _) = True
229isVarPat (PAny : ps , _) = True
230isVarPat _ = False
232isConPat (PCon _ _ _ : ps, _) = True
233isConPat (PConst _ : ps, _) = True
234isConPat _ = False
236partition :: [Clause] -> [Partition]
237partition [] = []
238partition ms@(m : _)
239 | isVarPat m = let (vars, rest) = span isVarPat ms in
240 Vars vars : partition rest
241 | isConPat m = let (cons, rest) = span isConPat ms in
242 Cons cons : partition rest
243partition xs = error $ "Partition " ++ show xs
245match :: [Name] -> [Clause] -> SC -- error case
246 -> State CS SC
247match [] (([], ret) : xs) err
248 = do (ts, v) <- get
249 put (ts ++ (map (fst.snd) xs), v)
250 case snd ret of
251 Impossible -> return ImpossibleCase
252 tm -> return $ STerm tm -- run out of arguments
253match vs cs err = do let ps = partition cs
254 cs <- mixture vs ps err
255 return cs
257mixture :: [Name] -> [Partition] -> SC -> State CS SC
258mixture vs [] err = return err
259mixture vs (Cons ms : ps) err = do fallthrough <- mixture vs ps err
260 conRule vs ms fallthrough
261mixture vs (Vars ms : ps) err = do fallthrough <- mixture vs ps err
262 varRule vs ms fallthrough
264data ConType = CName Name Int -- named constructor
265 | CConst Const -- constant, not implemented yet
267data Group = ConGroup ConType -- Constructor
268 [([Pat], Clause)] -- arguments and rest of alternative
270conRule :: [Name] -> [Clause] -> SC -> State CS SC
271conRule (v:vs) cs err = do groups <- groupCons cs
272 caseGroups (v:vs) groups err
274caseGroups :: [Name] -> [Group] -> SC -> State CS SC
275caseGroups (v:vs) gs err = do g <- altGroups gs
276 return $ Case v (sort g)
277 where
278 altGroups [] = return [DefaultCase err]
279 altGroups (ConGroup (CName n i) args : cs)
280 = do g <- altGroup n i args
281 rest <- altGroups cs
282 return (g : rest)
283 altGroups (ConGroup (CConst c) args : cs)
284 = do g <- altConstGroup c args
285 rest <- altGroups cs
286 return (g : rest)
288 altGroup n i gs = do (newArgs, nextCs) <- argsToAlt gs
289 matchCs <- match (newArgs ++ vs) nextCs err
290 return $ ConCase n i newArgs matchCs
291 altConstGroup n gs = do (_, nextCs) <- argsToAlt gs
292 matchCs <- match vs nextCs err
293 return $ ConstCase n matchCs
295argsToAlt :: [([Pat], Clause)] -> State CS ([Name], [Clause])
296argsToAlt [] = return ([], [])
297argsToAlt rs@((r, m) : rest)
298 = do newArgs <- getNewVars r
299 return (newArgs, addRs rs)
300 where
301 getNewVars [] = return []
302 getNewVars ((PV n) : ns) = do v <- getVar
303 nsv <- getNewVars ns
304 return (v : nsv)
305 getNewVars (_ : ns) = do v <- getVar
306 nsv <- getNewVars ns
307 return (v : nsv)
308 addRs [] = []
309 addRs ((r, (ps, res)) : rs) = ((r++ps, res) : addRs rs)
311 uniq i (UN n) = MN i n
312 uniq i n = n
314getVar :: State CS Name
315getVar = do (t, v) <- get; put (t, v+1); return (MN v "e")
317groupCons :: [Clause] -> State CS [Group]
318groupCons cs = gc [] cs
319 where
320 gc acc [] = return acc
321 gc acc ((p : ps, res) : cs) =
322 do acc' <- addGroup p ps res acc
323 gc acc' cs
324 addGroup p ps res acc = case p of
325 PCon con i args -> return $ addg con i args (ps, res) acc
326 PConst cval -> return $ addConG cval (ps, res) acc
327 pat -> fail $ show pat ++ " is not a constructor or constant (can't happen)"
329 addg con i conargs res [] = [ConGroup (CName con i) [(conargs, res)]]
330 addg con i conargs res (g@(ConGroup (CName n j) cs):gs)
331 | i == j = ConGroup (CName n i) (cs ++ [(conargs, res)]) : gs
332 | otherwise = g : addg con i conargs res gs
334 addConG con res [] = [ConGroup (CConst con) [([], res)]]
335 addConG con res (g@(ConGroup (CConst n) cs) : gs)
336 | con == n = ConGroup (CConst n) (cs ++ [([], res)]) : gs
337 | otherwise = g : addConG con res gs
339varRule :: [Name] -> [Clause] -> SC -> State CS SC
340varRule (v : vs) alts err =
341 do let alts' = map (repVar v) alts
342 match vs alts' err
343 where
344 repVar v (PV p : ps , (lhs, res)) = (ps, (lhs, subst p (P Bound v Erased) res))
345 repVar v (PAny : ps , res) = (ps, res)
347-- fix: case e of S k -> f (S k) ==> case e of S k -. f e
348depatt :: [Name] -> SC -> SC
349depatt ns tm = dp [] tm
350 where
351 dp ms (STerm tm) = STerm (applyMaps ms tm)
352 dp ms (Case x alts) = Case x (map (dpa ms x) alts)
353 dp ms sc = sc
355 dpa ms x (ConCase n i args sc)
356 = ConCase n i args (dp ((x, (n, args)) : ms) sc)
357 dpa ms x (ConstCase c sc) = ConstCase c (dp ms sc)
358 dpa ms x (DefaultCase sc) = DefaultCase (dp ms sc)
360 applyMaps ms f@(App _ _)
361 | (P nt cn pty, args) <- unApply f
362 = let args' = map (applyMaps ms) args in
363 applyMap ms nt cn pty args'
364 where
365 applyMap [] nt cn pty args' = mkApp (P nt cn pty) args'
366 applyMap ((x, (n, args)) : ms) nt cn pty args'
367 | and ((n == cn) : zipWith same args args') = P Ref x Erased
368 | otherwise = applyMap ms nt cn pty args'
369 same n (P _ n' _) = n == n'
370 same _ _ = False
371 applyMaps ms (App f a) = App (applyMaps ms f) (applyMaps ms a)
372 applyMaps ms t = t
374prune :: Bool -> -- ^ Convert single branches to projections (only useful at runtime)
375 SC -> SC
376prune proj (Case n alts)
377 = let alts' = filter notErased (map pruneAlt alts) in
378 case alts' of
379 [] -> ImpossibleCase
380 as@[ConCase cn i args sc] -> if proj then mkProj n 0 args sc
381 else Case n as
382 as -> Case n as
383 where pruneAlt (ConCase cn i ns sc) = ConCase cn i ns (prune proj sc)
384 pruneAlt (ConstCase c sc) = ConstCase c (prune proj sc)
385 pruneAlt (DefaultCase sc) = DefaultCase (prune proj sc)
387 notErased (DefaultCase (STerm Erased)) = False
388 notErased (DefaultCase ImpossibleCase) = False
389 notErased _ = True
391 mkProj n i [] sc = sc
392 mkProj n i (x : xs) sc = mkProj n (i + 1) xs (projRep x n i sc)
394 projRep :: Name -> Name -> Int -> SC -> SC
395 projRep arg n i (Case x alts)
396 | x == arg = ProjCase (Proj (P Bound n Erased) i)
397 (map (projRepAlt arg n i) alts)
398 | otherwise = Case x (map (projRepAlt arg n i) alts)
399 projRep arg n i (ProjCase t alts)
400 = ProjCase (projRepTm arg n i t) (map (projRepAlt arg n i) alts)
401 projRep arg n i (STerm t) = STerm (projRepTm arg n i t)
402 projRep arg n i c = c -- unmatched
404 projRepAlt arg n i (ConCase cn t args rhs)
405 = ConCase cn t args (projRep arg n i rhs)
406 projRepAlt arg n i (ConstCase t rhs)
407 = ConstCase t (projRep arg n i rhs)
408 projRepAlt arg n i (DefaultCase rhs)
409 = DefaultCase (projRep arg n i rhs)
411 projRepTm arg n i t = subst arg (Proj (P Bound n Erased) i) t
413prune _ t = t
415stripLambdas :: CaseDef -> CaseDef
416stripLambdas (CaseDef ns (STerm (Bind x (Lam _) sc)) tm)
417 = stripLambdas (CaseDef (ns ++ [x]) (STerm (instantiate (P Bound x Erased) sc)) tm)
418stripLambdas x = x