[Concepts] Constraint Satisfaction Caching

Add a simple cache for constraint satisfaction results. Whether or not this simple caching
would be permitted in final C++2a is currently being discussed but it is required for
acceptable performance so we use it in the meantime, with the possibility of adding some
cache invalidation mechanisms later.

Differential Revision: https://reviews.llvm.org/D72552

(cherry picked from commit b933d37cd3774e5431b35e82187eebb59b1ff59e)

10 files changed, 144 insertions, 55 deletions

diff --git a/clang/include/clang/AST/ASTConcept.h b/clang/include/clang/AST/ASTConcept.h
index 84a611c14e0..30c4706d2a1 100644
--- a/clang/include/clang/AST/ASTConcept.h
+++ b/clang/include/clang/AST/ASTConcept.h
@@ -24,9 +24,23 @@ namespace clang {
 class ConceptDecl;
 class ConceptSpecializationExpr;
 
-/// \brief The result of a constraint satisfaction check, containing the
-/// necessary information to diagnose an unsatisfied constraint.
-struct ConstraintSatisfaction {
+/// The result of a constraint satisfaction check, containing the necessary
+/// information to diagnose an unsatisfied constraint.
+class ConstraintSatisfaction : public llvm::FoldingSetNode {
+  // The template-like entity that 'owns' the constraint checked here (can be a
+  // constrained entity or a concept).
+  NamedDecl *ConstraintOwner = nullptr;
+  llvm::SmallVector<TemplateArgument, 4> TemplateArgs;
+
+public:
+
+  ConstraintSatisfaction() = default;
+
+  ConstraintSatisfaction(NamedDecl *ConstraintOwner,
+                         ArrayRef<TemplateArgument> TemplateArgs) :
+      ConstraintOwner(ConstraintOwner), TemplateArgs(TemplateArgs.begin(),
+                                                     TemplateArgs.end()) { }
+
   using SubstitutionDiagnostic = std::pair<SourceLocation, StringRef>;
   using Detail = llvm::PointerUnion<Expr *, SubstitutionDiagnostic *>;
 
@@ -38,9 +52,13 @@ struct ConstraintSatisfaction {
   /// invalid expression.
   llvm::SmallVector<std::pair<const Expr *, Detail>, 4> Details;
 
-  // This can leak if used in an AST node, use ASTConstraintSatisfaction
-  // instead.
-  void *operator new(size_t bytes, ASTContext &C) = delete;
+  void Profile(llvm::FoldingSetNodeID &ID, const ASTContext &C) {
+    Profile(ID, C, ConstraintOwner, TemplateArgs);
+  }
+
+  static void Profile(llvm::FoldingSetNodeID &ID, const ASTContext &C,
+                      NamedDecl *ConstraintOwner,
+                      ArrayRef<TemplateArgument> TemplateArgs);
 };
 
 /// Pairs of unsatisfied atomic constraint expressions along with the
diff --git a/clang/include/clang/Basic/LangOptions.def b/clang/include/clang/Basic/LangOptions.def
index 068f206f448..b8112eb2691 100644
--- a/clang/include/clang/Basic/LangOptions.def
+++ b/clang/include/clang/Basic/LangOptions.def
@@ -238,6 +238,7 @@ LANGOPT(AlignedAllocation , 1, 0, "aligned allocation")
 LANGOPT(AlignedAllocationUnavailable, 1, 0, "aligned allocation functions are unavailable")
 LANGOPT(NewAlignOverride  , 32, 0, "maximum alignment guaranteed by '::operator new(size_t)'")
 LANGOPT(ConceptsTS , 1, 0, "enable C++ Extensions for Concepts")
+LANGOPT(ConceptSatisfactionCaching , 1, 1, "enable satisfaction caching for C++2a Concepts")
 BENIGN_LANGOPT(ModulesCodegen , 1, 0, "Modules code generation")
 BENIGN_LANGOPT(ModulesDebugInfo , 1, 0, "Modules debug info")
 BENIGN_LANGOPT(ElideConstructors , 1, 1, "C++ copy constructor elision")
diff --git a/clang/include/clang/Driver/CC1Options.td b/clang/include/clang/Driver/CC1Options.td
index 1d550eb15ea..ac86cda81ee 100644
--- a/clang/include/clang/Driver/CC1Options.td
+++ b/clang/include/clang/Driver/CC1Options.td
@@ -554,6 +554,9 @@ def ftest_module_file_extension_EQ :
            "The argument is parsed as blockname:major:minor:hashed:user info">;
 def fconcepts_ts : Flag<["-"], "fconcepts-ts">,
   HelpText<"Enable C++ Extensions for Concepts.">;
+def fno_concept_satisfaction_caching : Flag<["-"],
+                                            "fno-concept-satisfaction-caching">,
+  HelpText<"Disable satisfaction caching for C++2a Concepts.">;
 
 let Group = Action_Group in {
 
diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index 0f02c0f9d2b..c94027cf56d 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -6200,6 +6200,9 @@ private:
   llvm::DenseMap<NamedDecl *, NormalizedConstraint *>
       NormalizationCache;
 
+  llvm::ContextualFoldingSet<ConstraintSatisfaction, const ASTContext &>
+      SatisfactionCache;
+
 public:
   const NormalizedConstraint *
   getNormalizedAssociatedConstraints(
@@ -6226,6 +6229,8 @@ public:
 
   /// \brief Check whether the given list of constraint expressions are
   /// satisfied (as if in a 'conjunction') given template arguments.
+  /// \param Template the template-like entity that triggered the constraints
+  /// check (either a concept or a constrained entity).
   /// \param ConstraintExprs a list of constraint expressions, treated as if
   /// they were 'AND'ed together.
   /// \param TemplateArgs the list of template arguments to substitute into the
@@ -6237,23 +6242,10 @@ public:
   /// expression.
   /// \returns true if an error occurred and satisfaction could not be checked,
   /// false otherwise.
-  bool CheckConstraintSatisfaction(TemplateDecl *Template,
-                                   ArrayRef<const Expr *> ConstraintExprs,
-                                   ArrayRef<TemplateArgument> TemplateArgs,
-                                   SourceRange TemplateIDRange,
-                                   ConstraintSatisfaction &Satisfaction);
-
-  bool CheckConstraintSatisfaction(ClassTemplatePartialSpecializationDecl *TD,
-                                   ArrayRef<const Expr *> ConstraintExprs,
-                                   ArrayRef<TemplateArgument> TemplateArgs,
-                                   SourceRange TemplateIDRange,
-                                   ConstraintSatisfaction &Satisfaction);
-
-  bool CheckConstraintSatisfaction(VarTemplatePartialSpecializationDecl *TD,
-                                   ArrayRef<const Expr *> ConstraintExprs,
-                                   ArrayRef<TemplateArgument> TemplateArgs,
-                                   SourceRange TemplateIDRange,
-                                   ConstraintSatisfaction &Satisfaction);
+  bool CheckConstraintSatisfaction(
+      NamedDecl *Template, ArrayRef<const Expr *> ConstraintExprs,
+      ArrayRef<TemplateArgument> TemplateArgs,
+      SourceRange TemplateIDRange, ConstraintSatisfaction &Satisfaction);
 
   /// \brief Check whether the given non-dependent constraint expression is
   /// satisfied. Returns false and updates Satisfaction with the satisfaction
diff --git a/clang/include/clang/Sema/TemplateDeduction.h b/clang/include/clang/Sema/TemplateDeduction.h
index b60939c9787..f787c2689d8 100644
--- a/clang/include/clang/Sema/TemplateDeduction.h
+++ b/clang/include/clang/Sema/TemplateDeduction.h
@@ -15,6 +15,7 @@
 #define LLVM_CLANG_SEMA_TEMPLATEDEDUCTION_H
 
 #include "clang/Sema/Ownership.h"
+#include "clang/Sema/SemaConcept.h"
 #include "clang/AST/ASTConcept.h"
 #include "clang/AST/DeclAccessPair.h"
 #include "clang/AST/DeclTemplate.h"
diff --git a/clang/lib/AST/ASTConcept.cpp b/clang/lib/AST/ASTConcept.cpp
index fc32e768d92..66d272da704 100644
--- a/clang/lib/AST/ASTConcept.cpp
+++ b/clang/lib/AST/ASTConcept.cpp
@@ -14,6 +14,7 @@
 
 #include "clang/AST/ASTConcept.h"
 #include "clang/AST/ASTContext.h"
+#include "clang/Sema/SemaConcept.h"
 using namespace clang;
 
 ASTConstraintSatisfaction::ASTConstraintSatisfaction(const ASTContext &C,
@@ -53,3 +54,12 @@ ASTConstraintSatisfaction::Create(const ASTContext &C,
   void *Mem = C.Allocate(size, alignof(ASTConstraintSatisfaction));
   return new (Mem) ASTConstraintSatisfaction(C, Satisfaction);
 }
+
+void ConstraintSatisfaction::Profile(
+    llvm::FoldingSetNodeID &ID, const ASTContext &C, NamedDecl *ConstraintOwner,
+    ArrayRef<TemplateArgument> TemplateArgs) {
+  ID.AddPointer(ConstraintOwner);
+  ID.AddInteger(TemplateArgs.size());
+  for (auto &Arg : TemplateArgs)
+    Arg.Profile(ID, C);
+}
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index e1e59565083..af5113f0aba 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -2853,6 +2853,8 @@ static void ParseLangArgs(LangOptions &Opts, ArgList &Args, InputKind IK,
     Opts.NewAlignOverride = 0;
   }
   Opts.ConceptsTS = Args.hasArg(OPT_fconcepts_ts);
+  Opts.ConceptSatisfactionCaching =
+      !Args.hasArg(OPT_fno_concept_satisfaction_caching);
   Opts.HeinousExtensions = Args.hasArg(OPT_fheinous_gnu_extensions);
   Opts.AccessControl = !Args.hasArg(OPT_fno_access_control);
   Opts.ElideConstructors = !Args.hasArg(OPT_fno_elide_constructors);
diff --git a/clang/lib/Sema/Sema.cpp b/clang/lib/Sema/Sema.cpp
index f8da1cb89b9..7eb8c8d2f76 100644
--- a/clang/lib/Sema/Sema.cpp
+++ b/clang/lib/Sema/Sema.cpp
@@ -153,10 +153,10 @@ Sema::Sema(Preprocessor &pp, ASTContext &ctxt, ASTConsumer &consumer,
       TUKind(TUKind), NumSFINAEErrors(0),
       FullyCheckedComparisonCategories(
           static_cast<unsigned>(ComparisonCategoryType::Last) + 1),
-      AccessCheckingSFINAE(false), InNonInstantiationSFINAEContext(false),
-      NonInstantiationEntries(0), ArgumentPackSubstitutionIndex(-1),
-      CurrentInstantiationScope(nullptr), DisableTypoCorrection(false),
-      TyposCorrected(0), AnalysisWarnings(*this),
+      SatisfactionCache(Context), AccessCheckingSFINAE(false),
+      InNonInstantiationSFINAEContext(false), NonInstantiationEntries(0),
+      ArgumentPackSubstitutionIndex(-1), CurrentInstantiationScope(nullptr),
+      DisableTypoCorrection(false), TyposCorrected(0), AnalysisWarnings(*this),
       ThreadSafetyDeclCache(nullptr), VarDataSharingAttributesStack(nullptr),
       CurScope(nullptr), Ident_super(nullptr), Ident___float128(nullptr) {
   TUScope = nullptr;
@@ -379,6 +379,14 @@ Sema::~Sema() {
   if (isMultiplexExternalSource)
     delete ExternalSource;
 
+  // Delete cached satisfactions.
+  std::vector<ConstraintSatisfaction *> Satisfactions;
+  Satisfactions.reserve(Satisfactions.size());
+  for (auto &Node : SatisfactionCache)
+    Satisfactions.push_back(&Node);
+  for (auto *Node : Satisfactions)
+    delete Node;
+
   threadSafety::threadSafetyCleanup(ThreadSafetyDeclCache);
 
   // Destroys data sharing attributes stack for OpenMP
diff --git a/clang/lib/Sema/SemaConcept.cpp b/clang/lib/Sema/SemaConcept.cpp
index 93e5b4511da..81601b09ce0 100755
--- a/clang/lib/Sema/SemaConcept.cpp
+++ b/clang/lib/Sema/SemaConcept.cpp
@@ -272,36 +272,56 @@ static bool CheckConstraintSatisfaction(Sema &S, TemplateDeclT *Template,
   return false;
 }
 
-bool Sema::CheckConstraintSatisfaction(TemplateDecl *Template,
-                                       ArrayRef<const Expr *> ConstraintExprs,
-                                       ArrayRef<TemplateArgument> TemplateArgs,
-                                       SourceRange TemplateIDRange,
-                                       ConstraintSatisfaction &Satisfaction) {
-  return ::CheckConstraintSatisfaction(*this, Template, ConstraintExprs,
-                                       TemplateArgs, TemplateIDRange,
-                                       Satisfaction);
-}
+bool Sema::CheckConstraintSatisfaction(
+    NamedDecl *Template, ArrayRef<const Expr *> ConstraintExprs,
+    ArrayRef<TemplateArgument> TemplateArgs, SourceRange TemplateIDRange,
+    ConstraintSatisfaction &OutSatisfaction) {
+  if (ConstraintExprs.empty()) {
+    OutSatisfaction.IsSatisfied = true;
+    return false;
+  }
 
-bool
-Sema::CheckConstraintSatisfaction(ClassTemplatePartialSpecializationDecl* Part,
-                                  ArrayRef<const Expr *> ConstraintExprs,
-                                  ArrayRef<TemplateArgument> TemplateArgs,
-                                  SourceRange TemplateIDRange,
-                                  ConstraintSatisfaction &Satisfaction) {
-  return ::CheckConstraintSatisfaction(*this, Part, ConstraintExprs,
-                                       TemplateArgs, TemplateIDRange,
-                                       Satisfaction);
-}
+  llvm::FoldingSetNodeID ID;
+  void *InsertPos;
+  ConstraintSatisfaction *Satisfaction = nullptr;
+  if (LangOpts.ConceptSatisfactionCaching) {
+    ConstraintSatisfaction::Profile(ID, Context, Template, TemplateArgs);
+    Satisfaction = SatisfactionCache.FindNodeOrInsertPos(ID, InsertPos);
+    if (Satisfaction) {
+      OutSatisfaction = *Satisfaction;
+      return false;
+    }
+    Satisfaction = new ConstraintSatisfaction(Template, TemplateArgs);
+  } else {
+    Satisfaction = &OutSatisfaction;
+  }
+  bool Failed;
+  if (auto *T = dyn_cast<TemplateDecl>(Template))
+    Failed = ::CheckConstraintSatisfaction(*this, T, ConstraintExprs,
+                                           TemplateArgs, TemplateIDRange,
+                                           *Satisfaction);
+  else if (auto *P =
+               dyn_cast<ClassTemplatePartialSpecializationDecl>(Template))
+    Failed = ::CheckConstraintSatisfaction(*this, P, ConstraintExprs,
+                                           TemplateArgs, TemplateIDRange,
+                                           *Satisfaction);
+  else
+    Failed = ::CheckConstraintSatisfaction(
+        *this, cast<VarTemplatePartialSpecializationDecl>(Template),
+        ConstraintExprs, TemplateArgs, TemplateIDRange, *Satisfaction);
+  if (Failed) {
+    if (LangOpts.ConceptSatisfactionCaching)
+      delete Satisfaction;
+    return true;
+  }
 
-bool
-Sema::CheckConstraintSatisfaction(VarTemplatePartialSpecializationDecl* Partial,
-                                  ArrayRef<const Expr *> ConstraintExprs,
-                                  ArrayRef<TemplateArgument> TemplateArgs,
-                                  SourceRange TemplateIDRange,
-                                  ConstraintSatisfaction &Satisfaction) {
-  return ::CheckConstraintSatisfaction(*this, Partial, ConstraintExprs,
-                                       TemplateArgs, TemplateIDRange,
-                                       Satisfaction);
+  if (LangOpts.ConceptSatisfactionCaching) {
+    // We cannot use InsertNode here because CheckConstraintSatisfaction might
+    // have invalidated it.
+    SatisfactionCache.InsertNode(Satisfaction);
+    OutSatisfaction = *Satisfaction;
+  }
+  return false;
 }
 
 bool Sema::CheckConstraintSatisfaction(const Expr *ConstraintExpr,
diff --git a/clang/test/SemaTemplate/cxx2a-constraint-caching.cpp b/clang/test/SemaTemplate/cxx2a-constraint-caching.cpp
new file mode 100644
index 00000000000..fd95c80ff61
--- /dev/null
+++ b/clang/test/SemaTemplate/cxx2a-constraint-caching.cpp
@@ -0,0 +1,34 @@
+// RUN:  %clang_cc1 -std=c++2a -fconcepts-ts -verify %s
+// RUN:  %clang_cc1 -std=c++2a -fconcepts-ts -verify %s -fno-concept-satisfaction-caching -DNO_CACHE
+// expected-no-diagnostics
+
+template<typename T>
+concept C = (f(T()), true);
+
+template<typename T>
+constexpr bool foo() { return false; }
+
+template<typename T>
+  requires (f(T()), true)
+constexpr bool foo() requires (f(T()), true) { return true; }
+
+namespace a {
+  struct A {};
+  void f(A a);
+}
+
+static_assert(C<a::A>);
+static_assert(foo<a::A>());
+
+namespace a {
+  // This makes calls to f ambiguous, but the second check will still succeed
+  // because the constraint satisfaction results are cached.
+  void f(A a, int = 2);
+}
+#ifdef NO_CACHE
+static_assert(!C<a::A>);
+static_assert(!foo<a::A>());
+#else
+static_assert(C<a::A>);
+static_assert(foo<a::A>());
+#endif
\ No newline at end of file