Do not clone reverse-mode leaf representations until read.

A reverse-mode StmtDiff carries three representations of a value --
forward, derivative, and reverse-sweep. The leaf visitors
(VisitDeclRefExpr, VisitIntegerLiteral, VisitFloatingLiteral) built all
three eagerly, cloning the underlying node once per representation. But
most parents read only a subset: a terminal product-rule leaf, having
emitted its own increment, never reads its adjoint, and a value whose
primal is not reconstructed never reads its forward. The unread clones
are allocated in the ASTContext and never parented -- wasted memory.

Give StmtDiff optional lazy representations: a slot may hold a deferred
clone of a source node, produced (and cached) only when that
representation is first read, written as LazyClone(N) in a constructor
argument. A representation no consumer reads is never cloned; two
representations off the same source still materialize distinct nodes, so
the no-shared-node invariant (findSharedNode) is preserved. Return the
reverse-mode DeclRef and literal leaves through LazyClone, and drop the
separate eager reverse-sweep copy, which defaults to the forward node.
The full lit suite is unchanged.

Author: vgvassilev

include/clad/Differentiator/VisitorBase.h

@@ -51,42 +51,114 @@ namespace clad {
   /// other (intermediate) statements, they are output to the current block.
   class StmtDiff {
   private:
-    std::array<clang::Stmt*, 2> data;
+    std::array<clang::Stmt*, 2> m_Data{};
     clang::Stmt* m_ValueForRevSweep;
 
+    // Lazy representations. When a source is set (via LazyClone() in a
+    // constructor argument), the matching slot is produced by cloning the
+    // source on first read and then cached, so a representation no consumer
+    // reads is never cloned (no orphaned "dead clone"), while two
+    // representations off the same source still materialize distinct nodes (no
+    // sharing). An unset source means the slot is eager -- the stored pointer
+    // (possibly null) is the value -- the default for every non-lazy argument.
+    const clang::Stmt* m_StmtSrc = nullptr;
+    const clang::Stmt* m_StmtDxSrc = nullptr;
+    const clang::Stmt* m_RevSweepSrc = nullptr;
+    utils::StmtClone* m_Cloner = nullptr;
+
+    // Clone Src into Slot on first read; a no-op when Src is null (eager slot).
+    clang::Stmt* materialize(clang::Stmt*& Slot, const clang::Stmt*& Src);
+
   public:
-    StmtDiff(clang::Stmt* orig = nullptr, clang::Stmt* diff = nullptr,
-             clang::Stmt* valueForRevSweep = nullptr)
-        : m_ValueForRevSweep(valueForRevSweep) {
-      data[1] = orig;
-      data[0] = diff;
+    /// A deferred clone marker: a representation cloned from Src on first read,
+    /// created by VisitorBase::LazyClone(). No member initializers, so In
+    /// (which embeds it) can be a `= {}` default argument while StmtDiff is
+    /// still being defined.
+    struct Lazy {
+      utils::StmtClone* Cloner;
+      const clang::Stmt* Src;
+    };
+
+    /// A constructor input for one representation: either an already-built node
+    /// (eager -- the node itself is the value) or a Lazy deferred clone. Node
+    /// and Deferred are set in the constructors rather than by default member
+    /// initializers, so In can be used as a `= {}` default argument here.
+    struct In {
+      clang::Stmt* Node;
+      Lazy Deferred;
+      In(clang::Stmt* N = nullptr) : Node(N), Deferred{nullptr, nullptr} {}
+      In(Lazy L) : Node(nullptr), Deferred(L) {}
+    };
+
+    /// Implicit single-representation constructor. Keeps the Expr*/Stmt* ->
+    /// StmtDiff conversion pervasive code relies on (return expr; sd = expr;),
+    /// which needs one user-defined conversion -- reaching the general
+    /// constructor below through In(Stmt*) would need two.
+    StmtDiff(clang::Stmt* orig) : m_ValueForRevSweep(nullptr) {
+      m_Data[1] = orig;
+      m_Data[0] = nullptr;
+    }
+
+    /// General constructor: each representation is eager (a node, the default
+    /// for every existing multi-argument call site) or lazy (LazyClone(src)).
+    /// Reads e.g. StmtDiff(fwd, LazyClone(dx)) or StmtDiff(LazyClone(fwd), dx).
+    StmtDiff(In orig = {}, In diff = {}, In valueForRevSweep = {})
+        : m_ValueForRevSweep(valueForRevSweep.Node),
+          m_StmtSrc(orig.Deferred.Src), m_StmtDxSrc(diff.Deferred.Src),
+          m_RevSweepSrc(valueForRevSweep.Deferred.Src),
+          // Every lazy slot shares the one cloner (VisitorBase::m_NodeCloner);
+          // take the first non-null.
+          m_Cloner([&] {
+            if (orig.Deferred.Cloner)
+              return orig.Deferred.Cloner;
+            if (diff.Deferred.Cloner)
+              return diff.Deferred.Cloner;
+            return valueForRevSweep.Deferred.Cloner;
+          }()) {
+      m_Data[1] = orig.Node;
+      m_Data[0] = diff.Node;
     }
 
-    clang::Stmt* getStmt() { return data[1]; }
-    clang::Stmt* getStmt_dx() { return data[0]; }
+    clang::Stmt* getStmt() { return materialize(m_Data[1], m_StmtSrc); }
+    clang::Stmt* getStmt_dx() { return materialize(m_Data[0], m_StmtDxSrc); }
     clang::Expr* getExpr() {
       return llvm::cast_or_null<clang::Expr>(getStmt());
     }
     clang::Expr* getExpr_dx() {
       return llvm::cast_or_null<clang::Expr>(getStmt_dx());
     }
 
-    void updateStmt(clang::Stmt* S) { data[1] = S; }
-    void updateStmtDx(clang::Stmt* S) { data[0] = S; }
-    void updateRevSweep(clang::Stmt* S) { m_ValueForRevSweep = S; }
+    void updateStmt(clang::Stmt* S) {
+      m_Data[1] = S;
+      m_StmtSrc = nullptr;
+    }
+    void updateStmtDx(clang::Stmt* S) {
+      m_Data[0] = S;
+      m_StmtDxSrc = nullptr;
+    }
+    void updateRevSweep(clang::Stmt* S) {
+      m_ValueForRevSweep = S;
+      m_RevSweepSrc = nullptr;
+    }
     // Stmt_dx goes first!
-    std::array<clang::Stmt*, 2>& getBothStmts() { return data; }
+    std::array<clang::Stmt*, 2>& getBothStmts() {
+      // A caller taking the array by reference parents both directions, so
+      // materialize both.
+      getStmt();
+      getStmt_dx();
+      return m_Data;
+    }
 
     clang::Expr* getRevSweepAsExpr() {
       return llvm::cast_or_null<clang::Expr>(getRevSweepStmt());
     }
 
     clang::Stmt* getRevSweepStmt() {
-      /// If there is no specific value for
-      /// the reverse sweep, use Stmt_dx.
-      if (!m_ValueForRevSweep)
-        return data[1];
-      return m_ValueForRevSweep;
+      if (clang::Stmt* R = materialize(m_ValueForRevSweep, m_RevSweepSrc))
+        return R;
+      // If there is no specific value for the reverse sweep, use the forward
+      // statement.
+      return getStmt();
     }
   };
 
@@ -703,6 +775,13 @@ namespace clad {
     clang::Expr* CloneNode(const clang::Expr* E);
     /// Statement overload of the structural copy above.
     clang::Stmt* CloneNode(const clang::Stmt* S);
+    /// A deferred CloneNode(\p N): the clone is produced only if the StmtDiff
+    /// representation it is stored in is actually read, so a representation no
+    /// consumer needs allocates no orphaned clone. Drop-in for CloneNode(N) in
+    /// a StmtDiff argument position.
+    StmtDiff::Lazy LazyClone(const clang::Stmt* N) {
+      return {m_Builder.m_NodeCloner.get(), N};
+    }
     /// Cloning types is necessary since VariableArrayType
     /// store a pointer to their size expression.
     clang::QualType CloneType(clang::QualType T);

lib/Differentiator/ReverseModeVisitor.cpp

@@ -1589,24 +1589,30 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
       // dfdx) statement.
       if (Expr* add_assign = BuildDiffIncrement(CloneNode(dExpr)))
         addToCurrentBlock(add_assign, direction::reverse);
-      // getExpr (forward rebuild) and the reverse-sweep value are consumed by
-      // different statements, so hand out distinct clones of the primal ref.
-      return StmtDiff(clonedDRE, CloneNode(dExpr), CloneNode(clonedDRE));
+      // The adjoint (getExpr_dx, e.g. *_d_x) is read only by parents that
+      // propagate a subexpression's derivative -- a conditional, a paren, an
+      // array-subscript base; a terminal product-rule leaf, having already
+      // emitted its own increment above, never reads it. Clone it lazily so
+      // those terminal leaves allocate no orphaned copy, while a consumer still
+      // materializes a distinct node (no sharing). The reverse-sweep value
+      // defaults to the forward node.
+      return StmtDiff(clonedDRE, LazyClone(dExpr));
     }
 
-    return StmtDiff(clonedDRE, /*diff=*/nullptr, CloneNode(clonedDRE));
+    return StmtDiff(clonedDRE);
   }
 
   StmtDiff ReverseModeVisitor::VisitIntegerLiteral(const IntegerLiteral* IL) {
     auto* Constant0 =
         ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, 0);
-    // Distinct forward and reverse-sweep copies so a parent consuming both
-    // does not share the literal node.
-    return StmtDiff(Clone(IL), Constant0, Clone(IL));
+    // The forward value (also the reverse-sweep default) is a copy of the
+    // literal, needed only if a parent reads it; clone it lazily so a leaf no
+    // parent reads allocates nothing. The derivative of a constant is 0.
+    return StmtDiff(LazyClone(IL), Constant0);
   }
 
   StmtDiff ReverseModeVisitor::VisitFloatingLiteral(const FloatingLiteral* FL) {
-    return StmtDiff(Clone(FL), getZeroInit(FL->getType()), Clone(FL));
+    return StmtDiff(LazyClone(FL), getZeroInit(FL->getType()));
   }
 
   static bool isNAT(QualType T) {

lib/Differentiator/VisitorBase.cpp

@@ -424,6 +424,17 @@ namespace clad {
     return S ? m_Builder.m_NodeCloner->Clone(S) : nullptr;
   }
 
+  clang::Stmt* StmtDiff::materialize(clang::Stmt*& Slot,
+                                     const clang::Stmt*& Src) {
+    if (!Slot && Src && m_Cloner) {
+      Slot = m_Cloner->Clone(Src);
+      // The representation is now cached; drop the recipe so later reads (and
+      // identity/null checks) return the same node.
+      Src = nullptr;
+    }
+    return Slot;
+  }
+
   QualType VisitorBase::CloneType(const QualType QT) {
     auto clonedType = m_Builder.m_NodeCloner->CloneType(QT);
     utils::ReferencesUpdater up(m_Sema, getCurrentScope(), m_DiffReq.Function,