[flang][hlfir] Handle box and non constant must_free in end_associate.

[clementval]

LGTM

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

if (isTrivialValue) {
       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");
   }
 };
 
-static void genFreeIfMustFree(mlir::Location loc,
-                              mlir::ConversionPatternRewriter &rewriter,
+static void genFreeIfMustFree(mlir::Location loc, fir::FirOpBuilder &builder,
                               mlir::Value var, mlir::Value mustFree) {
   auto genFree = [&]() {
-    if (var.getType().isa<fir::BaseBoxType>())
-      TODO(loc, "unbox");
-    if (!var.getType().isa<fir::HeapType>())
-      var = rewriter.create<fir::ConvertOp>(
-          loc, fir::HeapType::get(fir::unwrapRefType(var.getType())), var);
-    rewriter.create<fir::FreeMemOp>(loc, var);
+    // fir::FreeMemOp operand type must be a fir::HeapType.
+    mlir::Type heapType = fir::HeapType::get(
+        hlfir::getFortranElementOrSequenceType(var.getType()));
+    if (var.getType().isa<fir::BaseBoxType, fir::BoxCharType>())
+      var = builder.create<fir::BoxAddrOp>(loc, heapType, var);
+    else if (!var.getType().isa<fir::HeapType>())
+      var = builder.create<fir::ConvertOp>(loc, heapType, var);
+    builder.create<fir::FreeMemOp>(loc, var);
   };
   if (auto cstMustFree = fir::getIntIfConstant(mustFree)) {
     if (*cstMustFree != 0)
       genFree();
-    // else, nothing to do.
+    // else, mustFree is false, nothing to do.
     return;
   }
-  TODO(loc, "conditional free");
+  builder.genIfThen(loc, mustFree).genThen(genFree).end();
 }
 
 struct EndAssociateOpConversion
     : public mlir::OpConversionPattern<hlfir::EndAssociateOp> {
   using mlir::OpConversionPattern<hlfir::EndAssociateOp>::OpConversionPattern;
   explicit EndAssociateOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::EndAssociateOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::EndAssociateOp endAssociate, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = endAssociate->getLoc();
-    genFreeIfMustFree(loc, rewriter, adaptor.getVar(), adaptor.getMustFree());
+    auto module = endAssociate->getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    genFreeIfMustFree(loc, builder, adaptor.getVar(), adaptor.getMustFree());
     rewriter.eraseOp(endAssociate);
     return mlir::success();
   }
 };
 
 struct DestroyOpConversion
     : public mlir::OpConversionPattern<hlfir::DestroyOp> {
   using mlir::OpConversionPattern<hlfir::DestroyOp>::OpConversionPattern;
   explicit DestroyOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::DestroyOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::DestroyOp destroy, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     // If expr was bufferized on the heap, now is time to deallocate the buffer.
     mlir::Location loc = destroy->getLoc();
     mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     if (!fir::isa_trivial(bufferizedExpr.getType())) {
+      auto module = destroy->getParentOfType<mlir::ModuleOp>();
+      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getExpr());
       mlir::Value firBase = hlfir::Entity(bufferizedExpr).getFirBase();
-      genFreeIfMustFree(loc, rewriter, firBase, mustFree);
+      genFreeIfMustFree(loc, builder, firBase, mustFree);
     }
     rewriter.eraseOp(destroy);
     return mlir::success();
   }
 };
 
 struct NoReassocOpConversion
     : public mlir::OpConversionPattern<hlfir::NoReassocOp> {

flang/test/HLFIR/associate-codegen.fir

 // CHECK:  %[[VAL_22:.*]] = arith.constant false
 // CHECK:  %[[VAL_23:.*]] = fir.undefined tuple<!fir.boxchar<1>, i1>
 // CHECK:  %[[VAL_24:.*]] = fir.insert_value %[[VAL_23]], %[[VAL_22]], [1 : index] : (tuple<!fir.boxchar<1>, i1>, i1) -> tuple<!fir.boxchar<1>, i1>
 // CHECK:  %[[VAL_25:.*]] = fir.insert_value %[[VAL_24]], %[[VAL_21]]#0, [0 : index] : (tuple<!fir.boxchar<1>, i1>, !fir.boxchar<1>) -> tuple<!fir.boxchar<1>, i1>
 // CHECK:  fir.call @take_c(%[[VAL_21]]#0) : (!fir.boxchar<1>) -> ()
 // CHECK-NOT: fir.freemem
 
 
+func.func @test_end_associate_box(%var: !fir.box<!fir.array<?xf64>>) {
+  %true = arith.constant 1 : i1
+  hlfir.end_associate %var, %true : !fir.box<!fir.array<?xf64>>, i1
+  return
+}
+// CHECK-LABEL:   func.func @test_end_associate_box(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?xf64>>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant true
+// 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>>
+
+
+func.func @test_end_associate_boxchar(%var: !fir.boxchar<2>) {
+  %true = arith.constant 1 : i1
+  hlfir.end_associate %var, %true : !fir.boxchar<2>, i1
+  return
+}
+// CHECK-LABEL:   func.func @test_end_associate_boxchar(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.boxchar<2>) {
+// CHECK:  %[[VAL_1:.*]] = arith.constant true
+// CHECK:  %[[VAL_2:.*]] = fir.box_addr %[[VAL_0]] : (!fir.boxchar<2>) -> !fir.heap<!fir.char<2,?>>
+// CHECK:  fir.freemem %[[VAL_2]] : !fir.heap<!fir.char<2,?>>
+
+
+func.func @test_end_associate_box_dynamic(%var: !fir.box<!fir.array<?xf64>>, %must_free: i1) {
+  hlfir.end_associate %var, %must_free : !fir.box<!fir.array<?xf64>>, i1
+  return
+}
+// 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 @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>)