use RespectOriginallyDefinedIn when mangling extension contexts

include/swift/AST/Decl.h

@@ -2132,7 +2132,10 @@ class ExtensionDecl final : public GenericContext, public Decl,
 
   /// Determine whether this extension context is in the same defining module as
   /// the original nominal type context.
-  bool isInSameDefiningModule() const;
+  ///
+  /// \param RespectOriginallyDefinedIn Whether to respect
+  /// \c @_originallyDefinedIn attributes or the actual location of the decls.
+  bool isInSameDefiningModule(bool RespectOriginallyDefinedIn = true) const;
 
   /// Determine whether this extension is equivalent to one that requires at
   /// at least some constraints to be written in the source.

lib/AST/ASTMangler.cpp

@@ -946,6 +946,8 @@ std::string ASTMangler::mangleObjCRuntimeName(const NominalTypeDecl *Nominal) {
 
 std::string ASTMangler::mangleTypeAsContextUSR(const NominalTypeDecl *type) {
   beginManglingWithoutPrefix();
+  llvm::SaveAndRestore<bool> respectOriginallyDefinedInRAII(
+      RespectOriginallyDefinedIn, false);
   llvm::SaveAndRestore<bool> allowUnnamedRAII(AllowNamelessEntities, true);
   BaseEntitySignature base(type);
   appendContext(type, base, type->getAlternateModuleName());
@@ -1014,6 +1016,8 @@ ASTMangler::mangleAnyDecl(const ValueDecl *Decl,
 
 std::string ASTMangler::mangleDeclAsUSR(const ValueDecl *Decl,
                                         StringRef USRPrefix) {
+  llvm::SaveAndRestore<bool> respectOriginallyDefinedInRAII(
+      RespectOriginallyDefinedIn, false);
   return (llvm::Twine(USRPrefix) + mangleAnyDecl(Decl, false)).str();
 }
 
@@ -1022,6 +1026,8 @@ std::string ASTMangler::mangleAccessorEntityAsUSR(AccessorKind kind,
                                                   StringRef USRPrefix,
                                                   bool isStatic) {
   beginManglingWithoutPrefix();
+  llvm::SaveAndRestore<bool> respectOriginallyDefinedInRAII(
+      RespectOriginallyDefinedIn, false);
   llvm::SaveAndRestore<bool> allowUnnamedRAII(AllowNamelessEntities, true);
   Buffer << USRPrefix;
   appendAccessorEntity(getCodeForAccessorKind(kind), decl, isStatic);
@@ -3049,9 +3055,9 @@ void ASTMangler::appendExtension(const ExtensionDecl* ext,
   // "extension is to a protocol" would no longer be a reason to use the
   // extension mangling, because an extension method implementation could be
   // resiliently moved into the original protocol itself.
-  if (ext->isInSameDefiningModule() // case 1
-        && !sigParts.hasRequirements() // case 2
-        && !ext->getDeclaredInterfaceType()->isExistentialType()) { // case 3
+  if (ext->isInSameDefiningModule(RespectOriginallyDefinedIn)     // case 1
+      && !sigParts.hasRequirements()                              // case 2
+      && !ext->getDeclaredInterfaceType()->isExistentialType()) { // case 3
     // skip extension mangling
     return appendAnyGenericType(decl);
   }

lib/AST/Decl.cpp

@@ -2169,10 +2169,13 @@ bool ExtensionDecl::isWrittenWithConstraints() const {
   return false;
 }
 
-bool ExtensionDecl::isInSameDefiningModule() const {
+bool ExtensionDecl::isInSameDefiningModule(
+    bool RespectOriginallyDefinedIn) const {
   auto decl = getExtendedNominal();
-  auto extensionAlterName = getAlternateModuleName();
-  auto typeAlterName = decl->getAlternateModuleName();
+  auto extensionAlterName =
+      RespectOriginallyDefinedIn ? getAlternateModuleName() : "";
+  auto typeAlterName =
+      RespectOriginallyDefinedIn ? decl->getAlternateModuleName() : "";
 
   if (!extensionAlterName.empty()) {
     if (!typeAlterName.empty()) {

test/SymbolGraph/Symbols/OriginallyDefinedInExtension.swift

@@ -0,0 +1,22 @@
+// RUN: %empty-directory(%t)
+// RUN: %empty-directory(%t/macos)
+// RUN: %empty-directory(%t/ios)
+
+// RUN: %target-swift-frontend -target %target-cpu-apple-macos %s -module-name OriginallyDefinedInExtension -emit-module -emit-module-path %t/macos/OriginallyDefinedInExtension.swiftmodule -emit-symbol-graph -emit-symbol-graph-dir %t/macos/
+// RUN: %FileCheck %s --input-file %t/macos/OriginallyDefinedInExtension.symbols.json
+// RUN: %target-swift-frontend -target %target-cpu-apple-ios-simulator %s -module-name OriginallyDefinedInExtension -emit-module -emit-module-path %t/ios/OriginallyDefinedInExtension.swiftmodule -emit-symbol-graph -emit-symbol-graph-dir %t/ios/
+// RUN: %FileCheck %s --input-file %t/ios/OriginallyDefinedInExtension.symbols.json
+
+// CHECK:   "precise":"s:28OriginallyDefinedInExtension12SimpleStructV05InnerF0V"
+
+// REQUIRES: SWIFT_SDK=osx
+// REQUIRES: SWIFT_SDK=ios_simulator
+
+@available(macOS 10.8, *)
+@_originallyDefinedIn(module: "another", macOS 11.0)
+public struct SimpleStruct {}
+
+@available(macOS 12.0, iOS 13.0, *)
+public extension SimpleStruct {
+    struct InnerStruct {}
+}

test/lit.cfg

@@ -905,6 +905,11 @@ config.available_features.add("SWIFT_VERSION=" + swift_version)
 
 config.available_features.add("STDLIB_VARIANT={}".format(config.variant_suffix[1:]))
 
+if "target-same-as-host" in config.available_features:
+    # Only add SWIFT_SDKS features if we're building host tools
+    for sdk in config.swift_sdks:
+        config.available_features.add("SWIFT_SDK=" + sdk.lower())
+
 if "optimized_stdlib" in config.available_features:
   config.available_features.add("optimized_stdlib_" + run_cpu)
 

test/lit.site.cfg.in

@@ -173,6 +173,8 @@ config.freestanding_sdk_name = "@SWIFT_SDK_FREESTANDING_LIB_SUBDIR@"
 if '@SWIFT_BUILD_SWIFT_SYNTAX@' == 'TRUE':
   config.available_features.add('swift_swift_parser')
 
+config.swift_sdks = "@SWIFT_SDKS@".split(";")
+
 # Let the main config do the real work.
 if config.test_exec_root is None:
     config.test_exec_root = os.path.dirname(lit.util.abs_path_preserve_drive(__file__))

validation-test/lit.site.cfg.in

@@ -137,6 +137,8 @@ config.swift_stdlib_enable_objc_interop = "@SWIFT_STDLIB_ENABLE_OBJC_INTEROP@" =
 # Configured in DarwinSDKs.cmake
 config.freestanding_sdk_name = "@SWIFT_SDK_FREESTANDING_LIB_SUBDIR@"
 
+config.swift_sdks = "@SWIFT_SDKS@".split(";")
+
 # Let the main config do the real work.
 config.test_exec_root = os.path.dirname(os.path.realpath(__file__))
 lit_config.load_config(config, os.path.join(config.test_exec_root, "lit.swift-features.cfg"))