git.iko.soy/Idris-dev
on fac49fc help sign in
Idris-dev/src/Idris/DataOpts.hs
310 lines 13 kB history raw
1{-# LANGUAGE PatternGuards #-}
3module Idris.DataOpts where
5-- Forcing, detagging and collapsing
7import Idris.AbsSyntax
8import Idris.AbsSyntaxTree
9import Idris.Core.TT
11import Control.Applicative
12import qualified Data.IntMap as M
13import Data.List
14import Data.Maybe
15import Debug.Trace
17type ForceMap = M.IntMap Forceability
19-- Calculate the forceable arguments to a constructor
20-- and update the set of optimisations.
21forceArgs :: Name -> Name -> Type -> Idris ()
22forceArgs typeName n t = do
23 ist <- getIState
24 let fargs = getForcedArgs ist typeName t
25 copt = case lookupCtxt n (idris_optimisation ist) of
26 [] -> Optimise False False [] []
27 op:_ -> op
28 opts = addDef n (copt { forceable = M.toList fargs }) (idris_optimisation ist)
29 putIState (ist { idris_optimisation = opts })
30 addIBC (IBCOpt n)
31 iLOG $ "Forced: " ++ show n ++ " " ++ show fargs ++ "\n from " ++ show t
33getForcedArgs :: IState -> Name -> Type -> ForceMap
34getForcedArgs ist typeName t = addCollapsibleArgs 0 t $ forcedInTarget 0 t
35 where
36 maxUnion = M.unionWith max
38 -- Label all occurrences of the variable bound in Pi in the rest of
39 -- the term with the number i so that we can recognize them anytime later.
40 label i = instantiate $ P Bound (MN i "ctor_arg") Erased
42 addCollapsibleArgs :: Int -> Type -> ForceMap -> ForceMap
43 addCollapsibleArgs i (Bind vn (Pi ty) rest) alreadyForceable
44 = addCollapsibleArgs (i+1) (label i rest) (forceable $ unApply ty)
45 where
46 -- forceable takes an un-applied type of a ctor argument
47 forceable (P _ tn _, args)
48 -- if `ty' is collapsible, the argument is unconditionally forceable
49 | isCollapsible tn
50 = M.insert i Unconditional alreadyForceable
52 -- a recursive occurrence with known indices is conditionally forceable
53 | tn == typeName
54 = M.insertWith max i Conditional alreadyForceable
56 forceable _ = alreadyForceable
58 isCollapsible :: Name -> Bool
59 isCollapsible n = case lookupCtxt n (idris_optimisation ist) of
60 [oi] -> collapsible oi
61 _ -> False
63 addCollapsibleArgs _ _ fs = fs
65 forcedInTarget :: Int -> Type -> ForceMap
66 forcedInTarget i (Bind _ (Pi _) rest) = forcedInTarget (i+1) (label i rest)
67 forcedInTarget i t@(App f a) | (_, as) <- unApply t = unionMap guardedArgs as
68 forcedInTarget _ _ = M.empty
70 guardedArgs :: Term -> ForceMap
71 guardedArgs t@(App f a) | (P (DCon _ _) _ _, args) <- unApply t
72 = unionMap bareArg args `maxUnion` unionMap guardedArgs args
73 guardedArgs t = bareArg t
75 bareArg :: Term -> ForceMap
76 bareArg (P _ (MN i "ctor_arg") _) = M.singleton i Unconditional
77 bareArg _ = M.empty
79 unionMap :: (a -> ForceMap) -> [a] -> ForceMap
80 unionMap f = M.unionsWith max . map f
82-- Calculate whether a collection of constructors is collapsible
83-- and update the state accordingly.
84collapseCons :: Name -> [(Name, Type)] -> Idris ()
85collapseCons tn ctors = do
86 ist <- getIState
87 case ctors of
88 _
89 | all (ctorCollapsible ist) ctors
90 , disjointTerms ctorTargetArgs
91 -> mapM_ setCollapsible (tn : map fst ctors)
93 [(cn, ct)]
94 -> checkNewType ist cn ct
96 _ -> return () -- nothing can be done
97 where
98 ctorTargetArgs = map (snd . unApply . getRetTy . snd) ctors
100 ctorArity :: Type -> Int
101 ctorArity = length . getArgTys
103 ctorCollapsible :: IState -> (Name, Type) -> Bool
104 ctorCollapsible ist (n, t) = all (`M.member` forceMap) [0 .. ctorArity t - 1]
105 where
106 forceMap = case lookupCtxt n (idris_optimisation ist) of
107 oi:_ -> M.fromList $ forceable oi
108 _ -> M.empty
110 -- one constructor; if one remaining argument, treat as newtype
111 checkNewType :: IState -> Name -> Type -> Idris ()
112 checkNewType ist cn ct
113 | oi:_ <- lookupCtxt cn opt
114 , length (getArgTys ct) == 1 + forcedCnt (M.fromList $ forceable oi)
115 = putIState ist{ idris_optimisation = opt' oi }
116 | otherwise = return ()
117 where
118 opt = idris_optimisation ist
119 opt' oi = addDef cn oi{ isnewtype = True } opt
121 setCollapsible :: Name -> Idris ()
122 setCollapsible n
123 = do i <- getIState
124 iLOG $ show n ++ " collapsible"
125 case lookupCtxt n (idris_optimisation i) of
126 (oi:_) -> do let oi' = oi { collapsible = True }
127 let opts = addDef n oi' (idris_optimisation i)
128 putIState (i { idris_optimisation = opts })
129 [] -> do let oi = Optimise True False [] []
130 let opts = addDef n oi (idris_optimisation i)
131 putIState (i { idris_optimisation = opts })
132 addIBC (IBCOpt n)
134 disjointTerms :: [[Term]] -> Bool
135 disjointTerms [] = True
136 disjointTerms [xs] = True
137 disjointTerms (xs : xss) =
138 -- xs is disjoint with every pattern from xss
139 all (or . zipWith disjoint xs) xss
140 -- and xss is pairwise disjoint, too
141 && disjointTerms xss
143 -- Return True if the two patterns are provably disjoint.
144 -- Return False if they're not or if unsure.
145 disjoint :: Term -> Term -> Bool
146 disjoint x y = case (cx, cy) of
147 -- data constructors -> compare their names
148 (P (DCon _ _) nx _, P (DCon _ _) ny _)
149 | nx /= ny -> True
150 | otherwise -> or $ zipWith disjoint xargs yargs
151 _ -> False
152 where
153 (cx, xargs) = unApply x
154 (cy, yargs) = unApply y
156class Optimisable term where
157 applyOpts :: term -> Idris term
158 stripCollapsed :: term -> Idris term
160instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
161 applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
162 stripCollapsed (x, y) = (,) <$> stripCollapsed x <*> stripCollapsed y
164instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
165 applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
166 stripCollapsed (v, x, y) = (,,) v <$> stripCollapsed x <*> stripCollapsed y
168instance Optimisable a => Optimisable [a] where
169 applyOpts = mapM applyOpts
170 stripCollapsed = mapM stripCollapsed
172instance Optimisable a => Optimisable (Either a (a, a)) where
173 applyOpts (Left t) = Left <$> applyOpts t
174 applyOpts (Right t) = Right <$> applyOpts t
175 stripCollapsed (Left t) = Left <$> stripCollapsed t
176 stripCollapsed (Right t) = Right <$> stripCollapsed t
178-- Raw is for compile time optimisation (before type checking)
179-- Term is for run time optimisation (after type checking, collapsing allowed)
181-- Compile time: no collapsing
183instance Optimisable Raw where
184 applyOpts t@(RApp f a)
185 | (Var n, args) <- raw_unapply t -- MAGIC HERE
186 = do args' <- mapM applyOpts args
187 i <- getIState
188 return $ case lookupCtxt n (idris_optimisation i) of
189 oi:_ -> applyDataOpt oi n args'
190 _ -> raw_apply (Var n) args'
191 | otherwise = RApp <$> applyOpts f <*> applyOpts a
193 applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
194 applyOpts (RForce t) = applyOpts t
195 applyOpts t = return t
197 stripCollapsed t = return t
199-- Erase types (makes ibc smaller, and we don't need them)
200instance Optimisable (Binder (TT Name)) where
201 applyOpts (Let t v) = Let <$> return Erased <*> applyOpts v
202 applyOpts b = return (b { binderTy = Erased })
203 stripCollapsed (Let t v) = Let <$> return Erased <*> stripCollapsed v
204 stripCollapsed b = return (b { binderTy = Erased })
206instance Optimisable (Binder Raw) where
207 applyOpts b = do t' <- applyOpts (binderTy b)
208 return (b { binderTy = t' })
209 stripCollapsed (Let t v) = Let <$> stripCollapsed t <*> stripCollapsed v
210 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
211 return (b { binderTy = t' })
213forcedArgSeq :: OptInfo -> [Maybe Forceability]
214forcedArgSeq oi = map (\i -> M.lookup i forceMap) [0..]
215 where
216 forceMap = M.fromList $ forceable oi
218forcedCnt :: ForceMap -> Int
219forcedCnt = length . filter (== Unconditional) . M.elems
221applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
222applyDataOpt oi n args
223 = raw_apply (Var n) $ zipWith doForce (forcedArgSeq oi) args
224 where
225 doForce (Just Unconditional) a = RForce a
226 doForce _ a = a
228-- Run-time: do everything
230instance Optimisable (TT Name) where
231 applyOpts (P _ (NS (UN "plus") ["Nat","Prelude"]) _)
232 = return (P Ref (UN "prim__addBigInt") Erased)
233 applyOpts (P _ (NS (UN "mult") ["Nat","Prelude"]) _)
234 = return (P Ref (UN "prim__mulBigInt") Erased)
235 applyOpts (App (P _ (NS (UN "fromIntegerNat") ["Nat","Prelude"]) _) x)
236 = applyOpts x
237 applyOpts (P _ (NS (UN "fromIntegerNat") ["Nat","Prelude"]) _)
238 = return (App (P Ref (NS (UN "id") ["Basics","Prelude"]) Erased) Erased)
239 applyOpts (P _ (NS (UN "toIntegerNat") ["Nat","Prelude"]) _)
240 = return (App (P Ref (NS (UN "id") ["Basics","Prelude"]) Erased) Erased)
241 applyOpts c@(P (DCon t arity) n _)
242 = do i <- getIState
243 case lookupCtxt n (idris_optimisation i) of
244 (oi:_) -> return $ applyDataOptRT oi n t arity []
245 _ -> return c
246 applyOpts t@(App f a)
247 | (c@(P (DCon t arity) n _), args) <- unApply t
248 = do args' <- mapM applyOpts args
249 i <- getIState
250 case lookupCtxt n (idris_optimisation i) of
251 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
252 _ -> return (mkApp c args')
253 | otherwise = do f' <- applyOpts f
254 a' <- applyOpts a
255 return (App f' a')
256 applyOpts (Bind n b t) = do b' <- applyOpts b
257 t' <- applyOpts t
258 return (Bind n b' t')
259 applyOpts (Proj t i) = do t' <- applyOpts t
260 return (Proj t' i)
261 applyOpts t = return t
263 stripCollapsed (Bind n (PVar x) t) | (P _ ty _, _) <- unApply x
264 = do i <- getIState
265 -- NOTE: This assumes that 'ty' is in normal form, which it
266 -- has to be before now because we're not keeping track of
267 -- an environment so we can't do it here.
268 case lookupCtxt ty (idris_optimisation i) of
269 [oi] -> if collapsible oi
270 then do t' <- stripCollapsed t
271 return (Bind n (PVar x) (instantiate Erased t'))
272 else do t' <- stripCollapsed t
273 return (Bind n (PVar x) t')
274 _ -> do t' <- stripCollapsed t
275 return (Bind n (PVar x) t')
276 stripCollapsed (Bind n (PVar x) t)
277 = do t' <- stripCollapsed t
278 return (Bind n (PVar x) t')
279 stripCollapsed t = return t
281-- Need to saturate arguments first to ensure that erasure happens uniformly
283applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
284applyDataOptRT oi n tag arity args
285 | length args == arity = doOpts n args (collapsible oi) (M.fromList $ forceable oi)
286 | otherwise = let extra = satArgs (arity - length args)
287 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
288 (collapsible oi) (M.fromList $ forceable oi) in
289 bind extra tm
290 where
291 satArgs n = map (\i -> MN i "sat") [1..n]
293 bind [] tm = tm
294 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
296 -- Nat special cases
297 -- TODO: Would be nice if this was configurable in idris source!
298 doOpts (NS (UN "Z") ["Nat", "Prelude"]) [] _ _ = Constant (BI 0)
299 doOpts (NS (UN "S") ["Nat", "Prelude"]) [k] _ _
300 = App (App (P Ref (UN "prim__addBigInt") Erased) k) (Constant (BI 1))
302 doOpts n args True _ = Erased
303 doOpts n args _ forceMap
304 | isnewtype oi = case args' of
305 [val] -> val
306 _ -> error $ "Can't happen (newtype not a singleton): " ++ show args'
307 | otherwise = mkApp ctor' args'
308 where
309 ctor' = (P (DCon tag (arity - forcedCnt forceMap)) n Erased)
310 args' = [t | (f, t) <- zip (forcedArgSeq oi) args, f /= Just Unconditional]