[flang-commits] [flang] [flang] Implement conditional expressions lowering (F2023) (PR #186490)

[Caroline Newcombe via flang-commits]

================
@@ -1821,6 +1821,158 @@ class HlfirBuilder {
     llvm_unreachable("unknown descriptor inquiry");
   }
 
+  /// Build nested if-then-else chain by walking the right-skewed
+  /// ConditionalExpr tree. The assignValue callback generates and assigns
+  /// each value to avoid evaluating non-taken branches.
+  template <typename T, typename Callback>
+  void
+  buildConditionalIfChain(const Fortran::evaluate::ConditionalExpr<T> &condExpr,
+                          const Callback &assignValue) {
+    const mlir::Location loc{getLoc()};
+    fir::FirOpBuilder &builder{getBuilder()};
+    getStmtCtx().pushScope();
+    const hlfir::EntityWithAttributes condEntity{gen(condExpr.condition())};
+    mlir::Value condition{hlfir::loadTrivialScalar(loc, builder, condEntity)};
+    condition = builder.createConvert(loc, builder.getI1Type(), condition);
+    builder.genIfOp(loc, {}, condition, /*withElseRegion=*/true)
+        .genThen([&]() {
+          getStmtCtx().pushScope();
+          assignValue(condExpr.thenValue());
+          getStmtCtx().finalizeAndPop();
+        })
+        .genElse([&]() {
+          getStmtCtx().pushScope();
+          assignValue(condExpr.elseValue());
+          getStmtCtx().finalizeAndPop();
+        })
+        .end();
+    getStmtCtx().finalizeAndPop();
+  }
+
+  /// Generate scalar conditional with lazy evaluation using assignment.
+  /// Creates a temporary and assigns the selected branch value to it.
+  template <typename T>
+  hlfir::Entity
+  genScalarConditional(const Fortran::evaluate::ConditionalExpr<T> &condExpr,
+                       mlir::Type elementType,
+                       const llvm::SmallVector<mlir::Value, 1> &typeParams) {
+    const mlir::Location loc{getLoc()};
+    fir::FirOpBuilder &builder{getBuilder()};
+    const mlir::Value tempStorage{builder.createTemporary(
+        loc, elementType, ".cond.scalar",
+        /*shape=*/mlir::ValueRange{}, /*typeParams=*/typeParams)};
+    const hlfir::DeclareOp tempDecl{hlfir::DeclareOp::create(
+        builder, loc, tempStorage, ".cond.result",
+        /*shape=*/mlir::Value{}, /*typeParams=*/typeParams)};
+    const hlfir::Entity temp{tempDecl};
+    buildConditionalIfChain(
+        condExpr, [&](const Fortran::evaluate::Expr<T> &expr) {
+          hlfir::Entity entity{gen(expr)};
+          entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
+          hlfir::AssignOp::create(builder, loc, entity, temp);
+        });
+    return temp;
+  }
+
+  /// Generate conditional expression using an allocatable temporary with lazy
+  /// evaluation. Creates an unallocated allocatable, then uses assignment to
+  /// set the value from the chosen branch (allocation/reallocation handled by
+  /// runtime).
+  template <typename T>
+  hlfir::Entity genAllocatableConditional(
+      const Fortran::evaluate::ConditionalExpr<T> &condExpr,
+      mlir::Type resultType, llvm::StringRef debugName) {
+    const mlir::Location loc{getLoc()};
+    fir::FirOpBuilder &builder{getBuilder()};
+    const mlir::Type heapType{fir::HeapType::get(resultType)};
+    const mlir::Type boxHeapType{fir::BoxType::get(heapType)};
+    const mlir::Value tempStorage{
+        builder.createTemporary(loc, boxHeapType, debugName)};
+    const mlir::Value unallocBox{fir::factory::createUnallocatedBox(
+        builder, loc, boxHeapType, /*nonDeferredParams=*/{})};
+    builder.createStoreWithConvert(loc, unallocBox, tempStorage);
+    const hlfir::DeclareOp tempDecl{
+        hlfir::DeclareOp::create(builder, loc, tempStorage, ".cond.result")};
+    const hlfir::Entity temp{tempDecl};
+    // Lazy evaluation: only the selected branch is evaluated and assigned.
+    buildConditionalIfChain(
+        condExpr, [&](const Fortran::evaluate::Expr<T> &expr) {
+          const hlfir::Entity entity{gen(expr)};
+          hlfir::AssignOp::create(builder, loc, entity, temp,
+                                  /*isWholeAllocatableAssignment=*/true,
+                                  /*keepLhsLengthIfRealloc=*/false,
+                                  /*temporary_lhs=*/true);
+        });
+    fir::FirOpBuilder *const bldr{&builder};
+    getStmtCtx().attachCleanup([=]() {
+      fir::factory::genFreememIfAllocated(
+          *bldr, loc,
+          fir::MutableBoxValue{tempStorage, /*lenParams=*/{},
+                               fir::MutableProperties{}});
+    });
+    return temp;
+  }
+
+  /// Generate scalar CHARACTER conditional with proper length handling.
+  template <typename T>
+  std::optional<hlfir::EntityWithAttributes> genCharacterConditional(
+      const Fortran::evaluate::ConditionalExpr<T> &condExpr) {
+    const mlir::Location loc{getLoc()};
+    fir::FirOpBuilder &builder{getBuilder()};
+    const mlir::Type resultType{Fortran::lower::translateSomeExprToFIRType(
+        converter, toEvExpr(condExpr))};
+    const mlir::Type elementType{hlfir::getFortranElementType(resultType)};
+    if (auto charType = mlir::dyn_cast<fir::CharacterType>(elementType)) {
+      if (charType.hasConstantLen()) {
+        llvm::SmallVector<mlir::Value, 1> typeParams;
+        const mlir::Value len{builder.createIntegerConstant(
+            loc, builder.getCharacterLengthType(), charType.getLen())};
+        typeParams.push_back(len);
+        return hlfir::EntityWithAttributes{
+            genScalarConditional(condExpr, elementType, typeParams)};
+      }
+      // Non-constant/varying length: use allocatable conditional to get length
+      // from selected branch.
+      return hlfir::EntityWithAttributes{
+          genAllocatableConditional(condExpr, elementType, ".cond.char")};
+    }
+    return std::nullopt;
+  }
+
+  /// Conditional expression (Fortran 2023)
+  template <typename T>
+  hlfir::EntityWithAttributes
+  gen(const Fortran::evaluate::ConditionalExpr<T> &condExpr) {
+    const int rank{condExpr.Rank()};
+
+    // Arrays: handle early to avoid unnecessary type checks.
+    // Per F2023 10.1.4(7), the shape is determined by the chosen branch.
+    if (rank != 0) {
+      mlir::Type condResultType{Fortran::lower::translateSomeExprToFIRType(
+          converter, toEvExpr(condExpr))};
+      if (auto boxTy = mlir::dyn_cast<fir::BoxType>(condResultType))
+        condResultType = fir::unwrapRefType(boxTy.getEleTy());
+      return hlfir::EntityWithAttributes{
+          genAllocatableConditional(condExpr, condResultType, ".cond.array")};
----------------
cenewcombe wrote:

Thanks, this is definitely worth another look. I’m questioning my latest change and could use some guidance on which direction to go with this.

The simplification is  good for readability, but we do lose something with the allocation of additional temps for chained conditionals. Previously, I had only placed the optimization in the else-branch to handle the common right-skewed structure of a chained conditional. It would be less common to write a conditional with a parenthesized conditional in the then-branch, so for simplicity I did not include that case. I think something like `(a ? (b ? x : y) : z)` would be unusual.

If Jean (or anyone) could weigh in on the various options:
1. Leave as is, with allocation of additional temps.
2. Return to the original right-skewed optimization (avoids creating extra temps for chained conditionals).
3. Implement similar optimization for the then- and else- branches (slightly more complicated, but would cover the uncommon case as well).

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