Reverse switches on the stored condition instead of a control-flow tape

include/clad/Differentiator/ReverseModeVisitor.h

@@ -703,17 +703,18 @@ namespace clad {
 
     StmtDiff DifferentiateCanonicalLoop(const clang::ForStmt* S);
 
-    /// This class modifies forward and reverse blocks of the loop/switch
-    /// body so that `break` and `continue` statements are correctly
-    /// handled. `break` and `continue` statements are handled by
-    /// enclosing entire reverse block loop body in a switch statement
-    /// and only executing the statements, with the help of case labels,
-    /// that were executed in the associated forward iteration. This is
-    /// determined by keeping track of which `break`/`continue` statement
-    /// was hit in which iteration and that in turn helps to determine which
-    /// case label should be selected.
-    ///
-    /// Class usage:
+    /// Handles `break`/`continue` inside a differentiated loop. It owns a
+    /// control-flow tape recording which one fired in which iteration, so the
+    /// reverse loop body -- wrapped in a switch over that tape -- replays
+    /// exactly the statements the forward iteration executed. The members below
+    /// serve this tape.
+    ///
+    /// One handler is pushed per enclosing loop or switch, so the top of the
+    /// stack is the innermost construct a `break` binds to. A switch sets
+    /// m_IsInvokedBySwitchStmt and leaves the tape unused -- its reverse is
+    /// rebuilt from the stored condition instead (VisitSwitchStmt).
+    ///
+    /// Loop usage:
     ///
     /// ```cpp
     /// auto activeBreakContStmtHandler = PushBreakContStmtHandler();
@@ -760,6 +761,9 @@ namespace clad {
       clang::Expr* CreateCFTapePushExpr(std::size_t value);
 
     public:
+      /// True when the innermost breakable construct is a source-level switch,
+      /// whose `break`s are handled in VisitSwitchStmt rather than by the
+      /// control-flow tape. Lets VisitBreakStmt pick the right handling.
       bool m_IsInvokedBySwitchStmt = false;
 
       BreakContStmtHandler(ReverseModeVisitor& RMV, bool forSwitchStmt = false)
@@ -835,9 +839,18 @@ namespace clad {
 
     /// Stores data required for differentiating a switch statement.
     struct SwitchStmtInfo {
+      /// The forward-pass case/default labels, in source order.
       llvm::SmallVector<clang::SwitchCase*, 16> cases;
+      /// The stored switch condition (`_cond`), reused both as the reverse
+      /// switch discriminator and by every case guard.
       clang::Expr* switchStmtCond = nullptr;
       clang::IfStmt* defaultIfBreakExpr = nullptr;
+      /// Index into `cases` of the first label of the fall-through group not
+      /// yet closed by a `break`.
+      std::size_t groupStart = 0;
+      /// The reverse switch's entry labels, built from the original case values
+      /// rather than a control-flow tape counter.
+      llvm::SmallVector<clang::SwitchCase*, 16> reverseEntryCases;
     };
 
     /// Maintains a stack of `SwitchStmtInfo`.
@@ -854,6 +867,12 @@ namespace clad {
 
     void PopSwitchStmtInfo() { m_SwitchStmtsData.pop_back(); }
 
+    /// Closes the currently open fall-through group `cases[groupStart..)`:
+    /// builds its reverse-switch entry from the group's original case labels,
+    /// registers them in `reverseEntryCases`, advances `groupStart`, and
+    /// returns the label chain to prepend before the group's adjoint replay.
+    clang::Stmt* CloseReverseSwitchCaseGroup(SwitchStmtInfo& SSData);
+
   private:
     // When differentiating ArrayInitLoopExpr, we need to replace
     // ArrayInitIndexExpr with real indices. We need to both add and pop them in

lib/Differentiator/ReverseModeVisitor.cpp

@@ -4412,6 +4412,8 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     // statement body will be processed in both the forward and the reverse
     // pass. Thus, we do not need to add them in the differentiated function.
     if (!(SSData->cases.empty())) {
+      // The forward sweep is a clone of the original switch: control flow is
+      // recorded implicitly by the stored condition, so no tape is needed.
       Sema::ConditionResult condRes =
           m_Sema.ActOnCondition(getCurrentScope(), noLoc, CloneNode(condExpr),
                                 Sema::ConditionKind::Switch);
@@ -4421,18 +4423,40 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
                                       /*LParenLoc=*/noLoc, nullptr, condRes,
                                       /*RParenLoc=*/noLoc)
               .getAs<SwitchStmt>();
-      activeBreakContHandler->UpdateForwAndRevBlocks(bodyDiff);
-
-      // Registers all the cases to the switch statement.
       for (auto* SC : SSData->cases)
         forwardSS->addSwitchCase(SC);
-
       forwardSS =
           m_Sema.ActOnFinishSwitchStmt(noLoc, forwardSS, bodyDiff.getStmt())
               .getAs<SwitchStmt>();
 
+      // The reverse sweep re-switches on the stored condition. Each
+      // fall-through group's adjoint replay is entered through the original
+      // case values (the per-case `if (v == _cond) break` guards emitted by
+      // VisitCaseStmt peel off the cases that did not run). The trailing group
+      // is closed by the switch end rather than a break, so label it here; it
+      // is the topmost group in the bottom-up reverse block.
+      Stmt* revBody = bodyDiff.getStmt_dx();
+      if (SSData->groupStart < SSData->cases.size()) {
+        Stmt* finalEntry = CloseReverseSwitchCaseGroup(*SSData);
+        revBody =
+            utils::PrependAndCreateCompoundStmt(m_Context, revBody, finalEntry);
+      }
+      Sema::ConditionResult revCondRes =
+          m_Sema.ActOnCondition(getCurrentScope(), noLoc, CloneNode(condExpr),
+                                Sema::ConditionKind::Switch);
+      SwitchStmt* reverseSS =
+          m_Sema
+              .ActOnStartOfSwitchStmt(/*SwitchLoc=*/noLoc,
+                                      /*LParenLoc=*/noLoc, nullptr, revCondRes,
+                                      /*RParenLoc=*/noLoc)
+              .getAs<SwitchStmt>();
+      for (auto* SC : SSData->reverseEntryCases)
+        reverseSS->addSwitchCase(SC);
+      reverseSS = m_Sema.ActOnFinishSwitchStmt(noLoc, reverseSS, revBody)
+                      .getAs<SwitchStmt>();
+
       addToCurrentBlock(forwardSS, direction::forward);
-      addToCurrentBlock(bodyDiff.getStmt_dx(), direction::reverse);
+      addToCurrentBlock(reverseSS, direction::reverse);
     }
 
     PopBreakContStmtHandler();
@@ -4492,6 +4516,36 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     return {endBlock(direction::forward), endBlock(direction::reverse)};
   }
 
+  Stmt*
+  ReverseModeVisitor::CloseReverseSwitchCaseGroup(SwitchStmtInfo& SSData) {
+    // Build `case v1: case v2: ... [default:] ;` for the still-open group's
+    // original labels, innermost first. The order of the labels is irrelevant
+    // (they all fall into the same reverse replay); the shared null
+    // substatement lets the group's adjoints follow as siblings in the switch
+    // body.
+    Stmt* inner = m_Sema.ActOnNullStmt(noLoc).get();
+    for (std::size_t i = SSData.groupStart, e = SSData.cases.size(); i != e;
+         ++i) {
+      SwitchCase* rev = nullptr;
+      if (isa<DefaultStmt>(SSData.cases[i])) {
+        rev = new (m_Context) DefaultStmt(noLoc, noLoc, inner);
+      } else {
+        auto* fwd = cast<CaseStmt>(SSData.cases[i]);
+        // Clone the range high end too, matching the forward case, so a GNU
+        // `case a ... b:` keeps its extent in the reverse switch.
+        Expr* rhs = fwd->getRHS() ? CloneNode(fwd->getRHS()) : nullptr;
+        auto* caseStmt = CaseStmt::Create(m_Context, CloneNode(fwd->getLHS()),
+                                          rhs, noLoc, noLoc, noLoc);
+        caseStmt->setSubStmt(inner);
+        rev = caseStmt;
+      }
+      SSData.reverseEntryCases.push_back(rev);
+      inner = rev;
+    }
+    SSData.groupStart = SSData.cases.size();
+    return inner;
+  }
+
   static bool hasCheckpointingPragma(ASTContext& C, SourceLocation loopLoc,
                                      const DiffRequest& request) {
     if (!loopLoc.isValid())
@@ -4640,10 +4694,18 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     Stmt* newBS =
         clad_compat::ActOnBreakStmt(m_Sema, noLoc, getCurrentScope()).get();
     auto* activeBreakContHandler = GetActiveBreakContStmtHandler();
+    // A break in a source-level switch only closes the current fall-through
+    // group (VisitSwitchStmt turns each group into a reverse-switch entry).
+    // Unlike the loop path below, it records nothing on a control-flow tape.
+    if (activeBreakContHandler->m_IsInvokedBySwitchStmt) {
+      addToCurrentBlock(newBS);
+      Stmt* revEntry = CloseReverseSwitchCaseGroup(*GetActiveSwitchStmtInfo());
+      return {endBlock(direction::forward), revEntry};
+    }
     Stmt* CFCaseStmt = activeBreakContHandler->GetNextCFCaseStmt();
     Stmt* pushExprToCurrentCase = activeBreakContHandler
                                       ->CreateCFTapePushExprToCurrentCase();
-    if (isInsideLoop && !activeBreakContHandler->m_IsInvokedBySwitchStmt) {
+    if (isInsideLoop) {
       Expr* tapeBackExprForCurrentCase =
           activeBreakContHandler->CreateCFTapeBackExprForCurrentCase();
       if (m_CurrentBreakFlagExpr) {
@@ -4733,7 +4795,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
 
   void ReverseModeVisitor::BreakContStmtHandler::UpdateForwAndRevBlocks(
       StmtDiff& bodyDiff) {
-    if (m_SwitchCases.empty() && !m_IsInvokedBySwitchStmt)
+    // Only loops reach here; a loop with no break/continue needs no
+    // control-flow switch in its reverse body.
+    if (m_SwitchCases.empty())
       return;
 
     // Add case statement in the beginning of the reverse block