Introduce AnalysisDeclContext to avoid repetetive CFG construction

Author: Max Andriychuk

include/clad/Differentiator/DiffPlanner.h

@@ -8,15 +8,19 @@
 #include "clad/Differentiator/Timers.h"
 
 #include "clang/AST/Decl.h"
+#include "clang/AST/DeclBase.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <iterator>
+#include <memory>
 #include <set>
 
 namespace clang {
@@ -31,7 +35,8 @@ class Type;
 } // namespace clang
 
 namespace clad {
-
+using OwnedAnalysisContexts =
+    llvm::SmallVector<std::unique_ptr<clang::AnalysisDeclContext>, 4>;
 /// A struct containing information about request to differentiate a function.
 struct DiffRequest {
 private:
@@ -123,6 +128,8 @@ struct DiffRequest {
   /// This will be particularly useful for pushforward and pullback functions.
   bool DeclarationOnly = false;
 
+  clang::AnalysisDeclContext* m_AnalysisDC;
+
   /// Recomputes `DiffInputVarsInfo` using the current values of data members.
   ///
   /// Differentiation parameters info is computed by parsing the argument
@@ -146,6 +153,8 @@ struct DiffRequest {
     // Note that CallContext is always different and we should ignore it.
     // CustomDerivative is an Expr* and is not always equal even if
     // the set of overloads is the same.
+    // Including AnalysisDC would complicate constructing requests to find the
+    // existing once.
     return Function == other.Function &&
            BaseFunctionName == other.BaseFunctionName &&
            CurrentDerivativeOrder == other.CurrentDerivativeOrder &&
@@ -178,6 +187,10 @@ struct DiffRequest {
   std::string ComputeDerivativeName() const;
   bool HasIndependentParameter(const clang::ParmVarDecl* PVD) const;
 
+  std::set<clang::SourceLocation>& getToBeRecorded() const {
+    m_TbrRunInfo.HasAnalysisRun = true;
+    return m_TbrRunInfo.ToBeRecorded;
+  }
   void addVariedDecl(const clang::VarDecl* init) {
     m_ActivityRunInfo.VariedDecls.insert(init);
   }
@@ -190,6 +203,7 @@ struct DiffRequest {
   std::set<const clang::VarDecl*>& getUsefulDecls() const {
     return m_UsefulRunInfo.UsefulDecls;
   }
+  bool HasTbrAnalysisRun() const { return m_TbrRunInfo.HasAnalysisRun; }
 };
 
   using DiffInterval = std::vector<clang::SourceRange>;
@@ -210,7 +224,10 @@ struct DiffRequest {
     /// Graph to store the dependencies between different requests.
     ///
     clad::DynamicGraph<DiffRequest>& m_DiffRequestGraph;
-
+    /// Map that contains all AnalysisDeclContext for all declrations.
+    /// Essentially needed for prolonging the lifetime of
+    /// unique_ptr<clang::AnalysisDeclContext>.
+    OwnedAnalysisContexts& m_AllAnalysisDC;
     /// If set it means that we need to find the called functions and
     /// add them for implicit diff.
     ///
@@ -228,7 +245,7 @@ struct DiffRequest {
   public:
     DiffCollector(clang::DeclGroupRef DGR, DiffInterval& Interval,
                   clad::DynamicGraph<DiffRequest>& requestGraph, clang::Sema& S,
-                  RequestOptions& opts);
+                  RequestOptions& opts, OwnedAnalysisContexts& AllAnalysisDC);
     bool VisitCallExpr(clang::CallExpr* E);
     bool VisitDeclRefExpr(clang::DeclRefExpr* DRE);
     bool VisitCXXConstructExpr(clang::CXXConstructExpr* e);

lib/Differentiator/ActivityAnalyzer.cpp

@@ -6,14 +6,10 @@ using namespace clang;
 namespace clad {
 
 void VariedAnalyzer::Analyze(const FunctionDecl* FD) {
-  // Build the CFG (control-flow graph) of FD.
-  clang::CFG::BuildOptions Options;
-  m_CFG = clang::CFG::buildCFG(FD, FD->getBody(), &m_Context, Options);
-
-  m_BlockData.resize(m_CFG->size());
+  m_BlockData.resize(m_AnalysisDC->getCFG()->size());
   // Set current block ID to the ID of entry the block.
-  CFGBlock* entry = &m_CFG->getEntry();
-  m_CurBlockID = entry->getBlockID();
+  CFGBlock& entry = m_AnalysisDC->getCFG()->getEntry();
+  m_CurBlockID = entry.getBlockID();
   m_BlockData[m_CurBlockID] = createNewVarsData({});
   for (const VarDecl* i : m_VariedDecls)
     m_BlockData[m_CurBlockID]->insert(i);
@@ -38,7 +34,7 @@ void mergeVarsData(std::set<const clang::VarDecl*>* targetData,
 }
 
 CFGBlock* VariedAnalyzer::getCFGBlockByID(unsigned ID) {
-  return *(m_CFG->begin() + ID);
+  return *(m_AnalysisDC->getCFG()->begin() + ID);
 }
 
 void VariedAnalyzer::AnalyzeCFGBlock(const CFGBlock& block) {

lib/Differentiator/ActivityAnalyzer.h

@@ -2,6 +2,7 @@
 #define CLAD_DIFFERENTIATOR_ACTIVITYANALYZER_H
 
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
 #include "clang/Analysis/CFG.h"
 
 #include "clad/Differentiator/CladUtils.h"
@@ -35,8 +36,7 @@ class VariedAnalyzer : public clang::RecursiveASTVisitor<VariedAnalyzer> {
 
   clang::CFGBlock* getCFGBlockByID(unsigned ID);
 
-  clang::ASTContext& m_Context;
-  std::unique_ptr<clang::CFG> m_CFG;
+  clang::AnalysisDeclContext* m_AnalysisDC;
   std::vector<std::unique_ptr<VarsData>> m_BlockData;
   unsigned m_CurBlockID{};
   std::set<unsigned> m_CFGQueue;
@@ -55,9 +55,9 @@ class VariedAnalyzer : public clang::RecursiveASTVisitor<VariedAnalyzer> {
 
 public:
   /// Constructor
-  VariedAnalyzer(clang::ASTContext& Context,
+  VariedAnalyzer(clang::AnalysisDeclContext* AnalysisDC,
                  std::set<const clang::VarDecl*>& Decls)
-      : m_VariedDecls(Decls), m_Context(Context) {}
+      : m_VariedDecls(Decls), m_AnalysisDC(AnalysisDC) {}
 
   /// Destructor
   ~VariedAnalyzer() = default;

lib/Differentiator/DiffPlanner.cpp

@@ -21,7 +21,7 @@
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/ExprObjC.h"
 #include "clang/AST/RecursiveASTVisitor.h"
-#include "clang/Analysis/CallGraph.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
 #include "clang/Basic/IdentifierTable.h"
 #include "clang/Basic/LLVM.h" // isa, dyn_cast
 #include "clang/Basic/SourceLocation.h"
@@ -32,6 +32,7 @@
 #include "clang/Sema/TemplateDeduction.h"
 
 #include <algorithm>
+#include <memory>
 #include <string>
 #include <utility>
 
@@ -273,9 +274,10 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
 
   DiffCollector::DiffCollector(DeclGroupRef DGR, DiffInterval& Interval,
                                clad::DynamicGraph<DiffRequest>& requestGraph,
-                               clang::Sema& S, RequestOptions& opts)
-      : m_Interval(Interval), m_DiffRequestGraph(requestGraph), m_Sema(S),
-        m_Options(opts) {
+                               clang::Sema& S, RequestOptions& opts,
+                               OwnedAnalysisContexts& AllAnalysisDC)
+      : m_Interval(Interval), m_DiffRequestGraph(requestGraph),
+        m_AllAnalysisDC(AllAnalysisDC), m_Sema(S), m_Options(opts) {
 
     if (Interval.empty())
       return;
@@ -648,13 +650,11 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
     if (E->getType()->isPointerType())
       return true;
 
-    if (!m_TbrRunInfo.HasAnalysisRun && !isLambdaCallOperator(Function)) {
+    if (!m_TbrRunInfo.HasAnalysisRun && !isLambdaCallOperator(Function) &&
+        Function->isDefined()) {
       TimedAnalysisRegion R("TBR " + BaseFunctionName);
-
-      TBRAnalyzer analyzer(Function->getASTContext(),
-                           m_TbrRunInfo.ToBeRecorded);
-      analyzer.Analyze(Function);
-      m_TbrRunInfo.HasAnalysisRun = true;
+      TBRAnalyzer analyzer(m_AnalysisDC, getToBeRecorded());
+      analyzer.Analyze(*this);
     }
     auto found = m_TbrRunInfo.ToBeRecorded.find(E->getBeginLoc());
     return found != m_TbrRunInfo.ToBeRecorded.end();
@@ -1214,22 +1214,27 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
       request.Function = request.Function->getDefinition();
 
     if (!LookupCustomDerivativeDecl(request)) {
-      if (m_TopMostReq->EnableVariedAnalysis &&
-          m_TopMostReq->Mode == DiffMode::reverse) {
+      clang::CFG::BuildOptions Options;
+      std::unique_ptr<AnalysisDeclContext> AnalysisDC =
+          std::make_unique<AnalysisDeclContext>(
+              /*AnalysisDeclContextManager=*/nullptr, request.Function,
+              Options);
+
+      if (m_TopMostReq->EnableVariedAnalysis) {
         TimedAnalysisRegion R("VA " + request.BaseFunctionName);
-        VariedAnalyzer analyzer(request.Function->getASTContext(),
-                                request.getVariedDecls());
+        VariedAnalyzer analyzer(AnalysisDC.get(), request.getVariedDecls());
         analyzer.Analyze(request.Function);
       }
 
       if (m_TopMostReq->EnableUsefulAnalysis) {
         TimedAnalysisRegion R("UA " + request.BaseFunctionName);
-
-        UsefulAnalyzer analyzer(request.Function->getASTContext(),
-                                request.getUsefulDecls());
+        UsefulAnalyzer analyzer(AnalysisDC.get(), request.getUsefulDecls());
         analyzer.Analyze(request.Function);
       }
 
+      m_AllAnalysisDC.push_back(std::move(AnalysisDC));
+      request.m_AnalysisDC = m_AllAnalysisDC.back().get();
+
       // Recurse into call graph.
       TraverseFunctionDeclOnce(request.Function);
     }
@@ -1313,9 +1318,19 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
     if (m_Sema.isStdInitializerList(recordTy, /*elemType=*/nullptr))
       return true;
 
-    if (!LookupCustomDerivativeDecl(request))
+    if (!LookupCustomDerivativeDecl(request)) {
+      clang::CFG::BuildOptions Options;
+      std::unique_ptr<AnalysisDeclContext> AnalysisDC =
+          std::make_unique<AnalysisDeclContext>(
+              /*AnalysisDeclContextManager=*/nullptr, request.Function,
+              Options);
+      // FIXME: Add proper support for objects in VA and UA.
+      m_AllAnalysisDC.push_back(std::move(AnalysisDC));
+      request.m_AnalysisDC = m_AllAnalysisDC.back().get();
+
       // Recurse into call graph.
       TraverseFunctionDeclOnce(request.Function);
+    }
     m_DiffRequestGraph.addNode(request, /*isSource=*/true);
 
     return true;

lib/Differentiator/TBRAnalyzer.cpp

@@ -8,8 +8,12 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/OperationKinds.h"
+#include "clang/Analysis/CFG.h"
 #include "clang/Basic/LLVM.h"
 
+#include "clad/Differentiator/Compatibility.h"
+#include "clad/Differentiator/DiffPlanner.h"
+
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
@@ -251,6 +255,7 @@ void TBRAnalyzer::addVar(const clang::VarDecl* VD, bool forceNonRefType) {
     varType = arrayParam->getOriginalType();
   else
     varType = VD->getType();
+
   // If varType represents auto or auto*, get the type of init.
   if (utils::IsAutoOrAutoPtrType(varType))
     varType = VD->getInit()->getType();
@@ -294,19 +299,16 @@ TBRAnalyzer::getVarDataFromDecl(const clang::VarDecl* VD) {
   return nullptr;
 }
 
-void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
-  // Build the CFG (control-flow graph) of FD.
-  clang::CFG::BuildOptions Options;
-  m_CFG = clang::CFG::buildCFG(FD, FD->getBody(), &m_Context, Options);
-
-  m_BlockData.resize(m_CFG->size());
-  m_BlockPassCounter.resize(m_CFG->size(), 0);
+void TBRAnalyzer::Analyze(const DiffRequest& request) {
+  m_BlockData.resize(request.m_AnalysisDC->getCFG()->size());
+  m_BlockPassCounter.resize(request.m_AnalysisDC->getCFG()->size(), 0);
 
   // Set current block ID to the ID of entry the block.
-  auto* entry = &m_CFG->getEntry();
-  m_CurBlockID = entry->getBlockID();
+  CFGBlock& entry = request.m_AnalysisDC->getCFG()->getEntry();
+  m_CurBlockID = entry.getBlockID();
   m_BlockData[m_CurBlockID] = std::unique_ptr<VarsData>(new VarsData());
 
+  const FunctionDecl* FD = request.Function;
   // If we are analysing a non-static method, add a VarData for 'this' pointer
   // (it is represented with nullptr).
   const auto* MD = dyn_cast<CXXMethodDecl>(FD);
@@ -331,13 +333,13 @@ void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
     m_CurBlockID = *IDIter;
     m_CFGQueue.erase(IDIter);
 
-    CFGBlock& nextBlock = *getCFGBlockByID(m_CurBlockID);
+    CFGBlock& nextBlock = *getCFGBlockByID(request.m_AnalysisDC, m_CurBlockID);
     VisitCFGBlock(nextBlock);
   }
 #ifndef NDEBUG
   for (int id = m_CurBlockID; id >= 0; --id) {
     LLVM_DEBUG(llvm::dbgs() << "\n-----BLOCK" << id << "-----\n\n");
-    for (auto succ : getCFGBlockByID(id)->succs())
+    for (auto succ : getCFGBlockByID(request.m_AnalysisDC, id)->succs())
       if (succ)
         LLVM_DEBUG(llvm::dbgs() << "successor: " << succ->getBlockID() << "\n");
   }
@@ -409,8 +411,8 @@ void TBRAnalyzer::VisitCFGBlock(const CFGBlock& block) {
   LLVM_DEBUG(llvm::dbgs() << "Leaving block " << block.getBlockID() << "\n");
 }
 
-CFGBlock* TBRAnalyzer::getCFGBlockByID(unsigned ID) {
-  return *(m_CFG->begin() + ID);
+CFGBlock* TBRAnalyzer::getCFGBlockByID(AnalysisDeclContext* ADC, unsigned ID) {
+  return *(ADC->getCFG()->begin() + ID);
 }
 
 TBRAnalyzer::VarsData*
@@ -566,9 +568,11 @@ bool TBRAnalyzer::TraverseDeclStmt(DeclStmt* DS) {
     if (auto* VD = dyn_cast<VarDecl>(D)) {
       addVar(VD);
       if (clang::Expr* init = VD->getInit()) {
+
         setMode(Mode::kMarkingMode);
         TraverseStmt(init);
         resetMode();
+
         auto& VDExpr = getCurBlockVarsData()[VD];
         // if the declared variable is ref type attach its VarData to the
         // VarData of the RHS variable.