[flang] Add TypeAndShape::MeasureElementSizeInBytes()

[klausler]

Summary

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.

[PeteSteinfeld]

All builds, tests, and looks good.

flang/include/flang/Evaluate/characteristics.h

public:
   const Attrs &attrs() const { return attrs_; }
   int corank() const { return corank_; }
 
   int Rank() const { return GetRank(shape_); }
   bool IsCompatibleWith(parser::ContextualMessages &, const TypeAndShape &that,
       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;
 
   // called by Fold() to rewrite in place
   TypeAndShape &Rewrite(FoldingContext &);
 
   llvm::raw_ostream &Dump(llvm::raw_ostream &) const;
 

flang/lib/Evaluate/characteristics.cpp

if (!type_.IsTkCompatibleWith(that.type_)) {
     return false;
   }
   return isElemental ||
       CheckConformance(messages, shape_, that.shape_, thisIs, thatIs, false,
           false /* no scalar expansion */, thisIsDeferredShape,
           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));
     }
   }
   return std::nullopt;
 }
 
 void TypeAndShape::AcquireAttrs(const semantics::Symbol &symbol) {

flang/lib/Evaluate/fold-integer.cpp

Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
   } else if (name == "sizeof") { // in bytes; extension
     if (auto info{
             characteristics::TypeAndShape::Characterize(args[0], context)}) {
       if (auto bytes{info->MeasureSizeInBytes(context)}) {
         return Expr<T>{Fold(context, ConvertToType<T>(std::move(*bytes)))};
       }
     }
   } 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)}) {
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
   }
   // TODO:
   // cshift, dot_product, eoshift,
   // findloc, iall, iany, iparity, ibits, image_status, ishftc,
   // matmul, maxloc, maxval,
   // minloc, minval, not, pack, product, reduce,

flang/lib/Evaluate/shape.cpp

if (intrinsic->name == "shape" || intrinsic->name == "lbound" ||
             // SIZE= is absent and MOLD= is array: result is vector whose
             // length is determined by sizes of types.  See 16.9.193p4 case(ii).
             if (auto sourceTypeAndShape{
                     characteristics::TypeAndShape::Characterize(
                         call.arguments().at(0), *context_)}) {
               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) +
                         common::Clone(*moldElementBytes) - ExtentExpr{1}) /
                         common::Clone(*moldElementBytes))};
                 return Shape{MaybeExtentExpr{std::move(extent)}};
               }
             }

flang/lib/Semantics/compute-offsets.cpp

auto ComputeOffsetsHelper::GetSizeAndAlignment(
   if (auto chars{evaluate::characteristics::TypeAndShape::Characterize(
           symbol, foldingContext)}) {
     if (entire) {
       if (auto size{ToInt64(chars->MeasureSizeInBytes(foldingContext))}) {
         return {static_cast<std::size_t>(*size),
             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)};
       }
     }
   }
   return {};
 }