content: {-# 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 { getRoleMap :: M.Map Text Role }
  
  instance Semigroup RoleMap where
    (<>) =
      coerce @(M.Map Text Role -> _ -> _) @(RoleMap -> _ -> _) (M.unionWith min)
  
  instance Monoid RoleMap where
    mempty =
      RoleMap M.empty
  
  type RoleEnv = M.Map (Qualified (ProperName 'TypeName)) [Role]
  
  typeKindRoles :: TypeKind -> Maybe [Role]
  typeKindRoles = \case
    DataType _ args _ ->
      Just $ map (\(_, _, role) -> role) args
    ExternData roles ->
      Just roles
    _ ->
      Nothing
  
  getRoleEnv :: Environment -> RoleEnv
  getRoleEnv env =
    M.mapMaybe (typeKindRoles . snd) (types env)
  
  updateRoleEnv
    :: Qualified (ProperName 'TypeName)
    -> [Role]
    -> RoleEnv
    -> (Any, RoleEnv)
  updateRoleEnv qualTyName roles' roleEnv =
    let roles = fromMaybe (repeat Phantom) $ M.lookup qualTyName roleEnv
        mostRestrictiveRoles = zipWith min roles roles'
        didRolesChange = any (uncurry (<)) $ zip mostRestrictiveRoles roles
    in (Any didRolesChange, M.insert qualTyName mostRestrictiveRoles 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 env tyName =
    fromMaybe [] $ M.lookup tyName (types env) >>= typeKindRoles . 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 tyArgs declaredRoles = do
    let k (var, _, inf) dec =
          when (inf < dec) . throwError . errorMessage $ RoleMismatch var inf dec
    zipWithM_ k tyArgs declaredRoles
  
  checkRoleDeclarationArity
    :: forall m
     . (MonadError MultipleErrors m)
    => ProperName 'TypeName
    -> [Role]
    -> Int
    -> m ()
  checkRoleDeclarationArity tyName roles expected = do
    let actual = length roles
    unless (expected == actual) $
      throwError . errorMessage $
        RoleDeclarationArityMismatch tyName expected 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 env moduleName tyName tyArgs ctors =
    inferDataBindingGroupRoles env moduleName [] [(tyName, tyArgs, ctors)] tyName tyArgs
  
  inferDataBindingGroupRoles
    :: Environment
    -> ModuleName
    -> [RoleDeclarationData]
    -> [DataDeclaration]
    -> ProperName 'TypeName
    -> [(Text, Maybe SourceType)]
    -> [Role]
  inferDataBindingGroupRoles env moduleName roleDeclarations group =
    let declaredRoleEnv = M.fromList $ map (Qualified (ByModuleName moduleName) . rdeclIdent &&& rdeclRoles) roleDeclarations
        inferredRoleEnv = getRoleEnv env
        initialRoleEnv = declaredRoleEnv `M.union` inferredRoleEnv
        inferredRoleEnv' = inferDataBindingGroupRoles' moduleName group initialRoleEnv
    in \tyName tyArgs ->
          let qualTyName = Qualified (ByModuleName moduleName) tyName
              inferredRoles = M.lookup qualTyName inferredRoleEnv'
          in fromMaybe (Phantom <$ tyArgs) inferredRoles
  
  type DataDeclaration =
    ( ProperName 'TypeName
    , [(Text, Maybe SourceType)]
    , [DataConstructorDeclaration]
    )
  
  inferDataBindingGroupRoles'
    :: ModuleName
    -> [DataDeclaration]
    -> RoleEnv
    -> RoleEnv
  inferDataBindingGroupRoles' moduleName group roleEnv =
    let (Any didRolesChange, roleEnv') = flip runState roleEnv $
          mconcat <$> traverse (state . inferDataDeclarationRoles moduleName) group
    in if didRolesChange
       then inferDataBindingGroupRoles' moduleName group roleEnv'
       else 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 (tyName, tyArgs, ctors) roleEnv =
    let qualTyName = Qualified (ByModuleName moduleName) tyName
        ctorRoles = getRoleMap . foldMap (walk mempty . snd) $ ctors >>= dataCtorFields
        inferredRoles = map (\(arg, _) -> fromMaybe Phantom (M.lookup arg ctorRoles)) tyArgs
    in updateRoleEnv qualTyName inferredRoles 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 btvs (TypeVar _ 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)@).
      | S.member v btvs =
          mempty
      | otherwise =
          RoleMap $ M.singleton v Representational
    walk btvs (ForAll _ _ tv _ t _) =
      -- 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.
      walk (S.insert tv btvs) t
    walk btvs (ConstrainedType _ Constraint{..} t) =
      -- For constrained types, mark all free variables in the constraint
      -- arguments as nominal and recurse on the type beneath the constraint.
      walk btvs t <> foldMap (freeNominals btvs) constraintArgs
    walk btvs (RCons _ _ thead ttail) = do
      -- For row types, we just walk along them and collect the results.
      walk btvs thead <> walk btvs ttail
    walk btvs (KindedType _ t _k) =
      -- For kind-annotated types, discard the annotation and recurse on the
      -- type beneath.
      walk btvs t
    walk btvs t
      | (t1, _, t2s) <- unapplyTypes t
      , not $ null t2s =
          case 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 _ t1Name ->
              let
                t1Roles = fromMaybe (repeat Phantom) $ M.lookup t1Name roleEnv
                k role ti = case role of
                  Nominal ->
                    freeNominals btvs ti
                  Representational ->
                    go ti
                  Phantom ->
                    mempty
              in mconcat (zipWith k t1Roles 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.
            _ -> do
              go t1 <> foldMap (freeNominals btvs) t2s
      | otherwise =
          mempty
      where
        go = walk 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 btvs x =
    let ftvs = filter (flip S.notMember btvs) (freeTypeVariables x)
    in  RoleMap (M.fromList $ map (, Nominal) ftvs)
  