[flang][hlfir] Preserve hlfir.assign in the bufferize-hlfir pass

[jeanPerier]

Summary

hlfir.assign must be rewritten in the bufferize pass since its operands,
that can be expressions, may have been updated. This is just
an operand update rewrite. The previous code was replacing the
operation, but it was dropping all hlfir.assign on the floor doing
so. This broke allocatable assignment semantics that use attributes.

Update the operands in place instead to preserve the attributes, if any.

[vzakhari]

LGTM

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

 
 struct AssignOpConversion : public mlir::OpConversionPattern<hlfir::AssignOp> {
   using mlir::OpConversionPattern<hlfir::AssignOp>::OpConversionPattern;
   explicit AssignOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AssignOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AssignOp assign, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<hlfir::AssignOp>(
-        assign, getBufferizedExprStorage(adaptor.getOperands()[0]),
-        getBufferizedExprStorage(adaptor.getOperands()[1]));
+    llvm::SmallVector<mlir::Value> newOperands;
+    for (mlir::Value operand : adaptor.getOperands())
+      newOperands.push_back(getBufferizedExprStorage(operand));
+    rewriter.startRootUpdate(assign);
+    assign->setOperands(newOperands);
+    rewriter.finalizeRootUpdate(assign);
     return mlir::success();
   }
 };
 
 struct ConcatOpConversion : public mlir::OpConversionPattern<hlfir::ConcatOp> {
   using mlir::OpConversionPattern<hlfir::ConcatOp>::OpConversionPattern;
   explicit ConcatOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::ConcatOp>{ctx} {}

flang/test/HLFIR/assign-bufferize.fir

+// Test hlfir.assign rewrite in the bufferization pass.
+// Assign in itself is not transformed, but its operands may be
+// expressions that are bufferized and must be updated.
+// RUN: fir-opt %s -bufferize-hlfir | FileCheck %s
+
+func.func @keep_attributes(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,8>>>>>, %arg1: !fir.box<!fir.array<1x!fir.char<1,?>>>) {
+  %true = arith.constant true
+  %0 = hlfir.as_expr %arg1 move %true : (!fir.box<!fir.array<1x!fir.char<1,?>>>, i1) -> !hlfir.expr<1x!fir.char<1,?>>
+  hlfir.assign %0 to %arg0 realloc keep_lhs_len : !hlfir.expr<1x!fir.char<1,?>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,8>>>>>
+  return
+}
+
+// CHECK-LABEL: func.func @keep_attributes(
+// CHECK-SAME: %[[X:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,8>>>>>,
+// CHECK-SAME: %[[Y:.*]]: !fir.box<!fir.array<1x!fir.char<1,?>>>) {
+// CHECK: hlfir.assign %[[Y]] to %[[X]] realloc keep_lhs_len : !fir.box<!fir.array<1x!fir.char<1,?>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.char<1,8>>>>>