{- |
    Module      :  $Header$
    Description :  Type substitution
    Copyright   :  (c) 2003 Wolfgang Lux
                       2016 Finn Teegen
    License     :  BSD-3-clause

    Maintainer  :  bjp@informatik.uni-kiel.de
    Stability   :  experimental
    Portability :  portable

   This module implements substitutions on types.
-}

module Base.TypeSubst
  ( module Base.TypeSubst, idSubst, singleSubst, bindSubst, compose
  ) where

import           Data.List       (nub)
import           Data.Maybe      (fromMaybe)
import qualified Data.Set as Set (Set, map)

import Base.Subst
import Base.TopEnv
import Base.Types

import Env.Value (ValueInfo (..))

type TypeSubst = Subst Int Type

class SubstType a where
  subst :: TypeSubst -> a -> a

bindVar :: Int -> Type -> TypeSubst -> TypeSubst
bindVar :: Int -> Type -> TypeSubst -> TypeSubst
bindVar tv :: Int
tv ty :: Type
ty = TypeSubst -> TypeSubst -> TypeSubst
forall v e. Ord v => Subst v e -> Subst v e -> Subst v e
compose (Int -> Type -> TypeSubst -> TypeSubst
forall v e. Ord v => v -> e -> Subst v e -> Subst v e
bindSubst Int
tv Type
ty TypeSubst
forall a b. Subst a b
idSubst)

substVar :: TypeSubst -> Int -> Type
substVar :: TypeSubst -> Int -> Type
substVar = (Int -> Type)
-> (TypeSubst -> Type -> Type) -> TypeSubst -> Int -> Type
forall v e.
Ord v =>
(v -> e) -> (Subst v e -> e -> e) -> Subst v e -> v -> e
substVar' Int -> Type
TypeVariable TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst

instance (Ord a, SubstType a) => SubstType (Set.Set a) where
  subst :: TypeSubst -> Set a -> Set a
subst sigma :: TypeSubst
sigma = (a -> a) -> Set a -> Set a
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (TypeSubst -> a -> a
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma)

instance SubstType a => SubstType [a] where
  subst :: TypeSubst -> [a] -> [a]
subst sigma :: TypeSubst
sigma = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (TypeSubst -> a -> a
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma)

instance SubstType Type where
  subst :: TypeSubst -> Type -> Type
subst sigma :: TypeSubst
sigma ty :: Type
ty = TypeSubst -> Type -> [Type] -> Type
subst' TypeSubst
sigma Type
ty []

subst' :: TypeSubst -> Type -> [Type] -> Type
subst' :: TypeSubst -> Type -> [Type] -> Type
subst' _     tc :: Type
tc@(TypeConstructor   _) = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
TypeApply Type
tc
subst' sigma :: TypeSubst
sigma (TypeApply      ty1 :: Type
ty1 ty2 :: Type
ty2) = TypeSubst -> Type -> [Type] -> Type
subst' TypeSubst
sigma Type
ty1 ([Type] -> Type) -> ([Type] -> [Type]) -> [Type] -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty2 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:)
subst' sigma :: TypeSubst
sigma (TypeVariable        tv :: Int
tv) = Type -> [Type] -> Type
applyType (TypeSubst -> Int -> Type
substVar TypeSubst
sigma Int
tv)
subst' sigma :: TypeSubst
sigma (TypeArrow      ty1 :: Type
ty1 ty2 :: Type
ty2) =
  (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
TypeApply (Type -> Type -> Type
TypeArrow (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty1) (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty2))
subst' sigma :: TypeSubst
sigma (TypeConstrained tys :: [Type]
tys tv :: Int
tv) = case TypeSubst -> Int -> Type
substVar TypeSubst
sigma Int
tv of
  TypeVariable tv' :: Int
tv' -> (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
TypeApply ([Type] -> Int -> Type
TypeConstrained [Type]
tys Int
tv')
  ty :: Type
ty               -> (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
TypeApply Type
ty
subst' sigma :: TypeSubst
sigma (TypeForall      tvs :: [Int]
tvs ty :: Type
ty) =
  Type -> [Type] -> Type
applyType ([Int] -> Type -> Type
TypeForall [Int]
tvs (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty))

instance SubstType Pred where
  subst :: TypeSubst -> Pred -> Pred
subst sigma :: TypeSubst
sigma (Pred qcls :: QualIdent
qcls ty :: Type
ty) = QualIdent -> Type -> Pred
Pred QualIdent
qcls (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty)

instance SubstType PredType where
  subst :: TypeSubst -> PredType -> PredType
subst sigma :: TypeSubst
sigma (PredType ps :: PredSet
ps ty :: Type
ty) = PredSet -> Type -> PredType
PredType (TypeSubst -> PredSet -> PredSet
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma PredSet
ps) (TypeSubst -> Type -> Type
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
sigma Type
ty)

instance SubstType TypeScheme where
  subst :: TypeSubst -> TypeScheme -> TypeScheme
subst sigma :: TypeSubst
sigma (ForAll n :: Int
n ty :: PredType
ty) =
    Int -> PredType -> TypeScheme
ForAll Int
n (TypeSubst -> PredType -> PredType
forall a. SubstType a => TypeSubst -> a -> a
subst ((Int -> TypeSubst -> TypeSubst) -> TypeSubst -> [Int] -> TypeSubst
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int -> TypeSubst -> TypeSubst
forall v e. Ord v => v -> Subst v e -> Subst v e
unbindSubst TypeSubst
sigma [0 .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1]) PredType
ty)

instance SubstType ValueInfo where
  subst :: TypeSubst -> ValueInfo -> ValueInfo
subst _     dc :: ValueInfo
dc@(DataConstructor  _ _ _ _) = ValueInfo
dc
  subst _     nc :: ValueInfo
nc@(NewtypeConstructor _ _ _) = ValueInfo
nc
  subst theta :: TypeSubst
theta (Value             v :: QualIdent
v cm :: Maybe QualIdent
cm a :: Int
a ty :: TypeScheme
ty) = QualIdent -> Maybe QualIdent -> Int -> TypeScheme -> ValueInfo
Value QualIdent
v Maybe QualIdent
cm Int
a (TypeSubst -> TypeScheme -> TypeScheme
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
theta TypeScheme
ty)
  subst theta :: TypeSubst
theta (Label                l :: QualIdent
l r :: [QualIdent]
r ty :: TypeScheme
ty) = QualIdent -> [QualIdent] -> TypeScheme -> ValueInfo
Label QualIdent
l [QualIdent]
r (TypeSubst -> TypeScheme -> TypeScheme
forall a. SubstType a => TypeSubst -> a -> a
subst TypeSubst
theta TypeScheme
ty)

instance SubstType a => SubstType (TopEnv a) where
  subst :: TypeSubst -> TopEnv a -> TopEnv a
subst = (a -> a) -> TopEnv a -> TopEnv a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a) -> TopEnv a -> TopEnv a)
-> (TypeSubst -> a -> a) -> TypeSubst -> TopEnv a -> TopEnv a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubst -> a -> a
forall a. SubstType a => TypeSubst -> a -> a
subst

-- The class method 'expandAliasType' expands all occurrences of a
-- type synonym in its second argument.

class ExpandAliasType a where
  expandAliasType :: [Type] -> a -> a

instance ExpandAliasType a => ExpandAliasType [a] where
  expandAliasType :: [Type] -> [a] -> [a]
expandAliasType tys :: [Type]
tys = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ([Type] -> a -> a
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys)

instance (Ord a, ExpandAliasType a) => ExpandAliasType (Set.Set a) where
  expandAliasType :: [Type] -> Set a -> Set a
expandAliasType tys :: [Type]
tys = (a -> a) -> Set a -> Set a
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map ([Type] -> a -> a
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys)

instance ExpandAliasType Type where
  expandAliasType :: [Type] -> Type -> Type
expandAliasType tys :: [Type]
tys ty :: Type
ty = [Type] -> Type -> [Type] -> Type
expandAliasType' [Type]
tys Type
ty []

expandAliasType' :: [Type] -> Type -> [Type] -> Type
expandAliasType' :: [Type] -> Type -> [Type] -> Type
expandAliasType' _   tc :: Type
tc@(TypeConstructor   _) = Type -> [Type] -> Type
applyType Type
tc
expandAliasType' tys :: [Type]
tys (TypeApply      ty1 :: Type
ty1 ty2 :: Type
ty2) =
  [Type] -> Type -> [Type] -> Type
expandAliasType' [Type]
tys Type
ty1 ([Type] -> Type) -> ([Type] -> [Type]) -> [Type] -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty2 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:)
expandAliasType' tys :: [Type]
tys tv :: Type
tv@(TypeVariable      n :: Int
n)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 0    = Type -> [Type] -> Type
applyType ([Type]
tys [Type] -> Int -> Type
forall a. [a] -> Int -> a
!! Int
n)
  | Bool
otherwise = Type -> [Type] -> Type
applyType Type
tv
expandAliasType' _   tc :: Type
tc@(TypeConstrained _ _) = Type -> [Type] -> Type
applyType Type
tc
expandAliasType' tys :: [Type]
tys (TypeArrow      ty1 :: Type
ty1 ty2 :: Type
ty2) =
  Type -> [Type] -> Type
applyType (Type -> Type -> Type
TypeArrow ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty1) ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty2))
expandAliasType' tys :: [Type]
tys (TypeForall      tvs :: [Int]
tvs ty :: Type
ty) =
  Type -> [Type] -> Type
applyType ([Int] -> Type -> Type
TypeForall [Int]
tvs ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty))

instance ExpandAliasType Pred where
  expandAliasType :: [Type] -> Pred -> Pred
expandAliasType tys :: [Type]
tys (Pred qcls :: QualIdent
qcls ty :: Type
ty) = QualIdent -> Type -> Pred
Pred QualIdent
qcls ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty)

instance ExpandAliasType PredType where
  expandAliasType :: [Type] -> PredType -> PredType
expandAliasType tys :: [Type]
tys (PredType ps :: PredSet
ps ty :: Type
ty) =
    PredSet -> Type -> PredType
PredType ([Type] -> PredSet -> PredSet
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys PredSet
ps) ([Type] -> Type -> Type
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Type]
tys Type
ty)

-- After the expansion we have to reassign the type indices for all type
-- variables. Otherwise, expanding a type synonym like type 'Pair a b = (b,a)'
-- could break the invariant that the universally quantified type variables
-- are assigned indices in the order of their occurrence. This is handled by
-- the function 'normalize'. The function has a threshold parameter that allows
-- preserving the indices of type variables bound on the left hand side
-- of a type declaration and in the head of a type class declaration,
-- respectively.

normalize :: Int -> PredType -> PredType
normalize :: Int -> PredType -> PredType
normalize n :: Int
n ty :: PredType
ty = [Type] -> PredType -> PredType
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType [Int -> Type
TypeVariable (Int -> Int
occur Int
tv) | Int
tv <- [0..]] PredType
ty
  where tvs :: [(Int, Int)]
tvs = [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n) (PredType -> [Int]
forall t. IsType t => t -> [Int]
typeVars PredType
ty))) [Int
n..]
        occur :: Int -> Int
occur tv :: Int
tv = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
tv (Int -> [(Int, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
tv [(Int, Int)]
tvs)

-- The function 'instanceType' computes an instance of a polymorphic type by
-- substituting the first type argument for all occurrences of the type
-- variable with index 0 in the second argument. The function carefully
-- assigns new indices to all other type variables of the second argument
-- so that they do not conflict with the type variables of the first argument.

instanceType :: ExpandAliasType a => Type -> a -> a
instanceType :: Type -> a -> a
instanceType ty :: Type
ty = [Type] -> a -> a
forall a. ExpandAliasType a => [Type] -> a -> a
expandAliasType (Type
ty Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: (Int -> Type) -> [Int] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Type
TypeVariable [Int
n ..])
  where ForAll n :: Int
n _ = Type -> TypeScheme
polyType Type
ty