Requestify computation of break/continue targets

hamishknight · hamishknight · commit 791e2cd7ff6b · 2022-12-01T19:08:27.000Z
Introduce BreakTargetRequest and
ContinueTargetRequest that perform the ASTScope
lookup for the target of a break or continue.
diff --git a/include/swift/AST/Stmt.h b/include/swift/AST/Stmt.h
@@ -34,6 +34,7 @@ class AnyPattern;
 class ASTContext;
 class ASTWalker;
 class Decl;
+class DeclContext;
 class Expr;
 class FuncDecl;
 class Pattern;
@@ -1336,13 +1337,13 @@ class BreakStmt : public Stmt {
   SourceLoc Loc;
   Identifier TargetName; // Named target statement, if specified in the source.
   SourceLoc TargetLoc;
-  LabeledStmt *Target = nullptr;  // Target stmt, wired up by Sema.
+  DeclContext *DC;
+
 public:
   BreakStmt(SourceLoc Loc, Identifier TargetName, SourceLoc TargetLoc,
-            Optional<bool> implicit = None)
+            DeclContext *DC, Optional<bool> implicit = None)
     : Stmt(StmtKind::Break, getDefaultImplicitFlag(implicit, Loc)), Loc(Loc),
-      TargetName(TargetName), TargetLoc(TargetLoc) {
-  }
+      TargetName(TargetName), TargetLoc(TargetLoc), DC(DC) {}
 
   SourceLoc getLoc() const { return Loc; }
 
@@ -1351,16 +1352,16 @@ class BreakStmt : public Stmt {
   SourceLoc getTargetLoc() const { return TargetLoc; }
   void setTargetLoc(SourceLoc L) { TargetLoc = L; }
 
-  // Manipulate the target loop/switch that is bring broken out of.  This is set
-  // by sema during type checking.
-  void setTarget(LabeledStmt *LS) { Target = LS; }
-  LabeledStmt *getTarget() const { return Target; }
+  /// Retrieve the target statement being jumped out of.
+  LabeledStmt *getTarget() const;
 
   SourceLoc getStartLoc() const { return Loc; }
   SourceLoc getEndLoc() const {
     return (TargetLoc.isValid() ? TargetLoc : Loc);
   }
 
+  DeclContext *getDeclContext() const { return DC; }
+
   static bool classof(const Stmt *S) {
     return S->getKind() == StmtKind::Break;
   }
@@ -1371,24 +1372,21 @@ class ContinueStmt : public Stmt {
   SourceLoc Loc;
   Identifier TargetName; // Named target statement, if specified in the source.
   SourceLoc TargetLoc;
-  LabeledStmt *Target = nullptr;
+  DeclContext *DC;
 
 public:
   ContinueStmt(SourceLoc Loc, Identifier TargetName, SourceLoc TargetLoc,
-               Optional<bool> implicit = None)
+               DeclContext *DC, Optional<bool> implicit = None)
     : Stmt(StmtKind::Continue, getDefaultImplicitFlag(implicit, Loc)), Loc(Loc),
-      TargetName(TargetName), TargetLoc(TargetLoc) {
-  }
+      TargetName(TargetName), TargetLoc(TargetLoc), DC(DC) {}
 
   Identifier getTargetName() const { return TargetName; }
   void setTargetName(Identifier N) { TargetName = N; }
   SourceLoc getTargetLoc() const { return TargetLoc; }
   void setTargetLoc(SourceLoc L) { TargetLoc = L; }
 
-  // Manipulate the target loop that is bring continued.  This is set by sema
-  // during type checking.
-  void setTarget(LabeledStmt *LS) { Target = LS; }
-  LabeledStmt *getTarget() const { return Target; }
+  /// Retrieve the target statement being jumped to.
+  LabeledStmt *getTarget() const;
   
   SourceLoc getLoc() const { return Loc; }
   
@@ -1397,6 +1395,8 @@ class ContinueStmt : public Stmt {
     return (TargetLoc.isValid() ? TargetLoc : Loc);
   }
 
+  DeclContext *getDeclContext() const { return DC; }
+
   static bool classof(const Stmt *S) {
     return S->getKind() == StmtKind::Continue;
   }
@@ -1475,7 +1475,9 @@ class PoundAssertStmt : public Stmt {
   }
 };
 
-inline void simple_display(llvm::raw_ostream &out, Stmt *S) {
+SourceLoc extractNearestSourceLoc(const Stmt *stmt);
+
+inline void simple_display(llvm::raw_ostream &out, const Stmt *S) {
   if (S)
     out << Stmt::getKindName(S->getKind());
   else
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -40,11 +40,14 @@ namespace swift {
 class AbstractStorageDecl;
 class AccessorDecl;
 enum class AccessorKind;
+class BreakStmt;
 class ContextualPattern;
+class ContinueStmt;
 class DefaultArgumentExpr;
 class DefaultArgumentType;
 class ClosureExpr;
 class GenericParamList;
+class LabeledStmt;
 struct MacroDefinition;
 class PrecedenceGroupDecl;
 class PropertyWrapperInitializerInfo;
@@ -3699,6 +3702,40 @@ class SynthesizeLocalVariableForTypeWrapperStorage
   bool isCached() const { return true; }
 };
 
+/// Lookup the target of a break statement.
+class BreakTargetRequest
+    : public SimpleRequest<BreakTargetRequest,
+                           LabeledStmt *(const BreakStmt *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  LabeledStmt *evaluate(Evaluator &evaluator, const BreakStmt *BS) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
+/// Lookup the target of a continue statement.
+class ContinueTargetRequest
+    : public SimpleRequest<ContinueTargetRequest,
+                           LabeledStmt *(const ContinueStmt *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  LabeledStmt *evaluate(Evaluator &evaluator, const ContinueStmt *CS) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 class GetTypeWrapperInitializer
     : public SimpleRequest<GetTypeWrapperInitializer,
                            ConstructorDecl *(NominalTypeDecl *),
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -434,6 +434,12 @@ SWIFT_REQUEST(TypeChecker, SynthesizeTypeWrappedTypeStorageWrapperInitializer,
 SWIFT_REQUEST(TypeChecker, SynthesizeLocalVariableForTypeWrapperStorage,
               VarDecl *(ConstructorDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, BreakTargetRequest,
+              LabeledStmt *(const BreakStmt *),
+              Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ContinueTargetRequest,
+              LabeledStmt *(const ContinueStmt *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapperInitializer,
               ConstructorDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
diff --git a/lib/AST/Stmt.cpp b/lib/AST/Stmt.cpp
@@ -20,6 +20,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/Pattern.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/PointerUnion.h"
 
@@ -621,6 +622,20 @@ SwitchStmt *SwitchStmt::create(LabeledStmtInfo LabelInfo, SourceLoc SwitchLoc,
   return theSwitch;
 }
 
+LabeledStmt *BreakStmt::getTarget() const {
+  auto &eval = getDeclContext()->getASTContext().evaluator;
+  return evaluateOrDefault(eval, BreakTargetRequest{this}, nullptr);
+}
+
+LabeledStmt *ContinueStmt::getTarget() const {
+  auto &eval = getDeclContext()->getASTContext().evaluator;
+  return evaluateOrDefault(eval, ContinueTargetRequest{this}, nullptr);
+}
+
+SourceLoc swift::extractNearestSourceLoc(const Stmt *S) {
+  return S->getStartLoc();
+}
+
 // See swift/Basic/Statistic.h for declaration: this enables tracing Stmts, is
 // defined here to avoid too much layering violation / circular linkage
 // dependency.
diff --git a/lib/Parse/ParseStmt.cpp b/lib/Parse/ParseStmt.cpp
@@ -682,8 +682,8 @@ ParserResult<Stmt> Parser::parseStmtBreak() {
   Status |= parseOptionalControlTransferTarget(*this, Target, TargetLoc,
                                                StmtKind::Break);
 
-  return makeParserResult(Status,
-                          new (Context) BreakStmt(Loc, Target, TargetLoc));
+  auto *BS = new (Context) BreakStmt(Loc, Target, TargetLoc, CurDeclContext);
+  return makeParserResult(Status, BS);
 }
 
 /// parseStmtContinue
@@ -699,8 +699,8 @@ ParserResult<Stmt> Parser::parseStmtContinue() {
   Status |= parseOptionalControlTransferTarget(*this, Target, TargetLoc,
                                                StmtKind::Continue);
 
-  return makeParserResult(Status,
-                          new (Context) ContinueStmt(Loc, Target, TargetLoc));
+  auto *CS = new (Context) ContinueStmt(Loc, Target, TargetLoc, CurDeclContext);
+  return makeParserResult(Status, CS);
 }
 
 
diff --git a/lib/Sema/CSSyntacticElement.cpp b/lib/Sema/CSSyntacticElement.cpp
@@ -1342,27 +1342,14 @@ class SyntacticElementSolutionApplication
   }
 
   ASTNode visitBreakStmt(BreakStmt *breakStmt) {
-    auto *DC = context.getAsDeclContext();
-    if (auto target = findBreakOrContinueStmtTarget(
-            DC->getASTContext(), DC->getParentSourceFile(), breakStmt->getLoc(),
-            breakStmt->getTargetName(), breakStmt->getTargetLoc(),
-            /*isContinue=*/false, context.getAsDeclContext())) {
-      breakStmt->setTarget(target);
-    }
-
+    // Force the target to be computed in case it produces diagnostics.
+    (void)breakStmt->getTarget();
     return breakStmt;
   }
 
   ASTNode visitContinueStmt(ContinueStmt *continueStmt) {
-    auto *DC = context.getAsDeclContext();
-    if (auto target = findBreakOrContinueStmtTarget(
-            DC->getASTContext(), DC->getParentSourceFile(),
-            continueStmt->getLoc(), continueStmt->getTargetName(),
-            continueStmt->getTargetLoc(), /*isContinue=*/true,
-            context.getAsDeclContext())) {
-      continueStmt->setTarget(target);
-    }
-
+    // Force the target to be computed in case it produces diagnostics.
+    (void)continueStmt->getTarget();
     return continueStmt;
   }
 
diff --git a/lib/Sema/TypeCheckStmt.cpp b/lib/Sema/TypeCheckStmt.cpp
@@ -437,6 +437,23 @@ LabeledStmt *swift::findBreakOrContinueStmtTarget(
   return nullptr;
 }
 
+LabeledStmt *
+BreakTargetRequest::evaluate(Evaluator &evaluator, const BreakStmt *BS) const {
+  auto *DC = BS->getDeclContext();
+  return findBreakOrContinueStmtTarget(
+      DC->getASTContext(), DC->getParentSourceFile(), BS->getLoc(),
+      BS->getTargetName(), BS->getTargetLoc(), /*isContinue*/ false, DC);
+}
+
+LabeledStmt *
+ContinueTargetRequest::evaluate(Evaluator &evaluator,
+                                const ContinueStmt *CS) const {
+  auto *DC = CS->getDeclContext();
+  return findBreakOrContinueStmtTarget(
+      DC->getASTContext(), DC->getParentSourceFile(), CS->getLoc(),
+      CS->getTargetName(), CS->getTargetLoc(), /*isContinue*/ true, DC);
+}
+
 static Expr *getDeclRefProvidingExpressionForHasSymbol(Expr *E) {
   // Strip coercions, which are necessary in source to disambiguate overloaded
   // functions or generic functions, e.g.
@@ -1030,24 +1047,14 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   }
 
   Stmt *visitBreakStmt(BreakStmt *S) {
-    if (auto target = findBreakOrContinueStmtTarget(
-            getASTContext(), DC->getParentSourceFile(), S->getLoc(),
-            S->getTargetName(), S->getTargetLoc(), /*isContinue=*/false,
-            DC)) {
-      S->setTarget(target);
-    }
-
+    // Force the target to be computed in case it produces diagnostics.
+    (void)S->getTarget();
     return S;
   }
 
   Stmt *visitContinueStmt(ContinueStmt *S) {
-    if (auto target = findBreakOrContinueStmtTarget(
-            getASTContext(), DC->getParentSourceFile(), S->getLoc(),
-            S->getTargetName(), S->getTargetLoc(), /*isContinue=*/true,
-            DC)) {
-      S->setTarget(target);
-    }
-
+    // Force the target to be computed in case it produces diagnostics.
+    (void)S->getTarget();
     return S;
   }
 
