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

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.

Reviewed By: clementval

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

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

@@ -294,20 +294,31 @@ struct AssociateOpConversion
   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
@@ -321,8 +332,6 @@ struct AssociateOpConversion
       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);

flang/test/HLFIR/associate-codegen.fir

@@ -115,6 +115,33 @@ func.func @test_end_associate_box_dynamic(%var: !fir.box<!fir.array<?xf64>>, %mu
 // 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>)