{-# 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 qualified Data.Map as M
import Data.Maybe (fromMaybe)
import qualified Data.Set as S
import Data.Semigroup (Any(..))
import Data.Text (Text)

import Language.PureScript.Environment
import Language.PureScript.Errors
import Language.PureScript.Names
import Language.PureScript.Roles
import Language.PureScript.Types

-- |
-- 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)