Merge pull request #41172 from CodaFi/special-sauce

include/swift/AST/DiagnosticsSema.def

@@ -1904,6 +1904,8 @@ ERROR(extension_metatype,none,
 ERROR(extension_specialization,none,
       "constrained extension must be declared on the unspecialized generic "
       "type %0 with constraints specified by a 'where' clause", (Identifier))
+ERROR(extension_placeholder,none,
+      "cannot extend a type that contains placeholders", ())
 ERROR(extension_stored_property,none,
       "extensions must not contain stored properties", ())
 NOTE(extension_stored_property_fixit,none,

include/swift/Basic/LangOptions.h

@@ -477,6 +477,13 @@ namespace swift {
     // FrontendOptions.
     bool AllowModuleWithCompilerErrors = false;
 
+    /// Enable extensions of (sugared) bound generic types
+    ///
+    /// \code
+    /// extension [Int] { /**/ }
+    /// \endcode
+    bool EnableExperimentalBoundGenericExtensions = false;
+
     /// A helper enum to represent whether or not we customized the default
     /// ASTVerifier behavior via a frontend flag. By default, we do not
     /// customize.

include/swift/Option/FrontendOptions.td

@@ -309,6 +309,10 @@ def enable_resilience : Flag<["-"], "enable-resilience">,
 def enable_experimental_async_top_level :
   Flag<["-"], "enable-experimental-async-top-level">,
   HelpText<"Enable experimental concurrency in top-level code">;
+
+def enable_experimental_bound_generic_extensions :
+  Flag<["-"], "enable-experimental-bound-generic-extensions">,
+  HelpText<"Enable experimental support for extensions of bound generic types">;
 }
 
 // HIDDEN FLAGS

lib/Frontend/CompilerInvocation.cpp

@@ -503,6 +503,9 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   Opts.EnableExperimentalStringProcessing |=
       Args.hasArg(OPT_enable_experimental_string_processing);
 
+  Opts.EnableExperimentalBoundGenericExtensions |=
+    Args.hasArg(OPT_enable_experimental_bound_generic_extensions);
+
   Opts.DisableAvailabilityChecking |=
       Args.hasArg(OPT_disable_availability_checking);
   Opts.CheckAPIAvailabilityOnly |=

lib/Sema/TypeCheckDecl.cpp

@@ -2792,9 +2792,17 @@ ExtendedTypeRequest::evaluate(Evaluator &eval, ExtensionDecl *ext) const {
     return error();
   }
 
-  // Cannot extend a bound generic type, unless it's referenced via a
-  // non-generic typealias type.
-  if (extendedType->isSpecialized() &&
+  // Cannot extend types who contain placeholders.
+  if (extendedType->hasPlaceholder()) {
+    diags.diagnose(ext->getLoc(), diag::extension_placeholder)
+      .highlight(extendedRepr->getSourceRange());
+    return error();
+  }
+
+  // By default, the user cannot extend a bound generic type, unless it's
+  // referenced via a non-generic typealias type.
+  if (!ext->getASTContext().LangOpts.EnableExperimentalBoundGenericExtensions &&
+      extendedType->isSpecialized() &&
       !isNonGenericTypeAliasType(extendedType)) {
     diags.diagnose(ext->getLoc(), diag::extension_specialization,
                    extendedType->getAnyNominal()->getName())

lib/Sema/TypeCheckGeneric.cpp

@@ -562,8 +562,8 @@ static Type formExtensionInterfaceType(
   auto nominal = dyn_cast<NominalTypeDecl>(genericDecl);
   auto typealias = dyn_cast<TypeAliasDecl>(genericDecl);
   if (!nominal) {
-    Type underlyingType = typealias->getUnderlyingType();
-    nominal = underlyingType->getNominalOrBoundGenericNominal();
+    type = typealias->getUnderlyingType();
+    nominal = type->getNominalOrBoundGenericNominal();
   }
 
   // Form the result.

test/ModuleInterface/specialized-extension.swift

@@ -0,0 +1,40 @@
+// RUN: %target-swift-frontend -module-name Test -typecheck -emit-module-interface-path - -enable-experimental-bound-generic-extensions %s | %FileCheck %s
+
+public struct Tree<T> {
+  public struct Branch<B> {
+    public struct Nest<N> {
+      public struct Egg {}
+    }
+  }
+}
+
+// CHECK: extension Test.Tree.Branch.Nest.Egg {
+// CHECK:   public static func tweet()
+// CHECK: }
+extension Tree.Branch.Nest.Egg { public static func tweet() {} }
+
+// CHECK: extension Test.Tree.Branch.Nest.Egg where T == Swift.Int {
+// CHECK:   public static func twoot()
+// CHECK: }
+extension Tree<Int>.Branch.Nest.Egg { public static func twoot() {} }
+
+// CHECK: extension Test.Tree.Branch.Nest.Egg where T == Swift.Int, B == Swift.String {
+// CHECK:   public static func twote()
+// CHECK: }
+extension Tree<Int>.Branch<String>.Nest.Egg { public static func twote() {} }
+
+// CHECK: extension Test.Tree.Branch.Nest.Egg where T == Swift.Int, B == Swift.String, N == Swift.Void {
+// CHECK:   public static func twite()
+// CHECK: }
+extension Tree<Int>.Branch<String>.Nest<Void>.Egg { public static func twite() {} }
+
+// CHECK: extension Swift.Array where Element == Swift.String {
+// CHECK:   public func rejoinder() -> Swift.String
+// CHECK: }
+extension [String] { public func rejoinder() -> String { return self.joined() } }
+
+// CHECK: public typealias StringDict<T> = [Swift.String : T]
+public typealias StringDict<T> = [String: T]
+
+// CHECK: extension Swift.Dictionary where Key == Swift.String, Value == Swift.Int
+extension StringDict<Int> { public static func mark() {} }

test/SILGen/mangling_specialized_extensions.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-emit-silgen -enable-experimental-bound-generic-extensions %s | %FileCheck %s
+
+struct Tree<T> {
+  struct Branch<B> {
+    struct Nest<N> {
+      struct Egg {}
+    }
+  }
+}
+
+// CHECK: extension Tree.Branch.Nest.Egg {
+// CHECK:   static func tweet()
+// CHECK: }
+
+// CHECK: extension Tree.Branch.Nest.Egg where T == Int {
+// CHECK:   static func twoot()
+// CHECK: }
+
+// CHECK: extension Tree.Branch.Nest.Egg where T == Int, B == String {
+// CHECK:   static func twote()
+// CHECK: }
+
+// CHECK: extension Tree.Branch.Nest.Egg where T == Int, B == String, N == () {
+// CHECK:   static func twite()
+// CHECK: }
+
+// CHECK: @$s31mangling_specialized_extensions4TreeV6BranchV4NestV3EggV5tweetyyFZ : $@convention(method) <T><B><N> (@thin Tree<T>.Branch<B>.Nest<N>.Egg.Type) -> ()
+extension Tree.Branch.Nest.Egg { static func tweet() {} }
+
+// CHECK: @$s31mangling_specialized_extensions4TreeV6BranchV4NestV3EggVAASiRszrlE5twootyyFZ : $@convention(method) <T where T == Int><B><N> (@thin Tree<Int>.Branch<B>.Nest<N>.Egg.Type) -> ()
+extension Tree<Int>.Branch.Nest.Egg { static func twoot() {} }
+
+// CHECK: @$s31mangling_specialized_extensions4TreeV6BranchV4NestV3EggVAASiRszSSRsd__rlE5twoteyyFZ : $@convention(method) <T where T == Int><B where B == String><N> (@thin Tree<Int>.Branch<String>.Nest<N>.Egg.Type) -> ()
+extension Tree<Int>.Branch<String>.Nest.Egg { static func twote() {} }
+
+// CHECK: @$s31mangling_specialized_extensions4TreeV6BranchV4NestV3EggVAASiRszSSRsd__ytRsd0__rlE5twiteyyFZ : $@convention(method) (@thin Tree<Int>.Branch<String>.Nest<()>.Egg.Type) -> ()
+extension Tree<Int>.Branch<String>.Nest<Void>.Egg { static func twite() {} }

test/decl/ext/specialize.swift

@@ -0,0 +1,76 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-bound-generic-extensions
+
+extension Array<Int> {
+  func someIntFuncOnArray() {}
+}
+
+let _ = [0, 1, 2].someIntFuncOnArray()
+
+extension [Character] {
+  func makeString() -> String { fatalError() }
+}
+
+let _ = ["a", "b", "c"].makeString()
+let _ = [1, 2, 3].makeString() // expected-error 3 {{cannot convert value of type 'Int' to expected element type 'Character'}}
+
+extension Set<_> {} // expected-error {{cannot extend a type that contains placeholders}}
+
+// https://bugs.swift.org/browse/SR-4875
+
+struct Foo<T, U> {
+  var x: T
+  var y: U
+}
+
+typealias IntFoo<U> = Foo<Int, U>
+
+extension IntFoo where U == Int {
+  func hello() {
+    print("hello")
+  }
+}
+
+Foo(x: "test", y: 1).hello()
+
+struct MyType<TyA, TyB> {
+  var a : TyA, b : TyB
+}
+
+typealias A<T1, T2> = MyType<T2, T1>
+
+extension A {}
+
+extension A<Float, Int> {}
+extension A<Void, Void> {}
+
+struct Tree<T> {
+  struct Branch<B> {
+    struct Nest<N> {
+      struct Egg {}
+    }
+  }
+}
+
+extension Tree.Branch.Nest.Egg { static func tweet() {} }
+extension Tree<Int>.Branch.Nest.Egg { static func twoot() {} }
+extension Tree<Int>.Branch<String>.Nest.Egg { static func twote() {} }
+extension Tree<Int>.Branch<String>.Nest<Void>.Egg { static func twite() {} }
+
+func testNestedExtensions() {
+  do {
+    Tree<Void>.Branch<Void>.Nest<Void>.Egg.tweet()
+  }
+
+  do {
+    Tree<Int>.Branch<Void>.Nest<Void>.Egg.twoot()
+    Tree<Int>.Branch<Int>.Nest<Void>.Egg.twoot()
+    Tree<Int>.Branch<Int>.Nest<Int>.Egg.twoot()
+  }
+
+  do {
+    Tree<Int>.Branch<String>.Nest<Void>.Egg.twote()
+    Tree<Int>.Branch<String>.Nest<Float>.Egg.twote()
+  }
+
+  Tree<Int>.Branch<String>.Nest<Void>.Egg.twite()
+}