git.iko.soy/Idris-dev
on ff726fe help sign in
Idris-dev/src/Core/CaseTree.hs
554 lines 22 kB history raw
1{-# LANGUAGE PatternGuards, DeriveFunctor, TypeSynonymInstances #-}
3module Core.CaseTree(CaseDef(..), SC, SC'(..), CaseAlt, CaseAlt'(..),
4 Phase(..), CaseTree,
5 simpleCase, small, namesUsed, findCalls, findUsedArgs) where
7import Core.TT
9import Control.Monad.State
10import Data.Maybe
11import Data.List hiding (partition)
12import Debug.Trace
14data CaseDef = CaseDef [Name] SC [Term]
15 deriving Show
17data SC' t = Case Name [CaseAlt' t] -- ^ invariant: lowest tags first
18 | ProjCase t [CaseAlt' t] -- ^ special case for projections
19 | STerm t
20 | UnmatchedCase String -- ^ error message
21 | ImpossibleCase -- ^ already checked to be impossible
22 deriving (Eq, Ord, Functor)
23{-!
24deriving instance Binary SC
25!-}
27type SC = SC' Term
29data CaseAlt' t = ConCase Name Int [Name] (SC' t)
30 | FnCase Name [Name] (SC' t) -- ^ reflection function
31 | ConstCase Const (SC' t)
32 | SucCase Name (SC' t)
33 | DefaultCase (SC' t)
34 deriving (Show, Eq, Ord, Functor)
35{-!
36deriving instance Binary CaseAlt
37!-}
39type CaseAlt = CaseAlt' Term
41instance Show t => Show (SC' t) where
42 show sc = show' 1 sc
43 where
44 show' i (Case n alts) = "case " ++ show n ++ " of\n" ++ indent i ++
45 showSep ("\n" ++ indent i) (map (showA i) alts)
46 show' i (ProjCase tm alts) = "case " ++ show tm ++ " of " ++
47 showSep ("\n" ++ indent i) (map (showA i) alts)
48 show' i (STerm tm) = show tm
49 show' i (UnmatchedCase str) = "error " ++ show str
50 show' i ImpossibleCase = "impossible"
52 indent i = concat $ take i (repeat " ")
54 showA i (ConCase n t args sc)
55 = show n ++ "(" ++ showSep (", ") (map show args) ++ ") => "
56 ++ show' (i+1) sc
57 showA i (FnCase n args sc)
58 = "FN " ++ show n ++ "(" ++ showSep (", ") (map show args) ++ ") => "
59 ++ show' (i+1) sc
60 showA i (ConstCase t sc)
61 = show t ++ " => " ++ show' (i+1) sc
62 showA i (SucCase n sc)
63 = show n ++ "+1 => " ++ show' (i+1) sc
64 showA i (DefaultCase sc)
65 = "_ => " ++ show' (i+1) sc
68type CaseTree = SC
69type Clause = ([Pat], (Term, Term))
70type CS = ([Term], Int)
72instance TermSize SC where
73 termsize n (Case n' as) = termsize n as
74 termsize n (ProjCase n' as) = termsize n as
75 termsize n (STerm t) = termsize n t
76 termsize n _ = 1
78instance TermSize CaseAlt where
79 termsize n (ConCase _ _ _ s) = termsize n s
80 termsize n (FnCase _ _ s) = termsize n s
81 termsize n (ConstCase _ s) = termsize n s
82 termsize n (SucCase _ s) = termsize n s
83 termsize n (DefaultCase s) = termsize n s
85-- simple terms can be inlined trivially - good for primitives in particular
86-- To avoid duplicating work, don't inline something which uses one
87-- of its arguments in more than one place
89small :: Name -> [Name] -> SC -> Bool
90small n args t = let as = findAllUsedArgs t args in
91 length as == length (nub as) &&
92 termsize n t < 10
94namesUsed :: SC -> [Name]
95namesUsed sc = nub $ nu' [] sc where
96 nu' ps (Case n alts) = nub (concatMap (nua ps) alts) \\ [n]
97 nu' ps (ProjCase t alts) = nub $ (nut ps t ++
98 (concatMap (nua ps) alts))
99 nu' ps (STerm t) = nub $ nut ps t
100 nu' ps _ = []
102 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc) \\ args
103 nua ps (FnCase n args sc) = nub (nu' (ps ++ args) sc) \\ args
104 nua ps (ConstCase _ sc) = nu' ps sc
105 nua ps (SucCase _ sc) = nu' ps sc
106 nua ps (DefaultCase sc) = nu' ps sc
108 nut ps (P _ n _) | n `elem` ps = []
109 | otherwise = [n]
110 nut ps (App f a) = nut ps f ++ nut ps a
111 nut ps (Proj t _) = nut ps t
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-- Return all called functions, and which arguments are used in each argument position
117-- for the call, in order to help reduce compilation time, and trace all unused
118-- arguments
120findCalls :: SC -> [Name] -> [(Name, [[Name]])]
121findCalls sc topargs = nub $ nu' topargs sc where
122 nu' ps (Case n alts) = nub (concatMap (nua (n : ps)) alts)
123 nu' ps (ProjCase t alts) = nub (nut ps t ++ concatMap (nua ps) alts)
124 nu' ps (STerm t) = nub $ nut ps t
125 nu' ps _ = []
127 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc)
128 nua ps (FnCase n args sc) = nub (nu' (ps ++ args) sc)
129 nua ps (ConstCase _ sc) = nu' ps sc
130 nua ps (SucCase _ sc) = nu' ps sc
131 nua ps (DefaultCase sc) = nu' ps sc
133 nut ps (P Ref n _) | n `elem` ps = []
134 | otherwise = [(n, [])] -- tmp
135 nut ps fn@(App f a)
136 | (P Ref n _, args) <- unApply fn
137 = if n `elem` ps then nut ps f ++ nut ps a
138 else [(n, map argNames args)] ++ concatMap (nut ps) args
139 | otherwise = nut ps f ++ nut ps a
140 nut ps (Bind n (Let t v) sc) = nut ps v ++ nut (n:ps) sc
141 nut ps (Proj t _) = nut ps t
142 nut ps (Bind n b sc) = nut (n:ps) sc
143 nut ps _ = []
145 argNames tm = let ns = directUse tm in
146 filter (\x -> x `elem` ns) topargs
148-- Find names which are used directly (i.e. not in a function call) in a term
150directUse :: Eq n => TT n -> [n]
151directUse (P _ n _) = [n]
152directUse (Bind n (Let t v) sc) = nub $ directUse v ++ (directUse sc \\ [n])
153 ++ directUse t
154directUse (Bind n b sc) = nub $ directUse (binderTy b) ++ (directUse sc \\ [n])
155directUse fn@(App f a)
156 | (P Ref n _, args) <- unApply fn = [] -- need to know what n does with them
157 | otherwise = nub $ directUse f ++ directUse a
158directUse (Proj x i) = nub $ directUse x
159directUse _ = []
161-- Find all directly used arguments (i.e. used but not in function calls)
163findUsedArgs :: SC -> [Name] -> [Name]
164findUsedArgs sc topargs = nub (findAllUsedArgs sc topargs)
166findAllUsedArgs sc topargs = filter (\x -> x `elem` topargs) (nu' sc) where
167 nu' (Case n alts) = n : concatMap nua alts
168 nu' (ProjCase t alts) = directUse t ++ concatMap nua alts
169 nu' (STerm t) = directUse t
170 nu' _ = []
172 nua (ConCase n i args sc) = nu' sc
173 nua (FnCase n args sc) = nu' sc
174 nua (ConstCase _ sc) = nu' sc
175 nua (SucCase _ sc) = nu' sc
176 nua (DefaultCase sc) = nu' sc
178data Phase = CompileTime | RunTime
179 deriving (Show, Eq)
181-- Generate a simple case tree
182-- Work Right to Left
184simpleCase :: Bool -> Bool -> Bool ->
185 Phase -> FC -> [([Name], Term, Term)] ->
186 TC CaseDef
187simpleCase tc cover reflect phase fc cs
188 = sc' tc cover phase fc (filter (\(_, _, r) ->
189 case r of
190 Impossible -> False
191 _ -> True) cs)
192 where
193 sc' tc cover phase fc []
194 = return $ CaseDef [] (UnmatchedCase "No pattern clauses") []
195 sc' tc cover phase fc cs
196 = let proj = phase == RunTime
197 pats = map (\ (avs, l, r) ->
198 (avs, toPats reflect tc l, (l, r))) cs
199 chkPats = mapM chkAccessible pats in
200 case chkPats of
201 OK pats ->
202 let numargs = length (fst (head pats))
203 ns = take numargs args
204 (ns', ps') = order ns pats
205 (tree, st) = runState
206 (match ns' ps' (defaultCase cover)) ([], numargs)
207 t = CaseDef ns (prune proj (depatt ns' tree)) (fst st) in
208 if proj then return (stripLambdas t) else return t
209 Error err -> Error (At fc err)
210 where args = map (\i -> MN i "e") [0..]
211 defaultCase True = STerm Erased
212 defaultCase False = UnmatchedCase "Error"
214 chkAccessible (avs, l, c)
215 | phase == RunTime || reflect = return (l, c)
216 | otherwise = do mapM_ (acc l) avs
217 return (l, c)
219 acc [] n = Error (Inaccessible n)
220 acc (PV x : xs) n | x == n = OK ()
221 acc (PCon _ _ ps : xs) n = acc (ps ++ xs) n
222 acc (PSuc p : xs) n = acc (p : xs) n
223 acc (_ : xs) n = acc xs n
225data Pat = PCon Name Int [Pat]
226 | PConst Const
227 | PV Name
228 | PSuc Pat -- special case for n+1 on Integer
229 | PReflected Name [Pat]
230 | PAny
231 deriving Show
233-- If there are repeated variables, take the *last* one (could be name shadowing
234-- in a where clause, so take the most recent).
236toPats :: Bool -> Bool -> Term -> [Pat]
237toPats reflect tc f = reverse (toPat reflect tc (getArgs f)) where
238 getArgs (App f a) = a : getArgs f
239 getArgs _ = []
241toPat :: Bool -> Bool -> [Term] -> [Pat]
242toPat reflect tc tms = evalState (mapM (\x -> toPat' x []) tms) []
243 where
244 toPat' (P (DCon t a) n _) args = do args' <- mapM (\x -> toPat' x []) args
245 return $ PCon n t args'
246 -- n + 1
247 toPat' (P _ (UN "prim__addBigInt") _)
248 [p, Constant (BI 1)]
249 = do p' <- toPat' p []
250 return $ PSuc p'
251 -- Typecase
252 toPat' (P (TCon t a) n _) args | tc
253 = do args' <- mapM (\x -> toPat' x []) args
254 return $ PCon n t args'
255 toPat' (Constant (AType (ATInt ITNative))) []
256 | tc = return $ PCon (UN "Int") 1 []
257 toPat' (Constant (AType ATFloat)) [] | tc = return $ PCon (UN "Float") 2 []
258 toPat' (Constant (AType (ATInt ITChar))) [] | tc = return $ PCon (UN "Char") 3 []
259 toPat' (Constant StrType) [] | tc = return $ PCon (UN "String") 4 []
260 toPat' (Constant PtrType) [] | tc = return $ PCon (UN "Ptr") 5 []
261 toPat' (Constant (AType (ATInt ITBig))) []
262 | tc = return $ PCon (UN "Integer") 6 []
263 toPat' (Constant (AType (ATInt (ITFixed n)))) []
264 | tc = return $ PCon (UN (fixedN n)) (7 + fromEnum n) [] -- 7-10 inclusive
265 toPat' (P Bound n _) [] = do ns <- get
266 if n `elem` ns
267 then return PAny
268 else do put (n : ns)
269 return (PV n)
270 toPat' (App f a) args = toPat' f (a : args)
271 toPat' (Constant x) [] = return $ PConst x
272 toPat' (Bind n (Pi t) sc) [] | reflect && noOccurrence n sc
273 = do t' <- toPat' t []
274 sc' <- toPat' sc []
275 return $ PReflected (UN "->") (t':sc':[])
276 toPat' (P _ n _) args | reflect
277 = do args' <- mapM (\x -> toPat' x []) args
278 return $ PReflected n args'
279 toPat' _ _ = return PAny
281 fixedN IT8 = "Bits8"
282 fixedN IT16 = "Bits16"
283 fixedN IT32 = "Bits32"
284 fixedN IT64 = "Bits64"
287data Partition = Cons [Clause]
288 | Vars [Clause]
289 deriving Show
291isVarPat (PV _ : ps , _) = True
292isVarPat (PAny : ps , _) = True
293isVarPat _ = False
295isConPat (PCon _ _ _ : ps, _) = True
296isConPat (PReflected _ _ : ps, _) = True
297isConPat (PSuc _ : ps, _) = True
298isConPat (PConst _ : ps, _) = True
299isConPat _ = False
301partition :: [Clause] -> [Partition]
302partition [] = []
303partition ms@(m : _)
304 | isVarPat m = let (vars, rest) = span isVarPat ms in
305 Vars vars : partition rest
306 | isConPat m = let (cons, rest) = span isConPat ms in
307 Cons cons : partition rest
308partition xs = error $ "Partition " ++ show xs
310-- reorder the patterns so that the one with most distinct names
311-- comes next. Take rightmost first, otherwise (i.e. pick value rather
312-- than dependency)
314order :: [Name] -> [Clause] -> ([Name], [Clause])
315order [] cs = ([], cs)
316order ns [] = (ns, [])
317order ns cs = let patnames = transpose (map (zip ns) (map fst cs))
318 pats' = transpose (sortBy moreDistinct (reverse patnames)) in
319 (getNOrder pats', zipWith rebuild pats' cs)
320 where
321 getNOrder [] = error $ "Failed order on " ++ show (ns, cs)
322 getNOrder (c : _) = map fst c
324 rebuild patnames clause = (map snd patnames, snd clause)
326 moreDistinct xs ys = compare (numNames [] (map snd ys))
327 (numNames [] (map snd xs))
329 numNames xs (PCon n _ _ : ps)
330 | not (Left n `elem` xs) = numNames (Left n : xs) ps
331 numNames xs (PConst c : ps)
332 | not (Right c `elem` xs) = numNames (Right c : xs) ps
333 numNames xs (_ : ps) = numNames xs ps
334 numNames xs [] = length xs
336match :: [Name] -> [Clause] -> SC -- error case
337 -> State CS SC
338match [] (([], ret) : xs) err
339 = do (ts, v) <- get
340 put (ts ++ (map (fst.snd) xs), v)
341 case snd ret of
342 Impossible -> return ImpossibleCase
343 tm -> return $ STerm tm -- run out of arguments
344match vs cs err = do let ps = partition cs
345 mixture vs ps err
347mixture :: [Name] -> [Partition] -> SC -> State CS SC
348mixture vs [] err = return err
349mixture vs (Cons ms : ps) err = do fallthrough <- mixture vs ps err
350 conRule vs ms fallthrough
351mixture vs (Vars ms : ps) err = do fallthrough <- mixture vs ps err
352 varRule vs ms fallthrough
354data ConType = CName Name Int -- named constructor
355 | CFn Name -- reflected function name
356 | CSuc -- n+1
357 | CConst Const -- constant, not implemented yet
358 deriving (Show, Eq)
360data Group = ConGroup ConType -- Constructor
361 [([Pat], Clause)] -- arguments and rest of alternative
362 deriving Show
364conRule :: [Name] -> [Clause] -> SC -> State CS SC
365conRule (v:vs) cs err = do groups <- groupCons cs
366 caseGroups (v:vs) groups err
368caseGroups :: [Name] -> [Group] -> SC -> State CS SC
369caseGroups (v:vs) gs err = do g <- altGroups gs
370 return $ Case v (sort g)
371 where
372 altGroups [] = return [DefaultCase err]
373 altGroups (ConGroup (CName n i) args : cs)
374 = do g <- altGroup n i args
375 rest <- altGroups cs
376 return (g : rest)
377 altGroups (ConGroup (CFn n) args : cs)
378 = do g <- altFnGroup n args
379 rest <- altGroups cs
380 return (g : rest)
381 altGroups (ConGroup CSuc args : cs)
382 = do g <- altSucGroup args
383 rest <- altGroups cs
384 return (g : rest)
385 altGroups (ConGroup (CConst c) args : cs)
386 = do g <- altConstGroup c args
387 rest <- altGroups cs
388 return (g : rest)
390 altGroup n i gs = do (newArgs, nextCs) <- argsToAlt gs
391 matchCs <- match (newArgs ++ vs) nextCs err
392 return $ ConCase n i newArgs matchCs
393 altFnGroup n gs = do (newArgs, nextCs) <- argsToAlt gs
394 matchCs <- match (newArgs ++ vs) nextCs err
395 return $ FnCase n newArgs matchCs
396 altSucGroup gs = do ([newArg], nextCs) <- argsToAlt gs
397 matchCs <- match (newArg:vs) nextCs err
398 return $ SucCase newArg matchCs
399 altConstGroup n gs = do (_, nextCs) <- argsToAlt gs
400 matchCs <- match vs nextCs err
401 return $ ConstCase n matchCs
403argsToAlt :: [([Pat], Clause)] -> State CS ([Name], [Clause])
404argsToAlt [] = return ([], [])
405argsToAlt rs@((r, m) : rest)
406 = do newArgs <- getNewVars r
407 return (newArgs, addRs rs)
408 where
409 getNewVars [] = return []
410 getNewVars ((PV n) : ns) = do v <- getVar "e"
411 nsv <- getNewVars ns
412 return (v : nsv)
413 getNewVars (PAny : ns) = do v <- getVar "i"
414 nsv <- getNewVars ns
415 return (v : nsv)
416 getNewVars (_ : ns) = do v <- getVar "e"
417 nsv <- getNewVars ns
418 return (v : nsv)
419 addRs [] = []
420 addRs ((r, (ps, res)) : rs) = ((r++ps, res) : addRs rs)
422 uniq i (UN n) = MN i n
423 uniq i n = n
425getVar :: String -> State CS Name
426getVar b = do (t, v) <- get; put (t, v+1); return (MN v b)
428groupCons :: [Clause] -> State CS [Group]
429groupCons cs = gc [] cs
430 where
431 gc acc [] = return acc
432 gc acc ((p : ps, res) : cs) =
433 do acc' <- addGroup p ps res acc
434 gc acc' cs
435 addGroup p ps res acc = case p of
436 PCon con i args -> return $ addg (CName con i) args (ps, res) acc
437 PConst cval -> return $ addConG cval (ps, res) acc
438 PSuc n -> return $ addg CSuc [n] (ps, res) acc
439 PReflected fn args -> return $ addg (CFn fn) args (ps, res) acc
440 pat -> fail $ show pat ++ " is not a constructor or constant (can't happen)"
442 addg c conargs res []
443 = [ConGroup c [(conargs, res)]]
444 addg c conargs res (g@(ConGroup c' cs):gs)
445 | c == c' = ConGroup c (cs ++ [(conargs, res)]) : gs
446 | otherwise = g : addg c conargs res gs
448 addConG con res [] = [ConGroup (CConst con) [([], res)]]
449 addConG con res (g@(ConGroup (CConst n) cs) : gs)
450 | con == n = ConGroup (CConst n) (cs ++ [([], res)]) : gs
451-- | otherwise = g : addConG con res gs
452 addConG con res (g : gs) = g : addConG con res gs
454varRule :: [Name] -> [Clause] -> SC -> State CS SC
455varRule (v : vs) alts err =
456 do let alts' = map (repVar v) alts
457 match vs alts' err
458 where
459 repVar v (PV p : ps , (lhs, res))
460 = (ps, (lhs, subst p (P Bound v Erased) res))
461 repVar v (PAny : ps , res) = (ps, res)
463-- fix: case e of S k -> f (S k) ==> case e of S k -> f e
465depatt :: [Name] -> SC -> SC
466depatt ns tm = dp [] tm
467 where
468 dp ms (STerm tm) = STerm (applyMaps ms tm)
469 dp ms (Case x alts) = Case x (map (dpa ms x) alts)
470 dp ms sc = sc
472 dpa ms x (ConCase n i args sc)
473 = ConCase n i args (dp ((x, (n, args)) : ms) sc)
474 dpa ms x (FnCase n args sc)
475 = FnCase n args (dp ((x, (n, args)) : ms) sc)
476 dpa ms x (ConstCase c sc) = ConstCase c (dp ms sc)
477 dpa ms x (SucCase n sc) = SucCase n (dp ms sc)
478 dpa ms x (DefaultCase sc) = DefaultCase (dp ms sc)
480 applyMaps ms f@(App _ _)
481 | (P nt cn pty, args) <- unApply f
482 = let args' = map (applyMaps ms) args in
483 applyMap ms nt cn pty args'
484 where
485 applyMap [] nt cn pty args' = mkApp (P nt cn pty) args'
486 applyMap ((x, (n, args)) : ms) nt cn pty args'
487 | and ((length args == length args') :
488 (n == cn) : zipWith same args args') = P Ref x Erased
489 | otherwise = applyMap ms nt cn pty args'
490 same n (P _ n' _) = n == n'
491 same _ _ = False
493 applyMaps ms (App f a) = App (applyMaps ms f) (applyMaps ms a)
494 applyMaps ms t = t
496prune :: Bool -- ^ Convert single branches to projections (only useful at runtime)
497 -> SC -> SC
498prune proj (Case n alts)
499 = let alts' = filter notErased (map pruneAlt alts) in
500 case alts' of
501 [] -> ImpossibleCase
502 as@[ConCase cn i args sc] -> if proj then mkProj n 0 args sc
503 else Case n as
504 -- Bit of a hack here! The default case will always be 0, make sure
505 -- it gets caught first.
506 [s@(SucCase _ _), DefaultCase dc]
507 -> Case n [ConstCase (BI 0) dc, s]
508 as -> Case n as
509 where pruneAlt (ConCase cn i ns sc) = ConCase cn i ns (prune proj sc)
510 pruneAlt (FnCase cn ns sc) = FnCase cn ns (prune proj sc)
511 pruneAlt (ConstCase c sc) = ConstCase c (prune proj sc)
512 pruneAlt (SucCase n sc) = SucCase n (prune proj sc)
513 pruneAlt (DefaultCase sc) = DefaultCase (prune proj sc)
515 notErased (DefaultCase (STerm Erased)) = False
516 notErased (DefaultCase ImpossibleCase) = False
517 notErased _ = True
519 mkProj n i [] sc = sc
520 mkProj n i (x : xs) sc = mkProj n (i + 1) xs (projRep x n i sc)
522 projRep :: Name -> Name -> Int -> SC -> SC
523 projRep arg n i (Case x alts)
524 | x == arg = ProjCase (Proj (P Bound n Erased) i)
525 (map (projRepAlt arg n i) alts)
526 | otherwise = Case x (map (projRepAlt arg n i) alts)
527 projRep arg n i (ProjCase t alts)
528 = ProjCase (projRepTm arg n i t) (map (projRepAlt arg n i) alts)
529 projRep arg n i (STerm t) = STerm (projRepTm arg n i t)
530 projRep arg n i c = c -- unmatched
532 projRepAlt arg n i (ConCase cn t args rhs)
533 = ConCase cn t args (projRep arg n i rhs)
534 projRepAlt arg n i (FnCase cn args rhs)
535 = FnCase cn args (projRep arg n i rhs)
536 projRepAlt arg n i (ConstCase t rhs)
537 = ConstCase t (projRep arg n i rhs)
538 projRepAlt arg n i (SucCase sn rhs)
539 = SucCase sn (projRep arg n i rhs)
540 projRepAlt arg n i (DefaultCase rhs)
541 = DefaultCase (projRep arg n i rhs)
543 projRepTm arg n i t = subst arg (Proj (P Bound n Erased) i) t
545prune _ t = t
547stripLambdas :: CaseDef -> CaseDef
548stripLambdas (CaseDef ns (STerm (Bind x (Lam _) sc)) tm)
549 = stripLambdas (CaseDef (ns ++ [x]) (STerm (instantiate (P Bound x Erased) sc)) tm)
550stripLambdas x = x