[AST] Tighten up type of InitContext

Switch from promising a DeclContext to a
PatternBindingInitializer.

This has a couple of benefits:
- It eliminates a few places where we were force
`cast`'ing to PatternBindingInitializer.
- It improves the clarity of what's being stored,
it's not whatever the parent context of the
initializer is, it's specifically the
PatternBindingInitializer context if it exists.

Author: hamishknight

include/swift/AST/Decl.h

@@ -31,6 +31,7 @@
 #include "swift/AST/GenericParamKey.h"
 #include "swift/AST/IfConfigClause.h"
 #include "swift/AST/Import.h"
+#include "swift/AST/Initializer.h"
 #include "swift/AST/LayoutConstraint.h"
 #include "swift/AST/LifetimeAnnotation.h"
 #include "swift/AST/ReferenceCounting.h"
@@ -1962,7 +1963,7 @@ class PatternBindingEntry {
     IsFromDebugger = 1 << 2,
   };
   /// The initializer context used for this pattern binding entry.
-  llvm::PointerIntPair<DeclContext *, 3, OptionSet<PatternFlags>>
+  llvm::PointerIntPair<PatternBindingInitializer *, 3, OptionSet<PatternFlags>>
       InitContextAndFlags;
 
   /// Values captured by this initializer.
@@ -2003,7 +2004,7 @@ class PatternBindingEntry {
 public:
   /// \p E is the initializer as parsed.
   PatternBindingEntry(Pattern *P, SourceLoc EqualLoc, Expr *E,
-                      DeclContext *InitContext)
+                      PatternBindingInitializer *InitContext)
       : PatternAndFlags(P, {}),
         InitExpr({E, {E, InitializerStatus::NotChecked}, EqualLoc}),
         InitContextAndFlags({InitContext, llvm::None}) {}
@@ -2107,12 +2108,14 @@ class PatternBindingEntry {
   VarDecl *getAnchoringVarDecl() const;
 
   // Retrieve the declaration context for the initializer.
-  DeclContext *getInitContext() const {
+  PatternBindingInitializer *getInitContext() const {
     return InitContextAndFlags.getPointer();
   }
 
   /// Override the initializer context.
-  void setInitContext(DeclContext *dc) { InitContextAndFlags.setPointer(dc); }
+  void setInitContext(PatternBindingInitializer *init) {
+    InitContextAndFlags.setPointer(init);
+  }
 
   SourceLoc getStartLoc() const;
 
@@ -2310,15 +2313,16 @@ class PatternBindingDecl final : public Decl,
   Pattern *getPattern(unsigned i) const {
     return getPatternList()[i].getPattern();
   }
-  
-  void setPattern(unsigned i, Pattern *Pat, DeclContext *InitContext,
+
+  void setPattern(unsigned i, Pattern *Pat,
+                  PatternBindingInitializer *InitContext,
                   bool isFullyValidated = false);
 
   bool isFullyValidated(unsigned i) const {
     return getPatternList()[i].isFullyValidated();
   }
 
-  DeclContext *getInitContext(unsigned i) const {
+  PatternBindingInitializer *getInitContext(unsigned i) const {
     return getPatternList()[i].getInitContext();
   }
 

include/swift/AST/Initializer.h

@@ -21,9 +21,9 @@
 #define SWIFT_INITIALIZER_H
 
 #include "swift/AST/DeclContext.h"
-#include "swift/AST/Decl.h"
 
 namespace swift {
+class ParamDecl;
 class PatternBindingDecl;
 
 enum class InitializerKind : uint8_t {
@@ -90,12 +90,8 @@ class PatternBindingInitializer : public Initializer {
   static PatternBindingInitializer *createDeserialized(PatternBindingDecl *PBD,
                                                        unsigned index);
 
-  void setBinding(PatternBindingDecl *binding, unsigned bindingIndex) {
-    setParent(binding->getDeclContext());
-    Binding = binding;
-    SpareBits = bindingIndex;
-  }
-  
+  void setBinding(PatternBindingDecl *binding, unsigned bindingIndex);
+
   PatternBindingDecl *getBinding() const { return Binding; }
 
   unsigned getBindingIndex() const { return SpareBits; }

lib/AST/ASTScopeLookup.cpp

@@ -437,9 +437,7 @@ bool BraceStmtScope::lookupLocalsOrMembers(DeclConsumer consumer) const {
 bool PatternEntryInitializerScope::lookupLocalsOrMembers(
     DeclConsumer consumer) const {
   // 'self' is available within the pattern initializer of a 'lazy' variable.
-  auto *initContext = dyn_cast_or_null<PatternBindingInitializer>(
-      decl->getInitContext(0));
-  if (initContext) {
+  if (auto *initContext = decl->getInitContext(0)) {
     if (auto *selfParam = initContext->getImplicitSelfDecl()) {
       return consumer.consume({selfParam});
     }

lib/AST/Decl.cpp

@@ -1924,8 +1924,7 @@ PatternBindingDecl::create(ASTContext &Ctx, SourceLoc StaticLoc,
                           PBD->getTrailingObjects<PatternBindingEntry>());
 
   for (auto idx : range(PBD->getNumPatternEntries())) {
-    auto *initContext =
-        cast_or_null<PatternBindingInitializer>(PBD->getInitContext(idx));
+    auto *initContext = PBD->getInitContext(idx);
 
     // FIXME: We ought to reconsider this since it won't recontextualize any
     // closures/decls present in the initialization expr. This currently should
@@ -2004,6 +2003,13 @@ ParamDecl *PatternBindingInitializer::getImplicitSelfDecl() const {
   return SelfParam;
 }
 
+void PatternBindingInitializer::setBinding(PatternBindingDecl *binding,
+                                           unsigned bindingIndex) {
+  setParent(binding->getDeclContext());
+  Binding = binding;
+  SpareBits = bindingIndex;
+}
+
 VarDecl *PatternBindingInitializer::getInitializedLazyVar() const {
   if (auto binding = getBinding()) {
     if (auto var = binding->getSingleVar()) {
@@ -2232,7 +2238,7 @@ PatternBindingDecl::getCheckedPatternBindingEntry(unsigned i) const {
 }
 
 void PatternBindingDecl::setPattern(unsigned i, Pattern *P,
-                                    DeclContext *InitContext,
+                                    PatternBindingInitializer *InitContext,
                                     bool isFullyValidated) {
   auto PatternList = getMutablePatternList();
   PatternList[i].setPattern(P);
@@ -2251,7 +2257,6 @@ void PatternBindingDecl::setPattern(unsigned i, Pattern *P,
   }
 }
 
-
 VarDecl *PatternBindingDecl::getSingleVar() const {
   if (getNumPatternEntries() == 1)
     return getPatternList()[0].getPattern()->getSingleVar();

lib/Sema/CSApply.cpp

@@ -9476,7 +9476,8 @@ ExprWalker::rewriteTarget(SyntacticElementTarget target) {
       // Record that the pattern has been fully validated,
       // this is important for subsequent call to typeCheckDecl
       // because otherwise it would try to re-typecheck pattern.
-      patternBinding->setPattern(index, pattern, resultTarget->getDeclContext(),
+      patternBinding->setPattern(index, pattern,
+                                 patternBinding->getInitContext(index),
                                  /*isFullyValidated=*/true);
 
       if (patternBinding->isExplicitlyInitialized(index) ||

lib/Sema/TypeCheckConstraints.cpp

@@ -871,14 +871,8 @@ bool TypeChecker::typeCheckPatternBinding(PatternBindingDecl *PBD,
   Expr *init = PBD->getInit(patternNumber);
 
   // Enter an initializer context if necessary.
-  PatternBindingInitializer *initContext = nullptr;
-  DeclContext *DC = PBD->getDeclContext();
-  if (!DC->isLocalContext()) {
-    initContext = cast_or_null<PatternBindingInitializer>(
-        PBD->getInitContext(patternNumber));
-    if (initContext)
-      DC = initContext;
-  }
+  PatternBindingInitializer *initContext = PBD->getInitContext(patternNumber);
+  DeclContext *DC = initContext ? initContext : PBD->getDeclContext();
 
   // If we weren't given a pattern type, compute one now.
   if (!patternType) {

lib/Sema/TypeCheckStorage.cpp

@@ -608,9 +608,7 @@ static void checkAndContextualizePatternBindingInit(PatternBindingDecl *binding,
       return;
   }
 
-  auto *initContext =
-      cast_or_null<PatternBindingInitializer>(binding->getInitContext(i));
-  if (initContext) {
+  if (auto *initContext = binding->getInitContext(i)) {
     auto *init = binding->getInit(i);
     TypeChecker::contextualizeInitializer(initContext, init);
     (void)binding->getInitializerIsolation(i);
@@ -2235,10 +2233,7 @@ static AccessorDecl *createGetterPrototype(AbstractStorageDecl *storage,
       auto *varDecl = cast<VarDecl>(storage);
       auto *bindingDecl = varDecl->getParentPatternBinding();
       const auto i = bindingDecl->getPatternEntryIndexForVarDecl(varDecl);
-      auto *bindingInit = cast<PatternBindingInitializer>(
-        bindingDecl->getInitContext(i));
-
-      selfDecl = bindingInit->getImplicitSelfDecl();
+      selfDecl = bindingDecl->getInitContext(i)->getImplicitSelfDecl();
     }
   }
 

lib/Serialization/Deserialization.cpp

@@ -4363,7 +4363,8 @@ class DeclDeserializer {
 
     for (unsigned i = 0; i != patterns.size(); ++i) {
       DeclContext *initContext = MF.getDeclContext(patterns[i].second);
-      binding->setPattern(i, patterns[i].first, initContext);
+      binding->setPattern(i, patterns[i].first,
+                          cast_or_null<PatternBindingInitializer>(initContext));
     }
 
     return binding;