Merge pull request #81780 from eeckstein/optimize-enum-comparison

eeckstein · web-flow · commit 89bdb311abde · 2025-05-28T07:01:29.000+02:00
Optimize enum comparisons
diff --git a/SwiftCompilerSources/Sources/Basic/Utils.swift b/SwiftCompilerSources/Sources/Basic/Utils.swift
@@ -79,11 +79,15 @@ public extension NoReflectionChildren {
 //                              StringRef
 //===----------------------------------------------------------------------===//
 
-public struct StringRef : CustomStringConvertible, NoReflectionChildren {
+public struct StringRef : CustomStringConvertible, NoReflectionChildren, ExpressibleByStringLiteral {
   public let _bridged: BridgedStringRef
 
   public init(bridged: BridgedStringRef) { self._bridged = bridged }
 
+  public init(stringLiteral: StaticString) {
+    self._bridged = BridgedStringRef(data: stringLiteral.utf8Start, count: stringLiteral.utf8CodeUnitCount)
+  }
+
   public var string: String { String(_bridged)  }
   public var description: String { string }
 
diff --git a/SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift b/SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift
@@ -28,6 +28,9 @@ extension ApplyInst : OnoneSimplifiable, SILCombineSimplifiable {
     if tryRemoveArrayCast(apply: self, context) {
       return
     }
+    if tryOptimizeEnumComparison(apply: self, context) {
+      return
+    }
     if !context.preserveDebugInfo {
       _ = tryReplaceExistentialArchetype(of: self, context)
     }
@@ -110,6 +113,48 @@ private func tryRemoveArrayCast(apply: ApplyInst, _ context: SimplifyContext) ->
   return true
 }
 
+/// Optimize (the very inefficient) RawRepresentable comparison to a simple compare of enum tags.
+/// For example,
+/// ```
+///   enum E: String {
+///     case  a, b, c
+///   }
+/// ```
+/// is compared by getting the raw values of both operands and doing a string compare.
+/// This peephole optimizations replaces the call to such a comparison function with a direct compare of
+/// the enum tags, which boils down to a single integer comparison instruction.
+///
+private func tryOptimizeEnumComparison(apply: ApplyInst, _ context: SimplifyContext) -> Bool {
+  guard let callee = apply.referencedFunction,
+        apply.arguments.count == 2,
+        callee.hasSemanticsAttribute("rawrepresentable.is_equal"),
+        apply.type.isStruct
+  else {
+    return false
+  }
+  let lhs = apply.arguments[0]
+  let rhs = apply.arguments[1]
+  guard let enumDecl = lhs.type.nominal as? EnumDecl,
+        !enumDecl.isResilient(in: apply.parentFunction),
+        !enumDecl.hasClangNode,
+        lhs.type.isAddress,
+        lhs.type == rhs.type
+  else {
+    return false
+  }
+  let builder = Builder(before: apply, context)
+  let tagType = context.getBuiltinIntegerType(bitWidth: 32)
+  let lhsTag = builder.createBuiltin(name: "getEnumTag", type: tagType,
+                                     substitutions: apply.substitutionMap, arguments: [lhs])
+  let rhsTag = builder.createBuiltin(name: "getEnumTag", type: tagType,
+                                     substitutions: apply.substitutionMap, arguments: [rhs])
+  let builtinBoolType = context.getBuiltinIntegerType(bitWidth: 1)
+  let cmp = builder.createBuiltin(name: "cmp_eq_Int32", type: builtinBoolType, arguments: [lhsTag, rhsTag])
+  let booleanResult = builder.createStruct(type: apply.type, elements: [cmp])
+  apply.replace(with: booleanResult, context)
+  return true
+}
+
 /// If the apply uses an existential archetype (`@opened("...")`) and the concrete type is known,
 /// replace the existential archetype with the concrete type
 ///   1. in the apply's substitution map
diff --git a/include/swift/AST/SemanticAttrs.def b/include/swift/AST/SemanticAttrs.def
@@ -73,6 +73,8 @@ SEMANTICS_ATTR(OPTIMIZE_SIL_SPECIALIZE_GENERIC_PARTIAL_NEVER,
           "optimize.sil.specialize.generic.partial.never")
 SEMANTICS_ATTR(OPTIMIZE_SIL_INLINE_CONSTANT_ARGUMENTS,
           "optimize.sil.inline.constant.arguments")
+SEMANTICS_ATTR(DERIVED_ENUM_EQUALS,
+          "derived_enum_equals")
 SEMANTICS_ATTR(OPTIMIZE_SIL_SPECIALIZE_GENERIC_SIZE_NEVER,
           "optimize.sil.specialize.generic.size.never")
 SEMANTICS_ATTR(OPTIMIZE_SIL_SPECIALIZE_OWNED2GUARANTEE_NEVER,
diff --git a/lib/SILOptimizer/Transforms/PerformanceInliner.cpp b/lib/SILOptimizer/Transforms/PerformanceInliner.cpp
@@ -524,6 +524,14 @@ static bool hasConstantArguments(FullApplySite fas) {
   return true;
 }
 
+static bool hasConstantEnumArgument(FullApplySite fas) {
+  for (SILValue arg : fas.getArguments()) {
+    if (isa<EnumInst>(arg))
+      return true;
+  }
+  return false;
+}
+
 bool SILPerformanceInliner::isProfitableToInline(
     FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
     int &NumCallerBlocks,
@@ -597,6 +605,13 @@ bool SILPerformanceInliner::isProfitableToInline(
     return true;
   }
 
+  // If there is a "constant" enum argument to a synthesized enum comparison,
+  // we can always inline it, because most of it will be constant folded anyway.
+  if (Callee->hasSemanticsAttr(semantics::DERIVED_ENUM_EQUALS) &&
+      hasConstantEnumArgument(AI)) {
+    return true;
+  }
+
   // Bail out if this generic call can be optimized by means of
   // the generic specialization, because we prefer generic specialization
   // to inlining of generics.
diff --git a/lib/Sema/DerivedConformance/DerivedConformanceEquatableHashable.cpp b/lib/Sema/DerivedConformance/DerivedConformanceEquatableHashable.cpp
@@ -394,10 +394,12 @@ deriveEquatable_eq(
   auto boolTy = C.getBoolType();
 
   Identifier generatedIdentifier;
+  bool isDerivedEnumEquals = false;
   if (parentDC->getParentModule()->isResilient()) {
     generatedIdentifier = C.Id_EqualsOperator;
   } else if (selfIfaceTy->getEnumOrBoundGenericEnum()) {
     generatedIdentifier = C.Id_derived_enum_equals;
+    isDerivedEnumEquals = true;
   } else {
     assert(selfIfaceTy->getStructOrBoundGenericStruct());
     generatedIdentifier = C.Id_derived_struct_equals;
@@ -411,6 +413,9 @@ deriveEquatable_eq(
       /*GenericParams=*/nullptr, params, boolTy, parentDC);
   eqDecl->setUserAccessible(false);
   eqDecl->setSynthesized();
+  if (isDerivedEnumEquals) {
+    eqDecl->getAttrs().add(new (C) SemanticsAttr("derived_enum_equals", SourceLoc(), SourceRange(), /*Implicit=*/true));
+  }
 
   // Add the @_implements(Equatable, ==(_:_:)) attribute
   if (generatedIdentifier != C.Id_EqualsOperator) {
diff --git a/stdlib/public/core/CompilerProtocols.swift b/stdlib/public/core/CompilerProtocols.swift
@@ -149,6 +149,7 @@ public protocol RawRepresentable<RawValue> {
 ///   - lhs: A raw-representable instance.
 ///   - rhs: A second raw-representable instance.
 @inlinable // trivial-implementation
+@_semantics("rawrepresentable.is_equal")
 public func == <T: RawRepresentable>(lhs: T, rhs: T) -> Bool
   where T.RawValue: Equatable {
   return lhs.rawValue == rhs.rawValue
diff --git a/test/SILGen/protocol_operators_local_conformance.swift b/test/SILGen/protocol_operators_local_conformance.swift
@@ -64,7 +64,7 @@ func test6() {
 // CHECK: function_ref @$[[TEST6_EQUALS_WITNESS:[_0-9a-zA-Z]+]]
 // CHECK: }
 
-// CHECK: sil [serialized] @$[[TEST6_EQUALS_WITNESS]] : $@convention(thin) <τ_0_0 where τ_0_0 : RawRepresentable, τ_0_0.RawValue : Equatable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> Bool
+// CHECK: sil [serialized] {{.*}}@$[[TEST6_EQUALS_WITNESS]] : $@convention(thin) <τ_0_0 where τ_0_0 : RawRepresentable, τ_0_0.RawValue : Equatable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> Bool
 
 func test7() {
   struct Outer {
diff --git a/test/SILGen/synthesized_conformance_class.swift b/test/SILGen/synthesized_conformance_class.swift
@@ -56,7 +56,7 @@ class Nonfinal<T> {
 
 // CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}11stringValueAFyx_GSgSS_tcfC : $@convention(method) <T> (@owned String, @thin Final<T>.CodingKeys.Type) -> Optional<Final<T>.CodingKeys> {
 // CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}8intValueAFyx_GSgSi_tcfC : $@convention(method) <T> (Int, @thin Final<T>.CodingKeys.Type) -> Optional<Final<T>.CodingKeys> {
-// CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}21__derived_enum_equalsySbAFyx_G_AHtFZ : $@convention(method) <T> (Final<T>.CodingKeys, Final<T>.CodingKeys, @thin Final<T>.CodingKeys.Type) -> Bool {
+// CHECK-LABEL: sil private [_semantics "derived_enum_equals"] [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}21__derived_enum_equalsySbAFyx_G_AHtFZ : $@convention(method) <T> (Final<T>.CodingKeys, Final<T>.CodingKeys, @thin Final<T>.CodingKeys.Type) -> Bool {
 // CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}4hash4intoys6HasherVz_tF : $@convention(method) <T> (@inout Hasher, Final<T>.CodingKeys) -> () {
 // CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}9hashValueSivg : $@convention(method) <T> (Final<T>.CodingKeys) -> Int {
 // CHECK-LABEL: sil private [ossa] @$s29synthesized_conformance_class5FinalC10CodingKeys{{.*}}8intValueSiSgvg : $@convention(method) <T> (Final<T>.CodingKeys) -> Optional<Int> {
diff --git a/test/SILGen/synthesized_conformance_enum.swift b/test/SILGen/synthesized_conformance_enum.swift
@@ -37,7 +37,7 @@ enum NoValues {
 extension Enum: Equatable where T: Equatable {}
 // CHECK-FRAGILE-LABEL: // static Enum<A>.__derived_enum_equals(_:_:)
 // CHECK-FRAGILE-NEXT: // Isolation: unspecified
-// CHECK-FRAGILE-NEXT: sil hidden [ossa] @$s28synthesized_conformance_enum4EnumOAASQRzlE010__derived_C7_equalsySbACyxG_AEtFZ : $@convention(method) <T where T : Equatable> (@in_guaranteed Enum<T>, @in_guaranteed Enum<T>, @thin Enum<T>.Type) -> Bool {
+// CHECK-FRAGILE-NEXT: sil hidden [_semantics "derived_enum_equals"] [ossa] @$s28synthesized_conformance_enum4EnumOAASQRzlE010__derived_C7_equalsySbACyxG_AEtFZ : $@convention(method) <T where T : Equatable> (@in_guaranteed Enum<T>, @in_guaranteed Enum<T>, @thin Enum<T>.Type) -> Bool {
 // CHECK-RESILIENT-LABEL: // static Enum<A>.== infix(_:_:)
 // CHECK-RESILIENT-NEXT: // Isolation: unspecified
 // CHECK-RESILIENT-NEXT: sil hidden [ossa] @$s28synthesized_conformance_enum4EnumOAASQRzlE2eeoiySbACyxG_AEtFZ : $@convention(method) <T where T : Equatable> (@in_guaranteed Enum<T>, @in_guaranteed Enum<T>, @thin Enum<T>.Type) -> Bool {
diff --git a/test/SILOptimizer/enum-comparison.swift b/test/SILOptimizer/enum-comparison.swift
@@ -0,0 +1,118 @@
+// RUN: %empty-directory(%t) 
+// RUN: %target-build-swift -O -module-name=test %s -o %t/a.out
+// RUN: %target-build-swift -O -module-name=test %s -emit-ir | %FileCheck %s
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s --check-prefix=OUT
+
+// REQUIRES: executable_test,optimized_stdlib
+
+enum E: String {
+  case a, b, c, long_case_name_for_testing, d, e
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test9compareeqySbAA1EO_ADtF"(i8 %0, i8 %1)
+// CHECK:         %2 = icmp eq i8 %0, %1
+// CHECK-NEXT:    ret i1 %2
+@inline(never)
+func compareeq(_ a: E, _ b: E) -> Bool {
+  return a == b
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test9compareneySbAA1EO_ADtF"(i8 %0, i8 %1)
+// CHECK:         %2 = icmp ne i8 %0, %1
+// CHECK-NEXT:    ret i1 %2
+@inline(never)
+func comparene(_ a: E, _ b: E) -> Bool {
+  return a != b
+}
+
+enum LargeEnum: Equatable {
+  case a1, a2, a3, a4, a5, a6, a7, a8, a9
+  case b1, b2, b3, b4, b5, b6, b7, b8, b9
+  case c1, c2, c3, c4, c5, c6, c7, c8, c9
+  case d1, d2, d3, d4, d5, d6, d7, d8, d9
+  case e1(Int64), e2(Int64), e3(Int64), e4(Int64), e5(Int64), e6(Int64), e7(Int64), e8(Int64), e9(Int64)
+  case f1, f2, f3, f4, f5, f6, f7, f8, f9
+  case g1, g2, g3, g4, g5, g6, g7, g8, g9
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test8compare1ySbAA9LargeEnumOF"(i64 %0, i8 %1)
+// CHECK:       entry:
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    {{(and|select)}}
+// CHECK-NEXT:    ret
+@inline(never)
+func compare1(_ x: LargeEnum) -> Bool {
+  return x == .b2
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test8compare2ySbAA9LargeEnumOF"(i64 %0, i8 %1)
+// CHECK:       entry:
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    {{(and|select)}}
+// CHECK-NEXT:    ret
+@inline(never)
+func compare2(_ x: LargeEnum) -> Bool {
+  return .f2 == x
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test8compare3ySbAA9LargeEnumOF"(i64 %0, i8 %1)
+// CHECK:       entry:
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    {{(and|select)}}
+// CHECK-NEXT:    ret
+@inline(never)
+func compare3(_ x: LargeEnum) -> Bool {
+  return .e2(27) == x
+}
+
+// CHECK-LABEL: define {{.*}} i1 @"$s4test8compare4ySbAA9LargeEnumOF"(i64 %0, i8 %1)
+// CHECK:       entry:
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    icmp
+// CHECK-NEXT:    {{(and|select)}}
+// CHECK-NEXT:    ret
+@inline(never)
+func compare4(_ x: LargeEnum) -> Bool {
+  return x == .e3(28)
+}
+
+// OUT: 1: false
+print("1: \(compareeq(.c, .long_case_name_for_testing))")
+
+// OUT: 2: true
+print("2: \(compareeq(.c, .c))")
+
+// OUT: 3: true
+print("3: \(comparene(.c, .long_case_name_for_testing))")
+
+// OUT: 4: false
+print("4: \(comparene(.c, .c))")
+
+// OUT: 5: false
+print("5: \(compare1(.b1))")
+
+// OUT: 6: true
+print("6: \(compare1(.b2))")
+
+// OUT: 7: false
+print("7: \(compare2(.b1))")
+
+// OUT: 8: true
+print("8: \(compare2(.f2))")
+
+// OUT: 9: true
+print("9: \(compare3(.e2(27)))")
+
+// OUT: 10: false
+print("10: \(compare3(.e2(28)))")
+
+// OUT: 11: true
+print("11: \(compare4(.e3(28)))")
+
+// OUT: 12: false
+print("12: \(compare4(.e3(27)))")
+
diff --git a/test/SILOptimizer/simplify_apply.sil b/test/SILOptimizer/simplify_apply.sil
@@ -26,6 +26,10 @@ struct GenS<T> {
   var x: Int
 }
 
+enum E: String {
+  case a, b, c, d, e
+}
+
 sil @cl : $@convention(thin) () -> Int
 
 // CHECK-LABEL: sil [ossa] @thick_to_thin :
@@ -165,3 +169,31 @@ bb0(%0 : @guaranteed $Array<Any>):
   return %2
 }
 
+sil [_semantics "rawrepresentable.is_equal"] @rawrepresentable_is_equal : $@convention(thin) <T where T : RawRepresentable, T.RawValue : Equatable> (@in_guaranteed T, @in_guaranteed T) -> Bool
+sil [_semantics "rawrepresentable.is_equal"] @rawrepresentable_is_equal_wrong_convention : $@convention(thin) (E, E) -> Bool
+
+// CHECK-LABEL: sil [ossa] @string_enum_is_equal :
+// CHECK:         %2 = builtin "getEnumTag"<E>(%0) : $Builtin.Int32
+// CHECK:         %3 = builtin "getEnumTag"<E>(%1) : $Builtin.Int32
+// CHECK:         %4 = builtin "cmp_eq_Int32"(%2, %3) : $Builtin.Int1
+// CHECK:         %5 = struct $Bool (%4)
+// CHECK:         return %5
+// CHECK:       } // end sil function 'string_enum_is_equal'
+sil [ossa] @string_enum_is_equal : $@convention(thin) (@in_guaranteed E, @in_guaranteed E) -> Bool {
+bb0(%0 : $*E, %1 : $*E):
+  %2 = function_ref @rawrepresentable_is_equal : $@convention(thin) <T where T : RawRepresentable, T.RawValue : Equatable> (@in_guaranteed T, @in_guaranteed T) -> Bool
+  %3 = apply %2<E>(%0, %1) : $@convention(thin) <τ_0_0 where τ_0_0 : RawRepresentable, τ_0_0.RawValue : Equatable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> Bool
+  return %3
+}
+
+// CHECK-LABEL: sil [ossa] @string_enum_is_equal_wrong_convention :
+// CHECK:         function_ref
+// CHECK:         apply
+// CHECK:       } // end sil function 'string_enum_is_equal_wrong_convention'
+sil [ossa] @string_enum_is_equal_wrong_convention : $@convention(thin) (E, E) -> Bool {
+bb0(%0 : $E, %1 : $E):
+  %2 = function_ref @rawrepresentable_is_equal_wrong_convention : $@convention(thin) (E, E) -> Bool
+  %3 = apply %2(%0, %1) : $@convention(thin) (E, E) -> Bool
+  return %3
+}
+
diff --git a/test/api-digester/stability-stdlib-abi-without-asserts.test b/test/api-digester/stability-stdlib-abi-without-asserts.test
@@ -857,6 +857,10 @@ Var UnsafeMutableBufferPointer.indices has mangled name changing from 'Swift.Uns
 Var UnsafeMutableBufferPointer.indices is now with @_preInverseGenerics
 Func !=(_:_:) has been removed
 Func ==(_:_:) has been removed
+Func ==(_:_:) has generic signature change from to <T where T : Swift.RawRepresentable, T.RawValue : Swift.Equatable>
+Func ==(_:_:) has mangled name changing from 'Swift.== infix(Swift.Optional<Any.Type>, Swift.Optional<Any.Type>) -> Swift.Bool' to 'Swift.== infix<A where A: Swift.RawRepresentable, A.RawValue: Swift.Equatable>(A, A) -> Swift.Bool'
+Func ==(_:_:) has parameter 0 type change from (any Any.Type)? to τ_0_0
+Func ==(_:_:) has parameter 1 type change from (any Any.Type)? to τ_0_0                              
 Func type(of:) has been removed
 
 // *** DO NOT DISABLE OR XFAIL THIS TEST. *** (See comment above.)
