Do not splice the original function's AST nodes into derivatives.

A generated derivative must share no Stmt node with the primal it is
built from. The primal's AST outlives differentiation, so a node clad
reuses from the original exposes the user's own code to any later
in-place edit of the derivative. Several forward-mode paths handed the
primal node straight through as the "cloned" primal instead of cloning
it: the member `this`, non-differentiable decls, __null and __func__,
captured lambdas, the function-name string inside a PredefinedExpr, and
top-level lambda closures.

Clone each of them so the derivative owns a distinct subtree, and add
findPrimalSharedNode to keep it that way -- the cross-boundary sibling
of findSharedNode. It intersects the Stmt::children() reachability sets
of the derivative and its primal (child edges only, so Clang's
legitimate type, template-argument and default-argument shares are
excluded) and is enforced beside findSharedNode: assert in debug,
diagnostic in release.

Author: vgvassilev

lib/Differentiator/ASTIntegrity.cpp

@@ -13,12 +13,55 @@
 #include "clang/AST/StmtOpenMP.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace clang;
 
 namespace clad {
 
+// Collect the Stmt::children() reachability set of Root. Child edges only, NOT
+// RecursiveASTVisitor: RAV also follows type and template-argument edges, and
+// Clang legitimately shares nodes across those (a template-argument constant,
+// an array bound). children() also never descends into a CXXDefaultArgExpr's
+// stored default or a type, so those Clang-owned shares are excluded for free.
+static void collectChildEdges(const Stmt* Root,
+                              llvm::DenseSet<const Stmt*>& Set) {
+  if (!Root)
+    return;
+  llvm::SmallVector<const Stmt*, 64> Work;
+  Set.insert(Root);
+  Work.push_back(Root);
+  while (!Work.empty()) {
+    const Stmt* Cur = Work.pop_back_val();
+    for (const Stmt* Ch : Cur->children())
+      if (Ch && Set.insert(Ch).second)
+        Work.push_back(Ch);
+  }
+}
+
+const Stmt* findPrimalSharedNode(const Stmt* Derivative, const Stmt* Primal) {
+  if (!Derivative || !Primal)
+    return nullptr;
+  llvm::DenseSet<const Stmt*> PrimalNodes;
+  collectChildEdges(Primal, PrimalNodes);
+  llvm::SmallVector<const Stmt*, 64> Work;
+  llvm::DenseSet<const Stmt*> Seen;
+  Seen.insert(Derivative);
+  Work.push_back(Derivative);
+  while (!Work.empty()) {
+    const Stmt* Cur = Work.pop_back_val();
+    for (const Stmt* Ch : Cur->children())
+      if (Ch && Seen.insert(Ch).second) {
+        if (PrimalNodes.count(Ch))
+          return Ch;
+        Work.push_back(Ch);
+      }
+  }
+  return nullptr;
+}
+
 const Stmt* findSharedNode(const Stmt* Root) {
   if (!Root)
     return nullptr;

lib/Differentiator/ASTIntegrity.h

@@ -27,6 +27,16 @@ namespace clad {
 /// proper tree.
 const clang::Stmt* findSharedNode(const clang::Stmt* Root);
 
+/// Return the first node in \p Derivative that is also reachable from \p Primal
+/// through Stmt::children() edges, or nullptr if the derivative splices no node
+/// from its primal. The original function's AST outlives differentiation, so a
+/// generated derivative sharing one of its nodes would corrupt the user's own
+/// code if a later pass edits it. Child edges only: Clang legitimately shares
+/// type/template-argument constants and default-argument expressions, which
+/// Stmt::children() does not traverse.
+const clang::Stmt* findPrimalSharedNode(const clang::Stmt* Derivative,
+                                        const clang::Stmt* Primal);
+
 } // namespace clad
 
 #endif // CLAD_AST_INTEGRITY_H

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -1665,18 +1665,15 @@ StmtDiff BaseForwardModeVisitor::VisitDeclStmt(const DeclStmt* DS) {
     // supported.
     if (typeDecl && (clad::utils::hasNonDifferentiableAttribute(typeDecl) ||
                      typeDecl->isLambda())) {
-      // When reconstructing the primal inside a pushforward, a lambda copied
-      // here would share its operator() body with the primal lambda emitted in
-      // the enclosing derivative. Rebuild it with a fresh closure. The
-      // top-level primal (forward mode) must keep the original closure so its
-      // generated pushforward method still resolves, hence the mode guard.
-      const bool inPushforward = m_DiffReq.Mode == DiffMode::pushforward ||
-                                 m_DiffReq.Mode == DiffMode::vector_pushforward;
+      // A lambda copied verbatim would splice the primal's closure -- its
+      // operator() body -- into the derivative. Rebuild it with a fresh closure
+      // instead; the call's pushforward is regenerated for that closure in
+      // VisitCallExpr, so it still resolves.
       for (auto* D : DS->decls()) {
         assert(isa<VarDecl>(D) && "Mixed decl types in a single decl stmt is "
                                   "not standard c++ syntax");
         auto* VDecl = cast<VarDecl>(D);
-        if (inPushforward && typeDecl->isLambda() && VDecl->getInit())
+        if (typeDecl->isLambda() && VDecl->getInit())
           if (const auto* InnerLE =
                   dyn_cast<LambdaExpr>(VDecl->getInit()->IgnoreImplicit())) {
             Expr* ClonedLambda = buildClonedLambda(InnerLE);
@@ -1689,7 +1686,20 @@ StmtDiff BaseForwardModeVisitor::VisitDeclStmt(const DeclStmt* DS) {
             decls.push_back(NewVD);
             continue;
           }
-        decls.push_back(VDecl);
+        // A lambda without a direct initializer (or any other
+        // non-differentiable decl) is cloned so the derivative does not splice
+        // the primal's VarDecl and init into itself.
+        if (typeDecl->isLambda()) {
+          decls.push_back(VDecl);
+        } else {
+          Expr* clonedInit =
+              VDecl->getInit() ? Clone(VDecl->getInit()) : nullptr;
+          VarDecl* copyVD =
+              BuildVarDecl(VDecl->getType(), VDecl->getNameAsString(),
+                           clonedInit, VDecl->isDirectInit());
+          m_DeclReplacements[VDecl] = copyVD;
+          decls.push_back(copyVD);
+        }
       }
       Stmt* DSClone = BuildDeclStmt(decls);
       return StmtDiff(DSClone, nullptr);
@@ -1794,16 +1804,16 @@ BaseForwardModeVisitor::VisitCStyleCastExpr(const CStyleCastExpr* CSCE) {
 StmtDiff BaseForwardModeVisitor::VisitGNUNullExpr(const clang::GNUNullExpr* E) {
   auto* Constant0 =
       ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, /*val=*/0);
-  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-  return StmtDiff(const_cast<clang::GNUNullExpr*>(E), Constant0);
+  // Clone so the derivative owns its copy rather than the primal's node.
+  return StmtDiff(CloneNode(E), Constant0);
 }
 
 StmtDiff
 BaseForwardModeVisitor::VisitPredefinedExpr(const clang::PredefinedExpr* E) {
   auto* Constant0 =
       ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, /*val=*/0);
-  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-  return StmtDiff(const_cast<clang::PredefinedExpr*>(E), Constant0);
+  // Clone so the derivative owns its copy rather than the primal's node.
+  return StmtDiff(CloneNode(E), Constant0);
 }
 
 StmtDiff
@@ -2301,10 +2311,10 @@ StmtDiff BaseForwardModeVisitor::VisitCXXTemporaryObjectExpr(
 
 StmtDiff
 BaseForwardModeVisitor::VisitCXXThisExpr(const clang::CXXThisExpr* CTE) {
-  // m_ThisExprDerivative is a single cached `_d_this` ref; hand out a fresh
-  // clone so distinct uses do not share the same node.
-  return StmtDiff(const_cast<CXXThisExpr*>(CTE),
-                  CloneNode(m_ThisExprDerivative));
+  // Clone CTE so the derivative owns its `this` node rather than splicing the
+  // primal method's; m_ThisExprDerivative is a single cached `_d_this` ref, so
+  // clone it too so distinct uses do not share the same node.
+  return StmtDiff(CloneNode(CTE), CloneNode(m_ThisExprDerivative));
 }
 
 StmtDiff BaseForwardModeVisitor::VisitCXXNewExpr(const clang::CXXNewExpr* CNE) {

lib/Differentiator/DerivativeBuilder.cpp

@@ -664,6 +664,26 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
                "%1; this is a clad bug -- please report it at "
                "https://github.com/vgvassilev/clad")
               << Shared->getStmtClassName() << FD;
+
+        // A derivative must also not splice a node owned by its primal. The
+        // original function's AST outlives differentiation, so a shared node
+        // exposes the user's own code to any later in-place edit of the
+        // derivative -- the same corruption risk findSharedNode guards against,
+        // across the primal/derivative boundary it cannot see. Enforce it too.
+        if (const clang::FunctionDecl* PrimalFD = request.Function)
+          if (const clang::Stmt* PrimalBody = PrimalFD->getBody()) {
+            const clang::Stmt* FromPrimal =
+                findPrimalSharedNode(Body, PrimalBody);
+            assert(!FromPrimal &&
+                   "clad spliced a primal AST node into a derivative");
+            if (FromPrimal)
+              diag(DiagnosticsEngine::Warning, FD->getLocation(),
+                   "clad reused a '%0' AST node from the original function "
+                   "while "
+                   "differentiating %1; this is a clad bug -- please report it "
+                   "at https://github.com/vgvassilev/clad")
+                  << FromPrimal->getStmtClassName() << FD;
+          }
       }
 #endif
 

lib/Differentiator/ReverseModeVisitor.cpp

@@ -3597,16 +3597,16 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
   StmtDiff ReverseModeVisitor::VisitGNUNullExpr(const clang::GNUNullExpr* E) {
     auto* Constant0 = ConstantFolder::synthesizeLiteral(m_Context.IntTy,
                                                         m_Context, /*val=*/0);
-    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-    return StmtDiff(const_cast<clang::GNUNullExpr*>(E), Constant0);
+    // Clone so the derivative owns its copy rather than the primal's node.
+    return StmtDiff(CloneNode(E), Constant0);
   }
 
   StmtDiff
   ReverseModeVisitor::VisitPredefinedExpr(const clang::PredefinedExpr* E) {
     auto* Constant0 = ConstantFolder::synthesizeLiteral(m_Context.IntTy,
                                                         m_Context, /*val=*/0);
-    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-    return StmtDiff(const_cast<clang::PredefinedExpr*>(E), Constant0);
+    // Clone so the derivative owns its copy rather than the primal's node.
+    return StmtDiff(CloneNode(E), Constant0);
   }
 
   StmtDiff ReverseModeVisitor::VisitCXXFunctionalCastExpr(

lib/Differentiator/StmtClone.cpp

@@ -81,11 +81,17 @@ Stmt* StmtClone::VisitDeclRefExpr(DeclRefExpr *Node) {
 DEFINE_CREATE_EXPR(IntegerLiteral,
                    (Ctx, Node->getValue(), CloneType(Node->getType()),
                     Node->getLocation()))
-DEFINE_CLONE_EXPR_CO(PredefinedExpr,
-                     (Ctx, Node->getLocation(), CloneType(Node->getType()),
-                      Node->getIdentKind()
-                          CLAD_COMPAT_CLANG17_IsTransparent(Node),
-                      Node->getFunctionName()))
+Stmt* StmtClone::VisitPredefinedExpr(PredefinedExpr* Node) {
+  // Clone the inner function-name StringLiteral; passing
+  // Node->getFunctionName() would share it with the source.
+  StringLiteral* FN = Node->getFunctionName();
+  PredefinedExpr* result = PredefinedExpr::Create(
+      Ctx, Node->getLocation(), CloneType(Node->getType()),
+      Node->getIdentKind() CLAD_COMPAT_CLANG17_IsTransparent(Node),
+      FN ? cast<StringLiteral>(Clone(FN)) : nullptr);
+  clad_compat::ExprSetDeps(result, Node);
+  return result;
+}
 DEFINE_CLONE_EXPR(CharacterLiteral,
                   (Node->getValue(), Node->getKind(),
                    CloneType(Node->getType()), Node->getLocation()))

lib/Differentiator/VisitorBase.cpp

@@ -830,26 +830,37 @@ namespace clad {
   }
 
   Expr* VisitorBase::buildClonedLambda(const LambdaExpr* LE) {
-    // A primal copy needs a *fresh* closure type; a plain StmtClone reuses
-    // the original closure, so two clones end up sharing the operator() body
-    // -- violating the one-parent-per-node invariant. Captured lambdas would
-    // need capture rebinding we do not model here; fall back to a plain clone
-    // (they are not currently a source of node sharing).
+    // A primal copy needs a *fresh* closure type; a plain StmtClone reuses the
+    // original closure, so two clones share the operator() body -- both a
+    // one-parent-per-node violation and a node shared with the primal lambda.
+    //
+    // We reproduce only explicit by-copy/by-ref captures of a named variable
+    // (each re-resolves by name against the derivative's in-scope copy). Fall
+    // back to a plain clone for the kinds we do not model -- this-capture,
+    // init-capture, and packs.
 #if CLANG_VERSION_MAJOR < 17
     // Lambda differentiation is unsupported below clang-17 (Lambdas.C is
     // UNSUPPORTED there) and the Sema lambda-introduction entry points used
     // below do not exist yet. Never reached; keep the source compilable.
     return cast<Expr>(Clone(LE));
 #else
-    if (LE->capture_size() != 0)
-      return cast<Expr>(Clone(LE));
+    for (const clang::LambdaCapture& Cap : LE->explicit_captures())
+      if (!Cap.capturesVariable() ||
+          (Cap.getCaptureKind() != clang::LCK_ByCopy &&
+           Cap.getCaptureKind() != clang::LCK_ByRef))
+        return cast<Expr>(Clone(LE));
 
     const CXXMethodDecl* CallOp = LE->getCallOperator();
 
     // Mirror the lambda-introduction dance performed while parsing a lambda;
     // see buildDerivedLambda for the differentiating counterpart.
     LambdaIntroducer Intro;
-    Intro.Default = LCD_None;
+    Intro.Default = LE->getCaptureDefault();
+    for (const clang::LambdaCapture& Cap : LE->explicit_captures())
+      Intro.addCapture(Cap.getCaptureKind(), Cap.getLocation(),
+                       Cap.getCapturedVar()->getIdentifier(), /*EllipsisLoc=*/
+                       noLoc, clang::LambdaCaptureInitKind::NoInit,
+                       clang::ExprResult(), clang::ParsedType(), SourceRange());
     Intro.Range.setBegin(LE->getBeginLoc());
     Intro.Range.setEnd(LE->getEndLoc());
     AttributeFactory AttrFactory;