purescript · kl0tl · Jun 28, 2020
diff --git a/lib/purescript-ast/src/Language/PureScript/Environment.hs b/lib/purescript-ast/src/Language/PureScript/Environment.hs
@@ -35,7 +35,7 @@ data Environment = Environment
   , dataConstructors :: M.Map (Qualified (ProperName 'ConstructorName)) (DataDeclType, ProperName 'TypeName, SourceType, [Ident])
   -- ^ Data constructors currently in scope, along with their associated type
   -- constructor name, argument types and return type.
-  , roleDeclarations :: M.Map (Qualified (ProperName 'TypeName)) [Role]
+  , roleDeclarations :: RoleEnv
   -- ^ Explicit role declarations currently in scope. Note that this field is
   -- only used to store declared roles temporarily until they can be checked;
   -- to find a type's real checked and/or inferred roles, refer to the TypeKind
@@ -52,6 +52,21 @@ data Environment = Environment
 
 instance NFData Environment
 
+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)
+
 -- | Information about a type class
 data TypeClassData = TypeClassData
   { typeClassArguments :: [(Text, Maybe SourceType)]

diff --git a/src/Language/PureScript/TypeChecker.hs b/src/Language/PureScript/TypeChecker.hs
@@ -92,8 +92,9 @@ addExplicitRoleDeclaration
   -> [Role]
   -> m ()
 addExplicitRoleDeclaration moduleName name roles = do
-  env <- getEnv
-  putEnv $ env { roleDeclarations = M.insert (Qualified (Just moduleName) name) roles (roleDeclarations env) }
+  let insertRoles = M.insert (Qualified (Just moduleName) name) roles
+  modifyEnv $ \env -> env { roleDeclarations = insertRoles $ roleDeclarations env }
+  modifyRoleEnv insertRoles
 
 addTypeSynonym
   :: (MonadState CheckState m, MonadError MultipleErrors m, MonadWriter MultipleErrors m)
@@ -277,8 +278,8 @@ typeCheckAll moduleName _ = traverse go
       checkDuplicateTypeArguments $ map fst args
       (dataCtors, ctorKind) <- kindOfData moduleName (sa, name, args, dctors)
       let args' = args `withKinds` ctorKind
-      env <- getEnv
-      roles <- checkRoles env moduleName name args' dctors
+      inferDataBindingGroupRoles moduleName [(name, args', dctors)]
+      roles <- checkRoles moduleName name args'
       let args'' = args' `withRoles` roles
       addDataType moduleName dtype name args'' dataCtors ctorKind
     return $ DataDeclaration sa dtype name args dctors
@@ -290,22 +291,22 @@ typeCheckAll moduleName _ = traverse go
         bindingGroupNames = ordNub ((syns^..traverse._2) ++ (dataDecls^..traverse._2._2) ++ (fmap coerceProperName (clss^..traverse._2._2)))
         sss = fmap declSourceSpan tys
     warnAndRethrow (addHint (ErrorInDataBindingGroup bindingGroupNames) . addHint (PositionedError sss)) $ do
-      env <- getEnv
       (syn_ks, data_ks, cls_ks) <- kindsOfAll moduleName syns (fmap snd dataDecls) (fmap snd clss)
       let dataDeclsWithKinds = zipWith (\(dtype, (_, name, args, _)) (dataCtors, ctorKind) -> (dtype, name, args `withKinds` ctorKind, dataCtors, ctorKind)) dataDecls data_ks
-          checkRoles' = checkDataBindingGroupRoles env moduleName $
-            map (\(_, name, args, dataCtors, _) -> (name, args, map fst dataCtors)) dataDeclsWithKinds
+      inferDataBindingGroupRoles moduleName $
+        map (\(_, name, args, dataCtors, _) -> (name, args, map fst dataCtors)) dataDeclsWithKinds
       for_ dataDeclsWithKinds $ \(dtype, name, args', dataCtors, ctorKind) -> do
         when (dtype == Newtype) $ checkNewtype name (map fst dataCtors)
         checkDuplicateTypeArguments $ map fst args'
-        roles <- checkRoles' name args'
+        roles <- checkRoles moduleName name args'
         let args'' = args' `withRoles` roles
         addDataType moduleName dtype name args'' dataCtors ctorKind
       for_ (zip syns syn_ks) $ \((_, name, args, _), (elabTy, kind)) -> do
         checkDuplicateTypeArguments $ map fst args
         let args' = args `withKinds` kind
         addTypeSynonym moduleName name args' elabTy kind
       for_ (zip clss cls_ks) $ \((deps, (sa, pn, _, _, _)), (args', implies', tys', kind)) -> do
+        env <- getEnv
         let qualifiedClassName = Qualified (Just moduleName) pn
         guardWith (errorMessage (DuplicateTypeClass pn (fst sa))) $
           not (M.member qualifiedClassName (typeClasses env))

diff --git a/src/Language/PureScript/TypeChecker/Monad.hs b/src/Language/PureScript/TypeChecker/Monad.hs
@@ -60,6 +60,8 @@ emptySubstitution = Substitution M.empty M.empty
 data CheckState = CheckState
   { checkEnv :: Environment
   -- ^ The current @Environment@
+  , checkRoleEnv :: RoleEnv
+  -- ^ The current role environment used for role inference
   , checkNextType :: Int
   -- ^ The next type unification variable
   , checkNextSkolem :: Int
@@ -79,7 +81,7 @@ data CheckState = CheckState
 
 -- | Create an empty @CheckState@
 emptyCheckState :: Environment -> CheckState
-emptyCheckState env = CheckState env 0 0 0 Nothing emptySubstitution []
+emptyCheckState env = CheckState env (getRoleEnv env) 0 0 0 Nothing emptySubstitution []
 
 -- | Unification variables
 type Unknown = Int
@@ -296,6 +298,10 @@ putEnv env = modify (\s -> s { checkEnv = env })
 modifyEnv :: (MonadState CheckState m) => (Environment -> Environment) -> m ()
 modifyEnv f = modify (\s -> s { checkEnv = f (checkEnv s) })
 
+-- | Modify the @RoleEnv@
+modifyRoleEnv :: (MonadState CheckState m) => (RoleEnv -> RoleEnv) -> m ()
+modifyRoleEnv f = modify (\s -> s { checkRoleEnv = f (checkRoleEnv s) })
+
 -- | Run a computation in the typechecking monad, starting with an empty @Environment@
 runCheck :: (Functor m) => StateT CheckState m a -> m (a, Environment)
 runCheck = runCheck' (emptyCheckState initEnvironment)

diff --git a/src/Language/PureScript/TypeChecker/Roles.hs b/src/Language/PureScript/TypeChecker/Roles.hs
@@ -7,13 +7,14 @@
 module Language.PureScript.TypeChecker.Roles
   ( lookupRoles
   , checkRoles
-  , checkDataBindingGroupRoles
+  , inferDataBindingGroupRoles
   ) where
 
 import Prelude.Compat
 
+import Control.Monad
 import Control.Monad.Error.Class (MonadError(..))
-import Control.Monad.State (MonadState(..), runState, state)
+import Control.Monad.State (MonadState, gets, modify)
 import Data.Coerce (coerce)
 import qualified Data.Map as M
 import Data.Maybe (fromMaybe)
@@ -26,6 +27,7 @@ import Language.PureScript.Errors
 import Language.PureScript.Names
 import Language.PureScript.Roles
 import Language.PureScript.Types
+import Language.PureScript.TypeChecker.Monad
 
 -- |
 -- A map of a type's formal parameter names to their roles. This type's
@@ -43,31 +45,17 @@ 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)
+  :: MonadState CheckState m
+  => 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'
+  -> m Any
+updateRoleEnv qualTyName roles' = do
+  roles <- gets $ fromMaybe (repeat Phantom) . M.lookup qualTyName . checkRoleEnv
+  let mostRestrictiveRoles = zipWith min roles roles'
       didRolesChange = any (uncurry (<)) $ zip mostRestrictiveRoles roles
-  in (Any didRolesChange, M.insert qualTyName mostRestrictiveRoles roleEnv)
+  modify $ \st -> st { checkRoleEnv = M.insert qualTyName mostRestrictiveRoles $ checkRoleEnv st }
+  pure $ Any didRolesChange
 
 -- |
 -- Lookup the roles for a type in the environment. If the type does not have
@@ -81,102 +69,85 @@ lookupRoles
 lookupRoles env tyName =
   fromMaybe [] $ M.lookup tyName (types env) >>= typeKindRoles . snd
 
--- | This function does the following:
---
--- * Infers roles for the given data type declaration
---
--- * Compares the inferred roles to the explicitly declared roles (if any) and
---   ensures that the explicitly declared roles are not more permissive than
---   the inferred ones
+-- |
+-- Compares the inferred roles to the explicitly declared roles (if any) and
+-- ensures that the explicitly declared roles are not more permissive than
+-- the inferred ones
 --
 checkRoles
-  :: forall m
-   . (MonadError MultipleErrors m)
-  => Environment
-  -> ModuleName
+  :: MonadError MultipleErrors m
+  => MonadState CheckState m
+  => 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
   -> m [Role]
-checkRoles env moduleName tyName tyArgs ctors =
-  checkDataBindingGroupRoles env moduleName [(tyName, tyArgs, ctors)] tyName tyArgs
+checkRoles moduleName tyName tyArgs = do
+  let qualTyName = Qualified (Just moduleName) tyName
+  inferredRoles <- gets $ M.lookup qualTyName . checkRoleEnv
+  rethrow (addHint (ErrorInRoleDeclaration tyName)) $ do
+    env <- getEnv
+    case M.lookup qualTyName (roleDeclarations env) of
+      Just declaredRoles -> do
+        let
+          k (var, _) inf dec =
+            if inf < dec
+              then throwError . errorMessage $ RoleMismatch var inf dec
+              else pure dec
+        sequence $ zipWith3 k tyArgs (fromMaybe (repeat Phantom) inferredRoles) declaredRoles
+      Nothing ->
+        pure $ fromMaybe (Phantom <$ tyArgs) inferredRoles
 
 type DataDeclaration =
   ( ProperName 'TypeName
   , [(Text, Maybe SourceType)]
   , [DataConstructorDeclaration]
   )
 
-checkDataBindingGroupRoles
-  :: forall m
-   . (MonadError MultipleErrors m)
-  => Environment
-  -> ModuleName
-  -> [DataDeclaration]
-  -> ProperName 'TypeName
-  -> [(Text, Maybe SourceType)]
-  -> m [Role]
-checkDataBindingGroupRoles env moduleName group =
-  let initialRoleEnv = M.union (roleDeclarations env) (getRoleEnv env)
-      inferredRoleEnv = inferDataBindingGroupRoles moduleName group initialRoleEnv
-  in \tyName tyArgs -> do
-    let qualTyName = Qualified (Just moduleName) tyName
-        inferredRoles = M.lookup qualTyName inferredRoleEnv
-    rethrow (addHint (ErrorInRoleDeclaration tyName)) $ do
-      case M.lookup qualTyName (roleDeclarations env) of
-        Just declaredRoles -> do
-          let
-            k (var, _) inf dec =
-              if inf < dec
-                then throwError . errorMessage $ RoleMismatch var inf dec
-                else pure dec
-          sequence $ zipWith3 k tyArgs (fromMaybe (repeat Phantom) inferredRoles) declaredRoles
-        Nothing ->
-          pure $ fromMaybe (Phantom <$ tyArgs) inferredRoles
-
+-- |
+-- Infers roles for the given data type declarations.
+--
 inferDataBindingGroupRoles
-  :: ModuleName
+  :: MonadState CheckState m
+  => 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'
+  -> m ()
+inferDataBindingGroupRoles moduleName group = do
+  Any didRolesChange <- mconcat <$> traverse (inferDataDeclarationRoles moduleName) group
+  when didRolesChange $ inferDataBindingGroupRoles moduleName group
 
 -- |
 -- 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
+  :: forall m
+   . MonadState CheckState m
+  => ModuleName
   -> DataDeclaration
-  -> RoleEnv
-  -> (Any, RoleEnv)
-inferDataDeclarationRoles moduleName (tyName, tyArgs, ctors) roleEnv =
-  let qualTyName = Qualified (Just 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
+  -> m Any
+inferDataDeclarationRoles moduleName (tyName, tyArgs, ctors) = do
+  let qualName = Qualified (Just moduleName) tyName
+  roleMaps <- traverse (walk mempty . snd) $ ctors >>= dataCtorFields
+  let
+    ctorRoles = getRoleMap $ mconcat roleMaps
+    inferredRoles = map (\(arg, _) -> fromMaybe Phantom (M.lookup arg ctorRoles)) tyArgs
+  updateRoleEnv qualName inferredRoles
   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 :: S.Set Text -> SourceType -> m 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
+        pure mempty
     | otherwise =
-        RoleMap $ M.singleton v Representational
+        pure $ 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
@@ -189,7 +160,9 @@ inferDataDeclarationRoles moduleName (tyName, tyArgs, ctors) roleEnv =
     walk (S.insert tv btvs) t
   walk btvs (RCons _ _ thead ttail) = do
     -- For row types, we just walk along them and collect the results.
-    walk btvs thead <> walk btvs ttail
+    h <- walk btvs thead
+    t <- walk btvs ttail
+    pure (h <> t)
   walk btvs (KindedType _ t _k) =
     -- For kind-annotated types, discard the annotation and recurse on the
     -- type beneath.
@@ -209,24 +182,25 @@ inferDataDeclarationRoles moduleName (tyName, tyArgs, ctors) roleEnv =
           --   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 ->
+          TypeConstructor _ t1Name -> do
+            t1Roles <- gets $ fromMaybe (repeat Phantom) . M.lookup t1Name . checkRoleEnv
             let
-              t1Roles = fromMaybe (repeat Phantom) $ M.lookup t1Name roleEnv
               k role ti = case role of
                 Nominal ->
-                  freeNominals ti
+                  pure $ freeNominals ti
                 Representational ->
                   go ti
                 Phantom ->
-                  mempty
-            in mconcat (zipWith k t1Roles t2s)
+                  pure mempty
+            fmap mconcat (zipWithM 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 t2s
+            r <- go t1
+            pure (r <> foldMap freeNominals t2s)
     | otherwise =
-        mempty
+        pure mempty
     where
       go = walk btvs
       -- Given a type, computes the list of free variables in that type