WIP

PetroZarytskyi · Dec 12, 2023 · f672ce0 · f672ce0
1 parent ec5ea41
commit f672ce0
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h
@@ -41,6 +41,8 @@ namespace clad {
     /// the reverse mode we also accumulate Stmts for the reverse pass which
     /// will be executed on return.
     std::vector<Stmts> m_Reverse;
+    /// Accumulates local variables for all visited blocks.
+    std::vector<std::vector<clang::VarDecl*>> m_Locals;
     /// Stack is used to pass the arguments (dfdx) to further nodes
     /// in the Visit method.
     std::stack<clang::Expr*> m_Stack;
@@ -142,15 +144,23 @@ namespace clad {
     }
     /// Create new block.
     Stmts& beginBlock(direction d = direction::forward) {
-      if (d == direction::forward)
+      if (d == direction::forward) {
         m_Blocks.emplace_back();
-      else
+        m_Locals.emplace_back();
+      } else {
         m_Reverse.emplace_back();
+      }
       return getCurrentBlock(d);
     }
     /// Remove the block from the stack, wrap it in CompoundStmt and return it.
     clang::CompoundStmt* endBlock(direction d = direction::forward) {
       if (d == direction::forward) {
+        for (clang::VarDecl* VD : m_Locals.back()) {
+          clang::Stmt* declStmt = BuildDeclStmt(VD);
+          addToCurrentBlock(declStmt, direction::reverse);
+        }
+        m_Locals.pop_back();
+
         auto* CS = MakeCompoundStmt(getCurrentBlock(direction::forward));
         m_Blocks.pop_back();
         return CS;
@@ -164,9 +174,14 @@ namespace clad {
 
     Stmts EndBlockWithoutCreatingCS(direction d = direction::forward) {
       auto blk = getCurrentBlock(d);
-      if (d == direction::forward)
+      if (d == direction::forward) {
         m_Blocks.pop_back();
-      else
+        for (clang::VarDecl* VD : m_Locals.back()) {
+          clang::Stmt* declStmt = BuildDeclStmt(VD);
+          addToCurrentBlock(declStmt, direction::reverse);
+        }
+        m_Locals.pop_back();
+      } else
         m_Reverse.pop_back();
       return blk;
     }

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -2613,6 +2613,8 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
       }
     }
     m_Variables.emplace(VDClone, derivedVDE);
+    if (m_Blocks.size()!=1)
+      m_Locals.back().push_back(VDClone);
 
     return VarDeclDiff(VDClone, VDDerived);
   }