Add -warn-long-expression-type-checking=<limit> frontend option.

Generates a warning for any expression that takes longer than <limit>
milliseconds to type check. This compliments the existing
-warn-long-function-body=<limit> option.

Author: rudkx

include/swift/AST/DiagnosticsSema.def

@@ -3709,6 +3709,9 @@ WARNING(debug_long_function_body, none,
 WARNING(debug_long_closure_body, none,
         "closure took %0ms to type-check (limit: %1ms)",
         (unsigned, unsigned))
+WARNING(debug_long_expression, none,
+        "expression took %0ms to type-check (limit: %1ms)",
+        (unsigned, unsigned))
 
 //------------------------------------------------------------------------------
 // Pattern match diagnostics

include/swift/Frontend/FrontendOptions.h

@@ -138,6 +138,12 @@ class FrontendOptions {
   /// Intended for debugging purposes only.
   unsigned WarnLongFunctionBodies = 0;
 
+  /// If non-zero, warn when type-checking an expression takes longer
+  /// than this many milliseconds.
+  ///
+  /// Intended for debugging purposes only.
+  unsigned WarnLongExpressionTypeChecking = 0;
+
   enum ActionType {
     NoneAction, ///< No specific action
     Parse, ///< Parse only

include/swift/Option/FrontendOptions.td

@@ -320,6 +320,12 @@ def warn_long_function_bodies : Separate<["-"], "warn-long-function-bodies">,
 def warn_long_function_bodies_EQ : Joined<["-"], "warn-long-function-bodies=">,
   Alias<warn_long_function_bodies>;
 
+def warn_long_expression_type_checking : Separate<["-"], "warn-long-expression-type-checking">,
+  MetaVarName<"<n>">,
+  HelpText<"Warns when type-checking a function takes longer than <n> ms">;
+def warn_long_expression_type_checking_EQ : Joined<["-"], "warn-long-expression-type-checking=">,
+  Alias<warn_long_expression_type_checking>;
+
 def enable_source_import : Flag<["-"], "enable-source-import">,
   HelpText<"Enable importing of Swift source files">;
 

include/swift/Subsystems.h

@@ -191,7 +191,8 @@ namespace swift {
   void performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
                            OptionSet<TypeCheckingFlags> Options,
                            unsigned StartElem = 0,
-                           unsigned WarnLongFunctionBodies = 0);
+                           unsigned WarnLongFunctionBodies = 0,
+                           unsigned WarnLongExpressionTypeChecking = 0);
 
   /// Once type checking is complete, this walks protocol requirements
   /// to resolve default witnesses.

lib/Frontend/CompilerInvocation.cpp

@@ -216,6 +216,16 @@ static bool ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
     }
   }
 
+  if (const Arg *A = Args.getLastArg(OPT_warn_long_expression_type_checking)) {
+    unsigned attempt;
+    if (StringRef(A->getValue()).getAsInteger(10, attempt)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(Args), A->getValue());
+    } else {
+      Opts.WarnLongExpressionTypeChecking = attempt;
+    }
+  }
+
   Opts.PlaygroundTransform |= Args.hasArg(OPT_playground);
   if (Args.hasArg(OPT_disable_playground_transform))
     Opts.PlaygroundTransform = false;

lib/Frontend/Frontend.cpp

@@ -519,7 +519,8 @@ void CompilerInstance::performSema() {
       if (mainIsPrimary) {
         performTypeChecking(MainFile, PersistentState.getTopLevelContext(),
                             TypeCheckOptions, CurTUElem,
-                            options.WarnLongFunctionBodies);
+                            options.WarnLongFunctionBodies,
+                            options.WarnLongExpressionTypeChecking);
       }
       CurTUElem = MainFile.Decls.size();
     } while (!Done);
@@ -546,8 +547,9 @@ void CompilerInstance::performSema() {
     if (auto SF = dyn_cast<SourceFile>(File))
       if (PrimaryBufferID == NO_SUCH_BUFFER || SF == PrimarySourceFile)
         performTypeChecking(*SF, PersistentState.getTopLevelContext(),
-                            TypeCheckOptions, /*curElem*/0,
-                            options.WarnLongFunctionBodies);
+                            TypeCheckOptions, /*curElem*/ 0,
+                            options.WarnLongFunctionBodies,
+                            options.WarnLongExpressionTypeChecking);
 
   // Even if there were no source files, we should still record known
   // protocols.

lib/Sema/TypeCheckConstraints.cpp

@@ -1773,22 +1773,35 @@ namespace {
 
   class ExpressionTimer {
     Expr* E;
+    unsigned WarnLimit;
+    bool ShouldDump;
     ASTContext &Context;
     llvm::TimeRecord StartTime = llvm::TimeRecord::getCurrentTime();
 
   public:
-    ExpressionTimer(Expr *E, ASTContext &Context) : E(E), Context(Context) {}
+    ExpressionTimer(Expr *E, bool shouldDump, unsigned warnLimit,
+                    ASTContext &Context)
+        : E(E), WarnLimit(warnLimit), ShouldDump(shouldDump), Context(Context) {
+    }
 
     ~ExpressionTimer() {
       llvm::TimeRecord endTime = llvm::TimeRecord::getCurrentTime(false);
 
       auto elapsed = endTime.getProcessTime() - StartTime.getProcessTime();
+      unsigned elapsedMS = static_cast<unsigned>(elapsed * 1000);
+
+      if (ShouldDump) {
+        // Round up to the nearest 100th of a millisecond.
+        llvm::errs() << llvm::format("%0.2f", ceil(elapsed * 100000) / 100)
+                     << "ms\t";
+        E->getLoc().print(llvm::errs(), Context.SourceMgr);
+        llvm::errs() << "\n";
+      }
 
-      // Round up to the nearest 100th of a millisecond.
-      llvm::errs() << llvm::format("%0.2f", ceil(elapsed * 100000) / 100)
-                   << "ms\t";
-      E->getLoc().print(llvm::errs(), Context.SourceMgr);
-      llvm::errs() << "\n";
+      if (WarnLimit != 0 && elapsedMS >= WarnLimit && E->getLoc().isValid())
+        Context.Diags.diagnose(E->getLoc(), diag::debug_long_expression,
+                               elapsedMS, WarnLimit)
+          .highlight(E->getSourceRange());
     }
   };
 
@@ -1802,8 +1815,9 @@ bool TypeChecker::typeCheckExpression(Expr *&expr, DeclContext *dc,
                                       ExprTypeCheckListener *listener,
                                       ConstraintSystem *baseCS) {
   Optional<ExpressionTimer> timer;
-  if (DebugTimeExpressions)
-    timer.emplace(expr, Context);
+  if (DebugTimeExpressions || WarnLongExpressionTypeChecking)
+    timer.emplace(expr, DebugTimeExpressions, WarnLongExpressionTypeChecking,
+                  Context);
 
   PrettyStackTraceExpr stackTrace(Context, "type-checking", expr);
 

lib/Sema/TypeChecker.cpp

@@ -621,7 +621,8 @@ void swift::typeCheckExternalDefinitions(SourceFile &SF) {
 void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
                                 OptionSet<TypeCheckingFlags> Options,
                                 unsigned StartElem,
-                                unsigned WarnLongFunctionBodies) {
+                                unsigned WarnLongFunctionBodies,
+                                unsigned WarnLongExpressionTypeChecking) {
   if (SF.ASTStage == SourceFile::TypeChecked)
     return;
 
@@ -646,6 +647,7 @@ void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
 
     if (MyTC) {
       MyTC->setWarnLongFunctionBodies(WarnLongFunctionBodies);
+      MyTC->setWarnLongExpressionTypeChecking(WarnLongExpressionTypeChecking);
       if (Options.contains(TypeCheckingFlags::DebugTimeFunctionBodies))
         MyTC->enableDebugTimeFunctionBodies();
 

lib/Sema/TypeChecker.h

@@ -829,6 +829,12 @@ class TypeChecker final : public LazyResolver {
   /// Intended for debugging purposes only.
   unsigned WarnLongFunctionBodies = 0;
 
+  /// If \p timeInMS is non-zero, warn when type-chcking an expression
+  /// takes longer than this many milliseconds.
+  ///
+  /// Intended for debugging purposes only.
+  unsigned WarnLongExpressionTypeChecking = 0;
+
   /// If true, the time it takes to type-check each function will be dumped
   /// to llvm::errs().
   bool DebugTimeFunctionBodies = false;
@@ -872,6 +878,14 @@ class TypeChecker final : public LazyResolver {
     WarnLongFunctionBodies = timeInMS;
   }
 
+  /// If \p timeInMS is non-zero, warn when type-chcking an expression
+  /// takes longer than this many milliseconds.
+  ///
+  /// Intended for debugging purposes only.
+  void setWarnLongExpressionTypeChecking(unsigned timeInMS) {
+    WarnLongExpressionTypeChecking = timeInMS;
+  }
+
   bool getInImmediateMode() {
     return InImmediateMode;
   }

test/Constraints/warn_long_compile.swift

@@ -0,0 +1,9 @@
+// RUN: %target-typecheck-verify-swift -warn-long-expression-type-checking=1 -warn-long-function-bodies=1 %s
+func foo<T>(_ x: T) -> T { return x }
+func foo(_ x: Int) -> Int { return x }
+
+func test(m: Double) -> Int {
+  // expected-warning@-1 {{global function 'test(m:)' took}}
+  return Int(foo(Float(m) / 20) * 20 - 2) + 10
+  // expected-warning@-1 {{expression took}}
+}