Move name collision handling to ReferencesUpdater.

vgvassilev · Nov 28, 2023 · a3415b8 · a3415b8
1 parent 9d26a46
commit a3415b8
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 10 deletions.
diff --git a/include/clad/Differentiator/StmtClone.h b/include/clad/Differentiator/StmtClone.h
@@ -153,9 +153,12 @@ namespace utils {
     clang::Sema& m_Sema; // We don't own.
     clang::Scope* m_CurScope; // We don't own.
     const clang::FunctionDecl* m_Function; // We don't own.
+    const std::unordered_map<const clang::VarDecl*, clang::VarDecl*>& m_DeclReplacements; // We don't own.
   public:
     ReferencesUpdater(clang::Sema& SemaRef, clang::Scope* S,
-                      const clang::FunctionDecl* FD);
+                      const clang::FunctionDecl* FD,
+                      const std::unordered_map<const clang::VarDecl*,
+                        clang::VarDecl*>& DeclReplacements);
     bool VisitDeclRefExpr(clang::DeclRefExpr* DRE);
     bool VisitStmt(clang::Stmt* S);
     /// Used to update the size expression of QT

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -1328,11 +1328,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     // Check if referenced Decl was "replaced" with another identifier inside
     // the derivative
     if (const auto* VD = dyn_cast<VarDecl>(DRE->getDecl())) {
-      auto it = m_DeclReplacements.find(VD);
-      if (it != std::end(m_DeclReplacements))
-        clonedDRE = BuildDeclRef(it->second);
-      else
-        clonedDRE = cast<DeclRefExpr>(Clone(DRE));
+      clonedDRE = cast<DeclRefExpr>(Clone(DRE));
       // If current context is different than the context of the original
       // declaration (e.g. we are inside lambda), rebuild the DeclRefExpr
       // with Sema::BuildDeclRefExpr. This is required in some cases, e.g.

diff --git a/lib/Differentiator/StmtClone.cpp b/lib/Differentiator/StmtClone.cpp
@@ -502,15 +502,26 @@ Stmt* StmtClone::VisitStmt(Stmt*) {
 }
 
 ReferencesUpdater::ReferencesUpdater(Sema& SemaRef, Scope* S,
-                                     const FunctionDecl* FD)
-    : m_Sema(SemaRef), m_CurScope(S), m_Function(FD) {}
+                                     const FunctionDecl* FD,
+                                      const std::unordered_map<const clang::VarDecl*,
+                                      clang::VarDecl*>& DeclReplacements)
+    : m_Sema(SemaRef), m_CurScope(S), m_Function(FD), m_DeclReplacements(DeclReplacements) {}
 
 bool ReferencesUpdater::VisitDeclRefExpr(DeclRefExpr* DRE) {
   // We should only update references of the declarations that were inside
   // the original function declaration context.
   // Original function = function that we are currently differentiating.
   if (!DRE->getDecl()->getDeclContext()->Encloses(m_Function))
     return true;
+
+  // Replace the declaration if it is present in `m_DeclReplacements`.
+  if (VarDecl* VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+    auto it = m_DeclReplacements.find(VD);
+    if (it != std::end(m_DeclReplacements)) {
+      DRE->setDecl(it->second);
+    }
+  }
+
   DeclarationNameInfo DNI = DRE->getNameInfo();
 
   LookupResult R(m_Sema, DNI, Sema::LookupOrdinaryName);

diff --git a/lib/Differentiator/VisitorBase.cpp b/lib/Differentiator/VisitorBase.cpp
@@ -84,7 +84,7 @@ namespace clad {
   }
 
   void VisitorBase::updateReferencesOf(Stmt* InSubtree) {
-    utils::ReferencesUpdater up(m_Sema, getCurrentScope(), m_Function);
+    utils::ReferencesUpdater up(m_Sema, getCurrentScope(), m_Function, m_DeclReplacements);
     up.TraverseStmt(InSubtree);
   }
 
@@ -304,7 +304,7 @@ namespace clad {
 
   QualType VisitorBase::CloneType(const QualType QT) {
     auto clonedType = m_Builder.m_NodeCloner->CloneType(QT);
-    utils::ReferencesUpdater up(m_Sema, getCurrentScope(), m_Function);
+    utils::ReferencesUpdater up(m_Sema, getCurrentScope(), m_Function, m_DeclReplacements);
     up.updateType(clonedType);
     return clonedType;
   }