Merge pull request #62741 from DougGregor/local-named-discriminators

Rework computation of local discriminators for named entities.

Author: DougGregor

include/swift/AST/Decl.h

@@ -2292,10 +2292,14 @@ class GenericParameterReferenceInfo final {
 /// ValueDecl - All named decls that are values in the language.  These can
 /// have a type, etc.
 class ValueDecl : public Decl {
+public:
+  enum : unsigned { InvalidDiscriminator = 0xFFFF };
+
+private:
   DeclName Name;
   SourceLoc NameLoc;
   llvm::PointerIntPair<Type, 3, OptionalEnum<AccessLevel>> TypeAndAccess;
-  unsigned LocalDiscriminator = 0;
+  unsigned LocalDiscriminator = InvalidDiscriminator;
 
   struct {
     /// Whether the "IsObjC" bit has been computed yet.
@@ -2582,6 +2586,12 @@ class ValueDecl : public Decl {
   unsigned getLocalDiscriminator() const;
   void setLocalDiscriminator(unsigned index);
 
+  /// Whether this declaration has a local discriminator.
+  bool hasLocalDiscriminator() const;
+
+  /// Return the "raw" local discriminator, without computing it.
+  unsigned getRawLocalDiscriminator() const { return LocalDiscriminator; }
+
   /// Retrieve the declaration that this declaration overrides, if any.
   ValueDecl *getOverriddenDecl() const;
 

include/swift/AST/TypeCheckRequests.h

@@ -3850,7 +3850,7 @@ class SynthesizeRuntimeMetadataAttrGeneratorBody
 /// Compute the local discriminators for the given declaration context.
 ///
 /// This is a state-changing operation for closures within the context, which
-/// produces the number of assigned discriminators.
+/// produces the discriminator value that any subsequent requests should use.
 class LocalDiscriminatorsRequest
     : public SimpleRequest<LocalDiscriminatorsRequest,
                            unsigned(DeclContext *),

include/swift/Parse/LocalContext.h

@@ -27,7 +27,6 @@
 #include "swift/AST/Stmt.h"
 #include "swift/Basic/OptionSet.h"
 #include "swift/Parse/Lexer.h"
-#include "swift/Parse/LocalContext.h"
 #include "swift/Parse/PersistentParserState.h"
 #include "swift/Parse/Token.h"
 #include "swift/Parse/ParserPosition.h"
@@ -176,8 +175,6 @@ class Parser {
   bool InInactiveClauseEnvironment = false;
   bool InSwiftKeyPath = false;
 
-  LocalContext *CurLocalContext = nullptr;
-
   /// Whether we should delay parsing nominal type, extension, and function
   /// bodies.
   bool isDelayedParsingEnabled() const;
@@ -241,18 +238,15 @@ class Parser {
   protected:
     Parser &P;
     DeclContext *OldContext; // null signals that this has been popped
-    LocalContext *OldLocal;
 
     ContextChange(const ContextChange &) = delete;
     ContextChange &operator=(const ContextChange &) = delete;
 
   public:
-    ContextChange(Parser &P, DeclContext *DC,
-                  LocalContext *newLocal = nullptr)
-      : P(P), OldContext(P.CurDeclContext), OldLocal(P.CurLocalContext) {
+    ContextChange(Parser &P, DeclContext *DC)
+      : P(P), OldContext(P.CurDeclContext) {
       assert(DC && "pushing null context?");
       P.CurDeclContext = DC;
-      P.CurLocalContext = newLocal;
     }
 
     /// Prematurely pop the DeclContext installed by the constructor.
@@ -270,16 +264,15 @@ class Parser {
   private:
     void popImpl() {
       P.CurDeclContext = OldContext;
-      P.CurLocalContext = OldLocal;
     }
   };
 
   /// A RAII object for parsing a new local context.
-  class ParseFunctionBody : public LocalContext {
+  class ParseFunctionBody {
   private:
     ContextChange CC;
   public:
-    ParseFunctionBody(Parser &P, DeclContext *DC) : CC(P, DC, this) {
+    ParseFunctionBody(Parser &P, DeclContext *DC) : CC(P, DC) {
       assert(!isa<TopLevelCodeDecl>(DC) &&
              "top-level code should be parsed using TopLevelCodeContext!");
     }
@@ -990,8 +983,7 @@ class Parser {
   /// 'isLine = true' indicates parsing #line instead of #sourcelocation
   ParserStatus parseLineDirective(bool isLine = false);
 
-  void setLocalDiscriminator(ValueDecl *D);
-  void setLocalDiscriminatorToParamList(ParameterList *PL);
+  void recordLocalType(TypeDecl *TD);
 
   /// Skip an `#if` configuration block containing only attributes.
   ///

include/swift/Parse/Parser.h

@@ -18,7 +18,6 @@
 #define SWIFT_PARSE_PERSISTENTPARSERSTATE_H
 
 #include "swift/Basic/SourceLoc.h"
-#include "swift/Parse/LocalContext.h"
 
 namespace swift {
 
@@ -54,9 +53,6 @@ class IDEInspectionDelayedDeclState {
 class PersistentParserState {
   std::unique_ptr<IDEInspectionDelayedDeclState> IDEInspectionDelayedDeclStat;
 
-  /// The local context for all top-level code.
-  TopLevelContext TopLevelCode;
-
 public:
   PersistentParserState();
   PersistentParserState(ASTContext &ctx) : PersistentParserState() { }
@@ -88,10 +84,6 @@ class PersistentParserState {
     assert(hasIDEInspectionDelayedDeclState());
     return std::move(IDEInspectionDelayedDeclStat);
   }
-
-  TopLevelContext &getTopLevelContext() {
-    return TopLevelCode;
-  }
 };
 
 } // end namespace swift

include/swift/Parse/PersistentParserState.h

@@ -2696,13 +2696,55 @@ bool ValueDecl::isInstanceMember() const {
   llvm_unreachable("bad DeclKind");
 }
 
+bool ValueDecl::hasLocalDiscriminator() const {
+  // Generic parameters and unnamed parameters never have local discriminators.
+  if (isa<GenericTypeParamDecl>(this) ||
+      (isa<ParamDecl>(this) && !hasName()))
+    return false;
+
+  // Opaque types never have local discriminators.
+  if (isa<OpaqueTypeDecl>(this))
+    return false;
+
+  // Accessors never have local discriminators.
+  if (isa<AccessorDecl>(this))
+    return false;
+
+  // Implicit and unnamed declarations never have local discriminators.
+  if (getBaseName().isSpecial())
+    return false;
+
+  // If we are not in a local context, there's nothing to do.
+  if (!getDeclContext()->isLocalContext())
+    return false;
+
+  return true;
+}
+
 unsigned ValueDecl::getLocalDiscriminator() const {
+  // If we have already assigned a local discriminator, we're done.
+  if (LocalDiscriminator != InvalidDiscriminator)
+    return LocalDiscriminator;
+
+  // If this declaration does not have a local discriminator, use 0 as a
+  // stand-in.
+  if (!hasLocalDiscriminator())
+    return 0;
+
+  // Assign local discriminators in this context.
+  evaluateOrDefault(
+      getASTContext().evaluator,
+      LocalDiscriminatorsRequest{getDeclContext()}, InvalidDiscriminator);
+
+  assert(LocalDiscriminator != InvalidDiscriminator);
+
   return LocalDiscriminator;
 }
 
 void ValueDecl::setLocalDiscriminator(unsigned index) {
-  assert(getDeclContext()->isLocalContext());
-  assert(LocalDiscriminator == 0 && "LocalDiscriminator is set multiple times");
+  assert(hasLocalDiscriminator());
+  assert(LocalDiscriminator == InvalidDiscriminator &&
+         "LocalDiscriminator is set multiple times");
   LocalDiscriminator = index;
 }
 

lib/AST/Decl.cpp

@@ -448,6 +448,10 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
   StringScratchSpace stringStorage;
   ContainerTracker Containers;
 
+  // Contains a mapping for captures of the form [x], from the declared "x"
+  // to the captured "x" in the enclosing scope.
+  llvm::DenseMap<VarDecl *, VarDecl *> sameNamedCaptures;
+
   bool getNameAndUSR(ValueDecl *D, ExtensionDecl *ExtD,
                      StringRef &name, StringRef &USR) {
     auto &result = nameAndUSRCache[ExtD ? (Decl*)ExtD : D];
@@ -701,6 +705,26 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
   bool walkToExprPre(Expr *E) override {
     if (Cancelled)
       return false;
+
+    // Record any captures of the form [x], herso we can treat
+    if (auto captureList = dyn_cast<CaptureListExpr>(E)) {
+      for (const auto &capture : captureList->getCaptureList()) {
+        auto declaredVar = capture.getVar();
+        if (capture.PBD->getEqualLoc(0).isValid())
+          continue;
+
+        VarDecl *capturedVar = nullptr;
+        if (auto init = capture.PBD->getInit(0)) {
+          if (auto declRef = dyn_cast<DeclRefExpr>(init))
+            capturedVar = dyn_cast_or_null<VarDecl>(declRef->getDecl());
+        }
+
+        if (capturedVar) {
+          sameNamedCaptures[declaredVar] = capturedVar;
+        }
+      }
+    }
+
     ExprStack.push_back(E);
     Containers.activateContainersFor(E);
     handleMemberwiseInitRefs(E);
@@ -1611,6 +1635,15 @@ bool IndexSwiftASTWalker::initIndexSymbol(ValueDecl *D, SourceLoc Loc,
   if (auto *VD = dyn_cast<VarDecl>(D)) {
     // Always base the symbol information on the canonical VarDecl
     D = VD->getCanonicalVarDecl();
+
+    // Dig back to the original captured variable.
+    while (true) {
+      auto captured = sameNamedCaptures.find(cast<VarDecl>(D));
+      if (captured == sameNamedCaptures.end())
+        break;
+
+      D = captured->second;
+    }
   }
 
   Info.decl = D;