[Scope map] Add support for querying the immediate and nearest enclosing DeclContexts of a scope.

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -502,6 +502,21 @@ class ASTScope {
   /// Find the innermost enclosing scope that contains this source location.
   const ASTScope *findInnermostEnclosingScope(SourceLoc loc) const;
 
+  /// Retrieve the declaration context directly associated with this scope, or
+  /// NULL if there is no such declaration context.
+  ///
+  /// \seealso getInnermostEnclosingDeclContext().
+  DeclContext *getDeclContext() const;
+
+  /// Retrieve the innermost enclosing declaration context in which this
+  /// scope
+  ///
+  /// This is semantically equivalent to calling \c getDeclContext() on this
+  /// node and each of its parents until we get a non-null result.
+  ///
+  /// \seealso getDeclContext().
+  DeclContext *getInnermostEnclosingDeclContext() const;
+
   /// Expand the entire scope map.
   ///
   /// Normally, the scope map will be expanded only as needed by its queries,

lib/AST/ASTScope.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/Initializer.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
@@ -1539,6 +1540,71 @@ const ASTScope *ASTScope::findInnermostEnclosingScope(SourceLoc loc) const {
   };
 }
 
+DeclContext *ASTScope::getDeclContext() const {
+  switch (getKind()) {
+  case ASTScopeKind::SourceFile:
+    return sourceFile.file;
+
+  case ASTScopeKind::TypeOrExtensionBody:
+    if (auto nominal = dyn_cast<NominalTypeDecl>(iterableDeclContext))
+      return nominal;
+
+    return cast<ExtensionDecl>(iterableDeclContext);
+
+  case ASTScopeKind::AbstractFunctionDecl:
+    return abstractFunction;
+
+  case ASTScopeKind::DefaultArgument:
+    return parameter->getDefaultArgumentInitContext();
+
+  case ASTScopeKind::PatternInitializer:
+    return patternBinding.decl->getPatternList()[patternBinding.entry]
+             .getInitContext();
+
+  case ASTScopeKind::Closure:
+    return closure;
+
+  case ASTScopeKind::Accessors:
+    // FIXME: Somewhat odd modeling because Subscripts don't have their
+    // own nodes. Maybe that should.
+    if (auto subscript = dyn_cast<SubscriptDecl>(abstractStorageDecl))
+      return subscript;
+
+    return nullptr;
+
+  case ASTScopeKind::TopLevelCode:
+    return topLevelCode;
+
+  case ASTScopeKind::GenericParams:
+  case ASTScopeKind::AbstractFunctionParams:
+  case ASTScopeKind::PatternBinding:
+  case ASTScopeKind::AfterPatternBinding:
+  case ASTScopeKind::Preexpanded:
+  case ASTScopeKind::BraceStmt:
+  case ASTScopeKind::IfStmt:
+  case ASTScopeKind::ConditionalClause:
+  case ASTScopeKind::GuardStmt:
+  case ASTScopeKind::RepeatWhileStmt:
+  case ASTScopeKind::ForEachStmt:
+  case ASTScopeKind::ForEachPattern:
+  case ASTScopeKind::DoCatchStmt:
+  case ASTScopeKind::CatchStmt:
+  case ASTScopeKind::SwitchStmt:
+  case ASTScopeKind::CaseStmt:
+  case ASTScopeKind::ForStmt:
+  case ASTScopeKind::ForStmtInitializer:
+  case ASTScopeKind::LocalDeclaration:
+    return nullptr;
+  }
+}
+
+DeclContext *ASTScope::getInnermostEnclosingDeclContext() const {
+  for (const ASTScope *scope = this; ; scope = scope->getParent()) {
+    if (auto dc = scope->getDeclContext()) return dc;
+  }
+  llvm_unreachable("Top-most scope is a declaration context");
+}
+
 void ASTScope::expandAll() const {
   if (!isExpanded())
     expand();

lib/FrontendTool/FrontendTool.cpp

@@ -809,6 +809,11 @@ static bool performCompile(CompilerInstance &Instance,
             << lineColumn.second << "***\n";
           auto locScope = scope.findInnermostEnclosingScope(loc);
           locScope->print(llvm::errs(), 0, false, false);
+
+          // Dump the AST context, too.
+          if (auto dc = locScope->getDeclContext()) {
+            dc->printContext(llvm::errs());
+          }
         }
 
         llvm::errs() << "***Complete scope map***\n";

test/NameBinding/scope_map.swift

@@ -174,6 +174,11 @@ func defaultArguments(i: Int = 1,
   { $0 }(a)
 }
 
+struct PatternInitializers {
+  var (a, b) = (1, 2),
+      (c, d) = (1.5, 2.5)
+}
+
 // RUN: not %target-swift-frontend -dump-scope-maps expanded %s 2> %t.expanded
 // RUN: %FileCheck -check-prefix CHECK-EXPANDED %s < %t.expanded
 
@@ -365,15 +370,40 @@ func defaultArguments(i: Int = 1,
 // CHECK-EXPANDED-NEXT: {{^}}            `-BraceStmt {{.*}} [167:32 - 167:42] expanded
 // CHECK-EXPANDED-NEXT: {{^}}        `-AbstractFunctionParams {{.*}} defaultArguments(i:j:) param 0:1 [167:48 - 175:1] expanded
 
-// RUN: not %target-swift-frontend -dump-scope-maps 70:8,26:20 %s 2> %t.searches
+// RUN: not %target-swift-frontend -dump-scope-maps 70:8,26:20,5:18,166:32,179:18 %s 2> %t.searches
 // RUN: %FileCheck -check-prefix CHECK-SEARCHES %s < %t.searches
 
 // CHECK-SEARCHES-LABEL: ***Scope at 70:8***
 // CHECK-SEARCHES-NEXT: AfterPatternBinding {{.*}} entry 0 [69:13 - 71:3] expanded
+
 // CHECK-SEARCHES-LABEL: ***Scope at 26:20***
 // CHECK-SEARCHES-NEXT: AbstractFunctionParams {{.*}} init(t:u:) param 1:0 [26:17 - 27:3] expanded
+
+// CHECK-SEARCHES-LABEL: ***Scope at 5:18***
+// CHECK-SEARCHES-NEXT: TypeOrExtensionBody {{.*}} 'InnerC0' [5:17 - 5:19] expanded
+// CHECK-SEARCHES-NEXT: Module name=scope_map
+// CHECK-SEARCHES-NEXT:   FileUnit file="{{.*}}scope_map.swift"
+// CHECK-SEARCHES-NEXT:     StructDecl name=S0
+// CHECK-SEARCHES-NEXT:       ClassDecl name=InnerC0
+
+// CHECK-SEARCHES-LABEL: ***Scope at 166:32***
+// CHECK-SEARCHES-NEXT: DefaultArgument {{.*}} [166:32 - 166:32] expanded
+// CHECK-SEARCHES-NEXT: Module name=scope_map
+// CHECK-SEARCHES-NEXT:   FileUnit file="{{.*}}scope_map.swift"
+// CHECK-SEARCHES-NEXT:     AbstractFunctionDecl name=defaultArguments : (Int, Int) -> ()
+// CHECK-SEARCHES-NEXT:       {{.*}} Initializer DefaultArgument index=0
+
+// CHECK-SEARCHES-LABEL: ***Scope at 179:18***
+// CHECK-SEARCHES-NEXT: PatternInitializer {{.*}} entry 1 [179:16 - 179:25] expanded
+// CHECK-SEARCHES-NEXT: {{.*}} Module name=scope_map
+// CHECK-SEARCHES-NEXT:   {{.*}} FileUnit file="{{.*}}scope_map.swift"
+// CHECK-SEARCHES-NEXT:     {{.*}} StructDecl name=PatternInitializers
+// CHECK-SEARCHES-NEXT:       {{.*}} Initializer PatternBinding {{.*}} #1
+
 // CHECK-SEARCHES-LABEL: ***Complete scope map***
 // CHECK-SEARCHES-NEXT: SourceFile {{.*}} '{{.*}}scope_map.swift' [1:1 - {{.*}}:1] expanded
+// CHECK-SEARCHES: TypeOrExtensionBody {{.*}} 'S0' [4:11 - 6:1] expanded
+// CHECK-SEARCHES: -TypeOrExtensionBody {{.*}} 'InnerC0' [5:17 - 5:19] expanded
 // CHECK-SEARCHES-NOT: {{ expanded}}
 // CHECK-SEARCHES: |-TypeOrExtensionBody {{.*}} 'ContainsGenerics0' [25:25 - 31:1] expanded
 // CHECK-SEARCHES-NEXT:   |-AbstractFunctionDecl {{.*}} init(t:u:) [26:3 - 27:3] expanded
@@ -386,4 +416,11 @@ func defaultArguments(i: Int = 1,
 // CHECK-SEARCHES: |-AbstractFunctionDecl {{.*}} functionBodies1(a:b:) [41:1 - 100:1] expanded
 // CHECK-SEARCHES: `-AbstractFunctionParams {{.*}} functionBodies1(a:b:) param 0:0 [41:25 - 100:1] expanded
 // CHECK-SEARCHES: |-AbstractFunctionDecl {{.*}} throwing() [102:1 - 102:26] unexpanded
+// CHECK-SEARCHES: -AbstractFunctionDecl {{.*}} defaultArguments(i:j:) [166:1 - 175:1] expanded
+// CHECK-SEARCHES: DefaultArgument {{.*}} [166:32 - 166:32] expanded
+// CHECK-SEARCHES-NOT: {{ expanded}}
+// CHECK-SEARCHES: -TypeOrExtensionBody {{.*}} 'PatternInitializers' [177:28 - 180:1] expanded
+// CHECK-SEARCHES:    |-PatternBinding {{.*}} entry 0 [178:7 - 178:21] unexpanded
+// CHECK-SEARCHES:    `-PatternBinding {{.*}} entry 1 [179:7 - 179:25] expanded
+// CHECK-SEARCHES:      `-PatternInitializer {{.*}} entry 1 [179:16 - 179:25] expanded
 // CHECK-SEARCHES-NOT: {{ expanded}}