git.iko.soy/Idris-dev
on 047be9e help sign in
Idris-dev/src/Idris/DataOpts.hs
150 lines 5.5 kB history raw
1{-# LANGUAGE PatternGuards #-}
3module Idris.DataOpts where
5-- Forcing, detagging and collapsing
7import Idris.AbsSyntax
8import Core.TT
10import Data.List
11import Data.Maybe
12import Debug.Trace
14-- Calculate the forceable arguments to a constructor and update the set of
15-- optimisations
17forceArgs :: Name -> Type -> Idris ()
18forceArgs n t = do let fargs = force 0 t
19 i <- getIState
20 copt <- case lookupCtxt Nothing n (idris_optimisation i) of
21 [] -> return $ Optimise 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
28 where
29 force :: Int -> Term -> [Int]
30 force i (Bind _ (Pi _) sc)
31 = force (i + 1) $ instantiate (P Bound (MN i "?") Erased) sc
32 force _ sc@(App f a)
33 | (_, args) <- unApply sc
34 = nub $ concatMap guarded args
35 force _ _ = []
37 isF (P _ (MN force "?") _) = Just force
38 isF _ = Nothing
40 guarded :: Term -> [Int]
41 guarded t@(App f a)
42 | (P (TCon _ _) _ _, args) <- unApply t
43 = mapMaybe isF args ++ concatMap guarded args
44 | (P (DCon _ _) _ _, args) <- unApply t
45 = mapMaybe isF args ++ concatMap guarded args
46 guarded t = mapMaybe isF [t]
48-- Calculate whether a collection of constructors is collapsible
50collapseCons :: Name -> [(Name, Type)] -> Idris ()
51collapseCons ty cons =
52 do i <- getIState
53 return ()
55class Optimisable term where
56 applyOpts :: term -> Idris term
58instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
59 applyOpts (x, y) = do x' <- applyOpts x
60 y' <- applyOpts y
61 return (x', y')
63instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
64 applyOpts (v, x, y) = do x' <- applyOpts x
65 y' <- applyOpts y
66 return (v, x', y')
68instance Optimisable a => Optimisable [a] where
69 applyOpts = mapM applyOpts
71-- Raw is for compile time optimisation (before type checking)
72-- Term is for run time optimisation (after type checking, collapsing allowed)
74-- Compile time: no collapsing
76instance Optimisable Raw where
77 applyOpts t@(RApp f a)
78 | (Var n, args) <- raw_unapply t -- MAGIC HERE
79 = do args' <- mapM applyOpts args
80 i <- getIState
81 case lookupCtxt Nothing n (idris_optimisation i) of
82 (oi:_) -> return $ applyDataOpt oi n args'
83 _ -> return (raw_apply (Var n) args')
84 | otherwise = do f' <- applyOpts f
85 a' <- applyOpts a
86 return (RApp f' a')
87 applyOpts (RBind n b t) = do b' <- applyOpts b
88 t' <- applyOpts t
89 return (RBind n b' t')
90 applyOpts (RForce t) = applyOpts t
91 applyOpts t = return t
93instance Optimisable t => Optimisable (Binder t) where
94 applyOpts (Let t v) = do t' <- applyOpts t
95 v' <- applyOpts v
96 return (Let t' v')
97 applyOpts b = do t' <- applyOpts (binderTy b)
98 return (b { binderTy = t' })
101applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
102applyDataOpt oi n args
103 = let args' = zipWith doForce (map (\x -> x `elem` (forceable oi)) [0..])
104 args in
105 raw_apply (Var n) args'
106 where
107 doForce True a = RForce a
108 doForce False a = a
110-- Run-time: do everything
112instance Optimisable (TT Name) where
113 applyOpts t@(App f a)
114 | (c@(P (DCon t arity) n _), args) <- unApply t -- MAGIC HERE
115 = do args' <- mapM applyOpts args
116 i <- getIState
117 case lookupCtxt Nothing n (idris_optimisation i) of
118 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
119 _ -> return (mkApp c args')
120 | otherwise = do f' <- applyOpts f
121 a' <- applyOpts a
122 return (App f' a')
123 applyOpts (Bind n b t) = do b' <- applyOpts b
124 t' <- applyOpts t
125 return (Bind n b' t')
126 applyOpts t = return t
128-- Need to saturate arguments first to ensure that erasure happens uniformly
130applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
131applyDataOptRT oi n tag arity args
132 | length args == arity = doOpts n args (collapsible oi) (forceable oi)
133 | otherwise = let extra = satArgs (arity - length args)
134 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
135 (collapsible oi) (forceable oi) in
136 bind extra tm
137 where
138 satArgs n = map (\i -> MN i "sat") [1..n]
140 bind [] tm = tm
141 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
143 doOpts n args True f = Erased
144 doOpts n args _ forced
145 = let args' = filter keep (zip (map (\x -> x `elem` forced) [0..])
146 args) in
147 mkApp (P (DCon tag (arity - length forced)) n Erased) (map snd args')
149 keep (forced, _) = not forced