[flang] add helper to create descriptor with new base address (#161347)

There is currently no helper to create a descriptor for a copy of a
Fortran entity based on the descriptor of the original entity and the
base address of the copy (most places that are doing this currently are
also doing allocation of the copy at the same time or using the
runtime).
Add a helper for this with a unit test.

Author: jeanPerier

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -959,6 +959,15 @@ mlir::Value genLifetimeStart(mlir::OpBuilder &builder, mlir::Location loc,
 void genLifetimeEnd(mlir::OpBuilder &builder, mlir::Location loc,
                     mlir::Value mem);
 
+/// Given a fir.box or fir.class \p box describing an entity and a raw address
+/// \p newAddr for an entity with the same Fortran properties (rank, dynamic
+/// type, length parameters and bounds) and attributes (POINTER or ALLOCATABLE),
+/// create a box for \p newAddr with the same type as \p box. This assumes \p
+/// newAddr is for contiguous storage (\p box does not have to be contiguous).
+mlir::Value getDescriptorWithNewBaseAddress(fir::FirOpBuilder &builder,
+                                            mlir::Location loc, mlir::Value box,
+                                            mlir::Value newAddr);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -1974,3 +1974,25 @@ void fir::factory::genLifetimeEnd(mlir::OpBuilder &builder, mlir::Location loc,
                                   mlir::Value cast) {
   mlir::LLVM::LifetimeEndOp::create(builder, loc, cast);
 }
+
+mlir::Value fir::factory::getDescriptorWithNewBaseAddress(
+    fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value box,
+    mlir::Value newAddr) {
+  auto boxType = llvm::dyn_cast<fir::BaseBoxType>(box.getType());
+  assert(boxType &&
+         "expected a box type input in getDescriptorWithNewBaseAddress");
+  if (boxType.isAssumedRank())
+    TODO(loc, "changing descriptor base address for an assumed rank entity");
+  llvm::SmallVector<mlir::Value> lbounds;
+  fir::factory::genDimInfoFromBox(builder, loc, box, &lbounds,
+                                  /*extents=*/nullptr, /*strides=*/nullptr);
+  fir::BoxValue inputBoxValue(box, lbounds, /*explicitParams=*/{});
+  fir::ExtendedValue openedInput =
+      fir::factory::readBoxValue(builder, loc, inputBoxValue);
+  mlir::Value shape = fir::isArray(openedInput)
+                          ? builder.createShape(loc, openedInput)
+                          : mlir::Value{};
+  mlir::Value typeMold = fir::isPolymorphicType(boxType) ? box : mlir::Value{};
+  return builder.createBox(loc, boxType, newAddr, shape, /*slice=*/{},
+                           fir::getTypeParams(openedInput), typeMold);
+}

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

@@ -644,3 +644,50 @@ TEST_F(FIRBuilderTest, genArithIntegerOverflow) {
   auto op4_ioff = op4_iofi.getOverflowAttr().getValue();
   EXPECT_EQ(op4_ioff, nsw);
 }
+
+TEST_F(FIRBuilderTest, getDescriptorWithNewBaseAddress) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+
+  // Build an input fir.box for a 1-D array of i64 with constant extent 10.
+  auto i64Ty = builder.getI64Type();
+  auto seqTy = fir::SequenceType::get({10}, i64Ty);
+  auto refArrTy = fir::ReferenceType::get(seqTy);
+  auto ptrTy = fir::PointerType::get(seqTy);
+  auto boxTy = fir::BoxType::get(ptrTy);
+  // Create an undef box descriptor value (descriptor contents are unspecified).
+  mlir::Value inputBox = fir::UndefOp::create(builder, loc, boxTy);
+
+  // New base address (same element type and properties).
+  mlir::Value addr2 = fir::UndefOp::create(builder, loc, refArrTy);
+
+  mlir::Value newBox = fir::factory::getDescriptorWithNewBaseAddress(
+      builder, loc, inputBox, addr2);
+
+  // The returned descriptor must have the same type as the input box.
+  EXPECT_EQ(newBox.getType(), inputBox.getType());
+
+  // It must be constructed by an embox using the new base address.
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::EmboxOp>(newBox.getDefiningOp()));
+  auto embox = llvm::dyn_cast<fir::EmboxOp>(newBox.getDefiningOp());
+  EXPECT_EQ(embox.getMemref(), addr2);
+
+  // The shape should be derived from the input box; expect a fir.shape with one
+  // extent that comes from a fir.box_dims reading from the original input box.
+  mlir::Value shape = embox.getShape();
+  ASSERT_TRUE(shape);
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::ShapeShiftOp>(shape.getDefiningOp()));
+  auto shapeOp = llvm::dyn_cast<fir::ShapeShiftOp>(shape.getDefiningOp());
+  ASSERT_EQ(shapeOp.getExtents().size(), 1u);
+  mlir::Value extent0 = shapeOp.getExtents()[0];
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::BoxDimsOp>(extent0.getDefiningOp()));
+  auto dimOp = llvm::dyn_cast<fir::BoxDimsOp>(extent0.getDefiningOp());
+  EXPECT_EQ(dimOp.getVal(), inputBox);
+
+  // Also verify the origin comes from a BoxDims on the same input box.
+  ASSERT_EQ(shapeOp.getOrigins().size(), 1u);
+  mlir::Value origin0 = shapeOp.getOrigins()[0];
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::BoxDimsOp>(origin0.getDefiningOp()));
+  auto lbOp = llvm::dyn_cast<fir::BoxDimsOp>(origin0.getDefiningOp());
+  EXPECT_EQ(lbOp.getVal(), inputBox);
+}