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

[Caroline Newcombe via flang-commits]

================
@@ -1821,6 +1821,196 @@ 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([&]() {
+          // Recurse for nested ConditionalExpr; else assign terminal value.
+          if (const auto *nested =
+                  std::get_if<Fortran::evaluate::ConditionalExpr<T>>(
+                      &condExpr.elseValue().u)) {
+            buildConditionalIfChain(*nested, assignValue);
+          } else {
+            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 genScalarConditionalLazy(
+      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)};
+          auto [exv, cleanup] = hlfir::convertToValue(loc, builder, entity);
+          hlfir::Entity value{fir::getBase(exv)};
+          if (cleanup) {
+            getStmtCtx().attachCleanup(*cleanup);
+          }
+          hlfir::AssignOp::create(builder, loc, value, 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 genAllocatableConditionalLazy(
+      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 deferred-length using lazy
+  /// evaluation. Only the chosen branch is evaluated; the result gets that
+  /// branch's length.
+  template <typename T>
+  hlfir::Entity genScalarDeferredCharConditionalLazy(
+      const Fortran::evaluate::ConditionalExpr<T> &condExpr,
+      mlir::Type elementType) {
+    return genAllocatableConditionalLazy(condExpr, elementType, ".cond.char");
+  }
+
+  /// Generate array conditional expression with lazy evaluation using
+  /// assignment. Per F2023 10.1.4(7), the shape is determined by the chosen
+  /// branch.
+  template <typename T>
+  hlfir::Entity genArrayConditionalLazy(
+      const Fortran::evaluate::ConditionalExpr<T> &condExpr) {
+    mlir::Type condResultType{Fortran::lower::translateSomeExprToFIRType(
+        converter, toEvExpr(condExpr))};
+    if (auto boxTy = mlir::dyn_cast<fir::BoxType>(condResultType)) {
+      condResultType = fir::unwrapRefType(boxTy.getEleTy());
+    }
+    return genAllocatableConditionalLazy(condExpr, condResultType,
+                                         ".cond.array");
+  }
+
+  /// 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{
+            genScalarConditionalLazy(condExpr, elementType, typeParams)};
+      }
+      // Non-constant/varying length: use allocatable conditional to get length
+      // from selected branch.
+      return hlfir::EntityWithAttributes{
+          genScalarDeferredCharConditionalLazy(condExpr, elementType)};
+    }
+    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) {
+      return hlfir::EntityWithAttributes{genArrayConditionalLazy(condExpr)};
+    }
+    // CHARACTER scalars require special handling for type parameters.
+    if constexpr (T::category == Fortran::common::TypeCategory::Character) {
+      if (auto result = genCharacterConditional(condExpr)) {
+        return *result;
+      }
+    }
+    // Derived type scalars require special handling for type parameters.
+    if constexpr (T::category == Fortran::common::TypeCategory::Derived) {
+      const mlir::Type resultType{Fortran::lower::translateSomeExprToFIRType(
+          converter, toEvExpr(condExpr))};
+      const mlir::Type elementType{hlfir::getFortranElementType(resultType)};
+      return hlfir::EntityWithAttributes{
+          genScalarConditionalLazy(condExpr, elementType, {})};
+    }
+    // Other scalar types (INTEGER, REAL, COMPLEX, LOGICAL, UNSIGNED).
+    mlir::Type condResultType{Fortran::lower::translateSomeExprToFIRType(
+        converter, toEvExpr(condExpr))};
+    if (auto boxTy = mlir::dyn_cast<fir::BoxType>(condResultType)) {
+      condResultType = fir::unwrapRefType(boxTy.getEleTy());
+    }
+    const mlir::Type resultType{hlfir::getFortranElementType(condResultType)};
+    return hlfir::EntityWithAttributes{
+        genScalarConditionalLazy(condExpr, resultType, {})};
+  }
----------------
cenewcombe wrote:

Done

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