[ASTScope] Model "do" statements in the scope map

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -1826,6 +1826,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;

lib/AST/ASTScope.cpp

@@ -237,6 +237,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);

test/NameLookup/scope_map-astscopelookup.swift

@@ -194,6 +194,17 @@ class LazyProperties {
   lazy var prop: Int = self.value
 }
 
+func HasLabeledDo() {
+  label:
+  do {
+    for x in 0..<100 {
+      if x % 3 == 0 {
+        break label
+      }
+    }
+  }
+}
+
 // RUN: not %target-swift-frontend -dump-scope-maps expanded %s 2> %t.expanded
 // RUN: %FileCheck -check-prefix CHECK-EXPANDED %s < %t.expanded