[flang][hlfir] Fixed AssociateOp codegen for 0-dim variables.

The codegen tried to fir.convert !fir.box to !fir.ref for this case.
I used BoxAddr under a check for the type mismatch, but I am not sure
if this is the right fix. Maybe it has to be handled in the lowering.

2 files changed, 41 insertions, 2 deletions

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
index 3b30feb3b65f..7b12fad984f4 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -346,9 +346,27 @@ struct AssociateOpConversion
 
     auto replaceWith = [&](mlir::Value hlfirVar, mlir::Value firVar,
                            mlir::Value flag) {
-      hlfirVar =
-          builder.createConvert(loc, associate.getResultTypes()[0], hlfirVar);
+      // 0-dim variables may need special handling:
+      //   %0 = hlfir.as_expr %x move %true :
+      //       (!fir.box<!fir.heap<!fir.type<_T{y:i32}>>>, i1) ->
+      //       !hlfir.expr<!fir.type<_T{y:i32}>>
+      //   %1:3 = hlfir.associate %0 {uniq_name = "adapt.valuebyref"} :
+      //       (!hlfir.expr<!fir.type<_T{y:i32}>>) ->
+      //       (!fir.ref<!fir.type<_T{y:i32}>>,
+      //        !fir.ref<!fir.type<_T{y:i32}>>,
+      //        i1)
+      //
+      // !fir.box<!fir.heap<!fir.type<_T{y:i32}>>> value must be propagated
+      // as the box address !fir.ref<!fir.type<_T{y:i32}>>.
+      mlir::Type associateHlfirVarType = associate.getResultTypes()[0];
+      if (hlfirVar.getType().isa<fir::BaseBoxType>() &&
+          !associateHlfirVarType.isa<fir::BaseBoxType>())
+        hlfirVar = builder.create<fir::BoxAddrOp>(loc, associateHlfirVarType,
+                                                  hlfirVar);
+      else
+        hlfirVar = builder.createConvert(loc, associateHlfirVarType, hlfirVar);
       associate.getResult(0).replaceAllUsesWith(hlfirVar);
+
       mlir::Type associateFirVarType = associate.getResultTypes()[1];
       if ((firVar.getType().isa<fir::BaseBoxType>() &&
            !associateFirVarType.isa<fir::BaseBoxType>()) ||
diff --git a/flang/test/HLFIR/associate-codegen.fir b/flang/test/HLFIR/associate-codegen.fir
index 7b554a1fa882..5127f78e783c 100644
--- a/flang/test/HLFIR/associate-codegen.fir
+++ b/flang/test/HLFIR/associate-codegen.fir
@@ -173,6 +173,27 @@ func.func @test_result_convert(%x : !fir.heap<!fir.array<10xi32>>) {
 // CHECK: fir.call @bar2(%[[ADDR]]) : (!fir.ref<!fir.array<10xi32>>) -> ()
 
 
+func.func @test_0dim_box(%x : !fir.ref<!fir.box<!fir.heap<i32>>>) {
+  %0 = fir.load %x : !fir.ref<!fir.box<!fir.heap<i32>>>
+  %1:2 = hlfir.declare %0 {uniq_name = ".tmp.intrinsic_result"} : (!fir.box<!fir.heap<i32>>) -> (!fir.box<!fir.heap<i32>>, !fir.box<!fir.heap<i32>>)
+  %true = arith.constant true
+  %2 = hlfir.as_expr %1#0 move %true : (!fir.box<!fir.heap<i32>>, i1) -> !hlfir.expr<i32>
+  %3:3 = hlfir.associate %2 {uniq_name = "adapt.valuebyref"} : (!hlfir.expr<i32>) -> (!fir.ref<i32>, !fir.ref<i32>, i1)
+  return
+}
+// CHECK-LABEL:   func.func @test_0dim_box(
+// CHECK-SAME:                             %[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.heap<i32>>>) {
+// CHECK:           %[[VAL_1:.*]] = fir.load %[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<i32>>>
+// CHECK:           %[[VAL_2:.*]]:2 = hlfir.declare %[[VAL_1]] {uniq_name = ".tmp.intrinsic_result"} : (!fir.box<!fir.heap<i32>>) -> (!fir.box<!fir.heap<i32>>, !fir.box<!fir.heap<i32>>)
+// CHECK:           %[[VAL_3:.*]] = arith.constant true
+// CHECK:           %[[VAL_4:.*]] = fir.undefined tuple<!fir.box<!fir.heap<i32>>, i1>
+// CHECK:           %[[VAL_5:.*]] = fir.insert_value %[[VAL_4]], %[[VAL_3]], [1 : index] : (tuple<!fir.box<!fir.heap<i32>>, i1>, i1) -> tuple<!fir.box<!fir.heap<i32>>, i1>
+// CHECK:           %[[VAL_6:.*]] = fir.insert_value %[[VAL_5]], %[[VAL_2]]#0, [0 : index] : (tuple<!fir.box<!fir.heap<i32>>, i1>, !fir.box<!fir.heap<i32>>) -> tuple<!fir.box<!fir.heap<i32>>, i1>
+// CHECK:           %[[VAL_7:.*]] = fir.box_addr %[[VAL_2]]#0 : (!fir.box<!fir.heap<i32>>) -> !fir.ref<i32>
+// CHECK:           %[[VAL_8:.*]] = fir.box_addr %[[VAL_2]]#1 : (!fir.box<!fir.heap<i32>>) -> !fir.ref<i32>
+// CHECK:           return
+// CHECK:         }
+
 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>>)