idris-0.1.7: Idris/PartialEval.lhs
Various bits and pieces to help do partial evaluation better.
This module doesn't do PE itself - rather, it sets up the %transform, %spec
and %freeze annotations.
> module Idris.PartialEval(partialeval, addPEdefs, staticDecls,
> getNewDefs) where
> import Idris.AbsSyntax
> import Ivor.TT as TT
> import Maybe
> import List
> import Control.Monad.State
> import Debug.Trace
> type NewTrans = [(ViewTerm, ViewTerm)]
> type NewFreeze = [Name]
> type NewDefs = [(Name, ViewTerm, Patterns, NewFreeze, (ViewTerm, ViewTerm))]
> type StaticState = (NewDefs, StaticUsed, NewTrans, NewFreeze, Int)
For everything in the context, look for applications of PEable functions,
with 'findStatic' below.
For each one, add a new definition, %spec it, %freeze it and add a
%transform rule that applies the new definition backwards.
For a definition d which uses PEable functions, add the new definitions
to the new context *before* adding d, followed by the %transform and %freeze
for each new definition, followed by d.
addPEdefs :: Statics -> IdrisState -> Ctxt IvorFun -> Ctxt IvorFun
addPEdefs sts ist raw = addpes newCtxt [] (ctxtAlist raw)
where addpes acc stu [] = acc
addpes acc stu ((n,i):is) =
let (i', defs, stu', _, _) = getPEdefs n sts ist raw stu i in
addpes (addEntry acc [] n i') stu' is
> getPEdefs :: Id -> Statics -> IdrisState ->
> Ctxt IvorFun -> StaticUsed -> IvorFun ->
> (IvorFun, NewDefs, StaticUsed, NewTrans, NewFreeze)
> getPEdefs n sts ist raw stu i =
> case ivorDef i of
> Just d -> let (def', (nds, stused, ts, fs, _)) = runState (getPEdef n sts ist raw stu d) ([], stu, [], [], 0) in
> (i { ivorDef = Just def' }, nds, stused, ts, fs)
> Nothing -> (i, [], stu, [], [])
> getNewDefs :: Id -> Statics -> IdrisState ->
> Ctxt IvorFun -> StaticUsed -> IvorDef ->
> (IvorDef, NewDefs, StaticUsed, NewTrans, NewFreeze)
> getNewDefs n sts ist raw stu d =
> let (d', (sts', stu', ts, fs, _)) = runState (getPEdef n sts ist raw stu d) ([],stu,[],[],0) in
> (d', sts', stu', ts, fs)
> getPEdef :: Id -> Statics -> IdrisState -> Ctxt IvorFun ->
> StaticUsed -> IvorDef ->
> State StaticState IvorDef
> getPEdef n sts ist raw stu (PattDef (Patterns ps)) =
> do ps' <- mapM pepat ps
> return (PattDef (Patterns ps'))
> where
> pepat (PClause args bs ret)
> = do ret' <- findStatic n sts ist raw ret
> return (PClause args bs ret')
> pepat (PWithClause p args sc (Patterns ps))
> = do sc' <- findStatic n sts ist raw sc
> ps' <- mapM pepat ps
> return (PWithClause p args sc' (Patterns ps'))
> getPEdef n sts ist raw stu (SimpleDef p) =
> do p' <- findStatic n sts ist raw p
> return (SimpleDef p')
> getPEdef _ _ _ _ _ x = return x
Given a list of functions with static arguments, and a term...
Look for applications of that function in the term.
Replace them with a partially evaluated version (look in StaticsUsed first
to see if it's already been done.)
Return: the new term, the new definitions added, and an updated cached of PEed
functions.
> findStatic :: Id -> Statics -> IdrisState -> Ctxt IvorFun -> ViewTerm ->
> State StaticState ViewTerm
> findStatic n sts ist raw vt = fs [] vt
> where
> fs stk (App f a) = do a' <- fs [] a
> fs (a':stk) f
> fs stk (Lambda n ty sc) = do sc' <- fs [] sc
> freturn (Lambda n ty sc') stk
> fs stk (Let n ty v sc) = do v' <- fs [] v
> sc' <- fs [] sc
> freturn (Let n ty v' sc') stk
> fs stk (Annotation a vt) = do vt' <- fs stk vt
> return (Annotation a vt')
Don't bother with PE inside types
> fs stk x = freturn x stk
> freturn (Name _ f) args
> | Just (sts,arity,ty) <- lookup f sts
> = if length args == arity
> then papply sts f ty args
> else if length args > arity
> then do let (args', rest) = (take arity args,
> drop arity args)
> app <- papply sts f ty args'
> return (apply app rest)
> else return (apply (Name Unknown f) args)
> freturn f args = return (apply f args)
Check the arguments in static position are indeed statically known.
> papply sts f ty args
> | all (known args) sts = -- trace (show (f, args, ty)) $
> do let knownArgs = mkArgs sts 0 args
> let ty' = newTy ty knownArgs
> addDef f ty' knownArgs
> return (apply (Name Unknown f) args)
> papply sts f ty args = return (apply (Name Unknown f) args)
Pull out the arguments that are statically known
> mkArgs sts _ [] = []
> mkArgs sts i (a:args) | i `elem` sts = Right a : (mkArgs sts (i+1) args)
> | otherwise = Left a : mkArgs sts (i+1) args
> newTy (Forall n ty sc) (Left _ : rest)
> = Forall n ty (newTy sc rest)
> newTy (Forall n ty sc) (Right val : rest)
> = newTy (subst n val sc) rest
> newTy (Annotation a x) rest = Annotation a (newTy x rest)
> newTy x _ = x
> known args i = all nknown (namesTypesIn (args!!!(i,"known fail")))
> nknown x@(_,Free) = True
> nknown x@(_,DataCon) = True
> nknown x@(_,TypeCon) = True
> nknown x = False
> addDef f ty args =
> do (nds, used, ts, fs, name) <- get
> let defname = toIvorName $ MN ("PE"++show n) name
> let sargs = mapMaybe getRight args
> let idx = (f, sargs)
> if (idx `elem` used) then return () else
> do
> let dargs = getDargs args nameSupply
> let rhs = reImplicit ist raw $
> apply (Name Unknown f) (mkAppArgs args sargs dargs)
> let dargs' = map (dused (namesIn rhs)) dargs
> let used' = idx:used
> let transFrom = mktrans (getMVs args dargs') rhs
> let transTo = mktrans (getMVs args dargs') (apply (Name Unknown defname) dargs')
> let trans = (transFrom, transTo)
> let freeze = getFrozen transFrom
> let newdef = (defname, ty, Patterns [PClause dargs' [] rhs], freeze, trans)
> let nds' = newdef:nds
> put (nds', used', trans:ts, freeze ++ fs, name+1)
> nameSupply = map (toIvorName.(MN "parg")) [0..]
> dused xs (Name _ n) | not (n `elem` xs) = Placeholder
> dused xs x = x
> -- getDargs ((Left n@(Name _ _)):xs) ns = n : getDargs xs ns
> getDargs ((Left _):xs) (n:ns) = (Name Unknown n) : getDargs xs ns
> getDargs (_:xs) ns = getDargs xs ns
> getDargs [] _ = []
> getRight (Right x) = Just x
> getRight _ = Nothing
Make a list of arguments for the specialisable application
> mkAppArgs [] _ _ = []
> mkAppArgs (Left _:xs) ss (d:ds) = d:(mkAppArgs xs ss ds)
> mkAppArgs (Right v:xs) (s:ss) ds = s:(mkAppArgs xs ss ds)
> mkAppArgs _ _ _ = []
Make the LHS and RHS of a transformation rule for the new definition
> getMVs (Left _:xs) (Name _ n:ds) = n:(getMVs xs ds)
> getMVs (Left _:xs) (_:ds) = getMVs xs ds
> getMVs (Right v:xs) ds = getMVs xs ds
> getMVs _ _ = []
> mktrans (n:ns) tm = mktrans ns (subst n (Metavar n) tm)
> mktrans [] tm = tm
> getFrozen tm = map fst (filter (\ (n, ty) -> ty == Free) (namesTypesIn tm))
Re-add _s for implicit arguments in PE definitions (because the type checker
will make a better job of working out what they should be than we will...)
The terms will just be simple applications
> reImplicit :: IdrisState -> Ctxt IvorFun -> ViewTerm -> ViewTerm
> reImplicit ist raw tm = reImp tm []
> where reImp fn@(Name _ n) stk
> = case getName n of
> Right ifn -> let imps = implicitArgs ifn
> stk' = take imps (repeat Placeholder)
> ++ drop imps stk in
> apply fn stk'
> _ -> apply fn stk
> reImp (App f a) stk = reImp f ((reImp a []):stk)
> reImp (Annotation a t) stk = Annotation a (reImp t stk)
> reImp x stk = apply x stk
> names = mkNameMap raw
> getName n = case lookup n names of
> Just x -> ctxtLookup raw [] x
> Nothing -> fail "No name"
Get a list of functions with static arguments, and their arities.
> staticDecls :: Ctxt IvorFun -> Statics
> staticDecls ctx = mapMaybe getStatic (ctxtAlist ctx)
> where getStatic (n,i)
> = if null (staticArgs i) || fwdDef (rawDecl i)
> then Nothing
> else let statics = (map (+ (implicitArgs i)) (staticArgs i))
> ar = arity (ivorDef i)
> extras = getExtras statics ar (ivorFType i)
> in Just (toIvorName n,
> (nub (extras ++ statics), ar,
> getType (ivorFType i)))
> fwdDef (Fwd _ _ _) = True
> fwdDef _ = False
> getType (Just t) = t
> arity (Just (PattDef ps)) = parity ps
> arity _ = 0
> parity (Patterns []) = 0
> parity (Patterns ((PClause xs _ _):_)) = length xs
> parity (Patterns ((PWithClause _ xs _ _):_)) = length xs
Look for dependencies on the static arguments. Keep going until there are
no more. Add any dependencies on static arguments which are dependencies on
*no* dynamic arguments.
> getExtras _ _ Nothing = []
> getExtras ss ar (Just t)
> = let ds = [0..ar-1] \\ ss
> args = TT.getArgTypes t
> stypes = map (\x -> snd (args!!!(x,"extra 1"))) ss
> dtypes = map (\x -> snd (args!!!(x,"extra 2"))) ds
> stnames = concatMap namesIn stypes
> dnames = concatMap namesIn dtypes
> newss_in = nub (ss ++ (mapMaybe (\x -> lookupIdx 0 x args) stnames))
> newds = mapMaybe (\x -> lookupIdx 0 x args) dnames
> newss = sort $ newss_in \\ newds in
> if (ss==newss) then ss else getExtras newss ar (Just t)
> lookupIdx i x ((n,v):xs) | n==x = Just i
> | otherwise = lookupIdx (i+1) x xs
> lookupIdx i x [] = Nothing
> addPEdefs :: Ctxt IvorFun -> Context -> Statics -> UserOps -> NewDefs ->
> TTM (Context, UserOps, [Id])
> addPEdefs raw ctxt sts uo nds = tryAdd ctxt uo [] nds
> where tryAdd ctxt uo frz [] = return (ctxt, uo, frz)
> tryAdd ctxt uo frz (d:ds)
> = case addPEdef ctxt uo d of
> Right (ctxt', uo', frz') -> -- trace ("WIN: " ++ show d) $
> tryAdd ctxt' uo' (frz++frz') ds
> Left err -> trace ("FAIL: " ++ show d ++ "\n" ++ show err) $
> tryAdd ctxt uo frz ds
> addPEdef ctxt (UO fix trans frz syn) (n, ty, pdef, frz', trans') =
> do (ctxt, []) <- addPatternDef ctxt n ty pdef
> [Specialise (map (\x -> (x,1)) frz'),
> SpecStatic (map getst sts)]
> return (ctxt, UO fix (trans':trans)
> (map getName frz'++frz) syn,
> map getName frz')
> getst (n, (args, arity, ty)) = (n,(args,arity))
> names = mkNameMap raw
> getName n = case lookup n names of
> Just x -> x
> Nothing -> UN (show n)
Go through the context, evaluating enough to work out what the transform
rules need to be. Return a new set of transform rules, and specialised versions
of PEable functions.
Algorithm is:
In a definition f, PE(f):
For each static function I, with static positions s1...sn:
* If it appears in a pattern clause,
make a new function Ispec di = mkSpec (I di si)
repeat PE(Ispec)
add a transform I ?di si => Ispec ?di
if any si is a name, freeze it
mkSpec(I di si):
* Evaluate I di si. Any si which is a name, expand only once.
After the top level I, do not evaluate any static arguments of any I -
either they are all expanded, or we'll PE them separately.
> partialeval :: Ctxt IvorFun -> Context -> Statics -> UserOps ->
> (Context, UserOps)
> partialeval raw ctxt sts uos = (ctxt, uos)