Encode reverse-mode early returns with a named lambda

The reverse-mode early-return encoding emitted a label inside the master
reverse and a goto into it from each early-return point. That violates
[stmt.dcl]/2 whenever a local with a non-trivial initializer or
destructor sits between the goto and its target, which is why
DifferentiateWithClad had to bypass Sema::ActOnFinishFunctionBody on the
generated function.

Materialise the master reverse as a `[&]` lambda bound to an `auto`
VarDecl instead, called at each early return and at the tail. The body
is assembled outside the closure scope, so its references to captured
variables are rebuilt in scope (rebuildLambdaCaptures) where the `[&]`
default registers the capture; the reverse statements are emitted into
the closure directly. An early return is a NullStmt marker in the
forward block, patched with `{ _rev(); return; }` at finalization so no
subtree gains a second parent.

Whether a function has a non-tail return is a fact about the primal, so
DiffRequest carries it as HasEarlyReturns; the encoder reads it to pick
the lambda vs plain shape up front. A branch condition read by the
master reverse is zero-initialized so an early return that skips its
forward store leaves it false -- the same zero-initialized additive
state clad already relies on for adjoints and loop counters.

Fixes #367

Author: vgvassilev

include/clad/Differentiator/DiffPlanner.h

@@ -42,7 +42,11 @@ using OwnedAnalysisContexts =
     llvm::SmallVector<std::unique_ptr<clang::AnalysisDeclContext>, 4>;
 using ParamSet = std::set<const clang::ParmVarDecl*>;
 using ParamInfo = std::map<const clang::FunctionDecl*, ParamSet>;
-/// A struct containing information about request to differentiate a function.
+/// A read-only, AD-oriented view over the primal being differentiated: it
+/// wraps the primal FunctionDecl and surfaces the AD-relevant facts the
+/// FunctionDecl itself does not. Recording such facts here, rather than
+/// rediscovering them inside a visitor, keeps them available to every visitor
+/// and correct after a request is copied and re-pointed at another Function.
 struct DiffRequest {
 private:
   /// Based on To-Be-Recorded analysis performed before differentiation, tells
@@ -66,7 +70,28 @@ struct DiffRequest {
     bool HasAnalysisRun = false;
   } m_UsefulRunInfo;
 
+  /// Cache for hasEarlyReturns(): whether the primal body has a return that is
+  /// not in tail position. A property of the Function, computed once on demand.
+  mutable struct EarlyReturnInfo {
+    bool HasEarlyReturns = false;
+    bool HasAnalysisRun = false;
+  } m_EarlyReturnInfo;
+
 public:
+  /// The primal body's tail-position return -- the one an early-return encoder
+  /// lets control fall through to, as opposed to an early return that needs a
+  /// jump/call. Null when the body does not end in a return. O(1): reads
+  /// Function's body directly, so a copied request re-pointed at a new Function
+  /// (a lambda's operator(), a pullback callee) answers for its own Function.
+  const clang::ReturnStmt* getTailReturn() const;
+
+  /// Whether the primal body has a return that is not the tail return -- one
+  /// the reverse mode must encode with the early-return lambda. A fact about
+  /// the primal, read directly off Function's body (returns inside nested
+  /// lambdas belong to their own function and are skipped). A void function
+  /// seeds no return value, so its returns skip the encoding.
+  bool hasEarlyReturns() const;
+
   /// Function to be differentiated.
   const clang::FunctionDecl* Function = nullptr;
   /// Name of the base function to be differentiated. Can be different from
@@ -243,17 +268,17 @@ struct DiffRequest {
   bool HasTbrAnalysisRun() const { return m_TbrRunInfo.HasAnalysisRun; }
 };
 
-  using DiffInterval = std::vector<clang::SourceRange>;
+using DiffInterval = std::vector<clang::SourceRange>;
 
-  // FIXME: These are translation-unit-wide defaults taken from the compiler
-  // invocation, not the options of a request; rename to InvocationOptions.
-  struct RequestOptions {
-    /// This is a flag to indicate the default behaviour to enable/disable
-    /// TBR analysis during reverse-mode differentiation.
-    bool EnableTBRAnalysis = false;
-    bool EnableVariedAnalysis = false;
-    bool EnableUsefulAnalysis = false;
-  };
+// FIXME: These are translation-unit-wide defaults taken from the compiler
+// invocation, not the options of a request; rename to InvocationOptions.
+struct RequestOptions {
+  /// This is a flag to indicate the default behaviour to enable/disable
+  /// TBR analysis during reverse-mode differentiation.
+  bool EnableTBRAnalysis = false;
+  bool EnableVariedAnalysis = false;
+  bool EnableUsefulAnalysis = false;
+};
 
   class DiffCollector: public clang::RecursiveASTVisitor<DiffCollector> {
     /// The source interval where clad was activated.
@@ -298,6 +323,10 @@ struct DiffRequest {
     /// or constructor initializers. If we use Visit they would be processed
     /// under the parent DiffRequest which is not in the lambda scope.
     bool TraverseLambdaExpr(clang::LambdaExpr* LE);
+    /// Plan a lazily-scheduled nested request the static TU walk never reaches
+    /// (built in DerivativeBuilder::HandleNestedDiffRequest). Currently records
+    /// its early-return flag by walking the request's own body.
+    bool PlanNestedRequest(DiffRequest& request);
     bool TraverseFunctionDeclOnce(const clang::FunctionDecl* FD) {
       llvm::SaveAndRestore<bool> Saved(m_IsTraversingTopLevelDecl, false);
       if (m_Traversed.count(FD))

include/clad/Differentiator/DiffScheduler.h

@@ -38,6 +38,10 @@ class DiffScheduler {
 
   /// Static planning pass over a group of top-level declarations.
   void Plan(clang::DeclGroupRef DGR) { m_Collector.Walk(DGR); }
+
+  /// Plan a single lazily-scheduled request (a nested or higher-order
+  /// derivative) that the static walk never reached.
+  void Plan(DiffRequest& R) { m_Collector.PlanNestedRequest(R); }
 };
 
 } // namespace clad

include/clad/Differentiator/ReverseModeVisitor.h

@@ -29,6 +29,7 @@
 #include "clang/Sema/Sema.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <array>
@@ -75,6 +76,10 @@ namespace clad {
     /// the reverse mode we also accumulate Stmts for the reverse pass which
     /// will be executed on return.
     std::vector<Stmts> m_Reverse;
+    /// Markers emitted at early-return sites in the forward block, patched
+    /// with `{ _rev(); return; }` at finalization, where `_rev` is the lambda
+    /// wrapping the master reverse sweep.
+    llvm::SmallPtrSet<clang::Stmt*, 4> m_EarlyReturnMarkers;
     /// Storing expressions to delete/free memory in the reverse pass.
     Stmts m_DeallocExprs;
     /// A dfdx seed: eager expression, or a deferred builder materialized on
@@ -129,6 +134,15 @@ namespace clad {
     // Function to Differentiate with Enzyme as Backend
     void DifferentiateWithEnzyme();
 
+    /// Walk the current forward block and replace every Stmt in
+    /// m_EarlyReturnMarkers with a fresh Stmt built by `MakeReplacement`.
+    /// Each marker gets its own node — sharing one replacement across sites
+    /// would give it multiple parents and break the single-parent AST
+    /// invariant. Called once at function-body finalization, after the
+    /// forward sweep is fully assembled and before the body is handed to Sema.
+    void
+    patchEarlyReturnMarkers(llvm::function_ref<clang::Stmt*()> MakeReplacement);
+
   public:
     using direction = rmv::direction;
     virtual clang::Expr* dfdx() {
@@ -327,10 +341,10 @@ namespace clad {
     /// of the stack (clad::back(S)).
     clang::Expr* GlobalStoreAndRef(clang::Expr* E, clang::QualType Type,
                                    llvm::StringRef prefix = "_t",
-                                   bool force = false);
+                                   bool force = false, bool zeroInit = false);
     clang::Expr* GlobalStoreAndRef(clang::Expr* E,
                                    llvm::StringRef prefix = "_t",
-                                   bool force = false);
+                                   bool force = false, bool zeroInit = false);
     virtual StmtDiff StoreAndRestore(clang::Expr* E,
                                      llvm::StringRef prefix = "_t",
                                      bool moveToTape = false);

include/clad/Differentiator/VisitorBase.h

@@ -18,6 +18,7 @@
 #include "clang/AST/StmtVisitor.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/Lambda.h"
 #include "clang/Basic/OperatorKinds.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Sema/DeclSpec.h"
@@ -266,23 +267,56 @@ namespace clad {
     /// The currently visited statement. Useful for crash pretty-printing.
     const clang::Stmt* m_CurVisitedStmt = nullptr;
 
+    /// Resolves the captures of a lambda synthesized from a body that was
+    /// built outside the closure scope, so Sema never saw the uses. collect()
+    /// finds the body's free variables -- locals and parameters it references
+    /// but does not itself declare. resolve() re-creates those references in
+    /// the current scope, where BuildDeclRef under the `[&]` default registers
+    /// the capture and marks the reference for codegen; it must run with the
+    /// lambda scope active. contains() exposes the set for a caller that must
+    /// place the captured decls before the lambda.
+    class LambdaCaptures {
+      VisitorBase& m_V;
+      llvm::SmallPtrSet<clang::VarDecl*, 8> m_Captures;
+
+    public:
+      explicit LambdaCaptures(VisitorBase& V) : m_V(V) {}
+      void collect(llvm::ArrayRef<clang::Stmt*> Body);
+      /// A `[&]` capture binds a variable at the lambda's definition point, so
+      /// every captured decl must precede the lambda. \p Prefix and \p Suffix
+      /// are the forward block split at the lambda's insertion point; move each
+      /// captured DeclStmt from Suffix to the end of Prefix when its
+      /// initializer references only names already live there -- the function's
+      /// parameters, \p AlreadyLive (decls emitted earlier), and Prefix.
+      void orderCaptureDecls(llvm::SmallVectorImpl<clang::Stmt*>& Prefix,
+                             llvm::SmallVectorImpl<clang::Stmt*>& Suffix,
+                             llvm::ArrayRef<clang::Stmt*> AlreadyLive);
+      void resolve(llvm::ArrayRef<clang::Stmt*> Body);
+      bool contains(clang::VarDecl* VD) const { return m_Captures.count(VD); }
+    };
+
     /// Build a lambda whose body is produced by `func`. Returns the
     /// LambdaExpr without invoking it, so the caller can bind it to a VarDecl
     /// and call it from multiple sites. `func` is invoked inside the lambda's
     /// scope and block; statements are expected to be added via
     /// addToCurrentBlock from func's invocation.
+    ///
+    /// The lambda uses the `[&]` capture-default; Sema resolves captures from
+    /// the body's ODR-uses. A pre-built body whose DeclRefExprs were made
+    /// outside this scope must have those references rebuilt in scope so Sema
+    /// sees the uses (see LambdaCaptures::resolve).
     // FIXME: This will become problematic when we try to support C.
     template <typename F>
     static clang::Expr* buildLambda(VisitorBase& V, clang::Sema& S,
-                                    const clang::Expr* E, F&& func) {
+                                    const clang::Stmt* LocSrc, F&& func) {
       // FIXME: Here we use some of the things that are used from Parser, it
       // seems to be the easiest way to create lambda.
       clang::LambdaIntroducer Intro;
       Intro.Default = clang::LCD_ByRef;
       // FIXME: Using noLoc here results in assert failure. Any other valid
       // SourceLocation seems to work fine.
-      Intro.Range.setBegin(E->getBeginLoc());
-      Intro.Range.setEnd(E->getEndLoc());
+      Intro.Range.setBegin(LocSrc->getBeginLoc());
+      Intro.Range.setEnd(LocSrc->getEndLoc());
       clang::AttributeFactory AttrFactory;
       const clang::DeclSpec DS(AttrFactory);
       clang::Declarator D(
@@ -329,8 +363,8 @@ namespace clad {
     /// added by addToCurrentBlock from func invocation.
     template <typename F>
     static clang::Expr* wrapInLambda(VisitorBase& V, clang::Sema& S,
-                                     const clang::Expr* E, F&& func) {
-      clang::Expr* lambda = buildLambda(V, S, E, std::forward<F>(func));
+                                     const clang::Stmt* LocSrc, F&& func) {
+      clang::Expr* lambda = buildLambda(V, S, LocSrc, std::forward<F>(func));
       return S.ActOnCallExpr(V.getCurrentScope(), lambda, noLoc, {}, noLoc)
           .get();
     }
@@ -341,13 +375,29 @@ namespace clad {
     /// more sites via DeclRefExpr + ActOnCallExpr. Use this when the same
     /// lambda body must be invoked from multiple paths (e.g. a reverse-pass
     /// segment shared between an early-return path and the natural tail).
+    ///
+    /// The binding uses `auto` deduction so the pretty-printer renders it as
+    /// `auto X = [&] {...};` rather than the closure type's unspellable
+    /// `(lambda at ...)` form. Sema deduces the concrete closure type from
+    /// the initializer; the TypeSourceInfo retains the `auto` keyword.
+    ///
+    /// \p func emits the closure body; \p Captures then resolves that body's
+    /// references to enclosing variables (its collect() must have already run),
+    /// so callers hand over a pure body-emission callback.
     template <typename F>
-    clang::VarDecl* buildAndBindLambda(const clang::Expr* LocE,
-                                       llvm::StringRef NameHint, F&& func) {
-      clang::Expr* lambda =
-          buildLambda(*this, m_Sema, LocE, std::forward<F>(func));
+    clang::VarDecl* buildAndBindLambda(const clang::Stmt* LocSrc,
+                                       llvm::StringRef NameHint,
+                                       LambdaCaptures& Captures, F&& func) {
+      clang::Expr* lambda = buildLambda(*this, m_Sema, LocSrc, [&] {
+        std::forward<F>(func)();
+        // Resolve captures while the closure scope is active and its body is
+        // the current block.
+        Captures.resolve(getCurrentBlock());
+      });
       clang::IdentifierInfo* II = CreateUniqueIdentifier(NameHint);
-      return BuildVarDecl(lambda->getType(), II, lambda);
+      clang::QualType AutoTy = m_Context.getAutoDeductType();
+      clang::TypeSourceInfo* TSI = m_Context.getTrivialTypeSourceInfo(AutoTy);
+      return BuildVarDecl(AutoTy, II, lambda, /*DirectInit=*/false, TSI);
     }
 
     /// For a qualtype QT returns if it's type is Array or Pointer Type

lib/Differentiator/DerivativeBuilder.cpp

@@ -397,9 +397,11 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
 
   clang::FunctionDecl*
   DerivativeBuilder::HandleNestedDiffRequest(DiffRequest& request) {
-    // FIXME: Find a way to do this without accessing plugin namespace functions
     bool alreadyDerived = true;
     request.UpdateDiffParamsInfo(m_Sema);
+    // Plan this lazily-scheduled request statically, so it carries the planning
+    // info (currently the early-return flag) the static TU walk never produced.
+    m_Scheduler.Plan(request);
     FunctionDecl* derivative = this->FindDerivedFunction(request);
     if (!derivative) {
       alreadyDerived = false;

lib/Differentiator/DiffPlanner.cpp

@@ -348,6 +348,24 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     return true;
   }
 
+  bool DiffCollector::PlanNestedRequest(DiffRequest& request) {
+    // Lazily-scheduled requests (pushforward/pullback/higher-order, built in
+    // DerivativeBuilder::HandleNestedDiffRequest) never pass through the static
+    // TU walk, so their planning must be done here on demand. For now this only
+    // records the early-return flag: m_TopMostReq is left null, so
+    // VisitCallExpr and VisitDeclRefExpr bail and no sub-requests are spawned
+    // -- only VisitReturnStmt runs over this request's body.
+    // FIXME: This is where a nested request should also get its m_AnalysisDC
+    // and the analyses currently force-disabled in HandleNestedDiffRequest.
+    const FunctionDecl* Def =
+        request.Function ? request.Function->getDefinition() : nullptr;
+    if (!Def || !Def->hasBody())
+      return true;
+    llvm::SaveAndRestore<DiffRequest*> Saved(m_ParentReq, &request);
+    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
+    return TraverseStmt(const_cast<Stmt*>(Def->getBody()));
+  }
+
   bool DiffCollector::isInInterval(SourceLocation Loc) const {
     const SourceManager &SM = m_Sema.getSourceManager();
     for (size_t i = 0, e = m_Interval.size(); i < e; ++i) {
@@ -364,6 +382,43 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     return false;
   }
 
+  const ReturnStmt* DiffRequest::getTailReturn() const {
+    const FunctionDecl* Def = Function ? Function->getDefinition() : nullptr;
+    if (!Def || !Def->hasBody())
+      return nullptr;
+    const Stmt* Body = Def->getBody();
+    // The tail return is the body's last statement, when that is a return.
+    if (const auto* CS = dyn_cast<CompoundStmt>(Body))
+      return CS->body_empty() ? nullptr
+                              : dyn_cast<ReturnStmt>(*CS->body_rbegin());
+    return dyn_cast<ReturnStmt>(Body);
+  }
+
+  bool DiffRequest::hasEarlyReturns() const {
+    if (m_EarlyReturnInfo.HasAnalysisRun)
+      return m_EarlyReturnInfo.HasEarlyReturns;
+    const FunctionDecl* Def = Function ? Function->getDefinition() : nullptr;
+    if (!Def || !Def->hasBody() || Def->getReturnType()->isVoidType())
+      return false;
+    // An early-return body has a return that is not the tail return. Returns
+    // inside a nested lambda belong to that lambda's own function, so do not
+    // descend into one.
+    struct Finder : RecursiveASTVisitor<Finder> {
+      const ReturnStmt* Tail;
+      bool Found = false;
+      bool TraverseLambdaExpr(LambdaExpr*) { return true; }
+      bool VisitReturnStmt(ReturnStmt* RS) {
+        if (RS != Tail)
+          Found = true;
+        return !Found; // stop at the first early return
+      }
+    } F;
+    F.Tail = getTailReturn();
+    F.TraverseStmt(const_cast<Stmt*>(cast<Stmt>(Def->getBody())));
+    m_EarlyReturnInfo = {F.Found, /*HasAnalysisRun=*/true};
+    return F.Found;
+  }
+
   void DiffRequest::UpdateDiffParamsInfo(Sema& semaRef) {
     // Diff info for pullbacks is generated automatically,
     // its parameters are not provided by the user.
@@ -1504,7 +1559,10 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
   }
 
   bool DiffCollector::VisitDeclRefExpr(DeclRefExpr* DRE) {
-    if (!m_ParentReq)
+    // m_TopMostReq is dereferenced below; it is null when PlanNestedRequest
+    // walks a lazy request's body just to record its early-return flag, and no
+    // global-adjoint discovery is wanted there.
+    if (!m_ParentReq || !m_TopMostReq)
       return true;
     // FIXME: Add support for globals in other modes.
     if (m_ParentReq->Mode != DiffMode::reverse &&