[ConstraintSystem] ExpressionTimer: Accept locators as timer anchors

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -167,7 +167,9 @@ struct RestrictionOrFix {
 
 
 class ExpressionTimer {
-  Expr* E;
+  using AnchorType = llvm::PointerUnion<Expr *, ConstraintLocator *>;
+
+  AnchorType Anchor;
   ASTContext &Context;
   llvm::TimeRecord StartTime;
 
@@ -181,8 +183,8 @@ class ExpressionTimer {
 public:
   /// This constructor sets a default threshold defined for all expressions
   /// via compiler flag `solver-expression-time-threshold`.
-  ExpressionTimer(Expr *E, ConstraintSystem &CS);
-  ExpressionTimer(Expr *E, ConstraintSystem &CS, unsigned thresholdInMillis);
+  ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS);
+  ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS, unsigned thresholdInMillis);
 
   ~ExpressionTimer();
 

lib/Sema/ConstraintSystem.cpp

@@ -43,14 +43,14 @@ using namespace inference;
 
 #define DEBUG_TYPE "ConstraintSystem"
 
-ExpressionTimer::ExpressionTimer(Expr *E, ConstraintSystem &CS)
+ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS)
     : ExpressionTimer(
-          E, CS,
+          Anchor, CS,
           CS.getASTContext().TypeCheckerOpts.ExpressionTimeoutThreshold) {}
 
-ExpressionTimer::ExpressionTimer(Expr *E, ConstraintSystem &CS,
+ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS,
                                  unsigned thresholdInMillis)
-    : E(E), Context(CS.getASTContext()),
+    : Anchor(Anchor), Context(CS.getASTContext()),
       StartTime(llvm::TimeRecord::getCurrentTime()),
       ThresholdInMillis(thresholdInMillis),
       PrintDebugTiming(CS.getASTContext().TypeCheckerOpts.DebugTimeExpressions),
@@ -64,21 +64,39 @@ ExpressionTimer::~ExpressionTimer() {
     // 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);
+    if (auto *E = Anchor.dyn_cast<Expr *>()) {
+      E->getLoc().print(llvm::errs(), Context.SourceMgr);
+    } else {
+      auto *locator = Anchor.get<ConstraintLocator *>();
+      locator->dump(&Context.SourceMgr, llvm::errs());
+    }
     llvm::errs() << "\n";
   }
 
   if (!PrintWarning)
     return;
 
+  ASTNode anchor;
+  if (auto *locator = Anchor.dyn_cast<ConstraintLocator *>()) {
+    anchor = simplifyLocatorToAnchor(locator);
+    // If locator couldn't be simplified down to a single AST
+    // element, let's warn about its root.
+    if (!anchor)
+      anchor = locator->getAnchor();
+  } else {
+    anchor = Anchor.get<Expr *>();
+  }
+
   const auto WarnLimit = getWarnLimit();
-  if (WarnLimit != 0 && elapsedMS >= WarnLimit && E->getLoc().isValid())
-    Context.Diags.diagnose(E->getLoc(), diag::debug_long_expression,
-                           elapsedMS, WarnLimit)
-      .highlight(E->getSourceRange());
+  if (WarnLimit != 0 && elapsedMS >= WarnLimit &&
+      anchor.getStartLoc().isValid()) {
+    Context.Diags
+        .diagnose(anchor.getStartLoc(), diag::debug_long_expression, elapsedMS,
+                  WarnLimit)
+        .highlight(anchor.getSourceRange());
+  }
 }
 
-
 ConstraintSystem::ConstraintSystem(DeclContext *dc,
                                    ConstraintSystemOptions options)
   : Context(dc->getASTContext()), DC(dc), Options(options),