[flang][hlfir] Fix hlfir.set_length codegen

[jeanPerier]

Summary

The bufferization pass was propagating the raw alloca storage
(which may not allow to later retrieve the length) instead of
the hlfir variable value (which is guaranteed to hold the
character length).

Fix this and makes packageBufferizedExpr "storage" argument and
getBufferizedExprStorage return an hlfir::Entity to avoid similar
error in the future (the caller of packageBufferizedExpr will have
to think a bit when adding the explicit hlfir::Entity{} cast).

[vzakhari]

LGTM. Thank you for the fix!

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

 namespace hlfir {
 #define GEN_PASS_DEF_BUFFERIZEHLFIR
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 } // namespace hlfir
 
 namespace {
 
 /// Helper to create tuple from a bufferized expr storage and clean up
-/// instruction flag.
+/// instruction flag. The storage is an HLFIR variable so that it can
+/// be manipulated as a variable later (all shape and length information
+/// cam be retrieved from it).
 static mlir::Value packageBufferizedExpr(mlir::Location loc,
                                          fir::FirOpBuilder &builder,
-                                         mlir::Value storage,
+                                         hlfir::Entity storage,
                                          mlir::Value mustFree) {
   auto tupleType = mlir::TupleType::get(
       builder.getContext(),
       mlir::TypeRange{storage.getType(), mustFree.getType()});
   auto undef = builder.create<fir::UndefOp>(loc, tupleType);
   auto insert = builder.create<fir::InsertValueOp>(
       loc, tupleType, undef, mustFree,
       builder.getArrayAttr(
           {builder.getIntegerAttr(builder.getIndexType(), 1)}));
   return builder.create<fir::InsertValueOp>(
       loc, tupleType, insert, storage,
       builder.getArrayAttr(
           {builder.getIntegerAttr(builder.getIndexType(), 0)}));
 }
 
 /// Helper to create tuple from a bufferized expr storage and constant
 /// boolean clean-up flag.
 static mlir::Value packageBufferizedExpr(mlir::Location loc,
                                          fir::FirOpBuilder &builder,
-                                         mlir::Value storage, bool mustFree) {
+                                         hlfir::Entity storage, bool mustFree) {
   mlir::Value mustFreeValue = builder.createBool(loc, mustFree);
   return packageBufferizedExpr(loc, builder, storage, mustFreeValue);
 }
 
 /// Helper to extract the storage from a tuple created by packageBufferizedExpr.
 /// It assumes no tuples are used as HLFIR operation operands, which is
 /// currently enforced by the verifiers that only accept HLFIR value or
 /// variable types which do not include tuples.
-static mlir::Value getBufferizedExprStorage(mlir::Value bufferizedExpr) {
+static hlfir::Entity getBufferizedExprStorage(mlir::Value bufferizedExpr) {
   auto tupleType = bufferizedExpr.getType().dyn_cast<mlir::TupleType>();
   if (!tupleType)
-    return bufferizedExpr;
+    return hlfir::Entity{bufferizedExpr};
   assert(tupleType.size() == 2 && "unexpected tuple type");
   if (auto insert = bufferizedExpr.getDefiningOp<fir::InsertValueOp>())
     if (insert.getVal().getType() == tupleType.getType(0))
-      return insert.getVal();
+      return hlfir::Entity{insert.getVal()};
   TODO(bufferizedExpr.getLoc(), "general extract storage case");
 }
 
 /// Helper to extract the clean-up flag from a tuple created by
 /// packageBufferizedExpr.
 static mlir::Value getBufferizedExprMustFreeFlag(mlir::Value bufferizedExpr) {
   auto tupleType = bufferizedExpr.getType().dyn_cast<mlir::TupleType>();
   if (!tupleType)
struct AsExprOpConversion : public mlir::OpConversionPattern<hlfir::AsExprOp> {
   matchAndRewrite(hlfir::AsExprOp asExpr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = asExpr->getLoc();
     auto module = asExpr->getParentOfType<mlir::ModuleOp>();
     fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
     if (asExpr.isMove()) {
       // Move variable storage for the hlfir.expr buffer.
       mlir::Value bufferizedExpr = packageBufferizedExpr(
-          loc, builder, adaptor.getVar(), adaptor.getMustFree());
+          loc, builder, hlfir::Entity{adaptor.getVar()}, adaptor.getMustFree());
       rewriter.replaceOp(asExpr, bufferizedExpr);
       return mlir::success();
     }
     // Otherwise, create a copy in a new buffer.
     hlfir::Entity source = hlfir::Entity{adaptor.getVar()};
     auto [temp, cleanup] = createTempFromMold(loc, builder, source);
     builder.create<hlfir::AssignOp>(loc, source, temp);
     mlir::Value bufferizedExpr =
         packageBufferizedExpr(loc, builder, temp, cleanup);
     rewriter.replaceOp(asExpr, bufferizedExpr);
     return mlir::success();
   }
 };
 
 struct ApplyOpConversion : public mlir::OpConversionPattern<hlfir::ApplyOp> {
   using mlir::OpConversionPattern<hlfir::ApplyOp>::OpConversionPattern;
   explicit ApplyOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::ApplyOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::ApplyOp apply, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = apply->getLoc();
-    hlfir::Entity bufferizedExpr{getBufferizedExprStorage(adaptor.getExpr())};
+    hlfir::Entity bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     mlir::Type resultType = hlfir::getVariableElementType(bufferizedExpr);
     mlir::Value result = rewriter.create<hlfir::DesignateOp>(
         loc, resultType, bufferizedExpr, adaptor.getIndices(),
         adaptor.getTypeparams());
     if (fir::isa_trivial(apply.getType()))
       result = rewriter.create<fir::LoadOp>(loc, result);
     rewriter.replaceOp(apply, result);
     return mlir::success();
matchAndRewrite(hlfir::ConcatOp concat, OpAdaptor adaptor,
     mlir::Location loc = concat->getLoc();
     auto module = concat->getParentOfType<mlir::ModuleOp>();
     fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
     assert(adaptor.getStrings().size() >= 2 &&
            "must have at least two strings operands");
     if (adaptor.getStrings().size() > 2)
       TODO(loc, "codegen of optimized chained concatenation of more than two "
                 "strings");
-    hlfir::Entity lhs{getBufferizedExprStorage(adaptor.getStrings()[0])};
-    hlfir::Entity rhs{getBufferizedExprStorage(adaptor.getStrings()[1])};
+    hlfir::Entity lhs = getBufferizedExprStorage(adaptor.getStrings()[0]);
+    hlfir::Entity rhs = getBufferizedExprStorage(adaptor.getStrings()[1]);
     auto [lhsExv, c1] = hlfir::translateToExtendedValue(loc, builder, lhs);
     auto [rhsExv, c2] = hlfir::translateToExtendedValue(loc, builder, rhs);
     assert(!c1 && !c2 && "expected variables");
     fir::ExtendedValue res =
         fir::factory::CharacterExprHelper{builder, loc}.createConcatenate(
             *lhsExv.getCharBox(), *rhsExv.getCharBox());
     // Ensure the memory type is the same as the result type.
     mlir::Type addrType = fir::ReferenceType::get(
         hlfir::getFortranElementType(concat.getResult().getType()));
     mlir::Value cast = builder.createConvert(loc, addrType, fir::getBase(res));
     res = fir::substBase(res, cast);
-    auto hlfirTempRes = hlfir::genDeclare(loc, builder, res, "tmp",
-                                          fir::FortranVariableFlagsAttr{})
-                            .getBase();
+    hlfir::Entity hlfirTempRes =
+        hlfir::Entity{hlfir::genDeclare(loc, builder, res, "tmp",
+                          .getBase()};
     mlir::Value bufferizedExpr =
         packageBufferizedExpr(loc, builder, hlfirTempRes, false);
     rewriter.replaceOp(concat, bufferizedExpr);
     return mlir::success();
   }
 };
 
 struct SetLengthOpConversion
     : public mlir::OpConversionPattern<hlfir::SetLengthOp> {
   using mlir::OpConversionPattern<hlfir::SetLengthOp>::OpConversionPattern;
   explicit SetLengthOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::SetLengthOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::SetLengthOp setLength, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = setLength->getLoc();
     auto module = setLength->getParentOfType<mlir::ModuleOp>();
     fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
     // Create a temp with the new length.
-    hlfir::Entity string{getBufferizedExprStorage(adaptor.getString())};
+    hlfir::Entity string = getBufferizedExprStorage(adaptor.getString());
     auto charType = hlfir::getFortranElementType(setLength.getType());
     llvm::StringRef tmpName{".tmp"};
     llvm::SmallVector<mlir::Value, 1> lenParams{adaptor.getLength()};
     auto alloca = builder.createTemporary(loc, charType, tmpName,
                                           /*shape=*/std::nullopt, lenParams);
     auto declareOp = builder.create<hlfir::DeclareOp>(
         loc, alloca, tmpName, /*shape=*/mlir::Value{}, lenParams,
         fir::FortranVariableFlagsAttr{});
+    hlfir::Entity temp{declareOp.getBase()};
     // Assign string value to the created temp.
-    builder.create<hlfir::AssignOp>(loc, string, declareOp.getBase());
+    builder.create<hlfir::AssignOp>(loc, string, temp);
     mlir::Value bufferizedExpr =
-        packageBufferizedExpr(loc, builder, alloca, false);
+        packageBufferizedExpr(loc, builder, temp, false);
     rewriter.replaceOp(setLength, bufferizedExpr);
     return mlir::success();
   }
 };
 
 static bool allOtherUsesAreDestroys(mlir::Value value,
                                     mlir::Operation *currentUse) {
   for (mlir::Operation *useOp : value.getUsers())
struct DestroyOpConversion
   using mlir::OpConversionPattern<hlfir::DestroyOp>::OpConversionPattern;
   explicit DestroyOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::DestroyOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::DestroyOp destroy, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     // If expr was bufferized on the heap, now is time to deallocate the buffer.
     mlir::Location loc = destroy->getLoc();
-    mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
+    hlfir::Entity bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     if (!fir::isa_trivial(bufferizedExpr.getType())) {
       auto module = destroy->getParentOfType<mlir::ModuleOp>();
       fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getExpr());
-      mlir::Value firBase = hlfir::Entity(bufferizedExpr).getFirBase();
+      mlir::Value firBase = bufferizedExpr.getFirBase();
       genFreeIfMustFree(loc, builder, firBase, mustFree);
     }
     rewriter.eraseOp(destroy);
     return mlir::success();
   }
 };
 
 struct NoReassocOpConversion

flang/test/HLFIR/set_length-codegen.fir

 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.boxchar<1>) {
 // CHECK:  %[[VAL_1:.*]] = fir.alloca !fir.char<1,10> {bindc_name = ".tmp"}
 // CHECK:  %[[VAL_2:.*]] = arith.constant 10 : index
 // CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_1]] typeparams %[[VAL_2]] {uniq_name = ".tmp"} : (!fir.ref<!fir.char<1,10>>, index) -> (!fir.ref<!fir.char<1,10>>, !fir.ref<!fir.char<1,10>>)
 // CHECK:  hlfir.assign %[[VAL_0]] to %[[VAL_3]]#0 : !fir.boxchar<1>, !fir.ref<!fir.char<1,10>>
 // CHECK:  %[[VAL_4:.*]] = arith.constant false
 // CHECK:  %[[VAL_5:.*]] = fir.undefined tuple<!fir.ref<!fir.char<1,10>>, i1>
 // CHECK:  %[[VAL_6:.*]] = fir.insert_value %[[VAL_5]], %[[VAL_4]], [1 : index] : (tuple<!fir.ref<!fir.char<1,10>>, i1>, i1) -> tuple<!fir.ref<!fir.char<1,10>>, i1>
-// CHECK:  %[[VAL_7:.*]] = fir.insert_value %[[VAL_6]], %[[VAL_1]], [0 : index] : (tuple<!fir.ref<!fir.char<1,10>>, i1>, !fir.ref<!fir.char<1,10>>) -> tuple<!fir.ref<!fir.char<1,10>>, i1>
+// CHECK:  %[[VAL_7:.*]] = fir.insert_value %[[VAL_6]], %[[VAL_3]]#0, [0 : index] : (tuple<!fir.ref<!fir.char<1,10>>, i1>, !fir.ref<!fir.char<1,10>>) -> tuple<!fir.ref<!fir.char<1,10>>, i1>
 
 func.func @test_dyn_len(%str : !fir.ref<!fir.char<1,10>>, %len : index) {
   %0 = hlfir.set_length %str len %len : (!fir.ref<!fir.char<1,10>>, index) -> !hlfir.expr<!fir.char<1,?>>
   return
 }
 // CHECK-LABEL:   func.func @test_dyn_len(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<!fir.char<1,10>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: index) {
 // CHECK:  %[[VAL_2:.*]] = fir.alloca !fir.char<1,?>(%[[VAL_1]] : index) {bindc_name = ".tmp"}
 // CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_2]] typeparams %[[VAL_1]] {uniq_name = ".tmp"} : (!fir.ref<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
 // CHECK:  hlfir.assign %[[VAL_0]] to %[[VAL_3]]#0 : !fir.ref<!fir.char<1,10>>, !fir.boxchar<1>
 // CHECK:  %[[VAL_4:.*]] = arith.constant false
-// CHECK:  %[[VAL_5:.*]] = fir.undefined tuple<!fir.ref<!fir.char<1,?>>, i1>
-// CHECK:  %[[VAL_6:.*]] = fir.insert_value %[[VAL_5]], %[[VAL_4]], [1 : index] : (tuple<!fir.ref<!fir.char<1,?>>, i1>, i1) -> tuple<!fir.ref<!fir.char<1,?>>, i1>
-// CHECK:  %[[VAL_7:.*]] = fir.insert_value %[[VAL_6]], %[[VAL_2]], [0 : index] : (tuple<!fir.ref<!fir.char<1,?>>, i1>, !fir.ref<!fir.char<1,?>>) -> tuple<!fir.ref<!fir.char<1,?>>, i1>
+// CHECK:  %[[VAL_5:.*]] = fir.undefined tuple<!fir.boxchar<1>, i1>
+// CHECK:  %[[VAL_6:.*]] = fir.insert_value %[[VAL_5]], %[[VAL_4]], [1 : index] : (tuple<!fir.boxchar<1>, i1>, i1) -> tuple<!fir.boxchar<1>, i1>
+// CHECK:  %[[VAL_7:.*]] = fir.insert_value %[[VAL_6]], %[[VAL_3]]#0, [0 : index] : (tuple<!fir.boxchar<1>, i1>, !fir.boxchar<1>) -> tuple<!fir.boxchar<1>, i1>