git.iko.soy/Idris-dev
on caed480 help sign in
Idris-dev/src/Idris/DataOpts.hs
321 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 (sMN 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) _)
77 | ctor_arg == txt "ctor_arg" = M.singleton i Unconditional
78 bareArg _ = M.empty
80 unionMap :: (a -> ForceMap) -> [a] -> ForceMap
81 unionMap f = M.unionsWith max . map f
83-- Calculate whether a collection of constructors is collapsible
84-- and update the state accordingly.
85collapseCons :: Name -> [(Name, Type)] -> Idris ()
86collapseCons tn ctors = do
87 ist <- getIState
88 case ctors of
89 _
90 | all (ctorCollapsible ist) ctors
91 , disjointTerms ctorTargetArgs
92 -> mapM_ setCollapsible (tn : map fst ctors)
94 [(cn, ct)]
95 -> checkNewType ist cn ct
97 _ -> return () -- nothing can be done
98 where
99 ctorTargetArgs = map (snd . unApply . getRetTy . snd) ctors
101 ctorArity :: Type -> Int
102 ctorArity = length . getArgTys
104 ctorCollapsible :: IState -> (Name, Type) -> Bool
105 ctorCollapsible ist (n, t) = all (`M.member` forceMap) [0 .. ctorArity t - 1]
106 where
107 forceMap = case lookupCtxt n (idris_optimisation ist) of
108 oi:_ -> M.fromList $ forceable oi
109 _ -> M.empty
111 -- one constructor; if one remaining argument, treat as newtype
112 checkNewType :: IState -> Name -> Type -> Idris ()
113 checkNewType ist cn ct
114 | oi:_ <- lookupCtxt cn opt
115 , length (getArgTys ct) == 1 + forcedCnt (M.fromList $ forceable oi)
116 = putIState ist{ idris_optimisation = opt' oi }
117 | otherwise = return ()
118 where
119 opt = idris_optimisation ist
120 opt' oi = addDef cn oi{ isnewtype = True } opt
122 setCollapsible :: Name -> Idris ()
123 setCollapsible n
124 = do i <- getIState
125 iLOG $ show n ++ " collapsible"
126 case lookupCtxt n (idris_optimisation i) of
127 (oi:_) -> do let oi' = oi { collapsible = True }
128 let opts = addDef n oi' (idris_optimisation i)
129 putIState (i { idris_optimisation = opts })
130 [] -> do let oi = Optimise True False [] []
131 let opts = addDef n oi (idris_optimisation i)
132 putIState (i { idris_optimisation = opts })
133 addIBC (IBCOpt n)
135 disjointTerms :: [[Term]] -> Bool
136 disjointTerms [] = True
137 disjointTerms [xs] = True
138 disjointTerms (xs : xss) =
139 -- xs is disjoint with every pattern from xss
140 all (or . zipWith disjoint xs) xss
141 -- and xss is pairwise disjoint, too
142 && disjointTerms xss
144 -- Return True if the two patterns are provably disjoint.
145 -- Return False if they're not or if unsure.
146 disjoint :: Term -> Term -> Bool
147 disjoint x y = case (cx, cy) of
148 -- data constructors -> compare their names
149 (P (DCon _ _) nx _, P (DCon _ _) ny _)
150 | nx /= ny -> True
151 | otherwise -> or $ zipWith disjoint xargs yargs
152 _ -> False
153 where
154 (cx, xargs) = unApply x
155 (cy, yargs) = unApply y
157class Optimisable term where
158 applyOpts :: term -> Idris term
159 stripCollapsed :: term -> Idris term
161instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
162 applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
163 stripCollapsed (x, y) = (,) <$> stripCollapsed x <*> stripCollapsed y
165instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
166 applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
167 stripCollapsed (v, x, y) = (,,) v <$> stripCollapsed x <*> stripCollapsed y
169instance Optimisable a => Optimisable [a] where
170 applyOpts = mapM applyOpts
171 stripCollapsed = mapM stripCollapsed
173instance Optimisable a => Optimisable (Either a (a, a)) where
174 applyOpts (Left t) = Left <$> applyOpts t
175 applyOpts (Right t) = Right <$> applyOpts t
176 stripCollapsed (Left t) = Left <$> stripCollapsed t
177 stripCollapsed (Right t) = Right <$> stripCollapsed t
179-- Raw is for compile time optimisation (before type checking)
180-- Term is for run time optimisation (after type checking, collapsing allowed)
182-- Compile time: no collapsing
184instance Optimisable Raw where
185 applyOpts t@(RApp f a)
186 | (Var n, args) <- raw_unapply t -- MAGIC HERE
187 = do args' <- mapM applyOpts args
188 i <- getIState
189 return $ case lookupCtxt n (idris_optimisation i) of
190 oi:_ -> applyDataOpt oi n args'
191 _ -> raw_apply (Var n) args'
192 | otherwise = RApp <$> applyOpts f <*> applyOpts a
194 applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
195 applyOpts (RForce t) = applyOpts t
196 applyOpts t = return t
198 stripCollapsed t = return t
200-- Erase types (makes ibc smaller, and we don't need them)
201instance Optimisable (Binder (TT Name)) where
202 applyOpts (Let t v) = Let <$> return Erased <*> applyOpts v
203 applyOpts b = return (b { binderTy = Erased })
204 stripCollapsed (Let t v) = Let <$> return Erased <*> stripCollapsed v
205 stripCollapsed b = return (b { binderTy = Erased })
207instance Optimisable (Binder Raw) where
208 applyOpts b = do t' <- applyOpts (binderTy b)
209 return (b { binderTy = t' })
210 stripCollapsed (Let t v) = Let <$> stripCollapsed t <*> stripCollapsed v
211 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
212 return (b { binderTy = t' })
214forcedArgSeq :: OptInfo -> [Maybe Forceability]
215forcedArgSeq oi = map (\i -> M.lookup i forceMap) [0..]
216 where
217 forceMap = M.fromList $ forceable oi
219forcedCnt :: ForceMap -> Int
220forcedCnt = length . filter (== Unconditional) . M.elems
222applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
223applyDataOpt oi n args
224 = raw_apply (Var n) $ zipWith doForce (forcedArgSeq oi) args
225 where
226 doForce (Just Unconditional) a = RForce a
227 doForce _ a = a
229-- Run-time: do everything
231prel = [txt "Nat", txt "Prelude"]
233instance Optimisable (TT Name) where
234 applyOpts (P _ (NS (UN fn) mod) _)
235 | fn == txt "plus" && mod == prel
236 = return (P Ref (sUN "prim__addBigInt") Erased)
237 applyOpts (P _ (NS (UN fn) mod) _)
238 | fn == txt "mult" && mod == prel
239 = return (P Ref (sUN "prim__mulBigInt") Erased)
240 applyOpts (App (P _ (NS (UN fn) mod) _) x)
241 | fn == txt "fromIntegerNat" && mod == prel
242 = applyOpts x
243 applyOpts (P _ (NS (UN fn) mod) _)
244 | fn == txt "fromIntegerNat" && mod == prel
245 = return (App (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
246 applyOpts (P _ (NS (UN fn) mod) _)
247 | fn == txt "toIntegerNat" && mod == prel
248 = return (App (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
249 applyOpts c@(P (DCon t arity) n _)
250 = do i <- getIState
251 case lookupCtxt n (idris_optimisation i) of
252 (oi:_) -> return $ applyDataOptRT oi n t arity []
253 _ -> return c
254 applyOpts t@(App f a)
255 | (c@(P (DCon t arity) n _), args) <- unApply t
256 = do args' <- mapM applyOpts args
257 i <- getIState
258 case lookupCtxt n (idris_optimisation i) of
259 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
260 _ -> return (mkApp c args')
261 | otherwise = do f' <- applyOpts f
262 a' <- applyOpts a
263 return (App f' a')
264 applyOpts (Bind n b t) = do b' <- applyOpts b
265 t' <- applyOpts t
266 return (Bind n b' t')
267 applyOpts (Proj t i) = do t' <- applyOpts t
268 return (Proj t' i)
269 applyOpts t = return t
271 stripCollapsed (Bind n (PVar x) t) | (P _ ty _, _) <- unApply x
272 = do i <- getIState
273 -- NOTE: This assumes that 'ty' is in normal form, which it
274 -- has to be before now because we're not keeping track of
275 -- an environment so we can't do it here.
276 case lookupCtxt ty (idris_optimisation i) of
277 [oi] -> if collapsible oi
278 then do t' <- stripCollapsed t
279 return (Bind n (PVar x) (instantiate Erased t'))
280 else do t' <- stripCollapsed t
281 return (Bind n (PVar x) t')
282 _ -> do t' <- stripCollapsed t
283 return (Bind n (PVar x) t')
284 stripCollapsed (Bind n (PVar x) t)
285 = do t' <- stripCollapsed t
286 return (Bind n (PVar x) t')
287 stripCollapsed t = return t
289-- Need to saturate arguments first to ensure that erasure happens uniformly
291applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
292applyDataOptRT oi n tag arity args
293 | length args == arity = doOpts n args (collapsible oi) (M.fromList $ forceable oi)
294 | otherwise = let extra = satArgs (arity - length args)
295 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
296 (collapsible oi) (M.fromList $ forceable oi) in
297 bind extra tm
298 where
299 satArgs n = map (\i -> sMN i "sat") [1..n]
301 bind [] tm = tm
302 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
304 -- Nat special cases
305 -- TODO: Would be nice if this was configurable in idris source!
306 doOpts (NS (UN z) [nat, prelude]) [] _ _
307 | z == txt "Z" && nat == txt "Nat" && prelude == txt "Prelude"
308 = Constant (BI 0)
309 doOpts (NS (UN s) [nat, prelude]) [k] _ _
310 | s == txt "S" && nat == txt "Nat" && prelude == txt "Prelude"
311 = App (App (P Ref (sUN "prim__addBigInt") Erased) k) (Constant (BI 1))
313 doOpts n args True _ = Erased
314 doOpts n args _ forceMap
315 | isnewtype oi = case args' of
316 [val] -> val
317 _ -> error $ "Can't happen (newtype not a singleton): " ++ show args'
318 | otherwise = mkApp ctor' args'
319 where
320 ctor' = (P (DCon tag (arity - forcedCnt forceMap)) n Erased)
321 args' = [t | (f, t) <- zip (forcedArgSeq oi) args, f /= Just Unconditional]