{- |
Module      : Language.Egison.Type.Subst
Licence     : MIT

This module provides type substitution operations for the type system.
-}

{-# LANGUAGE DeriveGeneric #-}

module Language.Egison.Type.Subst
  ( Subst(..)
  , emptySubst
  , singletonSubst
  , composeSubst
  , applySubst
  , applySubstScheme
  , applySubstConstraint
  , SubstIndex
  , emptySubstIndex
  , singletonSubstIndex
  , applySubstIndex
  ) where

import           Data.Map.Strict            (Map)
import qualified Data.Map.Strict            as Map
import           GHC.Generics               (Generic)

import           Language.Egison.Type.Index (Index (..), IndexSpec, IndexTyVar (..))
import           Language.Egison.Type.Types (TyVar (..), Type (..), TypeScheme (..), Constraint(..))

-- | Type substitution: a mapping from type variables to types
newtype Subst = Subst { Subst -> Map TyVar Type
unSubst :: Map TyVar Type }
  deriving (Subst -> Subst -> Bool
(Subst -> Subst -> Bool) -> (Subst -> Subst -> Bool) -> Eq Subst
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Subst -> Subst -> Bool
== :: Subst -> Subst -> Bool
$c/= :: Subst -> Subst -> Bool
/= :: Subst -> Subst -> Bool
Eq, Int -> Subst -> ShowS
[Subst] -> ShowS
Subst -> String
(Int -> Subst -> ShowS)
-> (Subst -> String) -> ([Subst] -> ShowS) -> Show Subst
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Subst -> ShowS
showsPrec :: Int -> Subst -> ShowS
$cshow :: Subst -> String
show :: Subst -> String
$cshowList :: [Subst] -> ShowS
showList :: [Subst] -> ShowS
Show, (forall x. Subst -> Rep Subst x)
-> (forall x. Rep Subst x -> Subst) -> Generic Subst
forall x. Rep Subst x -> Subst
forall x. Subst -> Rep Subst x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Subst -> Rep Subst x
from :: forall x. Subst -> Rep Subst x
$cto :: forall x. Rep Subst x -> Subst
to :: forall x. Rep Subst x -> Subst
Generic)

-- | Empty substitution
emptySubst :: Subst
emptySubst :: Subst
emptySubst = Map TyVar Type -> Subst
Subst Map TyVar Type
forall k a. Map k a
Map.empty

-- | Create a substitution with a single binding
singletonSubst :: TyVar -> Type -> Subst
singletonSubst :: TyVar -> Type -> Subst
singletonSubst TyVar
v Type
t = Map TyVar Type -> Subst
Subst (Map TyVar Type -> Subst) -> Map TyVar Type -> Subst
forall a b. (a -> b) -> a -> b
$ TyVar -> Type -> Map TyVar Type
forall k a. k -> a -> Map k a
Map.singleton TyVar
v Type
t

-- | Compose two substitutions (s2 after s1)
-- (s2 `composeSubst` s1) x = s2 (s1 x)
composeSubst :: Subst -> Subst -> Subst
composeSubst :: Subst -> Subst -> Subst
composeSubst s2 :: Subst
s2@(Subst Map TyVar Type
m2) (Subst Map TyVar Type
m1) =
  Map TyVar Type -> Subst
Subst (Map TyVar Type -> Subst) -> Map TyVar Type -> Subst
forall a b. (a -> b) -> a -> b
$ (Type -> Type) -> Map TyVar Type -> Map TyVar Type
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map (Subst -> Type -> Type
applySubst Subst
s2) Map TyVar Type
m1 Map TyVar Type -> Map TyVar Type -> Map TyVar Type
forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union` Map TyVar Type
m2

-- | Apply a substitution to a type
applySubst :: Subst -> Type -> Type
applySubst :: Subst -> Type -> Type
applySubst Subst
_ Type
TInt             = Type
TInt
applySubst Subst
_ Type
TMathExpr        = Type
TMathExpr
applySubst Subst
_ Type
TPolyExpr        = Type
TPolyExpr
applySubst Subst
_ Type
TTermExpr        = Type
TTermExpr
applySubst Subst
_ Type
TSymbolExpr      = Type
TSymbolExpr
applySubst Subst
_ Type
TIndexExpr       = Type
TIndexExpr
applySubst Subst
_ Type
TFloat           = Type
TFloat
applySubst Subst
_ Type
TBool            = Type
TBool
applySubst Subst
_ Type
TChar            = Type
TChar
applySubst Subst
_ Type
TString          = Type
TString
applySubst (Subst Map TyVar Type
m) t :: Type
t@(TVar TyVar
v) = Type -> TyVar -> Map TyVar Type -> Type
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Type
t TyVar
v Map TyVar Type
m
applySubst Subst
s (TTuple [Type]
ts)      = [Type] -> Type
TTuple ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Type -> Type
applySubst Subst
s) [Type]
ts)
applySubst Subst
s (TCollection Type
t)  = Type -> Type
TCollection (Subst -> Type -> Type
applySubst Subst
s Type
t)
applySubst Subst
s (TInductive String
name [Type]
ts) = String -> [Type] -> Type
TInductive String
name ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Type -> Type
applySubst Subst
s) [Type]
ts)
applySubst Subst
s (TTensor Type
t)      = Type -> Type
TTensor (Subst -> Type -> Type
applySubst Subst
s Type
t)
applySubst Subst
s (THash Type
k Type
v)      = Type -> Type -> Type
THash (Subst -> Type -> Type
applySubst Subst
s Type
k) (Subst -> Type -> Type
applySubst Subst
s Type
v)
applySubst Subst
s (TMatcher Type
t)     = Type -> Type
TMatcher (Subst -> Type -> Type
applySubst Subst
s Type
t)
applySubst Subst
s (TFun Type
t1 Type
t2)     = Type -> Type -> Type
TFun (Subst -> Type -> Type
applySubst Subst
s Type
t1) (Subst -> Type -> Type
applySubst Subst
s Type
t2)
applySubst Subst
s (TIO Type
t)          = Type -> Type
TIO (Subst -> Type -> Type
applySubst Subst
s Type
t)
applySubst Subst
s (TIORef Type
t)       = Type -> Type
TIORef (Subst -> Type -> Type
applySubst Subst
s Type
t)
applySubst Subst
_ Type
TPort            = Type
TPort
applySubst Subst
_ Type
TAny             = Type
TAny

-- | Apply a substitution to a type scheme
applySubstScheme :: Subst -> TypeScheme -> TypeScheme
applySubstScheme :: Subst -> TypeScheme -> TypeScheme
applySubstScheme (Subst Map TyVar Type
m) (Forall [TyVar]
vs [Constraint]
cs Type
t) =
  let m' :: Map TyVar Type
m' = (TyVar -> Map TyVar Type -> Map TyVar Type)
-> Map TyVar Type -> [TyVar] -> Map TyVar Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TyVar -> Map TyVar Type -> Map TyVar Type
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Map TyVar Type
m [TyVar]
vs
      s' :: Subst
s' = Map TyVar Type -> Subst
Subst Map TyVar Type
m'
  in [TyVar] -> [Constraint] -> Type -> TypeScheme
Forall [TyVar]
vs ((Constraint -> Constraint) -> [Constraint] -> [Constraint]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Constraint -> Constraint
applySubstConstraint Subst
s') [Constraint]
cs) (Subst -> Type -> Type
applySubst Subst
s' Type
t)

-- | Apply a substitution to a constraint
applySubstConstraint :: Subst -> Constraint -> Constraint
applySubstConstraint :: Subst -> Constraint -> Constraint
applySubstConstraint Subst
s (Constraint String
cls Type
ty) = String -> Type -> Constraint
Constraint String
cls (Subst -> Type -> Type
applySubst Subst
s Type
ty)

-- | Index substitution: mapping from index variables to indices
newtype SubstIndex = SubstIndex { SubstIndex -> Map IndexTyVar Index
unSubstIndex :: Map IndexTyVar Index }
  deriving (SubstIndex -> SubstIndex -> Bool
(SubstIndex -> SubstIndex -> Bool)
-> (SubstIndex -> SubstIndex -> Bool) -> Eq SubstIndex
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SubstIndex -> SubstIndex -> Bool
== :: SubstIndex -> SubstIndex -> Bool
$c/= :: SubstIndex -> SubstIndex -> Bool
/= :: SubstIndex -> SubstIndex -> Bool
Eq, Int -> SubstIndex -> ShowS
[SubstIndex] -> ShowS
SubstIndex -> String
(Int -> SubstIndex -> ShowS)
-> (SubstIndex -> String)
-> ([SubstIndex] -> ShowS)
-> Show SubstIndex
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SubstIndex -> ShowS
showsPrec :: Int -> SubstIndex -> ShowS
$cshow :: SubstIndex -> String
show :: SubstIndex -> String
$cshowList :: [SubstIndex] -> ShowS
showList :: [SubstIndex] -> ShowS
Show, (forall x. SubstIndex -> Rep SubstIndex x)
-> (forall x. Rep SubstIndex x -> SubstIndex) -> Generic SubstIndex
forall x. Rep SubstIndex x -> SubstIndex
forall x. SubstIndex -> Rep SubstIndex x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. SubstIndex -> Rep SubstIndex x
from :: forall x. SubstIndex -> Rep SubstIndex x
$cto :: forall x. Rep SubstIndex x -> SubstIndex
to :: forall x. Rep SubstIndex x -> SubstIndex
Generic)

-- | Empty index substitution
emptySubstIndex :: SubstIndex
emptySubstIndex :: SubstIndex
emptySubstIndex = Map IndexTyVar Index -> SubstIndex
SubstIndex Map IndexTyVar Index
forall k a. Map k a
Map.empty

-- | Create an index substitution with a single binding
singletonSubstIndex :: IndexTyVar -> Index -> SubstIndex
singletonSubstIndex :: IndexTyVar -> Index -> SubstIndex
singletonSubstIndex IndexTyVar
v Index
i = Map IndexTyVar Index -> SubstIndex
SubstIndex (Map IndexTyVar Index -> SubstIndex)
-> Map IndexTyVar Index -> SubstIndex
forall a b. (a -> b) -> a -> b
$ IndexTyVar -> Index -> Map IndexTyVar Index
forall k a. k -> a -> Map k a
Map.singleton IndexTyVar
v Index
i

-- | Apply an index substitution to an index specification
applySubstIndex :: SubstIndex -> IndexSpec -> IndexSpec
applySubstIndex :: SubstIndex -> IndexSpec -> IndexSpec
applySubstIndex (SubstIndex Map IndexTyVar Index
m) = (Index -> Index) -> IndexSpec -> IndexSpec
forall a b. (a -> b) -> [a] -> [b]
map Index -> Index
apply
  where
    apply :: Index -> Index
apply i :: Index
i@(IndexVar String
s) = Index -> IndexTyVar -> Map IndexTyVar Index -> Index
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Index
i (String -> IndexTyVar
IndexTyVar String
s) Map IndexTyVar Index
m
    apply Index
i = Index
i