First cut at ASTScope expansion requestification

Author: David Ungar

include/swift/AST/ASTScope.h

@@ -30,6 +30,9 @@
 
 #include "swift/AST/ASTNode.h"
 #include "swift/AST/NameLookup.h" // for DeclVisibilityKind
+// extractNearestSourceLoc, simple_display declarations must come before this
+// Alternatively, declare these in SimpleDisplay.h or elsewhere
+//#include "swift/AST/SimpleRequest.h"
 #include "swift/Basic/Compiler.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/NullablePtr.h"
@@ -88,6 +91,12 @@ struct AnnotatedInsertionPoint {
   ASTScopeImpl *insertionPoint;
   const char *explanation;
 };
+} // namespace ast_scope
+
+SourceLoc extractNearestSourceLoc(
+    std::tuple<ast_scope::ASTScopeImpl *, ast_scope::ScopeCreator *>);
+
+namespace ast_scope {
 
 #pragma mark the root ASTScopeImpl class
 
@@ -333,6 +342,7 @@ class ASTScopeImpl {
 public:
   /// expandScope me, sending deferred nodes to my descendants.
   /// Return the scope into which to place subsequent decls
+  ASTScopeImpl *expandAndBeCurrentDetectingRecursion(ScopeCreator &);
   ASTScopeImpl *expandAndBeCurrent(ScopeCreator &);
 
   unsigned getASTAncestorScopeCount() const { return astAncestorScopeCount; }

include/swift/AST/NameLookupRequests.h

@@ -30,6 +30,10 @@ class GenericContext;
 class GenericParamList;
 class TypeAliasDecl;
 class TypeDecl;
+namespace ast_scope {
+class ASTScopeImpl;
+class ScopeCreator;
+} // namespace ast_scope
 
 /// Display a nominal type or extension thereof.
 void simple_display(
@@ -324,6 +328,28 @@ class LookupPrecedenceGroupRequest
   bool isCached() const { return true; }
 };
 
+/// Expand the given ASTScope. Requestified to detect recursion.
+class ExpandASTScopeRequest
+    : public SimpleRequest<ExpandASTScopeRequest,
+                           ast_scope::ASTScopeImpl *(ast_scope::ASTScopeImpl *,
+                                                     ast_scope::ScopeCreator *),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<ast_scope::ASTScopeImpl *>
+  evaluate(Evaluator &evaluator, ast_scope::ASTScopeImpl *,
+           ast_scope::ScopeCreator *) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -18,6 +18,10 @@
 SWIFT_REQUEST(NameLookup, CustomAttrNominalRequest,
               NominalTypeDecl *(CustomAttr *, DeclContext *), Cached,
               NoLocationInfo)
+SWIFT_REQUEST(NameLookup, ExpandASTScopeRequest,
+              ast_scope::ASTScopeImpl* (ast_scope::ASTScopeImpl*, ast_scope::ScopeCreator*),
+              Uncached,
+              NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ExtendedNominalRequest,
               NominalTypeDecl *(ExtensionDecl *), SeparatelyCached,
               NoLocationInfo)

include/swift/Basic/SimpleDisplay.h

@@ -54,6 +54,13 @@ namespace swift {
 
 #undef HAS_TRIVIAL_DISPLAY
 
+  namespace ast_scope {
+  class ASTScopeImpl;
+  class ScopeCreator;
+  } // namespace ast_scope
+  void simple_display(llvm::raw_ostream &out, const ast_scope::ASTScopeImpl *);
+  void simple_display(llvm::raw_ostream &out, const ast_scope::ScopeCreator *);
+
   template<typename T>
   typename std::enable_if<HasTrivialDisplay<T>::value>::type
   simple_display(llvm::raw_ostream &out, const T &value) {

lib/AST/ASTScopeCreation.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/Initializer.h"
 #include "swift/AST/LazyResolver.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
 #include "swift/AST/SourceFile.h"
@@ -340,7 +341,8 @@ class ScopeCreator final {
       if (auto *ip = child->insertionPointForDeferredExpansion().getPtrOrNull())
         return ip;
     }
-    ASTScopeImpl *insertionPoint = child->expandAndBeCurrent(*this);
+    ASTScopeImpl *insertionPoint =
+        child->expandAndBeCurrentDetectingRecursion(*this);
     ASTScopeAssert(child->verifyThatThisNodeComeAfterItsPriorSibling(),
                    "Ensure search will work");
     return insertionPoint;
@@ -1083,6 +1085,18 @@ void ASTScopeImpl::disownDescendants(ScopeCreator &scopeCreator) {
 
 #pragma mark implementations of expansion
 
+ASTScopeImpl *
+ASTScopeImpl::expandAndBeCurrentDetectingRecursion(ScopeCreator &scopeCreator) {
+  return evaluateOrDefault(scopeCreator.getASTContext().evaluator,
+                           ExpandASTScopeRequest{this, &scopeCreator}, nullptr);
+}
+
+llvm::Expected<ASTScopeImpl *>
+ExpandASTScopeRequest::evaluate(Evaluator &evaluator, ASTScopeImpl *parent,
+                                ScopeCreator *scopeCreator) const {
+  return parent->expandAndBeCurrent(*scopeCreator);
+}
+
 ASTScopeImpl *ASTScopeImpl::expandAndBeCurrent(ScopeCreator &scopeCreator) {
   auto *insertionPoint = expandSpecifically(scopeCreator);
   if (scopeCreator.shouldBeLazy()) {
@@ -1755,7 +1769,7 @@ void ASTScopeImpl::reexpand(ScopeCreator &scopeCreator) {
   disownDescendants(scopeCreator);
   // If the expansion recurses back into the tree for lookup, the ASTAncestor
   // scopes will have already been rescued and won't be found!
-  expandAndBeCurrent(scopeCreator);
+  expandAndBeCurrentDetectingRecursion(scopeCreator);
   replaceASTAncestorScopes(astAncestorScopes);
 }
 
@@ -2073,3 +2087,8 @@ ScopeCreator::findLocalizableDeclContextsInAST() const {
 bool ASTSourceFileScope::crossCheckWithAST() {
   return scopeCreator->containsAllDeclContextsFromAST();
 }
+
+void swift::simple_display(llvm::raw_ostream &out,
+                           const ast_scope::ScopeCreator *scopeCreator) {
+  scopeCreator->print(out);
+}

lib/AST/ASTScopePrinting.cpp

@@ -201,3 +201,8 @@ bool GenericTypeOrExtensionScope::doesDeclHaveABody() const { return false; }
 bool IterableTypeScope::doesDeclHaveABody() const {
   return getBraces().Start != getBraces().End;
 }
+
+void swift::simple_display(llvm::raw_ostream &out,
+                           const ast_scope::ASTScopeImpl *scope) {
+  scope->print(out);
+}

lib/AST/ASTScopeSourceRange.cpp

@@ -756,3 +756,9 @@ SourceLoc getLocAfterExtendedNominal(const ExtensionDecl *const ext) {
   return Lexer::getCharSourceRangeFromSourceRange(SM, etr->getSourceRange())
       .getEnd();
 }
+
+SourceLoc swift::extractNearestSourceLoc(
+    std::tuple<ASTScopeImpl *, ScopeCreator *> scopeAndCreator) {
+  const ASTScopeImpl *scope = std::get<0>(scopeAndCreator);
+  return scope->getSourceRangeOfThisASTNode().Start;
+}