Merge pull request swiftlang#33254 from DougGregor/stmt-checker-astscope

[Statement checker] Leverage ASTScope to eliminate recursive walk

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -195,7 +195,7 @@ class ASTScopeImpl {
 
   /// Get ride of descendants and remove them from scopedNodes so the scopes
   /// can be recreated. Needed because typechecking inserts a return statment
-  /// into intiailizers.
+  /// into initializers.
   void disownDescendants(ScopeCreator &);
 
 public: // for addReusedBodyScopes
@@ -415,10 +415,17 @@ class ASTScopeImpl {
   unqualifiedLookup(SourceFile *, DeclNameRef, SourceLoc,
                     const DeclContext *startingContext, DeclConsumer);
 
+  /// Entry point into ASTScopeImpl-land for labeled statement lookups.
+  static llvm::SmallVector<LabeledStmt *, 4>
+  lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc);
+
   static Optional<bool>
   computeIsCascadingUse(ArrayRef<const ASTScopeImpl *> history,
                         Optional<bool> initialIsCascadingUse);
 
+  static std::pair<CaseStmt *, CaseStmt *>
+  lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
+
 #pragma mark - - lookup- starting point
 private:
   static const ASTScopeImpl *findStartingScopeForLookup(SourceFile *,
@@ -538,6 +545,12 @@ class ASTScopeImpl {
 
   NullablePtr<const ASTScopeImpl>
   ancestorWithDeclSatisfying(function_ref<bool(const Decl *)> predicate) const;
+
+  /// Whether this scope terminates lookup of labeled statements in the
+  /// children below it, because one cannot perform a "break" or a "continue"
+  /// in a child that goes outside of this scope.
+  virtual bool isLabeledStmtLookupTerminator() const;
+
 }; // end of ASTScopeImpl
 
 #pragma mark - specific scope classes
@@ -1357,6 +1370,9 @@ class ConditionalClauseScope final : public ASTScopeImpl {
 private:
   ArrayRef<StmtConditionElement> getCond() const;
   const StmtConditionElement &getStmtConditionElement() const;
+
+protected:
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 /// If, while, & guard statements all start with a conditional clause, then some
@@ -1380,6 +1396,7 @@ class ConditionalClausePatternUseScope final : public ASTScopeImpl {
   bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
                              DeclConsumer) const override;
   void printSpecifics(llvm::raw_ostream &out) const override;
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 
@@ -1713,17 +1730,16 @@ class AbstractStmtScope : public ASTScopeImpl {
   virtual Stmt *getStmt() const = 0;
   NullablePtr<Stmt> getStmtIfAny() const override { return getStmt(); }
   NullablePtr<const void> getReferrent() const override;
+
+protected:
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 class LabeledConditionalStmtScope : public AbstractStmtScope {
 public:
   Stmt *getStmt() const override;
   virtual LabeledConditionalStmt *getLabeledConditionalStmt() const = 0;
 
-  /// If a condition is present, create the martuska.
-  /// Return the lookupParent for the use scope.
-  ASTScopeImpl *createCondScopes();
-
 protected:
   /// Return the lookupParent required to search these.
   ASTScopeImpl *createNestedConditionalClauseScopes(ScopeCreator &,
@@ -1807,6 +1823,7 @@ class LookupParentDiversionScope final : public ASTScopeImpl {
   NullablePtr<const ASTScopeImpl> getLookupParent() const override {
     return lookupParent;
   }
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 class RepeatWhileScope final : public AbstractStmtScope {
@@ -1826,6 +1843,23 @@ class RepeatWhileScope final : public AbstractStmtScope {
   Stmt *getStmt() const override { return stmt; }
 };
 
+class DoStmtScope final : public AbstractStmtScope {
+public:
+  DoStmt *const stmt;
+  DoStmtScope(DoStmt *e) : stmt(e) {}
+  virtual ~DoStmtScope() {}
+
+protected:
+  ASTScopeImpl *expandSpecifically(ScopeCreator &scopeCreator) override;
+
+private:
+  void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
+
+public:
+  std::string getClassName() const override;
+  Stmt *getStmt() const override { return stmt; }
+};
+
 class DoCatchStmtScope final : public AbstractStmtScope {
 public:
   DoCatchStmt *const stmt;
@@ -1897,6 +1931,7 @@ class ForEachPatternScope final : public ASTScopeImpl {
 protected:
   bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
                              DeclConsumer) const override;
+  bool isLabeledStmtLookupTerminator() const override;
 };
 
 class CaseStmtScope final : public AbstractStmtScope {

include/swift/AST/AnyFunctionRef.h

@@ -53,6 +53,20 @@ class AnyFunctionRef {
     }
   }
 
+  /// Construct an AnyFunctionRef from a decl context that might be
+  /// some sort of function.
+  static Optional<AnyFunctionRef> fromDeclContext(DeclContext *dc) {
+    if (auto fn = dyn_cast<AbstractFunctionDecl>(dc)) {
+      return AnyFunctionRef(fn);
+    }
+
+    if (auto ace = dyn_cast<AbstractClosureExpr>(dc)) {
+      return AnyFunctionRef(ace);
+    }
+
+    return None;
+  }
+
   CaptureInfo getCaptureInfo() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
       return AFD->getCaptureInfo();

include/swift/AST/NameLookup.h

@@ -706,6 +706,31 @@ class ASTScope {
   computeIsCascadingUse(ArrayRef<const ast_scope::ASTScopeImpl *> history,
                         Optional<bool> initialIsCascadingUse);
 
+  /// Entry point to record the visible statement labels from the given
+  /// point.
+  ///
+  /// This lookup only considers labels that are visible within the current
+  /// function, so it will not return any labels from lexical scopes that
+  /// are not reachable via labeled control flow.
+  ///
+  /// \returns the set of labeled statements visible from the given source
+  /// location, with the innermost labeled statement first and proceeding
+  /// to the outermost labeled statement.
+  static llvm::SmallVector<LabeledStmt *, 4>
+  lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc);
+
+  /// Look for the directly enclosing case statement and the next case
+  /// statement, which together act as the source and destination for a
+  /// 'fallthrough' statement within a switch case.
+  ///
+  /// \returns a pair (fallthrough source, fallthrough dest). If the location
+  /// is not within the body of a case statement at all, the fallthrough
+  /// source will be \c nullptr. If there is a fallthrough source that case is
+  /// the last one, the fallthrough destination will be \c nullptr. A
+  /// well-formed 'fallthrough' statement has both a source and destination.
+  static std::pair<CaseStmt *, CaseStmt *>
+  lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
+
   SWIFT_DEBUG_DUMP;
   void print(llvm::raw_ostream &) const;
   void dumpOneScopeMapLocation(std::pair<unsigned, unsigned>);

lib/AST/ASTScope.cpp

@@ -54,6 +54,16 @@ Optional<bool> ASTScope::computeIsCascadingUse(
   return ASTScopeImpl::computeIsCascadingUse(history, initialIsCascadingUse);
 }
 
+llvm::SmallVector<LabeledStmt *, 4> ASTScope::lookupLabeledStmts(
+    SourceFile *sourceFile, SourceLoc loc) {
+  return ASTScopeImpl::lookupLabeledStmts(sourceFile, loc);
+}
+
+std::pair<CaseStmt *, CaseStmt *> ASTScope::lookupFallthroughSourceAndDest(
+    SourceFile *sourceFile, SourceLoc loc) {
+  return ASTScopeImpl::lookupFallthroughSourceAndDest(sourceFile, loc);
+}
+
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(push)
 #pragma warning(disable : 4996)
@@ -237,6 +247,7 @@ DEFINE_GET_CLASS_NAME(WhileStmtScope)
 DEFINE_GET_CLASS_NAME(GuardStmtScope)
 DEFINE_GET_CLASS_NAME(LookupParentDiversionScope)
 DEFINE_GET_CLASS_NAME(RepeatWhileScope)
+DEFINE_GET_CLASS_NAME(DoStmtScope)
 DEFINE_GET_CLASS_NAME(DoCatchStmtScope)
 DEFINE_GET_CLASS_NAME(SwitchStmtScope)
 DEFINE_GET_CLASS_NAME(ForEachStmtScope)

lib/AST/ASTScopeCreation.cpp

@@ -866,6 +866,7 @@ class NodeAdder
   VISIT_AND_CREATE(IfStmt, IfStmtScope)
   VISIT_AND_CREATE(WhileStmt, WhileStmtScope)
   VISIT_AND_CREATE(RepeatWhileStmt, RepeatWhileScope)
+  VISIT_AND_CREATE(DoStmt, DoStmtScope)
   VISIT_AND_CREATE(DoCatchStmt, DoCatchStmtScope)
   VISIT_AND_CREATE(SwitchStmt, SwitchStmtScope)
   VISIT_AND_CREATE(ForEachStmt, ForEachStmtScope)
@@ -908,11 +909,6 @@ class NodeAdder
                                            ScopeCreator &scopeCreator) {
     return scopeCreator.ifUniqueConstructExpandAndInsert<GuardStmtScope>(p, e);
   }
-  NullablePtr<ASTScopeImpl> visitDoStmt(DoStmt *ds, ASTScopeImpl *p,
-                                        ScopeCreator &scopeCreator) {
-    scopeCreator.addToScopeTreeAndReturnInsertionPoint(ds->getBody(), p);
-    return p; // Don't put subsequent decls inside the "do"
-  }
   NullablePtr<ASTScopeImpl> visitTopLevelCodeDecl(TopLevelCodeDecl *d,
                                                   ASTScopeImpl *p,
                                                   ScopeCreator &scopeCreator) {
@@ -1204,6 +1200,7 @@ NO_NEW_INSERTION_POINT(CaptureListScope)
 NO_NEW_INSERTION_POINT(CaseStmtScope)
 NO_NEW_INSERTION_POINT(ClosureBodyScope)
 NO_NEW_INSERTION_POINT(DefaultArgumentInitializerScope)
+NO_NEW_INSERTION_POINT(DoStmtScope)
 NO_NEW_INSERTION_POINT(DoCatchStmtScope)
 NO_NEW_INSERTION_POINT(ForEachPatternScope)
 NO_NEW_INSERTION_POINT(ForEachStmtScope)
@@ -1470,6 +1467,11 @@ void RepeatWhileScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
   scopeCreator.addToScopeTree(stmt->getCond(), this);
 }
 
+void DoStmtScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
+    ScopeCreator &scopeCreator) {
+  scopeCreator.addToScopeTree(stmt->getBody(), this);
+}
+
 void DoCatchStmtScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
     ScopeCreator &scopeCreator) {
   scopeCreator.addToScopeTree(stmt->getBody(), this);

lib/AST/ASTScopeLookup.cpp

@@ -852,3 +852,108 @@ bool isLocWithinAnInactiveClause(const SourceLoc loc, SourceFile *SF) {
   SF->walk(tester);
   return tester.wasFoundWithinInactiveClause;
 }
+
+#pragma mark isLabeledStmtLookupTerminator implementations
+bool ASTScopeImpl::isLabeledStmtLookupTerminator() const {
+  return true;
+}
+
+bool LookupParentDiversionScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool ConditionalClauseScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool ConditionalClausePatternUseScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool AbstractStmtScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+bool ForEachPatternScope::isLabeledStmtLookupTerminator() const {
+  return false;
+}
+
+llvm::SmallVector<LabeledStmt *, 4>
+ASTScopeImpl::lookupLabeledStmts(SourceFile *sourceFile, SourceLoc loc) {
+  // Find the innermost scope from which to start our search.
+  auto *const fileScope = sourceFile->getScope().impl;
+  const auto *innermost = fileScope->findInnermostEnclosingScope(loc, nullptr);
+  ASTScopeAssert(innermost->getWasExpanded(),
+                 "If looking in a scope, it must have been expanded.");
+
+  llvm::SmallVector<LabeledStmt *, 4> labeledStmts;
+  for (auto scope = innermost; scope && !scope->isLabeledStmtLookupTerminator();
+       scope = scope->getParent().getPtrOrNull()) {
+    // If we have a labeled statement, record it.
+    auto stmt = scope->getStmtIfAny();
+    if (!stmt) continue;
+
+    auto labeledStmt = dyn_cast<LabeledStmt>(stmt.get());
+    if (!labeledStmt) continue;
+
+    // Skip guard statements; they aren't actually targets for break or
+    // continue.
+    if (isa<GuardStmt>(labeledStmt)) continue;
+
+    labeledStmts.push_back(labeledStmt);
+  }
+
+  return labeledStmts;
+}
+
+std::pair<CaseStmt *, CaseStmt *> ASTScopeImpl::lookupFallthroughSourceAndDest(
+    SourceFile *sourceFile, SourceLoc loc) {
+  // Find the innermost scope from which to start our search.
+  auto *const fileScope = sourceFile->getScope().impl;
+  const auto *innermost = fileScope->findInnermostEnclosingScope(loc, nullptr);
+  ASTScopeAssert(innermost->getWasExpanded(),
+                 "If looking in a scope, it must have been expanded.");
+
+  // Look for the enclosing case statement and its 'switch' statement.
+  CaseStmt *fallthroughSource = nullptr;
+  SwitchStmt *switchStmt = nullptr;
+  for (auto scope = innermost; scope && !scope->isLabeledStmtLookupTerminator();
+       scope = scope->getParent().getPtrOrNull()) {
+    // If we have a case statement, record it.
+    auto stmt = scope->getStmtIfAny();
+    if (!stmt) continue;
+
+    // If we've found the first case statement of a switch, record it as the
+    // fallthrough source. do-catch statements don't support fallthrough.
+    if (auto caseStmt = dyn_cast<CaseStmt>(stmt.get())) {
+      if (!fallthroughSource &&
+          caseStmt->getParentKind() == CaseParentKind::Switch)
+        fallthroughSource = caseStmt;
+
+      continue;
+    }
+
+    // If we've found the first switch statement, record it and we're done.
+    switchStmt = dyn_cast<SwitchStmt>(stmt.get());
+    if (switchStmt)
+      break;
+  }
+
+  // If we don't have both a fallthrough source and a switch statement
+  // enclosing it, the 'fallthrough' statement is ill-formed.
+  if (!fallthroughSource || !switchStmt)
+    return { nullptr, nullptr };
+
+  // Find this case in the list of cases for the switch. If we don't find it
+  // here, it means that the case isn't directly nested inside the switch, so
+  // the case and fallthrough are both ill-formed.
+  auto caseIter = llvm::find(switchStmt->getCases(), fallthroughSource);
+  if (caseIter == switchStmt->getCases().end())
+    return { nullptr, nullptr };
+
+  // Move along to the next case. This is the fallthrough destination.
+  ++caseIter;
+  auto fallthroughDest = caseIter == switchStmt->getCases().end() ? nullptr
+      : *caseIter;
+  return { fallthroughSource, fallthroughDest };
+}