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

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.

Differential Revision: https://reviews.llvm.org/D148310

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

@@ -196,9 +196,12 @@ struct AssignOpConversion : public mlir::OpConversionPattern<hlfir::AssignOp> {
   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();
   }
 };

flang/test/HLFIR/assign-bufferize.fir

@@ -0,0 +1,16 @@
+// 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>>>>>