[flang] Add TypeAndShape::MeasureElementSizeInBytes()

Split up MeasureSizeInBytes() so that array element sizes can be
calculated accurately; use the new API in some places where
DynamicType::MeasureSizeInBytes() was being used but the new
API performs better due to TypeAndShape having precise CHARACTER
length information.

Differential Revision: https://reviews.llvm.org/D95897

5 files changed, 23 insertions, 15 deletions

diff --git a/flang/include/flang/Evaluate/characteristics.h b/flang/include/flang/Evaluate/characteristics.h
index f18a220..7b37557 100644
--- a/flang/include/flang/Evaluate/characteristics.h
+++ b/flang/include/flang/Evaluate/characteristics.h
@@ -146,6 +146,8 @@ public:
       const char *thisIs = "pointer", const char *thatIs = "target",
       bool isElemental = false, bool thisIsDeferredShape = false,
       bool thatIsDeferredShape = false) const;
+  std::optional<Expr<SubscriptInteger>> MeasureElementSizeInBytes(
+      FoldingContext &, bool align) const;
   std::optional<Expr<SubscriptInteger>> MeasureSizeInBytes(
       FoldingContext &) const;
 
diff --git a/flang/lib/Evaluate/characteristics.cpp b/flang/lib/Evaluate/characteristics.cpp
index e53058a..92b93e7 100644
--- a/flang/lib/Evaluate/characteristics.cpp
+++ b/flang/lib/Evaluate/characteristics.cpp
@@ -168,19 +168,26 @@ bool TypeAndShape::IsCompatibleWith(parser::ContextualMessages &messages,
           thatIsDeferredShape);
 }
 
+std::optional<Expr<SubscriptInteger>> TypeAndShape::MeasureElementSizeInBytes(
+    FoldingContext &foldingContext, bool align) const {
+  if (LEN_) {
+    CHECK(type_.category() == TypeCategory::Character);
+    return Fold(foldingContext,
+        Expr<SubscriptInteger>{type_.kind()} * Expr<SubscriptInteger>{*LEN_});
+  }
+  if (auto elementBytes{type_.MeasureSizeInBytes(foldingContext, align)}) {
+    return Fold(foldingContext, std::move(*elementBytes));
+  }
+  return std::nullopt;
+}
+
 std::optional<Expr<SubscriptInteger>> TypeAndShape::MeasureSizeInBytes(
     FoldingContext &foldingContext) const {
   if (auto elements{GetSize(Shape{shape_})}) {
     // Sizes of arrays (even with single elements) are multiples of
     // their alignments.
-    if (LEN_) {
-      CHECK(type_.category() == TypeCategory::Character);
-      return Fold(foldingContext,
-          std::move(*elements) * Expr<SubscriptInteger>{type_.kind()} *
-              Expr<SubscriptInteger>{*LEN_});
-    }
     if (auto elementBytes{
-            type_.MeasureSizeInBytes(foldingContext, GetRank(shape_) > 0)}) {
+            MeasureElementSizeInBytes(foldingContext, GetRank(shape_) > 0)}) {
       return Fold(
           foldingContext, std::move(*elements) * std::move(*elementBytes));
     }
diff --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 8770519..9f3482d 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -598,12 +598,11 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
       }
     }
   } else if (name == "storage_size") { // in bits
-    if (const auto *expr{UnwrapExpr<Expr<SomeType>>(args[0])}) {
-      if (auto type{expr->GetType()}) {
-        if (auto bytes{type->MeasureSizeInBytes(context, true)}) {
-          return Expr<T>{
-              Fold(context, Expr<T>{8} * ConvertToType<T>(std::move(*bytes)))};
-        }
+    if (auto info{
+            characteristics::TypeAndShape::Characterize(args[0], context)}) {
+      if (auto bytes{info->MeasureElementSizeInBytes(context, true)}) {
+        return Expr<T>{
+            Fold(context, Expr<T>{8} * ConvertToType<T>(std::move(*bytes)))};
       }
     }
   } else if (name == "ubound") {
diff --git a/flang/lib/Evaluate/shape.cpp b/flang/lib/Evaluate/shape.cpp
index 9652d0c..217270e 100644
--- a/flang/lib/Evaluate/shape.cpp
+++ b/flang/lib/Evaluate/shape.cpp
@@ -724,7 +724,7 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
               auto sourceBytes{
                   sourceTypeAndShape->MeasureSizeInBytes(*context_)};
               auto moldElementBytes{
-                  moldTypeAndShape->type().MeasureSizeInBytes(*context_, true)};
+                  moldTypeAndShape->MeasureElementSizeInBytes(*context_, true)};
               if (sourceBytes && moldElementBytes) {
                 ExtentExpr extent{Fold(*context_,
                     (std::move(*sourceBytes) +
diff --git a/flang/lib/Semantics/compute-offsets.cpp b/flang/lib/Semantics/compute-offsets.cpp
index a11ec2b..bb2f4d9 100644
--- a/flang/lib/Semantics/compute-offsets.cpp
+++ b/flang/lib/Semantics/compute-offsets.cpp
@@ -323,7 +323,7 @@ auto ComputeOffsetsHelper::GetSizeAndAlignment(
             chars->type().GetAlignment(foldingContext)};
       }
     } else { // element size only
-      if (auto size{ToInt64(chars->type().MeasureSizeInBytes(
+      if (auto size{ToInt64(chars->MeasureElementSizeInBytes(
               foldingContext, true /*aligned*/))}) {
         return {static_cast<std::size_t>(*size),
             chars->type().GetAlignment(foldingContext)};