Always create initializer contexts for pattern binding entries in non-local scopes.

We were optimizing away unused pattern binding initializer contexts in
both the parser and in semantic analysis, which led to a
somewhat-unpredictable set of DeclContexts in the AST. Normalize
everything by always creating these contexts.

Author: DougGregor

include/swift/AST/ASTContext.h

@@ -549,13 +549,6 @@ class ASTContext {
     addCleanup([&object]{ object.~T(); });
   }
 
-  /// Create a context for the initializer of a non-local variable,
-  /// like a global or a field.  To reduce memory usage, if the
-  /// context goes unused, it should be returned to the ASTContext
-  /// with destroyPatternBindingContext.
-  PatternBindingInitializer *createPatternBindingContext(DeclContext *parent);
-  void destroyPatternBindingContext(PatternBindingInitializer *DC);
-
   /// Create a context for the initializer of the nth default argument
   /// of the given function.  To reduce memory usage, if the context
   /// goes unused, it should be returned to the ASTContext with

include/swift/AST/Decl.h

@@ -246,8 +246,11 @@ class alignas(1 << DeclAlignInBits) Decl {
 
     /// \brief Whether 'static' or 'class' was used.
     unsigned StaticSpelling : 2;
+
+    /// \brief The number of pattern binding declarations.
+    unsigned NumPatternEntries : 16;
   };
-  enum { NumPatternBindingDeclBits = NumDeclBits + 3 };
+  enum { NumPatternBindingDeclBits = NumDeclBits + 19 };
   static_assert(NumPatternBindingDeclBits <= 32, "fits in an unsigned");
 
   class ValueDeclBitfields {
@@ -1668,9 +1671,14 @@ class PatternBindingEntry {
   // initializer is ASTContext-allocated it is safe.
   llvm::PointerIntPair<Expr *, 2, OptionSet<Flags>> InitCheckedAndRemoved;
 
+  /// The initializer context used for this pattern binding entry.
+  DeclContext *InitContext = nullptr;
+
+  friend class PatternBindingInitializer;
+
 public:
-  PatternBindingEntry(Pattern *P, Expr *E)
-    : ThePattern(P), InitCheckedAndRemoved(E, {}) {}
+  PatternBindingEntry(Pattern *P, Expr *E, DeclContext *InitContext)
+    : ThePattern(P), InitCheckedAndRemoved(E, {}), InitContext(InitContext) {}
 
   Pattern *getPattern() const { return ThePattern; }
   void setPattern(Pattern *P) { ThePattern = P; }
@@ -1690,6 +1698,12 @@ class PatternBindingEntry {
 
   // Return the first variable initialized by this pattern.
   VarDecl *getAnchoringVarDecl() const;
+
+  // Retrieve the declaration context for the intializer.
+  DeclContext *getInitContext() const { return InitContext; }
+
+  /// Override the initializer context.
+  void setInitContext(DeclContext *dc) { InitContext = dc; }
 };
 
 /// \brief This decl contains a pattern and optional initializer for a set
@@ -1711,8 +1725,6 @@ class PatternBindingDecl final : public Decl,
   SourceLoc StaticLoc; ///< Location of the 'static/class' keyword, if present.
   SourceLoc VarLoc;    ///< Location of the 'var' keyword.
 
-  unsigned numPatternEntries;
-
   friend class Decl;
 
   PatternBindingDecl(SourceLoc StaticLoc, StaticSpellingKind StaticSpelling,
@@ -1732,13 +1744,22 @@ class PatternBindingDecl final : public Decl,
                                     Pattern *Pat, Expr *E,
                                     DeclContext *Parent);
 
+  static PatternBindingDecl *createDeserialized(
+                               ASTContext &Ctx, SourceLoc StaticLoc,
+                               StaticSpellingKind StaticSpelling,
+                               SourceLoc VarLoc,
+                               unsigned NumPatternEntries,
+                               DeclContext *Parent);
+
   SourceLoc getStartLoc() const {
     return StaticLoc.isValid() ? StaticLoc : VarLoc;
   }
   SourceLoc getLoc() const { return VarLoc; }
   SourceRange getSourceRange() const;
 
-  unsigned getNumPatternEntries() const { return numPatternEntries; }
+  unsigned getNumPatternEntries() const {
+    return PatternBindingDeclBits.NumPatternEntries;
+  }
 
   ArrayRef<PatternBindingEntry> getPatternList() const {
     return const_cast<PatternBindingDecl*>(this)->getMutablePatternList();
@@ -1760,7 +1781,7 @@ class PatternBindingDecl final : public Decl,
     return getPatternList()[i].getPattern();
   }
 
-  void setPattern(unsigned i, Pattern *Pat);
+  void setPattern(unsigned i, Pattern *Pat, DeclContext *InitContext);
 
   /// Given that this PBD is the parent pattern for the specified VarDecl,
   /// return the entry of the VarDecl in our PatternList.  For example, in:
@@ -1810,7 +1831,7 @@ class PatternBindingDecl final : public Decl,
 private:
   MutableArrayRef<PatternBindingEntry> getMutablePatternList() {
     // Pattern entries are tail allocated.
-    return {getTrailingObjects<PatternBindingEntry>(), numPatternEntries};
+    return {getTrailingObjects<PatternBindingEntry>(), getNumPatternEntries()};
   }
 };
 

include/swift/Serialization/ModuleFormat.h

@@ -53,7 +53,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 265; // Last change: pattern binding initializer index
+const uint16_t VERSION_MINOR = 266; // Last change: pattern binding init contexts
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -927,9 +927,9 @@ namespace decls_block {
     BCFixed<1>,  // implicit flag
     BCFixed<1>,  // static?
     StaticSpellingKindField, // spelling of 'static' or 'class'
-    BCVBR<3>    // numpatterns
-
-    // The pattern trails the record.
+    BCVBR<3>,    // numpatterns
+    BCArray<DeclContextIDField> // init contexts
+    // The patterns and decl-contexts trail the record.
   >;
 
   template <unsigned Code>

lib/AST/ASTContext.cpp

@@ -259,9 +259,6 @@ struct ASTContext::Implementation {
   llvm::DenseMap<Decl *, std::pair<LazyMemberLoader *, uint64_t>>
     ConformanceLoaders;
 
-  /// \brief A cached unused pattern-binding initializer context.
-  PatternBindingInitializer *UnusedPatternBindingContext = nullptr;
-
   /// \brief A cached unused default-argument initializer context.
   DefaultArgumentInitializer *UnusedDefaultArgumentContext = nullptr;
 
@@ -1563,22 +1560,6 @@ void ValueDecl::setLocalDiscriminator(unsigned index) {
   getASTContext().Impl.LocalDiscriminators.insert({this, index});
 }
 
-PatternBindingInitializer *
-ASTContext::createPatternBindingContext(DeclContext *parent) {
-  // Check for an existing context we can re-use.
-  if (auto existing = Impl.UnusedPatternBindingContext) {
-    Impl.UnusedPatternBindingContext = nullptr;
-    existing->reset(parent);
-    return existing;
-  }
-
-  return new (*this) PatternBindingInitializer(parent);
-}
-void ASTContext::destroyPatternBindingContext(PatternBindingInitializer *DC) {
-  // There isn't much value in caching more than one of these.
-  Impl.UnusedPatternBindingContext = DC;
-}
-
 DefaultArgumentInitializer *
 ASTContext::createDefaultArgumentContext(DeclContext *fn, unsigned index) {
   // Check for an existing context we can re-use.

lib/AST/ASTWalker.cpp

@@ -150,7 +150,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     for (auto entry : PBD->getPatternList()) {
       ++idx;
       if (Pattern *Pat = doIt(entry.getPattern()))
-        PBD->setPattern(idx, Pat);
+        PBD->setPattern(idx, Pat, entry.getInitContext());
       else
         return true;
       if (entry.getInit()) {

lib/AST/Decl.cpp

@@ -792,11 +792,11 @@ PatternBindingDecl::PatternBindingDecl(SourceLoc StaticLoc,
                                        unsigned NumPatternEntries,
                                        DeclContext *Parent)
   : Decl(DeclKind::PatternBinding, Parent),
-    StaticLoc(StaticLoc), VarLoc(VarLoc),
-    numPatternEntries(NumPatternEntries) {
+    StaticLoc(StaticLoc), VarLoc(VarLoc) {
   PatternBindingDeclBits.IsStatic = StaticLoc.isValid();
   PatternBindingDeclBits.StaticSpelling =
        static_cast<unsigned>(StaticSpelling);
+  PatternBindingDeclBits.NumPatternEntries = NumPatternEntries;
 }
 
 PatternBindingDecl *
@@ -805,9 +805,18 @@ PatternBindingDecl::create(ASTContext &Ctx, SourceLoc StaticLoc,
                            SourceLoc VarLoc,
                            Pattern *Pat, Expr *E,
                            DeclContext *Parent) {
-  return create(Ctx, StaticLoc, StaticSpelling, VarLoc,
-                PatternBindingEntry(Pat, E),
-                Parent);
+  DeclContext *BindingInitContext = nullptr;
+  if (!Parent->isLocalContext())
+    BindingInitContext = new (Ctx) PatternBindingInitializer(Parent);
+
+  auto Result = create(Ctx, StaticLoc, StaticSpelling, VarLoc,
+                       PatternBindingEntry(Pat, E, BindingInitContext),
+                       Parent);
+
+  if (BindingInitContext)
+    cast<PatternBindingInitializer>(BindingInitContext)->setBinding(Result, 0);
+
+  return Result;
 }
 
 PatternBindingDecl *
@@ -829,7 +838,31 @@ PatternBindingDecl::create(ASTContext &Ctx, SourceLoc StaticLoc,
     ++elt;
     auto &newEntry = entries[elt];
     newEntry = pe; // This should take care of initializer with flags
-    PBD->setPattern(elt, pe.getPattern());
+    DeclContext *initContext = pe.getInitContext();
+    if (!initContext && !Parent->isLocalContext()) {
+      auto pbi = new (Ctx) PatternBindingInitializer(Parent);
+      pbi->setBinding(PBD, elt);
+      initContext = pbi;
+    }
+
+    PBD->setPattern(elt, pe.getPattern(), initContext);
+  }
+  return PBD;
+}
+
+PatternBindingDecl *PatternBindingDecl::createDeserialized(
+                      ASTContext &Ctx, SourceLoc StaticLoc,
+                      StaticSpellingKind StaticSpelling,
+                      SourceLoc VarLoc,
+                      unsigned NumPatternEntries,
+                      DeclContext *Parent) {
+  size_t Size = totalSizeToAlloc<PatternBindingEntry>(NumPatternEntries);
+  void *D = allocateMemoryForDecl<PatternBindingDecl>(Ctx, Size,
+                                                      /*ClangNode*/false);
+  auto PBD = ::new (D) PatternBindingDecl(StaticLoc, StaticSpelling, VarLoc,
+                                          NumPatternEntries, Parent);
+  for (auto &entry : PBD->getMutablePatternList()) {
+    entry = PatternBindingEntry(nullptr, nullptr, nullptr);
   }
   return PBD;
 }
@@ -924,9 +957,11 @@ bool PatternBindingDecl::hasStorage() const {
   return false;
 }
 
-void PatternBindingDecl::setPattern(unsigned i, Pattern *P) {
+void PatternBindingDecl::setPattern(unsigned i, Pattern *P,
+                                    DeclContext *InitContext) {
   auto PatternList = getMutablePatternList();
   PatternList[i].setPattern(P);
+  PatternList[i].setInitContext(InitContext);
 
   // Make sure that any VarDecl's contained within the pattern know about this
   // PatternBindingDecl as their parent.

lib/Parse/ParseDecl.cpp

@@ -3974,10 +3974,6 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
   // so we can build our singular PatternBindingDecl at the end.
   SmallVector<PatternBindingEntry, 4> PBDEntries;
 
-  // The pattern binding initializer contexts, which have to be set up at the
-  // end.
-  SmallVector<PatternBindingInitializer *, 4> PBDInitContexts;
-
   // No matter what error path we take, make sure the
   // PatternBindingDecl/TopLevel code block are added.
   SWIFT_DEFER {
@@ -3992,10 +3988,9 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
                                           VarLoc, PBDEntries, CurDeclContext);
 
     // Wire up any initializer contexts we needed.
-    assert(PBDInitContexts.size() == PBDEntries.size());
-    for (unsigned i : indices(PBDInitContexts)) {
-      auto initContext = PBDInitContexts[i];
-      if (initContext) initContext->setBinding(PBD, i);
+    for (unsigned i : indices(PBDEntries)) {
+      if (auto initContext = PBDEntries[i].getInitContext())
+        cast<PatternBindingInitializer>(initContext)->setBinding(PBD, i);
     }
 
     // If we're setting up a TopLevelCodeDecl, configure it by setting up the
@@ -4041,8 +4036,7 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
 
     // Remember this pattern/init pair for our ultimate PatternBindingDecl. The
     // Initializer will be added later when/if it is parsed.
-    PBDEntries.push_back({pattern, nullptr});
-    PBDInitContexts.push_back(nullptr);
+    PBDEntries.push_back({pattern, nullptr, nullptr});
 
     Expr *PatternInit = nullptr;
 
@@ -4052,7 +4046,6 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
       // into (should we have one).  This happens for properties and global
       // variables in libraries.
       PatternBindingInitializer *initContext = nullptr;
-      bool usedInitContext = false;
 
       // Record the variables that we're trying to initialize.  This allows us
       // to cleanly reject "var x = x" when "x" isn't bound to an enclosing
@@ -4072,7 +4065,7 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
       // for the PBD we'll eventually create.  This allows us to have reasonable
       // DeclContexts for any closures that may live inside of initializers.
       if (!CurDeclContext->isLocalContext() && !topLevelDecl && !initContext)
-        initContext = Context.createPatternBindingContext(CurDeclContext);
+        initContext = new (Context) PatternBindingInitializer(CurDeclContext);
 
       // If we're using a local context (either a TopLevelCodeDecl or a
       // PatternBindingContext) install it now so that CurDeclContext is set
@@ -4109,21 +4102,9 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
       PatternInit = init.getPtrOrNull();
       PBDEntries.back().setInit(PatternInit);
 
-      // If we allocated an initContext and the expression had a closure in it,
-      // we'll need to keep the initContext around.
-      if (initContext)
-        usedInitContext |= initParser->hasClosures();
-
       // If we set up an initialization context for a property or module-level
-      // global, check to see if we needed it and wind it down.
-      if (initContext) {
-        // If we didn't need the context, "destroy" it, which recycles it for
-        // the next user.
-        if (!usedInitContext)
-          Context.destroyPatternBindingContext(initContext);
-        else
-          PBDInitContexts.back() = initContext;
-      }
+      // global, record it.
+      PBDEntries.back().setInitContext(initContext);
 
       // If we are doing second pass of code completion, we don't want to
       // suddenly cut off parsing and throw away the declaration.
@@ -4158,7 +4139,6 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
         // into (should we have one).  This happens for properties and global
         // variables in libraries.
         PatternBindingInitializer *initContext = nullptr;
-        bool usedInitContext = false;
 
         // Record the variables that we're trying to set up.  This allows us
         // to cleanly reject "var x = x" when "x" isn't bound to an enclosing
@@ -4178,7 +4158,7 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
         // This will be recontextualized to a method we synthesize during
         // type checking.
         if (!CurDeclContext->isLocalContext() && !topLevelDecl && !initContext)
-          initContext = Context.createPatternBindingContext(CurDeclContext);
+          initContext = new (Context) PatternBindingInitializer(CurDeclContext);
         Optional<ParseFunctionBody> initParser;
         Optional<ContextChange> topLevelParser;
         if (topLevelDecl)
@@ -4188,10 +4168,7 @@ ParserStatus Parser::parseDeclVar(ParseDeclOptions Flags,
           initParser.emplace(*this, initContext);
 
         auto closure = parseExprClosure();
-        if (initContext) {
-          usedInitContext = true;
-          PBDInitContexts.back() = initContext;
-        }
+        PBDEntries.back().setInitContext(initContext);
 
         if (closure.isParseError())
           return makeParserError();