From 34cc706b9664882c6b4564a42d3c108abd043009 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Thu, 31 Mar 2022 13:36:09 -0700
Subject: [PATCH] [flang] Fold IBITS() intrinsic function

Implement constant folding of IBITS(); add test.

Differential Revision: https://reviews.llvm.org/D123707
---
 flang/lib/Evaluate/fold-integer.cpp | 24 +++++++++++++++++++++++-
 flang/test/Evaluate/fold-ibits.f90  | 13 +++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 flang/test/Evaluate/fold-ibits.f90

diff --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 6cd0370e..d1a87dd 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -562,6 +562,28 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
           }
           return std::invoke(fptr, i, posVal);
         }));
+  } else if (name == "ibits") {
+    return FoldElementalIntrinsic<T, T, Int4, Int4>(context, std::move(funcRef),
+        ScalarFunc<T, T, Int4, Int4>([&](const Scalar<T> &i,
+                                         const Scalar<Int4> &pos,
+                                         const Scalar<Int4> &len) -> Scalar<T> {
+          auto posVal{static_cast<int>(pos.ToInt64())};
+          auto lenVal{static_cast<int>(len.ToInt64())};
+          if (posVal < 0) {
+            context.messages().Say(
+                "bit position for IBITS(POS=%d,LEN=%d) is negative"_err_en_US,
+                posVal, lenVal);
+          } else if (lenVal < 0) {
+            context.messages().Say(
+                "bit length for IBITS(POS=%d,LEN=%d) is negative"_err_en_US,
+                posVal, lenVal);
+          } else if (posVal + lenVal > i.bits) {
+            context.messages().Say(
+                "IBITS(POS=%d,LEN=%d) must have POS+LEN no greater than %d"_err_en_US,
+                posVal + lenVal, i.bits);
+          }
+          return i.IBITS(posVal, lenVal);
+        }));
   } else if (name == "index" || name == "scan" || name == "verify") {
     if (auto *charExpr{UnwrapExpr<Expr<SomeCharacter>>(args[0])}) {
       return std::visit(
@@ -949,7 +971,7 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
   }
-  // TODO: dot_product, ibits, ishftc, matmul, sign, transfer
+  // TODO: dot_product, ishftc, matmul, sign, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
diff --git a/flang/test/Evaluate/fold-ibits.f90 b/flang/test/Evaluate/fold-ibits.f90
new file mode 100644
index 0000000..b49fce9
--- /dev/null
+++ b/flang/test/Evaluate/fold-ibits.f90
@@ -0,0 +1,13 @@
+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of IBITS exhaustively over POS/LEN ranges
+module m1
+  implicit integer(a-z)
+  integer, parameter :: res1(*) = [((ibits(not(0),pos,len),len=0,31-pos),pos=0,31)]
+  integer, parameter :: expect1(*) = [((maskr(len),len=0,31-pos),pos=0,31)]
+  logical, parameter :: test1 = all(res1 == expect1)
+  logical, parameter :: test2 = all([((ibits(0,pos,len),len=0,31-pos),pos=0,31)] == 0)
+  integer, parameter :: mess = z'a5a55a5a'
+  integer, parameter :: res3(*) = [((ibits(mess,pos,len),len=0,31-pos),pos=0,31)]
+  integer, parameter :: expect3(*) = [((iand(shiftr(mess,pos),maskr(len)),len=0,31-pos),pos=0,31)]
+  logical, parameter :: test3 = all(res3 == expect3)
+end module
-- 
2.7.4