[Macros] Use the macro recorded in the constraint system for expansion.

Rather than lookup up the macro by name again during constraint
application, use the overload choice for the macro as recorded in the
constraint system to apply the correct macro.

Author: DougGregor

include/swift/Sema/ConstraintSystem.h

@@ -4906,14 +4906,6 @@ class ConstraintSystem {
                           ConstraintLocator *locator,
                           OpenedTypeMap *replacements = nullptr);
 
-#if SWIFT_SWIFT_PARSER
-  /// Retrieve the opened type of a macro with the given name.
-  ///
-  /// \returns The opened type of the macro with this name, or the null \c Type
-  /// if no such macro exists.
-  Type getTypeOfMacroReference(Identifier macroName, Expr *anchor);
-#endif
-
   /// Retrieve a list of generic parameter types solver has "opened" (replaced
   /// with a type variable) at the given location.
   ArrayRef<OpenedType> getOpenedTypes(ConstraintLocator *locator) const {

lib/Sema/CSApply.cpp

@@ -2975,14 +2975,20 @@ namespace {
 #if SWIFT_SWIFT_PARSER
       auto &ctx = cs.getASTContext();
       if (ctx.LangOpts.hasFeature(Feature::BuiltinMacros)) {
-        auto kind = MagicIdentifierLiteralExpr::getKindString(expr->getKind())
-            .drop_front();
         auto expandedType = solution.simplifyType(solution.getType(expr));
         cs.setType(expr, expandedType);
 
-        if (auto newExpr = expandMacroExpr(dc, expr, kind, expandedType)) {
+        auto locator = cs.getConstraintLocator(expr);
+        auto overload = solution.getOverloadChoice(locator);
+
+        auto macro = cast<MacroDecl>(overload.choice.getDecl());
+        // FIXME: use the macro for expansion!
+        // FIXME: form a proper ConcreteDeclRef for the macro expansion!
+
+        ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
+        if (auto newExpr = expandMacroExpr(dc, expr, macroRef, expandedType)) {
           auto expansion = new (ctx) MacroExpansionExpr(
-              expr->getStartLoc(), DeclNameRef(ctx.getIdentifier(kind)),
+              expr->getStartLoc(), DeclNameRef(macro->getName()),
               DeclNameLoc(expr->getLoc()), nullptr, /*isImplicit=*/true,
               expandedType);
           expansion->setRewritten(newExpr);
@@ -5392,11 +5398,16 @@ namespace {
 #if SWIFT_SWIFT_PARSER
       auto &ctx = cs.getASTContext();
       if (ctx.LangOpts.hasFeature(Feature::Macros)) {
-        auto macroIdent = E->getMacroName().getBaseIdentifier();
         auto expandedType = solution.simplifyType(solution.getType(E));
         cs.setType(E, expandedType);
-        if (auto newExpr = expandMacroExpr(
-                dc, E, macroIdent.str(), expandedType)) {
+
+        auto locator = cs.getConstraintLocator(E);
+        auto overload = solution.getOverloadChoice(locator);
+
+        auto macro = cast<MacroDecl>(overload.choice.getDecl());
+        ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
+
+        if (auto newExpr = expandMacroExpr(dc, E, macroRef, expandedType)) {
           E->setRewritten(newExpr);
           cs.cacheExprTypes(E);
           return E;

lib/Sema/CSGen.cpp

@@ -1200,17 +1200,23 @@ namespace {
         if (!protocol)
           return Type();
 
-        auto openedType = CS.getTypeOfMacroReference(
-            ctx.getIdentifier(kind), expr);
-        if (!openedType)
-          return Type();
-
-        CS.addConstraint(ConstraintKind::LiteralConformsTo, openedType,
-                         protocol->getDeclaredInterfaceType(),
-                         CS.getConstraintLocator(expr));
-
-        return openedType;
+        auto macroIdent = ctx.getIdentifier(kind);
+        auto macros = lookupMacros(macroIdent, FunctionRefKind::Unapplied);
+        if (!macros.empty()) {
+          // Introduce an overload set for the macro reference.
+          auto locator = CS.getConstraintLocator(expr);
+          auto macroRefType = Type(CS.createTypeVariable(locator, 0));
+          CS.addOverloadSet(macroRefType, macros, CurDC, locator);
+
+          // FIXME: Can this be encoded in the macro definition somehow?
+          CS.addConstraint(ConstraintKind::LiteralConformsTo, macroRefType,
+                           protocol->getDeclaredInterfaceType(),
+                           CS.getConstraintLocator(expr));
+
+          return macroRefType;
+        }
       }
+
       // Fall through to use old implementation.
 #endif
 
@@ -3629,41 +3635,72 @@ namespace {
       return resultTy;
     }
 
+    /// Lookup all macros with the given macro name.
+    SmallVector<OverloadChoice, 1>
+    lookupMacros(Identifier macroName, FunctionRefKind functionRefKind) {
+      auto req = MacroLookupRequest{macroName, CS.DC->getParentModule()};
+      auto macros = evaluateOrDefault(
+          CS.getASTContext().evaluator, req, { });
+      if (macros.empty())
+        return { };
+
+      // FIXME: At some point, we need to check for function-like macros without
+      // arguments and vice-versa.
+
+
+      SmallVector<OverloadChoice, 1> choices;
+      for (auto macro : macros) {
+        OverloadChoice choice = OverloadChoice(Type(), macro, functionRefKind);
+        choices.push_back(choice);
+      }
+      return choices;
+    }
+
     Type visitMacroExpansionExpr(MacroExpansionExpr *expr) {
 #if SWIFT_SWIFT_PARSER
       auto &ctx = CS.getASTContext();
       if (ctx.LangOpts.hasFeature(Feature::Macros)) {
+        // Look up the macros with this name.
         auto macroIdent = expr->getMacroName().getBaseIdentifier();
-        auto refType = CS.getTypeOfMacroReference(macroIdent, expr);
-        if (!refType) {
+        bool isCall = expr->getArgs() != nullptr;
+        FunctionRefKind functionRefKind = isCall ? FunctionRefKind::SingleApply
+                                                 : FunctionRefKind::Unapplied;
+        auto macros = lookupMacros(macroIdent, functionRefKind);
+        if (macros.empty()) {
           ctx.Diags.diagnose(expr->getMacroNameLoc(), diag::macro_undefined,
                              macroIdent)
               .highlight(expr->getMacroNameLoc().getSourceRange());
           return Type();
         }
-        if (expr->getArgs()) {
-          CS.associateArgumentList(CS.getConstraintLocator(expr), expr->getArgs());
-          // FIXME: Do we have object-like vs. function-like macros?
-          if (auto fnType = dyn_cast<FunctionType>(refType.getPointer())) {
-            SmallVector<AnyFunctionType::Param, 8> params;
-            getMatchingParams(expr->getArgs(), params);
-
-            Type resultType = CS.createTypeVariable(
-                CS.getConstraintLocator(expr, ConstraintLocator::FunctionResult),
-                TVO_CanBindToNoEscape);
-
-            CS.addConstraint(
-                ConstraintKind::ApplicableFunction,
-                FunctionType::get(params, resultType),
-                fnType,
-                CS.getConstraintLocator(
-                  expr, ConstraintLocator::ApplyFunction));
-
-            return resultType;
-          }
-        }
 
-        return refType;
+        // Introduce an overload set for the macro reference.
+        auto locator = CS.getConstraintLocator(expr);
+        auto macroRefType = Type(CS.createTypeVariable(locator, 0));
+        CS.addOverloadSet(macroRefType, macros, CurDC, locator);
+
+        // For non-calls, the type variable is the result.
+        if (!isCall)
+          return macroRefType;
+
+        // For calls, set up the argument list and form the applicable-function
+        // constraint. The result type is the result of that call.
+        CS.associateArgumentList(locator, expr->getArgs());
+
+        SmallVector<AnyFunctionType::Param, 8> params;
+        getMatchingParams(expr->getArgs(), params);
+
+        Type resultType = CS.createTypeVariable(
+            CS.getConstraintLocator(expr, ConstraintLocator::FunctionResult),
+            TVO_CanBindToNoEscape);
+
+        CS.addConstraint(
+            ConstraintKind::ApplicableFunction,
+            FunctionType::get(params, resultType),
+            macroRefType,
+            CS.getConstraintLocator(
+              expr, ConstraintLocator::ApplyFunction));
+
+        return resultType;
       }
 #endif
       return Type();

lib/Sema/ConstraintSystem.cpp

@@ -1606,6 +1606,29 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
     return { type, type, type, type };
   }
 
+  // Unqualified reference to a macro.
+  if (auto macro = dyn_cast<MacroDecl>(value)) {
+    Type macroType = macro->getInterfaceType();
+
+    // Remove argument labels when we have a function-like macro.
+    if (isa<AnyFunctionType>(macroType.getPointer()))
+      macroType = macroType->removeArgumentLabels(1);
+
+    // Open any the generic types.
+    OpenedTypeMap replacements;
+    openGeneric(macro->getParentModule(), macro->getGenericSignature(),
+                locator, replacements);
+
+    // If we opened up any type variables, record the replacements.
+    recordOpenedTypes(locator, replacements);
+
+    Type openedType = openType(macroType, replacements);
+
+    // FIXME: Should we use replaceParamErrorTypeByPlaceholder() here?
+
+    return { openedType, openedType, openedType, openedType };
+  }
+
   // Only remaining case: unqualified reference to a property.
   auto *varDecl = cast<VarDecl>(value);
 
@@ -2471,26 +2494,6 @@ ConstraintSystem::getTypeOfMemberReference(
   return { origOpenedType, openedType, origType, type };
 }
 
-#if SWIFT_SWIFT_PARSER
-Type ConstraintSystem::getTypeOfMacroReference(Identifier macroName,
-                                               Expr *anchor) {
-  auto req = MacroLookupRequest{macroName, DC->getParentModule()};
-  auto macros = evaluateOrDefault(getASTContext().evaluator, req, { });
-  if (macros.empty())
-    return Type();
-
-  // FIXME: Handle macro overloading.
-  auto macro = macros.front();
-
-  // Open any the generic types.
-  OpenedTypeMap replacements;
-  openGeneric(macro->getParentModule(), macro->getGenericSignature(),
-              getConstraintLocator(anchor), replacements);
-
-  return openType(macro->getInterfaceType(), replacements);
-}
-#endif
-
 Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
                                                 const OverloadChoice &overload,
                                                 bool allowMembers,

lib/Sema/TypeCheckMacros.cpp

@@ -331,7 +331,7 @@ ArrayRef<MacroDecl *> MacroLookupRequest::evaluate(
 
 #if SWIFT_SWIFT_PARSER
 Expr *swift::expandMacroExpr(
-    DeclContext *dc, Expr *expr, StringRef macroName, Type expandedType
+    DeclContext *dc, Expr *expr, ConcreteDeclRef macroRef, Type expandedType
 ) {
   ASTContext &ctx = dc->getASTContext();
   SourceManager &sourceMgr = ctx.SourceMgr;
@@ -348,53 +348,53 @@ Expr *swift::expandMacroExpr(
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
 
-  // FIXME: The caller should tell us what macro is being expanded, so we
-  // don't do this lookup again.
-
-  // Built-in macros go through `MacroSystem` in Swift Syntax linked to this
-  // compiler.
-  if (auto *macro = swift_ASTGen_lookupMacro(macroName.str().c_str())) {
-    auto astGenSourceFile = sourceFile->exportedSourceFile;
-    if (!astGenSourceFile)
-      return nullptr;
-  
-    const char *evaluatedSourceAddress;
-    ptrdiff_t evaluatedSourceLength;
-    swift_ASTGen_evaluateMacro(
-        astGenSourceFile, expr->getStartLoc().getOpaquePointerValue(),
-        &evaluatedSourceAddress, &evaluatedSourceLength);
-    if (!evaluatedSourceAddress)
-      return nullptr;
-    evaluatedSource = NullTerminatedStringRef(evaluatedSourceAddress,
-                                              (size_t)evaluatedSourceLength);
-    swift_ASTGen_destroyMacro(macro);
-  }
-  // Other macros go through a compiler plugin.
-  else {
-    auto mee = cast<MacroExpansionExpr>(expr);
-    auto *plugin = ctx.getLoadedPlugin(macroName);
-    assert(plugin && "Should have been checked during earlier type checking");
-    auto bufferID = sourceFile->getBufferID();
-    auto sourceFileText = sourceMgr.getEntireTextForBuffer(*bufferID);
-    auto evaluated = plugin->invokeRewrite(
-        /*targetModuleName*/ dc->getParentModule()->getName().str(),
-        /*filePath*/ sourceFile->getFilename(),
-        /*sourceFileText*/ sourceFileText,
-        /*range*/ Lexer::getCharSourceRangeFromSourceRange(
-            sourceMgr, mee->getSourceRange()),
-        ctx);
-    if (evaluated)
-      evaluatedSource = *evaluated;
-    else
-      return nullptr;
+  MacroDecl *macro = cast<MacroDecl>(macroRef.getDecl());
+
+  switch (macro->implementationKind) {
+    case MacroDecl::ImplementationKind::Builtin: {
+      // Builtin macros are handled via ASTGen.
+      auto astGenSourceFile = sourceFile->exportedSourceFile;
+      if (!astGenSourceFile)
+        return nullptr;
+
+      // FIXME: Tell ASTGen which macro to use.
+      const char *evaluatedSourceAddress;
+      ptrdiff_t evaluatedSourceLength;
+      swift_ASTGen_evaluateMacro(
+           astGenSourceFile, expr->getStartLoc().getOpaquePointerValue(),
+           &evaluatedSourceAddress, &evaluatedSourceLength);
+      if (!evaluatedSourceAddress)
+        return nullptr;
+      evaluatedSource = NullTerminatedStringRef(evaluatedSourceAddress,
+                                                (size_t)evaluatedSourceLength);
+      break;
+    }
+
+    case MacroDecl::ImplementationKind::Plugin: {
+      auto *plugin = (CompilerPlugin *)macro->opaqueHandle;
+      auto bufferID = sourceFile->getBufferID();
+      auto sourceFileText = sourceMgr.getEntireTextForBuffer(*bufferID);
+      auto evaluated = plugin->invokeRewrite(
+          /*targetModuleName*/ dc->getParentModule()->getName().str(),
+          /*filePath*/ sourceFile->getFilename(),
+          /*sourceFileText*/ sourceFileText,
+          /*range*/ Lexer::getCharSourceRangeFromSourceRange(
+              sourceMgr, expr->getSourceRange()),
+          ctx);
+      if (evaluated)
+        evaluatedSource = *evaluated;
+      else
+        return nullptr;
+      break;
+    }
   }
 
   // Figure out a reasonable name for the macro expansion buffer.
   std::string bufferName;
   {
     llvm::raw_string_ostream out(bufferName);
 
-    out << "Macro expansion of #" << macroName;
+    out << "Macro expansion of #" << macro->getName();
     if (auto bufferID = sourceFile->getBufferID()) {
       unsigned startLine, startColumn;
       std::tie(startLine, startColumn) =

lib/Sema/TypeCheckMacros.h

@@ -16,8 +16,8 @@
 #ifndef SWIFT_SEMA_TYPECHECKMACROS_H
 #define SWIFT_SEMA_TYPECHECKMACROS_H
 
+#include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/Type.h"
-#include "llvm/ADT/StringRef.h"
 
 namespace swift {
 
@@ -32,7 +32,7 @@ class TypeRepr;
 /// \returns the type-checked replacement expression, or NULL if the
 // macro could not be expanded.
 Expr *expandMacroExpr(
-    DeclContext *dc, Expr *expr, StringRef macroName, Type expandedType);
+    DeclContext *dc, Expr *expr, ConcreteDeclRef macroRef, Type expandedType);
 
 #endif // SWIFT_SWIFT_PARSER