Merge pull request #63247 from rxwei/resolve-macro-request

Author: rxwei

include/swift/AST/MacroDeclaration.h

@@ -54,6 +54,10 @@ enum class MacroRole: uint32_t {
 /// The contexts in which a particular macro declaration can be used.
 using MacroRoles = OptionSet<MacroRole>;
 
+void simple_display(llvm::raw_ostream &out, MacroRoles roles);
+bool operator==(MacroRoles lhs, MacroRoles rhs);
+llvm::hash_code hash_value(MacroRoles roles);
+
 /// Retrieve the string form of the given macro role, as written on the
 /// corresponding attribute.
 StringRef getMacroRoleString(MacroRole role);
@@ -62,10 +66,14 @@ StringRef getMacroRoleString(MacroRole role);
 /// written in the source code with the `#` syntax.
 bool isFreestandingMacro(MacroRoles contexts);
 
+MacroRoles getFreestandingMacroRoles();
+
 /// Whether a macro with the given set of macro contexts is attached, i.e.,
 /// written in the source code as an attribute with the `@` syntax.
 bool isAttachedMacro(MacroRoles contexts);
 
+MacroRoles getAttachedMacroRoles();
+
 enum class MacroIntroducedDeclNameKind {
   Named,
   Overloaded,

include/swift/AST/TypeCheckRequests.h

@@ -3128,10 +3128,53 @@ class CheckRedeclarationRequest
                            evaluator::SideEffect) const;
 };
 
+class UnresolvedMacroReference {
+private:
+  llvm::PointerUnion<MacroExpansionDecl *, MacroExpansionExpr *, CustomAttr *>
+    pointer;
+
+public:
+  UnresolvedMacroReference(MacroExpansionDecl *decl) : pointer(decl) {}
+  UnresolvedMacroReference(MacroExpansionExpr *expr) : pointer(expr) {}
+  UnresolvedMacroReference(CustomAttr *attr) : pointer(attr) {}
+
+  MacroExpansionDecl *getDecl() const {
+    return pointer.dyn_cast<MacroExpansionDecl *>();
+  }
+  MacroExpansionExpr *getExpr() const {
+    return pointer.dyn_cast<MacroExpansionExpr *>();
+  }
+  CustomAttr *getAttr() const {
+    return pointer.dyn_cast<CustomAttr *>();
+  }
+  void *getOpaqueValue() const {
+    return pointer.getOpaqueValue();
+  }
+
+  DeclNameRef getMacroName() const;
+  DeclNameLoc getMacroNameLoc() const;
+  SourceRange getGenericArgsRange() const;
+  ArrayRef<TypeRepr *> getGenericArgs() const;
+  ArgumentList *getArgs() const;
+
+  friend bool operator==(const UnresolvedMacroReference &lhs,
+                         const UnresolvedMacroReference &rhs) {
+    return lhs.getOpaqueValue() == rhs.getOpaqueValue();
+  }
+
+  friend llvm::hash_code hash_value(const UnresolvedMacroReference &ref) {
+    return reinterpret_cast<ptrdiff_t>(ref.pointer.getOpaqueValue());
+  }
+};
+
+void simple_display(llvm::raw_ostream &out,
+                    const UnresolvedMacroReference &ref);
+
 /// Resolve a given custom attribute to an attached macro declaration.
-class ResolveAttachedMacroRequest
-    : public SimpleRequest<ResolveAttachedMacroRequest,
-                           MacroDecl *(CustomAttr *, DeclContext *),
+class ResolveMacroRequest
+    : public SimpleRequest<ResolveMacroRequest,
+                           MacroDecl *(UnresolvedMacroReference, MacroRoles,
+                                       DeclContext *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -3140,7 +3183,8 @@ class ResolveAttachedMacroRequest
   friend SimpleRequest;
 
   MacroDecl *
-  evaluate(Evaluator &evaluator, CustomAttr *attr, DeclContext *dc) const;
+  evaluate(Evaluator &evaluator, UnresolvedMacroReference macroRef,
+           MacroRoles roles, DeclContext *dc) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -340,7 +340,7 @@ SWIFT_REQUEST(TypeChecker, PreCheckResultBuilderRequest,
 SWIFT_REQUEST(TypeChecker, ResolveImplicitMemberRequest,
               evaluator::SideEffect(NominalTypeDecl *, ImplicitMemberAction),
               Uncached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, ResolveAttachedMacroRequest,
+SWIFT_REQUEST(TypeChecker, ResolveMacroRequest,
               MacroDecl *(CustomAttr *, DeclContext *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveTypeEraserTypeRequest,

lib/AST/Attr.cpp

@@ -2356,7 +2356,8 @@ bool CustomAttr::isAttachedMacro(const Decl *decl) const {
 
   auto *macroDecl = evaluateOrDefault(
       ctx.evaluator,
-      ResolveAttachedMacroRequest{const_cast<CustomAttr *>(this), dc},
+      ResolveMacroRequest{const_cast<CustomAttr *>(this),
+                          getAttachedMacroRoles(), dc},
       nullptr);
 
   return macroDecl != nullptr;

lib/AST/Decl.cpp

@@ -384,7 +384,7 @@ void Decl::forEachAttachedMacro(MacroRole role,
     auto customAttr = const_cast<CustomAttr *>(customAttrConst);
     auto *macroDecl = evaluateOrDefault(
         ctx.evaluator,
-        ResolveAttachedMacroRequest{customAttr, dc},
+        ResolveMacroRequest{customAttr, getAttachedMacroRoles(), dc},
         nullptr);
 
     if (!macroDecl)
@@ -9777,10 +9777,18 @@ bool swift::isFreestandingMacro(MacroRoles contexts) {
   return bool(contexts & freestandingMacroRoles);
 }
 
+MacroRoles swift::getFreestandingMacroRoles() {
+  return freestandingMacroRoles;
+}
+
 bool swift::isAttachedMacro(MacroRoles contexts) {
   return bool(contexts & attachedMacroRoles);
 }
 
+MacroRoles swift::getAttachedMacroRoles() {
+  return attachedMacroRoles;
+}
+
 MacroDecl::MacroDecl(
     SourceLoc macroLoc, DeclName name, SourceLoc nameLoc,
     GenericParamList *genericParams,

lib/Sema/TypeCheckAttr.cpp

@@ -3570,7 +3570,8 @@ void AttributeChecker::visitCustomAttr(CustomAttr *attr) {
   if (!nominal) {
     // Try resolving an attached macro attribute.
     auto *macro = evaluateOrDefault(
-        Ctx.evaluator, ResolveAttachedMacroRequest{attr, dc}, nullptr);
+        Ctx.evaluator, ResolveMacroRequest{attr, getAttachedMacroRoles(), dc},
+        nullptr);
     if (macro || !attr->isValid())
       return;
 

lib/Sema/TypeCheckDecl.cpp

@@ -1602,14 +1602,16 @@ TypeChecker::lookupPrecedenceGroup(DeclContext *dc, Identifier name,
 
 SmallVector<MacroDecl *, 1>
 TypeChecker::lookupMacros(DeclContext *dc, DeclNameRef macroName,
-                          SourceLoc loc, MacroRoles contexts) {
-  auto result = lookupUnqualified(dc, DeclNameRef(macroName), loc,
-                                  (defaultUnqualifiedLookupOptions |
-                                      NameLookupFlags::IncludeOuterResults));
+                          SourceLoc loc, MacroRoles roles) {
   SmallVector<MacroDecl *, 1> choices;
-  for (const auto &found : result.allResults())
+  auto moduleScopeDC = dc->getModuleScopeContext();
+  ASTContext &ctx = moduleScopeDC->getASTContext();
+  UnqualifiedLookupDescriptor descriptor(macroName, moduleScopeDC);
+  auto lookup = evaluateOrDefault(
+      ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
+  for (const auto &found : lookup.allResults())
     if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl()))
-      if (contexts.contains(macro->getMacroRoles()))
+      if (roles.contains(macro->getMacroRoles()))
         choices.push_back(macro);
   return choices;
 }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -3685,9 +3685,12 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
   }
   // Resolve macro candidates.
   MacroDecl *macro;
-  // Non-call declaration macros cannot be overloaded.
-  auto *args = MED->getArgs();
-  if (!args) {
+  if (auto *args = MED->getArgs()) {
+    macro = evaluateOrDefault(
+        ctx.evaluator, ResolveMacroRequest{MED, MacroRole::Declaration, dc},
+        nullptr);
+  }
+  else {
     if (foundMacros.size() > 1) {
       MED->diagnose(diag::ambiguous_decl_ref, MED->getMacro())
           .highlight(MED->getMacroLoc().getSourceRange());
@@ -3697,63 +3700,8 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
     }
     macro = foundMacros.front();
   }
-    // Call-like macros need to be resolved.
-  else {
-    using namespace constraints;
-    ConstraintSystem cs(dc, ConstraintSystemOptions());
-    // Type-check macro arguments.
-    for (auto *arg : args->getArgExprs())
-      cs.setType(arg, TypeChecker::typeCheckExpression(arg, dc));
-    auto choices = map<SmallVector<OverloadChoice, 1>>(
-        foundMacros,
-        [](MacroDecl *md) {
-          return OverloadChoice(Type(), md, FunctionRefKind::SingleApply);
-        });
-    auto locator = cs.getConstraintLocator(MED);
-    auto macroRefType = Type(cs.createTypeVariable(locator, 0));
-    cs.addOverloadSet(macroRefType, choices, dc, locator);
-    if (MED->getGenericArgsRange().isValid()) {
-      // FIXME: Deal with generic args.
-      MED->diagnose(diag::macro_unsupported);
-      return {};
-    }
-    auto getMatchingParams = [&](
-        ArgumentList *argList,
-        SmallVectorImpl<AnyFunctionType::Param> &result) {
-      for (auto arg : *argList) {
-        ParameterTypeFlags flags;
-        auto ty = cs.getType(arg.getExpr()); if (arg.isInOut()) {
-          ty = ty->getInOutObjectType();
-          flags = flags.withInOut(true);
-        }
-        if (arg.isConst()) {
-          flags = flags.withCompileTimeConst(true);
-        }
-        result.emplace_back(ty, arg.getLabel(), flags);
-      }
-    };
-    SmallVector<AnyFunctionType::Param, 8> params;
-    getMatchingParams(args, params);
-    cs.associateArgumentList(locator, args);
-    cs.addConstraint(
-        ConstraintKind::ApplicableFunction,
-        FunctionType::get(params, ctx.getVoidType()),
-        macroRefType,
-        cs.getConstraintLocator(MED, ConstraintLocator::ApplyFunction));
-    // Solve.
-    auto solution = cs.solveSingle();
-    if (!solution) {
-      MED->diagnose(diag::no_overloads_match_exactly_in_call,
-                    /*reference|call*/false, DescriptiveDeclKind::Macro,
-                    false, MED->getMacro().getBaseName())
-          .highlight(MED->getMacroLoc().getSourceRange());
-      for (auto *candidate : foundMacros)
-        candidate->diagnose(diag::found_candidate);
-      return {};
-    }
-    auto choice = solution->getOverloadChoice(locator).choice;
-    macro = cast<MacroDecl>(choice.getDecl());
-  }
+  if (!macro)
+    return {};
   MED->setMacroRef(macro);
 
   // Expand the macro.