[Concepts] Correctly form initial parameter mapping for parameter packs, support substitution into SubstNonTypeTemplateParmExpr

We previously would not correctly for the initial parameter mapping for variadic template parameters in Concepts.
Testing this lead to the discovery that with the normalization process we would need to substitute into already-substituted-into
template arguments, which means we need to add NonTypeTemplateParmExpr support to TemplateInstantiator.
We do that by substituting into the replacement and the type separately, and then re-checking the expression against the NTTP
with the new type, in order to form any new required implicit casts (for cases where the type of the NTTP was dependent).

(cherry picked from commit ba1f3db4b0729ad932aa4f091e9578132d98a0c8)

7 files changed, 136 insertions, 26 deletions

diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index 72d49c119e8..697d1911be8 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -6997,7 +6997,7 @@ public:
   /// Get a template argument mapping the given template parameter to itself,
   /// e.g. for X in \c template<int X>, this would return an expression template
   /// argument referencing X.
-  TemplateArgumentLoc getIdentityTemplateArgumentLoc(Decl *Param,
+  TemplateArgumentLoc getIdentityTemplateArgumentLoc(NamedDecl *Param,
                                                      SourceLocation Location);
 
   void translateTemplateArguments(const ASTTemplateArgsPtr &In,
diff --git a/clang/include/clang/Sema/SemaConcept.h b/clang/include/clang/Sema/SemaConcept.h
index 7fc42a4816e..c5f9fc45612 100644
--- a/clang/include/clang/Sema/SemaConcept.h
+++ b/clang/include/clang/Sema/SemaConcept.h
@@ -43,11 +43,15 @@ struct AtomicConstraint {
     if (ParameterMapping->size() != Other.ParameterMapping->size())
       return false;
 
-    for (unsigned I = 0, S = ParameterMapping->size(); I < S; ++I)
-      if (!C.getCanonicalTemplateArgument((*ParameterMapping)[I].getArgument())
-               .structurallyEquals(C.getCanonicalTemplateArgument(
-                  (*Other.ParameterMapping)[I].getArgument())))
+    for (unsigned I = 0, S = ParameterMapping->size(); I < S; ++I) {
+      llvm::FoldingSetNodeID IDA, IDB;
+      C.getCanonicalTemplateArgument((*ParameterMapping)[I].getArgument())
+          .Profile(IDA, C);
+      C.getCanonicalTemplateArgument((*Other.ParameterMapping)[I].getArgument())
+          .Profile(IDB, C);
+      if (IDA != IDB)
         return false;
+    }
     return true;
   }
 
diff --git a/clang/lib/Sema/SemaConcept.cpp b/clang/lib/Sema/SemaConcept.cpp
index 8fdc6023040..39169664dad 100755
--- a/clang/lib/Sema/SemaConcept.cpp
+++ b/clang/lib/Sema/SemaConcept.cpp
@@ -676,6 +676,10 @@ static bool substituteParameterMappings(Sema &S, NormalizedConstraint &N,
                   ArgsAsWritten->arguments().back().getSourceRange().getEnd()));
   if (S.SubstTemplateArguments(*Atomic.ParameterMapping, MLTAL, SubstArgs))
     return true;
+  Atomic.ParameterMapping.emplace(
+        MutableArrayRef<TemplateArgumentLoc>(
+            new (S.Context) TemplateArgumentLoc[SubstArgs.size()],
+            SubstArgs.size()));
   std::copy(SubstArgs.arguments().begin(), SubstArgs.arguments().end(),
             N.getAtomicConstraint()->ParameterMapping->begin());
   return false;
diff --git a/clang/lib/Sema/SemaTemplateDeduction.cpp b/clang/lib/Sema/SemaTemplateDeduction.cpp
index 394c81c8279..1a71f270679 100644
--- a/clang/lib/Sema/SemaTemplateDeduction.cpp
+++ b/clang/lib/Sema/SemaTemplateDeduction.cpp
@@ -2488,7 +2488,7 @@ Sema::getTrivialTemplateArgumentLoc(const TemplateArgument &Arg,
     case TemplateArgument::Template:
     case TemplateArgument::TemplateExpansion: {
       NestedNameSpecifierLocBuilder Builder;
-      TemplateName Template = Arg.getAsTemplate();
+      TemplateName Template = Arg.getAsTemplateOrTemplatePattern();
       if (DependentTemplateName *DTN = Template.getAsDependentTemplateName())
         Builder.MakeTrivial(Context, DTN->getQualifier(), Loc);
       else if (QualifiedTemplateName *QTN =
@@ -2514,27 +2514,10 @@ Sema::getTrivialTemplateArgumentLoc(const TemplateArgument &Arg,
 }
 
 TemplateArgumentLoc
-Sema::getIdentityTemplateArgumentLoc(Decl *TemplateParm,
+Sema::getIdentityTemplateArgumentLoc(NamedDecl *TemplateParm,
                                      SourceLocation Location) {
-  if (auto *TTP = dyn_cast<TemplateTypeParmDecl>(TemplateParm))
-    return getTrivialTemplateArgumentLoc(
-        TemplateArgument(
-            Context.getTemplateTypeParmType(TTP->getDepth(), TTP->getIndex(),
-                                            TTP->isParameterPack(), TTP)),
-        QualType(), Location.isValid() ? Location : TTP->getLocation());
-  else if (auto *TTP = dyn_cast<TemplateTemplateParmDecl>(TemplateParm))
-    return getTrivialTemplateArgumentLoc(TemplateArgument(TemplateName(TTP)),
-                                         QualType(),
-                                         Location.isValid() ? Location :
-                                         TTP->getLocation());
-  auto *NTTP = cast<NonTypeTemplateParmDecl>(TemplateParm);
-  CXXScopeSpec SS;
-  DeclarationNameInfo Info(NTTP->getDeclName(),
-                           Location.isValid() ? Location : NTTP->getLocation());
-  Expr *E = BuildDeclarationNameExpr(SS, Info, NTTP).get();
-  return getTrivialTemplateArgumentLoc(TemplateArgument(E), NTTP->getType(),
-                                       Location.isValid() ? Location :
-                                       NTTP->getLocation());
+  return getTrivialTemplateArgumentLoc(
+      Context.getInjectedTemplateArg(TemplateParm), QualType(), Location);
 }
 
 /// Convert the given deduced template argument and add it to the set of
diff --git a/clang/lib/Sema/SemaTemplateInstantiate.cpp b/clang/lib/Sema/SemaTemplateInstantiate.cpp
index 26dc5d92f23..dff336f2ff2 100644
--- a/clang/lib/Sema/SemaTemplateInstantiate.cpp
+++ b/clang/lib/Sema/SemaTemplateInstantiate.cpp
@@ -1057,6 +1057,8 @@ namespace {
                                             NonTypeTemplateParmDecl *D);
     ExprResult TransformSubstNonTypeTemplateParmPackExpr(
                                            SubstNonTypeTemplateParmPackExpr *E);
+    ExprResult TransformSubstNonTypeTemplateParmExpr(
+                                           SubstNonTypeTemplateParmExpr *E);
 
     /// Rebuild a DeclRefExpr for a VarDecl reference.
     ExprResult RebuildVarDeclRefExpr(VarDecl *PD, SourceLocation Loc);
@@ -1535,6 +1537,44 @@ TemplateInstantiator::TransformSubstNonTypeTemplateParmPackExpr(
                                          Arg);
 }
 
+ExprResult
+TemplateInstantiator::TransformSubstNonTypeTemplateParmExpr(
+                                          SubstNonTypeTemplateParmExpr *E) {
+  ExprResult SubstReplacement = TransformExpr(E->getReplacement());
+  if (SubstReplacement.isInvalid())
+    return true;
+  QualType SubstType = TransformType(E->getType());
+  if (SubstType.isNull())
+    return true;
+  // The type may have been previously dependent and not now, which means we
+  // might have to implicit cast the argument to the new type, for example:
+  // template<auto T, decltype(T) U>
+  // concept C = sizeof(U) == 4;
+  // void foo() requires C<2, 'a'> { }
+  // When normalizing foo(), we first form the normalized constraints of C:
+  // AtomicExpr(sizeof(U) == 4,
+  //            U=SubstNonTypeTemplateParmExpr(Param=U,
+  //                                           Expr=DeclRef(U),
+  //                                           Type=decltype(T)))
+  // Then we substitute T = 2, U = 'a' into the parameter mapping, and need to
+  // produce:
+  // AtomicExpr(sizeof(U) == 4,
+  //            U=SubstNonTypeTemplateParmExpr(Param=U,
+  //                                           Expr=ImpCast(
+  //                                               decltype(2),
+  //                                               SubstNTTPE(Param=U, Expr='a',
+  //                                                          Type=char)),
+  //                                           Type=decltype(2)))
+  // The call to CheckTemplateArgument here produces the ImpCast.
+  TemplateArgument Converted;
+  if (SemaRef.CheckTemplateArgument(E->getParameter(), SubstType,
+                                    SubstReplacement.get(),
+                                    Converted).isInvalid())
+    return true;
+  return transformNonTypeTemplateParmRef(E->getParameter(),
+                                         E->getExprLoc(), Converted);
+}
+
 ExprResult TemplateInstantiator::RebuildVarDeclRefExpr(VarDecl *PD,
                                                        SourceLocation Loc) {
   DeclarationNameInfo NameInfo(PD->getDeclName(), Loc);
diff --git a/clang/test/CXX/temp/temp.constr/temp.constr.normal/p1.cpp b/clang/test/CXX/temp/temp.constr/temp.constr.normal/p1.cpp
index 95fe35b4591..153d4a56bea 100644
--- a/clang/test/CXX/temp/temp.constr/temp.constr.normal/p1.cpp
+++ b/clang/test/CXX/temp/temp.constr/temp.constr.normal/p1.cpp
@@ -16,3 +16,54 @@ template<typename T> requires Bar2<T> struct S2 { };
 template<typename T> requires Bar2<T> && true struct S2<T> { };
 // expected-error@-1{{class template partial specialization is not more specialized than the primary template}}
 // expected-note@-2{{while calculating associated constraint of template 'S2' here}}
+
+namespace type_pack {
+  template<typename... Args>
+  concept C1 = ((sizeof(Args) >= 0) && ...);
+
+  template<typename A, typename... B>
+  concept C2 = C1<A, B...>;
+
+  template<typename T>
+  constexpr void foo() requires C2<T, char, T> { }
+
+  template<typename T>
+  constexpr void foo() requires C1<T, char, T> && true { }
+
+  static_assert((foo<int>(), true));
+}
+
+namespace template_pack {
+  template<typename T> struct S1 {};
+  template<typename T> struct S2 {};
+
+  template<template<typename> typename... Args>
+  concept C1 = ((sizeof(Args<int>) >= 0) && ...);
+
+  template<template<typename> typename A, template<typename> typename... B>
+  concept C2 = C1<A, B...>;
+
+  template<template<typename> typename T>
+  constexpr void foo() requires C2<T, S1, T> { }
+
+  template<template<typename> typename T>
+  constexpr void foo() requires C1<T, S1, T> && true { }
+
+  static_assert((foo<S2>(), true));
+}
+
+namespace non_type_pack {
+  template<int... Args>
+  concept C1 = ((Args >= 0) && ...);
+
+  template<int A, int... B>
+  concept C2 = C1<A, B...>;
+
+  template<int T>
+  constexpr void foo() requires C2<T, 2, T> { }
+
+  template<int T>
+  constexpr void foo() requires C1<T, 2, T> && true { }
+
+  static_assert((foo<1>(), true));
+}
diff --git a/clang/test/SemaTemplate/instantiate-template-argument.cpp b/clang/test/SemaTemplate/instantiate-template-argument.cpp
new file mode 100644
index 00000000000..43d5d00c8cb
--- /dev/null
+++ b/clang/test/SemaTemplate/instantiate-template-argument.cpp
@@ -0,0 +1,28 @@
+// RUN: %clang_cc1 -std=c++2a -x c++ %s -verify
+
+template<auto T, decltype(T) U>
+concept C1 = sizeof(U) >= 4;
+// sizeof(U) >= 4 [U = U (decltype(T))]
+
+template<typename Y, char V>
+concept C2 = C1<Y{}, V>;
+// sizeof(U) >= 4 [U = V (decltype(Y{}))]
+
+template<char W>
+constexpr int foo() requires C2<int, W> { return 1; }
+// sizeof(U) >= 4 [U = W (decltype(int{}))]
+
+template<char X>
+// expected-note@+1{{candidate function}}
+constexpr int foo() requires C1<1, X> && true { return 2; }
+// sizeof(U) >= 4 [U = X (decltype(1))]
+
+static_assert(foo<'a'>() == 2);
+
+template<char Z>
+// expected-note@+1{{candidate function}}
+constexpr int foo() requires C2<long long, Z> && true { return 3; }
+// sizeof(U) >= 4 [U = Z (decltype(long long{}))]
+
+static_assert(foo<'a'>() == 3);
+// expected-error@-1{{call to 'foo' is ambiguous}}
\ No newline at end of file