Optimize replaceAllTypeSynonyms with per-node TypeFlags

src/Language/PureScript/TypeChecker/Skolems.hs

@@ -10,6 +10,7 @@ module Language.PureScript.TypeChecker.Skolems
 
 import Prelude
 
+import Control.Exception (assert)
 import Control.Monad.Error.Class (MonadError(..))
 import Control.Monad.State.Class (MonadState(..), gets, modify)
 import Data.Foldable (traverse_)
@@ -21,7 +22,7 @@ import Language.PureScript.Crash (internalError)
 import Language.PureScript.Errors (ErrorMessage(..), MultipleErrors, SimpleErrorMessage(..), positionedError, singleError)
 import Language.PureScript.Traversals (defS)
 import Language.PureScript.TypeChecker.Monad (CheckState(..))
-import Language.PureScript.Types (SkolemScope(..), SourceType, Type(..), everythingOnTypes, everywhereOnTypesM, replaceTypeVars)
+import Language.PureScript.Types (SkolemScope(..), SourceType, Type(..), everythingOnTypes, everywhereOnTypesM, hasFlag, replaceTypeVars, tfHasUnscopedForAlls, typeFlags)
 
 -- | Generate a new skolem constant
 newSkolemConstant :: MonadState CheckState m => m Int
@@ -30,13 +31,27 @@ newSkolemConstant = do
   modify $ \st -> st { checkNextSkolem = s + 1 }
   return s
 
--- | Introduce skolem scope at every occurrence of a ForAll
+-- | Introduce skolem scope at every occurrence of a ForAll.
+-- Short-circuits if the type has no unscoped ForAlls.
 introduceSkolemScope :: MonadState CheckState m => Type a -> m (Type a)
-introduceSkolemScope = everywhereOnTypesM go
+introduceSkolemScope ty
+  -- Sanity check in debug builds: the flag says no unscoped ForAlls exist,
+  -- so a scan should agree. 'assert' is compiled away with -O.
+  | not (hasFlag tfHasUnscopedForAlls (typeFlags ty)) =
+      return $! assert (not (containsUnscopedForAlls ty)) ty
+  | otherwise = everywhereOnTypesM go ty
   where
-  go (ForAll ann vis ident mbK ty Nothing) = ForAll ann vis ident mbK ty <$> (Just <$> newSkolemScope)
+  go (ForAll ann vis ident mbK t Nothing) = ForAll ann vis ident mbK t <$> (Just <$> newSkolemScope)
   go other = return other
 
+-- | Scan a type for ForAll nodes missing a SkolemScope.
+-- Used as a correctness check for the 'tfHasUnscopedForAlls' flag.
+containsUnscopedForAlls :: Type a -> Bool
+containsUnscopedForAlls = everythingOnTypes (||) isUnscoped
+  where
+    isUnscoped (ForAll _ _ _ _ _ Nothing) = True
+    isUnscoped _ = False
+
 -- | Generate a new skolem scope
 newSkolemScope :: MonadState CheckState m => m SkolemScope
 newSkolemScope = do

src/Language/PureScript/TypeChecker/Synonyms.hs

@@ -11,6 +11,7 @@ module Language.PureScript.TypeChecker.Synonyms
 
 import Prelude
 
+import Control.Exception (assert)
 import Control.Monad.Error.Class (MonadError(..))
 import Data.Maybe (fromMaybe)
 import Data.Map qualified as M
@@ -19,22 +20,71 @@ import Language.PureScript.Environment (Environment(..), TypeKind)
 import Language.PureScript.Errors (MultipleErrors, SimpleErrorMessage(..), SourceSpan, errorMessage')
 import Language.PureScript.Names (ProperName, ProperNameType(..), Qualified)
 import Language.PureScript.TypeChecker.Monad (getEnv, TypeCheckM)
-import Language.PureScript.Types (SourceType, Type(..), completeBinderList, everywhereOnTypesTopDownM, getAnnForType, replaceAllTypeVars)
+import Language.PureScript.Types
+  ( SourceType, Type(..), TypeFlags
+  , combineFlags, completeBinderList, constraintNodeFlags, everythingOnTypes, forAllNodeFlags
+  , getAnnForType, hasFlag, overConstraintArgsAll, replaceAllTypeVars
+  , setFlag, skolemNodeFlags, tfSynonymsFree, typeFlags
+  )
 
 -- | Type synonym information (arguments with kinds, aliased type), indexed by name
 type SynonymMap = M.Map (Qualified (ProperName 'TypeName)) ([(Text, Maybe SourceType)], SourceType)
 
 type KindMap = M.Map (Qualified (ProperName 'TypeName)) (SourceType, TypeKind)
 
+-- | Replace fully applied type synonyms and mark every output node
+-- with 'tfSynonymsFree'. Uses a custom traversal that:
+--
+-- 1. Short-circuits on subtrees already marked synonym-free
+-- 2. Only tries synonym expansion on potential application heads
+-- 3. Sets 'tfSynonymsFree' on every output node in a single pass
 replaceAllTypeSynonyms'
   :: SynonymMap
   -> KindMap
   -> SourceType
   -> Either MultipleErrors SourceType
-replaceAllTypeSynonyms' syns kinds = everywhereOnTypesTopDownM try
+replaceAllTypeSynonyms' syns kinds
+  | M.null syns = Right . markSF
+  | otherwise = walk
   where
-  try :: SourceType -> Either MultipleErrors SourceType
-  try t = fromMaybe t <$> go (fst $ getAnnForType t) 0 [] [] t
+  sf :: TypeFlags -> TypeFlags
+  sf = setFlag tfSynonymsFree
+
+  -- Mark a single node as synonym-free (no recursion)
+  markSF :: SourceType -> SourceType
+  markSF (TUnknown_ f a b) = TUnknown_ (sf f) a b
+  markSF (TypeVar_ f a b) = TypeVar_ (sf f) a b
+  markSF (TypeLevelString_ f a b) = TypeLevelString_ (sf f) a b
+  markSF (TypeLevelInt_ f a b) = TypeLevelInt_ (sf f) a b
+  markSF (TypeWildcard_ f a b) = TypeWildcard_ (sf f) a b
+  markSF (TypeConstructor_ f a b) = TypeConstructor_ (sf f) a b
+  markSF (TypeOp_ f a b) = TypeOp_ (sf f) a b
+  markSF (TypeApp_ f a t1 t2) = TypeApp_ (sf f) a t1 t2
+  markSF (KindApp_ f a t1 t2) = KindApp_ (sf f) a t1 t2
+  markSF (ForAll_ f a v i k t s) = ForAll_ (sf f) a v i k t s
+  markSF (ConstrainedType_ f a c t) = ConstrainedType_ (sf f) a c t
+  markSF (Skolem_ f a n k i s) = Skolem_ (sf f) a n k i s
+  markSF (REmpty_ f a) = REmpty_ (sf f) a
+  markSF (RCons_ f a l t r) = RCons_ (sf f) a l t r
+  markSF (KindedType_ f a t k) = KindedType_ (sf f) a t k
+  markSF (BinaryNoParensType_ f a t1 t2 t3) = BinaryNoParensType_ (sf f) a t1 t2 t3
+  markSF (ParensInType_ f a t) = ParensInType_ (sf f) a t
+
+  -- Main walk: try synonym expansion at potential application sites,
+  -- then recurse into children. Sets tfSynonymsFree on all output nodes.
+  walk :: SourceType -> Either MultipleErrors SourceType
+  walk t | hasFlag tfSynonymsFree (typeFlags t) = Right t
+  walk t@(TypeApp_ _ _ _ _) = trySyn t >>= walkChildren
+  walk t@(KindApp_ _ _ _ _) = trySyn t >>= walkChildren
+  walk t@(TypeConstructor_ _ _ _) = trySyn t >>= \t' -> case t' of
+    TypeConstructor_ _ _ _ -> Right (markSF t')  -- leaf
+    _ -> walkChildren t'  -- synonym expanded to non-leaf
+  walk t = walkChildren t
+
+  -- Try to expand a synonym application at the root.
+  -- Uses the original 'go' logic to peel TypeApp/KindApp and find the constructor.
+  trySyn :: SourceType -> Either MultipleErrors SourceType
+  trySyn t = fromMaybe t <$> go (fst $ getAnnForType t) 0 [] [] t
 
   go :: SourceSpan -> Int -> [SourceType] -> [SourceType] -> SourceType -> Either MultipleErrors (Maybe SourceType)
   go ss c kargs args (TypeConstructor _ ctor)
@@ -43,19 +93,67 @@ replaceAllTypeSynonyms' syns kinds = everywhereOnTypesTopDownM try
     , kindArgs <- lookupKindArgs ctor
     , length kargs == length kindArgs
     = let repl = replaceAllTypeVars (zip (map fst synArgs) args <> zip kindArgs kargs) body
-      in Just <$> try repl
+      in Just <$> trySyn repl
     | Just (synArgs, _) <- M.lookup ctor syns
     , length synArgs > c
     = throwError . errorMessage' ss $ PartiallyAppliedSynonym ctor
   go ss c kargs args (TypeApp _ f arg) = go ss (c + 1) kargs (arg : args) f
   go ss c kargs args (KindApp _ f arg) = go ss c (arg : kargs) args f
   go _ _ _ _ _ = return Nothing
 
+  -- Walk children and reconstruct with recomputed structural flags + tfSynonymsFree.
+  -- Uses raw constructors to set flags in a single allocation.
+  walkChildren :: SourceType -> Either MultipleErrors SourceType
+  walkChildren (TypeApp_ _ ann t1 t2) = do
+    t1' <- walk t1; t2' <- walk t2
+    return $! TypeApp_ (sf (typeFlags t1' `combineFlags` typeFlags t2')) ann t1' t2'
+  walkChildren (KindApp_ _ ann t1 t2) = do
+    t1' <- walk t1; t2' <- walk t2
+    return $! KindApp_ (sf (typeFlags t1' `combineFlags` typeFlags t2')) ann t1' t2'
+  walkChildren (ForAll_ _ ann vis ident mbK ty sco) = do
+    mbK' <- traverse walk mbK; ty' <- walk ty
+    return $! ForAll_ (sf (forAllNodeFlags mbK' ty' sco)) ann vis ident mbK' ty' sco
+  walkChildren (ConstrainedType_ _ ann c ty) = do
+    c' <- overConstraintArgsAll (mapM walk) c; ty' <- walk ty
+    return $! ConstrainedType_ (sf (constraintNodeFlags c' ty')) ann c' ty'
+  walkChildren (Skolem_ _ ann name mbK i sc) = do
+    mbK' <- traverse walk mbK
+    return $! Skolem_ (sf (skolemNodeFlags mbK')) ann name mbK' i sc
+  walkChildren (RCons_ _ ann name ty rest) = do
+    ty' <- walk ty; rest' <- walk rest
+    return $! RCons_ (sf (typeFlags ty' `combineFlags` typeFlags rest')) ann name ty' rest'
+  walkChildren (KindedType_ _ ann ty k) = do
+    ty' <- walk ty; k' <- walk k
+    return $! KindedType_ (sf (typeFlags ty' `combineFlags` typeFlags k')) ann ty' k'
+  walkChildren (BinaryNoParensType_ _ ann t1 t2 t3) = do
+    t1' <- walk t1; t2' <- walk t2; t3' <- walk t3
+    return $! BinaryNoParensType_ (sf (typeFlags t1' `combineFlags` typeFlags t2' `combineFlags` typeFlags t3')) ann t1' t2' t3'
+  walkChildren (ParensInType_ _ ann t) = do
+    t' <- walk t
+    return $! ParensInType_ (sf (typeFlags t')) ann t'
+  walkChildren other = return $! markSF other
+
   lookupKindArgs :: Qualified (ProperName 'TypeName) -> [Text]
   lookupKindArgs ctor = fromMaybe [] $ fmap (fmap (fst . snd) . fst) . completeBinderList . fst =<< M.lookup ctor kinds
 
--- | Replace fully applied type synonyms
+-- | Replace fully applied type synonyms.
+-- Short-circuits if the type is already marked as synonym-free.
 replaceAllTypeSynonyms :: SourceType -> TypeCheckM SourceType
-replaceAllTypeSynonyms d = do
-  env <- getEnv
-  either throwError return $ replaceAllTypeSynonyms' (typeSynonyms env) (types env) d
+replaceAllTypeSynonyms d
+  | hasFlag tfSynonymsFree (typeFlags d) = do
+      env <- getEnv
+      -- Sanity check in debug builds: the flag says this type is synonym-free,
+      -- so scanning should confirm no TypeConstructor in it refers to a synonym.
+      -- 'assert' is compiled away with -O, so this is a no-op in production.
+      return $! assert (not (containsTypeSynonyms (typeSynonyms env) d)) d
+  | otherwise = do
+      env <- getEnv
+      either throwError return $ replaceAllTypeSynonyms' (typeSynonyms env) (types env) d
+
+-- | Scan a type for TypeConstructors that are type synonyms.
+-- Used as a correctness check for the 'tfSynonymsFree' flag.
+containsTypeSynonyms :: SynonymMap -> Type a -> Bool
+containsTypeSynonyms syns = everythingOnTypes (||) isSyn
+  where
+    isSyn (TypeConstructor _ ctor) = M.member ctor syns
+    isSyn _ = False

src/Language/PureScript/TypeChecker/Unify.hs

@@ -16,6 +16,7 @@ module Language.PureScript.TypeChecker.Unify
 
 import Prelude
 
+import Control.Exception (assert)
 import Control.Monad (forM_, void, when)
 import Control.Monad.Error.Class (MonadError(..))
 import Control.Monad.State.Class (MonadState(..), gets, modify, state)
@@ -32,7 +33,7 @@ import Language.PureScript.Errors (ErrorMessageHint(..), SimpleErrorMessage(..),
 import Language.PureScript.TypeChecker.Kinds (elaborateKind, instantiateKind, unifyKinds')
 import Language.PureScript.TypeChecker.Monad (CheckState(..), Substitution(..), UnkLevel(..), Unknown, getLocalContext, guardWith, lookupUnkName, withErrorMessageHint, TypeCheckM)
 import Language.PureScript.TypeChecker.Skolems (newSkolemConstant, skolemize)
-import Language.PureScript.Types (Constraint(..), pattern REmptyKinded, RowListItem(..), SourceType, Type(..), WildcardData(..), alignRowsWith, everythingOnTypes, everywhereOnTypes, everywhereOnTypesM, getAnnForType, mkForAll, rowFromList, srcTUnknown)
+import Language.PureScript.Types (Constraint(..), pattern REmptyKinded, RowListItem(..), SourceType, Type(..), WildcardData(..), alignRowsWith, everythingOnTypes, everywhereOnTypes, everywhereOnTypesM, getAnnForType, hasFlag, mkForAll, rowFromList, srcTUnknown, tfHasWildcards, typeFlags)
 import Data.Set qualified as S
 
 -- | Generate a fresh type variable with an unknown kind. Avoid this if at all possible.
@@ -192,10 +193,16 @@ unifyRows r1 r2 = sequence_ matches *> uncurry unifyTails rest where
     throwError . errorMessage $ TypesDoNotUnify r1 r2
 
 -- |
--- Replace type wildcards with unknowns
+-- Replace type wildcards with unknowns.
+-- Short-circuits if the type has no wildcards.
 --
 replaceTypeWildcards :: SourceType -> TypeCheckM SourceType
-replaceTypeWildcards = everywhereOnTypesM replace
+replaceTypeWildcards ty
+  -- Sanity check in debug builds: the flag says no wildcards, so a scan
+  -- should agree. 'assert' is compiled away with -O.
+  | not (hasFlag tfHasWildcards (typeFlags ty)) =
+      return $! assert (not (containsTypeWildcards ty)) ty
+  | otherwise = everywhereOnTypesM replace ty
   where
   replace (TypeWildcard ann wdata) = do
     t <- freshType
@@ -208,6 +215,14 @@ replaceTypeWildcards = everywhereOnTypesM replace
     return t
   replace other = return other
 
+-- | Scan a type for TypeWildcard nodes.
+-- Used as a correctness check for the 'tfHasWildcards' flag.
+containsTypeWildcards :: Type a -> Bool
+containsTypeWildcards = everythingOnTypes (||) isWild
+  where
+    isWild (TypeWildcard _ _) = True
+    isWild _ = False
+
 -- |
 -- Replace outermost unsolved unification variables with named type variables
 --