ConstantFolding: fix wrong constant folding for float-infinity compares

lib/SILOptimizer/Utils/ConstantFolding.cpp

@@ -381,6 +381,13 @@ static SILValue constantFoldIntrinsic(BuiltinInst *BI, llvm::Intrinsic::ID ID,
   return nullptr;
 }
 
+static bool isFiniteFloatLiteral(SILValue v) {
+  if (auto *lit = dyn_cast<FloatLiteralInst>(v)) {
+    return lit->getValue().isFinite();
+  }
+  return false;
+}
+
 static SILValue constantFoldCompareFloat(BuiltinInst *BI, BuiltinValueKind ID) {
   static auto hasIEEEFloatNanBitRepr = [](const APInt val) -> bool {
     auto bitWidth = val.getBitWidth();
@@ -640,17 +647,11 @@ static SILValue constantFoldCompareFloat(BuiltinInst *BI, BuiltinValueKind ID) {
                 m_BuiltinInst(BuiltinValueKind::FCMP_ULE, 
                               m_SILValue(Other), m_BitCast(m_IntegerLiteralInst(builtinArg)))))) {
     APInt val = builtinArg->getValue();
-    if (hasIEEEFloatPosInfBitRepr(val)) {
-      // One of the operands is infinity, but unless the other operand is not
-      // fully visible we cannot definitively say what it is. It can be anything, 
-      // including NaN and infinity itself. Therefore, we cannot fold the comparison
-      // just yet.
-      if (isa<StructExtractInst>(Other) || isa<TupleExtractInst>(Other)) {
-        return nullptr;
-      } else {  
-        SILBuilderWithScope B(BI);
-        return B.createIntegerLiteral(BI->getLoc(), BI->getType(), APInt(1, 1));
-      }
+    if (hasIEEEFloatPosInfBitRepr(val) &&
+        // Only if `Other` is a literal we can be sure that it's not Inf or NaN.
+        isFiniteFloatLiteral(Other)) {
+      SILBuilderWithScope B(BI);
+      return B.createIntegerLiteral(BI->getLoc(), BI->getType(), APInt(1, 1));
     }
   }
 
@@ -682,17 +683,11 @@ static SILValue constantFoldCompareFloat(BuiltinInst *BI, BuiltinValueKind ID) {
                 m_BuiltinInst(BuiltinValueKind::FCMP_ULE, 
                               m_BitCast(m_IntegerLiteralInst(builtinArg)), m_SILValue(Other))))) {
     APInt val = builtinArg->getValue();
-    if (hasIEEEFloatPosInfBitRepr(val)) {
-      // One of the operands is infinity, but unless the other operand is not
-      // fully visible we cannot definitively say what it is. It can be anything, 
-      // including NaN and infinity itself. Therefore, we cannot fold the comparison
-      // just yet.
-      if (isa<StructExtractInst>(Other) || isa<TupleExtractInst>(Other)) {
-        return nullptr;
-      } else {  
-        SILBuilderWithScope B(BI);
-        return B.createIntegerLiteral(BI->getLoc(), BI->getType(), APInt(1, 0));
-      }
+    if (hasIEEEFloatPosInfBitRepr(val) &&
+        // Only if `Other` is a literal we can be sure that it's not Inf or NaN.
+        isFiniteFloatLiteral(Other)) {
+      SILBuilderWithScope B(BI);
+      return B.createIntegerLiteral(BI->getLoc(), BI->getType(), APInt(1, 0));
     }
   }
 

test/SILOptimizer/constant_fold_float.swift

@@ -0,0 +1,18 @@
+// RUN: %target-swift-frontend -parse-as-library -module-name test %s -O -emit-sil | %FileCheck %s
+
+// REQUIRES: swift_stdlib_no_asserts,optimized_stdlib
+
+// CHECK-LABEL: sil @$s4test17dont_fold_inf_cmpySbSfF :
+// CHECK:         builtin "fcmp_olt_FPIEEE32"
+// CHECK:       } // end sil function '$s4test17dont_fold_inf_cmpySbSfF'
+public func dont_fold_inf_cmp(_ f: Float) -> Bool {
+  (f + 0) < .infinity
+}
+
+// CHECK-LABEL: sil @$s4test014dont_fold_inf_D4_cmpSbyF :
+// CHECK:         builtin "fcmp_olt_FPIEEE32"
+// CHECK:       } // end sil function '$s4test014dont_fold_inf_D4_cmpSbyF'
+public func dont_fold_inf_inf_cmp() -> Bool {
+  0x1.0p128 < Float.infinity
+}
+

test/SILOptimizer/constant_propagation.sil

@@ -870,6 +870,35 @@ bb0:
 // CHECK-NEXT: } // end sil function 'fold_double_comparison_with_inf'
 }
 
+sil @dont_fold_comparison_with_inf : $@convention(thin) (Builtin.FPIEEE32) -> Builtin.Int1 {
+bb0(%0 : $Builtin.FPIEEE32):
+  %2 = float_literal $Builtin.FPIEEE32, 0x0
+  %4 = builtin "fadd_FPIEEE32"(%0 : $Builtin.FPIEEE32, %2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32
+  %5 = integer_literal $Builtin.Int32, 2139095040
+  %6 = builtin "bitcast_Int32_FPIEEE32"(%5 : $Builtin.Int32) : $Builtin.FPIEEE32
+  %9 = builtin "fcmp_olt_FPIEEE32"(%4 : $Builtin.FPIEEE32, %6 : $Builtin.FPIEEE32) : $Builtin.Int1
+  return %9 : $Builtin.Int1
+
+// CHECK-LABEL: sil @dont_fold_comparison_with_inf :
+// CHECK:         [[R:%.*]] = builtin "fcmp_olt_FPIEEE32"
+// CHECK:         return [[R]]
+// CHECK:       } // end sil function 'dont_fold_comparison_with_inf'
+}
+
+sil @dont_fold_comparison_with_inf2 : $@convention(thin) () -> Builtin.Int1 {
+bb0:
+  %2 = float_literal $Builtin.FPIEEE32, 0x7F800000 // +Inf // user: %3
+  %5 = integer_literal $Builtin.Int32, 2139095040
+  %6 = builtin "bitcast_Int32_FPIEEE32"(%5 : $Builtin.Int32) : $Builtin.FPIEEE32
+  %9 = builtin "fcmp_olt_FPIEEE32"(%2 : $Builtin.FPIEEE32, %6 : $Builtin.FPIEEE32) : $Builtin.Int1
+  return %9 : $Builtin.Int1
+
+// CHECK-LABEL: sil @dont_fold_comparison_with_inf2 :
+// CHECK:         [[R:%.*]] = builtin "fcmp_olt_FPIEEE32"
+// CHECK:         return [[R]]
+// CHECK:       } // end sil function 'dont_fold_comparison_with_inf2'
+}
+
 // fold float comparison operations with Infinity/NaN when the other argument is not constant
 sil @fold_float_comparison_with_non_constant_arg : $@convention(thin) (Float) -> () {
 bb0(%0: $Float):