[clang] [Clang][Sema] Refactor collection of multi-level template argument lists (PR #106585)

[Younan Zhang via cfe-commits]

================
@@ -171,374 +139,379 @@ bool isLambdaEnclosedByTypeAliasDecl(
               .TraverseType(Underlying);
 }
 
-// Add template arguments from a variable template instantiation.
-Response
-HandleVarTemplateSpec(const VarTemplateSpecializationDecl *VarTemplSpec,
-                      MultiLevelTemplateArgumentList &Result,
-                      bool SkipForSpecialization) {
-  // For a class-scope explicit specialization, there are no template arguments
-  // at this level, but there may be enclosing template arguments.
-  if (VarTemplSpec->isClassScopeExplicitSpecialization())
-    return Response::DontClearRelativeToPrimaryNextDecl(VarTemplSpec);
-
-  // We're done when we hit an explicit specialization.
-  if (VarTemplSpec->getSpecializationKind() == TSK_ExplicitSpecialization &&
-      !isa<VarTemplatePartialSpecializationDecl>(VarTemplSpec))
-    return Response::Done();
-
-  // If this variable template specialization was instantiated from a
-  // specialized member that is a variable template, we're done.
-  assert(VarTemplSpec->getSpecializedTemplate() && "No variable template?");
-  llvm::PointerUnion<VarTemplateDecl *, VarTemplatePartialSpecializationDecl *>
-      Specialized = VarTemplSpec->getSpecializedTemplateOrPartial();
-  if (VarTemplatePartialSpecializationDecl *Partial =
-          Specialized.dyn_cast<VarTemplatePartialSpecializationDecl *>()) {
-    if (!SkipForSpecialization)
-      Result.addOuterTemplateArguments(
-          Partial, VarTemplSpec->getTemplateInstantiationArgs().asArray(),
-          /*Final=*/false);
-    if (Partial->isMemberSpecialization())
-      return Response::Done();
-  } else {
-    VarTemplateDecl *Tmpl = Specialized.get<VarTemplateDecl *>();
-    if (!SkipForSpecialization)
-      Result.addOuterTemplateArguments(
-          Tmpl, VarTemplSpec->getTemplateInstantiationArgs().asArray(),
-          /*Final=*/false);
-    if (Tmpl->isMemberSpecialization())
-      return Response::Done();
+struct TemplateInstantiationArgumentCollecter
+    : DeclVisitor<TemplateInstantiationArgumentCollecter, Decl *> {
+  Sema &S;
+  MultiLevelTemplateArgumentList &Result;
+  std::optional<ArrayRef<TemplateArgument>> Innermost;
+  bool RelativeToPrimary;
+  bool ForConstraintInstantiation;
+
+  TemplateInstantiationArgumentCollecter(
+      Sema &S, MultiLevelTemplateArgumentList &Result,
+      std::optional<ArrayRef<TemplateArgument>> Innermost,
+      bool RelativeToPrimary, bool ForConstraintInstantiation)
+      : S(S), Result(Result), Innermost(Innermost),
+        RelativeToPrimary(RelativeToPrimary),
+        ForConstraintInstantiation(ForConstraintInstantiation) {}
+
+  Decl *Done() { return nullptr; }
+
+  Decl *ChangeDecl(const Decl *D) {
+    RelativeToPrimary = false;
+    return const_cast<Decl *>(D);
   }
-  return Response::DontClearRelativeToPrimaryNextDecl(VarTemplSpec);
-}
 
-// If we have a template template parameter with translation unit context,
-// then we're performing substitution into a default template argument of
-// this template template parameter before we've constructed the template
-// that will own this template template parameter. In this case, we
-// use empty template parameter lists for all of the outer templates
-// to avoid performing any substitutions.
-Response
-HandleDefaultTempArgIntoTempTempParam(const TemplateTemplateParmDecl *TTP,
-                                      MultiLevelTemplateArgumentList &Result) {
-  for (unsigned I = 0, N = TTP->getDepth() + 1; I != N; ++I)
-    Result.addOuterTemplateArguments(std::nullopt);
-  return Response::Done();
-}
+  Decl *ChangeDecl(const DeclContext *DC) {
+    return ChangeDecl(Decl::castFromDeclContext(DC));
+  }
 
-Response HandlePartialClassTemplateSpec(
-    const ClassTemplatePartialSpecializationDecl *PartialClassTemplSpec,
-    MultiLevelTemplateArgumentList &Result, bool SkipForSpecialization) {
-  if (!SkipForSpecialization)
-      Result.addOuterRetainedLevels(PartialClassTemplSpec->getTemplateDepth());
-  return Response::Done();
-}
+  Decl *UseNextDecl(const Decl *D) { return ChangeDecl(D->getDeclContext()); }
 
-// Add template arguments from a class template instantiation.
-Response
-HandleClassTemplateSpec(const ClassTemplateSpecializationDecl *ClassTemplSpec,
-                        MultiLevelTemplateArgumentList &Result,
-                        bool SkipForSpecialization) {
-  if (!ClassTemplSpec->isClassScopeExplicitSpecialization()) {
-    // We're done when we hit an explicit specialization.
-    if (ClassTemplSpec->getSpecializationKind() == TSK_ExplicitSpecialization &&
-        !isa<ClassTemplatePartialSpecializationDecl>(ClassTemplSpec))
-      return Response::Done();
+  Decl *DontClearRelativeToPrimaryNextDecl(const Decl *D) {
+    return const_cast<Decl *>(Decl::castFromDeclContext(D->getDeclContext()));
+  }
 
-    if (!SkipForSpecialization)
-      Result.addOuterTemplateArguments(
-          const_cast<ClassTemplateSpecializationDecl *>(ClassTemplSpec),
-          ClassTemplSpec->getTemplateInstantiationArgs().asArray(),
-          /*Final=*/false);
+  void AddInnermostTemplateArguments(const Decl *D) {
+    assert(Innermost);
+    Result.addOuterTemplateArguments(const_cast<Decl *>(D), *Innermost,
+                                     /*Final=*/false);
+    Innermost.reset();
+  }
 
-    // If this class template specialization was instantiated from a
-    // specialized member that is a class template, we're done.
-    assert(ClassTemplSpec->getSpecializedTemplate() && "No class template?");
-    if (ClassTemplSpec->getSpecializedTemplate()->isMemberSpecialization())
-      return Response::Done();
-
-    // If this was instantiated from a partial template specialization, we need
-    // to get the next level of declaration context from the partial
-    // specialization, as the ClassTemplateSpecializationDecl's
-    // DeclContext/LexicalDeclContext will be for the primary template.
-    if (auto *InstFromPartialTempl = ClassTemplSpec->getSpecializedTemplateOrPartial()
-                      .dyn_cast<ClassTemplatePartialSpecializationDecl *>())
-      return Response::ChangeDecl(InstFromPartialTempl->getLexicalDeclContext());
+  void AddOuterTemplateArguments(const Decl *D, ArrayRef<TemplateArgument> Args,
+                                 bool Final) {
+    Result.addOuterTemplateArguments(const_cast<Decl *>(D), Args, Final);
   }
-  return Response::UseNextDecl(ClassTemplSpec);
-}
 
-Response HandleFunction(Sema &SemaRef, const FunctionDecl *Function,
-                        MultiLevelTemplateArgumentList &Result,
-                        const FunctionDecl *Pattern, bool RelativeToPrimary,
-                        bool ForConstraintInstantiation,
-                        bool ForDefaultArgumentSubstitution) {
-  // Add template arguments from a function template specialization.
-  if (!RelativeToPrimary &&
-      Function->getTemplateSpecializationKindForInstantiation() ==
-          TSK_ExplicitSpecialization)
-    return Response::Done();
-
-  if (!RelativeToPrimary &&
-      Function->getTemplateSpecializationKind() == TSK_ExplicitSpecialization) {
-    // This is an implicit instantiation of an explicit specialization. We
-    // don't get any template arguments from this function but might get
-    // some from an enclosing template.
-    return Response::UseNextDecl(Function);
-  } else if (const TemplateArgumentList *TemplateArgs =
-                 Function->getTemplateSpecializationArgs()) {
-    // Add the template arguments for this specialization.
-    Result.addOuterTemplateArguments(const_cast<FunctionDecl *>(Function),
-                                     TemplateArgs->asArray(),
-                                     /*Final=*/false);
+  Decl *VisitTemplateTemplateParmDecl(TemplateTemplateParmDecl *TTPD) {
+    if (Innermost)
+      AddInnermostTemplateArguments(TTPD);
----------------
zyn0217 wrote:

So... it turns out we don't have to move to the parent semantic declaration of TTPD even if Innermost is present?
Note that we still want to traverse its parents in the previous implementation anyway.

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