[flang-commits] [flang] [flang] Fix type of array constructor triplet (extension) (PR #92970)

[Peter Klausler via flang-commits]

https://github.com/klausler created https://github.com/llvm/llvm-project/pull/92970

The type of "[lo:hi:str]" is being forced to INTEGER(8), ignoring the types of the operands to the triplet.  The type of a triplet should be whatever the type of "lo+hi+str" would be.

(Array constructor triplets are a syntactic sugar extension for an implied DO loop.)

>From 710ab7078eba9d46106e1d6d7d01aa45169dc429 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Tue, 21 May 2024 15:17:29 -0700
Subject: [PATCH] [flang] Fix type of array constructor triplet (extension)

The type of "[lo:hi:str]" is being forced to INTEGER(8), ignoring
the types of the operands to the triplet.  The type of a triplet
should be whatever the type of "lo+hi+str" would be.

(Array constructor triplets are a syntactic sugar extension for
an implied DO loop.)
---
 flang/lib/Semantics/expression.cpp | 74 ++++++++++++++++++++----------
 flang/test/Evaluate/triplets01.f90 | 11 +++++
 2 files changed, 61 insertions(+), 24 deletions(-)
 create mode 100644 flang/test/Evaluate/triplets01.f90

diff --git a/flang/lib/Semantics/expression.cpp b/flang/lib/Semantics/expression.cpp
index 06e38da6626a9..50e2b41212d7d 100644
--- a/flang/lib/Semantics/expression.cpp
+++ b/flang/lib/Semantics/expression.cpp
@@ -1600,16 +1600,23 @@ class ArrayConstructorContext {
       parser::CharBlock name, std::int64_t lower, std::int64_t upper,
       std::int64_t stride);
 
-  template <int KIND, typename A>
-  std::optional<Expr<Type<TypeCategory::Integer, KIND>>> GetSpecificIntExpr(
-      const A &x) {
-    if (MaybeExpr y{exprAnalyzer_.Analyze(x)}) {
+  template <int KIND>
+  std::optional<Expr<Type<TypeCategory::Integer, KIND>>> ToSpecificInt(
+      MaybeExpr &&y) {
+    if (y) {
       Expr<SomeInteger> *intExpr{UnwrapExpr<Expr<SomeInteger>>(*y)};
       return Fold(exprAnalyzer_.GetFoldingContext(),
           ConvertToType<Type<TypeCategory::Integer, KIND>>(
               std::move(DEREF(intExpr))));
+    } else {
+      return std::nullopt;
     }
-    return std::nullopt;
+  }
+
+  template <int KIND, typename A>
+  std::optional<Expr<Type<TypeCategory::Integer, KIND>>> GetSpecificIntExpr(
+      const A &x) {
+    return ToSpecificInt<KIND>(exprAnalyzer_.Analyze(x));
   }
 
   // Nested array constructors all reference the same ExpressionAnalyzer,
@@ -1772,26 +1779,45 @@ void ArrayConstructorContext::Add(const parser::AcValue &x) {
 
 // Transforms l:u(:s) into (_,_=l,u(,s)) with an anonymous index '_'
 void ArrayConstructorContext::Add(const parser::AcValue::Triplet &triplet) {
-  std::optional<Expr<ImpliedDoIntType>> lower{
-      GetSpecificIntExpr<ImpliedDoIntType::kind>(std::get<0>(triplet.t))};
-  std::optional<Expr<ImpliedDoIntType>> upper{
-      GetSpecificIntExpr<ImpliedDoIntType::kind>(std::get<1>(triplet.t))};
-  std::optional<Expr<ImpliedDoIntType>> stride{
-      GetSpecificIntExpr<ImpliedDoIntType::kind>(std::get<2>(triplet.t))};
-  if (lower && upper) {
-    if (!stride) {
-      stride = Expr<ImpliedDoIntType>{1};
-    }
-    if (!type_) {
-      type_ = DynamicTypeWithLength{ImpliedDoIntType::GetType()};
+  MaybeExpr lowerExpr{exprAnalyzer_.Analyze(std::get<0>(triplet.t))};
+  MaybeExpr upperExpr{exprAnalyzer_.Analyze(std::get<1>(triplet.t))};
+  MaybeExpr strideExpr{exprAnalyzer_.Analyze(std::get<2>(triplet.t))};
+  if (lowerExpr && upperExpr) {
+    auto lowerType{lowerExpr->GetType()};
+    auto upperType{upperExpr->GetType()};
+    auto strideType{strideExpr ? strideExpr->GetType() : lowerType};
+    if (lowerType && upperType && strideType) {
+      int kind{lowerType->kind()};
+      if (upperType->kind() > kind) {
+        kind = upperType->kind();
+      }
+      if (strideType->kind() > kind) {
+        kind = strideType->kind();
+      }
+      auto lower{ToSpecificInt<ImpliedDoIntType::kind>(std::move(lowerExpr))};
+      auto upper{ToSpecificInt<ImpliedDoIntType::kind>(std::move(upperExpr))};
+      if (lower && upper) {
+        auto stride{
+            ToSpecificInt<ImpliedDoIntType::kind>(std::move(strideExpr))};
+        if (!stride) {
+          stride = Expr<ImpliedDoIntType>{1};
+        }
+        DynamicType type{TypeCategory::Integer, kind};
+        if (!type_) {
+          type_ = DynamicTypeWithLength{type};
+        }
+        parser::CharBlock anonymous;
+        if (auto converted{ConvertToType(type,
+                AsGenericExpr(
+                    Expr<ImpliedDoIntType>{ImpliedDoIndex{anonymous}}))}) {
+          auto v{std::move(values_)};
+          Push(std::move(converted));
+          std::swap(v, values_);
+          values_.Push(ImpliedDo<SomeType>{anonymous, std::move(*lower),
+              std::move(*upper), std::move(*stride), std::move(v)});
+        }
+      }
     }
-    auto v{std::move(values_)};
-    parser::CharBlock anonymous;
-    Push(Expr<SomeType>{
-        Expr<SomeInteger>{Expr<ImpliedDoIntType>{ImpliedDoIndex{anonymous}}}});
-    std::swap(v, values_);
-    values_.Push(ImpliedDo<SomeType>{anonymous, std::move(*lower),
-        std::move(*upper), std::move(*stride), std::move(v)});
   }
 }
 
diff --git a/flang/test/Evaluate/triplets01.f90 b/flang/test/Evaluate/triplets01.f90
new file mode 100644
index 0000000000000..aba9772f6b953
--- /dev/null
+++ b/flang/test/Evaluate/triplets01.f90
@@ -0,0 +1,11 @@
+! RUN: %python %S/test_folding.py %s %flang_fc1
+module m
+  logical, parameter :: test01 = all([1:10:2] == [(j, j=1,10,2)])
+  logical, parameter :: test02 = kind([1:20:2]) == kind(1)
+  logical, parameter :: test03 = all([10:1:-3,123] == [(j, j=10,1,-3),123])
+  logical, parameter :: test04 = kind([10:1:-3,123]) == kind(1)
+  logical, parameter :: test05 = kind([10_2:1_2:-3_2,123_2]) == 2
+  logical, parameter :: test06 = all([10_2:1_2:-3_2,123_2] == [(j, integer(2)::j=10,1,-3),123_2])
+  logical, parameter :: test07 = kind([10_2:1_4:-3_2]) == 4
+  logical, parameter :: test08 = kind([10_2:1_4]) == 4
+end