[swift_newtype] Special handling for Notifications

We now specially import global decls who we identify as fitting the
notification pattern: extern NSStrings whose name ends in
"Notification". When we see them, we import them as a member of
NSNotificationName and, if NSNotificationName is swift_newtype-ed, we
use that new type.

Test cases included.

lib/ClangImporter/ClangImporter.cpp

@@ -2091,6 +2091,52 @@ hasOrInheritsSwiftBridgeAttr(const clang::ObjCInterfaceDecl *objcClass) {
   return false;
 }
 
+/// Skip a leading 'k' in a 'kConstant' pattern
+static StringRef stripLeadingK(StringRef name) {
+  if (name.size() >= 2 && name[0] == 'k' &&
+      clang::isUppercase(name[1]))
+    return name.drop_front(1);
+  return name;
+}
+
+/// Strips a trailing "Notification", if present. Returns {} if name doesn't end
+/// in "Notification", or it there would be nothing left.
+static StringRef stripNotification(StringRef name) {
+  name = stripLeadingK(name);
+  StringRef notification = "Notification";
+  if (name.size() <= notification.size() || !name.endswith(notification))
+    return {};
+  return name.drop_back(notification.size());
+}
+
+bool ClangImporter::Implementation::isNSNotificationGlobal(
+    const clang::NamedDecl *decl) {
+  // Looking for: extern NSString *fooNotification;
+
+  // Must be extern global variable
+  auto vDecl = dyn_cast<clang::VarDecl>(decl);
+  if (!vDecl || !vDecl->hasExternalFormalLinkage())
+    return false;
+
+  // No explicit swift_name
+  if (decl->getAttr<clang::SwiftNameAttr>())
+    return false;
+
+  // Must end in Notification
+  if (!vDecl->getDeclName().isIdentifier())
+    return false;
+  if (StringRef() == stripNotification(vDecl->getName()))
+    return false;
+
+  // Must be NSString *
+  if (!isNSString(vDecl->getType()))
+    return false;
+
+  // We're a match!
+  return true;
+}
+
+
 /// Whether the decl is from a module who requested import-as-member inference
 static bool moduleIsInferImportAsMember(const clang::NamedDecl *decl,
                                         clang::Sema &clangSema) {
@@ -2123,14 +2169,41 @@ static bool moduleIsInferImportAsMember(const clang::NamedDecl *decl,
 
 // If this decl is associated with a swift_newtype typedef, return it, otherwise
 // null
-static clang::TypedefNameDecl *findSwiftNewtype(
-     ClangImporter::Implementation &impl,
-     const clang::Decl *decl,
-     bool useSwift2Name) {
-  if (auto varDecl = dyn_cast<clang::VarDecl>(decl))
-    if (auto typedefTy = varDecl->getType()->getAs<clang::TypedefType>())
-      if (impl.getSwiftNewtypeAttr(typedefTy->getDecl(), useSwift2Name))
-        return typedefTy->getDecl();
+clang::TypedefNameDecl *ClangImporter::Implementation::findSwiftNewtype(
+    const clang::NamedDecl *decl, clang::Sema &clangSema, bool useSwift2Name) {
+  // If we aren't honoring the swift_newtype attribute, don't even
+  // bother looking. Similarly for swift2 names
+  if (!HonorSwiftNewtypeAttr || useSwift2Name)
+    return nullptr;
+
+  auto varDecl = dyn_cast<clang::VarDecl>(decl);
+  if (!varDecl)
+    return nullptr;
+
+  if (auto typedefTy = varDecl->getType()->getAs<clang::TypedefType>())
+    if (getSwiftNewtypeAttr(typedefTy->getDecl(), false))
+      return typedefTy->getDecl();
+
+  // Special case: "extern NSString * fooNotification" adopts
+  // NSNotificationName type, and is a member of NSNotificationName
+  if (ClangImporter::Implementation::isNSNotificationGlobal(decl)) {
+    clang::IdentifierInfo *notificationName =
+        &clangSema.getASTContext().Idents.get("NSNotificationName");
+    clang::LookupResult lookupResult(clangSema, notificationName,
+                                     clang::SourceLocation(),
+                                     clang::Sema::LookupOrdinaryName);
+    if (!clangSema.LookupName(lookupResult, nullptr))
+      return nullptr;
+    auto nsDecl = lookupResult.getAsSingle<clang::TypedefNameDecl>();
+    if (!nsDecl)
+      return nullptr;
+
+    // Make sure it also has a newtype decl on it
+    if (getSwiftNewtypeAttr(nsDecl, false))
+      return nsDecl;
+
+    return nullptr;
+  }
 
   return nullptr;
 }
@@ -2243,6 +2316,24 @@ static clang::SwiftNameAttr *findSwiftNameAttr(const clang::Decl *decl,
   return nullptr;
 }
 
+/// Prepare global name for importing onto a swift_newtype.
+static StringRef determineSwiftNewtypeBaseName(StringRef baseName,
+                                               StringRef newtypeName) {
+  baseName = stripLeadingK(baseName);
+
+  bool nonIdentifier = false;
+  auto pre = getCommonWordPrefix(newtypeName, baseName, nonIdentifier);
+  if (pre.size())
+    baseName = baseName.drop_front(pre.size());
+
+  // Special case: Strip Notification for NSNotificationName
+  auto stripped = stripNotification(baseName);
+  if (stripped != StringRef())
+    baseName = stripped;
+
+  return baseName;
+}
+
 auto ClangImporter::Implementation::importFullName(
        const clang::NamedDecl *D,
        ImportNameOptions options,
@@ -2281,7 +2372,7 @@ auto ClangImporter::Implementation::importFullName(
       break;
     }
   // Import onto a swift_newtype if present
-  } else if (auto newtypeDecl = findSwiftNewtype(*this, D, swift2Name)) {
+  } else if (auto newtypeDecl = findSwiftNewtype(D, clangSema, swift2Name)) {
     result.EffectiveContext = newtypeDecl;
   // Everything else goes into its redeclaration context.
   } else {
@@ -2730,22 +2821,9 @@ auto ClangImporter::Implementation::importFullName(
 
   // swift_newtype-ed declarations may have common words with the type name
   // stripped.
-  if (auto newtypeDecl = findSwiftNewtype(*this, D, swift2Name)) {
-    // Skip a leading 'k' in a 'kConstant' pattern
-    if (baseName.size() >= 2 && baseName[0] == 'k' &&
-        clang::isUppercase(baseName[1])) {
-      baseName = baseName.drop_front(1);
-      strippedPrefix = true;
-    }
-
-    bool nonIdentifier = false;
-    auto pre =
-        getCommonWordPrefix(newtypeDecl->getName(), baseName, nonIdentifier);
-
-    if (pre.size()) {
-      baseName = baseName.drop_front(pre.size());
-      strippedPrefix = true;
-    }
+  if (auto newtypeDecl = findSwiftNewtype(D, clangSema, swift2Name)) {
+    baseName = determineSwiftNewtypeBaseName(baseName, newtypeDecl->getName());
+    strippedPrefix = true;
   }
 
   if (!result.isSubscriptAccessor() && !swift2Name) {
@@ -4305,7 +4383,8 @@ ClangImporter::Implementation::SwiftNameLookupExtension::hashExtension(
                             SWIFT_LOOKUP_TABLE_VERSION_MAJOR,
                             SWIFT_LOOKUP_TABLE_VERSION_MINOR,
                             Impl.InferImportAsMember,
-                            Impl.SwiftContext.LangOpts.StripNSPrefix);
+                            Impl.SwiftContext.LangOpts.StripNSPrefix,
+                            Impl.HonorSwiftNewtypeAttr);
 }
 
 std::unique_ptr<clang::ModuleFileExtensionWriter>

lib/ClangImporter/ImportDecl.cpp

@@ -3202,7 +3202,6 @@ namespace {
       if (!importedName) return nullptr;
 
       auto name = importedName.Imported.getBaseName();
-
       auto dc = Impl.importDeclContextOf(decl, importedName.EffectiveContext);
       if (!dc)
         return nullptr;
@@ -3212,7 +3211,15 @@ namespace {
       // involve an ownership transfer.
       bool isAudited = decl->getType().isConstQualified();
 
-      Type type = Impl.importType(decl->getType(),
+      auto declType = decl->getType();
+
+      // Special case: NS Notifications
+      if (ClangImporter::Implementation::isNSNotificationGlobal(decl))
+        if (auto newtypeDecl =
+                Impl.findSwiftNewtype(decl, Impl.getClangSema(), false))
+          declType = Impl.getClangASTContext().getTypedefType(newtypeDecl);
+
+      Type type = Impl.importType(declType,
                                   (isAudited ? ImportTypeKind::AuditedVariable
                                    : ImportTypeKind::Variable),
                                   isInSystemModule(dc),

lib/ClangImporter/ImportType.cpp

@@ -1261,6 +1261,18 @@ Type ClangImporter::Implementation::importType(clang::QualType type,
                                      optionality);
 }
 
+bool ClangImporter::Implementation::isNSString(const clang::Type *type) {
+  if (auto ptrType = type->getAs<clang::ObjCObjectPointerType>())
+    if (auto interfaceType = ptrType->getInterfaceType())
+      if (interfaceType->getDecl()->getName() == "NSString")
+        return true;
+  return false;
+}
+
+bool ClangImporter::Implementation::isNSString(clang::QualType qt) {
+  return qt.getTypePtrOrNull() && isNSString(qt.getTypePtrOrNull());
+}
+
 bool ClangImporter::Implementation::shouldImportGlobalAsLet(
        clang::QualType type)
 {
@@ -1269,13 +1281,9 @@ bool ClangImporter::Implementation::shouldImportGlobalAsLet(
     return true;
   }
   // Globals of type NSString * should be imported as 'let'.
-  if (auto ptrType = type->getAs<clang::ObjCObjectPointerType>()) {
-    if (auto interfaceType = ptrType->getInterfaceType()) {
-      if (interfaceType->getDecl()->getName() == "NSString") {
-        return true;
-      }
-    }
-  }
+  if (isNSString(type))
+    return true;
+
   return false;
 }
 

lib/ClangImporter/ImporterImpl.h

@@ -1459,6 +1459,20 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
 
   /// Dump the Swift-specific name lookup tables we generate.
   void dumpSwiftLookupTables();
+
+  /// Whether the given decl is a global Notification
+  static bool isNSNotificationGlobal(const clang::NamedDecl *);
+
+  // If this decl is associated with a swift_newtype (and we're honoring
+  // swift_newtype), return it, otherwise null
+  clang::TypedefNameDecl *findSwiftNewtype(const clang::NamedDecl *decl,
+                                           clang::Sema &clangSema,
+                                           bool useSwift2Name);
+
+  /// Whether the passed type is NSString *
+  static bool isNSString(const clang::Type *);
+  static bool isNSString(clang::QualType);
+
 };
 
 }

test/IDE/Inputs/custom-modules/Newtype.h

@@ -36,3 +36,16 @@ typedef NSString * NSURLResourceKey __attribute((swift_newtype(struct)));
 extern NSURLResourceKey const NSURLIsRegularFileKey;
 extern NSURLResourceKey const NSURLIsDirectoryKey;
 extern NSURLResourceKey const NSURLLocalizedNameKey;
+
+// Special case: Notifications
+extern const NSString *FooNotification;
+extern const NSString *kBarNotification;
+
+// But not just 'Notification'
+extern const NSString *kNotification;
+extern const NSString *Notification;
+
+// Nor when explicitly swift_name-ed
+extern const NSString *kSNNotification
+    __attribute((swift_name("swiftNamedNotification")));
+

test/IDE/newtype.swift

@@ -50,8 +50,15 @@
 // PRINT-NEXT:    static let isDirectoryKey: NSURLResourceKey
 // PRINT-NEXT:    static let localizedNameKey: NSURLResourceKey
 // PRINT-NEXT:  }
+// PRINT-NEXT:  extension NSNotificationName {
+// PRINT-NEXT:    static let foo: NSNotificationName
+// PRINT-NEXT:    static let bar: NSNotificationName
+// PRINT-NEXT:  }
+// PRINT-NEXT:  let kNotification: String
+// PRINT-NEXT:  let Notification: String
+// PRINT-NEXT:  let swiftNamedNotification: String
 
-// RUN: %target-parse-verify-swift -I %S/Inputs/custom-modules -enable-swift-newtype
+// RUN: %target-parse-verify-swift -sdk %clang-importer-sdk -I %S/Inputs/custom-modules -enable-swift-newtype
 import Newtype
 
 func tests() {
@@ -69,4 +76,6 @@ func tests() {
 	let _ = ErrorDomain(rawValue: thirdEnum.rawValue!)
 	let _ = ClosedEnum(rawValue: errOne.rawValue)
 
+	let _ = NSNotificationName.foo
+	let _ = NSNotificationName.bar
 }

test/IRGen/newtype.swift

@@ -1,4 +1,4 @@
-// RUN// RUN: %target-swift-frontend -I %S/../IDE/Inputs/custom-modules %s -emit-ir -enable-swift-newtype | FileCheck %s
+// RUN// RUN: %target-swift-frontend -sdk %clang-importer-sdk -I %S/../IDE/Inputs/custom-modules %s -emit-ir -enable-swift-newtype | FileCheck %s
 import Newtype
 
 // REQUIRES: objc_interop
@@ -9,3 +9,24 @@ public func getErrorDomain() -> ErrorDomain {
   return .one
 }
 
+// CHECK-LABEL: _TF7newtype6getFooFT_VSC18NSNotificationName
+public func getFoo() -> NSNotificationName {
+  return NSNotificationName.foo
+  // CHECK: load {{.*}} @FooNotification
+  // CHECK: ret
+}
+
+// CHECK-LABEL: _TF7newtype21getGlobalNotificationFSiSS
+public func getGlobalNotification(_ x: Int) -> String {
+  switch x {
+    case 1: return kNotification
+    // CHECK: load {{.*}} @kNotification
+    case 2: return Notification
+    // CHECK: load {{.*}} @Notification
+    case 3: return swiftNamedNotification
+    // CHECK: load {{.*}} @kSNNotification
+    default: return NSNotificationName.bar.rawValue
+    // CHECK: load {{.*}} @kBarNotification
+  }
+// CHECK: ret
+}

test/Inputs/clang-importer-sdk/usr/include/Foundation.h

@@ -764,6 +764,9 @@ extern void CGColorRelease(CGColorRef color) __attribute__((availability(macosx,
 @property (readonly) Class classForPortCoder NS_SWIFT_UNAVAILABLE("Use NSXPCConnection instead");
 @end
 
+typedef NSString *__nonnull NSNotificationName
+    __attribute((swift_newtype(struct)));
+
 NS_SWIFT_UNAVAILABLE("Use NSXPCConnection instead")
 extern NSString * const NSConnectionReplyMode;
 NS_SWIFT_UNAVAILABLE("Use NSXPCConnection instead")