git.iko.soy/Idris-dev
on d66da21 help sign in
Idris-dev/src/Idris/DataOpts.hs
302 lines 12 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 n (idris_optimisation i) of
21 [] -> return $ Optimise False 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 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 checkNewType cons'
66 where
67 -- one constructor; if one remaining argument, treat as newtype
68 checkNewType [(n, ts)] = do
69 i <- getIState
70 case lookupCtxt n (idris_optimisation i) of
71 (oi:_) -> do let remaining = length ts - length (forceable oi)
72 if remaining == 1 then
73 do let oi' = oi { isnewtype = True }
74 let opts = addDef n oi' (idris_optimisation i)
75 putIState (i { idris_optimisation = opts })
76 else return ()
77 _ -> return ()
79 checkNewType _ = return ()
81 setCollapsible :: Name -> Idris ()
82 setCollapsible n
83 = do i <- getIState
84 iLOG $ show n ++ " collapsible"
85 case lookupCtxt n (idris_optimisation i) of
86 (oi:_) -> do let oi' = oi { collapsible = True }
87 let opts = addDef n oi' (idris_optimisation i)
88 putIState (i { idris_optimisation = opts })
89 [] -> do let oi = Optimise True False [] []
90 let opts = addDef n oi (idris_optimisation i)
91 putIState (i { idris_optimisation = opts })
92 addIBC (IBCOpt n)
94 forceRec :: IState -> (Name, [Type]) -> Idris Bool
95 forceRec i (n, ts)
96 = case lookupCtxt n (idris_optimisation i) of
97 (oi:_) -> checkFR (forceable oi) 0 ts
98 _ -> return False
99 checkFR fs i [] = return True
100 checkFR fs i (_ : xs) | i `elem` fs = checkFR fs (i + 1) xs
101 checkFR fs i (t : xs)
102 -- must be recursive or type is not collapsible
103 = do let (rtf, rta) = unApply $ getRetTy t
104 if (ty `elem` freeNames rtf)
105 then checkFR fs (i+1) xs
106 else return False
108 detaggable :: [Type] -> Idris ()
109 detaggable rtys
110 = do let rtyArgs = map (snd . unApply) rtys
111 -- if every rtyArgs is disjoint with every other, it's detaggable,
112 -- therefore also collapsible given forceable/recursive check
113 if disjoint rtyArgs
114 then mapM_ setCollapsible (ty : map fst cons)
115 else return ()
117 disjoint :: [[Term]] -> Bool
118 disjoint [] = True
119 disjoint [x] = True
120 disjoint (x : xs) = anyDisjoint x xs && disjoint xs
122 anyDisjoint x [] = True
123 anyDisjoint x (y : ys) = disjointCons x y
125 disjointCons [] [] = False
126 disjointCons [] y = False
127 disjointCons x [] = False
128 disjointCons (x : xs) (y : ys)
129 = disjointCon x y || disjointCons xs ys
131 disjointCon x y = let (cx, _) = unApply x
132 (cy, _) = unApply y in
133 case (cx, cy) of
134 (P (DCon _ _) nx _, P (DCon _ _) ny _) -> nx /= ny
135 _ -> False
137class Optimisable term where
138 applyOpts :: term -> Idris term
139 stripCollapsed :: term -> Idris term
141instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
142 applyOpts (x, y) = do x' <- applyOpts x
143 y' <- applyOpts y
144 return (x', y')
145 stripCollapsed (x, y) = do x' <- stripCollapsed x
146 y' <- stripCollapsed y
147 return (x', y')
150instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
151 applyOpts (v, x, y) = do x' <- applyOpts x
152 y' <- applyOpts y
153 return (v, x', y')
154 stripCollapsed (v, x, y) = do x' <- stripCollapsed x
155 y' <- stripCollapsed y
156 return (v, x', y')
158instance Optimisable a => Optimisable [a] where
159 applyOpts = mapM applyOpts
160 stripCollapsed = mapM stripCollapsed
162instance Optimisable a => Optimisable (Either a (a, a)) where
163 applyOpts (Left t) = do t' <- applyOpts t; return $ Left t'
164 applyOpts (Right t) = do t' <- applyOpts t; return $ Right t'
165 stripCollapsed (Left t) = do t' <- stripCollapsed t; return $ Left t'
166 stripCollapsed (Right t) = do t' <- stripCollapsed t; return $ Right t'
168-- Raw is for compile time optimisation (before type checking)
169-- Term is for run time optimisation (after type checking, collapsing allowed)
171-- Compile time: no collapsing
173instance Optimisable Raw where
174 applyOpts t@(RApp f a)
175 | (Var n, args) <- raw_unapply t -- MAGIC HERE
176 = do args' <- mapM applyOpts args
177 i <- getIState
178 case lookupCtxt n (idris_optimisation i) of
179 (oi:_) -> return $ applyDataOpt oi n args'
180 _ -> return (raw_apply (Var n) args')
181 | otherwise = do f' <- applyOpts f
182 a' <- applyOpts a
183 return (RApp f' a')
184 applyOpts (RBind n b t) = do b' <- applyOpts b
185 t' <- applyOpts t
186 return (RBind n b' t')
187 applyOpts (RForce t) = applyOpts t
188 applyOpts t = return t
190 stripCollapsed t = return t
192instance Optimisable t => Optimisable (Binder t) where
193 applyOpts (Let t v) = do t' <- applyOpts t
194 v' <- applyOpts v
195 return (Let t' v')
196 applyOpts b = do t' <- applyOpts (binderTy b)
197 return (b { binderTy = t' })
198 stripCollapsed (Let t v) = do t' <- stripCollapsed t
199 v' <- stripCollapsed v
200 return (Let t' v')
201 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
202 return (b { binderTy = t' })
205applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
206applyDataOpt oi n args
207 = let args' = zipWith doForce (map (\x -> x `elem` (forceable oi)) [0..])
208 args in
209 raw_apply (Var n) args'
210 where
211 doForce True a = RForce a
212 doForce False a = a
214-- Run-time: do everything
216instance Optimisable (TT Name) where
217 applyOpts (P _ (NS (UN "plus") ["Nat","Prelude"]) _)
218 = return (P Ref (UN "prim__addBigInt") Erased)
219 applyOpts (P _ (NS (UN "mult") ["Nat","Prelude"]) _)
220 = return (P Ref (UN "prim__mulBigInt") Erased)
221 applyOpts (App (P _ (NS (UN "fromIntegerNat") ["Nat","Prelude"]) _) x)
222 = applyOpts x
223 applyOpts (P _ (NS (UN "fromIntegerNat") ["Nat","Prelude"]) _)
224 = return (App (P Ref (NS (UN "id") ["Basics","Prelude"]) Erased) Erased)
225 applyOpts (P _ (NS (UN "toIntegerNat") ["Nat","Prelude"]) _)
226 = return (App (P Ref (NS (UN "id") ["Basics","Prelude"]) Erased) Erased)
227 applyOpts c@(P (DCon t arity) n _)
228 = do i <- getIState
229 case lookupCtxt n (idris_optimisation i) of
230 (oi:_) -> return $ applyDataOptRT oi n t arity []
231 _ -> return c
232 applyOpts t@(App f a)
233 | (c@(P (DCon t arity) n _), args) <- unApply t
234 = do args' <- mapM applyOpts args
235 i <- getIState
236 case lookupCtxt n (idris_optimisation i) of
237 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
238 _ -> return (mkApp c args')
239 | otherwise = do f' <- applyOpts f
240 a' <- applyOpts a
241 return (App f' a')
242 applyOpts (Bind n b t) = do b' <- applyOpts b
243 t' <- applyOpts t
244 return (Bind n b' t')
245 applyOpts (Proj t i) = do t' <- applyOpts t
246 return (Proj t' i)
247 applyOpts t = return t
249 stripCollapsed (Bind n (PVar x) t) | (P _ ty _, _) <- unApply x
250 = do i <- getIState
251 -- NOTE: This assumes that 'ty' is in normal form, which it
252 -- has to be before now because we're not keeping track of
253 -- an environment so we can't do it here.
254 case lookupCtxt ty (idris_optimisation i) of
255 [oi] -> if collapsible oi
256 then do t' <- stripCollapsed t
257 return (Bind n (PVar x) (instantiate Erased t'))
258 else do t' <- stripCollapsed t
259 return (Bind n (PVar x) t')
260 _ -> do t' <- stripCollapsed t
261 return (Bind n (PVar x) t')
262 stripCollapsed (Bind n (PVar x) t)
263 = do t' <- stripCollapsed t
264 return (Bind n (PVar x) t')
265 stripCollapsed t = return t
267-- Need to saturate arguments first to ensure that erasure happens uniformly
269applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
270applyDataOptRT oi n tag arity args
271 | length args == arity = doOpts n args (collapsible oi) (forceable oi)
272 | otherwise = let extra = satArgs (arity - length args)
273 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
274 (collapsible oi) (forceable oi) in
275 bind extra tm
276 where
277 satArgs n = map (\i -> MN i "sat") [1..n]
279 bind [] tm = tm
280 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
282 -- Nat special cases
283 -- TODO: Would be nice if this was configurable in idris source!
284 doOpts (NS (UN "Z") ["Nat", "Prelude"]) [] _ _ = Constant (BI 0)
285 doOpts (NS (UN "S") ["Nat", "Prelude"]) [k] _ _
286 = App (App (P Ref (UN "prim__addBigInt") Erased) k) (Constant (BI 1))
288 doOpts n args True f = Erased
289 doOpts n args _ forced
290 = let args' = filter keep (zip (map (\x -> x `elem` forced) [0..])
291 args) in
292 if isnewtype oi
293 then case args' of
294 [(_, val)] -> val
295 _ -> error "Can't happen (not isnewtype)"
296 else
297 mkApp (P (DCon tag (arity - length forced)) n Erased)
298 (map snd args')
300 keep (forced, _) = not forced