@@ -5,78 +5,118 @@ module Idris.DataOpts where
55-- Forcing, detagging and collapsing
66
77import Idris.AbsSyntax
8+import Idris.AbsSyntaxTree
89import Core.TT
910
11+import Control.Applicative
12+import qualified Data.IntMap.Strict as M
1013import Data.List
1114import Data.Maybe
1215import Debug.Trace
1316
14−-- Calculate the forceable arguments to a constructor and update the set of
15−-- optimisations
16−
17−forceArgs :: Name -> Type -> Idris ()
18−forceArgs 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
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 -> Type -> Idris ()
22+forceArgs 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
32+
33+getForcedArgs :: IState -> Name -> Type -> ForceMap
34+getForcedArgs ist typeName t = addCollapsibleArgs 0 t $ forcedInTarget 0 t
2835 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 _ _ _ = []
38−
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
45−
46− isF (P _ (MN force "?") _) = Just force
47− isF _ = Nothing
48−
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]
36+ maxUnion = M.unionWith max
5637
57−-- Calculate whether a collection of constructors is collapsible
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
41+
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
51+
52+ -- a recursive occurrence with known indices is conditionally forceable
53+ | tn == typeName
54+ = M.insertWith max i Conditional alreadyForceable
55+
56+ forceable _ = alreadyForceable
57+
58+ isCollapsible :: Name -> Bool
59+ isCollapsible n = case lookupCtxt n (idris_optimisation ist) of
60+ [oi] -> collapsible oi
61+ _ -> False
62+
63+ addCollapsibleArgs _ _ fs = fs
5864
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
69+
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
74+
75+ bareArg :: Term -> ForceMap
76+ bareArg (P _ (MN i "ctor_arg") _) = M.singleton i Unconditional
77+ bareArg _ = M.empty
78+
79+ unionMap :: (a -> ForceMap) -> [a] -> ForceMap
80+ unionMap f = M.unionsWith max . map f
81+
82+-- Calculate whether a collection of constructors is collapsible
83+-- and update the state accordingly.
5984collapseCons :: Name -> [(Name, Type)] -> Idris ()
60−collapseCons 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'
85+collapseCons 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)
92+
93+ [(cn, ct)]
94+ -> checkNewType ist cn ct
95+
96+ _ -> return () -- nothing can be done
6697 where
98+ ctorTargetArgs = map (snd . unApply . getRetTy . snd) ctors
99+
100+ ctorArity :: Type -> Int
101+ ctorArity = length . getArgTys
102+
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
109+
67110 -- 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 ()
78−
79− checkNewType _ = return ()
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
80120
81121 setCollapsible :: Name -> Idris ()
82122 setCollapsible n
@@ -91,79 +131,49 @@ collapseCons ty cons =
91131 putIState (i { idris_optimisation = opts })
92132 addIBC (IBCOpt n)
93133
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
107−
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 ()
116−
117− disjoint :: [[Term]] -> Bool
118− disjoint [] = True
119− disjoint [x] = True
120− disjoint (x : xs) = anyDisjoint x xs && disjoint xs
121−
122− anyDisjoint x [] = True
123− anyDisjoint x (y : ys) = disjointCons x y
124−
125− disjointCons [] [] = False
126− disjointCons [] y = False
127− disjointCons x [] = False
128− disjointCons (x : xs) (y : ys)
129− = disjointCon x y || disjointCons xs ys
130−
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
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
142+
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
136155
137156class Optimisable term where
138157 applyOpts :: term -> Idris term
139158 stripCollapsed :: term -> Idris term
140159
141160instance (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')
148−
161+ applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
162+ stripCollapsed (x, y) = (,) <$> stripCollapsed x <*> stripCollapsed y
149163
150164instance (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')
165+ applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
166+ stripCollapsed (v, x, y) = (,,) v <$> stripCollapsed x <*> stripCollapsed y
157167
158168instance Optimisable a => Optimisable [a] where
159169 applyOpts = mapM applyOpts
160170 stripCollapsed = mapM stripCollapsed
161171
162172instance 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'
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
167177
168178-- Raw is for compile time optimisation (before type checking)
169179-- Term is for run time optimisation (after type checking, collapsing allowed)
@@ -175,41 +185,39 @@ instance Optimisable Raw where
175185 | (Var n, args) <- raw_unapply t -- MAGIC HERE
176186 = do args' <- mapM applyOpts args
177187 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+ 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
192+
193+ applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
194+ applyOpts (RForce t) = applyOpts t
188195 applyOpts t = return t
189196
190197 stripCollapsed t = return t
191198
192199instance Optimisable t => Optimisable (Binder t) where
193− applyOpts (Let t v) = do t' <- applyOpts t
194− v' <- applyOpts v
195− return (Let t' v')
200+ applyOpts (Let t v) = Let <$> applyOpts t <*> applyOpts v
196201 applyOpts b = do t' <- applyOpts (binderTy b)
197202 return (b { binderTy = t' })
198− stripCollapsed (Let t v) = do t' <- stripCollapsed t
199− v' <- stripCollapsed v
200− return (Let t' v')
203+ stripCollapsed (Let t v) = Let <$> stripCollapsed t <*> stripCollapsed v
201204 stripCollapsed b = do t' <- stripCollapsed (binderTy b)
202205 return (b { binderTy = t' })
203206
207+forcedArgSeq :: OptInfo -> [Maybe Forceability]
208+forcedArgSeq oi = map (\i -> M.lookup i forceMap) [0..]
209+ where
210+ forceMap = M.fromList $ forceable oi
211+
212+forcedCnt :: ForceMap -> Int
213+forcedCnt = length . filter (== Unconditional) . M.elems
204214
205215applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
206−applyDataOpt oi n args
207− = let args' = zipWith doForce (map (\x -> x `elem` (forceable oi)) [0..])
208− args in
209− raw_apply (Var n) args'
216+applyDataOpt oi n args
217+ = raw_apply (Var n) $ zipWith doForce (forcedArgSeq oi) args
210218 where
211− doForce True a = RForce a
212− doForce False a = a
219+ doForce (Just Unconditional) a = RForce a
220+ doForce _ a = a
213221
214222-- Run-time: do everything
215223
@@ -268,10 +276,10 @@ instance Optimisable (TT Name) where
268276
269277applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
270278applyDataOptRT oi n tag arity args
271− | length args == arity = doOpts n args (collapsible oi) (forceable oi)
279+ | length args == arity = doOpts n args (collapsible oi) (M.fromList $ forceable oi)
272280 | otherwise = let extra = satArgs (arity - length args)
273281 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
274− (collapsible oi) (forceable oi) in
282+ (collapsible oi) (M.fromList $ forceable oi) in
275283 bind extra tm
276284 where
277285 satArgs n = map (\i -> MN i "sat") [1..n]
@@ -285,18 +293,12 @@ applyDataOptRT oi n tag arity args
285293 doOpts (NS (UN "S") ["Nat", "Prelude"]) [k] _ _
286294 = App (App (P Ref (UN "prim__addBigInt") Erased) k) (Constant (BI 1))
287295
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')
299−
300− keep (forced, _) = not forced
301−
302−
296+ doOpts n args True _ = Erased
297+ doOpts n args _ forceMap
298+ | isnewtype oi = case args' of
299+ [val] -> val
300+ _ -> error $ "Can't happen (newtype not a singleton): " ++ show args'
301+ | otherwise = mkApp ctor' args'
302+ where
303+ ctor' = (P (DCon tag (arity - forcedCnt forceMap)) n Erased)
304+ args' = [t | (f, t) <- zip (forcedArgSeq oi) args, f /= Just Unconditional]