[ImportResolution] Override module imports for macros in namespaces

Contents in namespaces are imported to the bridging header module
`__ObjC`. When macros attached to a decl imported from a namespace are
expanded, they also end up in `__ObjC`. This prevents them from
inheriting imports from the module they originated in. This patch
implements a workaround by adding an alternative constructor to
`ImportResolver` that takes an explicit origin module as input,
overriding which module counts as the "parent" module.

Author: hnrklssn

include/swift/Subsystems.h

@@ -125,7 +125,8 @@ namespace swift {
 
   /// Resolve imports for a source file generated to adapt a given
   /// Clang module.
-  void performImportResolutionForClangMacroBuffer(SourceFile &SF);
+  void performImportResolutionForClangMacroBuffer(
+      SourceFile &SF, ModuleDecl *explicitOriginModule);
 
   /// Once type-checking is complete, this instruments code with calls to an
   /// intrinsic that record the expected values of local variables so they can

lib/Sema/ImportResolution.cpp

@@ -197,6 +197,14 @@ class ImportResolver final : public DeclVisitor<ImportResolver> {
     addImplicitImports();
   }
 
+  // FIXME: Remove this ctor once namespace contents are imported to their
+  // corresponding module instead of the __ObjC module
+  ImportResolver(SourceFile &SF, ModuleDecl *explicitUnderlyingClangModule)
+      : SF(SF), ctx(SF.getASTContext()),
+        underlyingClangModule(explicitUnderlyingClangModule) {
+    addImplicitImports();
+  }
+
   void addImplicitImports();
 
   void addImplicitImport(ModuleDecl *module) {
@@ -326,12 +334,13 @@ void swift::performImportResolution(SourceFile &SF) {
   verify(SF);
 }
 
-void swift::performImportResolutionForClangMacroBuffer(SourceFile &SF) {
+void swift::performImportResolutionForClangMacroBuffer(
+    SourceFile &SF, ModuleDecl *explicitOriginModule) {
   assert(SF.ASTStage == SourceFile::Unprocessed);
 
   // `getWrapperForModule` has already declared all the implicit clang module
   // imports we need
-  ImportResolver resolver(SF);
+  ImportResolver resolver(SF, explicitOriginModule);
   SF.setImports(resolver.getFinishedImports());
   SF.setImportedUnderlyingModule(resolver.getUnderlyingClangModule());
 
@@ -597,7 +606,19 @@ ModuleImplicitImportsRequest::evaluate(Evaluator &evaluator,
 }
 
 void ImportResolver::addImplicitImports() {
-  auto implicitImports = SF.getParentModule()->getImplicitImports();
+  // This is a workaround for the fact that namespaces are imported into the
+  // bridging header module. To allow correct lookup, we expose a ctor with an
+  // explicit underlying clang module input, which is the clang module that this
+  // macro expansion actually originates in.
+  const ModuleDecl *moduleToInherit = nullptr;
+  if (underlyingClangModule) {
+    moduleToInherit = underlyingClangModule;
+    boundImports.push_back(
+        AttributedImport(ImportedModule(underlyingClangModule)));
+  } else {
+    moduleToInherit = SF.getParentModule();
+  }
+  auto implicitImports = moduleToInherit->getImplicitImports();
 
   // TODO: Support cross-module imports.
   for (auto &import : implicitImports.imports) {

lib/Sema/TypeCheckMacros.cpp

@@ -1061,7 +1061,12 @@ createMacroSourceFile(std::unique_ptr<llvm::MemoryBuffer> buffer,
   if (auto parentSourceFile = dc->getParentSourceFile())
     macroSourceFile->setImports(parentSourceFile->getImports());
   else if (isa<ClangModuleUnit>(dc->getModuleScopeContext())) {
-    performImportResolutionForClangMacroBuffer(*macroSourceFile);
+    ModuleDecl *originModule = nullptr;
+    // FIXME: remove this workaround once namespace contents are imported into
+    // their corresponding modules
+    if (macroSourceFile->getParentModule()->isClangHeaderImportModule())
+      originModule = cast<Decl *>(target)->getModuleContextForNameLookup();
+    performImportResolutionForClangMacroBuffer(*macroSourceFile, originModule);
   }
   return macroSourceFile;
 }

test/Interop/Cxx/swiftify-import/counted-by-in-namespace.swift

@@ -1,28 +1,66 @@
 // REQUIRES: swift_feature_SafeInteropWrappers
 
-// RUN: rm -rf %t
+// RUN: %empty-directory(%t)
 // RUN: split-file %s %t
-// RUN: %target-swift-ide-test -plugin-path %swift-plugin-dir -I %t/Inputs -cxx-interoperability-mode=upcoming-swift -enable-experimental-feature SafeInteropWrappers -print-module -module-to-print=Namespace -source-filename=x | %FileCheck %s
-// RUN: %target-swift-frontend -plugin-path %swift-plugin-dir -I %t/Inputs -cxx-interoperability-mode=default -enable-experimental-feature SafeInteropWrappers %t/namespace.swift -dump-macro-expansions -typecheck -verify
+// RUN: %target-swift-frontend -plugin-path %swift-plugin-dir -I %t/Inputs -cxx-interoperability-mode=default -enable-experimental-feature SafeInteropWrappers %t/namespace.swift -emit-module -verify
+// RUN: %target-swift-ide-test -plugin-path %swift-plugin-dir -I %t/Inputs -cxx-interoperability-mode=upcoming-swift -enable-experimental-feature SafeInteropWrappers -print-module -module-to-print=Namespace -source-filename=x > %t/interface.txt
+// RUN: diff %t/interface.txt %t/interface.txt.expected
 
-// CHECK: enum foo {
-// CHECK:  static func bar(_ p: UnsafeMutableBufferPointer<Float>)
+//--- interface.txt.expected
+@_exported import Namespace.Baz
 
+enum foo {
+  static func foo_func(_ p: UnsafeMutablePointer<Float>!, _ len: Int32)
+  /// This is an auto-generated wrapper for safer interop
+  @_alwaysEmitIntoClient @_disfavoredOverload public static func foo_func(_ p: UnsafeMutableBufferPointer<Float>)
+  static func foo_func2(_ p: UnsafePointer<baz_t>!, _ len: baz_t)
+  /// This is an auto-generated wrapper for safer interop
+  @available(visionOS 1.0, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+  @_alwaysEmitIntoClient @_disfavoredOverload public static func foo_func2(_ p: Span<baz_t>)
+  enum bar {
+    static func bar_func(_ p: UnsafePointer<baz_t>!, _ len: baz_t)
+    /// This is an auto-generated wrapper for safer interop
+    @available(visionOS 1.0, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+    @_alwaysEmitIntoClient @_disfavoredOverload public static func bar_func(_ p: Span<baz_t>)
+  }
+}
 //--- Inputs/module.modulemap
 module Namespace {
     header "namespace.h"
     requires cplusplus
+    module Baz {
+      header "baz.h"
+    }
 }
 
 //--- Inputs/namespace.h
+#include "baz.h"
+
+#define __counted_by(x) __attribute__((__counted_by__(x)))
+#define __noescape __attribute__((noescape))
 
 namespace foo {
-   __attribute__((swift_attr("@_SwiftifyImport(.countedBy(pointer: .param(1), count: \"len\"))"))) void bar(float *p, int len);
+  __attribute__((swift_attr("@_SwiftifyImport(.countedBy(pointer: .param(1), count: \"len\"))"))) void foo_func(float *p, int len);
+  void foo_func2(const baz_t * __counted_by(len) p __noescape, baz_t len);
+  namespace bar {
+    void bar_func(const baz_t * __counted_by(len) p __noescape, baz_t len);
+  }
 }
 
+//--- Inputs/baz.h
+typedef int baz_t;
+
 //--- namespace.swift
 import Namespace
-  
+
 func test(s: UnsafeMutableBufferPointer<Float>) {
-  foo.bar(s)
-}    
+  foo.foo_func(s)
+}
+
+func test2(s: Span<CInt>) {
+  foo.foo_func2(s)
+}
+
+func test3(s: Span<CInt>) {
+  foo.bar.bar_func(s)
+}