git.iko.soy/Idris-dev
on aa8aefd help sign in
Idris-dev/src/Idris/DataOpts.hs
263 lines 10 kB history raw
1{-# LANGUAGE PatternGuards #-}
3module Idris.DataOpts where
5-- Forcing, detagging and collapsing
7import Idris.AbsSyntax
8import Core.TT
10import Data.List
11import Data.Maybe
12import Debug.Trace
14-- Calculate the forceable arguments to a constructor and update the set of
15-- optimisations
17forceArgs :: Name -> Type -> Idris ()
18forceArgs n t = do i <- getIState
19 let fargs = force i 0 t
20 copt <- case lookupCtxt Nothing n (idris_optimisation i) of
21 [] -> return $ Optimise False [] []
22 (op:_) -> return op
23 let opts = addDef n (copt { forceable = fargs }) (idris_optimisation i)
24 putIState (i { idris_optimisation = opts })
25 addIBC (IBCOpt n)
26 iLOG $ "Forced: " ++ show n ++ " " ++ show fargs ++ "\n from " ++
27 show t
28 where
29 force :: IState -> Int -> Term -> [Int]
30 force ist i (Bind _ (Pi ty) sc)
31 | collapsibleIn ist ty
32 = nub $ i : (force ist (i + 1) $ instantiate (P Bound (MN i "?") Erased) sc)
33 | otherwise = force ist (i + 1) $ instantiate (P Bound (MN i "?") Erased) sc
34 force _ _ sc@(App f a)
35 | (_, args) <- unApply sc
36 = nub $ concatMap guarded args
37 force _ _ _ = []
39 collapsibleIn i t
40 | (P _ tn _, _) <- unApply t
41 = case lookupCtxt Nothing tn (idris_optimisation i) of
42 [oi] -> collapsible oi
43 _ -> False
44 | otherwise = False
46 isF (P _ (MN force "?") _) = Just force
47 isF _ = Nothing
49 guarded :: Term -> [Int]
50 guarded t@(App f a)
51 | (P (TCon _ _) _ _, args) <- unApply t
52 = mapMaybe isF args ++ concatMap guarded args
53 | (P (DCon _ _) _ _, args) <- unApply t
54 = mapMaybe isF args ++ concatMap guarded args
55 guarded t = mapMaybe isF [t]
57-- Calculate whether a collection of constructors is collapsible
59collapseCons :: Name -> [(Name, Type)] -> Idris ()
60collapseCons ty cons =
61 do i <- getIState
62 let cons' = map (\ (n, t) -> (n, map snd (getArgTys t))) cons
63 allFR <- mapM (forceRec i) cons'
64 if and allFR then detaggable (map getRetTy (map snd cons))
65 else return () -- not collapsible as not detaggable
66 where
67 setCollapsible :: Name -> Idris ()
68 setCollapsible n
69 = do i <- getIState
70 iLOG $ show n ++ " collapsible"
71 case lookupCtxt Nothing n (idris_optimisation i) of
72 (oi:_) -> do let oi' = oi { collapsible = True }
73 let opts = addDef n oi' (idris_optimisation i)
74 putIState (i { idris_optimisation = opts })
75 [] -> do let oi = Optimise True [] []
76 let opts = addDef n oi (idris_optimisation i)
77 putIState (i { idris_optimisation = opts })
78 addIBC (IBCOpt n)
80 forceRec :: IState -> (Name, [Type]) -> Idris Bool
81 forceRec i (n, ts)
82 = case lookupCtxt Nothing n (idris_optimisation i) of
83 (oi:_) -> checkFR (forceable oi) 0 ts
84 _ -> return False
85 checkFR fs i [] = return True
86 checkFR fs i (_ : xs) | i `elem` fs = checkFR fs (i + 1) xs
87 checkFR fs i (t : xs)
88 -- must be recursive or type is not collapsible
89 = do let rt = getRetTy t
90 if (ty `elem` freeNames rt)
91 then checkFR fs (i+1) xs
92 else return False
94 detaggable :: [Type] -> Idris ()
95 detaggable rtys
96 = do let rtyArgs = map (snd . unApply) rtys
97 -- if every rtyArgs is disjoint with every other, it's detaggable,
98 -- therefore also collapsible given forceable/recursive check
99 if disjoint rtyArgs
100 then mapM_ setCollapsible (ty : map fst cons)
101 else return ()
103 disjoint :: [[Term]] -> Bool
104 disjoint [] = True
105 disjoint [x] = True
106 disjoint (x : xs) = anyDisjoint x xs && disjoint xs
108 anyDisjoint x [] = True
109 anyDisjoint x (y : ys) = disjointCons x y
111 disjointCons [] [] = False
112 disjointCons [] y = False
113 disjointCons x [] = False
114 disjointCons (x : xs) (y : ys)
115 = disjointCon x y || disjointCons xs ys
117 disjointCon x y = let (cx, _) = unApply x
118 (cy, _) = unApply y in
119 case (cx, cy) of
120 (P (DCon _ _) nx _, P (DCon _ _) ny _) -> nx /= ny
121 _ -> False
123class Optimisable term where
124 applyOpts :: term -> Idris term
125 stripCollapsed :: term -> Idris term
127instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
128 applyOpts (x, y) = do x' <- applyOpts x
129 y' <- applyOpts y
130 return (x', y')
131 stripCollapsed (x, y) = do x' <- stripCollapsed x
132 y' <- stripCollapsed y
133 return (x', y')
136instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
137 applyOpts (v, x, y) = do x' <- applyOpts x
138 y' <- applyOpts y
139 return (v, x', y')
140 stripCollapsed (v, x, y) = do x' <- stripCollapsed x
141 y' <- stripCollapsed y
142 return (v, x', y')
144instance Optimisable a => Optimisable [a] where
145 applyOpts = mapM applyOpts
146 stripCollapsed = mapM stripCollapsed
148instance Optimisable a => Optimisable (Either a (a, a)) where
149 applyOpts (Left t) = do t' <- applyOpts t; return $ Left t'
150 applyOpts (Right t) = do t' <- applyOpts t; return $ Right t'
151 stripCollapsed (Left t) = do t' <- stripCollapsed t; return $ Left t'
152 stripCollapsed (Right t) = do t' <- stripCollapsed t; return $ Right t'
154-- Raw is for compile time optimisation (before type checking)
155-- Term is for run time optimisation (after type checking, collapsing allowed)
157-- Compile time: no collapsing
159instance Optimisable Raw where
160 applyOpts t@(RApp f a)
161 | (Var n, args) <- raw_unapply t -- MAGIC HERE
162 = do args' <- mapM applyOpts args
163 i <- getIState
164 case lookupCtxt Nothing n (idris_optimisation i) of
165 (oi:_) -> return $ applyDataOpt oi n args'
166 _ -> return (raw_apply (Var n) args')
167 | otherwise = do f' <- applyOpts f
168 a' <- applyOpts a
169 return (RApp f' a')
170 applyOpts (RBind n b t) = do b' <- applyOpts b
171 t' <- applyOpts t
172 return (RBind n b' t')
173 applyOpts (RForce t) = applyOpts t
174 applyOpts t = return t
176 stripCollapsed t = return t
178instance Optimisable t => Optimisable (Binder t) where
179 applyOpts (Let t v) = do t' <- applyOpts t
180 v' <- applyOpts v
181 return (Let t' v')
182 applyOpts b = do t' <- applyOpts (binderTy b)
183 return (b { binderTy = t' })
184 stripCollapsed (Let t v) = do t' <- stripCollapsed t
185 v' <- stripCollapsed v
186 return (Let t' v')
187 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
188 return (b { binderTy = t' })
191applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
192applyDataOpt oi n args
193 = let args' = zipWith doForce (map (\x -> x `elem` (forceable oi)) [0..])
194 args in
195 raw_apply (Var n) args'
196 where
197 doForce True a = RForce a
198 doForce False a = a
200-- Run-time: do everything
202instance Optimisable (TT Name) where
203 applyOpts c@(P (DCon t arity) n _)
204 = do i <- getIState
205 case lookupCtxt Nothing n (idris_optimisation i) of
206 (oi:_) -> return $ applyDataOptRT oi n t arity []
207 _ -> return c
208 applyOpts t@(App f a)
209 | (c@(P (DCon t arity) n _), args) <- unApply t -- MAGIC HERE
210 = do args' <- mapM applyOpts args
211 i <- getIState
212 case lookupCtxt Nothing n (idris_optimisation i) of
213 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
214 _ -> return (mkApp c args')
215 | otherwise = do f' <- applyOpts f
216 a' <- applyOpts a
217 return (App f' a')
218 applyOpts (Bind n b t) = do b' <- applyOpts b
219 t' <- applyOpts t
220 return (Bind n b' t')
221 applyOpts (Proj t i) = do t' <- applyOpts t
222 return (Proj t' i)
223 applyOpts t = return t
225 stripCollapsed (Bind n (PVar x) t) | (P _ ty _, _) <- unApply x
226 = do i <- getIState
227 case lookupCtxt Nothing ty (idris_optimisation i) of
228 [oi] -> if collapsible oi
229 then do t' <- stripCollapsed t
230 return (Bind n (PVar x) (instantiate Erased t'))
231 else do t' <- stripCollapsed t
232 return (Bind n (PVar x) t')
233 _ -> do t' <- stripCollapsed t
234 return (Bind n (PVar x) t')
235 stripCollapsed (Bind n (PVar x) t)
236 = do t' <- stripCollapsed t
237 return (Bind n (PVar x) t')
238 stripCollapsed t = return t
240-- Need to saturate arguments first to ensure that erasure happens uniformly
242applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
243applyDataOptRT oi n tag arity args
244 | length args == arity = doOpts n args (collapsible oi) (forceable oi)
245 | otherwise = let extra = satArgs (arity - length args)
246 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
247 (collapsible oi) (forceable oi) in
248 bind extra tm
249 where
250 satArgs n = map (\i -> MN i "sat") [1..n]
252 bind [] tm = tm
253 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
255 doOpts n args True f = Erased
256 doOpts n args _ forced
257 = let args' = filter keep (zip (map (\x -> x `elem` forced) [0..])
258 args) in
259 mkApp (P (DCon tag (arity - length forced)) n Erased) (map snd args')
261 keep (forced, _) = not forced