[flang] harden TypeAndShape for assumed-ranks (llvm#96234)

SIZEOF and C_SIZEOF were broken for assumed-ranks because
`TypeAndShape::MeasureSizeInBytes` behaved as a scalar because the
`TypeAndShape::shape_` member was the same for scalar and assumed-ranks.

The easy fix would have been to add special handling in
`MeasureSizeInBytes` for assumed-ranks using the TypeAndShape
attributes, but I think this solution would leave `TypeAndShape::shape_`
manipulation fragile to future developers. Hence, I went for the
solution that turn shape_ into a `std::optional<Shape>`.

Author: jeanPerier

flang/include/flang/Evaluate/characteristics.h

@@ -55,26 +55,26 @@ std::optional<bool> DistinguishableOpOrAssign(
 // Shapes of function results and dummy arguments have to have
 // the same rank, the same deferred dimensions, and the same
 // values for explicit dimensions when constant.
-bool ShapesAreCompatible(
-    const Shape &, const Shape &, bool *possibleWarning = nullptr);
+bool ShapesAreCompatible(const std::optional<Shape> &,
+    const std::optional<Shape> &, bool *possibleWarning = nullptr);
 
 class TypeAndShape {
 public:
   ENUM_CLASS(
       Attr, AssumedRank, AssumedShape, AssumedSize, DeferredShape, Coarray)
   using Attrs = common::EnumSet<Attr, Attr_enumSize>;
 
-  explicit TypeAndShape(DynamicType t) : type_{t} { AcquireLEN(); }
-  TypeAndShape(DynamicType t, int rank) : type_{t}, shape_(rank) {
+  explicit TypeAndShape(DynamicType t) : type_{t}, shape_{Shape{}} {
+    AcquireLEN();
+  }
+  TypeAndShape(DynamicType t, int rank) : type_{t}, shape_{Shape(rank)} {
     AcquireLEN();
   }
   TypeAndShape(DynamicType t, Shape &&s) : type_{t}, shape_{std::move(s)} {
     AcquireLEN();
   }
   TypeAndShape(DynamicType t, std::optional<Shape> &&s) : type_{t} {
-    if (s) {
-      shape_ = std::move(*s);
-    }
+    shape_ = std::move(s);
     AcquireLEN();
   }
   DEFAULT_CONSTRUCTORS_AND_ASSIGNMENTS(TypeAndShape)
@@ -172,11 +172,12 @@ class TypeAndShape {
     LEN_ = std::move(len);
     return *this;
   }
-  const Shape &shape() const { return shape_; }
+  const std::optional<Shape> &shape() const { return shape_; }
   const Attrs &attrs() const { return attrs_; }
   int corank() const { return corank_; }
 
-  int Rank() const { return GetRank(shape_); }
+  // Return -1 for assumed-rank as a safety.
+  int Rank() const { return shape_ ? GetRank(*shape_) : -1; }
 
   // Can sequence association apply to this argument?
   bool CanBeSequenceAssociated() const {
@@ -211,7 +212,7 @@ class TypeAndShape {
 protected:
   DynamicType type_;
   std::optional<Expr<SubscriptInteger>> LEN_;
-  Shape shape_;
+  std::optional<Shape> shape_;
   Attrs attrs_;
   int corank_{0};
 };

flang/include/flang/Evaluate/shape.h

@@ -46,6 +46,13 @@ Constant<ExtentType> AsConstantShape(const ConstantSubscripts &);
 ConstantSubscripts AsConstantExtents(const Constant<ExtentType> &);
 std::optional<ConstantSubscripts> AsConstantExtents(
     FoldingContext &, const Shape &);
+inline std::optional<ConstantSubscripts> AsConstantExtents(
+    FoldingContext &foldingContext, const std::optional<Shape> &maybeShape) {
+  if (maybeShape) {
+    return AsConstantExtents(foldingContext, *maybeShape);
+  }
+  return std::nullopt;
+}
 Shape AsShape(const ConstantSubscripts &);
 std::optional<Shape> AsShape(const std::optional<ConstantSubscripts> &);
 
@@ -121,6 +128,12 @@ MaybeExtentExpr CountTrips(
 // Computes SIZE() == PRODUCT(shape)
 MaybeExtentExpr GetSize(Shape &&);
 ConstantSubscript GetSize(const ConstantSubscripts &);
+inline MaybeExtentExpr GetSize(const std::optional<Shape> &maybeShape) {
+  if (maybeShape) {
+    return GetSize(Shape(*maybeShape));
+  }
+  return std::nullopt;
+}
 
 // Utility predicate: does an expression reference any implied DO index?
 bool ContainsAnyImpliedDoIndex(const ExtentExpr &);

flang/lib/Evaluate/characteristics.cpp

@@ -39,13 +39,16 @@ static void CopyAttrs(const semantics::Symbol &src, A &dst,
 // Shapes of function results and dummy arguments have to have
 // the same rank, the same deferred dimensions, and the same
 // values for explicit dimensions when constant.
-bool ShapesAreCompatible(
-    const Shape &x, const Shape &y, bool *possibleWarning) {
-  if (x.size() != y.size()) {
+bool ShapesAreCompatible(const std::optional<Shape> &x,
+    const std::optional<Shape> &y, bool *possibleWarning) {
+  if (!x || !y) {
+    return !x && !y;
+  }
+  if (x->size() != y->size()) {
     return false;
   }
-  auto yIter{y.begin()};
-  for (const auto &xDim : x) {
+  auto yIter{y->begin()};
+  for (const auto &xDim : *x) {
     const auto &yDim{*yIter++};
     if (xDim && yDim) {
       if (auto equiv{AreEquivalentInInterface(*xDim, *yDim)}) {
@@ -178,9 +181,11 @@ bool TypeAndShape::IsCompatibleWith(parser::ContextualMessages &messages,
         thatIs, that.AsFortran(), thisIs, AsFortran());
     return false;
   }
-  return omitShapeConformanceCheck ||
-      CheckConformance(messages, shape_, that.shape_, flags, thisIs, thatIs)
-          .value_or(true /*fail only when nonconformance is known now*/);
+  return omitShapeConformanceCheck || (!shape_ && !that.shape_) ||
+      (shape_ && that.shape_ &&
+          CheckConformance(
+              messages, *shape_, *that.shape_, flags, thisIs, thatIs)
+              .value_or(true /*fail only when nonconformance is known now*/));
 }
 
 std::optional<Expr<SubscriptInteger>> TypeAndShape::MeasureElementSizeInBytes(
@@ -201,11 +206,11 @@ std::optional<Expr<SubscriptInteger>> TypeAndShape::MeasureElementSizeInBytes(
 
 std::optional<Expr<SubscriptInteger>> TypeAndShape::MeasureSizeInBytes(
     FoldingContext &foldingContext) const {
-  if (auto elements{GetSize(Shape{shape_})}) {
+  if (auto elements{GetSize(shape_)}) {
     // Sizes of arrays (even with single elements) are multiples of
     // their alignments.
     if (auto elementBytes{
-            MeasureElementSizeInBytes(foldingContext, GetRank(shape_) > 0)}) {
+            MeasureElementSizeInBytes(foldingContext, Rank() > 0)}) {
       return Fold(
           foldingContext, std::move(*elements) * std::move(*elementBytes));
     }
@@ -254,10 +259,12 @@ std::string TypeAndShape::AsFortran() const {
 llvm::raw_ostream &TypeAndShape::Dump(llvm::raw_ostream &o) const {
   o << type_.AsFortran(LEN_ ? LEN_->AsFortran() : "");
   attrs_.Dump(o, EnumToString);
-  if (!shape_.empty()) {
+  if (!shape_) {
+    o << " dimension(..)";
+  } else if (!shape_->empty()) {
     o << " dimension";
     char sep{'('};
-    for (const auto &expr : shape_) {
+    for (const auto &expr : *shape_) {
       o << sep;
       sep = ',';
       if (expr) {
@@ -1112,6 +1119,7 @@ bool FunctionResult::CanBeReturnedViaImplicitInterface(
 
 static std::optional<std::string> AreIncompatibleFunctionResultShapes(
     const Shape &x, const Shape &y) {
+  // Function results cannot be assumed-rank, hence the non optional arguments.
   int rank{GetRank(x)};
   if (int yrank{GetRank(y)}; yrank != rank) {
     return "rank "s + std::to_string(rank) + " vs " + std::to_string(yrank);
@@ -1147,7 +1155,8 @@ bool FunctionResult::IsCompatibleWith(
         }
       } else if (!attrs.test(Attr::Allocatable) && !attrs.test(Attr::Pointer) &&
           (details = AreIncompatibleFunctionResultShapes(
-               ifaceTypeShape->shape(), actualTypeShape->shape()))) {
+               ifaceTypeShape->shape().value(),
+               actualTypeShape->shape().value()))) {
         if (whyNot) {
           *whyNot = "function results have distinct extents (" + *details + ')';
         }

flang/lib/Evaluate/check-expression.cpp

@@ -419,7 +419,7 @@ std::optional<Expr<SomeType>> NonPointerInitializationExpr(const Symbol &symbol,
     if (converted) {
       auto folded{Fold(context, std::move(*converted))};
       if (IsActuallyConstant(folded)) {
-        int symRank{GetRank(symTS->shape())};
+        int symRank{symTS->Rank()};
         if (IsImpliedShape(symbol)) {
           if (folded.Rank() == symRank) {
             return ArrayConstantBoundChanger{
@@ -442,7 +442,8 @@ std::optional<Expr<SomeType>> NonPointerInitializationExpr(const Symbol &symbol,
                     context, GetRawLowerBounds(context, NamedEntity{symbol}))}
                 .Expand(std::move(folded));
           } else if (auto resultShape{GetShape(context, folded)}) {
-            if (CheckConformance(context.messages(), symTS->shape(),
+            CHECK(symTS->shape()); // Assumed-ranks cannot be initialized.
+            if (CheckConformance(context.messages(), *symTS->shape(),
                     *resultShape, CheckConformanceFlags::None,
                     "initialized object", "initialization expression")
                     .value_or(false /*fail if not known now to conform*/)) {

flang/lib/Evaluate/shape.cpp

@@ -1037,7 +1037,7 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
       } else if (context_) {
         if (auto moldTypeAndShape{characteristics::TypeAndShape::Characterize(
                 call.arguments().at(1), *context_)}) {
-          if (GetRank(moldTypeAndShape->shape()) == 0) {
+          if (moldTypeAndShape->Rank() == 0) {
             // SIZE= is absent and MOLD= is scalar: result is scalar
             return ScalarShape();
           } else {

flang/lib/Lower/CallInterface.cpp

@@ -177,9 +177,10 @@ mlir::Location Fortran::lower::CallerInterface::getCalleeLocation() const {
 // explicit.
 static Fortran::evaluate::characteristics::DummyDataObject
 asImplicitArg(Fortran::evaluate::characteristics::DummyDataObject &&dummy) {
-  Fortran::evaluate::Shape shape =
-      dummy.type.attrs().none() ? dummy.type.shape()
-                                : Fortran::evaluate::Shape(dummy.type.Rank());
+  std::optional<Fortran::evaluate::Shape> shape =
+      dummy.type.attrs().none()
+          ? dummy.type.shape()
+          : std::make_optional<Fortran::evaluate::Shape>(dummy.type.Rank());
   return Fortran::evaluate::characteristics::DummyDataObject(
       Fortran::evaluate::characteristics::TypeAndShape(dummy.type.type(),
                                                        std::move(shape)));
@@ -1308,18 +1309,17 @@ class Fortran::lower::CallInterfaceImpl {
   // with the shape (may contain unknown extents) for arrays.
   std::optional<fir::SequenceType::Shape> getBounds(
       const Fortran::evaluate::characteristics::TypeAndShape &typeAndShape) {
-    using ShapeAttr = Fortran::evaluate::characteristics::TypeAndShape::Attr;
-    if (typeAndShape.shape().empty() &&
-        !typeAndShape.attrs().test(ShapeAttr::AssumedRank))
+    if (typeAndShape.shape() && typeAndShape.shape()->empty())
       return std::nullopt;
     fir::SequenceType::Shape bounds;
-    for (const std::optional<Fortran::evaluate::ExtentExpr> &extent :
-         typeAndShape.shape()) {
-      fir::SequenceType::Extent bound = fir::SequenceType::getUnknownExtent();
-      if (std::optional<std::int64_t> i = toInt64(extent))
-        bound = *i;
-      bounds.emplace_back(bound);
-    }
+    if (typeAndShape.shape())
+      for (const std::optional<Fortran::evaluate::ExtentExpr> &extent :
+           *typeAndShape.shape()) {
+        fir::SequenceType::Extent bound = fir::SequenceType::getUnknownExtent();
+        if (std::optional<std::int64_t> i = toInt64(extent))
+          bound = *i;
+        bounds.emplace_back(bound);
+      }
     return bounds;
   }
   std::optional<std::int64_t>

flang/lib/Semantics/check-call.cpp

@@ -143,8 +143,8 @@ static void CheckCharacterActual(evaluate::Expr<evaluate::SomeType> &actual,
         bool canAssociate{CanAssociateWithStorageSequence(dummy)};
         if (dummy.type.Rank() > 0 && canAssociate) {
           // Character storage sequence association (F'2023 15.5.2.12p4)
-          if (auto dummySize{evaluate::ToInt64(evaluate::Fold(foldingContext,
-                  evaluate::GetSize(evaluate::Shape{dummy.type.shape()})))}) {
+          if (auto dummySize{evaluate::ToInt64(evaluate::Fold(
+                  foldingContext, evaluate::GetSize(dummy.type.shape())))}) {
             auto dummyChars{*dummySize * *dummyLength};
             if (actualType.Rank() == 0) {
               evaluate::DesignatorFolder folder{
@@ -183,8 +183,7 @@ static void CheckCharacterActual(evaluate::Expr<evaluate::SomeType> &actual,
               }
             } else { // actual.type.Rank() > 0
               if (auto actualSize{evaluate::ToInt64(evaluate::Fold(
-                      foldingContext,
-                      evaluate::GetSize(evaluate::Shape(actualType.shape()))))};
+                      foldingContext, evaluate::GetSize(actualType.shape())))};
                   actualSize &&
                   *actualSize * *actualLength < *dummySize * *dummyLength &&
                   (extentErrors ||
@@ -251,7 +250,7 @@ static void ConvertIntegerActual(evaluate::Expr<evaluate::SomeType> &actual,
   if (dummyType.type().category() == TypeCategory::Integer &&
       actualType.type().category() == TypeCategory::Integer &&
       dummyType.type().kind() != actualType.type().kind() &&
-      GetRank(dummyType.shape()) == 0 && GetRank(actualType.shape()) == 0 &&
+      dummyType.Rank() == 0 && actualType.Rank() == 0 &&
       !evaluate::IsVariable(actual)) {
     auto converted{
         evaluate::ConvertToType(dummyType.type(), std::move(actual))};
@@ -387,10 +386,10 @@ static void CheckExplicitDataArg(const characteristics::DummyDataObject &dummy,
       // if the actual argument is an array or array element designator,
       // and the dummy is an array, but not assumed-shape or an INTENT(IN)
       // pointer that's standing in for an assumed-shape dummy.
-    } else {
+    } else if (dummy.type.shape() && actualType.shape()) {
       // Let CheckConformance accept actual scalars; storage association
       // cases are checked here below.
-      CheckConformance(messages, dummy.type.shape(), actualType.shape(),
+      CheckConformance(messages, *dummy.type.shape(), *actualType.shape(),
           dummyIsAllocatableOrPointer
               ? evaluate::CheckConformanceFlags::None
               : evaluate::CheckConformanceFlags::RightScalarExpandable,
@@ -579,8 +578,8 @@ static void CheckExplicitDataArg(const characteristics::DummyDataObject &dummy,
         CanAssociateWithStorageSequence(dummy) &&
         !dummy.attrs.test(
             characteristics::DummyDataObject::Attr::DeducedFromActual)) {
-      if (auto dummySize{evaluate::ToInt64(evaluate::Fold(foldingContext,
-              evaluate::GetSize(evaluate::Shape{dummy.type.shape()})))}) {
+      if (auto dummySize{evaluate::ToInt64(evaluate::Fold(
+              foldingContext, evaluate::GetSize(dummy.type.shape())))}) {
         if (actualRank == 0 && !actualIsAssumedRank) {
           if (evaluate::IsArrayElement(actual)) {
             // Actual argument is a scalar array element
@@ -622,8 +621,8 @@ static void CheckExplicitDataArg(const characteristics::DummyDataObject &dummy,
             }
           }
         } else { // actualRank > 0 || actualIsAssumedRank
-          if (auto actualSize{evaluate::ToInt64(evaluate::Fold(foldingContext,
-                  evaluate::GetSize(evaluate::Shape(actualType.shape()))))};
+          if (auto actualSize{evaluate::ToInt64(evaluate::Fold(
+                  foldingContext, evaluate::GetSize(actualType.shape())))};
               actualSize && *actualSize < *dummySize &&
               (extentErrors ||
                   context.ShouldWarn(common::UsageWarning::ShortArrayActual))) {

flang/lib/Semantics/check-declarations.cpp

@@ -1325,6 +1325,13 @@ class SubprogramMatchHelper {
   bool CheckSameAttrs(const Symbol &, const Symbol &, ATTRS, ATTRS);
   bool ShapesAreCompatible(const DummyDataObject &, const DummyDataObject &);
   evaluate::Shape FoldShape(const evaluate::Shape &);
+  std::optional<evaluate::Shape> FoldShape(
+      const std::optional<evaluate::Shape> &shape) {
+    if (shape) {
+      return FoldShape(*shape);
+    }
+    return std::nullopt;
+  }
   std::string AsFortran(DummyDataObject::Attr attr) {
     return parser::ToUpperCaseLetters(DummyDataObject::EnumToString(attr));
   }

flang/lib/Semantics/pointer-assignment.cpp

@@ -333,8 +333,8 @@ bool PointerAssignmentChecker::Check(const evaluate::Designator<T> &d) {
 
       } else if (!isBoundsRemapping_ &&
           !lhsType_->attrs().test(TypeAndShape::Attr::AssumedRank)) {
-        int lhsRank{evaluate::GetRank(lhsType_->shape())};
-        int rhsRank{evaluate::GetRank(rhsType->shape())};
+        int lhsRank{lhsType_->Rank()};
+        int rhsRank{rhsType->Rank()};
         if (lhsRank != rhsRank) {
           msg = MessageFormattedText{
               "Pointer has rank %d but target has rank %d"_err_en_US, lhsRank,

flang/lib/Semantics/runtime-type-info.cpp

@@ -748,7 +748,7 @@ evaluate::StructureConstructor RuntimeTableBuilder::DescribeComponent(
       symbol, foldingContext)};
   CHECK(typeAndShape.has_value());
   auto dyType{typeAndShape->type()};
-  const auto &shape{typeAndShape->shape()};
+  int rank{typeAndShape->Rank()};
   AddValue(values, componentSchema_, "name"s,
       SaveNameAsPointerTarget(scope, symbol.name().ToString()));
   AddValue(values, componentSchema_, "category"s,
@@ -830,7 +830,6 @@ evaluate::StructureConstructor RuntimeTableBuilder::DescribeComponent(
         SomeExpr{evaluate::NullPointer{}});
   }
   // Shape information
-  int rank{evaluate::GetRank(shape)};
   AddValue(values, componentSchema_, "rank"s, IntExpr<1>(rank));
   if (rank > 0 && !IsAllocatable(symbol) && !IsPointer(symbol)) {
     std::vector<evaluate::StructureConstructor> bounds;
@@ -1143,7 +1142,7 @@ void RuntimeTableBuilder::DescribeSpecialProc(
           isArgDescriptorSet |= 1;
         } else {
           which = scalarFinalEnum_;
-          if (int rank{evaluate::GetRank(typeAndShape.shape())}; rank > 0) {
+          if (int rank{typeAndShape.Rank()}; rank > 0) {
             which = IntExpr<1>(ToInt64(which).value() + rank);
             if (dummyData.IsPassedByDescriptor(proc->IsBindC())) {
               argThatMightBeDescriptor = 1;

flang/test/Evaluate/rewrite06.f90

@@ -31,3 +31,9 @@ subroutine test(k)
     print *, storage_size(return_pdt(k))
   end subroutine
 end module
+
+subroutine test_assumed_rank(x)
+  real :: x(..)
+  !CHECK: PRINT *, sizeof(x)
+  print *, sizeof(x)
+end subroutine