git.iko.soy/Idris-dev
on 3c0f4dc help sign in
Idris-dev/src/Core/CaseTree.hs
488 lines 19 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 | ConstCase Const (SC' t)
31 | DefaultCase (SC' t)
32 deriving (Show, Eq, Ord, Functor)
33{-!
34deriving instance Binary CaseAlt
35!-}
37type CaseAlt = CaseAlt' Term
39instance Show t => Show (SC' t) where
40 show sc = show' 1 sc
41 where
42 show' i (Case n alts) = "case " ++ show n ++ " of\n" ++ indent i ++
43 showSep ("\n" ++ indent i) (map (showA i) alts)
44 show' i (ProjCase tm alts) = "case " ++ show tm ++ " of " ++
45 showSep ("\n" ++ indent i) (map (showA i) alts)
46 show' i (STerm tm) = show tm
47 show' i (UnmatchedCase str) = "error " ++ show str
48 show' i ImpossibleCase = "impossible"
50 indent i = concat $ take i (repeat " ")
52 showA i (ConCase n t args sc)
53 = show n ++ "(" ++ showSep (", ") (map show args) ++ ") => "
54 ++ show' (i+1) sc
55 showA i (ConstCase t sc)
56 = show t ++ " => " ++ show' (i+1) sc
57 showA i (DefaultCase sc)
58 = "_ => " ++ show' (i+1) sc
61type CaseTree = SC
62type Clause = ([Pat], (Term, Term))
63type CS = ([Term], Int)
65instance TermSize SC where
66 termsize n (Case n' as) = termsize n as
67 termsize n (ProjCase n' as) = termsize n as
68 termsize n (STerm t) = termsize n t
69 termsize n _ = 1
71instance TermSize CaseAlt where
72 termsize n (ConCase _ _ _ s) = termsize n s
73 termsize n (ConstCase _ s) = termsize n s
74 termsize n (DefaultCase s) = termsize n s
76-- simple terms can be inlined trivially - good for primitives in particular
77-- To avoid duplicating work, don't inline something which uses one
78-- of its arguments in more than one place
80small :: Name -> [Name] -> SC -> Bool
81small n args t = let as = findAllUsedArgs t args in
82 length as == length (nub as) &&
83 termsize n t < 10
85namesUsed :: SC -> [Name]
86namesUsed sc = nub $ nu' [] sc where
87 nu' ps (Case n alts) = nub (concatMap (nua ps) alts) \\ [n]
88 nu' ps (ProjCase t alts) = nub $ (nut ps t ++
89 (concatMap (nua ps) alts))
90 nu' ps (STerm t) = nub $ nut ps t
91 nu' ps _ = []
93 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc) \\ args
94 nua ps (ConstCase _ sc) = nu' ps sc
95 nua ps (DefaultCase sc) = nu' ps sc
97 nut ps (P _ n _) | n `elem` ps = []
98 | otherwise = [n]
99 nut ps (App f a) = nut ps f ++ nut ps a
100 nut ps (Proj t _) = nut ps t
101 nut ps (Bind n (Let t v) sc) = nut ps v ++ nut (n:ps) sc
102 nut ps (Bind n b sc) = nut (n:ps) sc
103 nut ps _ = []
105-- Return all called functions, and which arguments are used in each argument position
106-- for the call, in order to help reduce compilation time, and trace all unused
107-- arguments
109findCalls :: SC -> [Name] -> [(Name, [[Name]])]
110findCalls sc topargs = nub $ nu' topargs sc where
111 nu' ps (Case n alts) = nub (concatMap (nua (n : ps)) alts)
112 nu' ps (ProjCase t alts) = nub (nut ps t ++ concatMap (nua ps) alts)
113 nu' ps (STerm t) = nub $ nut ps t
114 nu' ps _ = []
116 nua ps (ConCase n i args sc) = nub (nu' (ps ++ args) sc)
117 nua ps (ConstCase _ sc) = nu' ps sc
118 nua ps (DefaultCase sc) = nu' ps sc
120 nut ps (P Ref n _) | n `elem` ps = []
121 | otherwise = [(n, [])] -- tmp
122 nut ps fn@(App f a)
123 | (P Ref n _, args) <- unApply fn
124 = if n `elem` ps then nut ps f ++ nut ps a
125 else [(n, map argNames args)] ++ concatMap (nut ps) args
126 | otherwise = nut ps f ++ nut ps a
127 nut ps (Bind n (Let t v) sc) = nut ps v ++ nut (n:ps) sc
128 nut ps (Proj t _) = nut ps t
129 nut ps (Bind n b sc) = nut (n:ps) sc
130 nut ps _ = []
132 argNames tm = let ns = directUse tm in
133 filter (\x -> x `elem` ns) topargs
135-- Find names which are used directly (i.e. not in a function call) in a term
137directUse :: Eq n => TT n -> [n]
138directUse (P _ n _) = [n]
139directUse (Bind n (Let t v) sc) = nub $ directUse v ++ (directUse sc \\ [n])
140 ++ directUse t
141directUse (Bind n b sc) = nub $ directUse (binderTy b) ++ (directUse sc \\ [n])
142directUse fn@(App f a)
143 | (P Ref n _, args) <- unApply fn = [] -- need to know what n does with them
144 | otherwise = nub $ directUse f ++ directUse a
145directUse (Proj x i) = nub $ directUse x
146directUse _ = []
148-- Find all directly used arguments (i.e. used but not in function calls)
150findUsedArgs :: SC -> [Name] -> [Name]
151findUsedArgs sc topargs = nub (findAllUsedArgs sc topargs)
153findAllUsedArgs sc topargs = filter (\x -> x `elem` topargs) (nu' sc) where
154 nu' (Case n alts) = n : concatMap nua alts
155 nu' (ProjCase t alts) = directUse t ++ concatMap nua alts
156 nu' (STerm t) = directUse t
157 nu' _ = []
159 nua (ConCase n i args sc) = nu' sc
160 nua (ConstCase _ sc) = nu' sc
161 nua (DefaultCase sc) = nu' sc
163data Phase = CompileTime | RunTime
164 deriving (Show, Eq)
166-- Generate a simple case tree
167-- Work Right to Left
169simpleCase :: Bool -> Bool -> Phase -> FC -> [([Name], Term, Term)] ->
170 TC CaseDef
171simpleCase tc cover phase fc cs
172 = sc' tc cover phase fc (filter (\(_, _, r) ->
173 case r of
174 Impossible -> False
175 _ -> True) cs)
176 where
177 sc' tc cover phase fc []
178 = return $ CaseDef [] (UnmatchedCase "No pattern clauses") []
179 sc' tc cover phase fc cs
180 = let proj = phase == RunTime
181 pats = map (\ (avs, l, r) ->
182 (avs, toPats tc l, (l, r))) cs
183 chkPats = mapM chkAccessible pats in
184 case chkPats of
185 OK pats ->
186 let numargs = length (fst (head pats))
187 ns = take numargs args
188 (ns', ps') = order ns pats
189 (tree, st) = runState
190 (match ns' ps' (defaultCase cover)) ([], numargs)
191 t = CaseDef ns (prune proj (depatt ns' tree)) (fst st) in
192 if proj then return (stripLambdas t) else return t
193 Error err -> Error (At fc err)
194 where args = map (\i -> MN i "e") [0..]
195 defaultCase True = STerm Erased
196 defaultCase False = UnmatchedCase "Error"
198 chkAccessible (avs, l, c)
199 | phase == RunTime = return (l, c)
200 | otherwise = do mapM_ (acc l) avs
201 return (l, c)
203 acc [] n = Error (Inaccessible n)
204 acc (PV x : xs) n | x == n = OK ()
205 acc (PCon _ _ ps : xs) n = acc (ps ++ xs) n
206 acc (_ : xs) n = acc xs n
208data Pat = PCon Name Int [Pat]
209 | PConst Const
210 | PV Name
211 | PAny
212 deriving Show
214-- If there are repeated variables, take the *last* one (could be name shadowing
215-- in a where clause, so take the most recent).
217toPats :: Bool -> Term -> [Pat]
218toPats tc f = reverse (toPat tc (getArgs f)) where
219 getArgs (App f a) = a : getArgs f
220 getArgs _ = []
222toPat :: Bool -> [Term] -> [Pat]
223toPat tc tms = evalState (mapM (\x -> toPat' x []) tms) []
224 where
225 toPat' (P (DCon t a) n _) args = do args' <- mapM (\x -> toPat' x []) args
226 return $ PCon n t args'
227 -- Typecase
228 toPat' (P (TCon t a) n _) args | tc
229 = do args' <- mapM (\x -> toPat' x []) args
230 return $ PCon n t args'
231 toPat' (Constant (AType (ATInt ITNative))) []
232 | tc = return $ PCon (UN "Int") 1 []
233 toPat' (Constant (AType ATFloat)) [] | tc = return $ PCon (UN "Float") 2 []
234 toPat' (Constant ChType) [] | tc = return $ PCon (UN "Char") 3 []
235 toPat' (Constant StrType) [] | tc = return $ PCon (UN "String") 4 []
236 toPat' (Constant PtrType) [] | tc = return $ PCon (UN "Ptr") 5 []
237 toPat' (Constant (AType (ATInt ITBig))) []
238 | tc = return $ PCon (UN "Integer") 6 []
239 toPat' (Constant (AType (ATInt (ITFixed n)))) []
240 | tc = return $ PCon (UN (fixedN n)) (7 + fromEnum n) [] -- 7-10 inclusive
241 toPat' (Constant (AType (ATInt (ITVec ity count)))) []
242 | tc = return $ PCon (UN (fixedN ity)) ((fromEnum ity + 1) * 1000 + count) [] -- 1000-5000 inclusive
243 toPat' (P Bound n _) [] = do ns <- get
244 if n `elem` ns
245 then return PAny
246 else do put (n : ns)
247 return (PV n)
248 toPat' (App f a) args = toPat' f (a : args)
249 toPat' (Constant x) [] = return $ PConst x
250 toPat' _ _ = return PAny
252 fixedN IT8 = "Bits8"
253 fixedN IT16 = "Bits16"
254 fixedN IT32 = "Bits32"
255 fixedN IT64 = "Bits64"
258data Partition = Cons [Clause]
259 | Vars [Clause]
260 deriving Show
262isVarPat (PV _ : ps , _) = True
263isVarPat (PAny : ps , _) = True
264isVarPat _ = False
266isConPat (PCon _ _ _ : ps, _) = True
267isConPat (PConst _ : ps, _) = True
268isConPat _ = False
270partition :: [Clause] -> [Partition]
271partition [] = []
272partition ms@(m : _)
273 | isVarPat m = let (vars, rest) = span isVarPat ms in
274 Vars vars : partition rest
275 | isConPat m = let (cons, rest) = span isConPat ms in
276 Cons cons : partition rest
277partition xs = error $ "Partition " ++ show xs
279-- reorder the patterns so that the one with most distinct names
280-- comes next. Take rightmost first, otherwise (i.e. pick value rather
281-- than dependency)
283order :: [Name] -> [Clause] -> ([Name], [Clause])
284order [] cs = ([], cs)
285order ns [] = (ns, [])
286order ns cs = let patnames = transpose (map (zip ns) (map fst cs))
287 pats' = transpose (sortBy moreDistinct (reverse patnames)) in
288 (getNOrder pats', zipWith rebuild pats' cs)
289 where
290 getNOrder [] = error $ "Failed order on " ++ show (ns, cs)
291 getNOrder (c : _) = map fst c
293 rebuild patnames clause = (map snd patnames, snd clause)
295 moreDistinct xs ys = compare (numNames [] (map snd ys))
296 (numNames [] (map snd xs))
298 numNames xs (PCon n _ _ : ps)
299 | not (Left n `elem` xs) = numNames (Left n : xs) ps
300 numNames xs (PConst c : ps)
301 | not (Right c `elem` xs) = numNames (Right c : xs) ps
302 numNames xs (_ : ps) = numNames xs ps
303 numNames xs [] = length xs
305match :: [Name] -> [Clause] -> SC -- error case
306 -> State CS SC
307match [] (([], ret) : xs) err
308 = do (ts, v) <- get
309 put (ts ++ (map (fst.snd) xs), v)
310 case snd ret of
311 Impossible -> return ImpossibleCase
312 tm -> return $ STerm tm -- run out of arguments
313match vs cs err = do let ps = partition cs
314 mixture vs ps err
316mixture :: [Name] -> [Partition] -> SC -> State CS SC
317mixture vs [] err = return err
318mixture vs (Cons ms : ps) err = do fallthrough <- mixture vs ps err
319 conRule vs ms fallthrough
320mixture vs (Vars ms : ps) err = do fallthrough <- mixture vs ps err
321 varRule vs ms fallthrough
323data ConType = CName Name Int -- named constructor
324 | CConst Const -- constant, not implemented yet
326data Group = ConGroup ConType -- Constructor
327 [([Pat], Clause)] -- arguments and rest of alternative
329conRule :: [Name] -> [Clause] -> SC -> State CS SC
330conRule (v:vs) cs err = do groups <- groupCons cs
331 caseGroups (v:vs) groups err
333caseGroups :: [Name] -> [Group] -> SC -> State CS SC
334caseGroups (v:vs) gs err = do g <- altGroups gs
335 return $ Case v (sort g)
336 where
337 altGroups [] = return [DefaultCase err]
338 altGroups (ConGroup (CName n i) args : cs)
339 = do g <- altGroup n i args
340 rest <- altGroups cs
341 return (g : rest)
342 altGroups (ConGroup (CConst c) args : cs)
343 = do g <- altConstGroup c args
344 rest <- altGroups cs
345 return (g : rest)
347 altGroup n i gs = do (newArgs, nextCs) <- argsToAlt gs
348 matchCs <- match (newArgs ++ vs) nextCs err
349 return $ ConCase n i newArgs matchCs
350 altConstGroup n gs = do (_, nextCs) <- argsToAlt gs
351 matchCs <- match vs nextCs err
352 return $ ConstCase n matchCs
354argsToAlt :: [([Pat], Clause)] -> State CS ([Name], [Clause])
355argsToAlt [] = return ([], [])
356argsToAlt rs@((r, m) : rest)
357 = do newArgs <- getNewVars r
358 return (newArgs, addRs rs)
359 where
360 getNewVars [] = return []
361 getNewVars ((PV n) : ns) = do v <- getVar "e"
362 nsv <- getNewVars ns
363 return (v : nsv)
364 getNewVars (PAny : ns) = do v <- getVar "i"
365 nsv <- getNewVars ns
366 return (v : nsv)
367 getNewVars (_ : ns) = do v <- getVar "e"
368 nsv <- getNewVars ns
369 return (v : nsv)
370 addRs [] = []
371 addRs ((r, (ps, res)) : rs) = ((r++ps, res) : addRs rs)
373 uniq i (UN n) = MN i n
374 uniq i n = n
376getVar :: String -> State CS Name
377getVar b = do (t, v) <- get; put (t, v+1); return (MN v b)
379groupCons :: [Clause] -> State CS [Group]
380groupCons cs = gc [] cs
381 where
382 gc acc [] = return acc
383 gc acc ((p : ps, res) : cs) =
384 do acc' <- addGroup p ps res acc
385 gc acc' cs
386 addGroup p ps res acc = case p of
387 PCon con i args -> return $ addg con i args (ps, res) acc
388 PConst cval -> return $ addConG cval (ps, res) acc
389 pat -> fail $ show pat ++ " is not a constructor or constant (can't happen)"
391 addg con i conargs res [] = [ConGroup (CName con i) [(conargs, res)]]
392 addg con i conargs res (g@(ConGroup (CName n j) cs):gs)
393 | i == j = ConGroup (CName n i) (cs ++ [(conargs, res)]) : gs
394 | otherwise = g : addg con i conargs res gs
396 addConG con res [] = [ConGroup (CConst con) [([], res)]]
397 addConG con res (g@(ConGroup (CConst n) cs) : gs)
398 | con == n = ConGroup (CConst n) (cs ++ [([], res)]) : gs
399 | otherwise = g : addConG con res gs
401varRule :: [Name] -> [Clause] -> SC -> State CS SC
402varRule (v : vs) alts err =
403 do let alts' = map (repVar v) alts
404 match vs alts' err
405 where
406 repVar v (PV p : ps , (lhs, res))
407 = (ps, (lhs, subst p (P Bound v Erased) res))
408 repVar v (PAny : ps , res) = (ps, res)
410-- fix: case e of S k -> f (S k) ==> case e of S k -> f e
412depatt :: [Name] -> SC -> SC
413depatt ns tm = dp [] tm
414 where
415 dp ms (STerm tm) = STerm (applyMaps ms tm)
416 dp ms (Case x alts) = Case x (map (dpa ms x) alts)
417 dp ms sc = sc
419 dpa ms x (ConCase n i args sc)
420 = ConCase n i args (dp ((x, (n, args)) : ms) sc)
421 dpa ms x (ConstCase c sc) = ConstCase c (dp ms sc)
422 dpa ms x (DefaultCase sc) = DefaultCase (dp ms sc)
424 applyMaps ms f@(App _ _)
425 | (P nt cn pty, args) <- unApply f
426 = let args' = map (applyMaps ms) args in
427 applyMap ms nt cn pty args'
428 where
429 applyMap [] nt cn pty args' = mkApp (P nt cn pty) args'
430 applyMap ((x, (n, args)) : ms) nt cn pty args'
431 | and ((length args == length args') :
432 (n == cn) : zipWith same args args') = P Ref x Erased
433 | otherwise = applyMap ms nt cn pty args'
434 same n (P _ n' _) = n == n'
435 same _ _ = False
437 applyMaps ms (App f a) = App (applyMaps ms f) (applyMaps ms a)
438 applyMaps ms t = t
440prune :: Bool -- ^ Convert single branches to projections (only useful at runtime)
441 -> SC -> SC
442prune proj (Case n alts)
443 = let alts' = filter notErased (map pruneAlt alts) in
444 case alts' of
445 [] -> ImpossibleCase
446 as@[ConCase cn i args sc] -> if proj then mkProj n 0 args sc
447 else Case n as
448 as -> Case n as
449 where pruneAlt (ConCase cn i ns sc) = ConCase cn i ns (prune proj sc)
450 pruneAlt (ConstCase c sc) = ConstCase c (prune proj sc)
451 pruneAlt (DefaultCase sc) = DefaultCase (prune proj sc)
453 notErased (DefaultCase (STerm Erased)) = False
454 notErased (DefaultCase ImpossibleCase) = False
455 notErased _ = True
457 mkProj n i [] sc = sc
458 mkProj n i (x : xs) sc = mkProj n (i + 1) xs (projRep x n i sc)
460 projRep :: Name -> Name -> Int -> SC -> SC
461 projRep arg n i (Case x alts)
462 | x == arg = ProjCase (Proj (P Bound n Erased) i)
463 (map (projRepAlt arg n i) alts)
464 | otherwise = Case x (map (projRepAlt arg n i) alts)
465 projRep arg n i (ProjCase t alts)
466 = ProjCase (projRepTm arg n i t) (map (projRepAlt arg n i) alts)
467 projRep arg n i (STerm t) = STerm (projRepTm arg n i t)
468 projRep arg n i c = c -- unmatched
470 projRepAlt arg n i (ConCase cn t args rhs)
471 = ConCase cn t args (projRep arg n i rhs)
472 projRepAlt arg n i (ConstCase t rhs)
473 = ConstCase t (projRep arg n i rhs)
474 projRepAlt arg n i (DefaultCase rhs)
475 = DefaultCase (projRep arg n i rhs)
477 projRepTm arg n i t = subst arg (Proj (P Bound n Erased) i) t
479prune _ t = t
481stripLambdas :: CaseDef -> CaseDef
482stripLambdas (CaseDef ns (STerm (Bind x (Lam _) sc)) tm)
483 = stripLambdas (CaseDef (ns ++ [x]) (STerm (instantiate (P Bound x Erased) sc)) tm)
484stripLambdas x = x