{-# LANGUAGE TypeApplications #-}

-- |
-- Role inference
--
module Language.PureScript.TypeChecker.Roles
  ( lookupRoles
  , checkRoles
  , checkRoleDeclarationArity
  , inferRoles
  , inferDataBindingGroupRoles
  ) where

import Prelude

import Control.Arrow ((&&&))
import Control.Monad (unless, when, zipWithM_)
import Control.Monad.Error.Class (MonadError(..))
import Control.Monad.State (MonadState(..), runState, state)
import Data.Coerce (coerce)
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Data.Semigroup (Any(..))
import Data.Text (Text)

import Language.PureScript.Environment (Environment(..), TypeKind(..))
import Language.PureScript.Errors (DataConstructorDeclaration(..), MultipleErrors, RoleDeclarationData(..), SimpleErrorMessage(..), errorMessage)
import Language.PureScript.Names (ModuleName, ProperName, ProperNameType(..), Qualified(..), QualifiedBy(..))
import Language.PureScript.Roles (Role(..))
import Language.PureScript.Types (Constraint(..), SourceType, Type(..), freeTypeVariables, unapplyTypes)

-- |
-- A map of a type's formal parameter names to their roles. This type's
-- @Semigroup@ and @Monoid@ instances preserve the least-permissive role
-- ascribed to any given variable, as defined by the @Role@ type's @Ord@
-- instance. That is, a variable that has been marked as @Nominal@ can not
-- later be marked @Representational@, and so on.
newtype RoleMap = RoleMap { RoleMap -> Map Text Role
getRoleMap :: M.Map Text Role }

instance Semigroup RoleMap where
  <> :: RoleMap -> RoleMap -> RoleMap
(<>) =
    coerce :: forall a b. Coercible a b => a -> b
coerce @(M.Map Text Role -> _ -> _) @(RoleMap -> _ -> _) (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith forall a. Ord a => a -> a -> a
min)

instance Monoid RoleMap where
  mempty :: RoleMap
mempty =
    Map Text Role -> RoleMap
RoleMap forall k a. Map k a
M.empty

type RoleEnv = M.Map (Qualified (ProperName 'TypeName)) [Role]

typeKindRoles :: TypeKind -> Maybe [Role]
typeKindRoles :: TypeKind -> Maybe [Role]
typeKindRoles = \case
  DataType DataDeclType
_ [(Text, Maybe SourceType, Role)]
args [(ProperName 'ConstructorName, [SourceType])]
_ ->
    forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(Text
_, Maybe SourceType
_, Role
role) -> Role
role) [(Text, Maybe SourceType, Role)]
args
  ExternData [Role]
roles ->
    forall a. a -> Maybe a
Just [Role]
roles
  TypeKind
_ ->
    forall a. Maybe a
Nothing

getRoleEnv :: Environment -> RoleEnv
getRoleEnv :: Environment -> RoleEnv
getRoleEnv Environment
env =
  forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (TypeKind -> Maybe [Role]
typeKindRoles forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (Environment
-> Map (Qualified (ProperName 'TypeName)) (SourceType, TypeKind)
types Environment
env)

updateRoleEnv
  :: Qualified (ProperName 'TypeName)
  -> [Role]
  -> RoleEnv
  -> (Any, RoleEnv)
updateRoleEnv :: Qualified (ProperName 'TypeName)
-> [Role] -> RoleEnv -> (Any, RoleEnv)
updateRoleEnv Qualified (ProperName 'TypeName)
qualTyName [Role]
roles' RoleEnv
roleEnv =
  let roles :: [Role]
roles = forall a. a -> Maybe a -> a
fromMaybe (forall a. a -> [a]
repeat Role
Phantom) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Qualified (ProperName 'TypeName)
qualTyName RoleEnv
roleEnv
      mostRestrictiveRoles :: [Role]
mostRestrictiveRoles = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Ord a => a -> a -> a
min [Role]
roles [Role]
roles'
      didRolesChange :: Bool
didRolesChange = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Ord a => a -> a -> Bool
(<)) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Role]
mostRestrictiveRoles [Role]
roles
  in (Bool -> Any
Any Bool
didRolesChange, forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Qualified (ProperName 'TypeName)
qualTyName [Role]
mostRestrictiveRoles RoleEnv
roleEnv)

-- |
-- Lookup the roles for a type in the environment. If the type does not have
-- roles (e.g. is a type synonym or a type variable), then this function
-- returns an empty list.
--
lookupRoles
  :: Environment
  -> Qualified (ProperName 'TypeName)
  -> [Role]
lookupRoles :: Environment -> Qualified (ProperName 'TypeName) -> [Role]
lookupRoles Environment
env Qualified (ProperName 'TypeName)
tyName =
  forall a. a -> Maybe a -> a
fromMaybe [] forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Qualified (ProperName 'TypeName)
tyName (Environment
-> Map (Qualified (ProperName 'TypeName)) (SourceType, TypeKind)
types Environment
env) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TypeKind -> Maybe [Role]
typeKindRoles forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd

-- |
-- Compares the inferred roles to the explicitly declared roles and ensures
-- that the explicitly declared roles are not more permissive than the
-- inferred ones.
--
checkRoles
  :: forall m
   . (MonadError MultipleErrors m)
  => [(Text, Maybe SourceType, Role)]
    -- ^ type parameters for the data type whose roles we are checking
  -> [Role]
    -- ^ roles declared for the data type
  -> m ()
checkRoles :: forall (m :: * -> *).
MonadError MultipleErrors m =>
[(Text, Maybe SourceType, Role)] -> [Role] -> m ()
checkRoles [(Text, Maybe SourceType, Role)]
tyArgs [Role]
declaredRoles = do
  let k :: (Text, b, Role) -> Role -> f ()
k (Text
var, b
_, Role
inf) Role
dec =
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Role
inf forall a. Ord a => a -> a -> Bool
< Role
dec) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpleErrorMessage -> MultipleErrors
errorMessage forall a b. (a -> b) -> a -> b
$ Text -> Role -> Role -> SimpleErrorMessage
RoleMismatch Text
var Role
inf Role
dec
  forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ forall {f :: * -> *} {b}.
MonadError MultipleErrors f =>
(Text, b, Role) -> Role -> f ()
k [(Text, Maybe SourceType, Role)]
tyArgs [Role]
declaredRoles

checkRoleDeclarationArity
  :: forall m
   . (MonadError MultipleErrors m)
  => ProperName 'TypeName
  -> [Role]
  -> Int
  -> m ()
checkRoleDeclarationArity :: forall (m :: * -> *).
MonadError MultipleErrors m =>
ProperName 'TypeName -> [Role] -> Int -> m ()
checkRoleDeclarationArity ProperName 'TypeName
tyName [Role]
roles Int
expected = do
  let actual :: Int
actual = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Role]
roles
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
expected forall a. Eq a => a -> a -> Bool
== Int
actual) forall a b. (a -> b) -> a -> b
$
    forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpleErrorMessage -> MultipleErrors
errorMessage forall a b. (a -> b) -> a -> b
$
      ProperName 'TypeName -> Int -> Int -> SimpleErrorMessage
RoleDeclarationArityMismatch ProperName 'TypeName
tyName Int
expected Int
actual

-- |
-- Infers roles for the given data type declaration.
--
inferRoles
  :: Environment
  -> ModuleName
  -> ProperName 'TypeName
    -- ^ The name of the data type whose roles we are checking
  -> [(Text, Maybe SourceType)]
    -- ^ type parameters for the data type whose roles we are checking
  -> [DataConstructorDeclaration]
    -- ^ constructors of the data type whose roles we are checking
  -> [Role]
inferRoles :: Environment
-> ModuleName
-> ProperName 'TypeName
-> [(Text, Maybe SourceType)]
-> [DataConstructorDeclaration]
-> [Role]
inferRoles Environment
env ModuleName
moduleName ProperName 'TypeName
tyName [(Text, Maybe SourceType)]
tyArgs [DataConstructorDeclaration]
ctors =
  Environment
-> ModuleName
-> [RoleDeclarationData]
-> [DataDeclaration]
-> ProperName 'TypeName
-> [(Text, Maybe SourceType)]
-> [Role]
inferDataBindingGroupRoles Environment
env ModuleName
moduleName [] [(ProperName 'TypeName
tyName, [(Text, Maybe SourceType)]
tyArgs, [DataConstructorDeclaration]
ctors)] ProperName 'TypeName
tyName [(Text, Maybe SourceType)]
tyArgs

inferDataBindingGroupRoles
  :: Environment
  -> ModuleName
  -> [RoleDeclarationData]
  -> [DataDeclaration]
  -> ProperName 'TypeName
  -> [(Text, Maybe SourceType)]
  -> [Role]
inferDataBindingGroupRoles :: Environment
-> ModuleName
-> [RoleDeclarationData]
-> [DataDeclaration]
-> ProperName 'TypeName
-> [(Text, Maybe SourceType)]
-> [Role]
inferDataBindingGroupRoles Environment
env ModuleName
moduleName [RoleDeclarationData]
roleDeclarations [DataDeclaration]
group =
  let declaredRoleEnv :: RoleEnv
declaredRoleEnv = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. QualifiedBy -> a -> Qualified a
Qualified (ModuleName -> QualifiedBy
ByModuleName ModuleName
moduleName) forall b c a. (b -> c) -> (a -> b) -> a -> c
. RoleDeclarationData -> ProperName 'TypeName
rdeclIdent forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& RoleDeclarationData -> [Role]
rdeclRoles) [RoleDeclarationData]
roleDeclarations
      inferredRoleEnv :: RoleEnv
inferredRoleEnv = Environment -> RoleEnv
getRoleEnv Environment
env
      initialRoleEnv :: RoleEnv
initialRoleEnv = RoleEnv
declaredRoleEnv forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` RoleEnv
inferredRoleEnv
      inferredRoleEnv' :: RoleEnv
inferredRoleEnv' = ModuleName -> [DataDeclaration] -> RoleEnv -> RoleEnv
inferDataBindingGroupRoles' ModuleName
moduleName [DataDeclaration]
group RoleEnv
initialRoleEnv
  in \ProperName 'TypeName
tyName [(Text, Maybe SourceType)]
tyArgs ->
        let qualTyName :: Qualified (ProperName 'TypeName)
qualTyName = forall a. QualifiedBy -> a -> Qualified a
Qualified (ModuleName -> QualifiedBy
ByModuleName ModuleName
moduleName) ProperName 'TypeName
tyName
            inferredRoles :: Maybe [Role]
inferredRoles = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Qualified (ProperName 'TypeName)
qualTyName RoleEnv
inferredRoleEnv'
        in forall a. a -> Maybe a -> a
fromMaybe (Role
Phantom forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [(Text, Maybe SourceType)]
tyArgs) Maybe [Role]
inferredRoles

type DataDeclaration =
  ( ProperName 'TypeName
  , [(Text, Maybe SourceType)]
  , [DataConstructorDeclaration]
  )

inferDataBindingGroupRoles'
  :: ModuleName
  -> [DataDeclaration]
  -> RoleEnv
  -> RoleEnv
inferDataBindingGroupRoles' :: ModuleName -> [DataDeclaration] -> RoleEnv -> RoleEnv
inferDataBindingGroupRoles' ModuleName
moduleName [DataDeclaration]
group RoleEnv
roleEnv =
  let (Any Bool
didRolesChange, RoleEnv
roleEnv') = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> (a, s)
runState RoleEnv
roleEnv forall a b. (a -> b) -> a -> b
$
        forall a. Monoid a => [a] -> a
mconcat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModuleName -> DataDeclaration -> RoleEnv -> (Any, RoleEnv)
inferDataDeclarationRoles ModuleName
moduleName) [DataDeclaration]
group
  in if Bool
didRolesChange
     then ModuleName -> [DataDeclaration] -> RoleEnv -> RoleEnv
inferDataBindingGroupRoles' ModuleName
moduleName [DataDeclaration]
group RoleEnv
roleEnv'
     else RoleEnv
roleEnv'

-- |
-- Infers roles for the given data type declaration, along with a flag to tell
-- if more restrictive roles were added to the environment.
--
inferDataDeclarationRoles
  :: ModuleName
  -> DataDeclaration
  -> RoleEnv
  -> (Any, RoleEnv)
inferDataDeclarationRoles :: ModuleName -> DataDeclaration -> RoleEnv -> (Any, RoleEnv)
inferDataDeclarationRoles ModuleName
moduleName (ProperName 'TypeName
tyName, [(Text, Maybe SourceType)]
tyArgs, [DataConstructorDeclaration]
ctors) RoleEnv
roleEnv =
  let qualTyName :: Qualified (ProperName 'TypeName)
qualTyName = forall a. QualifiedBy -> a -> Qualified a
Qualified (ModuleName -> QualifiedBy
ByModuleName ModuleName
moduleName) ProperName 'TypeName
tyName
      ctorRoles :: Map Text Role
ctorRoles = RoleMap -> Map Text Role
getRoleMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set Text -> SourceType -> RoleMap
walk forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$ [DataConstructorDeclaration]
ctors forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DataConstructorDeclaration -> [(Ident, SourceType)]
dataCtorFields
      inferredRoles :: [Role]
inferredRoles = forall a b. (a -> b) -> [a] -> [b]
map (\(Text
arg, Maybe SourceType
_) -> forall a. a -> Maybe a -> a
fromMaybe Role
Phantom (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
arg Map Text Role
ctorRoles)) [(Text, Maybe SourceType)]
tyArgs
  in Qualified (ProperName 'TypeName)
-> [Role] -> RoleEnv -> (Any, RoleEnv)
updateRoleEnv Qualified (ProperName 'TypeName)
qualTyName [Role]
inferredRoles RoleEnv
roleEnv
  where
  -- This function is named @walk@ to match the specification given in the
  -- "Role inference" section of the paper "Safe Zero-cost Coercions for
  -- Haskell".
  walk :: S.Set Text -> SourceType -> RoleMap
  walk :: Set Text -> SourceType -> RoleMap
walk Set Text
btvs (TypeVar SourceAnn
_ Text
v)
    -- A type variable standing alone (e.g. @a@ in @data D a b = D a@) is
    -- representational, _unless_ it has been bound by a quantifier, in which
    -- case it is not actually a parameter to the type (e.g. @z@ in
    -- @data T z = T (forall z. z -> z)@).
    | forall a. Ord a => a -> Set a -> Bool
S.member Text
v Set Text
btvs =
        forall a. Monoid a => a
mempty
    | Bool
otherwise =
        Map Text Role -> RoleMap
RoleMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton Text
v Role
Representational
  walk Set Text
btvs (ForAll SourceAnn
_ Text
tv Maybe SourceType
_ SourceType
t Maybe SkolemScope
_) =
    -- We can walk under universal quantifiers as long as we make note of the
    -- variables that they bind. For instance, given a definition
    -- @data T z = T (forall z. z -> z)@, we will make note that @z@ is bound
    -- by a quantifier so that we do not mark @T@'s parameter as
    -- representational later on. Similarly, given a definition like
    -- @data D a = D (forall r. r -> a)@, we'll mark @r@ as bound so that it
    -- doesn't appear as a spurious parameter to @D@ when we complete
    -- inference.
    Set Text -> SourceType -> RoleMap
walk (forall a. Ord a => a -> Set a -> Set a
S.insert Text
tv Set Text
btvs) SourceType
t
  walk Set Text
btvs (ConstrainedType SourceAnn
_ Constraint{[SourceType]
Maybe ConstraintData
SourceAnn
Qualified (ProperName 'ClassName)
constraintData :: forall a. Constraint a -> Maybe ConstraintData
constraintArgs :: forall a. Constraint a -> [Type a]
constraintKindArgs :: forall a. Constraint a -> [Type a]
constraintClass :: forall a. Constraint a -> Qualified (ProperName 'ClassName)
constraintAnn :: forall a. Constraint a -> a
constraintData :: Maybe ConstraintData
constraintArgs :: [SourceType]
constraintKindArgs :: [SourceType]
constraintClass :: Qualified (ProperName 'ClassName)
constraintAnn :: SourceAnn
..} SourceType
t) =
    -- For constrained types, mark all free variables in the constraint
    -- arguments as nominal and recurse on the type beneath the constraint.
    Set Text -> SourceType -> RoleMap
walk Set Text
btvs SourceType
t forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set Text -> SourceType -> RoleMap
freeNominals Set Text
btvs) [SourceType]
constraintArgs
  walk Set Text
btvs (RCons SourceAnn
_ Label
_ SourceType
thead SourceType
ttail) = do
    -- For row types, we just walk along them and collect the results.
    Set Text -> SourceType -> RoleMap
walk Set Text
btvs SourceType
thead forall a. Semigroup a => a -> a -> a
<> Set Text -> SourceType -> RoleMap
walk Set Text
btvs SourceType
ttail
  walk Set Text
btvs (KindedType SourceAnn
_ SourceType
t SourceType
_k) =
    -- For kind-annotated types, discard the annotation and recurse on the
    -- type beneath.
    Set Text -> SourceType -> RoleMap
walk Set Text
btvs SourceType
t
  walk Set Text
btvs SourceType
t
    | (SourceType
t1, [SourceType]
_, [SourceType]
t2s) <- forall a. Type a -> (Type a, [Type a], [Type a])
unapplyTypes SourceType
t
    , Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null [SourceType]
t2s =
        case SourceType
t1 of
          -- If the type is an application of a type constructor to some
          -- arguments, recursively infer the roles of the type constructor's
          -- arguments. For each (role, argument) pair:
          --
          -- - If the role is nominal, mark all free variables in the argument
          -- as nominal also, since they cannot be coerced if the
          -- argument's nominality is to be preserved.
          --
          -- - If the role is representational, recurse on the argument, since
          -- its use of our parameters is important.
          --
          -- - If the role is phantom, terminate, since the argument's use of
          -- our parameters is unimportant.
          TypeConstructor SourceAnn
_ Qualified (ProperName 'TypeName)
t1Name ->
            let
              t1Roles :: [Role]
t1Roles = forall a. a -> Maybe a -> a
fromMaybe (forall a. a -> [a]
repeat Role
Phantom) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Qualified (ProperName 'TypeName)
t1Name RoleEnv
roleEnv
              k :: Role -> SourceType -> RoleMap
k Role
role SourceType
ti = case Role
role of
                Role
Nominal ->
                  Set Text -> SourceType -> RoleMap
freeNominals Set Text
btvs SourceType
ti
                Role
Representational ->
                  SourceType -> RoleMap
go SourceType
ti
                Role
Phantom ->
                  forall a. Monoid a => a
mempty
            in forall a. Monoid a => [a] -> a
mconcat (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Role -> SourceType -> RoleMap
k [Role]
t1Roles [SourceType]
t2s)
          -- If the type is an application of any other type-level term, walk
          -- that term to collect its roles and mark all free variables in
          -- its argument as nominal.
          SourceType
_ -> do
            SourceType -> RoleMap
go SourceType
t1 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set Text -> SourceType -> RoleMap
freeNominals Set Text
btvs) [SourceType]
t2s
    | Bool
otherwise =
        forall a. Monoid a => a
mempty
    where
      go :: SourceType -> RoleMap
go = Set Text -> SourceType -> RoleMap
walk Set Text
btvs

-- Given a type, computes the list of free variables in that type
-- (taking into account those bound in @walk@) and returns a @RoleMap@
-- ascribing a nominal role to each of those variables.
freeNominals :: S.Set Text -> SourceType -> RoleMap
freeNominals :: Set Text -> SourceType -> RoleMap
freeNominals Set Text
btvs SourceType
x =
  let ftvs :: [Text]
ftvs = forall a. (a -> Bool) -> [a] -> [a]
filter (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Ord a => a -> Set a -> Bool
S.notMember Set Text
btvs) (forall a. Type a -> [Text]
freeTypeVariables SourceType
x)
  in  Map Text Role -> RoleMap
RoleMap (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (, Role
Nominal) [Text]
ftvs)