Merge pull request #70148 from xedin/fix-perf-regression-with-keypath-inference

[TypeChecker] Improve performance of key path expression inference

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -519,6 +519,10 @@ class TypeVariableType::Implementation {
   /// Determine whether this type variable represents an opened opaque type.
   bool isOpaqueType() const;
 
+  /// Determine whether this type variable represents a type of an array literal
+  /// (represented by `ArrayExpr` in AST).
+  bool isArrayLiteralType() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///

lib/Sema/CSBindings.cpp

@@ -1253,10 +1253,21 @@ bool BindingSet::favoredOverDisjunction(Constraint *disjunction) const {
     return boundType->lookThroughAllOptionalTypes()->is<TypeVariableType>();
   }
 
+  // If this is an array literal type, it's preferrable to bind it
+  // early (unless it's delayed) to connect all of its elements even
+  // if it doesn't have any bindings.
+  if (TypeVar->getImpl().isArrayLiteralType())
+    return !involvesTypeVariables();
+
   // Don't prioritize type variables that don't have any direct bindings.
   if (Bindings.empty())
     return false;
 
+  // Always prefer key path type if it has bindings and is not delayed
+  // because that means that it was possible to infer its capability.
+  if (TypeVar->getImpl().isKeyPathType())
+    return true;
+
   return !involvesTypeVariables();
 }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -191,6 +191,10 @@ bool TypeVariableType::Implementation::isOpaqueType() const {
   return false;
 }
 
+bool TypeVariableType::Implementation::isArrayLiteralType() const {
+  return locator && locator->directlyAt<ArrayExpr>();
+}
+
 void *operator new(size_t bytes, ConstraintSystem& cs,
                    size_t alignment) {
   return cs.getAllocator().Allocate(bytes, alignment);

test/Constraints/casts.swift

@@ -349,8 +349,8 @@ func test_compatibility_coercions(_ arr: [Int], _ optArr: [Int]?, _ dict: [Strin
   _  = [i: stringAnyDict] as [String: Any] // expected-error {{cannot convert value of type 'Int' to expected dictionary key type 'String'}}
 
   // These are currently not peepholed.
-  _ = [i].self as Magic as [String] // expected-error {{cannot convert value of type 'Int' to expected element type 'String'}}
-  _ = (try [i]) as Magic as [String] // expected-error {{cannot convert value of type 'Int' to expected element type 'String'}}
+  _ = [i].self as Magic as [String] // expected-warning {{coercion from '[Int]' to '[String]' may fail; use 'as?' or 'as!' instead}}
+  _ = (try [i]) as Magic as [String] // expected-warning {{coercion from '[Int]' to '[String]' may fail; use 'as?' or 'as!' instead}}
   // expected-warning@-1 {{no calls to throwing functions occur within 'try' expression}}
 
   // These are wrong, but make sure we don't warn about the value cast always succeeding.

test/SILGen/collection_downcast.swift

@@ -297,4 +297,15 @@ func testCollectionCompatibilityCoercions(_ arr: [Int], _ optArr: [Any]?, _ set:
   // CHECK: [[CAST_FN:%.+]] = function_ref @$ss17_dictionaryUpCastySDyq0_q1_GSDyxq_GSHRzSHR0_r2_lF : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2, τ_0_3 where τ_0_0 : Hashable, τ_0_2 : Hashable> (@guaranteed Dictionary<τ_0_0, τ_0_1>) -> @owned Dictionary<τ_0_2, τ_0_3>
   // CHECK: apply [[CAST_FN]]<String, Int, Int, String>([[DICT]]) : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2, τ_0_3 where τ_0_0 : Hashable, τ_0_2 : Hashable> (@guaranteed Dictionary<τ_0_0, τ_0_1>) -> @owned Dictionary<τ_0_2, τ_0_3>
   _ = promote(promote(promote(dict))) as [Int: String]
+
+  typealias Magic<T> = T
+
+  // These are currently not peepholed.
+  // CHECK: [[CAST_FN:%.+]] = function_ref @$ss15_arrayForceCastySayq_GSayxGr0_lF : $@convention(thin) <τ_0_0, τ_0_1> (@guaranteed Array<τ_0_0>) -> @owned Array<τ_0_1>
+  // CHECK: apply [[CAST_FN]]<Int, String>
+  [i].self as Magic as [String]
+
+  // CHECK: [[CAST_FN:%.+]] = function_ref @$ss15_arrayForceCastySayq_GSayxGr0_lF : $@convention(thin) <τ_0_0, τ_0_1> (@guaranteed Array<τ_0_0>) -> @owned Array<τ_0_1>
+  // CHECK: apply [[CAST_FN]]<Int, String>
+  (try [i]) as Magic as [String]
 }

validation-test/Sema/type_checker_perf/fast/rdar118993030-literal.swift.gyb

@@ -0,0 +1,60 @@
+// RUN: %scale-test --begin 1 --end 5 --step 1 --select NumConstraintScopes %s
+// REQUIRES: asserts,no_asan
+
+struct Root {
+  enum E: Int, Comparable {
+    case test = 0
+    static func < (_ lhs: Self, _ rhs: Self) -> Bool { false }
+  }
+
+  let word: String
+  let maybeWord: String?
+  let e: E
+  let maybeE: E?
+}
+
+enum Order {
+  case forward
+}
+
+protocol P {
+}
+
+struct Descriptor<Base> {
+  public init<Value>(_ keyPath: KeyPath<Base, Value>, order: Order = .forward) where Value: Comparable { fatalError() }
+  public init<Value>(_ keyPath: KeyPath<Base, Value?>, order: Order = .forward) where Value: Comparable { fatalError() }
+
+  public init(_ keyPath: KeyPath<Base, Double>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Double?>, order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Float>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Float?>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, String>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, String?>, order: Order = .forward) { fatalError() }
+
+  public init(_ keyPath: KeyPath<Base, Int8>,    order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int8?>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int16>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int16?>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int32>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int32?>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int64>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, Int64?>,  order: Order = .forward) { fatalError() }
+
+  public init(_ keyPath: KeyPath<Base, UInt8>,   order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt8?>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt16>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt16?>, order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt32>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt32?>, order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt64>,  order: Order = .forward) { fatalError() }
+  public init(_ keyPath: KeyPath<Base, UInt64?>, order: Order = .forward) { fatalError() }
+}
+
+let _ = [
+%for i in range(0, N):
+  Descriptor(\Root.word),
+  Descriptor(\Root.maybeWord),
+  Descriptor(\Root.e),
+  Descriptor(\Root.maybeE),
+%end
+]

validation-test/Sema/type_checker_perf/fast/rdar118993030.swift.gyb

@@ -0,0 +1,14 @@
+// RUN: %scale-test --begin 2 --end 10 --step 1 --select NumConstraintScopes %s
+// REQUIRES: asserts,no_asan
+
+struct Test {
+  var id: String
+}
+
+func test(data: [Test]) -> Set<String> {
+  Set(data.map(\.id) +
+%for i in range(0, N):
+      data.map(\.id) +
+%end
+    data.map(\.id))
+}