Register a cloned lambda body's nested declarations before remapping.

Author: Shubham Shukla

lib/Differentiator/VisitorBase.cpp

@@ -878,6 +878,28 @@ namespace clad {
     return newPVD;
   }
 
+  /// Register every (original, clone) VarDecl pair, at any depth.
+  static void
+  registerClonedDecls(const Stmt* Orig, Stmt* Cloned,
+                      std::unordered_map<const VarDecl*, VarDecl*>& Repls) {
+    if (!Orig || !Cloned)
+      return;
+    if (const auto* DS = dyn_cast<DeclStmt>(Orig))
+      if (auto* ClonedDS = dyn_cast<DeclStmt>(Cloned)) {
+        auto O = DS->decl_begin();
+        auto C = ClonedDS->decl_begin();
+        for (; O != DS->decl_end() && C != ClonedDS->decl_end(); ++O, ++C)
+          if (const auto* OVD = dyn_cast<VarDecl>(*O))
+            if (auto* CVD = dyn_cast<VarDecl>(*C))
+              if (OVD != CVD)
+                Repls[OVD] = CVD;
+      }
+    auto O = Orig->child_begin();
+    auto C = Cloned->child_begin();
+    for (; O != Orig->child_end() && C != Cloned->child_end(); ++O, ++C)
+      registerClonedDecls(*O, *C, Repls);
+  }
+
   Expr* VisitorBase::buildClonedLambda(const LambdaExpr* LE) {
     // A primal copy needs a *fresh* closure type; a plain StmtClone reuses the
     // original closure, so two clones share the operator() body -- both a
@@ -1026,23 +1048,11 @@ namespace clad {
       // lambda's call operator rather than m_DiffReq.Function (the outer
       // function being differentiated), which would reject the lambda locals.
       Stmt* clonedS = CloneNode(S);
+      // Register before remapping, so references in this statement update too.
+      registerClonedDecls(S, clonedS, m_DeclReplacements);
       utils::ReferencesUpdater up(m_Sema, getCurrentScope(), CallOp,
                                   m_DeclReplacements);
       up.TraverseStmt(clonedS);
-      // Cloning a DeclStmt produces a fresh VarDecl, but StmtClone records that
-      // only in its own decl mapping. Register it here too, or a later
-      // statement's clone keeps referring to the original lambda's variable --
-      // a reference into the user's AST that outlives differentiation.
-      if (const auto* DS = dyn_cast<DeclStmt>(S))
-        if (auto* ClonedDS = dyn_cast<DeclStmt>(clonedS)) {
-          auto O = DS->decl_begin();
-          auto C = ClonedDS->decl_begin();
-          for (; O != DS->decl_end() && C != ClonedDS->decl_end(); ++O, ++C)
-            if (const auto* OVD = dyn_cast<VarDecl>(*O))
-              if (auto* CVD = dyn_cast<VarDecl>(*C))
-                if (OVD != CVD)
-                  m_DeclReplacements[OVD] = CVD;
-        }
       addToCurrentBlock(clonedS);
     }
     CompoundStmt* ClonedBody = endBlock();

test/ForwardMode/Lambdas.C

@@ -42,6 +42,16 @@ double fn4(double x) {
     return _f(x);
 }
 
+double fn5(double x) {
+    auto _f = [](double _x) {
+        double s = 0;
+        for (int k = 0; k < 3; ++k)
+            s += _x * _x;
+        return s;
+    };
+    return _f(x);
+}
+
 int main() {
     auto fn0_dx = clad::differentiate(fn0, 0);
     printf("Result is = %.2f\n", fn0_dx.execute(7)); // CHECK-EXEC: Result is = 14.00
@@ -62,4 +72,8 @@ int main() {
     auto fn4_dx = clad::differentiate(fn4, 0);
     printf("Result is = %.2f\n", fn4_dx.execute(7)); // CHECK-EXEC: Result is = 28.00
     printf("Result is = %.2f\n", fn4_dx.execute(-1)); // CHECK-EXEC: Result is = -4.00
+
+    auto fn5_dx = clad::differentiate(fn5, 0);
+    printf("Result is = %.2f\n", fn5_dx.execute(7)); // CHECK-EXEC: Result is = 42.00
+    printf("Result is = %.2f\n", fn5_dx.execute(-1)); // CHECK-EXEC: Result is = -6.00
 }