[flang-commits] [flang] [flang] Region-based HLFIR operation for conditional expressions lowering (PR #194411)

[via flang-commits]

================
@@ -857,6 +858,159 @@ struct CharExtremumOpConversion
   }
 };
 
+struct ConditionalOpConversion
+    : public mlir::OpConversionPattern<hlfir::ConditionalOp> {
+  using mlir::OpConversionPattern<hlfir::ConditionalOp>::OpConversionPattern;
+  explicit ConditionalOpConversion(mlir::MLIRContext *ctx)
+      : mlir::OpConversionPattern<hlfir::ConditionalOp>{ctx} {
+    // This pattern recursively converts nested ConditionalOp's
+    // by cloning and then converting them, so we have to allow
+    // for recursive pattern application. The recursion is bounded
+    // by the nesting level of ConditionalOp's.
+    setHasBoundedRewriteRecursion();
+  }
+  llvm::LogicalResult
+  matchAndRewrite(hlfir::ConditionalOp condOp, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    const mlir::Location loc{condOp->getLoc()};
+    fir::FirOpBuilder builder(rewriter, condOp.getOperation());
+    HLFIRListener listener{builder, rewriter};
+    builder.setListener(&listener);
+    // Use ExprType to ensure both branches produce identical MLIR temp types.
+    const auto exprType{
+        mlir::cast<hlfir::ExprType>(condOp.getResult().getType())};
+    const mlir::Type tempBaseType{hlfir::getVariableType(exprType)};
+
+    // Emit one branch: clone ops, create temp, assign, run deferred
+    // destroys, yield (temp, mustFree).
+    auto emitBranch = [&](mlir::Region &region) {
+      mlir::IRMapping mapper;
+      auto yield{mlir::cast<hlfir::YieldOp>(region.front().getTerminator())};
+      const mlir::Value yieldedEntity{yield.getEntity()};
+      // Try to forward the yielded entity's storage directly to avoid using a
+      // temp.
+      hlfir::AsExprOp asExprToForward;
+      if (yieldedEntity.getType() == exprType) {
+        if (auto asExpr{yieldedEntity.getDefiningOp<hlfir::AsExprOp>()};
+            asExpr && asExpr.isMove() &&
+            asExpr->getBlock() == &region.front() &&
+            asExpr.getVar().getType() == tempBaseType &&
+            llvm::all_of(yieldedEntity.getUsers(), [&](mlir::Operation *user) {
+              return user == yield.getOperation() ||
+                     mlir::isa<hlfir::DestroyOp>(user);
+            }))
+          asExprToForward = asExpr;
+      }
+      // Ops after yieldDefOp are cleanups that must be deferred until after
+      // the temp assign to avoid use-after-free.
+      mlir::Operation *const yieldDefOp{yieldedEntity.getDefiningOp()};
+      const bool hasDeferredCleanups{!asExprToForward && yieldDefOp &&
+                                     yieldDefOp->getBlock() == &region.front()};
+      // Clone ops up to and including yieldDefOp, collecting deferred cleanups
+      // and the destroy of yieldedEntity for later replay.
+      llvm::SmallVector<mlir::Operation *> deferredOps;
+      mlir::Operation *destroyOfYielded{nullptr};
+      bool pastYieldDef{false};
+      for (auto &op : region.front().without_terminator()) {
+        if (auto destroy{mlir::dyn_cast<hlfir::DestroyOp>(op)};
+            destroy && destroy.getExpr() == yieldedEntity) {
+          destroyOfYielded = &op;
+          continue;
+        }
+        if (asExprToForward && &op == asExprToForward.getOperation())
+          continue;
+        if (hasDeferredCleanups && &op == yieldDefOp) {
+          builder.clone(op, mapper);
+          pastYieldDef = true;
+          continue;
+        }
+        if (pastYieldDef) {
+          deferredOps.push_back(&op);
+          continue;
+        }
+        builder.clone(op, mapper);
+      }
+      if (asExprToForward) {
+        // Forward storage and mustFree directly; ownership transfers to
+        // the caller so no destroy is needed.
+        const mlir::Value fwdVar{
+            mapper.lookupOrDefault(asExprToForward.getVar())};
+        const mlir::Value fwdMustFree{
+            mapper.lookupOrDefault(asExprToForward.getMustFree())};
+        fir::ResultOp::create(builder, loc,
+                              mlir::ValueRange{fwdVar, fwdMustFree});
+        return;
+      }
+      // If the yielded entity is an hlfir.as_expr, unwrap to its underlying
+      // variable — it carries type descriptor info needed for polymorphic
+      // temp allocation.  The value is never a mutable box because lowering
+      // applied derefPointersAndAllocatables before yielding.
+      const mlir::Value mappedYield{mapper.lookupOrDefault(yieldedEntity)};
+      auto asExprOp{mappedYield.getDefiningOp<hlfir::AsExprOp>()};
+      const hlfir::Entity val{asExprOp ? asExprOp.getVar() : mappedYield};
+      auto [temp, isHeapAlloc]{hlfir::createTempFromMold(loc, builder, val)};
+      hlfir::AssignOp::create(builder, loc, val, temp,
+                              /*realloc=*/false,
+                              /*keep_lhs_length_if_realloc=*/false,
+                              /*temporary_lhs=*/true);
+      // Replay deferred cleanups after the assign; destroy last.
+      for (auto *op : deferredOps)
+        builder.clone(*op, mapper);
+      if (destroyOfYielded)
+        builder.clone(*destroyOfYielded, mapper);
+      // Cast temp to match the fir.if result type if needed (e.g. the mold
+      // may have static extents while ExprType prescribes dynamic ones).
+      if (temp.getType() != tempBaseType) {
+        if (mlir::isa<fir::BoxCharType>(tempBaseType)) {
+          llvm::SmallVector<mlir::Value> lenParams;
+          hlfir::genLengthParameters(loc, builder, temp, lenParams);
+          assert(!lenParams.empty() && "character must have length");
+          const mlir::Value ref{builder.createConvert(
+              loc, builder.getRefType(temp.getFortranElementType()), temp)};
+          temp = hlfir::Entity{fir::EmboxCharOp::create(
+              builder, loc, tempBaseType, ref, lenParams[0])};
+        } else if (mlir::isa<fir::BaseBoxType>(tempBaseType) &&
+                   !mlir::isa<fir::BaseBoxType>(temp.getType())) {
+          // Unboxed temp needs emboxing to satisfy the fir.if result type.
+          const mlir::Value shape{temp.isArray()
+                                      ? hlfir::genShape(loc, builder, temp)
+                                      : mlir::Value{}};
+          // Only pass length params for dynamic-length character elements;
+          // fir.embox rejects redundant length operands.
+          llvm::SmallVector<mlir::Value> lenParams;
+          hlfir::genLengthParameters(loc, builder, temp, lenParams);
+          const mlir::Type elemType{hlfir::getFortranElementType(tempBaseType)};
+          if (auto charTy{mlir::dyn_cast<fir::CharacterType>(elemType)};
+              charTy && charTy.hasConstantLen())
+            lenParams.clear();
+          temp = hlfir::Entity{
+              fir::EmboxOp::create(builder, loc, tempBaseType, temp, shape,
+                                   /*slice=*/nullptr, lenParams)};
+        } else {
+          temp = hlfir::Entity{builder.createConvert(loc, tempBaseType, temp)};
+        }
----------------
jeanPerier wrote:

This could be moved into a helper in this file so that it can be re-used and that the code here is more readable.

https://github.com/llvm/llvm-project/pull/194411