[CIR][Transforms] Fix flattening for TernaryOp to properly handle locations

Author: bcardosolopes

clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp

@@ -438,9 +438,13 @@ class CIRTernaryOpFlattening
     auto *condBlock = rewriter.getInsertionBlock();
     auto opPosition = rewriter.getInsertionPoint();
     auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);
-    auto *continueBlock = rewriter.createBlock(
-        remainingOpsBlock, op->getResultTypes(),
-        SmallVector<mlir::Location>(/* result number always 1 */ 1, loc));
+    SmallVector<mlir::Location, 2> locs;
+    // Ternary result is optional, make sure to populate the location only
+    // when relevant.
+    if (op->getResultTypes().size())
+      locs.push_back(loc);
+    auto *continueBlock =
+        rewriter.createBlock(remainingOpsBlock, op->getResultTypes(), locs);
     rewriter.create<mlir::cir::BrOp>(loc, remainingOpsBlock);
 
     auto &trueRegion = op.getTrueRegion();

clang/test/CIR/Transforms/ternary.cir

@@ -44,4 +44,25 @@ module {
 // CHECK:    cir.return %8 : !s32i
 // CHECK:  }
 
+  cir.func @foo2(%arg0: !cir.bool) {
+    cir.ternary(%arg0, true {
+      cir.yield
+    }, false {
+      cir.yield
+    }) : (!cir.bool) -> ()
+    cir.return
+  }
+
+// CHECK: cir.func @foo2(%arg0: !cir.bool) {
+// CHECK:   cir.brcond %arg0 ^bb1, ^bb2
+// CHECK: ^bb1:  // pred: ^bb0
+// CHECK:   cir.br ^bb3
+// CHECK: ^bb2:  // pred: ^bb0
+// CHECK:   cir.br ^bb3
+// CHECK: ^bb3:  // 2 preds: ^bb1, ^bb2
+// CHECK:   cir.br ^bb4
+// CHECK: ^bb4:  // pred: ^bb3
+// CHECK:   cir.return
+// CHECK: }
+
 }