@@ -9,224 +9,48 @@ import Idris.AbsSyntaxTree
99import Idris.Core.TT
1010
1111import Control.Applicative
12−import qualified Data.IntMap as M
1312import Data.List
1413import Data.Maybe
1514import Debug.Trace
1615
17−type ForceMap = M.IntMap Forceability
18−
19−-- Calculate the forceable arguments to a constructor
20−-- and update the set of optimisations.
21−forceArgs :: Name -> Name -> [Int] -> Type -> Idris ()
22−forceArgs 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 [] [] [] 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
34−
35−getForcedArgs :: IState -> Name -> Type -> ForceMap
36−getForcedArgs ist typeName t = M.empty -- addCollapsibleArgs 0 t $ forcedInTarget 0 t
37− where
38− maxUnion = M.unionWith max
39−
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
43−
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
53−
54− -- a recursive occurrence with known indices is conditionally forceable
55− | tn == typeName
56− = M.insertWith max i Conditional alreadyForceable
57−
58− forceable _ = alreadyForceable
59−
60− isCollapsible :: Name -> Bool
61− isCollapsible n = case lookupCtxt n (idris_optimisation ist) of
62− [oi] -> collapsible oi
63− _ -> False
64−
65− addCollapsibleArgs _ _ fs = fs
66−
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
71−
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
76−
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
81−
82− unionMap :: (a -> ForceMap) -> [a] -> ForceMap
83− unionMap f = M.unionsWith max . map f
84−
85−-- Calculate whether a collection of constructors is collapsible
86−-- and update the state accordingly.
87−collapseCons :: Name -> [(Name, Type)] -> Idris ()
88−collapseCons tn ctors = do
89− ist <- getIState
90− case ctors of
91− _
92− | all (ctorCollapsible ist) ctors
93− , disjointTerms ctorTargetArgs
94− -> return () -- mapM_ setCollapsible (tn : map fst ctors)
95−
96− [(cn, ct)]
97− -> checkNewType ist cn ct
98−
99− _ -> return () -- nothing can be done
100− where
101− ctorTargetArgs = map (snd . unApply . getRetTy . snd) ctors
102−
103− ctorArity :: Type -> Int
104− ctorArity = length . getArgTys
105−
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
112−
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
123−
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 [] [] [] False
133− let opts = addDef n oi (idris_optimisation i)
134− putIState (i { idris_optimisation = opts })
135− addIBC (IBCOpt n)
136−
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
145−
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
158−
15916class Optimisable term where
16017 applyOpts :: term -> Idris term
161− stripCollapsed :: term -> Idris term
16218
16319instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
16420 applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
165− stripCollapsed (x, y) = (,) <$> stripCollapsed x <*> stripCollapsed y
16621
16722instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
16823 applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
169− stripCollapsed (v, x, y) = (,,) v <$> stripCollapsed x <*> stripCollapsed y
17024
17125instance Optimisable a => Optimisable [a] where
17226 applyOpts = mapM applyOpts
173− stripCollapsed = mapM stripCollapsed
17427
17528instance Optimisable a => Optimisable (Either a (a, a)) where
17629 applyOpts (Left t) = Left <$> applyOpts t
17730 applyOpts (Right t) = Right <$> applyOpts t
178− stripCollapsed (Left t) = Left <$> stripCollapsed t
179− stripCollapsed (Right t) = Right <$> stripCollapsed t
18031
18132-- Raw is for compile time optimisation (before type checking)
18233-- Term is for run time optimisation (after type checking, collapsing allowed)
183−
18434-- Compile time: no collapsing
18535
18636instance Optimisable Raw where
18737 applyOpts t@(RApp f a)
18838 | (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'
39+ = raw_apply (Var n) <$> mapM applyOpts args
19440 | otherwise = RApp <$> applyOpts f <*> applyOpts a
19541
19642 applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
19743 applyOpts (RForce t) = applyOpts t
19844 applyOpts t = return t
19945
200− stripCollapsed t = return t
201−
20246-- Erase types (makes ibc smaller, and we don't need them)
20347instance Optimisable (Binder (TT Name)) where
20448 applyOpts (Let t v) = Let <$> return Erased <*> applyOpts v
20549 applyOpts b = return (b { binderTy = Erased })
206− stripCollapsed (Let t v) = Let <$> return Erased <*> stripCollapsed v
207− stripCollapsed b = return (b { binderTy = Erased })
20850
20951instance Optimisable (Binder Raw) where
21052 applyOpts b = do t' <- applyOpts (binderTy b)
21153 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' })
215−
216−forcedArgSeq :: OptInfo -> [Maybe Forceability]
217−forcedArgSeq oi = map (\i -> M.lookup i forceMap) [0..]
218− where
219− forceMap = M.fromList $ forceable oi
220−
221−forcedCnt :: ForceMap -> Int
222−forcedCnt = length . filter (== Unconditional) . M.elems
223−
224−applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
225−applyDataOpt 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
23054
23155-- Run-time: do everything
23256
@@ -249,54 +73,22 @@ instance Optimisable (TT Name) where
24973 | fn == txt "toIntegerNat" && mod == prel
25074 = return (App (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
25175 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
76+ = return $ applyDataOptRT n t arity []
25677 applyOpts t@(App f a)
25778 | (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)
79+ = applyDataOptRT n t arity <$> mapM applyOpts args
80+ | otherwise = App <$> applyOpts f <*> applyOpts a
81+ applyOpts (Bind n b t) = Bind n <$> applyOpts b <*> applyOpts t
82+ applyOpts (Proj t i) = Proj <$> applyOpts t <*> pure i
27183 applyOpts t = return t
27284
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
290−
291−-- Need to saturate arguments first to ensure that erasure happens uniformly
292−
293−applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
294−applyDataOptRT oi n tag arity args
295− | length args == arity = doOpts n args (collapsible oi) (M.fromList $ forceable oi)
85+-- Need to saturate arguments first to ensure that optimisation happens uniformly
86+applyDataOptRT :: Name -> Int -> Int -> [Term] -> Term
87+applyDataOptRT n tag arity args
88+ | length args == arity = doOpts n args
29689 | otherwise = let extra = satArgs (arity - length args)
29790 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
298− (collapsible oi) (M.fromList $ forceable oi) in
299− bind extra tm
91+ in bind extra tm
30092 where
30193 satArgs n = map (\i -> sMN i "sat") [1..n]
30294
@@ -305,20 +97,11 @@ applyDataOptRT oi n tag arity args
30597
30698 -- Nat special cases
30799 -- TODO: Would be nice if this was configurable in idris source!
308− doOpts (NS (UN z) [nat, prelude]) [] _ _
100+ doOpts (NS (UN z) [nat, prelude]) []
309101 | z == txt "Z" && nat == txt "Nat" && prelude == txt "Prelude"
310102 = Constant (BI 0)
311− doOpts (NS (UN s) [nat, prelude]) [k] _ _
103+ doOpts (NS (UN s) [nat, prelude]) [k]
312104 | s == txt "S" && nat == txt "Nat" && prelude == txt "Prelude"
313105 = App (App (P Ref (sUN "prim__addBigInt") Erased) k) (Constant (BI 1))
314106
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]
324−
107+ doOpts n args = mkApp (P (DCon tag arity) n Erased) args