git.iko.soy/Idris-dev
on bf4e535 help sign in
Idris-dev/src/Idris/PartialEval.hs
310 lines 13 kB history raw
1{-# LANGUAGE PatternGuards #-}
3module Idris.PartialEval(partial_eval, getSpecApps, specType,
4 mkPE_TyDecl, mkPE_TermDecl, PEArgType(..),
5 pe_app, pe_def, pe_clauses, pe_simple) where
7import Idris.AbsSyntax
8import Idris.Delaborate
10import Idris.Core.TT
11import Idris.Core.Evaluate
13import Control.Monad.State
14import Control.Applicative
15import Data.Maybe
16import Debug.Trace
18-- | Data type representing binding-time annotations for partial evaluation of arguments
19data PEArgType = ImplicitS -- ^ Implicit static argument
20 | ImplicitD -- ^ Implicit dynamic argument
21 | ExplicitS -- ^ Explicit static argument
22 | ExplicitD -- ^ Explicit dynamic argument
23 | UnifiedD -- ^ Erasable dynamic argument (found under unification)
24 deriving (Eq, Show)
26-- | A partially evaluated function. pe_app captures the lhs of the
27-- new definition, pe_def captures the rhs, and pe_clauses is the
28-- specialised implementation.
29-- pe_simple is set if the result is always reducible, because in such a
30-- case we'll also need to reduce the static argument
31data PEDecl = PEDecl { pe_app :: PTerm, -- new application
32 pe_def :: PTerm, -- old application
33 pe_clauses :: [(PTerm, PTerm)], -- clauses of new application
34 pe_simple :: Bool -- if just one reducible clause
35 }
37-- | Partially evaluates given terms under the given context.
38-- It is an error if partial evaluation fails to make any progress.
39-- Making progress is defined as: all of the names given with explicit
40-- reduction limits (in practice, the function being specialised)
41-- must have reduced at least once.
42-- If we don't do this, we might end up making an infinite function after
43-- applying the transformation.
44partial_eval :: Context ->
45 [(Name, Maybe Int)] ->
46 [Either Term (Term, Term)] ->
47 Maybe [Either Term (Term, Term)]
48partial_eval ctxt ns_in tms = mapM peClause tms where
49 ns = squash ns_in
50 squash ((n, Just x) : ns)
51 | Just (Just y) <- lookup n ns
52 = squash ((n, Just (x + y)) : drop n ns)
53 | otherwise = (n, Just x) : squash ns
54 squash (n : ns) = n : squash ns
55 squash [] = []
57 drop n ((m, _) : ns) | n == m = ns
58 drop n (x : ns) = x : drop n ns
59 drop n [] = []
61 -- If the term is not a clause, it is simply kept as is
62 peClause (Left t) = Just $ Left t
63 -- If the term is a clause, specialise the right hand side
64 peClause (Right (lhs, rhs))
65 = let (rhs', reductions) = specialise ctxt [] (map toLimit ns) rhs in
66 do when (length tms == 1) $ checkProgress ns reductions
67 return (Right (lhs, rhs'))
69 -- TMP HACK until I do PE by WHNF rather than using main evaluator
70 toLimit (n, Nothing) | isTCDict n ctxt = (n, 2)
71 toLimit (n, Nothing) = (n, 65536) -- somewhat arbitrary reduction limit
72 toLimit (n, Just l) = (n, l)
74 checkProgress ns [] = return ()
75 checkProgress ns ((n, r) : rs)
76 | Just (Just start) <- lookup n ns
77 = if start <= 1 || r < start then checkProgress ns rs else Nothing
78 | otherwise = checkProgress ns rs
80-- | Specialises the type of a partially evaluated TT function returning
81-- a pair of the specialised type and the types of expected arguments.
82specType :: [(PEArgType, Term)] -> Type -> (Type, [(PEArgType, Term)])
83specType args ty = let (t, args') = runState (unifyEq args ty) [] in
84 (st (map fst args') t, map fst args')
85 where
86 -- Specialise static argument in type by let-binding provided value instead
87 -- of expecting it as a function argument
88 st ((ExplicitS, v) : xs) (Bind n (Pi _ t _) sc)
89 = Bind n (Let t v) (st xs sc)
90 st ((ImplicitS, v) : xs) (Bind n (Pi _ t _) sc)
91 = Bind n (Let t v) (st xs sc)
92 -- Erase argument from function type
93 st ((UnifiedD, _) : xs) (Bind n (Pi _ t _) sc)
94 = st xs sc
95 -- Keep types as is
96 st (_ : xs) (Bind n (Pi i t k) sc)
97 = Bind n (Pi i t k) (st xs sc)
98 st _ t = t
100 -- Erase implicit dynamic argument if existing argument shares it value,
101 -- by substituting the value of previous argument
102 unifyEq (imp@(ImplicitD, v) : xs) (Bind n (Pi i t k) sc)
103 = do amap <- get
104 case lookup imp amap of
105 Just n' ->
106 do put (amap ++ [((UnifiedD, Erased), n)])
107 sc' <- unifyEq xs (subst n (P Bound n' Erased) sc)
108 return (Bind n (Pi i t k) sc') -- erase later
109 _ -> do put (amap ++ [(imp, n)])
110 sc' <- unifyEq xs sc
111 return (Bind n (Pi i t k) sc')
112 unifyEq (x : xs) (Bind n (Pi i t k) sc)
113 = do args <- get
114 put (args ++ [(x, n)])
115 sc' <- unifyEq xs sc
116 return (Bind n (Pi i t k) sc')
117 unifyEq xs t = do args <- get
118 put (args ++ (zip xs (repeat (sUN "_"))))
119 return t
121-- | Creates an Idris type declaration given current state and a
122-- specialised TT function application type.
123-- Can be used in combination with the output of 'specType'.
124--
125-- This should: specialise any static argument position, then generalise
126-- over any function applications in the result.
127mkPE_TyDecl :: IState -> [(PEArgType, Term)] -> Type -> PTerm
128mkPE_TyDecl ist args ty = mkty args ty
129 where
130 mkty ((ExplicitD, v) : xs) (Bind n (Pi _ t k) sc)
131 = PPi expl n (delab ist (generaliseIn t)) (mkty xs sc)
132 mkty ((ImplicitD, v) : xs) (Bind n (Pi _ t k) sc)
133 | concreteClass ist t = mkty xs sc
134 | classConstraint ist t
135 = PPi constraint n (delab ist (generaliseIn t)) (mkty xs sc)
136 | otherwise = PPi impl n (delab ist (generaliseIn t)) (mkty xs sc)
138 mkty (_ : xs) t
139 = mkty xs t
140 mkty [] t = delab ist t
142 generaliseIn tm = evalState (gen tm) 0
144 gen tm | (P _ fn _, args) <- unApply tm,
145 isFnName fn (tt_ctxt ist)
146 = do nm <- get
147 put (nm + 1)
148 return (P Bound (sMN nm "spec") Erased)
149 gen (App f a) = App <$> gen f <*> gen a
150 gen tm = return tm
152-- | Checks if a given argument is a type class constraint argument
153classConstraint :: Idris.AbsSyntax.IState -> TT Name -> Bool
154classConstraint ist v
155 | (P _ c _, args) <- unApply v = case lookupCtxt c (idris_classes ist) of
156 [_] -> True
157 _ -> False
158 | otherwise = False
160-- | Checks if the given arguments of a type class constraint are all either constants
161-- or references (i.e. that it doesn't contain any complex terms).
162concreteClass :: IState -> TT Name -> Bool
163concreteClass ist v
164 | not (classConstraint ist v) = False
165 | (P _ c _, args) <- unApply v = all concrete args
166 | otherwise = False
167 where concrete (Constant _) = True
168 concrete tm | (P _ n _, args) <- unApply tm
169 = case lookupTy n (tt_ctxt ist) of
170 [_] -> all concrete args
171 _ -> False
172 | otherwise = False
174mkNewPats :: IState ->
175 [(Term, Term)] -> -- definition to specialise
176 [(PEArgType, Term)] -> -- arguments to specialise with
177 Name -> -- New name
178 Name -> -- Specialised function name
179 PTerm -> -- Default lhs
180 PTerm -> -- Default rhs
181 PEDecl
182-- If all of the dynamic positions on the lhs are variables (rather than
183-- patterns or constants) then we can just make a simple definition
184-- directly applying the specialised function, since we know the
185-- definition isn't going to block on any of the dynamic arguments
186-- in this case
187mkNewPats ist d ns newname sname lhs rhs | all dynVar (map fst d)
188 = PEDecl lhs rhs [(lhs, rhs)] True
189 where dynVar ap = case unApply ap of
190 (_, args) -> dynArgs ns args
191 dynArgs _ [] = True -- can definitely reduce from here
192 -- if Static, doesn't matter what the argument is
193 dynArgs ((ImplicitS, _) : ns) (a : as) = dynArgs ns as
194 dynArgs ((ExplicitS, _) : ns) (a : as) = dynArgs ns as
195 -- if Dynamic, it had better be a variable or we'll need to
196 -- do some more work
197 dynArgs (_ : ns) (V _ : as) = dynArgs ns as
198 dynArgs (_ : ns) (P _ _ _ : as) = dynArgs ns as
199 dynArgs _ _ = False -- and now we'll get stuck
200mkNewPats ist d ns newname sname lhs rhs =
201 PEDecl lhs rhs (map mkClause d) False
202 where
203 mkClause :: (Term, Term) -> (PTerm, PTerm)
204 mkClause (oldlhs, oldrhs)
205 = let (_, as) = unApply oldlhs
206 lhsargs = mkLHSargs [] ns as
207 lhs = PApp emptyFC (PRef emptyFC newname) lhsargs
208 rhs = PApp emptyFC (PRef emptyFC sname)
209 (mkRHSargs ns lhsargs) in
210 (lhs, rhs)
212 mkLHSargs _ [] _ = []
213 -- dynamics don't appear if they're implicit
214 mkLHSargs sub ((ExplicitD, t) : ns) (a : as)
215 = pexp (delab ist (substNames sub a)) : mkLHSargs sub ns as
216 mkLHSargs sub ((ImplicitD, _) : ns) (a : as)
217 = mkLHSargs sub ns as
218 mkLHSargs sub ((UnifiedD, _) : ns) (a : as)
219 = mkLHSargs sub ns as
220 -- statics get dropped in any case
221 mkLHSargs sub ((ImplicitS, t) : ns) (a : as)
222 = mkLHSargs (extend a t sub) ns as
223 mkLHSargs sub ((ExplicitS, t) : ns) (a : as)
224 = mkLHSargs (extend a t sub) ns as
225 mkLHSargs sub _ [] = [] -- no more LHS
227 extend (P _ n _) t sub = (n, t) : sub
228 extend _ _ sub = sub
230 mkRHSargs ((ExplicitS, t) : ns) as = pexp (delab ist t) : mkRHSargs ns as
231 mkRHSargs ((ExplicitD, t) : ns) (a : as) = a : mkRHSargs ns as
232 mkRHSargs (_ : ns) as = mkRHSargs ns as
233 mkRHSargs _ _ = []
235 mkSubst :: (Term, Term) -> Maybe (Name, Term)
236 mkSubst (P _ n _, t) = Just (n, t)
237 mkSubst _ = Nothing
239-- | Creates a new declaration for a specialised function application.
240-- Simple version at the moment: just create a version which is a direct
241-- application of the function to be specialised.
242-- More complex version to do: specialise the definition clause by clause
243mkPE_TermDecl :: IState -> Name -> Name ->
244 [(PEArgType, Term)] -> PEDecl
245mkPE_TermDecl ist newname sname ns
246 = let lhs = PApp emptyFC (PRef emptyFC newname) (map pexp (mkp ns))
247 rhs = eraseImps $ delab ist (mkApp (P Ref sname Erased) (map snd ns))
248 patdef = lookupCtxtExact sname (idris_patdefs ist)
249 newpats = case patdef of
250 Nothing -> PEDecl lhs rhs [(lhs, rhs)] True
251 Just d -> mkNewPats ist (getPats d) ns
252 newname sname lhs rhs in
253 newpats where
255 getPats (ps, _) = map (\(_, lhs, rhs) -> (lhs, rhs)) ps
257 mkp [] = []
258 mkp ((ExplicitD, tm) : tms) = delab ist tm : mkp tms
259 mkp (_ : tms) = mkp tms
261 eraseImps tm = mapPT deImp tm
263 deImp (PApp fc t as) = PApp fc t (map deImpArg as)
264 deImp t = t
266 deImpArg a@(PImp _ _ _ _ _) = a { getTm = Placeholder }
267 deImpArg a = a
269-- | Get specialised applications for a given function
270getSpecApps :: IState -> [Name] -> Term ->
271 [(Name, [(PEArgType, Term)])]
272getSpecApps ist env tm = ga env (explicitNames tm) where
274-- staticArg env True _ tm@(P _ n _) _ | n `elem` env = Just (True, tm)
275-- staticArg env True _ tm@(App f a) _ | (P _ n _, args) <- unApply tm,
276-- n `elem` env = Just (True, tm)
277 staticArg env x imp tm n
278 | x && imparg imp = (ImplicitS, tm)
279 | x = (ExplicitS, tm)
280 | imparg imp = (ImplicitD, tm)
281 | otherwise = (ExplicitD, (P Ref (sUN (show n ++ "arg")) Erased))
283 imparg (PExp _ _ _ _) = False
284 imparg _ = True
286 buildApp env [] [] _ _ = []
287 buildApp env (s:ss) (i:is) (a:as) (n:ns)
288 = let s' = staticArg env s i a n
289 ss' = buildApp env ss is as ns in
290 (s' : ss')
292 -- if we have a *defined* function that has static arguments,
293 -- it will become a specialised application
294 ga env tm@(App f a) | (P _ n _, args) <- unApply tm,
295 n `notElem` map fst (idris_metavars ist) =
296 ga env f ++ ga env a ++
297 case (lookupCtxtExact n (idris_statics ist),
298 lookupCtxtExact n (idris_implicits ist)) of
299 (Just statics, Just imps) ->
300 if (length statics == length args && or statics) then
301 case buildApp env statics imps args [0..] of
302 args -> [(n, args)]
303-- _ -> []
304 else []
305 _ -> []
306 ga env (Bind n (Let t v) sc) = ga env v ++ ga (n : env) sc
307 ga env (Bind n t sc) = ga (n : env) sc
308 ga env t = []