[clang] Reapply "[Clang] Improve diagnostics for expansion length mismatch" (PR #121044)

[via cfe-commits]

================
@@ -749,132 +759,124 @@ ExprResult Sema::CheckPackExpansion(Expr *Pattern, SourceLocation EllipsisLoc,
     PackExpansionExpr(Context.DependentTy, Pattern, EllipsisLoc, NumExpansions);
 }
 
+static bool IsUnexpandedPackExpansion(const TemplateArgument &TA) {
+  if (!TA.isPackExpansion())
+    return false;
+
+  if (TA.getKind() == TemplateArgument::Type)
+    return !TA.getAsType()->getAs<PackExpansionType>()->getNumExpansions();
+
+  if (TA.getKind() == TemplateArgument::Expression)
+    return !cast<PackExpansionExpr>(TA.getAsExpr())->getNumExpansions();
+
+  return !TA.getNumTemplateExpansions();
+}
+
 bool Sema::CheckParameterPacksForExpansion(
     SourceLocation EllipsisLoc, SourceRange PatternRange,
     ArrayRef<UnexpandedParameterPack> Unexpanded,
     const MultiLevelTemplateArgumentList &TemplateArgs, bool &ShouldExpand,
     bool &RetainExpansion, std::optional<unsigned> &NumExpansions) {
   ShouldExpand = true;
   RetainExpansion = false;
-  std::pair<IdentifierInfo *, SourceLocation> FirstPack;
-  bool HaveFirstPack = false;
-  std::optional<unsigned> NumPartialExpansions;
-  SourceLocation PartiallySubstitutedPackLoc;
-  typedef LocalInstantiationScope::DeclArgumentPack DeclArgumentPack;
+  std::pair<const IdentifierInfo *, SourceLocation> FirstPack;
+  std::optional<std::pair<unsigned, SourceLocation>> PartialExpansion;
+  std::optional<unsigned> CurNumExpansions, CurMaximumOfLeastExpansions;
 
-  for (UnexpandedParameterPack ParmPack : Unexpanded) {
+  for (auto [P, Loc] : Unexpanded) {
     // Compute the depth and index for this parameter pack.
-    unsigned Depth = 0, Index = 0;
-    IdentifierInfo *Name;
-    bool IsVarDeclPack = false;
-    FunctionParmPackExpr *BindingPack = nullptr;
-
-    if (const TemplateTypeParmType *TTP =
-            ParmPack.first.dyn_cast<const TemplateTypeParmType *>()) {
-      Depth = TTP->getDepth();
-      Index = TTP->getIndex();
-      Name = TTP->getIdentifier();
-    } else {
-      NamedDecl *ND = cast<NamedDecl *>(ParmPack.first);
-      if (isa<VarDecl>(ND))
-        IsVarDeclPack = true;
-      else if (isa<BindingDecl>(ND)) {
+    std::optional<std::pair<unsigned, unsigned>> Pos;
+    unsigned NewPackSize, PendingPackExpansionSize = 0;
+    const auto *ND = dyn_cast_if_present<const NamedDecl *>(P);
+    if (ND) {
+      if (isa<VarDecl>(ND)) {
+        auto *DAP = dyn_cast<LocalInstantiationScope::DeclArgumentPack *>(
+            *CurrentInstantiationScope->findInstantiationOf(ND));
+        if (!DAP) {
+          // We can't expand this function parameter pack, so we can't expand
+          // the pack expansion.
+          ShouldExpand = false;
+          continue;
+        }
+        NewPackSize = DAP->size();
+      } else if (isa<BindingDecl>(ND)) {
         // Find the instantiated BindingDecl and check it for a resolved pack.
-        llvm::PointerUnion<Decl *, DeclArgumentPack *> *Instantiation =
-            CurrentInstantiationScope->findInstantiationOf(ND);
+        llvm::PointerUnion<Decl *, LocalInstantiationScope::DeclArgumentPack *>
+            *Instantiation = CurrentInstantiationScope->findInstantiationOf(ND);
         Decl *B = cast<Decl *>(*Instantiation);
         Expr *BindingExpr = cast<BindingDecl>(B)->getBinding();
-        BindingPack = cast_if_present<FunctionParmPackExpr>(BindingExpr);
+        auto *BindingPack = cast_if_present<FunctionParmPackExpr>(BindingExpr);
         if (!BindingPack) {
           ShouldExpand = false;
           continue;
         }
+        NewPackSize = BindingPack->getNumExpansions();
       } else
-        std::tie(Depth, Index) = getDepthAndIndex(ND);
-
-      Name = ND->getIdentifier();
+        Pos = getDepthAndIndex(ND);
+    } else if (const auto *TTP = dyn_cast<const TemplateTypeParmType *>(P)) {
+      Pos = {TTP->getDepth(), TTP->getIndex()};
+      ND = TTP->getDecl();
+      // FIXME: We either should have some fallback for canonical TTP, or
+      //        never have canonical TTP here.
+    } else if (const auto *STP =
+                   dyn_cast<const SubstTemplateTypeParmPackType *>(P)) {
+      NewPackSize = STP->getNumArgs();
+      PendingPackExpansionSize = llvm::count_if(
+          STP->getArgumentPack().getPackAsArray(), IsUnexpandedPackExpansion);
+      ND = STP->getReplacedParameter();
+    } else {
+      const auto *SEP = cast<const SubstNonTypeTemplateParmPackExpr *>(P);
+      NewPackSize = SEP->getArgumentPack().pack_size();
+      PendingPackExpansionSize = llvm::count_if(
+          SEP->getArgumentPack().getPackAsArray(), IsUnexpandedPackExpansion);
+      ND = SEP->getParameterPack();
----------------
cor3ntin wrote:

That makes a lot of sense. Can you add a comment along these lines?

https://github.com/llvm/llvm-project/pull/121044