{-# 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
import Control.Monad.Reader
import Control.Monad.Writer
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 :: * -> *}.
(IxValue (f TVarBinding) ~ TVarBinding,
 Index (f TVarBinding) ~ TVar, At (f TVarBinding),
 MonadReader (f TVarBinding) f, 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'