git.iko.soy/Idris-dev
on db306ca help sign in
Idris-dev/src/Idris/DataOpts.hs
324 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 -> [Int] -> Type -> Idris ()
22forceArgs typeName n expforce 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 ++
29 zip expforce (repeat Unconditional) })
30 (idris_optimisation ist)
31 putIState (ist { idris_optimisation = opts })
32 addIBC (IBCOpt n)
33 iLOG $ "Forced: " ++ show n ++ " " ++ show fargs ++ "\n from " ++ show t
35getForcedArgs :: IState -> Name -> Type -> ForceMap
36getForcedArgs ist typeName t = addCollapsibleArgs 0 t $ forcedInTarget 0 t
37 where
38 maxUnion = M.unionWith max
40 -- Label all occurrences of the variable bound in Pi in the rest of
41 -- the term with the number i so that we can recognize them anytime later.
42 label i = instantiate $ P Bound (sMN i "ctor_arg") Erased
44 addCollapsibleArgs :: Int -> Type -> ForceMap -> ForceMap
45 addCollapsibleArgs i (Bind vn (Pi ty) rest) alreadyForceable
46 = addCollapsibleArgs (i+1) (label i rest) (forceable $ unApply ty)
47 where
48 -- forceable takes an un-applied type of a ctor argument
49 forceable (P _ tn _, args)
50 -- if `ty' is collapsible, the argument is unconditionally forceable
51 | isCollapsible tn
52 = M.insert i Unconditional alreadyForceable
54 -- a recursive occurrence with known indices is conditionally forceable
55 | tn == typeName
56 = M.insertWith max i Conditional alreadyForceable
58 forceable _ = alreadyForceable
60 isCollapsible :: Name -> Bool
61 isCollapsible n = case lookupCtxt n (idris_optimisation ist) of
62 [oi] -> collapsible oi
63 _ -> False
65 addCollapsibleArgs _ _ fs = fs
67 forcedInTarget :: Int -> Type -> ForceMap
68 forcedInTarget i (Bind _ (Pi _) rest) = forcedInTarget (i+1) (label i rest)
69 forcedInTarget i t@(App f a) | (_, as) <- unApply t = unionMap guardedArgs as
70 forcedInTarget _ _ = M.empty
72 guardedArgs :: Term -> ForceMap
73 guardedArgs t@(App f a) | (P (DCon _ _) _ _, args) <- unApply t
74 = unionMap bareArg args `maxUnion` unionMap guardedArgs args
75 guardedArgs t = bareArg t
77 bareArg :: Term -> ForceMap
78 bareArg (P _ (MN i ctor_arg) _)
79 | ctor_arg == txt "ctor_arg" = M.singleton i Unconditional
80 bareArg _ = M.empty
82 unionMap :: (a -> ForceMap) -> [a] -> ForceMap
83 unionMap f = M.unionsWith max . map f
85-- Calculate whether a collection of constructors is collapsible
86-- and update the state accordingly.
87collapseCons :: Name -> [(Name, Type)] -> Idris ()
88collapseCons tn ctors = do
89 ist <- getIState
90 case ctors of
91 _
92 | all (ctorCollapsible ist) ctors
93 , disjointTerms ctorTargetArgs
94 -> mapM_ setCollapsible (tn : map fst ctors)
96 [(cn, ct)]
97 -> checkNewType ist cn ct
99 _ -> return () -- nothing can be done
100 where
101 ctorTargetArgs = map (snd . unApply . getRetTy . snd) ctors
103 ctorArity :: Type -> Int
104 ctorArity = length . getArgTys
106 ctorCollapsible :: IState -> (Name, Type) -> Bool
107 ctorCollapsible ist (n, t) = all (`M.member` forceMap) [0 .. ctorArity t - 1]
108 where
109 forceMap = case lookupCtxt n (idris_optimisation ist) of
110 oi:_ -> M.fromList $ forceable oi
111 _ -> M.empty
113 -- one constructor; if one remaining argument, treat as newtype
114 checkNewType :: IState -> Name -> Type -> Idris ()
115 checkNewType ist cn ct
116 | oi:_ <- lookupCtxt cn opt
117 , length (getArgTys ct) == 1 + forcedCnt (M.fromList $ forceable oi)
118 = putIState ist{ idris_optimisation = opt' oi }
119 | otherwise = return ()
120 where
121 opt = idris_optimisation ist
122 opt' oi = addDef cn oi{ isnewtype = True } opt
124 setCollapsible :: Name -> Idris ()
125 setCollapsible n
126 = do i <- getIState
127 iLOG $ show n ++ " collapsible"
128 case lookupCtxt n (idris_optimisation i) of
129 (oi:_) -> do let oi' = oi { collapsible = True }
130 let opts = addDef n oi' (idris_optimisation i)
131 putIState (i { idris_optimisation = opts })
132 [] -> do let oi = Optimise True False [] []
133 let opts = addDef n oi (idris_optimisation i)
134 putIState (i { idris_optimisation = opts })
135 addIBC (IBCOpt n)
137 disjointTerms :: [[Term]] -> Bool
138 disjointTerms [] = True
139 disjointTerms [xs] = True
140 disjointTerms (xs : xss) =
141 -- xs is disjoint with every pattern from xss
142 all (or . zipWith disjoint xs) xss
143 -- and xss is pairwise disjoint, too
144 && disjointTerms xss
146 -- Return True if the two patterns are provably disjoint.
147 -- Return False if they're not or if unsure.
148 disjoint :: Term -> Term -> Bool
149 disjoint x y = case (cx, cy) of
150 -- data constructors -> compare their names
151 (P (DCon _ _) nx _, P (DCon _ _) ny _)
152 | nx /= ny -> True
153 | otherwise -> or $ zipWith disjoint xargs yargs
154 _ -> False
155 where
156 (cx, xargs) = unApply x
157 (cy, yargs) = unApply y
159class Optimisable term where
160 applyOpts :: term -> Idris term
161 stripCollapsed :: term -> Idris term
163instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
164 applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
165 stripCollapsed (x, y) = (,) <$> stripCollapsed x <*> stripCollapsed y
167instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
168 applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
169 stripCollapsed (v, x, y) = (,,) v <$> stripCollapsed x <*> stripCollapsed y
171instance Optimisable a => Optimisable [a] where
172 applyOpts = mapM applyOpts
173 stripCollapsed = mapM stripCollapsed
175instance Optimisable a => Optimisable (Either a (a, a)) where
176 applyOpts (Left t) = Left <$> applyOpts t
177 applyOpts (Right t) = Right <$> applyOpts t
178 stripCollapsed (Left t) = Left <$> stripCollapsed t
179 stripCollapsed (Right t) = Right <$> stripCollapsed t
181-- Raw is for compile time optimisation (before type checking)
182-- Term is for run time optimisation (after type checking, collapsing allowed)
184-- Compile time: no collapsing
186instance Optimisable Raw where
187 applyOpts t@(RApp f a)
188 | (Var n, args) <- raw_unapply t -- MAGIC HERE
189 = do args' <- mapM applyOpts args
190 i <- getIState
191 return $ case lookupCtxt n (idris_optimisation i) of
192 oi:_ -> applyDataOpt oi n args'
193 _ -> raw_apply (Var n) args'
194 | otherwise = RApp <$> applyOpts f <*> applyOpts a
196 applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
197 applyOpts (RForce t) = applyOpts t
198 applyOpts t = return t
200 stripCollapsed t = return t
202-- Erase types (makes ibc smaller, and we don't need them)
203instance Optimisable (Binder (TT Name)) where
204 applyOpts (Let t v) = Let <$> return Erased <*> applyOpts v
205 applyOpts b = return (b { binderTy = Erased })
206 stripCollapsed (Let t v) = Let <$> return Erased <*> stripCollapsed v
207 stripCollapsed b = return (b { binderTy = Erased })
209instance Optimisable (Binder Raw) where
210 applyOpts b = do t' <- applyOpts (binderTy b)
211 return (b { binderTy = t' })
212 stripCollapsed (Let t v) = Let <$> stripCollapsed t <*> stripCollapsed v
213 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
214 return (b { binderTy = t' })
216forcedArgSeq :: OptInfo -> [Maybe Forceability]
217forcedArgSeq oi = map (\i -> M.lookup i forceMap) [0..]
218 where
219 forceMap = M.fromList $ forceable oi
221forcedCnt :: ForceMap -> Int
222forcedCnt = length . filter (== Unconditional) . M.elems
224applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
225applyDataOpt oi n args
226 = raw_apply (Var n) $ zipWith doForce (forcedArgSeq oi) args
227 where
228 doForce (Just Unconditional) a = RForce a
229 doForce _ a = a
231-- Run-time: do everything
233prel = [txt "Nat", txt "Prelude"]
235instance Optimisable (TT Name) where
236 applyOpts (P _ (NS (UN fn) mod) _)
237 | fn == txt "plus" && mod == prel
238 = return (P Ref (sUN "prim__addBigInt") Erased)
239 applyOpts (P _ (NS (UN fn) mod) _)
240 | fn == txt "mult" && mod == prel
241 = return (P Ref (sUN "prim__mulBigInt") Erased)
242 applyOpts (App (P _ (NS (UN fn) mod) _) x)
243 | fn == txt "fromIntegerNat" && mod == prel
244 = applyOpts x
245 applyOpts (P _ (NS (UN fn) mod) _)
246 | fn == txt "fromIntegerNat" && mod == prel
247 = return (App (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
248 applyOpts (P _ (NS (UN fn) mod) _)
249 | fn == txt "toIntegerNat" && mod == prel
250 = return (App (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
251 applyOpts c@(P (DCon t arity) n _)
252 = do i <- getIState
253 case lookupCtxt n (idris_optimisation i) of
254 (oi:_) -> return $ applyDataOptRT oi n t arity []
255 _ -> return c
256 applyOpts t@(App f a)
257 | (c@(P (DCon t arity) n _), args) <- unApply t
258 = do args' <- mapM applyOpts args
259 i <- getIState
260 case lookupCtxt n (idris_optimisation i) of
261 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
262 _ -> return (mkApp c args')
263 | otherwise = do f' <- applyOpts f
264 a' <- applyOpts a
265 return (App f' a')
266 applyOpts (Bind n b t) = do b' <- applyOpts b
267 t' <- applyOpts t
268 return (Bind n b' t')
269 applyOpts (Proj t i) = do t' <- applyOpts t
270 return (Proj t' i)
271 applyOpts t = return t
273 stripCollapsed (Bind n (PVar x) t) | (P _ ty _, _) <- unApply x
274 = do i <- getIState
275 -- NOTE: This assumes that 'ty' is in normal form, which it
276 -- has to be before now because we're not keeping track of
277 -- an environment so we can't do it here.
278 case lookupCtxt ty (idris_optimisation i) of
279 [oi] -> if collapsible oi
280 then do t' <- stripCollapsed t
281 return (Bind n (PVar x) (instantiate Erased t'))
282 else do t' <- stripCollapsed t
283 return (Bind n (PVar x) t')
284 _ -> do t' <- stripCollapsed t
285 return (Bind n (PVar x) t')
286 stripCollapsed (Bind n (PVar x) t)
287 = do t' <- stripCollapsed t
288 return (Bind n (PVar x) t')
289 stripCollapsed t = return t
291-- Need to saturate arguments first to ensure that erasure happens uniformly
293applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
294applyDataOptRT oi n tag arity args
295 | length args == arity = doOpts n args (collapsible oi) (M.fromList $ forceable oi)
296 | otherwise = let extra = satArgs (arity - length args)
297 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
298 (collapsible oi) (M.fromList $ forceable oi) in
299 bind extra tm
300 where
301 satArgs n = map (\i -> sMN i "sat") [1..n]
303 bind [] tm = tm
304 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
306 -- Nat special cases
307 -- TODO: Would be nice if this was configurable in idris source!
308 doOpts (NS (UN z) [nat, prelude]) [] _ _
309 | z == txt "Z" && nat == txt "Nat" && prelude == txt "Prelude"
310 = Constant (BI 0)
311 doOpts (NS (UN s) [nat, prelude]) [k] _ _
312 | s == txt "S" && nat == txt "Nat" && prelude == txt "Prelude"
313 = App (App (P Ref (sUN "prim__addBigInt") Erased) k) (Constant (BI 1))
315 doOpts n args True _ = Erased
316 doOpts n args _ forceMap
317 | isnewtype oi = case args' of
318 [val] -> val
319 _ -> error $ "Can't happen (newtype not a singleton): " ++ show args'
320 | otherwise = mkApp ctor' args'
321 where
322 ctor' = (P (DCon tag (arity - forcedCnt forceMap)) n Erased)
323 args' = [t | (f, t) <- zip (forcedArgSeq oi) args, f /= Just Unconditional]