git.iko.soy/Idris-dev
on b63e9bb help sign in
Idris-dev/src/Idris/DataOpts.hs
145 lines 5.3 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 [a] where
64 applyOpts = mapM applyOpts
66-- Raw is for compile time optimisation (before type checking)
67-- Term is for run time optimisation (after type checking, collapsing allowed)
69-- Compile time: no collapsing
71instance Optimisable Raw where
72 applyOpts t@(RApp f a)
73 | (Var n, args) <- raw_unapply t -- MAGIC HERE
74 = do args' <- mapM applyOpts args
75 i <- getIState
76 case lookupCtxt Nothing n (idris_optimisation i) of
77 (oi:_) -> return $ applyDataOpt oi n args'
78 _ -> return (raw_apply (Var n) args')
79 | otherwise = do f' <- applyOpts f
80 a' <- applyOpts a
81 return (RApp f' a')
82 applyOpts (RBind n b t) = do b' <- applyOpts b
83 t' <- applyOpts t
84 return (RBind n b' t')
85 applyOpts (RForce t) = applyOpts t
86 applyOpts t = return t
88instance Optimisable t => Optimisable (Binder t) where
89 applyOpts (Let t v) = do t' <- applyOpts t
90 v' <- applyOpts v
91 return (Let t' v')
92 applyOpts b = do t' <- applyOpts (binderTy b)
93 return (b { binderTy = t' })
96applyDataOpt :: OptInfo -> Name -> [Raw] -> Raw
97applyDataOpt oi n args
98 = let args' = zipWith doForce (map (\x -> x `elem` (forceable oi)) [0..])
99 args in
100 raw_apply (Var n) args'
101 where
102 doForce True a = RForce a
103 doForce False a = a
105-- Run-time: do everything
107instance Optimisable (TT Name) where
108 applyOpts t@(App f a)
109 | (c@(P (DCon t arity) n _), args) <- unApply t -- MAGIC HERE
110 = do args' <- mapM applyOpts args
111 i <- getIState
112 case lookupCtxt Nothing n (idris_optimisation i) of
113 (oi:_) -> do return $ applyDataOptRT oi n t arity args'
114 _ -> return (mkApp c args')
115 | otherwise = do f' <- applyOpts f
116 a' <- applyOpts a
117 return (App f' a')
118 applyOpts (Bind n b t) = do b' <- applyOpts b
119 t' <- applyOpts t
120 return (Bind n b' t')
121 applyOpts t = return t
123-- Need to saturate arguments first to ensure that erasure happens uniformly
125applyDataOptRT :: OptInfo -> Name -> Int -> Int -> [Term] -> Term
126applyDataOptRT oi n tag arity args
127 | length args == arity = doOpts n args (collapsible oi) (forceable oi)
128 | otherwise = let extra = satArgs (arity - length args)
129 tm = doOpts n (args ++ map (\n -> P Bound n Erased) extra)
130 (collapsible oi) (forceable oi) in
131 bind extra tm
132 where
133 satArgs n = map (\i -> MN i "sat") [1..n]
135 bind [] tm = tm
136 bind (n:ns) tm = Bind n (Lam Erased) (pToV n (bind ns tm))
138 doOpts n args True f = Erased
139 doOpts n args _ forced
140 = let args' = filter keep (zip (map (\x -> x `elem` forced) [0..])
141 args) in
142 mkApp (P (DCon tag (arity - length forced)) n Erased) (map snd args')
144 keep (forced, _) = not forced