Add support for ArrayInitLoopExpr/ArrayInitIndexExpr in reverse mode

Whenever Clang generates an implicit copy/move constructor of a class with a static array member, it uses ArrayInitLoopExpr to express the array copy. For example
```
struct arrWrapper {
    double arr[2];
};
```
The implicit copy-constructor will look like
```
arrWrapper(const arrWrapper& other) : arr(ArrayInitLoopExpr(other[ArrayInitIndexExpr])) {}
```
The problem is that we cannot do the same with explicit functions. Instead, we have to replicate this behaviour using loops.
This PR adds support for ``ArrayInitLoopExpr`` in constructors, including the ones that initialize high-dimensional arrays and arrays of copiable objects.
Also, the last error in #791 is caused by this.

Fixes #791.

Author: PetroZarytskyi

include/clad/Differentiator/ReverseModeVisitor.h

@@ -14,6 +14,7 @@
 #include "clad/Differentiator/ReverseModeVisitorDirectionKinds.h"
 #include "clad/Differentiator/VisitorBase.h"
 
+#include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
@@ -29,6 +30,7 @@
 #include <array>
 #include <limits>
 #include <memory>
+#include <queue>
 #include <stack>
 #include <unordered_map>
 
@@ -397,6 +399,9 @@ namespace clad {
     StmtDiff VisitIntegerLiteral(const clang::IntegerLiteral* IL);
     StmtDiff VisitMemberExpr(const clang::MemberExpr* ME);
     StmtDiff VisitParenExpr(const clang::ParenExpr* PE);
+    StmtDiff VisitArrayInitLoopExpr(const clang::ArrayInitLoopExpr* AILE);
+    StmtDiff VisitArrayInitIndexExpr(const clang::ArrayInitIndexExpr* AIIE);
+    StmtDiff VisitOpaqueValueExpr(const clang::OpaqueValueExpr* OVE);
     virtual StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS);
     StmtDiff VisitStmt(const clang::Stmt* S);
     virtual StmtDiff VisitUnaryOperator(const clang::UnaryOperator* UnOp);
@@ -718,6 +723,13 @@ namespace clad {
     void PopSwitchStmtInfo() { m_SwitchStmtsData.pop_back(); }
 
   private:
+    // When differentiating ArrayInitLoopExpr, we need to replace
+    // ArrayInitIndexExpr with real indices. We need to both add and pop them in
+    // the right order, so we use std::queue. For example, for `arr[i][j]`, we
+    // add `i` first, then `j`, and then pop them in the same order to generate
+    // the subscript expr.
+    std::queue<clang::VarDecl*> m_ArrayInitLoopIdx;
+
     // FIXME: This variable is used to track
     // whether we're currently visiting an init of a var decl.
     // This is only necessary because we don't create constructors
@@ -726,6 +738,7 @@ namespace clad {
     // other cases, we have to use InitListExpr and change the constructor
     // style. Remove this once we generate constructors explicitly.
     bool m_TrackVarDeclConstructor = false;
+
     /// A flag indicating if the Stmt is contained in a checkpointed loop.
     bool m_IsInsideCheckpointedLoop = false;
   };

include/clad/Differentiator/VisitorBase.h

@@ -13,6 +13,7 @@
 
 #include "clang/AST/Expr.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/AST/Stmt.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/Diagnostic.h"
@@ -28,6 +29,7 @@
 
 #include <array>
 #include <cassert>
+#include <cstddef>
 #include <stack>
 #include <unordered_map>
 
@@ -294,6 +296,14 @@ namespace clad {
     clang::Expr* BuildOperatorCall(clang::OverloadedOperatorKind OOK,
                                    llvm::MutableArrayRef<clang::Expr*> ArgExprs,
                                    clang::SourceLocation OpLoc = noLoc);
+
+    /// A shorthand to generage a standard loop of form
+    /// ```
+    /// for (type loopCounter = 0; loopCounter < N; ++loopCounter)
+    ///   body;
+    /// ```
+    clang::ForStmt* BuildStandardForLoop(clang::VarDecl* loopCounter, size_t N,
+                                         clang::Stmt* body);
     /// Function to resolve Unary Minus. If the leftmost operand
     /// has a Unary Minus then adds parens before adding the unary minus.
     /// \param[in] E Expression fed to the recursive call.

lib/Differentiator/DiffPlanner.cpp

@@ -709,6 +709,10 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
           return false;
         return true;
       }
+      // The sub-stmt of OpaqueValueExpr is not visited automatically
+      bool VisitOpaqueValueExpr(const clang::OpaqueValueExpr* OVE) {
+        return TraverseStmt(OVE->getSourceExpr());
+      }
       // FIXME: This is a temporary measure until we add support for
       // `this` in varied analysis.
       bool VisitCXXThisExpr(const clang::CXXThisExpr* TE) { return false; }

lib/Differentiator/ReverseModeVisitor.cpp

@@ -723,6 +723,49 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     return Visit(ILE->getSubExpr(), dfdx());
   }
 
+  StmtDiff
+  ReverseModeVisitor::VisitArrayInitLoopExpr(const ArrayInitLoopExpr* AILE) {
+    // Since ArrayInitLoopExpr is not possible to express with regular syntax,
+    // we have to replicate it with loops.
+    // The code we're differentiated is of the form
+    // res = ArrayInitLoopExpr(arr[ArrayInitIndexExpr])
+    // We have to replace ArrayInitIndexExpr with an actual index `i`
+    // and wrap the code in a for loop to compute the derivative as follows:
+    // for (int i = 0; i < N; ++i)
+    //   _d_arr[i] += _d_res[i];
+    beginScope(Scope::DeclScope);
+    VarDecl* idxDecl = BuildVarDecl(m_Context.UnsignedIntTy, "i",
+                                    getZeroInit(m_Context.IntTy));
+    // Push the index to the queue so that we can replace ArrayInitIndexExpr
+    // when we encounter it.
+    m_ArrayInitLoopIdx.push(idxDecl);
+    Expr* idx = BuildDeclRef(idxDecl);
+    // Build `_d_res[i]`
+    Expr* diff = BuildArraySubscript(dfdx(), {idx});
+    beginBlock(direction::reverse);
+    Visit(AILE->getSubExpr(), diff);
+    Stmt* block = utils::unwrapIfSingleStmt(endBlock(direction::reverse));
+    Stmt* loopDiff = BuildStandardForLoop(
+        idxDecl, AILE->getArraySize().getZExtValue(), block);
+    addToCurrentBlock(loopDiff, direction::reverse);
+    endScope();
+    // We cannot clone ArrayInitLoopExpr because it's not possible to express
+    // with standard c++ syntax.
+    return {};
+  }
+
+  StmtDiff
+  ReverseModeVisitor::VisitArrayInitIndexExpr(const ArrayInitIndexExpr* AIIE) {
+    VarDecl* idxDecl = m_ArrayInitLoopIdx.front();
+    m_ArrayInitLoopIdx.pop();
+    return {BuildDeclRef(idxDecl)};
+  }
+
+  StmtDiff
+  ReverseModeVisitor::VisitOpaqueValueExpr(const OpaqueValueExpr* OVE) {
+    return Visit(OVE->getSourceExpr(), dfdx());
+  }
+
   StmtDiff ReverseModeVisitor::VisitStmt(const Stmt* S) {
     diagUnsupported(S);
     // Unknown stmt, just clone it.

lib/Differentiator/VisitorBase.cpp

@@ -23,6 +23,7 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/OperationKinds.h"
+#include "clang/AST/Stmt.h"
 #include "clang/AST/TemplateBase.h"
 #include "clang/Basic/OperatorKinds.h"
 #include "clang/Basic/SourceManager.h"
@@ -41,6 +42,7 @@
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <numeric>
 
 #include "clad/Differentiator/Compatibility.h"
@@ -264,6 +266,17 @@ namespace clad {
     return new (m_Context) DeclStmt(DGR, noLoc, noLoc);
   }
 
+  ForStmt* VisitorBase::BuildStandardForLoop(VarDecl* loopCounter, size_t N,
+                                             Stmt* body) {
+    Stmt* init = BuildDeclStmt(loopCounter);
+    Expr* numExpr =
+        ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, N);
+    Expr* cond = BuildOp(BO_LT, BuildDeclRef(loopCounter), numExpr);
+    Expr* inc = BuildOp(UO_PreInc, BuildDeclRef(loopCounter));
+    return new (m_Context) ForStmt(m_Context, init, cond, /*CondVar=*/nullptr,
+                                   inc, body, noLoc, noLoc, noLoc);
+  }
+
   DeclRefExpr* VisitorBase::BuildDeclRef(DeclaratorDecl* D,
                                          NestedNameSpecifier* NNS /*=nullptr*/,
                                          ExprValueKind VK /*=VK_LValue*/) {

test/Gradient/Constructors.C

@@ -445,6 +445,123 @@ double fn9(S& s) {
 // CHECK-NEXT:      }
 // CHECK-NEXT:  }
 
+struct arrWrapper {
+    double arr[2];
+};
+
+// CHECK:  static inline constexpr void constructor_pullback(const arrWrapper &arg, arrWrapper *_d_this, arrWrapper *_d_arg) noexcept {
+// CHECK-NEXT:      for (unsigned {{int|long}} i = 0; i < 2; ++i)
+// CHECK-NEXT:          (*_d_arg).arr[i] += _d_this->arr[i];
+// CHECK-NEXT:  }
+
+double fn10(double x, double y) {
+    arrWrapper a = {x, y};
+    arrWrapper b = a;
+    return b.arr[0] + b.arr[1];
+}
+
+// CHECK:  void fn10_grad(double x, double y, double *_d_x, double *_d_y) {
+// CHECK-NEXT:      arrWrapper _d_a = {{.*0.*}};
+// CHECK-NEXT:      arrWrapper a = {{.*x, y.*}};
+// CHECK-NEXT:      arrWrapper _d_b = _d_a;
+// CHECK-NEXT:      arrWrapper b = a;
+// CHECK-NEXT:      {
+// CHECK-NEXT:          _d_b.arr[0] += 1;
+// CHECK-NEXT:          _d_b.arr[1] += 1;
+// CHECK-NEXT:      }
+// CHECK-NEXT:      arrWrapper::constructor_pullback(a, &_d_b, &_d_a);
+// CHECK-NEXT:      {
+// CHECK-NEXT:          *_d_x += _d_a.arr[0];
+// CHECK-NEXT:          *_d_y += _d_a.arr[1];
+// CHECK-NEXT:      }
+// CHECK-NEXT:  }
+
+struct arr2DWrapper {
+    double arr[1][2];
+};
+
+// CHECK:  static inline constexpr void constructor_pullback(const arr2DWrapper &arg, arr2DWrapper *_d_this, arr2DWrapper *_d_arg) noexcept {
+// CHECK-NEXT:      for (unsigned {{int|long}} i = 0; i < 1; ++i)
+// CHECK-NEXT:          for (unsigned {{int|long}} i0 = 0; i0 < 2; ++i0)
+// CHECK-NEXT:              (*_d_arg).arr[i][i0] += _d_this->arr[i][i0];
+// CHECK-NEXT:  }
+
+double fn11(double x, double y) {
+    arr2DWrapper a = {x, y};
+    arr2DWrapper b = a;
+    return b.arr[0][0] + b.arr[0][1];
+}
+
+// CHECK:  void fn11_grad(double x, double y, double *_d_x, double *_d_y) {
+// CHECK-NEXT:      arr2DWrapper _d_a = {{.*0.*}};
+// CHECK-NEXT:      arr2DWrapper a = {{.*x, y.*}};
+// CHECK-NEXT:      arr2DWrapper _d_b = _d_a;
+// CHECK-NEXT:      arr2DWrapper b = a;
+// CHECK-NEXT:      {
+// CHECK-NEXT:          _d_b.arr[0][0] += 1;
+// CHECK-NEXT:          _d_b.arr[0][1] += 1;
+// CHECK-NEXT:      }
+// CHECK-NEXT:      arr2DWrapper::constructor_pullback(a, &_d_b, &_d_a);
+// CHECK-NEXT:      {
+// CHECK-NEXT:          *_d_x += _d_a.arr[0][0];
+// CHECK-NEXT:          *_d_y += _d_a.arr[0][1];
+// CHECK-NEXT:      }
+// CHECK-NEXT:  }
+
+struct cust_double {
+    cust_double(double x = 0): val(x) {} 
+    double val;
+};
+
+// CHECK:  static void constructor_pullback(double x, cust_double *_d_this, double *_d_x) {
+// CHECK-NEXT:      {
+// CHECK-NEXT:          *_d_x += _d_this->val;
+// CHECK-NEXT:          _d_this->val = 0.;
+// CHECK-NEXT:      }
+// CHECK-NEXT:  }
+
+// CHECK:  static inline constexpr void constructor_pullback(const cust_double &arg, cust_double *_d_this, cust_double *_d_arg) noexcept {
+// CHECK-NEXT:      {
+// CHECK-NEXT:          (*_d_arg).val += _d_this->val;
+// CHECK-NEXT:          _d_this->val = 0.;
+// CHECK-NEXT:      }
+// CHECK-NEXT:  }
+
+struct arrStructWrapper {
+    cust_double arr[2];
+};
+
+// CHECK:  static inline constexpr void constructor_pullback(const arrStructWrapper &arg, arrStructWrapper *_d_this, arrStructWrapper *_d_arg) noexcept {
+// CHECK-NEXT:      for (unsigned {{int|long}} i = 0; i < 2; ++i)
+// CHECK-NEXT:          cust_double::constructor_pullback(arg.arr[i], &_d_this->arr[i], &(*_d_arg).arr[i]);
+// CHECK-NEXT:  }
+
+double fn12(double x, double y) {
+    arrStructWrapper a = {x, y};
+    arrStructWrapper b = a;
+    return b.arr[0].val + b.arr[1].val;
+}
+
+// CHECK:  void fn12_grad(double x, double y, double *_d_x, double *_d_y) {
+// CHECK-NEXT:      arrStructWrapper _d_a = {{.*}};
+// CHECK-NEXT:      arrStructWrapper a = {{.*x, y.*}};
+// CHECK-NEXT:      arrStructWrapper _d_b = _d_a;
+// CHECK-NEXT:      arrStructWrapper b = a;
+// CHECK-NEXT:      {
+// CHECK-NEXT:          _d_b.arr[0].val += 1;
+// CHECK-NEXT:          _d_b.arr[1].val += 1;
+// CHECK-NEXT:      }
+// CHECK-NEXT:      arrStructWrapper::constructor_pullback(a, &_d_b, &_d_a);
+// CHECK-NEXT:      {
+// CHECK-NEXT:          double _r0 = 0.;
+// CHECK-NEXT:          cust_double::constructor_pullback(x, &_d_a.arr[0], &_r0);
+// CHECK-NEXT:          *_d_x += _r0;
+// CHECK-NEXT:          double _r1 = 0.;
+// CHECK-NEXT:          cust_double::constructor_pullback(y, &_d_a.arr[1], &_r1);
+// CHECK-NEXT:          *_d_y += _r1;
+// CHECK-NEXT:      }
+// CHECK-NEXT:  }
+
 int main() {
     double d_i, d_j;
 
@@ -474,6 +591,14 @@ int main() {
     S s{new double[3]{5, 6, 7}}, _d_s{new double[3]{0}};
     auto dfn9 = clad::gradient(fn9);
     dfn9.execute(s, &_d_s);
-    printf("{%.2f, %.2f, %.2f}\n", _d_s.a[0], _d_s.a[1], _d_s.a[2]);
-    // TEST_GRADIENT(fn9, /*numOfDerivativeArgs=*/1, s, &d_s);    // CHECK-EXEC: {0.00, 0.00, 1.00}
+    printf("{%.2f, %.2f, %.2f}\n", _d_s.a[0], _d_s.a[1], _d_s.a[2]);   // CHECK-EXEC: {0.00, 0.00, 1.00}
+
+    INIT_GRADIENT(fn10);
+    TEST_GRADIENT(fn10, /*numOfDerivativeArgs=*/2, 7, 2, &d_i, &d_j);    // CHECK-EXEC: {1.00, 1.00}
+
+    INIT_GRADIENT(fn11);
+    TEST_GRADIENT(fn11, /*numOfDerivativeArgs=*/2, 9, -1, &d_i, &d_j);    // CHECK-EXEC: {1.00, 1.00}
+
+    INIT_GRADIENT(fn12);
+    TEST_GRADIENT(fn12, /*numOfDerivativeArgs=*/2, 3, 6, &d_i, &d_j);    // CHECK-EXEC: {1.00, 1.00}
 }