Reset CladUtils AST caches at plugin initialization

Author: Elvand-Lie

include/clad/Differentiator/CladUtils.h

@@ -397,6 +397,8 @@ namespace clad {
 
     /// Find namespace clad declaration.
     clang::NamespaceDecl* GetCladNamespace(clang::Sema& S);
+    /// Reset AST-owned lookup caches for a new translation unit.
+    void ResetCladUtilsASTCache(clang::ASTContext& C);
     /// Create clad::array<T> type.
     clang::QualType GetCladArrayOfType(clang::Sema& S, clang::QualType T);
     /// Create clad::matrix<T> type.

lib/Differentiator/CladUtils.cpp

@@ -892,17 +892,34 @@ namespace clad {
       return true;
     }
 
+    namespace {
+      NamespaceDecl* CachedCladNS = nullptr;
+      QualType CachedRestoreTrackerTy;
+      TemplateDecl* CachedMatrixDecl = nullptr;
+      TemplateDecl* CachedArrayDecl = nullptr;
+      TemplateDecl* CachedArrayRefDecl = nullptr;
+      TemplateDecl* CachedTagDecl = nullptr;
+    }
+
+    void ResetCladUtilsASTCache(ASTContext&) {
+      CachedCladNS = nullptr;
+      CachedRestoreTrackerTy = QualType();
+      CachedMatrixDecl = nullptr;
+      CachedArrayDecl = nullptr;
+      CachedArrayRefDecl = nullptr;
+      CachedTagDecl = nullptr;
+    }
+
     NamespaceDecl* GetCladNamespace(Sema& S) {
-      static NamespaceDecl* Result = nullptr;
-      if (Result)
-        return Result;
+      if (CachedCladNS)
+        return CachedCladNS;
       DeclarationName CladName = &S.getASTContext().Idents.get("clad");
       LookupResult CladR(S, CladName, noLoc, Sema::LookupNamespaceName,
                          CLAD_COMPAT_Sema_ForVisibleRedeclaration);
       S.LookupQualifiedName(CladR, S.getASTContext().getTranslationUnitDecl());
       assert(!CladR.empty() && "cannot find clad namespace");
-      Result = cast<NamespaceDecl>(CladR.getFoundDecl());
-      return Result;
+      CachedCladNS = cast<NamespaceDecl>(CladR.getFoundDecl());
+      return CachedCladNS;
     }
 
     Expr* getZeroInit(QualType T, Sema& S) {
@@ -945,9 +962,8 @@ namespace clad {
     }
 
     clang::QualType GetRestoreTrackerType(clang::Sema& S) {
-      static QualType T;
-      if (!T.isNull())
-        return T;
+      if (!CachedRestoreTrackerTy.isNull())
+        return CachedRestoreTrackerTy;
       NamespaceDecl* CladNS = GetCladNamespace(S);
       CXXScopeSpec CSS;
       CSS.Extend(S.getASTContext(), CladNS, noLoc, noLoc);
@@ -961,15 +977,16 @@ namespace clad {
       // This will instantiate restore_tracker<T> type and return it.
       auto* RD = cast<RecordDecl>(TrackerR.getFoundDecl());
       ASTContext& C = S.getASTContext();
-      T = clad_compat::getRecordType(C, RD);
+      CachedRestoreTrackerTy = clad_compat::getRecordType(C, RD);
       // Get clad namespace and its identifier clad::.
       clad_compat::NestedNameSpecifierTy NS = CSS.getScopeRep();
 
       // Create elaborated type with namespace specifier,
       // i.e. class<T> -> clad::class<T>
-      T = clad_compat::getElaboratedType(
-          C, clad_compat::ElaboratedTypeKeyword_None, NS, T);
-      return T;
+      CachedRestoreTrackerTy = clad_compat::getElaboratedType(
+          C, clad_compat::ElaboratedTypeKeyword_None, NS,
+          CachedRestoreTrackerTy);
+      return CachedRestoreTrackerTy;
     }
 
     TemplateDecl* LookupTemplateDeclInCladNamespace(Sema& S,
@@ -1134,19 +1151,18 @@ namespace clad {
     }
 
     QualType GetCladMatrixOfType(Sema& S, clang::QualType T) {
-      static TemplateDecl* matrixDecl = nullptr;
-      if (!matrixDecl)
-        matrixDecl =
+      if (!CachedMatrixDecl)
+        CachedMatrixDecl =
             utils::LookupTemplateDeclInCladNamespace(S,
                                                      /*ClassName=*/"matrix");
-      return InstantiateTemplate(S, matrixDecl, {T});
+      return InstantiateTemplate(S, CachedMatrixDecl, {T});
     }
 
     QualType GetCladArrayOfType(Sema& S, clang::QualType T) {
-      static TemplateDecl* arrayDecl = nullptr;
-      if (!arrayDecl)
-        arrayDecl = LookupTemplateDeclInCladNamespace(S, /*ClassName=*/"array");
-      return utils::InstantiateTemplate(S, arrayDecl, {T});
+      if (!CachedArrayDecl)
+        CachedArrayDecl =
+            LookupTemplateDeclInCladNamespace(S, /*ClassName=*/"array");
+      return utils::InstantiateTemplate(S, CachedArrayDecl, {T});
     }
 
     bool IsDifferentiableType(QualType T) {
@@ -1188,11 +1204,10 @@ namespace clad {
     }
 
     QualType GetCladArrayRefOfType(Sema& S, QualType T) {
-      static TemplateDecl* arrayRefDecl = nullptr;
-      if (!arrayRefDecl)
-        arrayRefDecl = utils::LookupTemplateDeclInCladNamespace(
+      if (!CachedArrayRefDecl)
+        CachedArrayRefDecl = utils::LookupTemplateDeclInCladNamespace(
             S, /*ClassName=*/"array_ref");
-      return utils::InstantiateTemplate(S, arrayRefDecl, {T});
+      return utils::InstantiateTemplate(S, CachedArrayRefDecl, {T});
     }
 
     QualType GetParameterDerivativeType(Sema& S, DiffMode Mode, QualType Type) {
@@ -1392,10 +1407,9 @@ namespace clad {
     }
 
     QualType GetCladTagOfType(Sema& S, QualType T) {
-      static clang::TemplateDecl* CladTag = nullptr;
-      if (!CladTag)
-        CladTag = utils::LookupTemplateDeclInCladNamespace(S, "Tag");
-      return utils::InstantiateTemplate(S, CladTag, {T});
+      if (!CachedTagDecl)
+        CachedTagDecl = utils::LookupTemplateDeclInCladNamespace(S, "Tag");
+      return utils::InstantiateTemplate(S, CachedTagDecl, {T});
     }
 
     Expr* GetCladTagExpr(Sema& S, QualType T) {

tools/ClangPlugin.cpp

@@ -175,6 +175,7 @@ void InitTimers();
                  std::vector<std::unique_ptr<ASTConsumer>>);
 
     void CladPlugin::Initialize(clang::ASTContext& C) {
+      utils::ResetCladUtilsASTCache(C);
       // We know we have a multiplexer. We commit a sin here by stealing it and
       // making the consumer pass-through so that we can delay all operations
       // until clad is happy.