Guard statement scopes with a ScopeRAII handle. NFC

clad's statement visitors open a clang::Scope with beginScope and have to
close it with a matching endScope on every exit path. A forgotten endScope
leaks the Scope object and strands the function-spanning scopes opened by
Derive(); this already bit the forward range-for and reverse switch
handlers, each since repaired by adding the missing endScope by hand.

Introduce a ScopeRAII guard that opens a scope on construction and closes
it on destruction, and adopt it in the statement handlers whose scope both
opens and closes within a single function body. The balance is now enforced
by the type rather than by hand, so an early return can no longer strand a
scope. The scopes that do not coincide with a single C++ block stay as
explicit begin/endScope pairs: Derive()'s function-spanning function and
body scopes, the derived-lambda construction scopes, the forward switch
body scope that closes in DeriveSwitchStmtBodyHelper, and the OpenMP
canonical-loop scope.

Author: vgvassilev

include/clad/Differentiator/VisitorBase.h

@@ -251,6 +251,25 @@ namespace clad {
     void beginScope(unsigned ScopeFlags);
     void endScope();
 
+    /// RAII guard that opens a scope on construction and closes it on
+    /// destruction, so a beginScope is always balanced by an endScope even
+    /// when a statement handler returns early. Use it for statement scopes
+    /// that nest within a single function; Derive()'s function-spanning scope
+    /// stays explicit because it is captured into m_DerivativeFnScope and
+    /// interleaves with Push/PopDeclContext.
+    class [[nodiscard]] ScopeRAII {
+      VisitorBase& m_Visitor;
+
+    public:
+      ScopeRAII(VisitorBase& Visitor, unsigned ScopeFlags)
+          : m_Visitor(Visitor) {
+        m_Visitor.beginScope(ScopeFlags);
+      }
+      ~ScopeRAII() { m_Visitor.endScope(); }
+      ScopeRAII(const ScopeRAII&) = delete;
+      ScopeRAII& operator=(const ScopeRAII&) = delete;
+    };
+
     /// A shorthand to simplify syntax for creation of new expressions.
     /// This function uses m_Sema.BuildUnOp internally to build unary
     /// operations. Typical usage of this function looks like the following:

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -466,15 +466,14 @@ StmtDiff BaseForwardModeVisitor::VisitStmt(const Stmt* S) {
 }
 
 StmtDiff BaseForwardModeVisitor::VisitCompoundStmt(const CompoundStmt* CS) {
-  beginScope(Scope::DeclScope);
+  ScopeRAII compoundScope(*this, Scope::DeclScope);
   beginBlock();
   for (Stmt* S : CS->body()) {
     StmtDiff SDiff = Visit(S);
     addToCurrentBlock(SDiff.getStmt_dx());
     addToCurrentBlock(SDiff.getStmt());
   }
   CompoundStmt* Result = endBlock();
-  endScope();
   // Differentation of CompundStmt produces another CompoundStmt with both
   // original and derived statements, i.e. Stmt() is Result and Stmt_dx() is
   // null.
@@ -484,7 +483,7 @@ StmtDiff BaseForwardModeVisitor::VisitCompoundStmt(const CompoundStmt* CS) {
 StmtDiff BaseForwardModeVisitor::VisitIfStmt(const IfStmt* If) {
   // Control scope of the IfStmt. E.g., in if (double x = ...) {...}, x goes
   // to this scope.
-  beginScope(Scope::DeclScope | Scope::ControlScope);
+  ScopeRAII ifScope(*this, Scope::DeclScope | Scope::ControlScope);
   // Create a block "around" if statement, e.g:
   // {
   //   ...
@@ -532,12 +531,11 @@ StmtDiff BaseForwardModeVisitor::VisitIfStmt(const IfStmt* If) {
       return BranchDiff.getStmt();
     } else {
       beginBlock();
-      beginScope(Scope::DeclScope);
+      ScopeRAII branchScope(*this, Scope::DeclScope);
       StmtDiff BranchDiff = Visit(Branch);
       for (Stmt* S : BranchDiff.getBothStmts())
         addToCurrentBlock(S);
       CompoundStmt* Block = endBlock();
-      endScope();
       if (Block->size() == 1)
         return Block->body_front();
       else
@@ -554,7 +552,6 @@ StmtDiff BaseForwardModeVisitor::VisitIfStmt(const IfStmt* If) {
   addToCurrentBlock(ifDiff);
   CompoundStmt* Block = endBlock();
   // If IfStmt is the only statement in the block, remove the block:
-  endScope();
   // {
   //   if (...) {...}
   // }
@@ -596,8 +593,8 @@ StmtDiff BaseForwardModeVisitor::VisitConditionalOperator(
 
 StmtDiff
 BaseForwardModeVisitor::VisitCXXForRangeStmt(const CXXForRangeStmt* FRS) {
-  beginScope(Scope::DeclScope | Scope::ControlScope | Scope::BreakScope |
-             Scope::ContinueScope);
+  ScopeRAII rangeScope(*this, Scope::DeclScope | Scope::ControlScope |
+                                  Scope::BreakScope | Scope::ContinueScope);
   // Visiting for range-based ststement produces __range1, __begin1 and __end1
   // variables, so for(auto i: a){
   //      ...
@@ -656,13 +653,12 @@ BaseForwardModeVisitor::VisitCXXForRangeStmt(const CXXForRangeStmt* FRS) {
   Stmt* forStmtDiff = new (m_Context)
       ForStmt(m_Context, nullptr, cond, /*condVar=*/nullptr, Inc, bodyResult,
               FRS->getForLoc(), FRS->getBeginLoc(), FRS->getEndLoc());
-  endScope();
   return StmtDiff(forStmtDiff);
 }
 
 StmtDiff BaseForwardModeVisitor::VisitForStmt(const ForStmt* FS) {
-  beginScope(Scope::DeclScope | Scope::ControlScope | Scope::BreakScope |
-             Scope::ContinueScope);
+  ScopeRAII forScope(*this, Scope::DeclScope | Scope::ControlScope |
+                                Scope::BreakScope | Scope::ContinueScope);
   beginBlock();
   const Stmt* init = FS->getInit();
   StmtDiff initDiff = init ? Visit(init) : StmtDiff{};
@@ -760,22 +756,22 @@ StmtDiff BaseForwardModeVisitor::VisitForStmt(const ForStmt* FS) {
 
   // Build the derived for loop body.
   const Stmt* body = FS->getBody();
-  beginScope(Scope::DeclScope);
   Stmt* bodyResult = nullptr;
-  beginBlock();
-  StmtDiff bodyVisited = Visit(body);
-  for (Stmt* S : bodyVisited.getBothStmts())
-    addToCurrentBlock(S);
-  bodyResult = utils::unwrapIfSingleStmt(endBlock());
-  endScope();
+  {
+    ScopeRAII bodyScope(*this, Scope::DeclScope);
+    beginBlock();
+    StmtDiff bodyVisited = Visit(body);
+    for (Stmt* S : bodyVisited.getBothStmts())
+      addToCurrentBlock(S);
+    bodyResult = utils::unwrapIfSingleStmt(endBlock());
+  }
 
   Stmt* forStmtDiff = new (m_Context)
       ForStmt(m_Context, initDiff.getStmt(), cond, /*condVar=*/nullptr,
               incResult, bodyResult, noLoc, noLoc, noLoc);
 
   addToCurrentBlock(forStmtDiff);
   CompoundStmt* Block = endBlock();
-  endScope();
 
   StmtDiff Result =
       (Block->size() == 1) ? StmtDiff(forStmtDiff) : StmtDiff(Block);
@@ -1823,9 +1819,9 @@ StmtDiff BaseForwardModeVisitor::VisitStringLiteral(const StringLiteral* SL) {
 }
 
 StmtDiff BaseForwardModeVisitor::VisitWhileStmt(const WhileStmt* WS) {
-  // begin scope for while loop
-  beginScope(Scope::ContinueScope | Scope::BreakScope | Scope::DeclScope |
-             Scope::ControlScope);
+  // Scope for the whole while loop.
+  ScopeRAII whileScope(*this, Scope::ContinueScope | Scope::BreakScope |
+                                  Scope::DeclScope | Scope::ControlScope);
 
   const VarDecl* condVar = WS->getConditionVariable();
   VarDecl* condVarClone = nullptr;
@@ -1882,13 +1878,12 @@ StmtDiff BaseForwardModeVisitor::VisitWhileStmt(const WhileStmt* WS) {
   if (isa<CompoundStmt>(body)) {
     bodyResult = Visit(body).getStmt();
   } else {
-    beginScope(Scope::DeclScope);
+    ScopeRAII bodyScope(*this, Scope::DeclScope);
     beginBlock();
     StmtDiff Result = Visit(body);
     for (Stmt* S : Result.getBothStmts())
       addToCurrentBlock(S);
     CompoundStmt* Block = endBlock();
-    endScope();
     bodyResult = Block;
   }
 
@@ -1897,8 +1892,6 @@ StmtDiff BaseForwardModeVisitor::VisitWhileStmt(const WhileStmt* WS) {
           .ActOnWhileStmt(/*WhileLoc=*/noLoc, /*LParenLoc=*/noLoc, condRes,
                           /*RParenLoc=*/noLoc, bodyResult)
           .get();
-  // end scope for while loop
-  endScope();
   return StmtDiff(WSDiff);
 }
 
@@ -1908,22 +1901,21 @@ BaseForwardModeVisitor::VisitContinueStmt(const ContinueStmt* ContStmt) {
 }
 
 StmtDiff BaseForwardModeVisitor::VisitDoStmt(const DoStmt* DS) {
-  // begin scope for do-while statement
-  beginScope(Scope::ContinueScope | Scope::BreakScope);
+  // Scope for the whole do-while statement.
+  ScopeRAII doScope(*this, Scope::ContinueScope | Scope::BreakScope);
   Expr* clonedCond = DS->getCond() ? Clone(DS->getCond()) : nullptr;
   const Stmt* body = DS->getBody();
 
   Stmt* bodyResult = nullptr;
   if (isa<CompoundStmt>(body)) {
     bodyResult = Visit(body).getStmt();
   } else {
-    beginScope(Scope::DeclScope);
+    ScopeRAII bodyScope(*this, Scope::DeclScope);
     beginBlock();
     StmtDiff Result = Visit(body);
     for (Stmt* S : Result.getBothStmts())
       addToCurrentBlock(S);
     CompoundStmt* Block = endBlock();
-    endScope();
     bodyResult = Block;
   }
 
@@ -1933,8 +1925,6 @@ StmtDiff BaseForwardModeVisitor::VisitDoStmt(const DoStmt* DS) {
                              /*CondRParen=*/noLoc)
                 .get();
 
-  // end scope for do-while statement
-  endScope();
   return StmtDiff(S);
 }
 

lib/Differentiator/ReverseModeForwPassVisitor.cpp

@@ -236,14 +236,13 @@ StmtDiff ReverseModeForwPassVisitor::ProcessSingleStmt(const clang::Stmt* S) {
 
 StmtDiff
 ReverseModeForwPassVisitor::VisitCompoundStmt(const clang::CompoundStmt* CS) {
-  beginScope(Scope::DeclScope);
+  ScopeRAII compoundScope(*this, Scope::DeclScope);
   beginBlock();
   for (Stmt* S : CS->body()) {
     StmtDiff SDiff = ProcessSingleStmt(S);
     addToCurrentBlock(SDiff.getStmt());
   }
   CompoundStmt* forward = endBlock();
-  endScope();
   return {forward};
 }