module Data.Rewriting.HigherOrder where
import Control.Monad.Reader
import Control.Monad.Writer
import qualified Data.Foldable as Foldable
import Data.Function (on)
import Data.List (groupBy)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Comp
import Data.Comp.Derive
import Data.Comp.Ops
import Data.Comp.Render
import Data.Rewriting.Rules
class Bind r
where
var :: Var r a -> r a
lam :: (Var r a -> r b) -> r (a -> b)
newtype VAR a = Var Name
deriving (Eq, Show, Ord, Num, Enum, Real, Integral, Functor, Foldable, Traversable)
data LAM a = Lam Name a
deriving (Eq, Show, Functor, Foldable, Traversable)
data APP a = App a a
deriving (Eq, Show, Functor, Foldable, Traversable)
derive [makeEqF, makeShowF, makeShowConstr] [''VAR]
derive [makeEqF, makeShowF, makeShowConstr] [''LAM]
derive [makeEqF, makeShowF, makeShowConstr] [''APP]
instance Render VAR
instance Render LAM
instance Render APP
fresh :: (LAM :<: f, Functor f, Foldable f) => Term f -> Name
fresh f
| Just (Lam v _) <- project f = v+1
| otherwise = maximum $ (0:) $ Foldable.toList $ fmap fresh $ unTerm f
mkLam
:: (Rep r, VAR :<: PF r, LAM :<: PF r, Functor (PF r), Foldable (PF r))
=> (VAR a -> Var r a) -> (Var r a -> r b) -> r (a -> b)
mkLam mkVar f = toRep $ inject $ Lam v $ fromRep body
where
body = f (mkVar $ Var v)
v = fresh (fromRep body)
app :: (APP :<: f) => Term (f :&: Set Name) -> Term (f :&: Set Name) -> Term (f :&: Set Name)
app f@(Term (_ :&: fv)) a@(Term (_ :&: av)) = Term (inj (App f a) :&: Set.union fv av)
type instance Var (LHS f) = VAR
type instance Var (RHS f) = VAR
instance (VAR :<: PF (LHS f), LAM :<: PF (LHS f), Functor f, Foldable f) =>
Bind (LHS f)
where
var = LHS . inject . Var . toInteger
lam = mkLam id
instance (VAR :<: PF (RHS f), LAM :<: PF (RHS f), Functor f, Foldable f) =>
Bind (RHS f)
where
var = RHS . inject . Var . toInteger
lam = mkLam id
type OneToOne a b = (Map a b, Map b a)
oEmpty :: OneToOne a b
oEmpty = (Map.empty, Map.empty)
oMember :: (Ord a, Ord b) => (a,b) -> OneToOne a b -> Bool
oMember (a,b) (ab,_) = case Map.lookup a ab of
Just b' -> b == b'
Nothing -> False
oMemberEither :: (Ord a, Ord b) => (a,b) -> OneToOne a b -> Bool
oMemberEither (a,b) (ab,ba) = Map.member a ab || Map.member b ba
oLookupL :: Ord a => a -> OneToOne a b -> Maybe b
oLookupL a (ab,_) = Map.lookup a ab
oInsert :: (Ord a, Ord b) => (a,b) -> OneToOne a b -> OneToOne a b
oInsert (a,b) (ab,ba) = (Map.insert a b ab', Map.insert b a ba')
where
ab' = case Map.lookup b ba of
Just a' -> Map.delete a' ab
Nothing -> ab
ba' = case Map.lookup a ab of
Just b' -> Map.delete b' ba
Nothing -> ba
getAnn :: Term (f :&: a) -> a
getAnn (Term (_ :&: a)) = a
type AlphaEnv = OneToOne Name Name
matchM :: forall f a
. ( VAR :<: f
, LAM :<: f
, VAR :<: PF (LHS f)
, LAM :<: PF (LHS f)
, Functor f, Foldable f, EqF f
)
=> LHS f a
-> Term (f :&: Set Name)
-> ReaderT AlphaEnv (WriterT (Subst (f :&: Set Name)) Maybe) ()
matchM (LHS lhs) t = go lhs t
where
go (Term (Inl WildCard)) _ = return ()
go (Term (Inr (Inl (Meta mv)))) t = ReaderT $ \env -> goo env mv t
where
goo :: AlphaEnv
-> MetaExp (LHS f) b
-> Term (f :&: Set Name)
-> WriterT (Subst (f :&: Set Name)) Maybe ()
goo env (MVar (MetaId m)) t
| Set.null (Set.intersection boundInPatt freeIn_t) = tell [(m,t)]
| otherwise = fail "Variables would escape"
where
boundInPatt = Map.keysSet $ snd env
freeIn_t = getAnn t
goo env (MApp mv (Var v)) t = do
let Just w = oLookupL v env
goo env mv (Term (inj (Lam w t) :&: Set.delete w (getAnn t)))
go p (Term (g :&: _))
| Just (Var v) <- project p
, Just (Var w) <- proj g
= do
env <- ask
guard ((v,w) `oMember` env)
go p (Term (g :&: _))
| Just (Lam v a) <- project p
, Just (Lam w b) <- proj g
= local (oInsert (v,w)) $ go a b
go (Term (Inr (Inr f))) (Term (g :&: _))
| Just subs <- eqMod f g
= mapM_ (uncurry go) subs
go _ _ = fail "No match"
alphaEq :: (VAR :<: f, LAM :<: f, Functor f, Foldable f, EqF f) => Term f -> Term f -> Bool
alphaEq a b = runReader (go a b) oEmpty
where
go t u
| Just (Var v) <- project t
, Just (Var w) <- project u
= reader $ \env -> oMember (v,w) env || (not (oMemberEither (v,w) env) && v==w)
go t u
| Just (Lam v a) <- project t
, Just (Lam w b) <- project u
= local (oInsert (v,w)) $ go a b
go (Term f) (Term g)
| Just subs <- eqMod f g
= fmap and $ mapM (uncurry go) subs
go _ _ = return False
solveTermAlpha :: (VAR :<: f, LAM :<: f, Functor f, Foldable f, EqF f) =>
[Term (f :&: a)] -> Maybe (Term (f :&: a))
solveTermAlpha (t:ts) = guard (all (alphaEq (stripA t)) (map stripA ts)) >> return t
solveTermAlpha _ = Nothing
solveSubstAlpha :: (VAR :<: f, LAM :<: f, Functor f, Foldable f, EqF f) =>
Subst (f :&: a) -> Maybe (Subst (f :&: a))
solveSubstAlpha s = sequence [fmap (v,) $ solveTermAlpha ts | g <- gs, let (v:_,ts) = unzip g]
where
gs = groupBy ((==) `on` fst) s
match
:: ( VAR :<: f
, LAM :<: f
, VAR :<: PF (LHS f)
, LAM :<: PF (LHS f)
, Functor f, Foldable f, EqF f
)
=> LHS f a -> Term (f :&: Set Name) -> Maybe (Subst (f :&: Set Name))
match lhs = solveSubstAlpha <=< execWriterT . flip runReaderT oEmpty . matchM lhs
annFreeVars :: (VAR :<: f, LAM :<: f, Functor f, Foldable f) =>
f (Term (f :&: Set Name)) -> Term (f :&: Set Name)
annFreeVars f
| Just (Var v) <- proj f = Term (inj (Var v) :&: Set.singleton v)
annFreeVars f
| Just (Lam v a) <- proj f
= Term (inj (Lam v a) :&: Set.delete v (getAnn a))
annFreeVars f = Term (f :&: Foldable.foldMap getAnn f)
type Aliases = (Map Name Name, Name)
initAliases :: Set Name -> Aliases
initAliases ns = (mp,next)
where
mp = Map.fromList [(n,n) | n <- Set.toList ns]
next = Set.findMax (Set.insert (1) ns) + 1
rename :: Name -> Aliases -> (Name,Aliases)
rename n (mp,next) = case Map.lookup n mp of
Nothing -> (n, (Map.insert n n mp, max next (n+1)))
Just _ -> (next, (Map.insert n next mp, next + 1))
renameNaive :: Name -> Aliases -> (Name,Aliases)
renameNaive n (mp,next) = (n, (Map.insert n n mp, max next (n+1)))
lookAlias :: Name -> Aliases -> Maybe Name
lookAlias n (mp,_) = Map.lookup n mp
substitute :: forall f g a
. ( VAR :<: f
, LAM :<: f
, VAR :<: PF (RHS f)
, LAM :<: PF (RHS f)
, Traversable f
, g ~ (f :&: Set Name)
)
=> (Term g -> Term g -> Term g)
-> Subst g
-> RHS f a
-> Maybe (Term g)
substitute app subst rhs = go (initAliases (Set.union fvS fvR)) (unRHS rhs)
where
fvS = Foldable.fold [fv | (_, Term (_ :&: fv)) <- subst]
Term (_ :&: fvR) = cata annFreeVars $ unRHS rhs
go :: Aliases -> Term (PF (RHS f)) -> Maybe (Term g)
go aliases (Term (Inl (Meta mv))) = goo mv
where
goo :: MetaExp (RHS f) b -> Maybe (Term g)
goo (MVar (MetaId v)) = lookup v subst
goo (MApp mv t) = liftM2 app (goo mv) $ go aliases (unRHS t)
go aliases t
| Just (Var v) <- project t
, Just v' <- lookAlias v aliases
= Just $ Term (inj (Var v') :&: Set.singleton v')
go aliases t
| Just (Lam v body) <- project t = do
let (v',aliases') = rename v aliases
body'@(Term (_ :&: fv)) <- go aliases' body
return $ Term (inj (Lam v' body') :&: Set.delete v' fv)
go aliases (Term (Inr f)) = fmap annFreeVars $ traverse (go aliases) f
rewrite
:: ( VAR :<: f
, LAM :<: f
, VAR :<: PF (LHS f)
, LAM :<: PF (LHS f)
, VAR :<: PF (RHS f)
, LAM :<: PF (RHS f)
, Traversable f, EqF f
, g ~ (f :&: Set Name)
)
=> (Term g -> Term g -> Term g)
-> Rule (LHS f) (RHS f)
-> Term g
-> Maybe (Term g)
rewrite app (Rule lhs rhs) t = do
subst <- match lhs t
substitute app subst rhs
applyFirst
:: ( VAR :<: f
, LAM :<: f
, VAR :<: PF (LHS f)
, LAM :<: PF (LHS f)
, VAR :<: PF (RHS f)
, LAM :<: PF (RHS f)
, Traversable f, EqF f
, g ~ (f :&: Set Name)
)
=> (Term g -> Term g -> Term g)
-> [Rule (LHS f) (RHS f)]
-> Term g
-> Term g
applyFirst app rs t = case [t' | r <- rs, Just t' <- [rewrite app r t]] of
t':_ -> t'
_ -> t
prepare :: (VAR :<: f, LAM :<: f, Functor f, Foldable f) => Term f -> Term (f :&: Set Name)
prepare = cata annFreeVars
stripAnn :: Functor f => Term (f :&: a) -> Term f
stripAnn = cata (\(f :&: _) -> Term f)
rewriteWith
:: ( VAR :<: f
, LAM :<: f
, Functor f
, Foldable f
, g ~ (f :&: Set Name)
)
=> (Term g -> Term g) -> Term f -> Term f
rewriteWith rew = stripAnn . rew . prepare