[const-expr] Teach ConstExpr how to handle switch_enum/unchecked_enum…

lib/SILOptimizer/Utils/ConstExpr.cpp

@@ -15,6 +15,7 @@
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/Basic/Defer.h"
+#include "swift/Basic/NullablePtr.h"
 #include "swift/Demangling/Demangle.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/FormalLinkage.h"
@@ -213,6 +214,19 @@ SymbolicValue ConstExprFunctionState::computeConstantValue(SILValue value) {
     return val;
   }
 
+  // If this is an unchecked_enum_data from a fragile type, then we can return
+  // the enum case value.
+  if (auto *uedi = dyn_cast<UncheckedEnumDataInst>(value)) {
+    auto aggValue = uedi->getOperand();
+    auto val = getConstantValue(aggValue);
+    if (val.isConstant()) {
+      assert(val.getKind() == SymbolicValue::EnumWithPayload);
+      return val.getEnumPayloadValue();
+    }
+    // Not a const.
+    return val;
+  }
+
   // If this is a destructure_result, then we can return the element being
   // extracted.
   if (isa<DestructureStructResult>(value) ||
@@ -1359,20 +1373,34 @@ static llvm::Optional<SymbolicValue> evaluateAndCacheCall(
       }
       if (!value.isConstant())
         return value;
+
       assert(value.getKind() == SymbolicValue::Enum ||
              value.getKind() == SymbolicValue::EnumWithPayload);
-      // Set up basic block arguments.
+
       auto *caseBB = switchInst->getCaseDestination(value.getEnumValue());
-      if (caseBB->getNumArguments() > 0) {
-        assert(value.getKind() == SymbolicValue::EnumWithPayload);
-        // When there are multiple payload components, they form a single
-        // tuple-typed argument.
-        assert(caseBB->getNumArguments() == 1);
-        auto argument = value.getEnumPayloadValue();
-        assert(argument.isConstant());
-        state.setValue(caseBB->getArgument(0), argument);
-      }
+
+      // Prepare to subsequently visit the case blocks instructions.
       nextInst = caseBB->begin();
+      // Then set up the arguments.
+      if (caseBB->getParent()->hasOwnership() &&
+          switchInst->getDefaultBBOrNull() == caseBB) {
+        // If we are visiting the default block and we are in ossa, then we may
+        // have uses of the failure parameter. That means we need to map the
+        // original value to the argument.
+        state.setValue(caseBB->getArgument(0), value);
+        continue;
+      }
+
+      if (caseBB->getNumArguments() == 0)
+        continue;
+
+      assert(value.getKind() == SymbolicValue::EnumWithPayload);
+      // When there are multiple payload components, they form a single
+      // tuple-typed argument.
+      assert(caseBB->getNumArguments() == 1);
+      auto argument = value.getEnumPayloadValue();
+      assert(argument.isConstant());
+      state.setValue(caseBB->getArgument(0), argument);
       continue;
     }
 

test/SILOptimizer/pound_assert_ossa.sil

@@ -535,3 +535,70 @@ bb0(%arg : $BiggerStruct):
   %ret = tuple ()
   return %ret : $()
 }
+
+enum Either {
+case left(Builtin.Int64)
+case right(Builtin.Int64)
+}
+
+// Make sure that we properly handle failure default cases.
+sil [ossa] @switch_enum_test_callee_1 : $@convention(thin) () -> Builtin.Int64 {
+bb0:
+  %0 = integer_literal $Builtin.Int64, 0
+  %1 = enum $Either, #Either.left!enumelt.1, %0 : $Builtin.Int64
+  switch_enum %1 : $Either, case #Either.left!enumelt.1: bb1, default bb2
+
+bb1(%2 : $Builtin.Int64):
+  br bb3(%2 : $Builtin.Int64)
+
+bb2(%3 : $Either):
+  %4 = integer_literal $Builtin.Int64, 1
+  br bb3(%4 : $Builtin.Int64)
+
+bb3(%5 : $Builtin.Int64):
+  return %5 : $Builtin.Int64
+}
+
+sil [ossa] @switch_enum_test_callee_2 : $@convention(thin) () -> Builtin.Int64 {
+bb0:
+  %0 = integer_literal $Builtin.Int64, 0
+  %1 = enum $Either, #Either.left!enumelt.1, %0 : $Builtin.Int64
+  // Make sure we go down the bad path.
+  switch_enum %1 : $Either, case #Either.right!enumelt.1: bb4, default bb5
+
+bb4(%7 : $Builtin.Int64):
+  br bb6(%7 : $Builtin.Int64)
+
+bb5(%8 : $Either):
+  %9 = unchecked_enum_data %8 : $Either, #Either.right!enumelt.1
+  br bb6(%9 : $Builtin.Int64)
+
+bb6(%10 : $Builtin.Int64):
+  return %10 : $Builtin.Int64
+}
+
+sil [ossa] @switch_enum_test_caller : $@convention(thin) () -> () {
+bb0:
+  %0 = function_ref @switch_enum_test_callee_1 : $@convention(thin) () -> Builtin.Int64
+  %0a = function_ref @switch_enum_test_callee_2 : $@convention(thin) () -> Builtin.Int64
+  %2 = apply %0() : $@convention(thin) () -> Builtin.Int64
+  %3 = apply %0a() : $@convention(thin) () -> Builtin.Int64
+  %str = string_literal utf8 ""
+  %resultPositive = integer_literal $Builtin.Int64, 0
+  %resultNegative = integer_literal $Builtin.Int64, 1
+  %cmp1Positive = builtin "cmp_eq_Int64"(%2 : $Builtin.Int64, %resultPositive : $Builtin.Int64) : $Builtin.Int1
+  builtin "poundAssert"(%cmp1Positive : $Builtin.Int1, %str : $Builtin.RawPointer) : $()
+  // Make sure we simplified down the bb1 path.
+  %cmp2Positive = builtin "cmp_eq_Int64"(%3 : $Builtin.Int64, %resultPositive : $Builtin.Int64) : $Builtin.Int1
+  builtin "poundAssert"(%cmp2Positive : $Builtin.Int1, %str : $Builtin.RawPointer) : $()
+
+  %cmp1Negative = builtin "cmp_eq_Int64"(%2 : $Builtin.Int64, %resultNegative : $Builtin.Int64) : $Builtin.Int1
+  // expected-error @+1 {{assertion failed}}
+  builtin "poundAssert"(%cmp1Negative : $Builtin.Int1, %str : $Builtin.RawPointer) : $()
+  %cmp2Negative = builtin "cmp_eq_Int64"(%3 : $Builtin.Int64, %resultNegative : $Builtin.Int64) : $Builtin.Int1
+  // expected-error @+1 {{assertion failed}}
+  builtin "poundAssert"(%cmp2Negative : $Builtin.Int1, %str : $Builtin.RawPointer) : $()
+
+  %9999 = tuple()
+  return %9999 : $()
+}