[flang-commits] [flang] [flang] handle scalars in getDescriptorWithNewBaseAddress (PR #161515)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/161515

Follow up on #161347 to allow scalar fir.box/class reconstruction (at least required for polymorphic types).
The assert in genDimInfoFromBox was rejecting scalars while there is no functional reason for that (only assumed-rank are an issue there).

>From e3b1d2cac63be74ad649f0231e5331354f76fc22 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 1 Oct 2025 05:57:56 -0700
Subject: [PATCH] [flang] handle scalars in getDescriptorWithNewBaseAddress

---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp    |  2 +-
 .../Optimizer/Builder/FIRBuilderTest.cpp      | 37 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index 5e6e20861fd85..5da27d1713825 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -1943,7 +1943,7 @@ void fir::factory::genDimInfoFromBox(
     return;
 
   unsigned rank = fir::getBoxRank(boxType);
-  assert(rank != 0 && "must be an array of known rank");
+  assert(!boxType.isAssumedRank() && "must be an array of known rank");
   mlir::Type idxTy = builder.getIndexType();
   for (unsigned i = 0; i < rank; ++i) {
     mlir::Value dim = builder.createIntegerConstant(loc, idxTy, i);
diff --git a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
index fffd4ab5446ca..10a7ddf339133 100644
--- a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
+++ b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
@@ -691,3 +691,40 @@ TEST_F(FIRBuilderTest, getDescriptorWithNewBaseAddress) {
   auto lbOp = llvm::dyn_cast<fir::BoxDimsOp>(origin0.getDefiningOp());
   EXPECT_EQ(lbOp.getVal(), inputBox);
 }
+
+TEST_F(FIRBuilderTest, getDescriptorWithNewBaseAddress_PolymorphicScalar) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+
+  // Build a polymorphic scalar: fir.class<ptr<!fir.type<rec>>>.
+  auto recTy = fir::RecordType::get(builder.getContext(), "poly_rec");
+  auto ptrRecTy = fir::PointerType::get(recTy);
+  auto classTy = fir::ClassType::get(ptrRecTy);
+
+  // Input descriptor is an undefined fir.class value.
+  mlir::Value inputBox = fir::UndefOp::create(builder, loc, classTy);
+
+  // New base address of the same element type (reference to the record).
+  auto refRecTy = fir::ReferenceType::get(recTy);
+  mlir::Value newAddr = fir::UndefOp::create(builder, loc, refRecTy);
+
+  mlir::Value newBox = fir::factory::getDescriptorWithNewBaseAddress(
+      builder, loc, inputBox, newAddr);
+
+  // Same descriptor type must be preserved.
+  EXPECT_EQ(newBox.getType(), inputBox.getType());
+
+  // Must be an embox using the new base address and carrying the original box
+  // as mold.
+  ASSERT_TRUE(llvm::isa_and_nonnull<fir::EmboxOp>(newBox.getDefiningOp()));
+  auto embox = llvm::dyn_cast<fir::EmboxOp>(newBox.getDefiningOp());
+  EXPECT_EQ(embox.getMemref(), newAddr);
+
+  // Polymorphic scalar should have no shape operand.
+  mlir::Value shape = embox.getShape();
+  EXPECT_TRUE(shape == nullptr);
+
+  // The type descriptor/mold must be the original input box.
+  mlir::Value tdesc = embox.getSourceBox();
+  EXPECT_EQ(tdesc, inputBox);
+}