{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
-- | The rules for checking types.
module Glam.Rules.Type where

import Data.Map (Map)
import Data.Map qualified as Map
import Data.Monoid
import Control.Monad.Except hiding (guard)
import Control.Monad.Reader hiding (guard)
import Control.Monad.Writer hiding (guard)
import Control.Lens

import Glam.Utils
import Glam.Type

-- | Things a type variable can refer to
data TVarBinding =
    -- | A type synonym (with a list of arguments)
      Syn [TVar] Type
    -- | A fixed point
    | Self { TVarBinding -> Guardedness
_guardedness :: Guardedness
           , TVarBinding -> Maybe [TVar]
arguments :: Maybe [TVar] -- ^ If this is the type synonym we're defining, contains its arguments
           }
    -- | Anything else
    | Bound

makeLenses ''TVarBinding

type TEnvironment = Map TVar TVarBinding

type MonadCheckType m = (MonadReader TEnvironment m, MonadError String m)

-- | Going under a @▸@
guard :: Guardedness -> Guardedness
guard :: Guardedness -> Guardedness
guard Guardedness
Unguarded = Guardedness
Guarded
guard Guardedness
x = Guardedness
x

-- | Check that a type is well-formed and expand type synonyms.
-- This handles desugaring recursive uses of type synonyms to fixed points:
-- if the current type synonym is used (and applied to the same arguments),
-- the second return value contains the fixed point variable to abstract over.
checkType :: MonadCheckType m => Type -> m (Type, First TVar)
checkType :: forall (m :: * -> *).
MonadCheckType m =>
Type -> m (Type, First TVar)
checkType = WriterT (First TVar) m Type -> m (Type, First TVar)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT (First TVar) m Type -> m (Type, First TVar))
-> (Type -> WriterT (First TVar) m Type)
-> Type
-> m (Type, First TVar)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> WriterT (First TVar) m Type
forall {f :: * -> *} {f :: * -> *} {f :: * -> *}.
(Index (f TVarBinding) ~ TVar,
 IxValue (f TVarBinding) ~ TVarBinding,
 MonadReader (f TVarBinding) f, At (f TVarBinding),
 MonadError TVar f, MonadWriter (f TVar) f, Applicative f,
 Functor f) =>
Type -> f Type
go where
  go :: Type -> f Type
go ty :: Type
ty@TVar{}     = Type -> [Type] -> f Type
apply Type
ty []
  go ty :: Type
ty@TApp{}     = Type -> [Type] -> f Type
apply Type
ty []
  go (Type
ta :*: Type
tb)   = Type -> Type -> Type
(:*:) (Type -> Type -> Type) -> f Type -> f (Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
ta f (Type -> Type) -> f Type -> f Type
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> f Type
go Type
tb
  go (Type
ta :+: Type
tb)   = Type -> Type -> Type
(:+:) (Type -> Type -> Type) -> f Type -> f (Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
ta f (Type -> Type) -> f Type -> f Type
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> f Type
go Type
tb
  go (Type
ta :->: Type
tb)  = Type -> Type -> Type
(:->:) (Type -> Type -> Type) -> f Type -> f (Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
ta f (Type -> Type) -> f Type -> f Type
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> f Type
go Type
tb
  go (Later Type
ty)    = (TVarBinding -> Identity TVarBinding)
-> f TVarBinding -> Identity (f TVarBinding)
Setter (f TVarBinding) (f TVarBinding) TVarBinding TVarBinding
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped((TVarBinding -> Identity TVarBinding)
 -> f TVarBinding -> Identity (f TVarBinding))
-> ((Guardedness -> Identity Guardedness)
    -> TVarBinding -> Identity TVarBinding)
-> (Guardedness -> Identity Guardedness)
-> f TVarBinding
-> Identity (f TVarBinding)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Guardedness -> Identity Guardedness)
-> TVarBinding -> Identity TVarBinding
Traversal' TVarBinding Guardedness
guardedness ((Guardedness -> Identity Guardedness)
 -> f TVarBinding -> Identity (f TVarBinding))
-> (Guardedness -> Guardedness) -> f TVarBinding -> f TVarBinding
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Guardedness -> Guardedness
guard (f TVarBinding -> f TVarBinding) -> f Type -> f Type
forall {r} {m :: * -> *} {a}.
MonadReader r m =>
(r -> r) -> m a -> m a
|- Type -> Type
Later (Type -> Type) -> f Type -> f Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
ty
  go (Constant Type
ty) = (TVarBinding -> Identity TVarBinding)
-> f TVarBinding -> Identity (f TVarBinding)
Setter (f TVarBinding) (f TVarBinding) TVarBinding TVarBinding
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped((TVarBinding -> Identity TVarBinding)
 -> f TVarBinding -> Identity (f TVarBinding))
-> ((Guardedness -> Identity Guardedness)
    -> TVarBinding -> Identity TVarBinding)
-> (Guardedness -> Identity Guardedness)
-> f TVarBinding
-> Identity (f TVarBinding)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Guardedness -> Identity Guardedness)
-> TVarBinding -> Identity TVarBinding
Traversal' TVarBinding Guardedness
guardedness ((Guardedness -> Identity Guardedness)
 -> f TVarBinding -> Identity (f TVarBinding))
-> Guardedness -> f TVarBinding -> f TVarBinding
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Guardedness
Forbidden (f TVarBinding -> f TVarBinding) -> f Type -> f Type
forall {r} {m :: * -> *} {a}.
MonadReader r m =>
(r -> r) -> m a -> m a
|- Type -> Type
Constant (Type -> Type) -> f Type -> f Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
ty
  go (TFix TVar
x Type
tf)   = Index (f TVarBinding)
-> Lens' (f TVarBinding) (Maybe (IxValue (f TVarBinding)))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at TVar
Index (f TVarBinding)
x ((Maybe (IxValue (f TVarBinding)) -> Identity (Maybe TVarBinding))
 -> f TVarBinding -> Identity (f TVarBinding))
-> TVarBinding -> f TVarBinding -> f TVarBinding
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Guardedness -> Maybe [TVar] -> TVarBinding
Self Guardedness
Unguarded Maybe [TVar]
forall a. Maybe a
Nothing (f TVarBinding -> f TVarBinding) -> f Type -> f Type
forall {r} {m :: * -> *} {a}.
MonadReader r m =>
(r -> r) -> m a -> m a
|- TVar -> Type -> Type
TFix TVar
x (Type -> Type) -> f Type -> f Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> f Type
go Type
tf
  go Type
ty            = Type -> f Type
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty
  apply :: Type -> [Type] -> f Type
apply (TApp Type
t1 Type
t2) [Type]
tys = do
      Type
t2' <- Type -> f Type
go Type
t2
      Type -> [Type] -> f Type
apply Type
t1 (Type
t2'Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
tys)
  apply (TVar TVar
x) [Type]
tys = Getting (Maybe TVarBinding) (f TVarBinding) (Maybe TVarBinding)
-> f (Maybe TVarBinding)
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Index (f TVarBinding)
-> Lens' (f TVarBinding) (Maybe (IxValue (f TVarBinding)))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at TVar
Index (f TVarBinding)
x) f (Maybe TVarBinding) -> (Maybe TVarBinding -> f Type) -> f Type
forall a b. f a -> (a -> f b) -> f b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= f Type -> (TVarBinding -> f Type) -> Maybe TVarBinding -> f Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$ TVar
"unbound type variable " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x) \case
      Syn [TVar]
ys Type
ty
          | [TVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TVar]
ys Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys -> Type -> f Type
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> f Type) -> Type -> f Type
forall a b. (a -> b) -> a -> b
$ TSubst -> Type -> Type
substituteType ([(TVar, Type)] -> TSubst
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([TVar] -> [Type] -> [(TVar, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TVar]
ys [Type]
tys)) Type
ty
          | Bool
otherwise -> TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$
              TVar
"wrong number of arguments for type constructor " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
": got " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++
              Int -> TVar
forall a. Show a => a -> TVar
show ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys) TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
", expecting " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ Int -> TVar
forall a. Show a => a -> TVar
show ([TVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TVar]
ys)
      Self Guardedness
_ (Just [TVar]
ys) | [Type]
tys [Type] -> [Type] -> Bool
forall a. Eq a => a -> a -> Bool
/= (TVar -> Type) -> [TVar] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TVar -> Type
TVar [TVar]
ys -> TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$
          TVar
"recursive type constructor " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
" must be applied to the same arguments"
      Self Guardedness
Unguarded Maybe [TVar]
_ -> TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$ TVar
"unguarded fixed point variable " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x
      Self Guardedness
Forbidden Maybe [TVar]
_ -> TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$ TVar
"fixed point variable " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
" cannot appear under #"
      Self Guardedness
Guarded (Just [TVar]
_) -> TVar -> Type
TVar TVar
x Type -> f () -> f Type
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f TVar -> f ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (TVar -> f TVar
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TVar
x)
      Self Guardedness
Guarded Maybe [TVar]
Nothing -> Type -> f Type
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar -> Type
TVar TVar
x)
      TVarBinding
Bound | [Type] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
tys -> Type -> f Type
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar -> Type
TVar TVar
x)
            | Bool
otherwise -> TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$ TVar
"not a type constructor: " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ TVar
x
  apply Type
ty [Type]
_ = TVar -> f Type
forall a. TVar -> f a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (TVar -> f Type) -> TVar -> f Type
forall a b. (a -> b) -> a -> b
$ TVar
"not a type constructor: " TVar -> TVar -> TVar
forall a. [a] -> [a] -> [a]
++ Type -> TVar
forall a. Show a => a -> TVar
show Type
ty

-- | Check a polytype
checkPolytype :: MonadCheckType m => Polytype -> m Polytype
checkPolytype :: forall (m :: * -> *). MonadCheckType m => Polytype -> m Polytype
checkPolytype (Forall [(TVar, Bool)]
as Type
ty) = do
    let scope :: Map TVar TVarBinding
scope = [(TVar, TVarBinding)] -> Map TVar TVarBinding
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(TVar
a, TVarBinding
Bound) | (TVar
a, Bool
_) <- [(TVar, Bool)]
as]
    (Type
ty', First TVar
_) <- Map TVar TVarBinding
-> Map TVar TVarBinding -> Map TVar TVarBinding
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map TVar TVarBinding
scope (Map TVar TVarBinding -> Map TVar TVarBinding)
-> m (Type, First TVar) -> m (Type, First TVar)
forall {r} {m :: * -> *} {a}.
MonadReader r m =>
(r -> r) -> m a -> m a
|- Type -> m (Type, First TVar)
forall (m :: * -> *).
MonadCheckType m =>
Type -> m (Type, First TVar)
checkType Type
ty
    Polytype -> m Polytype
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Polytype -> m Polytype) -> Polytype -> m Polytype
forall a b. (a -> b) -> a -> b
$ [(TVar, Bool)] -> Type -> Polytype
Forall [(TVar, Bool)]
as Type
ty'

-- | Check a type synonym definition
checkTypeSynonym :: MonadCheckType m => TVar -> [TVar] -> Type -> m Type
checkTypeSynonym :: forall (m :: * -> *).
MonadCheckType m =>
TVar -> [TVar] -> Type -> m Type
checkTypeSynonym TVar
x [TVar]
ys Type
ty = do
    let scope :: Map TVar TVarBinding
scope = [(TVar, TVarBinding)] -> Map TVar TVarBinding
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(TVar, TVarBinding)] -> Map TVar TVarBinding)
-> [(TVar, TVarBinding)] -> Map TVar TVarBinding
forall a b. (a -> b) -> a -> b
$ (TVar
x, Guardedness -> Maybe [TVar] -> TVarBinding
Self Guardedness
Unguarded ([TVar] -> Maybe [TVar]
forall a. a -> Maybe a
Just [TVar]
ys)) (TVar, TVarBinding)
-> [(TVar, TVarBinding)] -> [(TVar, TVarBinding)]
forall a. a -> [a] -> [a]
: [(TVar
y, TVarBinding
Bound) | TVar
y <- [TVar]
ys]
    (Type
ty', First Maybe TVar
autofix) <- Map TVar TVarBinding
-> Map TVar TVarBinding -> Map TVar TVarBinding
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map TVar TVarBinding
scope (Map TVar TVarBinding -> Map TVar TVarBinding)
-> m (Type, First TVar) -> m (Type, First TVar)
forall {r} {m :: * -> *} {a}.
MonadReader r m =>
(r -> r) -> m a -> m a
|- Type -> m (Type, First TVar)
forall (m :: * -> *).
MonadCheckType m =>
Type -> m (Type, First TVar)
checkType Type
ty
    Type -> m Type
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type)
-> (TVar -> Type -> Type) -> Maybe TVar -> Type -> Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Type -> Type
forall a. a -> a
id TVar -> Type -> Type
TFix Maybe TVar
autofix Type
ty'