[flang] Fix for array upper bounds with *

[clementval]

Summary

Even though the array is declared with '*' upper bounds, it has an
initial value that has a statically known shape. Use the shape from
the type of the initializer when the declared size is '*'.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Eric Schweitz <[e-mail]>

[jeanPerier]

LGTM

flang/lib/Lower/ConvertVariable.cpp

 static mlir::Value genExtentValue(fir::FirOpBuilder &builder,
                                   mlir::Location loc, mlir::Type idxTy,
                                   long frontEndExtent) {
   if (frontEndExtent >= 0)
     return builder.createIntegerConstant(loc, idxTy, frontEndExtent);
   return builder.create<fir::UndefOp>(loc, idxTy);
 }
 
+/// If a symbol is an array, it may have been declared with unknown extent
+/// parameters (e.g., `*`), but if it has an initial value then the actual size
+/// may be available from the initial array value's type.
+inline static llvm::SmallVector<std::int64_t>
+recoverShapeVector(llvm::ArrayRef<std::int64_t> shapeVec, mlir::Value initVal) {
+  llvm::SmallVector<std::int64_t> result;
+  if (initVal) {
+    if (auto seqTy = fir::unwrapUntilSeqType(initVal.getType())) {
+      for (auto [fst, snd] : llvm::zip(shapeVec, seqTy.getShape()))
+        result.push_back(fst == fir::SequenceType::getUnknownExtent() ? snd
+                                                                      : fst);
+      return result;
+    }
+  }
+  result.assign(shapeVec.begin(), shapeVec.end());
+  return result;
+}
+
 /// Lower specification expressions and attributes of variable \p var and
 /// add it to the symbol map.  For a global or an alias, the address must be
 /// pre-computed and provided in \p preAlloc.  A dummy argument for the current
 /// entry point has already been mapped to an mlir block argument in
 /// mapDummiesAndResults.  Its mapping may be updated here.
 void Fortran::lower::mapSymbolAttributes(
     AbstractConverter &converter, const Fortran::lower::pft::Variable &var,
     Fortran::lower::SymMap &symMap, Fortran::lower::StatementContext &stmtCtx,
ba.match(
         // object shape is constant, not a character
         mlir::Type castTy = builder.getRefType(converter.genType(var));
         mlir::Value addr = symMap.lookupSymbol(sym).getAddr();
         if (addr)
           addr = builder.createConvert(loc, castTy, addr);
         if (x.lboundAllOnes()) {
           // if lower bounds are all ones, build simple shaped object
           llvm::SmallVector<mlir::Value> shape;
-          for (int64_t i : x.shapes)
+          for (int64_t i : recoverShapeVector(x.shapes, preAlloc))
             shape.push_back(genExtentValue(builder, loc, idxTy, i));
           mlir::Value local =
               isDummy ? addr : createNewLocal(converter, loc, var, preAlloc);
           symMap.addSymbolWithShape(sym, local, shape, isDummy);
           return;
         }
         // If object is an array process the lower bound and extent values by
         // constructing constants and populating the lbounds and extents.
         llvm::SmallVector<mlir::Value> extents;
         llvm::SmallVector<mlir::Value> lbounds;
-        for (auto [fst, snd] : llvm::zip(x.lbounds, x.shapes)) {
+        for (auto [fst, snd] :
+             llvm::zip(x.lbounds, recoverShapeVector(x.shapes, preAlloc))) {
           lbounds.emplace_back(builder.createIntegerConstant(loc, idxTy, fst));
           extents.emplace_back(genExtentValue(builder, loc, idxTy, snd));
         }
         mlir::Value local =
             isDummy ? addr
                     : createNewLocal(converter, loc, var, preAlloc, extents);
-        assert(isDummy || Fortran::lower::isExplicitShape(sym));
+        // Must be a dummy argument, have an explicit shape, or be a PARAMETER.
+        assert(isDummy || Fortran::lower::isExplicitShape(sym) ||
+               Fortran::semantics::IsNamedConstant(sym));
         symMap.addSymbolWithBounds(sym, local, extents, lbounds, isDummy);
       },
 
       //===--------------------------------------------------------------===//
 
       [&](const Fortran::lower::details::DynamicArray &x) {
         if (genUnusedEntryPointBox())
           return;
ba.match(
         // object shape is constant
         mlir::Type castTy = builder.getRefType(converter.genType(var));
         if (addr)
           addr = builder.createConvert(loc, castTy, addr);
 
         if (x.lboundAllOnes()) {
           // if lower bounds are all ones, build simple shaped object
           llvm::SmallVector<mlir::Value> shape;
-          for (int64_t i : x.shapes)
+          for (int64_t i : recoverShapeVector(x.shapes, preAlloc))
             shape.push_back(genExtentValue(builder, loc, idxTy, i));
           mlir::Value local =
               isDummy ? addr : createNewLocal(converter, loc, var, preAlloc);
           symMap.addCharSymbolWithShape(sym, local, len, shape, isDummy);
           return;
         }
 
         // if object is an array process the lower bound and extent values
         llvm::SmallVector<mlir::Value> extents;
         llvm::SmallVector<mlir::Value> lbounds;
         // construct constants and populate `bounds`
-        for (auto [fst, snd] : llvm::zip(x.lbounds, x.shapes)) {
+        for (auto [fst, snd] :
+             llvm::zip(x.lbounds, recoverShapeVector(x.shapes, preAlloc))) {
           lbounds.emplace_back(builder.createIntegerConstant(loc, idxTy, fst));
           extents.emplace_back(genExtentValue(builder, loc, idxTy, snd));
         }
 
         if (isDummy) {
           symMap.addCharSymbolWithBounds(sym, addr, len, extents, lbounds,
                                          true);
           return;
ba.match(
         // cast to the known constant parts from the declaration
         mlir::Type castTy = builder.getRefType(converter.genType(var));
         if (addr)
           addr = builder.createConvert(loc, castTy, addr);
 
         if (x.lboundAllOnes()) {
           // if lower bounds are all ones, build simple shaped object
           llvm::SmallVector<mlir::Value> shape;
-          for (int64_t i : x.shapes)
+          for (int64_t i : recoverShapeVector(x.shapes, preAlloc))
             shape.push_back(genExtentValue(builder, loc, idxTy, i));
           if (isDummy) {
             symMap.addCharSymbolWithShape(sym, addr, len, shape, true);
             return;
           }
           // local CHARACTER array with constant size
           mlir::Value local = createNewLocal(converter, loc, var, preAlloc,
                                              llvm::None, lengths);
           symMap.addCharSymbolWithShape(sym, local, len, shape);
           return;
         }
 
         // if object is an array process the lower bound and extent values
         llvm::SmallVector<mlir::Value> extents;
         llvm::SmallVector<mlir::Value> lbounds;
 
         // construct constants and populate `bounds`
-        for (auto [fst, snd] : llvm::zip(x.lbounds, x.shapes)) {
+        for (auto [fst, snd] :
+             llvm::zip(x.lbounds, recoverShapeVector(x.shapes, preAlloc))) {
           lbounds.emplace_back(builder.createIntegerConstant(loc, idxTy, fst));
           extents.emplace_back(genExtentValue(builder, loc, idxTy, snd));
         }
         if (isDummy) {
           symMap.addCharSymbolWithBounds(sym, addr, len, extents, lbounds,
                                          true);
           return;
         }

flang/test/Lower/memory-alloc.f90

+! RUN: bbc -o - %s | FileCheck %s
+
+! CHECK-LABEL: func @_QMw0bPtest1(
+! CHECK: %[[TWO:.*]] = arith.constant 2 : index
+! CHECK: %[[HEAP:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[TWO]] {uniq_name = ".array.expr"}
+! CHECK: fir.freemem %[[HEAP]] : !fir.heap<!fir.array<?x!fir.logical<4>>>
+
+Module w0b
+  Integer,Parameter :: a(*,*) = Reshape( [ 1,2,3,4 ], [ 2,2 ])
+contains
+  Subroutine test1(i,expect)
+    Integer,Intent(In) :: i,expect(:)
+    Logical :: ok = .True.
+    If (Any(a(:,i)/=expect)) Then
+      !Print *,'FAIL 1:',a(:,i),'/=',expect
+      ok = .False.
+    End If
+  End Subroutine
+End Module