From dc59cd2043f001eec89cc2aa5704b2dc48f7d357 Mon Sep 17 00:00:00 2001
From: Jordan Rose <jordan_rose@apple.com>
Date: Thu, 12 Sep 2019 20:21:22 -0700
Subject: [PATCH] [Sema] Stop visiting existential exprs an extra time checking
 for uses (#27150)

This was causing an exponential amount of time traversing the AST with
deeply chained protocol extension methods, such as in the
TestCodableRouter.swift test in Kitura.

If the OpaqueValueExpr is referenced more than once within the
OpenExistentialExpr it'll still get visited more than once, but that
doesn't seem to happen in practice. If it turns out to be a problem,
we can weaken the assertion I'm adding here.

https://bugs.swift.org/browse/SR-11012
---
 lib/Sema/MiscDiagnostics.cpp                  | 21 +++++++++++++++++--
 .../var_decl_usage_checker.swift.gyb          | 17 +++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 validation-test/compiler_scale/var_decl_usage_checker.swift.gyb

diff --git a/lib/Sema/MiscDiagnostics.cpp b/lib/Sema/MiscDiagnostics.cpp
index 9e3e2f158ebbc..e79c864e408d0 100644
--- a/lib/Sema/MiscDiagnostics.cpp
+++ b/lib/Sema/MiscDiagnostics.cpp
@@ -22,6 +22,7 @@
 #include "swift/AST/Pattern.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/SourceManager.h"
+#include "swift/Basic/Statistic.h"
 #include "swift/Basic/StringExtras.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Parse/Parser.h"
@@ -29,6 +30,8 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/SaveAndRestore.h"
+
+#define DEBUG_TYPE "Sema"
 using namespace swift;
 
 /// Return true if this expression is an implicit promotion from T to T?.
@@ -1992,6 +1995,10 @@ class VarDeclUsageChecker : public ASTWalker {
   /// This is a mapping from VarDecls to the if/while/guard statement that they
   /// occur in, when they are in a pattern in a StmtCondition.
   llvm::SmallDenseMap<VarDecl*, LabeledConditionalStmt*> StmtConditionForVD;
+
+#ifndef NDEBUG
+  llvm::SmallPtrSet<Expr*, 32> AllExprsSeen;
+#endif
   
   bool sawError = false;
   
@@ -2732,6 +2739,10 @@ void VarDeclUsageChecker::markStoredOrInOutExpr(Expr *E, unsigned Flags) {
 
 /// The heavy lifting happens when visiting expressions.
 std::pair<bool, Expr *> VarDeclUsageChecker::walkToExprPre(Expr *E) {
+  STATISTIC(VarDeclUsageCheckerExprVisits,
+            "# of times VarDeclUsageChecker::walkToExprPre is called");
+  ++VarDeclUsageCheckerExprVisits;
+
   // Sema leaves some subexpressions null, which seems really unfortunate.  It
   // should replace them with ErrorExpr.
   if (E == nullptr || !E->getType() || E->getType()->hasError()) {
@@ -2739,6 +2750,8 @@ std::pair<bool, Expr *> VarDeclUsageChecker::walkToExprPre(Expr *E) {
     return { false, E };
   }
 
+  assert(AllExprsSeen.insert(E).second && "duplicate traversal");
+
   // If this is a DeclRefExpr found in a random place, it is a load of the
   // vardecl.
   if (auto *DRE = dyn_cast<DeclRefExpr>(E)) {
@@ -2783,9 +2796,13 @@ std::pair<bool, Expr *> VarDeclUsageChecker::walkToExprPre(Expr *E) {
     return { false, E };
   }
   
-  // If we see an OpenExistentialExpr, remember the mapping for its OpaqueValue.
-  if (auto *oee = dyn_cast<OpenExistentialExpr>(E))
+  // If we see an OpenExistentialExpr, remember the mapping for its OpaqueValue
+  // and only walk the subexpr.
+  if (auto *oee = dyn_cast<OpenExistentialExpr>(E)) {
     OpaqueValueMap[oee->getOpaqueValue()] = oee->getExistentialValue();
+    oee->getSubExpr()->walk(*this);
+    return { false, E };
+  }
 
   // Visit bindings.
   if (auto ove = dyn_cast<OpaqueValueExpr>(E)) {
diff --git a/validation-test/compiler_scale/var_decl_usage_checker.swift.gyb b/validation-test/compiler_scale/var_decl_usage_checker.swift.gyb
new file mode 100644
index 0000000000000..158eb2f333fd0
--- /dev/null
+++ b/validation-test/compiler_scale/var_decl_usage_checker.swift.gyb
@@ -0,0 +1,17 @@
+// RUN: %scale-test --begin 1 --end 5 --step 1 --select VarDeclUsageCheckerExprVisits %s
+// REQUIRES: OS=macosx
+// REQUIRES: asserts
+
+protocol Proto {}
+extension Proto {
+  var selfish: Proto { return self }
+}
+
+struct X: Proto {}
+
+func test(_ x: X) -> Proto {
+  return x
+%for i in range(0, N*5):
+    .selfish
+%end
+}