From 81003744374dd2a75ae8dd190cdf2188cc086820 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Tue, 4 Oct 2022 10:17:04 -0700
Subject: [PATCH] [flang] Don't force SET_EXPONENT(I=...) argument to
 integer(4)

The implementation of the folding code for SET_EXPONENT() was written
in such a fashion as to convert the I= actual argument value to a 32-bit
integer.  Which is usually not a problem, but it's not always correct
and a test case ran into trouble with it.  Fix to allow any kind of
INTEGER without conversion.

Differential Revision: https://reviews.llvm.org/D135203
---
 flang/include/flang/Evaluate/real.h | 16 ++++++++++------
 flang/lib/Evaluate/fold-real.cpp    | 17 ++++++++++++-----
 flang/lib/Evaluate/real.cpp         |  4 ++--
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/flang/include/flang/Evaluate/real.h b/flang/include/flang/Evaluate/real.h
index e727ac1..d07b0b6 100644
--- a/flang/include/flang/Evaluate/real.h
+++ b/flang/include/flang/Evaluate/real.h
@@ -170,7 +170,7 @@ public:
   static constexpr int MINEXPONENT{2 - exponentBias};
   Real RRSPACING() const;
   Real SPACING() const;
-  Real SET_EXPONENT(int) const;
+  Real SET_EXPONENT(std::int64_t) const;
   Real FRACTION() const;
 
   // SCALE(); also known as IEEE_SCALB and (in IEEE-754 '08) ScaleB.
@@ -182,16 +182,20 @@ public:
     //  be subnormal.)
     auto adjust{exponentBias + binaryPrecision - 1};
     auto expo{adjust + by.ToInt64()};
+    Real twoPow;
+    RealFlags flags;
+    int rMask{1};
     if (IsZero()) {
       expo = exponentBias; // ignore by, don't overflow
     } else if (by > INT{maxExponent}) {
       expo = maxExponent + binaryPrecision - 1;
-    } else if (by < INT{-adjust}) {
-      expo = -1;
+    } else if (by < INT{-adjust}) { // underflow
+      expo = 0;
+      rMask = 0;
+      flags.set(RealFlag::Underflow);
     }
-    Real twoPow;
-    RealFlags flags{
-        twoPow.Normalize(false, static_cast<int>(expo), Fraction::MASKR(1))};
+    flags |=
+        twoPow.Normalize(false, static_cast<int>(expo), Fraction::MASKR(rMask));
     ValueWithRealFlags<Real> result{Multiply(twoPow, rounding)};
     result.flags |= flags;
     return result;
diff --git a/flang/lib/Evaluate/fold-real.cpp b/flang/lib/Evaluate/fold-real.cpp
index 59b7637..5e10deb 100644
--- a/flang/lib/Evaluate/fold-real.cpp
+++ b/flang/lib/Evaluate/fold-real.cpp
@@ -257,11 +257,18 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
           byExpr->u);
     }
   } else if (name == "set_exponent") {
-    return FoldElementalIntrinsic<T, T, Int4>(context, std::move(funcRef),
-        ScalarFunc<T, T, Int4>(
-            [&](const Scalar<T> &x, const Scalar<Int4> &i) -> Scalar<T> {
-              return x.SET_EXPONENT(i.ToInt64());
-            }));
+    if (const auto *iExpr{UnwrapExpr<Expr<SomeInteger>>(args[1])}) {
+      return common::visit(
+          [&](const auto &iVal) {
+            using TY = ResultType<decltype(iVal)>;
+            return FoldElementalIntrinsic<T, T, TY>(context, std::move(funcRef),
+                ScalarFunc<T, T, TY>(
+                    [&](const Scalar<T> &x, const Scalar<TY> &i) -> Scalar<T> {
+                      return x.SET_EXPONENT(i.ToInt64());
+                    }));
+          },
+          iExpr->u);
+    }
   } else if (name == "sign") {
     return FoldElementalIntrinsic<T, T, T>(
         context, std::move(funcRef), &Scalar<T>::SIGN);
diff --git a/flang/lib/Evaluate/real.cpp b/flang/lib/Evaluate/real.cpp
index ca0f4fd..b5e92b7 100644
--- a/flang/lib/Evaluate/real.cpp
+++ b/flang/lib/Evaluate/real.cpp
@@ -756,7 +756,7 @@ template <typename W, int P> Real<W, P> Real<W, P>::SPACING() const {
 
 // 16.9.171
 template <typename W, int P>
-Real<W, P> Real<W, P>::SET_EXPONENT(int expo) const {
+Real<W, P> Real<W, P>::SET_EXPONENT(std::int64_t expo) const {
   if (IsNotANumber()) {
     return *this;
   } else if (IsInfinite()) {
@@ -764,7 +764,7 @@ Real<W, P> Real<W, P>::SET_EXPONENT(int expo) const {
   } else if (IsZero()) {
     return *this;
   } else {
-    return SCALE(Integer<32>(expo - UnbiasedExponent() - 1)).value;
+    return SCALE(Integer<64>(expo - UnbiasedExponent() - 1)).value;
   }
 }
 
-- 
2.7.4