Merge pull request #62090 from DougGregor/macros-import

[Macros] Import owning and supplemental modules when processing macro signatures

Author: DougGregor

lib/Sema/TypeCheckMacros.cpp

@@ -102,7 +102,9 @@ static Optional<std::pair<GenericSignature, Type>>
 getMacroSignature(
     ModuleDecl *mod, Identifier macroName,
     Optional<StringRef> genericSignature,
-    StringRef typeSignature
+    StringRef typeSignature,
+    ModuleDecl *owningModule,
+    ArrayRef<ModuleDecl *> supplementalImportModules
 ) {
   // Form a buffer containing the macro signature context.
   ASTContext &ctx = mod->getASTContext();
@@ -140,6 +142,21 @@ getMacroSignature(
   if (!decl)
     return None;
 
+  /// Add an import to the module.
+  auto addImport = [&](ModuleDecl *importedModule) {
+    ImportPath::Builder importPath;
+    importPath.push_back(importedModule->getName(), SourceLoc());
+    auto importDecl = ImportDecl::create(
+        ctx, macroSourceFile, SourceLoc(), ImportKind::Module,
+        SourceLoc(), importPath.get());
+    importDecl->setImplicit();
+    macroSourceFile->addTopLevelDecl(importDecl);
+  };
+  addImport(owningModule);
+  std::for_each(supplementalImportModules.begin(),
+                supplementalImportModules.end(),
+                addImport);
+
   // Make sure imports are resolved in this file.
   performImportResolution(*macroSourceFile);
 
@@ -183,7 +200,8 @@ static MacroDecl *createMacro(
 
   // Get the type signature of the macro.
   auto signature = getMacroSignature(
-      mod, macroName, genericSignature, typeSignature);
+      mod, macroName, genericSignature, typeSignature, owningModule,
+      supplementalImportModules);
   if (!signature) {
     // FIXME: Swap in ErrorTypes, perhaps?
     return nullptr;

test/Macros/Inputs/ColorLib.swift

@@ -0,0 +1,4 @@
+public struct MyColor: _ExpressibleByColorLiteral {
+  public init(_colorLiteralRed red: Float, green: Float, blue: Float, alpha: Float) { }
+  public init(_colorLiteralHue hue: Float, saturation: Float, value: Float) { }
+}

test/Macros/Inputs/macro_definition.swift

@@ -130,7 +130,7 @@ struct ColorLiteralMacro: _CompilerPlugin {
   }
 
   static func _owningModule() -> (UnsafePointer<UInt8>, count: Int) {
-    var swiftModule = "Swift"
+    var swiftModule = "ColorLib"
     return swiftModule.withUTF8 { buffer in
       let result = UnsafeMutablePointer<UInt8>.allocate(capacity: buffer.count)
       result.initialize(from: buffer.baseAddress!, count: buffer.count)
@@ -191,20 +191,15 @@ struct HSVColorLiteralMacro: _CompilerPlugin {
   }
 
   static func _genericSignature() -> (UnsafePointer<UInt8>?, count: Int) {
-    var genSig = "<T>"
-    return genSig.withUTF8 { buffer in
-      let result = UnsafeMutablePointer<UInt8>.allocate(capacity: buffer.count)
-      result.initialize(from: buffer.baseAddress!, count: buffer.count)
-      return (UnsafePointer(result), count: buffer.count)
-    }
+    return (nil, count: 0)
   }
 
   static func _typeSignature() -> (UnsafePointer<UInt8>, count: Int) {
     var typeSig =
       """
       (
        hue hue: Float, saturation saturation: Float, value value: Float
-      ) -> T
+      ) -> MyColor
       """
     return typeSig.withUTF8 { buffer in
       let result = UnsafeMutablePointer<UInt8>.allocate(capacity: buffer.count)
@@ -214,7 +209,7 @@ struct HSVColorLiteralMacro: _CompilerPlugin {
   }
 
   static func _owningModule() -> (UnsafePointer<UInt8>, count: Int) {
-    var swiftModule = "Swift"
+    var swiftModule = "ColorLib"
     return swiftModule.withUTF8 { buffer in
       let result = UnsafeMutablePointer<UInt8>.allocate(capacity: buffer.count)
       result.initialize(from: buffer.baseAddress!, count: buffer.count)

test/Macros/macro_plugin.swift

@@ -1,12 +1,14 @@
 // RUN: %empty-directory(%t)
 // RUN: %target-build-swift -Xfrontend -disable-availability-checking -I%platform-module-dir/../.. -L%platform-dylib-dir/../.. -emit-library -emit-library-path=%t/%target-library-name(MacroDefinition) -working-directory=%t -module-name=MacroDefinition %S/Inputs/macro_definition.swift
-// RUN: %target-swift-frontend -L%platform-dylib-dir/../.. -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -disable-availability-checking -dump-ast -primary-file %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-module -o %t %S/Inputs/ColorLib.swift
+// RUN: %target-swift-frontend -L%platform-dylib-dir/../.. -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -I %t -disable-availability-checking -dump-ast -primary-file %s | %FileCheck %s
 
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
 // rdar://102160067
 // UNSUPPORTED: CPU=arm64e
+import ColorLib
 
 let _ = #customStringify(1.byteSwapped + 2.advanced(by: 10))
 
@@ -35,11 +37,6 @@ let _ = #customStringify(["a", "b", "c"] + ["d", "e", "f"])
 // CHECK:     (binary_expr type='[String]'
 // CHECK:     (string_literal_expr type='String'
 
-struct MyColor: _ExpressibleByColorLiteral {
-  init(_colorLiteralRed red: Float, green: Float, blue: Float, alpha: Float) { }
-  init(_colorLiteralHue hue: Float, saturation: Float, value: Float) { }
-}
-
 // CHECK: (macro_expansion_expr type='MyColor' {{.*}} name=customColorLiteral
 let _: MyColor = #customColorLiteral(red: 0.5, green: 0.5, blue: 0.2, alpha: 0.9)