[flang][hlfir] Insert casts in hlfir.associate rewrites

[jeanPerier]

Summary

When the associated expression came from a moved variable, the
type of the moved variable may not exactly match the hlfir.associate
result and cannot be re-used directly. Insert fir.convert/fir.box_addr
as needed.

[clementval]

LGTM

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

struct AssociateOpConversion
     : public mlir::OpConversionPattern<hlfir::AssociateOp> {
   using mlir::OpConversionPattern<hlfir::AssociateOp>::OpConversionPattern;
   explicit AssociateOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AssociateOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AssociateOp associate, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = associate->getLoc();
-    // If this is the last use of the expression value and this is an hlfir.expr
-    // that was bufferized, re-use the storage.
-    // Otherwise, create a temp and assign the storage to it.
+    auto module = associate->getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
     mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getSource());
     const bool isTrivialValue = fir::isa_trivial(bufferizedExpr.getType());
 
     auto replaceWith = [&](mlir::Value hlfirVar, mlir::Value firVar,
                            mlir::Value flag) {
+      hlfirVar =
+          builder.createConvert(loc, associate.getResultTypes()[0], hlfirVar);
       associate.getResult(0).replaceAllUsesWith(hlfirVar);
+      mlir::Type associateFirVarType = associate.getResultTypes()[1];
+      if (firVar.getType().isa<fir::BaseBoxType>() &&
+          !associateFirVarType.isa<fir::BaseBoxType>())
+        firVar =
+            builder.create<fir::BoxAddrOp>(loc, associateFirVarType, firVar);
+      else
+        firVar = builder.createConvert(loc, associateFirVarType, firVar);
       associate.getResult(1).replaceAllUsesWith(firVar);
       associate.getResult(2).replaceAllUsesWith(flag);
       rewriter.replaceOp(associate, {hlfirVar, firVar, flag});
     };
 
+    // If this is the last use of the expression value and this is an hlfir.expr
+    // that was bufferized, re-use the storage.
+    // Otherwise, create a temp and assign the storage to it.
     if (!isTrivialValue && allOtherUsesAreDestroys(associate.getSource(),
                                                    associate.getOperation())) {
       // Re-use hlfir.expr buffer if this is the only use of the hlfir.expr
       // outside of the hlfir.destroy. Take on the cleaning-up responsibility
       // for the related hlfir.end_associate, and erase the hlfir.destroy (if
       // any).
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getSource());
       mlir::Value firBase = hlfir::Entity{bufferizedExpr}.getFirBase();
       replaceWith(bufferizedExpr, firBase, mustFree);
       eraseAllUsesInDestroys(associate.getSource(), rewriter);
       return mlir::success();
     }
     if (isTrivialValue) {
-      auto module = associate->getParentOfType<mlir::ModuleOp>();
-      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
       auto temp = builder.createTemporary(loc, bufferizedExpr.getType(),
                                           associate.getUniqName());
       builder.create<fir::StoreOp>(loc, bufferizedExpr, temp);
       mlir::Value mustFree = builder.createBool(loc, false);
       replaceWith(temp, temp, mustFree);
       return mlir::success();
     }
     TODO(loc, "hlfir.associate of hlfir.expr with more than one use");

flang/test/HLFIR/associate-codegen.fir

 // CHECK-LABEL:   func.func @test_end_associate_box_dynamic(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xf64>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: i1) {
 // CHECK:  fir.if %[[VAL_1]] {
 // CHECK:    %[[VAL_2:.*]] = fir.box_addr %[[VAL_0]] : (!fir.box<!fir.array<?xf64>>) -> !fir.heap<!fir.array<?xf64>>
 // CHECK:    fir.freemem %[[VAL_2]] : !fir.heap<!fir.array<?xf64>>
 // CHECK:  }
 
+func.func private @bar(!fir.ref<!fir.array<?xi32>>) -> ()
+func.func @test_result_box_addr(%x : !fir.box<!fir.array<?xi32>>) {
+  %true = arith.constant 1 : i1
+  %expr = hlfir.as_expr %x move %true : (!fir.box<!fir.array<?xi32>>, i1) -> !hlfir.expr<?xi32>
+  %y:3 = hlfir.associate %expr {uniq_name = "y"}: (!hlfir.expr<?xi32>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, i1)
+  fir.call @bar(%y#1) : (!fir.ref<!fir.array<?xi32>>) -> ()
+  return
+}
+// CHECK-LABEL: func.func @test_result_box_addr(
+// CHECK-SAME: %[[X:.*]]: !fir.box<!fir.array<?xi32>>) {
+// CHECK:  %[[ADDR:.*]] = fir.box_addr %[[X]] : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+// CHECK: fir.call @bar(%[[ADDR]]) : (!fir.ref<!fir.array<?xi32>>) -> ()
+
+func.func private @bar2(!fir.ref<!fir.array<10xi32>>) -> ()
+func.func @test_result_convert(%x : !fir.heap<!fir.array<10xi32>>) {
+  %true = arith.constant 1 : i1
+  %expr = hlfir.as_expr %x move %true : (!fir.heap<!fir.array<10xi32>>, i1) -> !hlfir.expr<10xi32>
+  %y:3 = hlfir.associate %expr {uniq_name = "y"}: (!hlfir.expr<10xi32>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>, i1)
+  fir.call @bar2(%y#1) : (!fir.ref<!fir.array<10xi32>>) -> ()
+  return
+}
+// CHECK-LABEL: func.func @test_result_convert(
+// CHECK-SAME: %[[X:.*]]: !fir.heap<!fir.array<10xi32>>) {
+// CHECK: fir.convert
+// CHECK:  %[[ADDR:.*]] = fir.convert %[[X]] : (!fir.heap<!fir.array<10xi32>>) -> !fir.ref<!fir.array<10xi32>>
+// CHECK: fir.call @bar2(%[[ADDR]]) : (!fir.ref<!fir.array<10xi32>>) -> ()
+
 
 func.func private @take_i4(!fir.ref<i32>)
 func.func private @take_r4(!fir.ref<f32>)
 func.func private @take_l4(!fir.ref<!fir.logical<4>>)
 func.func private @take_c(!fir.boxchar<1>)