idris-0.99: src/Idris/PartialEval.hs
{-|
Module : Idris.PartialEval
Description : Implementation of a partial evaluator.
Copyright :
License : BSD3
Maintainer : The Idris Community.
-}
{-# LANGUAGE PatternGuards #-}
module Idris.PartialEval(
partial_eval, getSpecApps, specType
, mkPE_TyDecl, mkPE_TermDecl, PEArgType(..)
, pe_app, pe_def, pe_clauses, pe_simple
) where
import Idris.AbsSyntax
import Idris.Core.CaseTree
import Idris.Core.Evaluate
import Idris.Core.TT
import Idris.Delaborate
import Control.Applicative
import Control.Monad.State
import Data.Maybe
import Debug.Trace
-- | Data type representing binding-time annotations for partial evaluation of arguments
data PEArgType = ImplicitS Name -- ^ Implicit static argument
| ImplicitD Name -- ^ Implicit dynamic argument
| ConstraintS -- ^ Implementation constraint
| ConstraintD -- ^ Implementation constraint
| ExplicitS -- ^ Explicit static argument
| ExplicitD -- ^ Explicit dynamic argument
| UnifiedD -- ^ Erasable dynamic argument (found under unification)
deriving (Eq, Show)
-- | A partially evaluated function. pe_app captures the lhs of the
-- new definition, pe_def captures the rhs, and pe_clauses is the
-- specialised implementation.
--
-- pe_simple is set if the result is always reducible, because in such
-- a case we'll also need to reduce the static argument
data PEDecl = PEDecl { pe_app :: PTerm, -- new application
pe_def :: PTerm, -- old application
pe_clauses :: [(PTerm, PTerm)], -- clauses of new application
pe_simple :: Bool -- if just one reducible clause
}
-- | Partially evaluates given terms under the given context.
-- It is an error if partial evaluation fails to make any progress.
-- Making progress is defined as: all of the names given with explicit
-- reduction limits (in practice, the function being specialised)
-- must have reduced at least once.
-- If we don't do this, we might end up making an infinite function after
-- applying the transformation.
partial_eval :: Context
-> [(Name, Maybe Int)]
-> [Either Term (Term, Term)]
-> Maybe [Either Term (Term, Term)]
partial_eval ctxt ns_in tms = mapM peClause tms where
ns = squash ns_in
squash ((n, Just x) : ns)
| Just (Just y) <- lookup n ns
= squash ((n, Just (x + y)) : drop n ns)
| otherwise = (n, Just x) : squash ns
squash (n : ns) = n : squash ns
squash [] = []
drop n ((m, _) : ns) | n == m = ns
drop n (x : ns) = x : drop n ns
drop n [] = []
-- If the term is not a clause, it is simply kept as is
peClause (Left t) = Just $ Left t
-- If the term is a clause, specialise the right hand side
peClause (Right (lhs, rhs))
= let (rhs', reductions) = specialise ctxt [] (map toLimit ns) rhs in
do when (length tms == 1) $ checkProgress ns reductions
return (Right (lhs, rhs'))
-- TMP HACK until I do PE by WHNF rather than using main evaluator
toLimit (n, Nothing) | isTCDict n ctxt = (n, 2)
toLimit (n, Nothing) = (n, 65536) -- somewhat arbitrary reduction limit
toLimit (n, Just l) = (n, l)
checkProgress ns [] = return ()
checkProgress ns ((n, r) : rs)
| Just (Just start) <- lookup n ns
= if start <= 1 || r < start then checkProgress ns rs else Nothing
| otherwise = checkProgress ns rs
-- | Specialises the type of a partially evaluated TT function returning
-- a pair of the specialised type and the types of expected arguments.
specType :: [(PEArgType, Term)] -> Type -> (Type, [(PEArgType, Term)])
specType args ty = let (t, args') = runState (unifyEq args ty) [] in
(st (map fst args') t, map fst args')
where
-- Specialise static argument in type by let-binding provided value instead
-- of expecting it as a function argument
st ((ExplicitS, v) : xs) (Bind n (Pi _ t _) sc)
= Bind n (Let t v) (st xs sc)
st ((ImplicitS _, v) : xs) (Bind n (Pi _ t _) sc)
= Bind n (Let t v) (st xs sc)
st ((ConstraintS, v) : xs) (Bind n (Pi _ t _) sc)
= Bind n (Let t v) (st xs sc)
-- Erase argument from function type
st ((UnifiedD, _) : xs) (Bind n (Pi _ t _) sc)
= st xs sc
-- Keep types as is
st (_ : xs) (Bind n (Pi i t k) sc)
= Bind n (Pi i t k) (st xs sc)
st _ t = t
-- Erase implicit dynamic argument if existing argument shares it value,
-- by substituting the value of previous argument
unifyEq (imp@(ImplicitD _, v) : xs) (Bind n (Pi i t k) sc)
= do amap <- get
case lookup imp amap of
Just n' ->
do put (amap ++ [((UnifiedD, Erased), n)])
sc' <- unifyEq xs (subst n (P Bound n' Erased) sc)
return (Bind n (Pi i t k) sc') -- erase later
_ -> do put (amap ++ [(imp, n)])
sc' <- unifyEq xs sc
return (Bind n (Pi i t k) sc')
unifyEq (x : xs) (Bind n (Pi i t k) sc)
= do args <- get
put (args ++ [(x, n)])
sc' <- unifyEq xs sc
return (Bind n (Pi i t k) sc')
unifyEq xs t = do args <- get
put (args ++ (zip xs (repeat (sUN "_"))))
return t
-- | Creates an Idris type declaration given current state and a
-- specialised TT function application type.
-- Can be used in combination with the output of 'specType'.
--
-- This should: specialise any static argument position, then generalise
-- over any function applications in the result.
mkPE_TyDecl :: IState -> [(PEArgType, Term)] -> Type -> PTerm
mkPE_TyDecl ist args ty = mkty args ty
where
mkty ((ExplicitD, v) : xs) (Bind n (Pi _ t k) sc)
= PPi expl n NoFC (delab ist (generaliseIn t)) (mkty xs sc)
mkty ((ConstraintD, v) : xs) (Bind n (Pi _ t k) sc)
| concreteInterface ist t = mkty xs sc
| interfaceConstraint ist t
= PPi constraint n NoFC (delab ist (generaliseIn t)) (mkty xs sc)
mkty ((ImplicitD _, v) : xs) (Bind n (Pi _ t k) sc)
= PPi impl n NoFC (delab ist (generaliseIn t)) (mkty xs sc)
mkty (_ : xs) t
= mkty xs t
mkty [] t = delab ist t
generaliseIn tm = evalState (gen tm) 0
gen tm | (P _ fn _, args) <- unApply tm,
isFnName fn (tt_ctxt ist)
= do nm <- get
put (nm + 1)
return (P Bound (sMN nm "spec") Erased)
gen (App s f a) = App s <$> gen f <*> gen a
gen tm = return tm
-- | Checks if a given argument is an interface constraint argument
interfaceConstraint :: Idris.AbsSyntax.IState -> TT Name -> Bool
interfaceConstraint ist v
| (P _ c _, args) <- unApply v = case lookupCtxt c (idris_interfaces ist) of
[_] -> True
_ -> False
| otherwise = False
-- | Checks if the given arguments of an interface constraint are all either constants
-- or references (i.e. that it doesn't contain any complex terms).
concreteInterface :: IState -> TT Name -> Bool
concreteInterface ist v
| not (interfaceConstraint ist v) = False
| (P _ c _, args) <- unApply v = all concrete args
| otherwise = False
where concrete (Constant _) = True
concrete tm | (P _ n _, args) <- unApply tm
= case lookupTy n (tt_ctxt ist) of
[_] -> all concrete args
_ -> False
| otherwise = False
mkNewPats :: IState
-> [(Term, Term)] -- ^ definition to specialise
-> [(PEArgType, Term)] -- ^ arguments to specialise with
-> Name -- ^ New name
-> Name -- ^ Specialised function name
-> PTerm -- ^ Default lhs
-> PTerm -- ^ Default rhs
-> PEDecl
-- If all of the dynamic positions on the lhs are variables (rather than
-- patterns or constants) then we can just make a simple definition
-- directly applying the specialised function, since we know the
-- definition isn't going to block on any of the dynamic arguments
-- in this case
mkNewPats ist d ns newname sname lhs rhs | all dynVar (map fst d)
= PEDecl lhs rhs [(lhs, rhs)] True
where dynVar ap = case unApply ap of
(_, args) -> dynArgs ns args
dynArgs _ [] = True -- can definitely reduce from here
-- if Static, doesn't matter what the argument is
dynArgs ((ImplicitS _, _) : ns) (a : as) = dynArgs ns as
dynArgs ((ConstraintS, _) : ns) (a : as) = dynArgs ns as
dynArgs ((ExplicitS, _) : ns) (a : as) = dynArgs ns as
-- if Dynamic, it had better be a variable or we'll need to
-- do some more work
dynArgs (_ : ns) (V _ : as) = dynArgs ns as
dynArgs (_ : ns) (P _ _ _ : as) = dynArgs ns as
dynArgs _ _ = False -- and now we'll get stuck
mkNewPats ist d ns newname sname lhs rhs =
PEDecl lhs rhs (map mkClause d) False
where
mkClause :: (Term, Term) -> (PTerm, PTerm)
mkClause (oldlhs, oldrhs)
= let (_, as) = unApply oldlhs
lhsargs = mkLHSargs [] ns as
lhs = PApp emptyFC (PRef emptyFC [] newname) lhsargs
rhs = PApp emptyFC (PRef emptyFC [] sname)
(mkRHSargs ns lhsargs) in
(lhs, rhs)
mkLHSargs _ [] _ = []
-- dynamics don't appear on the LHS if they're implicit
mkLHSargs sub ((ExplicitD, t) : ns) (a : as)
= pexp (delab ist (substNames sub a)) : mkLHSargs sub ns as
mkLHSargs sub ((ImplicitD n, t) : ns) (a : as)
= pimp n (delab ist (substNames sub a)) True : mkLHSargs sub ns as
mkLHSargs sub ((ConstraintD, t) : ns) (a : as)
= pconst (delab ist (substNames sub a)) : mkLHSargs sub ns as
mkLHSargs sub ((UnifiedD, _) : ns) (a : as)
= mkLHSargs sub ns as
-- statics get dropped in any case
mkLHSargs sub ((ImplicitS _, t) : ns) (a : as)
= mkLHSargs (extend a t sub) ns as
mkLHSargs sub ((ExplicitS, t) : ns) (a : as)
= mkLHSargs (extend a t sub) ns as
mkLHSargs sub ((ConstraintS, t) : ns) (a : as)
= mkLHSargs (extend a t sub) ns as
mkLHSargs sub _ [] = [] -- no more LHS
extend (P _ n _) t sub = (n, t) : sub
extend _ _ sub = sub
--- 'as' are the LHS arguments
mkRHSargs ((ExplicitS, t) : ns) as = pexp (delab ist t) : mkRHSargs ns as
mkRHSargs ((ExplicitD, t) : ns) (a : as) = a : mkRHSargs ns as
-- Keep the implicits on the RHS, in case they got matched on
mkRHSargs ((ImplicitD n, t) : ns) (a : as) = a : mkRHSargs ns as
mkRHSargs ((ImplicitS n, t) : ns) as -- Dropped from LHS
= pimp n (delab ist t) True : mkRHSargs ns as
mkRHSargs ((ConstraintD, t) : ns) (a : as) = a : mkRHSargs ns as
mkRHSargs ((ConstraintS, t) : ns) as -- Dropped from LHS
= pconst (delab ist t) : mkRHSargs ns as
mkRHSargs (_ : ns) as = mkRHSargs ns as
mkRHSargs _ _ = []
mkSubst :: (Term, Term) -> Maybe (Name, Term)
mkSubst (P _ n _, t) = Just (n, t)
mkSubst _ = Nothing
-- | Creates a new declaration for a specialised function application.
-- Simple version at the moment: just create a version which is a direct
-- application of the function to be specialised.
-- More complex version to do: specialise the definition clause by clause
mkPE_TermDecl :: IState
-> Name
-> Name
-> PTerm -- ^ Type of specialised function
-> [(PEArgType, Term)]
-> PEDecl
mkPE_TermDecl ist newname sname specty ns
{- We need to erase the *dynamic* arguments
where their *name* appears in the *type* of a later argument
in specty.
i.e. if a later dynamic argument depends on an earlier dynamic
argument, we should infer the earlier one.
Then we need to erase names from the LHS which no longer appear
on the RHS.
-}
= let deps = getDepNames (eraseRet specty)
lhs = eraseDeps deps $
PApp emptyFC (PRef emptyFC [] newname) (mkp ns)
rhs = eraseDeps deps $
delab ist (mkApp (P Ref sname Erased) (map snd ns))
patdef = -- trace (showTmImpls specty ++ "\n" ++ showTmImpls lhs ++ "\n"
-- ++ showTmImpls rhs) $
lookupCtxtExact sname (idris_patdefs ist)
newpats = case patdef of
Nothing -> PEDecl lhs rhs [(lhs, rhs)] True
Just d -> mkNewPats ist (getPats d) ns
newname sname lhs rhs in
newpats where
getPats (ps, _) = map (\(_, lhs, rhs) -> (lhs, rhs)) ps
eraseRet (PPi p n fc ty sc) = PPi p n fc ty (eraseRet sc)
eraseRet _ = Placeholder
-- Get names used in later arguments; assume we've called eraseRet so there's
-- no names going to appear in return type
getDepNames (PPi _ n _ _ sc)
| n `elem` allNamesIn sc = n : getDepNames sc
| otherwise = getDepNames sc
getDepNames tm = []
mkp [] = []
mkp ((ExplicitD, tm) : tms) = pexp (delab ist tm) : mkp tms
mkp ((ImplicitD n, tm) : tms) = pimp n (delab ist tm) True : mkp tms
mkp (_ : tms) = mkp tms
eraseDeps ns tm = mapPT (deImp ns) tm
deImp ns (PApp fc t as) = PApp fc t (map (deImpArg ns) as)
deImp ns t = t
deImpArg ns a | pname a `elem` ns = a { getTm = Placeholder }
| otherwise = a
-- | Get specialised applications for a given function
getSpecApps :: IState
-> [Name]
-> Term
-> [(Name, [(PEArgType, Term)])]
getSpecApps ist env tm = ga env (explicitNames tm) where
-- staticArg env True _ tm@(P _ n _) _ | n `elem` env = Just (True, tm)
-- staticArg env True _ tm@(App f a) _ | (P _ n _, args) <- unApply tm,
-- n `elem` env = Just (True, tm)
staticArg env True imp tm n
| Just n <- imparg imp = (ImplicitS n, tm)
| constrarg imp = (ConstraintS, tm)
| otherwise = (ExplicitS, tm)
staticArg env False imp tm n
| Just nm <- imparg imp = (ImplicitD nm, (P Ref (sUN (show n ++ "arg")) Erased))
| constrarg imp = (ConstraintD, tm)
| otherwise = (ExplicitD, (P Ref (sUN (show n ++ "arg")) Erased))
imparg (PExp _ _ _ _) = Nothing
imparg (PConstraint _ _ _ _) = Nothing
imparg arg = Just (pname arg)
constrarg (PConstraint _ _ _ _) = True
constrarg arg = False
buildApp env [] [] _ _ = []
buildApp env (s:ss) (i:is) (a:as) (n:ns)
= let s' = staticArg env s i a n
ss' = buildApp env ss is as ns in
(s' : ss')
-- if we have a *defined* function that has static arguments,
-- it will become a specialised application
ga env tm@(App _ f a) | (P _ n _, args) <- unApply tm,
n `notElem` map fst (idris_metavars ist) =
ga env f ++ ga env a ++
case (lookupCtxtExact n (idris_statics ist),
lookupCtxtExact n (idris_implicits ist)) of
(Just statics, Just imps) ->
if (length statics == length args && or statics
&& specialisable (tt_ctxt ist) n) then
case buildApp env statics imps args [0..] of
args -> [(n, args)]
-- _ -> []
else []
_ -> []
ga env (Bind n (Let t v) sc) = ga env v ++ ga (n : env) sc
ga env (Bind n t sc) = ga (n : env) sc
ga env t = []
-- A function is only specialisable if there are no overlapping
-- cases in the case tree (otherwise the partial evaluation could
-- easily get stuck)
specialisable :: Context -> Name -> Bool
specialisable ctxt n = case lookupDefExact n ctxt of
Just (CaseOp _ _ _ _ _ cds) ->
noOverlap (snd (cases_compiletime cds))
_ -> False
noOverlap :: SC -> Bool
noOverlap (Case _ _ [DefaultCase sc]) = noOverlap sc
noOverlap (Case _ _ alts) = noOverlapAlts alts
noOverlap _ = True
-- There's an overlap if the case tree has a default case along with
-- some other cases. It's fine if there's a default case on its own.
noOverlapAlts (ConCase _ _ _ sc : rest)
= noOverlapAlts rest && noOverlap sc
noOverlapAlts (FnCase _ _ sc : rest) = noOverlapAlts rest
noOverlapAlts (ConstCase _ sc : rest)
= noOverlapAlts rest && noOverlap sc
noOverlapAlts (SucCase _ sc : rest)
= noOverlapAlts rest && noOverlap sc
noOverlapAlts (DefaultCase _ : _) = False
noOverlapAlts _ = True