[AST] Eagerly wire up VarDecl parents when creating CaseStmt

Rather than waiting until type-checking, we can set the parents
immediately when we create the CaseStmt. This requires fixing up
NamingPatternRequest to look at the recursive parent statement as
now the VarDecl may have a variable parent.

Author: hamishknight

lib/AST/Stmt.cpp

@@ -773,33 +773,63 @@ CaseStmt::CaseStmt(CaseParentKind parentKind, SourceLoc itemIntroducerLoc,
   MutableArrayRef<CaseLabelItem> items{getTrailingObjects<CaseLabelItem>(),
                                        static_cast<size_t>(Bits.CaseStmt.NumPatterns)};
 
-  // At the beginning mark all of our var decls as being owned by this
-  // statement. In the typechecker we wireup the case stmt var decl list since
-  // we know everything is lined up/typechecked then.
   for (unsigned i : range(Bits.CaseStmt.NumPatterns)) {
     new (&items[i]) CaseLabelItem(caseLabelItems[i]);
-    items[i].getPattern()->markOwnedByStatement(this);
-  }
-  for (auto *vd : getCaseBodyVariables()) {
-    vd->setParentPatternStmt(this);
+    // Mark the CaseStmt as the parent for any canonical VarDecls in the
+    // pattern.
+    items[i].getPattern()->forEachVariable([&](VarDecl *VD) {
+      if (!VD->getParentVarDecl())
+        VD->setParentPatternStmt(this);
+    });
   }
 }
 
 namespace {
-static ArrayRef<VarDecl *>
-getCaseVarDecls(ASTContext &ctx, ArrayRef<CaseLabelItem> labelItems) {
-  // Grab the first case label item pattern and use it to initialize the case
-  // body var decls.
-  SmallVector<VarDecl *, 4> tmp;
-  labelItems.front().getPattern()->collectVariables(tmp);
-  return ctx.AllocateTransform<VarDecl *>(
-      llvm::ArrayRef(tmp), [&](VarDecl *vOld) -> VarDecl * {
-        auto *vNew = new (ctx) VarDecl(
-            /*IsStatic*/ false, vOld->getIntroducer(), vOld->getNameLoc(),
-            vOld->getName(), vOld->getDeclContext());
-        vNew->setImplicit();
-        return vNew;
-      });
+/// Produces an array of internal case body variables, and binds all of the
+/// pattern variables that occur within the case to their "parent" pattern
+/// variables, forming chains of variables with the same name.
+///
+/// Given a case such as:
+/// \code
+/// case .a(let x), .b(let x), .c(let x):
+/// \endcode
+///
+/// Each case item contains a (different) pattern variable named `x`. This
+/// function will set the "parent" variable of the second and third `x`
+/// variables to the `x` variable immediately to its left. A fourth `x` will be
+/// the body case variable, whose parent will be set to the `x` within the final
+/// case item.
+static ArrayRef<VarDecl *> getCaseVarDecls(ASTContext &ctx,
+                                           ArrayRef<CaseLabelItem> labelItems) {
+  SmallVector<VarDecl *, 4> caseVars;
+  llvm::SmallDenseMap<Identifier, VarDecl *, 4> allVars;
+
+  auto foundVar = [&](VarDecl *VD) {
+    if (!VD->hasName())
+      return;
+
+    auto &entry = allVars[VD->getName()];
+    if (entry) {
+      VD->setParentVarDecl(entry);
+    } else {
+      auto *caseVar = new (ctx) VarDecl(
+          /*IsStatic*/ false, VD->getIntroducer(), VD->getNameLoc(),
+          VD->getName(), VD->getDeclContext());
+      caseVar->setImplicit();
+      caseVars.push_back(caseVar);
+    }
+    entry = VD;
+  };
+
+  for (auto &caseItem : labelItems)
+    caseItem.getPattern()->forEachVariable(foundVar);
+
+  // Now that we've collected the case variables, ensure they're parented to
+  // the last pattern variables we saw.
+  for (auto caseVar : caseVars)
+    foundVar(caseVar);
+
+  return ctx.AllocateCopy(caseVars);
 }
 
 struct FallthroughFinder : ASTWalker {

lib/Sema/TypeCheckDecl.cpp

@@ -2715,7 +2715,7 @@ NamingPatternRequest::evaluate(Evaluator &evaluator, VarDecl *VD) const {
   }
 
   if (!namingPattern) {
-    if (auto parentStmt = VD->getParentPatternStmt()) {
+    if (auto parentStmt = VD->getRecursiveParentPatternStmt()) {
       // Try type checking parent control statement.
       if (auto condStmt = dyn_cast<LabeledConditionalStmt>(parentStmt)) {
         // The VarDecl is defined inside a condition of a `if` or `while` stmt.