From 6a42b46be3a01d091f86d26980e4a980d4f6f012 Mon Sep 17 00:00:00 2001
From: PetroZarytskyi <119341518+PetroZarytskyi@users.noreply.github.com>
Date: Tue, 12 Dec 2023 16:30:50 +0200
Subject: [PATCH] Make const loop variables global and drop const. (#680)

We didn't take into account const variables inside loops can be re-initialized on different loop iterations so we have to store them just like non-const variables. const should be dropped to allow us to replace initializations with assignments.
Fixes #667.
---
 lib/Differentiator/ReverseModeVisitor.cpp | 13 ++++--
 test/Gradient/Loops.C                     | 50 +++++++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
index f2443e4ee..0b3d53d0f 100644
--- a/lib/Differentiator/ReverseModeVisitor.cpp
+++ b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -2692,9 +2692,15 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
         // assignments. This is a temporary measure to avoid the bug that arises
         // from overwriting local variables on different loop passes.
         if (isInsideLoop) {
-          if (VD->getType()->isBuiltinType() &&
-              !VD->getType().isConstQualified()) {
+          if (VD->getType()->isBuiltinType()) {
             auto* decl = VDDiff.getDecl();
+            /// The same variable will be assigned with new values every
+            /// loop iteration so the const qualifier must be dropped.
+            if (decl->getType().isConstQualified()) {
+              QualType nonConstType =
+                  getNonConstType(decl->getType(), m_Context, m_Sema);
+              decl->setType(nonConstType);
+            }
             if (decl->getInit()) {
               auto* declRef = BuildDeclRef(decl);
               auto pushPop =
@@ -2744,8 +2750,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     /// overwriting local variables on different loop passes.
     if (isInsideLoop) {
       if (auto* VD = dyn_cast<VarDecl>(decls[0])) {
-        if (VD->getType()->isBuiltinType() &&
-            !VD->getType().isConstQualified()) {
+        if (VD->getType()->isBuiltinType()) {
           addToBlock(DSClone, m_Globals);
           Stmt* initAssignments = MakeCompoundStmt(inits);
           initAssignments = unwrapIfSingleStmt(initAssignments);
diff --git a/test/Gradient/Loops.C b/test/Gradient/Loops.C
index 7eac9b1e4..525d8693c 100644
--- a/test/Gradient/Loops.C
+++ b/test/Gradient/Loops.C
@@ -198,6 +198,55 @@ double f5(double x){
 //CHECK-NEXT:       }
 //CHECK-NEXT:   }
 
+double f_const_local(double x) {
+  double res = 0;
+  for (int i = 0; i < 3; ++i) {
+    const double n = x + i;
+    res += x * n;
+  }
+  return res;
+} // == 3x^2 + 3x
+
+//CHECK:   void f_const_local_grad(double x, clad::array_ref<double> _d_x) {
+//CHECK-NEXT:    double _d_res = 0;
+//CHECK-NEXT:    unsigned long _t0;
+//CHECK-NEXT:    int _d_i = 0;
+//CHECK-NEXT:    clad::tape<double> _t1 = {};
+//CHECK-NEXT:    double _d_n = 0;
+//CHECK-NEXT:    double n = 0;
+//CHECK-NEXT:    clad::tape<double> _t2 = {};
+//CHECK-NEXT:    double res = 0;
+//CHECK-NEXT:    _t0 = 0;
+//CHECK-NEXT:    for (int i = 0; i < 3; ++i) {
+//CHECK-NEXT:        _t0++;
+//CHECK-NEXT:        clad::push(_t1, n) , n = x + i;
+//CHECK-NEXT:        clad::push(_t2, res);
+//CHECK-NEXT:        res += x * n;
+//CHECK-NEXT:    }
+//CHECK-NEXT:    goto _label0;
+//CHECK-NEXT:  _label0:
+//CHECK-NEXT:    _d_res += 1;
+//CHECK-NEXT:    for (; _t0; _t0--) {
+//CHECK-NEXT:        --i;
+//CHECK-NEXT:        {
+//CHECK-NEXT:            res = clad::pop(_t2);
+//CHECK-NEXT:            double _r_d0 = _d_res;
+//CHECK-NEXT:            _d_res += _r_d0;
+//CHECK-NEXT:            double _r0 = _r_d0 * n;
+//CHECK-NEXT:            * _d_x += _r0;
+//CHECK-NEXT:            double _r1 = x * _r_d0;
+//CHECK-NEXT:            _d_n += _r1;
+//CHECK-NEXT:            _d_res -= _r_d0;
+//CHECK-NEXT:        }
+//CHECK-NEXT:        {
+//CHECK-NEXT:            * _d_x += _d_n;
+//CHECK-NEXT:            _d_i += _d_n;
+//CHECK-NEXT:            _d_n = 0;
+//CHECK-NEXT:            n = clad::pop(_t1);
+//CHECK-NEXT:        }
+//CHECK-NEXT:    }
+//CHECK-NEXT:}
+
 double f_sum(double *p, int n) {
   double s = 0;
   for (int i = 0; i < n; i++)
@@ -1715,6 +1764,7 @@ int main() {
   TEST(f3, 3); // CHECK-EXEC: {6.00}
   TEST(f4, 3); // CHECK-EXEC: {27.00}
   TEST(f5, 3); // CHECK-EXEC: {1.00}
+  TEST(f_const_local, 3); // CHECK-EXEC: {21.00}
 
   double p[] = { 1, 2, 3, 4, 5 };