Merge pull request #77590 from eeckstein/fix-enum-data-simplification

Optimizer: fix simplification of unchecked_enum_data

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyUncheckedEnumData.swift

@@ -12,7 +12,7 @@
 
 import SIL
 
-extension UncheckedEnumDataInst : OnoneSimplifyable {
+extension UncheckedEnumDataInst : OnoneSimplifyable, SILCombineSimplifyable {
   func simplify(_ context: SimplifyContext) {
     guard let enumInst = self.enum as? EnumInst else {
       return

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -113,6 +113,7 @@ private func registerSwiftPasses() {
   registerForSILCombine(DestructureTupleInst.self, { run(DestructureTupleInst.self, $0) })
   registerForSILCombine(TypeValueInst.self, { run(TypeValueInst.self, $0) })
   registerForSILCombine(ClassifyBridgeObjectInst.self, { run(ClassifyBridgeObjectInst.self, $0) })
+  registerForSILCombine(UncheckedEnumDataInst.self, { run(UncheckedEnumDataInst.self, $0) })
 
   // Test passes
   registerPass(aliasInfoDumper, { aliasInfoDumper.run($0) })

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -576,9 +576,18 @@ extension SimplifyContext {
     let singleUse = preserveDebugInfo ? first.uses.singleUse : first.uses.ignoreDebugUses.singleUse
     let canEraseFirst = singleUse?.instruction == second
 
-    if !canEraseFirst && first.parentFunction.hasOwnership && replacement.ownership == .owned {
-      // We cannot add more uses to `replacement` without inserting a copy.
-      return
+    if !canEraseFirst && first.parentFunction.hasOwnership {
+      if replacement.ownership == .owned {
+        // We cannot add more uses to `replacement` without inserting a copy.
+        return
+      }
+      if first.ownership == .owned {
+        // We have to insert a compensating destroy because we are deleting the second instruction but
+        // not the first one. This can happen if the first instruction is an `enum` which constructs a
+        // non-trivial enum from a trivial payload.
+        let builder = Builder(before: second, self)
+        builder.createDestroyValue(operand: first)
+      }
     }
 
     second.uses.replaceAll(with: replacement, self)

include/swift/SILOptimizer/PassManager/Passes.def

@@ -531,6 +531,7 @@ SWIFT_SILCOMBINE_PASS(DestroyValueInst)
 SWIFT_SILCOMBINE_PASS(DestructureStructInst)
 SWIFT_SILCOMBINE_PASS(DestructureTupleInst)
 SWIFT_SILCOMBINE_PASS(TypeValueInst)
+SWIFT_SILCOMBINE_PASS(UncheckedEnumDataInst)
 
 #undef IRGEN_PASS
 #undef SWIFT_MODULE_PASS

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -47,7 +47,6 @@ namespace {
     SILValue visitStructExtractInst(StructExtractInst *SEI);
     SILValue visitEnumInst(EnumInst *EI);
     SILValue visitSelectEnumInst(SelectEnumInst *SEI);
-    SILValue visitUncheckedEnumDataInst(UncheckedEnumDataInst *UEDI);
     SILValue visitAddressToPointerInst(AddressToPointerInst *ATPI);
     SILValue visitPointerToAddressInst(PointerToAddressInst *PTAI);
     SILValue visitRefToRawPointerInst(RefToRawPointerInst *RRPI);
@@ -171,24 +170,6 @@ SILValue InstSimplifier::visitStructExtractInst(StructExtractInst *sei) {
   return SILValue();
 }
 
-SILValue
-InstSimplifier::visitUncheckedEnumDataInst(UncheckedEnumDataInst *uedi) {
-  if (uedi->getOperand()->getType().isValueTypeWithDeinit())
-    return SILValue();
-  // (unchecked_enum_data (enum payload)) -> payload
-  auto opt = lookThroughOwnershipInsts(uedi->getOperand());
-  if (auto *ei = dyn_cast<EnumInst>(opt)) {
-    if (ei->getElement() != uedi->getElement())
-      return SILValue();
-
-    assert(ei->hasOperand() &&
-           "Should only get data from an enum with payload.");
-    return lookThroughOwnershipInsts(ei->getOperand());
-  }
-
-  return SILValue();
-}
-
 // Simplify:
 //   %1 = unchecked_enum_data %0 : $Optional<C>, #Optional.Some!enumelt
 //   %2 = enum $Optional<C>, #Optional.Some!enumelt, %1 : $C

test/SILOptimizer/sil_combine_enums_ossa.sil

@@ -18,7 +18,13 @@ enum Numerals {
   case Four
 }
 
+enum E2 {
+  case int(Int)
+  case str(String)
+}
+
 sil [ossa] @external_func: $@convention(thin) () -> ()
+sil @useit : $@convention(thin) (@owned E2) -> ()
 
 // CHECK-LABEL: sil [ossa] @eliminate_sw_enum_addr : $@convention(thin) () -> Int {
 // CHECK-NOT: switch_enum_addr
@@ -679,3 +685,33 @@ bb0(%0 : $*Optional<NC>):
   return %1 : $Optional<NC>
 }
 
+// CHECK-LABEL: sil [ossa] @trivial_replacement
+// CHECK:       bb1({{.*}}:
+// CHECK-NEXT:    end_borrow
+// CHECK-NEXT:    destroy_value %2
+// CHECK:         apply {{%[0-9]+}}(%1)
+// CHECK:       } // end sil function 'trivial_replacement'
+sil [ossa] @trivial_replacement : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %1 = enum $E2, #E2.int!enumelt, %0 : $Int
+  %2 = enum $Optional<E2>, #Optional.some!enumelt, %1 : $E2, forwarding: @owned
+  %3 = begin_borrow %2 : $Optional<E2>
+  switch_enum %3 : $Optional<E2>, case #Optional.some!enumelt: bb1, case #Optional.none!enumelt: bb2
+
+bb1(%5 : @guaranteed $E2):
+  end_borrow %3 : $Optional<E2>
+  %7 = unchecked_enum_data %2 : $Optional<E2>, #Optional.some!enumelt
+  %8 = function_ref @useit : $@convention(thin) (@owned E2) -> ()
+  %9 = apply %8(%7) : $@convention(thin) (@owned E2) -> ()
+  br bb3
+  
+bb2:
+  end_borrow %3 : $Optional<E2>
+  destroy_value %2 : $Optional<E2>
+  br bb3
+
+bb3:
+  %r = tuple ()
+  return %r : $()
+}
+

test/SILOptimizer/sil_combine_ossa.sil

@@ -5202,25 +5202,6 @@ bb0(%0 : @guaranteed $Function):
   return %13 : $Double
 }
 
-// CHECK-LABEL: sil [ossa] @test_borrow_unchecked_enum_data : $@convention(thin) (@guaranteed Function) -> Double {
-// CHECK-NOT: unchecked_enum_data
-// CHECK: } // end sil function 'test_borrow_unchecked_enum_data'
-sil [ossa] @test_borrow_unchecked_enum_data : $@convention(thin) (@guaranteed Function) -> Double {
-bb0(%0 : @guaranteed $Function):
-  %3 = function_ref @negateFunction : $@convention(thin) (MyInt, @guaranteed Function) -> Double
-  %4 = copy_value %0 : $Function
-  %5 = partial_apply [callee_guaranteed] %3(%4) : $@convention(thin) (MyInt, @guaranteed Function) -> Double
-  %7 = integer_literal $Builtin.Int32, 0
-  %8 = struct $MyInt (%7 : $Builtin.Int32)
-  %6 = enum $Optional<@callee_guaranteed (MyInt) -> Double>, #Optional.some!enumelt, %5 : $@callee_guaranteed (MyInt) -> Double
-  %9 = begin_borrow %6 : $Optional<@callee_guaranteed (MyInt) -> Double>
-  %12 = unchecked_enum_data %9 : $Optional<@callee_guaranteed (MyInt) -> Double>, #Optional.some!enumelt
-  %13 = apply %12(%8) : $@callee_guaranteed (MyInt) -> Double
-  end_borrow %9 : $Optional<@callee_guaranteed (MyInt) -> Double>
-  destroy_value %6 : $Optional<@callee_guaranteed (MyInt) -> Double>
-  return %13 : $Double
-}
-
 sil [reabstraction_thunk] @thunk : $@convention(thin) (@guaranteed Klass, @guaranteed Klass) -> ()
 
 // CHECK-LABEL: sil [ossa] @test_partial_apply_apply_opt1 :

test/SILOptimizer/simplify_cfg_args.sil

@@ -121,10 +121,15 @@ enum X {
   case C(Int32)
 }
 
-// CHECK-LABEL: simplify_switch_include_args
+// CHECK-LABEL: sil @simplify_switch_include_args :
+// CHECK:         tuple
+// CHECK-NEXT:    enum $X
+// CHECK-NEXT:    unchecked_enum_data
+// CHECK-NEXT:    tuple_extract
+// CHECK-NEXT:    return
+// CHECK:       } // end sil function 'simplify_switch_include_args'
 sil @simplify_switch_include_args : $@convention(thin) (Int32, Int32) -> Int32 {
 bb0(%0 : $Int32, %1 : $Int32):
-// CHECK: return %0 : $Int32
   %2 = tuple (%0 : $Int32, %1 : $Int32)
   %3 = enum $X, #X.B!enumelt, %2 : $(Int32, Int32)
   switch_enum %3 : $X, case #X.A!enumelt: bb1, case #X.B!enumelt: bb3, case #X.C!enumelt: bb5

test/SILOptimizer/simplify_cfg_args_ossa.sil

@@ -93,10 +93,15 @@ enum X {
   case C(Int32)
 }
 
-// CHECK-LABEL: simplify_switch_include_args
+// CHECK-LABEL: sil [ossa] @simplify_switch_include_args :
+// CHECK:         tuple
+// CHECK-NEXT:    enum $X
+// CHECK-NEXT:    unchecked_enum_data
+// CHECK-NEXT:    tuple_extract
+// CHECK-NEXT:    return
+// CHECK:       } // end sil function 'simplify_switch_include_args'
 sil [ossa] @simplify_switch_include_args : $@convention(thin) (Int32, Int32) -> Int32 {
 bb0(%0 : $Int32, %1 : $Int32):
-// CHECK: return %0 : $Int32
   %2 = tuple (%0 : $Int32, %1 : $Int32)
   %3 = enum $X, #X.B!enumelt, %2 : $(Int32, Int32)
   switch_enum %3 : $X, case #X.A!enumelt: bb1, case #X.B!enumelt: bb3, case #X.C!enumelt: bb5

test/SILOptimizer/simplify_cfg_ossa_dom_jumpthread.sil

@@ -511,11 +511,12 @@ bb7(%r : $Builtin.Int32, %carg2 : $Builtin.Int32):
 
 // CHECK-LABEL: sil [ossa] @jumpthread_switch_enum5
 // CHECK:      bb0:
-// CHECK:        br bb2
+// CHECK:        br bb1
 // CHECK:      bb1:
-// CHECK-NEXT:   br bb2
+// CHECK-NEXT:   cond_br undef, bb2, bb3
 // CHECK:      bb2:
-// CHECK:        cond_br undef, bb1, bb3
+// CHECK-NEXT:   unchecked_enum_data
+// CHECK-NEXT:   br bb1
 // CHECK:      bb3:
 // CHECK-NEXT:   tuple
 // CHECK-NEXT:   return