{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeOperators         #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Variables
-- Copyright   :  (c) 2010-2011 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk> and Tom Hvitved <hvitved@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- This module defines an abstract notion of (bound) variables in compositional
-- data types, and scoped substitution. Capture-avoidance is /not/ taken into
-- account.
--
--------------------------------------------------------------------------------

module Data.Comp.Variables
    (
     HasVars(..),
     Subst,
     CxtSubst,
     varsToHoles,
     containsVar,
     variables,
     variableList,
     variables',
     substVars,
     appSubst,
     compSubst,
     getBoundVars,
    (&),
    (|->),
    empty
    ) where

import Data.Comp.Algebra
import Data.Comp.Derive
import Data.Comp.Mapping
import Data.Comp.Term
import Data.Comp.Ops
import Data.Foldable hiding (elem, notElem)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as Set
import Prelude hiding (foldl, or)

-- | This type represents substitutions of contexts, i.e. finite
-- mappings from variables to contexts.
type CxtSubst h a f v = Map v (Cxt h f a)

-- | This type represents substitutions of terms, i.e. finite mappings
-- from variables to terms.
type Subst f v = CxtSubst NoHole () f v

{-| This multiparameter class defines functors with variables. An instance
  @HasVar f v@ denotes that values over @f@ might contain and bind variables of
  type @v@. -}
class HasVars f v where
    -- | Indicates whether the @f@ constructor is a variable. The
    -- default implementation returns @Nothing@.
    isVar :: f a -> Maybe v
    isVar f a
_ = Maybe v
forall a. Maybe a
Nothing

    -- | Indicates the set of variables bound by the @f@ constructor
    -- for each argument of the constructor. For example for a
    -- non-recursive let binding:
    -- 
    -- @
    -- data Let e = Let Var e e
    -- instance HasVars Let Var where
    --   bindsVars (Let v x y) = y |-> Set.singleton v
    -- @
    -- 
    -- If, instead, the let binding is recursive, the methods has to
    -- be implemented like this:
    -- 
    -- @
    --   bindsVars (Let v x y) = x |-> Set.singleton v &
    --                           y |-> Set.singleton v
    -- @
    -- 
    -- This indicates that the scope of the bound variable also
    -- extends to the right-hand side of the variable binding.
    --
    -- The default implementation returns the empty map.
    bindsVars :: Mapping m a => f a -> m (Set v)
    bindsVars f a
_ = m (Set v)
forall (m :: * -> *) k v. Mapping m k => m v
empty


$(derive [liftSum] [''HasVars])

instance HasVars f v => HasVars (f :&: a) v where
  isVar :: (:&:) f a a -> Maybe v
isVar (f a
f :&: a
_)     = f a -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar f a
f
  bindsVars :: (:&:) f a a -> m (Set v)
bindsVars (f a
f :&: a
_) = f a -> m (Set v)
forall (f :: * -> *) v (m :: * -> *) a.
(HasVars f v, Mapping m a) =>
f a -> m (Set v)
bindsVars f a
f

-- | Same as 'isVar' but it returns Nothing@ instead of @Just v@ if
-- @v@ is contained in the given set of variables.

isVar' :: (HasVars f v, Ord v) => Set v -> f a -> Maybe v
isVar' :: Set v -> f a -> Maybe v
isVar' Set v
b f a
t = do v
v <- f a -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar f a
t
                if v
v v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set v
b
                   then Maybe v
forall a. Maybe a
Nothing
                   else v -> Maybe v
forall (m :: * -> *) a. Monad m => a -> m a
return v
v

-- | This combinator pairs every argument of a given constructor with
-- the set of (newly) bound variables according to the corresponding
-- 'HasVars' type class instance.
getBoundVars :: (HasVars f v, Traversable f) => f a -> f (Set v, a)
getBoundVars :: f a -> f (Set v, a)
getBoundVars f a
t = let n :: f (Numbered a)
n = f a -> f (Numbered a)
forall (f :: * -> *) a. Traversable f => f a -> f (Numbered a)
number f a
t
                     m :: NumMap a (Set v)
m = f (Numbered a) -> NumMap a (Set v)
forall (f :: * -> *) v (m :: * -> *) a.
(HasVars f v, Mapping m a) =>
f a -> m (Set v)
bindsVars f (Numbered a)
n
                     trans :: Numbered a -> (Set v, a)
trans (Numbered Int
i a
x) = (Set v -> Int -> NumMap a (Set v) -> Set v
forall a t. a -> Int -> NumMap t a -> a
lookupNumMap Set v
forall a. Set a
Set.empty Int
i NumMap a (Set v)
m, a
x)
                 in (Numbered a -> (Set v, a)) -> f (Numbered a) -> f (Set v, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Numbered a -> (Set v, a)
trans f (Numbered a)
n

-- | This combinator combines 'getBoundVars' with the generic 'fmap' function.
fmapBoundVars :: (HasVars f v, Traversable f) => (Set v -> a -> b) -> f a -> f b
fmapBoundVars :: (Set v -> a -> b) -> f a -> f b
fmapBoundVars Set v -> a -> b
f f a
t = let n :: f (Numbered a)
n = f a -> f (Numbered a)
forall (f :: * -> *) a. Traversable f => f a -> f (Numbered a)
number f a
t
                        m :: NumMap a (Set v)
m = f (Numbered a) -> NumMap a (Set v)
forall (f :: * -> *) v (m :: * -> *) a.
(HasVars f v, Mapping m a) =>
f a -> m (Set v)
bindsVars f (Numbered a)
n
                        trans :: Numbered a -> b
trans (Numbered Int
i a
x) = Set v -> a -> b
f (Set v -> Int -> NumMap a (Set v) -> Set v
forall a t. a -> Int -> NumMap t a -> a
lookupNumMap Set v
forall a. Set a
Set.empty Int
i NumMap a (Set v)
m) a
x
                    in (Numbered a -> b) -> f (Numbered a) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Numbered a -> b
trans f (Numbered a)
n

-- | This combinator combines 'getBoundVars' with the generic 'foldl' function.
foldlBoundVars :: (HasVars f v, Traversable f) => (b -> Set v -> a -> b) -> b -> f a -> b
foldlBoundVars :: (b -> Set v -> a -> b) -> b -> f a -> b
foldlBoundVars b -> Set v -> a -> b
f b
e f a
t = let n :: f (Numbered a)
n = f a -> f (Numbered a)
forall (f :: * -> *) a. Traversable f => f a -> f (Numbered a)
number f a
t
                           m :: NumMap a (Set v)
m = f (Numbered a) -> NumMap a (Set v)
forall (f :: * -> *) v (m :: * -> *) a.
(HasVars f v, Mapping m a) =>
f a -> m (Set v)
bindsVars f (Numbered a)
n
                           trans :: b -> Numbered a -> b
trans b
x (Numbered Int
i a
y) = b -> Set v -> a -> b
f b
x (Set v -> Int -> NumMap a (Set v) -> Set v
forall a t. a -> Int -> NumMap t a -> a
lookupNumMap Set v
forall a. Set a
Set.empty Int
i NumMap a (Set v)
m) a
y
                       in (b -> Numbered a -> b) -> b -> f (Numbered a) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> Numbered a -> b
trans b
e f (Numbered a)
n

-- | Convert variables to holes, except those that are bound.
varsToHoles :: (Traversable f, HasVars f v, Ord v) => Term f -> Context f v
varsToHoles :: Term f -> Context f v
varsToHoles Term f
t = Alg f (Set v -> Context f v) -> Term f -> Set v -> Context f v
forall (f :: * -> *) a. Functor f => Alg f a -> Term f -> a
cata Alg f (Set v -> Context f v)
forall (f :: * -> *) v.
(Traversable f, HasVars f v, Ord v) =>
Alg f (Set v -> Context f v)
alg Term f
t Set v
forall a. Set a
Set.empty
    where alg :: (Traversable f, HasVars f v, Ord v) => Alg f (Set v -> Context f v)
          alg :: Alg f (Set v -> Context f v)
alg f (Set v -> Context f v)
t Set v
vars = case f (Set v -> Context f v) -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar f (Set v -> Context f v)
t of
            Just v
v | Bool -> Bool
not (v
v v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set v
vars) -> v -> Context f v
forall a (f :: * -> *). a -> Cxt Hole f a
Hole v
v
            Maybe v
_  -> f (Context f v) -> Context f v
forall (f :: * -> *) h a. f (Cxt h f a) -> Cxt h f a
Term (f (Context f v) -> Context f v) -> f (Context f v) -> Context f v
forall a b. (a -> b) -> a -> b
$ (Set v -> (Set v -> Context f v) -> Context f v)
-> f (Set v -> Context f v) -> f (Context f v)
forall (f :: * -> *) v a b.
(HasVars f v, Traversable f) =>
(Set v -> a -> b) -> f a -> f b
fmapBoundVars Set v -> (Set v -> Context f v) -> Context f v
run f (Set v -> Context f v)
t
              where
                run :: Set v -> (Set v -> Context f v) -> Context f v
run Set v
newVars Set v -> Context f v
f = Set v -> Context f v
f (Set v -> Context f v) -> Set v -> Context f v
forall a b. (a -> b) -> a -> b
$ Set v
newVars Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
vars

-- |Algebra for checking whether a variable is contained in a term, except those
-- that are bound.
containsVarAlg :: (Eq v, HasVars f v, Traversable f, Ord v) => v -> Alg f Bool
containsVarAlg :: v -> Alg f Bool
containsVarAlg v
v f Bool
t = (Bool -> Set v -> Bool -> Bool) -> Bool -> Alg f Bool
forall (f :: * -> *) v b a.
(HasVars f v, Traversable f) =>
(b -> Set v -> a -> b) -> b -> f a -> b
foldlBoundVars Bool -> Set v -> Bool -> Bool
run Bool
local f Bool
t
    where local :: Bool
local = case f Bool -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar f Bool
t of
                    Just v
v' -> v
v v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v'
                    Maybe v
Nothing -> Bool
False
          run :: Bool -> Set v -> Bool -> Bool
run Bool
acc Set v
vars Bool
b = Bool
acc Bool -> Bool -> Bool
|| (Bool -> Bool
not (v
v v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set v
vars) Bool -> Bool -> Bool
&& Bool
b)

{-| This function checks whether a variable is contained in a context. -}
containsVar :: (Eq v, HasVars f v, Traversable f, Ord v)
            => v -> Cxt h f a -> Bool
containsVar :: v -> Cxt h f a -> Bool
containsVar v
v = Alg f Bool -> (a -> Bool) -> Cxt h f a -> Bool
forall (f :: * -> *) h a b.
Functor f =>
Alg f b -> (a -> b) -> Cxt h f a -> b
free (v -> Alg f Bool
forall v (f :: * -> *).
(Eq v, HasVars f v, Traversable f, Ord v) =>
v -> Alg f Bool
containsVarAlg v
v) (Bool -> a -> Bool
forall a b. a -> b -> a
const Bool
False)

-- |Algebra for generating a set of variables contained in a term, except those
-- that are bound.
variablesAlg :: (Ord v, HasVars f v, Traversable f) => Alg f (Set v)
variablesAlg :: Alg f (Set v)
variablesAlg f (Set v)
t = (Set v -> Set v -> Set v -> Set v) -> Set v -> Alg f (Set v)
forall (f :: * -> *) v b a.
(HasVars f v, Traversable f) =>
(b -> Set v -> a -> b) -> b -> f a -> b
foldlBoundVars Set v -> Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a -> Set a
run Set v
local f (Set v)
t
    where local :: Set v
local = case f (Set v) -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar f (Set v)
t of
                    Just v
v -> v -> Set v
forall a. a -> Set a
Set.singleton v
v
                    Maybe v
Nothing -> Set v
forall a. Set a
Set.empty
          run :: Set a -> Set a -> Set a -> Set a
run Set a
acc Set a
bvars Set a
vars = Set a
acc Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` (Set a
vars Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set a
bvars)


{-| This function computes the list of variables occurring in a context. -}
variableList :: (Ord v, HasVars f v, Traversable f) => Cxt h f a -> [v]
variableList :: Cxt h f a -> [v]
variableList = Set v -> [v]
forall a. Set a -> [a]
Set.toList (Set v -> [v]) -> (Cxt h f a -> Set v) -> Cxt h f a -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cxt h f a -> Set v
forall v (f :: * -> *) h a.
(Ord v, HasVars f v, Traversable f) =>
Cxt h f a -> Set v
variables

{-| This function computes the set of variables occurring in a context. -}
variables :: (Ord v, HasVars f v, Traversable f) => Cxt h f a -> Set v
variables :: Cxt h f a -> Set v
variables = Alg f (Set v) -> (a -> Set v) -> Cxt h f a -> Set v
forall (f :: * -> *) h a b.
Functor f =>
Alg f b -> (a -> b) -> Cxt h f a -> b
free Alg f (Set v)
forall v (f :: * -> *).
(Ord v, HasVars f v, Traversable f) =>
Alg f (Set v)
variablesAlg (Set v -> a -> Set v
forall a b. a -> b -> a
const Set v
forall a. Set a
Set.empty)

{-| This function computes the set of variables occurring in a constant. -}
variables' :: (Ord v, HasVars f v, Foldable f, Functor f) => Const f -> Set v
variables' :: Const f -> Set v
variables' Const f
c = case Const f -> Maybe v
forall (f :: * -> *) v a. HasVars f v => f a -> Maybe v
isVar Const f
c of
                 Maybe v
Nothing -> Set v
forall a. Set a
Set.empty
                 Just v
v -> v -> Set v
forall a. a -> Set a
Set.singleton v
v

{-| This multiparameter class defines substitution of values of type @t@ for
  variables of type @v@ in values of type @a@. -}
class SubstVars v t a where
    substVars :: (v -> Maybe t) -> a -> a

-- |Apply the given substitution.
appSubst :: (Ord v, SubstVars v t a) => Map v t -> a -> a
appSubst :: Map v t -> a -> a
appSubst Map v t
subst = (v -> Maybe t) -> a -> a
forall v t a. SubstVars v t a => (v -> Maybe t) -> a -> a
substVars v -> Maybe t
f
    where f :: v -> Maybe t
f v
v = v -> Map v t -> Maybe t
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup v
v Map v t
subst

instance  {-# OVERLAPPABLE #-} (Ord v, HasVars f v, Traversable f)
    => SubstVars v (Cxt h f a) (Cxt h f a) where
        -- have to use explicit GADT pattern matching!!
        -- subst f = free (substAlg f) Hole
  substVars :: (v -> Maybe (Cxt h f a)) -> Cxt h f a -> Cxt h f a
substVars v -> Maybe (Cxt h f a)
subst = Set v -> Cxt h f a -> Cxt h f a
doSubst Set v
forall a. Set a
Set.empty
    where doSubst :: Set v -> Cxt h f a -> Cxt h f a
doSubst Set v
_ (Hole a
a) = a -> Cxt Hole f a
forall a (f :: * -> *). a -> Cxt Hole f a
Hole a
a
          doSubst Set v
b (Term f (Cxt h f a)
t) = case Set v -> f (Cxt h f a) -> Maybe v
forall (f :: * -> *) v a.
(HasVars f v, Ord v) =>
Set v -> f a -> Maybe v
isVar' Set v
b f (Cxt h f a)
t Maybe v -> (v -> Maybe (Cxt h f a)) -> Maybe (Cxt h f a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v -> Maybe (Cxt h f a)
subst of
            Just Cxt h f a
new -> Cxt h f a
new
            Maybe (Cxt h f a)
Nothing  -> f (Cxt h f a) -> Cxt h f a
forall (f :: * -> *) h a. f (Cxt h f a) -> Cxt h f a
Term (f (Cxt h f a) -> Cxt h f a) -> f (Cxt h f a) -> Cxt h f a
forall a b. (a -> b) -> a -> b
$ (Set v -> Cxt h f a -> Cxt h f a) -> f (Cxt h f a) -> f (Cxt h f a)
forall (f :: * -> *) v a b.
(HasVars f v, Traversable f) =>
(Set v -> a -> b) -> f a -> f b
fmapBoundVars Set v -> Cxt h f a -> Cxt h f a
run f (Cxt h f a)
t
              where run :: Set v -> Cxt h f a -> Cxt h f a
run Set v
vars = Set v -> Cxt h f a -> Cxt h f a
doSubst (Set v
b Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
vars)

instance  {-# OVERLAPPABLE #-} (SubstVars v t a, Functor f) => SubstVars v t (f a) where
    substVars :: (v -> Maybe t) -> f a -> f a
substVars v -> Maybe t
f = (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((v -> Maybe t) -> a -> a
forall v t a. SubstVars v t a => (v -> Maybe t) -> a -> a
substVars v -> Maybe t
f)

{-| This function composes two substitutions @s1@ and @s2@. That is,
applying the resulting substitution is equivalent to first applying
@s2@ and then @s1@. -}
compSubst :: (Ord v, HasVars f v, Traversable f)
          => CxtSubst h a f v -> CxtSubst h a f v -> CxtSubst h a f v
compSubst :: CxtSubst h a f v -> CxtSubst h a f v -> CxtSubst h a f v
compSubst CxtSubst h a f v
s1 CxtSubst h a f v
s2 = (Cxt h f a -> Cxt h f a) -> CxtSubst h a f v -> CxtSubst h a f v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CxtSubst h a f v -> Cxt h f a -> Cxt h f a
forall v t a. (Ord v, SubstVars v t a) => Map v t -> a -> a
appSubst CxtSubst h a f v
s1) CxtSubst h a f v
s2 CxtSubst h a f v -> CxtSubst h a f v -> CxtSubst h a f v
forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union` CxtSubst h a f v
s1