[flang] Correct folding of CSHIFT and EOSHIFT for DIM>1

[klausler]

Summary

The algorithm was wrong for higher dimensions, and so were
the expected test results. Rework.

[klausler]

Update some test results.

flang/lib/Evaluate/fold-implementation.h

if (shift->Rank() > 0) {
             ok = false;
           }
           ++k;
         }
       }
     }
     if (ok) {
       std::vector<Scalar<T>> resultElements;
-      ConstantSubscripts arrayAt{array->lbounds()};
-      ConstantSubscript dimLB{arrayAt[zbDim]};
+      ConstantSubscripts arrayLB{array->lbounds()};
+      ConstantSubscripts arrayAt{arrayLB};
+      ConstantSubscript &dimIndex{arrayAt[zbDim]};
+      ConstantSubscript dimLB{dimIndex}; // initial value
       ConstantSubscript dimExtent{array->shape()[zbDim]};
-      ConstantSubscripts shiftAt{shift->lbounds()};
-      for (auto n{GetSize(array->shape())}; n > 0; n -= dimExtent) {
-        ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
-        ConstantSubscript zbDimIndex{shiftCount % dimExtent};
-        if (zbDimIndex < 0) {
-          zbDimIndex += dimExtent;
+      ConstantSubscripts shiftLB{shift->lbounds()};
+      for (auto n{GetSize(array->shape())}; n > 0; --n) {
+        ConstantSubscript origDimIndex{dimIndex};
+        ConstantSubscripts shiftAt;
+        if (shift->Rank() > 0) {
+          int k{0};
+          for (int j{0}; j < rank; ++j) {
+            if (j != zbDim) {
+              shiftAt.emplace_back(shiftLB[k++] + arrayAt[j] - arrayLB[j]);
             }
-        for (ConstantSubscript j{0}; j < dimExtent; ++j) {
-          arrayAt[zbDim] = dimLB + zbDimIndex;
-          resultElements.push_back(array->At(arrayAt));
-          if (++zbDimIndex == dimExtent) {
-            zbDimIndex = 0;
           }
         }
-        arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
+        ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
+        dimIndex = dimLB + ((dimIndex - dimLB + shiftCount) % dimExtent);
+        if (dimIndex < dimLB) {
+          dimIndex += dimExtent;
+        } else if (dimIndex >= dimLB + dimExtent) {
+          dimIndex -= dimExtent;
+        }
+        resultElements.push_back(array->At(arrayAt));
+        dimIndex = origDimIndex;
         array->IncrementSubscripts(arrayAt);
-        shift->IncrementSubscripts(shiftAt);
       }
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};
     }
   }
   // Invalid, prevent re-folding
   return MakeInvalidIntrinsic(std::move(funcRef));
 }
if (boundary && boundary->Rank() > 0) {
             ok = false;
           }
           ++k;
         }
       }
     }
     if (ok) {
       std::vector<Scalar<T>> resultElements;
-      ConstantSubscripts arrayAt{array->lbounds()};
-      ConstantSubscript dimLB{arrayAt[zbDim]};
+      ConstantSubscripts arrayLB{array->lbounds()};
+      ConstantSubscripts arrayAt{arrayLB};
+      ConstantSubscript &dimIndex{arrayAt[zbDim]};
+      ConstantSubscript dimLB{dimIndex}; // initial value
       ConstantSubscript dimExtent{array->shape()[zbDim]};
-      ConstantSubscripts shiftAt{shift->lbounds()};
-      ConstantSubscripts boundaryAt;
+      ConstantSubscripts shiftLB{shift->lbounds()};
+      ConstantSubscripts boundaryLB;
       if (boundary) {
-        boundaryAt = boundary->lbounds();
+        boundaryLB = boundary->lbounds();
+      }
+      for (auto n{GetSize(array->shape())}; n > 0; --n) {
+        ConstantSubscript origDimIndex{dimIndex};
+        ConstantSubscripts shiftAt;
+        if (shift->Rank() > 0) {
+          int k{0};
+          for (int j{0}; j < rank; ++j) {
+            if (j != zbDim) {
+              shiftAt.emplace_back(shiftLB[k++] + arrayAt[j] - arrayLB[j]);
+            }
+          }
         }
-      for (auto n{GetSize(array->shape())}; n > 0; n -= dimExtent) {
         ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
-        for (ConstantSubscript j{0}; j < dimExtent; ++j) {
-          ConstantSubscript zbAt{shiftCount + j};
-          if (zbAt >= 0 && zbAt < dimExtent) {
-            arrayAt[zbDim] = dimLB + zbAt;
-            resultElements.push_back(array->At(arrayAt));
-          } else if (boundary) {
-            resultElements.push_back(boundary->At(boundaryAt));
-          } else if constexpr (T::category == TypeCategory::Integer ||
-              T::category == TypeCategory::Real ||
-              T::category == TypeCategory::Complex ||
-              T::category == TypeCategory::Logical) {
-            resultElements.emplace_back();
-          } else if constexpr (T::category == TypeCategory::Character) {
-            auto len{static_cast<std::size_t>(array->LEN())};
-            typename Scalar<T>::value_type space{' '};
-            resultElements.emplace_back(len, space);
-          } else {
-            DIE("no derived type boundary");
-          }
-        }
-        arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
+        dimIndex += shiftCount;
+        if (dimIndex >= dimLB && dimIndex < dimLB + dimExtent) {
+          ConstantSubscripts boundaryAt;
+          if (boundary->Rank() > 0) {
+            for (int j{0}; j < rank; ++j) {
+              int k{0};
+              if (j != zbDim) {
+                boundaryAt.emplace_back(
+                    boundaryLB[k++] + arrayAt[j] - arrayLB[j]);
+              }
+            }
+          }
+        dimIndex = origDimIndex;
         array->IncrementSubscripts(arrayAt);
-        shift->IncrementSubscripts(shiftAt);
-        if (boundary) {
-          boundary->IncrementSubscripts(boundaryAt);
-        }
       }
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};
     }
   }
   // Invalid, prevent re-folding
   return MakeInvalidIntrinsic(std::move(funcRef));
 }

flang/test/Evaluate/folding23.f90

 ! RUN: %python %S/test_folding.py %s %flang_fc1
 ! Tests folding of EOSHIFT (valid cases)
 module m
   integer, parameter :: arr(2,3) = reshape([1, 2, 3, 4, 5, 6], shape(arr))
   logical, parameter :: test_sanity = all([arr] == [1, 2, 3, 4, 5, 6])
   logical, parameter :: test_eoshift_0 = all(eoshift([1, 2, 3], 0) == [1, 2, 3])
   logical, parameter :: test_eoshift_1 = all(eoshift([1, 2, 3], 1) == [2, 3, 0])
   logical, parameter :: test_eoshift_2 = all(eoshift([1, 2, 3], -1) == [0, 1, 2])
   logical, parameter :: test_eoshift_3 = all(eoshift([1., 2., 3.], 1) == [2., 3., 0.])
   logical, parameter :: test_eoshift_4 = all(eoshift(['ab', 'cd', 'ef'], -1, 'x') == ['x ', 'ab', 'cd'])
   logical, parameter :: test_eoshift_5 = all([eoshift(arr, 1, dim=1)] == [2, 0, 4, 0, 6, 0])
-  logical, parameter :: test_eoshift_6 = all([eoshift(arr, 1, dim=2)] == [3, 5, 0, 4, 6, 0])
+  logical, parameter :: test_eoshift_6 = all([eoshift(arr, 1, dim=2)] == [3, 4, 5, 6, 0, 0])
   logical, parameter :: test_eoshift_7 = all([eoshift(arr, [1, -1, 0])] == [2, 0, 0, 3, 5, 6])
-  logical, parameter :: test_eoshift_8 = all([eoshift(arr, [1, -1], dim=2)] == [3, 5, 0, 0, 2, 4])
+  logical, parameter :: test_eoshift_8 = all([eoshift(arr, [1, -1], dim=2)] == [3, 0, 5, 2, 0, 4])
 end module

flang/test/Evaluate/folding27.f90

 ! RUN: %python %S/test_folding.py %s %flang_fc1
 ! Tests folding of CSHIFT (valid cases)
 module m
   integer, parameter :: arr(2,3) = reshape([1, 2, 3, 4, 5, 6], shape(arr))
   logical, parameter :: test_sanity = all([arr] == [1, 2, 3, 4, 5, 6])
   logical, parameter :: test_cshift_0 = all(cshift([1, 2, 3], 0) == [1, 2, 3])
   logical, parameter :: test_cshift_1 = all(cshift([1, 2, 3], 1) == [2, 3, 1])
   logical, parameter :: test_cshift_2 = all(cshift([1, 2, 3], 3) == [1, 2, 3])
   logical, parameter :: test_cshift_3 = all(cshift([1, 2, 3], 4) == [2, 3, 1])
   logical, parameter :: test_cshift_4 = all(cshift([1, 2, 3], -1) == [3, 1, 2])
   logical, parameter :: test_cshift_5 = all([cshift(arr, 1, dim=1)] == [2, 1, 4, 3, 6, 5])
-  logical, parameter :: test_cshift_6 = all([cshift(arr, 1, dim=2)] == [3, 5, 1, 4, 6, 2])
+  logical, parameter :: test_cshift_6 = all([cshift(arr, 1, dim=2)] == [3, 4, 5, 6, 1, 2])
   logical, parameter :: test_cshift_7 = all([cshift(arr, [1, 2, 3])] == [2, 1, 3, 4, 6, 5])
-  logical, parameter :: test_cshift_8 = all([cshift(arr, [1, 2], dim=2)] == [3, 5, 1, 6, 2, 4])
+  logical, parameter :: test_cshift_8 = all([cshift(arr, [1, 2], dim=2)] == [3, 6, 5, 2, 1, 4])
 end module