{-# LANGUAGE ExistentialQuantification, Rank2Types #-}


-- | Defines standard functions for modifying type expressions or retrieving
--   information from type expressions.

module Language.Haskell.FreeTheorems.Frontend.TypeExpressions (
    freeTypeVariables
  , allTypeVariables
  , createNewTypeVariableNotIn
  , alphaConversion
  , substituteTypeVariables
  , replaceAllTypeSynonyms
  , closeTypeExpressions
  , closureFor
) where



import Data.Generics
    ( Typeable, Data, synthesize, mkQ, everywhere, mkT, gmapT, GenericQ
    , GenericT )
import Data.List (find)
import Data.Set as Set
    ( Set, empty, union, insert, delete, fold, unions, difference, singleton
    , member )
import Data.Map as Map (Map, empty, lookup, delete, insert)
import Data.Maybe (maybe, fromJust)

import Language.Haskell.FreeTheorems.Syntax
import Language.Haskell.FreeTheorems.NameStores (typeNameStore)





------- Functions for type variables ------------------------------------------


-- | Returns all type variables occurring in a type expression.

allTypeVariables :: TypeExpression -> Set.Set TypeVariable
allTypeVariables :: TypeExpression -> Set TypeVariable
allTypeVariables = forall s t. s -> (t -> s -> s) -> GenericQ (s -> t) -> GenericQ t
synthesize forall a. Set a
Set.empty forall a. Ord a => Set a -> Set a -> Set a
Set.union (forall a. a -> a
id forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
`mkQ` TypeExpression -> Set TypeVariable -> Set TypeVariable
update)
  where
    update :: TypeExpression -> Set TypeVariable -> Set TypeVariable
update TypeExpression
t Set TypeVariable
s = case TypeExpression
t of
      TypeVar TypeVariable
v        -> forall a. Ord a => a -> Set a -> Set a
Set.insert TypeVariable
v Set TypeVariable
s
      TypeAbs TypeVariable
v [TypeClass]
_ TypeExpression
_    -> forall a. Ord a => a -> Set a -> Set a
Set.insert TypeVariable
v Set TypeVariable
s
      TypeAbsLab TypeVariable
v [TypeClass]
_ TypeExpression
_ -> forall a. Ord a => a -> Set a -> Set a
Set.insert TypeVariable
v Set TypeVariable
s
      TypeExpression
otherwise        -> Set TypeVariable
s



-- | Returns the free type variables of a type expression.

freeTypeVariables :: TypeExpression -> Set.Set TypeVariable
freeTypeVariables :: TypeExpression -> Set TypeVariable
freeTypeVariables = forall s t. s -> (t -> s -> s) -> GenericQ (s -> t) -> GenericQ t
synthesize forall a. Set a
Set.empty forall a. Ord a => Set a -> Set a -> Set a
Set.union (forall a. a -> a
id forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
`mkQ` TypeExpression -> Set TypeVariable -> Set TypeVariable
update)
  where
    update :: TypeExpression -> Set TypeVariable -> Set TypeVariable
update TypeExpression
t Set TypeVariable
s = case TypeExpression
t of
      TypeVar TypeVariable
v        -> forall a. Ord a => a -> Set a -> Set a
Set.insert TypeVariable
v Set TypeVariable
s
      TypeAbs TypeVariable
v [TypeClass]
_ TypeExpression
_    -> forall a. Ord a => a -> Set a -> Set a
Set.delete TypeVariable
v Set TypeVariable
s
      TypeAbsLab TypeVariable
v [TypeClass]
_ TypeExpression
_ -> forall a. Ord a => a -> Set a -> Set a
Set.delete TypeVariable
v Set TypeVariable
s
      TypeExpression
otherwise        -> Set TypeVariable
s



-- | Substitutes every free occurrence of all type variables given in the map
--   with the type expressions they are mapped to.

substituteTypeVariables ::
    Map.Map TypeVariable TypeExpression 
    -> TypeExpression 
    -> TypeExpression

substituteTypeVariables :: Map TypeVariable TypeExpression -> TypeExpression -> TypeExpression
substituteTypeVariables Map TypeVariable TypeExpression
env TypeExpression
t = 
  forall u. GenericQ (u -> u) -> GenericTw u -> u -> GenericT
everywhereWith (forall a. a -> a
id forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
`mkQ` forall {a}.
TypeExpression -> Map TypeVariable a -> Map TypeVariable a
update) (forall a b u.
(Typeable a, Typeable b) =>
(u -> a -> a) -> u -> b -> b
mkTw Map TypeVariable TypeExpression -> TypeExpression -> TypeExpression
substitute) Map TypeVariable TypeExpression
env TypeExpression
t
  where
    -- Updates the environment used in the top-down traversal.
    -- Removes bound type variables from the mapping. Thus, these variables
    -- won't be replaced in the second stage.
    update :: TypeExpression -> Map TypeVariable a -> Map TypeVariable a
update TypeExpression
t Map TypeVariable a
env = case TypeExpression
t of
      TypeAbs TypeVariable
v [TypeClass]
_ TypeExpression
_    -> forall k a. Ord k => k -> Map k a -> Map k a
Map.delete TypeVariable
v Map TypeVariable a
env
      TypeAbsLab TypeVariable
v [TypeClass]
_ TypeExpression
_ -> forall k a. Ord k => k -> Map k a -> Map k a
Map.delete TypeVariable
v Map TypeVariable a
env
      TypeExpression
otherwise        -> Map TypeVariable a
env
    
    -- Replaces a type variable by a type expression, if the type variable is
    -- contained in the environment.
    substitute :: Map TypeVariable TypeExpression -> TypeExpression -> TypeExpression
substitute Map TypeVariable TypeExpression
env TypeExpression
t = case TypeExpression
t of
      TypeVar TypeVariable
v -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeExpression
t forall a. a -> a
id (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TypeVariable
v Map TypeVariable TypeExpression
env)
      TypeExpression
otherwise -> TypeExpression
t



-- | Creates a new type variable not occurring in the given set of type
--   variables.
--
--   A type variable can either be named by the letters \'a\' to \'e\' or, if
--   that causes conflicts, by the letter \'a\' concatenated with a number
--   starting from 1.

createNewTypeVariableNotIn :: Set.Set TypeVariable -> TypeVariable
createNewTypeVariableNotIn :: Set TypeVariable -> TypeVariable
createNewTypeVariableNotIn Set TypeVariable
forbiddenVars =
  let vars :: [TypeVariable]
vars = forall a b. (a -> b) -> [a] -> [b]
map (Identifier -> TypeVariable
TV forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Identifier
Ident) [String]
typeNameStore
   in forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\TypeVariable
v -> Bool -> Bool
not (TypeVariable
v forall a. Ord a => a -> Set a -> Bool
`Set.member` Set TypeVariable
forbiddenVars)) [TypeVariable]
vars



-- | Replaces every bound occurrence of the given type variable in the given
--   type expression with a new type variable which is created by
--   'createNewTypeVariableNotIn'.
--
--   Several bound occurrences of the given type variable are replaced with the
--   same new type variable. Only one new type variable is created altogether.

alphaConversion :: TypeVariable -> TypeExpression -> TypeExpression
alphaConversion :: TypeVariable -> TypeExpression -> TypeExpression
alphaConversion TypeVariable
old TypeExpression
t =
  forall u. GenericQ (u -> u) -> GenericTw u -> u -> GenericT
everywhereWith (forall a. a -> a
id forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
`mkQ` TypeExpression
-> (TypeVariable -> TypeVariable) -> TypeVariable -> TypeVariable
change) (forall a b u.
(Typeable a, Typeable b) =>
(u -> a -> a) -> u -> b -> b
mkTw (TypeVariable -> TypeVariable) -> TypeExpression -> TypeExpression
replace) forall a. a -> a
id TypeExpression
t
  where
    -- The new type variable which will replace 'old' everywhere in the given
    -- type expression.
    new :: TypeVariable
new = Set TypeVariable -> TypeVariable
createNewTypeVariableNotIn (TypeExpression -> Set TypeVariable
allTypeVariables TypeExpression
t)
    
    -- This function replaces any type variable equal to 'old' with 'new' and
    -- keeps all other type variables unchanged.
    rep :: TypeVariable -> TypeVariable
rep TypeVariable
v = if (TypeVariable
v forall a. Eq a => a -> a -> Bool
== TypeVariable
old) then TypeVariable
new else TypeVariable
v

    -- Modifies the function which replaces type variables.
    -- If we are at the type abstraction where 'old' is bound, then 'old' has
    -- to be replaced in every subexpression by the new type variable.
    change :: TypeExpression
-> (TypeVariable -> TypeVariable) -> TypeVariable -> TypeVariable
change TypeExpression
t TypeVariable -> TypeVariable
f = case TypeExpression
t of
      TypeAbs    TypeVariable
v [TypeClass]
_ TypeExpression
_ -> if (TypeVariable
v forall a. Eq a => a -> a -> Bool
== TypeVariable
old) then TypeVariable -> TypeVariable
rep else TypeVariable -> TypeVariable
f
      TypeAbsLab TypeVariable
v [TypeClass]
_ TypeExpression
_ -> if (TypeVariable
v forall a. Eq a => a -> a -> Bool
== TypeVariable
old) then TypeVariable -> TypeVariable
rep else TypeVariable -> TypeVariable
f
      TypeExpression
otherwise         -> TypeVariable -> TypeVariable
f

    -- Applies the current replacement function to type variables.
    -- In type abstractions, the static function 'rep' is used to replace
    -- 'old' with 'new' or otherwise keep the type variable.
    -- Note that - independent of the usage of 'rep' - the replacement function
    -- 'r' will be modified by 'change' when advancing to subexpressions.
    replace :: (TypeVariable -> TypeVariable) -> TypeExpression -> TypeExpression
replace TypeVariable -> TypeVariable
r TypeExpression
t = case TypeExpression
t of
      TypeVar    TypeVariable
v       -> TypeVariable -> TypeExpression
TypeVar    (TypeVariable -> TypeVariable
r TypeVariable
v)
      TypeAbs    TypeVariable
v [TypeClass]
cs TypeExpression
t' -> TypeVariable -> [TypeClass] -> TypeExpression -> TypeExpression
TypeAbs    (TypeVariable -> TypeVariable
rep TypeVariable
v) [TypeClass]
cs TypeExpression
t'
      TypeAbsLab TypeVariable
v [TypeClass]
cs TypeExpression
t' -> TypeVariable -> [TypeClass] -> TypeExpression -> TypeExpression
TypeAbsLab (TypeVariable -> TypeVariable
rep TypeVariable
v) [TypeClass]
cs TypeExpression
t'
      TypeExpression
otherwise       -> TypeExpression
t





------- Generic helper definitions --------------------------------------------


-- | Generic transformations using a value of fixed type.

type GenericTw u = forall a . Data a => u -> a -> a



-- | Make a generic transformation which uses a value of fixed type.
--   This function takes a specific case into a general case, such that no
--   transformation is applied for types not covered by the specific case.

mkTw :: (Typeable a, Typeable b) => (u -> a -> a) -> u -> b -> b
mkTw :: forall a b u.
(Typeable a, Typeable b) =>
(u -> a -> a) -> u -> b -> b
mkTw u -> a -> a
f u
u = forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
mkT forall a b. (a -> b) -> a -> b
$ u -> a -> a
f u
u



-- | Pushes a value in a top-down fashion trough a tree and applies that value
--   from bottom to top to every node.
--
--   More detailed, the expression
--
-- >   everywhereWith update apply v
--
--   is evaluated as follows:
--   The value @v@ is transfered through the tree from top to bottom while,
--   at every node, the function @update@ is applied to it. This allows the
--   initial value to be changed like, for example, an environment gathering
--   information from the root to the leaf while moving through the tree.
--   Thereafter, the transformation function @apply@ is applied to every node
--   from bottom to top. It might use the value distributed to that node.

everywhereWith :: GenericQ (u -> u) -> GenericTw u -> u -> GenericT
everywhereWith :: forall u. GenericQ (u -> u) -> GenericTw u -> u -> GenericT
everywhereWith GenericQ (u -> u)
k GenericTw u
f u
u a
x = (GenericTw u
f u
u) forall a b. (a -> b) -> a -> b
$ forall a. Data a => GenericT -> a -> a
gmapT (forall u. GenericQ (u -> u) -> GenericTw u -> u -> GenericT
everywhereWith GenericQ (u -> u)
k GenericTw u
f (GenericQ (u -> u)
k a
x u
u)) a
x





------- Replacing type synonyms -----------------------------------------------


-- | Replaces all type synonyms in an arbitrary tree.
--   The first argument gives the list of known type synonyms and their 
--   declarations. Every occurrence of one of those type synonyms in the second
--   argument is replaced by the according right-hand side of the declaration.
--   
--   Note that the type synonym declarations given in the first argument may
--   themselves contain type synonyms. However, type synonym declarations must
--   not be recursive nor mutually recursive.

replaceAllTypeSynonyms :: (Typeable a, Data a) => [TypeDeclaration] -> a -> a
replaceAllTypeSynonyms :: forall a. (Typeable a, Data a) => [TypeDeclaration] -> a -> a
replaceAllTypeSynonyms [TypeDeclaration]
knownTypes = GenericT -> GenericT
everywhere (forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
mkT TypeExpression -> TypeExpression
replace)
    -- This functions replaces all type synonyms in a bottom-up manner.
    -- Thus, when applying 'replace', all type synonyms are already removed
    -- from all children of the node.

  where
    -- Replacing type synonyms only affects type constructors.
    -- Check if there is a type synonym declaration for the given type 
    -- constructor. If not, just return the unchanged type expression.
    -- Otherwise replace the type synonym by its definition.
    replace :: TypeExpression -> TypeExpression
replace TypeExpression
t = case TypeExpression
t of
      TypeCon TypeConstructor
c [TypeExpression]
ts -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe TypeExpression
t ([TypeExpression] -> TypeDeclaration -> TypeExpression
applyTypeSynonym [TypeExpression]
ts) ([TypeDeclaration] -> TypeConstructor -> Maybe TypeDeclaration
findTypeDecl [TypeDeclaration]
knownTypes TypeConstructor
c)
      TypeExpression
otherwise    -> TypeExpression
t

    -- Applies the declaration of a type synonym to a list of type expressions.
    -- The type expression composed in this way is returned.
    -- Note that the structure of 'replaceTypeSynonyms' guarantees that there is
    -- no type synonym in any of the type expressions of 'ts'.
    applyTypeSynonym :: [TypeExpression] -> TypeDeclaration -> TypeExpression
applyTypeSynonym [TypeExpression]
ts (Type Identifier
_ [TypeVariable]
vs TypeExpression
t) =
      let 
          -- First, remove all type synonyms from the declaration's right-hand
          -- side. Note that this terminates because type expressions cannot be
          -- declared recursively nor mutually recursively.
          t1 :: TypeExpression
t1 = forall a. (Typeable a, Data a) => [TypeDeclaration] -> a -> a
replaceAllTypeSynonyms [TypeDeclaration]
knownTypes TypeExpression
t

          -- Construct an environment to be used to substitute every free
          -- variable in 't' by the appropiate type expression of 'ts'.
          env :: Map TypeVariable TypeExpression
env = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert) forall k a. Map k a
Map.empty (forall a b. [a] -> [b] -> [(a, b)]
zip [TypeVariable]
vs [TypeExpression]
ts)

          -- Rename all bound variables in 't' such that no free variables of
          -- any type expression in 'ts' will get bound.
          allFreeVariables :: Set TypeVariable
allFreeVariables = forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TypeExpression -> Set TypeVariable
freeTypeVariables [TypeExpression]
ts
          t2 :: TypeExpression
t2 = forall a b. (a -> b -> b) -> b -> Set a -> b
Set.fold TypeVariable -> TypeExpression -> TypeExpression
alphaConversion TypeExpression
t1 Set TypeVariable
allFreeVariables

          -- Finally, apply the declaration's right-hand side to 'ts' and return
          -- the constructed type expression.
       in Map TypeVariable TypeExpression -> TypeExpression -> TypeExpression
substituteTypeVariables Map TypeVariable TypeExpression
env TypeExpression
t2
  


-- | Looks up the declaration for a type synonym constructor.
--   If the given type constructor is not a type synonym or there is no
--   declaration for this type constructor in the declarations list, then
--   @Nothing@ is returned.

findTypeDecl :: [TypeDeclaration] -> TypeConstructor -> Maybe TypeDeclaration
findTypeDecl :: [TypeDeclaration] -> TypeConstructor -> Maybe TypeDeclaration
findTypeDecl [TypeDeclaration]
decls TypeConstructor
con = case TypeConstructor
con of
  Con Identifier
name  -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\TypeDeclaration
d -> TypeDeclaration -> Identifier
typeName TypeDeclaration
d forall a. Eq a => a -> a -> Bool
== Identifier
name) [TypeDeclaration]
decls 
  TypeConstructor
otherwise -> forall a. Maybe a
Nothing





------- Closing type expressions ----------------------------------------------


-- | Closes all type expressions in type signature declarations.
--   Class methods are also modified, in that every free variable expect the
--   class variable is explicitly bound.

closeTypeExpressions :: [Declaration] -> [Declaration]
closeTypeExpressions :: [Declaration] -> [Declaration]
closeTypeExpressions = forall a b. (a -> b) -> [a] -> [b]
map Declaration -> Declaration
closeDecl
  where
    closeDecl :: Declaration -> Declaration
closeDecl Declaration
d = case Declaration
d of
      ClassDecl ClassDeclaration
d -> ClassDeclaration -> Declaration
ClassDecl (ClassDeclaration -> ClassDeclaration
closeClassDecl ClassDeclaration
d)
      TypeSig Signature
sig -> Signature -> Declaration
TypeSig (Signature -> Signature
closeSignature Signature
sig)
      Declaration
otherwise   -> Declaration
d



-- | Closes type signatures of class methods. Afterwards, the class variable is
--   still free in all class methods.

closeClassDecl :: ClassDeclaration -> ClassDeclaration
closeClassDecl :: ClassDeclaration -> ClassDeclaration
closeClassDecl ClassDeclaration
d =
  ClassDeclaration
d { classFuns :: [Signature]
classFuns = forall a b. (a -> b) -> [a] -> [b]
map (TypeVariable -> Signature -> Signature
closureWithout (ClassDeclaration -> TypeVariable
classVar ClassDeclaration
d)) (ClassDeclaration -> [Signature]
classFuns ClassDeclaration
d) }
  where
    -- Close the class method signature while keeping the class variable free.
    closureWithout :: TypeVariable -> Signature -> Signature
closureWithout TypeVariable
v Signature
s = 
      let t :: TypeExpression
t = Signature -> TypeExpression
signatureType Signature
s
          freeVars :: Set TypeVariable
freeVars = TypeExpression -> Set TypeVariable
freeTypeVariables TypeExpression
t forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` forall a. a -> Set a
Set.singleton TypeVariable
v
       in Signature
s { signatureType :: TypeExpression
signatureType = Set TypeVariable -> TypeExpression -> TypeExpression
closureFor Set TypeVariable
freeVars TypeExpression
t }



-- | Close a type signature declaration.

closeSignature :: Signature -> Signature
closeSignature :: Signature -> Signature
closeSignature Signature
s =
  let t :: TypeExpression
t = Signature -> TypeExpression
signatureType Signature
s
   in Signature
s { signatureType :: TypeExpression
signatureType = Set TypeVariable -> TypeExpression -> TypeExpression
closureFor (TypeExpression -> Set TypeVariable
freeTypeVariables TypeExpression
t) TypeExpression
t }



-- | Explicitly binds all type variables of the first argument by a type 
--   abstraction in the given type expression.

closureFor :: Set.Set TypeVariable -> TypeExpression -> TypeExpression
closureFor :: Set TypeVariable -> TypeExpression -> TypeExpression
closureFor Set TypeVariable
vs TypeExpression
t = forall a b. (a -> b -> b) -> b -> Set a -> b
Set.fold (\TypeVariable
v -> TypeVariable -> [TypeClass] -> TypeExpression -> TypeExpression
TypeAbs TypeVariable
v []) TypeExpression
t Set TypeVariable
vs