[Macros] Add support for explicit generic arguments of macros.

Enable type checking support for explicitly specifying generic arguments to
a macro, e.g., `#stringify<Double>(1 + 2)`. To do so, introduce a new
kind of constraint that performs explicit argument matching against the
generic parameters of a macro only after the overload is chosen.

Author: DougGregor

include/swift/Sema/Constraint.h

@@ -226,6 +226,13 @@ enum class ConstraintKind : char {
   BindTupleOfFunctionParams,
   /// The first type is a type pack, and the second type is its reduced shape.
   ShapeOf,
+  /// Represents explicit generic arguments provided for a reference to
+  /// a declaration.
+  ///
+  /// The first type is the type variable describing the bound type of
+  /// an overload. The second type is a PackType containing the explicit
+  /// generic arguments.
+  ExplicitGenericArguments,
 };
 
 /// Classification of the different kinds of constraints.
@@ -708,6 +715,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::Defaultable:
     case ConstraintKind::BindTupleOfFunctionParams:
     case ConstraintKind::ShapeOf:
+    case ConstraintKind::ExplicitGenericArguments:
       return ConstraintClassification::TypeProperty;
 
     case ConstraintKind::Disjunction:

include/swift/Sema/ConstraintSystem.h

@@ -5611,6 +5611,13 @@ class ConstraintSystem {
       Type type1, Type type2, TypeMatchOptions flags,
       ConstraintLocatorBuilder locator);
 
+  /// Simplify an explicit generic argument constraint by equating the
+  /// opened generic types of the bound left-hand type variable to the
+  /// pack type on the right-hand side.
+  SolutionKind simplifyExplicitGenericArgumentsConstraint(
+      Type type1, Type type2, TypeMatchOptions flags,
+      ConstraintLocatorBuilder locator);
+
 public: // FIXME: Public for use by static functions.
   /// Simplify a conversion constraint with a fix applied to it.
   SolutionKind simplifyFixConstraint(ConstraintFix *fix, Type type1, Type type2,

lib/Sema/CSBindings.cpp

@@ -1406,6 +1406,7 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     // Constraints from which we can't do anything.
     break;
 

lib/Sema/CSGen.cpp

@@ -1642,6 +1642,30 @@ namespace {
                                      expr->getOuterAlternatives());
     }
 
+    /// Given a set of specialization arguments, resolve those arguments and
+    /// introduce them as an explicit generic arguments constraint.
+    void addSpecializationConstraint(
+        ConstraintLocator *locator, Type boundType,
+        ArrayRef<TypeRepr *> specializationArgs) {
+      // Resolve each type.
+      SmallVector<Type, 2> specializationArgTypes;
+      const auto options =
+          TypeResolutionOptions(TypeResolverContext::InExpression);
+      for (auto specializationArg : specializationArgs) {
+        const auto result = TypeResolution::resolveContextualType(
+            specializationArg, CurDC, options,
+            // Introduce type variables for unbound generics.
+            OpenUnboundGenericType(CS, locator),
+            HandlePlaceholderType(CS, locator));
+        specializationArgTypes.push_back(result);
+      }
+
+      CS.addConstraint(
+          ConstraintKind::ExplicitGenericArguments, boundType,
+          PackType::get(CS.getASTContext(), specializationArgTypes),
+          locator);
+    }
+
     Type visitUnresolvedSpecializeExpr(UnresolvedSpecializeExpr *expr) {
       auto baseTy = CS.getType(expr->getSubExpr());
 
@@ -3678,6 +3702,13 @@ namespace {
         auto macroRefType = Type(CS.createTypeVariable(locator, 0));
         CS.addOverloadSet(macroRefType, macros, CurDC, locator);
 
+        // Add explicit generic arguments, if there were any.
+        if (expr->getGenericArgsRange().isValid()) {
+          addSpecializationConstraint(
+              CS.getConstraintLocator(expr), macroRefType,
+              expr->getGenericArgs());
+        }
+
         // For non-calls, the type variable is the result.
         if (!isCall)
           return macroRefType;

lib/Sema/CSSimplify.cpp

@@ -2306,6 +2306,7 @@ ConstraintSystem::matchTupleTypes(TupleType *tuple1, TupleType *tuple2,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     llvm_unreachable("Bad constraint kind in matchTupleTypes()");
   }
 
@@ -2666,6 +2667,7 @@ static bool matchFunctionRepresentations(FunctionType::ExtInfo einfo1,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     return true;
   }
 
@@ -3084,6 +3086,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     llvm_unreachable("Not a relational constraint");
   }
 
@@ -6459,6 +6462,7 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
     case ConstraintKind::BindTupleOfFunctionParams:
     case ConstraintKind::PackElementOf:
     case ConstraintKind::ShapeOf:
+    case ConstraintKind::ExplicitGenericArguments:
       llvm_unreachable("Not a relational constraint");
     }
   }
@@ -12644,6 +12648,95 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyShapeOfConstraint(
   return SolutionKind::Solved;
 }
 
+ConstraintSystem::SolutionKind
+ConstraintSystem::simplifyExplicitGenericArgumentsConstraint(
+    Type type1, Type type2, TypeMatchOptions flags,
+    ConstraintLocatorBuilder locator) {
+  auto formUnsolved = [&]() {
+    // If we're supposed to generate constraints, do so.
+    if (flags.contains(TMF_GenerateConstraints)) {
+      auto *shapeOf = Constraint::create(
+          *this, ConstraintKind::ShapeOf, type1, type2,
+          getConstraintLocator(locator));
+
+      addUnsolvedConstraint(shapeOf);
+      return SolutionKind::Solved;
+    }
+
+    return SolutionKind::Unsolved;
+  };
+
+  // Bail out if we haven't selected an overload yet.
+  auto simplifiedBoundType = simplifyType(type1, flags);
+  if (simplifiedBoundType->isTypeVariableOrMember())
+    return formUnsolved();
+
+  // Determine the overload locator for this constraint.
+  ConstraintLocator *overloadLocator = nullptr;
+  if (auto anchorExpr = locator.getAnchor().dyn_cast<Expr *>()) {
+    if (auto expansion = dyn_cast<MacroExpansionExpr>(anchorExpr)) {
+      overloadLocator = getConstraintLocator(expansion);
+    } else if (auto specialize =
+                   dyn_cast<UnresolvedSpecializeExpr>(anchorExpr)) {
+      overloadLocator = getConstraintLocator(
+          specialize->getSubExpr()->getSemanticsProvidingExpr());
+    }
+  } else if (auto anchorDecl = locator.getAnchor().dyn_cast<Decl *>()) {
+    if (auto expansion = dyn_cast<MacroExpansionDecl>(anchorDecl)) {
+      overloadLocator = getConstraintLocator(expansion);
+    }
+  }
+  assert(overloadLocator && "Specialize expression has the wrong form");
+
+  // If the overload hasn't been resolved, we can't simplify this constraint.
+  auto resolvedOverloadIter = getResolvedOverloads().find(overloadLocator);
+  if (resolvedOverloadIter == getResolvedOverloads().end())
+    return formUnsolved();
+
+  auto selectedOverload = resolvedOverloadIter->second;
+  auto overloadChoice = selectedOverload.choice;
+  if (!overloadChoice.isDecl()) {
+    return SolutionKind::Error;
+  }
+
+  auto decl = overloadChoice.getDecl();
+  auto genericContext = decl->getAsGenericContext();
+  if (!genericContext)
+    return SolutionKind::Error;
+
+  auto genericParams = genericContext->getGenericParams();
+  if (!genericParams || genericParams->size() == 0) {
+    // FIXME: Record an error here that we're ignoring the parameters.
+    return SolutionKind::Solved;
+  }
+
+  // Map the generic parameters we have over to their opened types.
+  SmallVector<Type, 2> openedGenericParams;
+  auto genericParamDepth = genericParams->getParams()[0]->getDepth();
+  for (const auto &openedType : getOpenedTypes(overloadLocator)) {
+    if (openedType.first->getDepth() == genericParamDepth) {
+      openedGenericParams.push_back(Type(openedType.second));
+    }
+  }
+  assert(openedGenericParams.size() == genericParams->size());
+
+  // Match the opened generic parameters to the specialized arguments.
+  auto specializedArgs = type2->castTo<PackType>()->getElementTypes();
+  PackMatcher matcher(openedGenericParams, specializedArgs, getASTContext());
+  if (matcher.match())
+    return SolutionKind::Error;
+
+  // Bind the opened generic parameters to the specialization arguments.
+  for (const auto &pair : matcher.pairs) {
+    addConstraint(
+        ConstraintKind::Bind, pair.lhs, pair.rhs,
+        getConstraintLocator(
+            locator, LocatorPathElt::GenericArgument(pair.idx)));
+  }
+
+  return SolutionKind::Solved;
+}
+
 static llvm::PointerIntPair<Type, 3, unsigned>
 getBaseTypeForPointer(TypeBase *type) {
   unsigned unwrapCount = 0;
@@ -13993,6 +14086,10 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
   case ConstraintKind::ShapeOf:
     return simplifyShapeOfConstraint(first, second, subflags, locator);
 
+  case ConstraintKind::ExplicitGenericArguments:
+    return simplifyExplicitGenericArgumentsConstraint(
+        first, second, subflags, locator);
+
   case ConstraintKind::ValueMember:
   case ConstraintKind::UnresolvedValueMember:
   case ConstraintKind::ValueWitness:
@@ -14583,6 +14680,11 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
     return simplifyShapeOfConstraint(
         constraint.getFirstType(), constraint.getSecondType(), /*flags*/ None,
         constraint.getLocator());
+
+  case ConstraintKind::ExplicitGenericArguments:
+    return simplifyExplicitGenericArgumentsConstraint(
+        constraint.getFirstType(), constraint.getSecondType(),
+        /*flags*/ None, constraint.getLocator());
   }
 
   llvm_unreachable("Unhandled ConstraintKind in switch.");

lib/Sema/Constraint.cpp

@@ -81,6 +81,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     assert(!First.isNull());
     assert(!Second.isNull());
     break;
@@ -169,6 +170,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second, Type Third,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     llvm_unreachable("Wrong constructor");
 
   case ConstraintKind::KeyPath:
@@ -316,6 +318,7 @@ Constraint *Constraint::clone(ConstraintSystem &cs) const {
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     return create(cs, getKind(), getFirstType(), getSecondType(), getLocator());
 
   case ConstraintKind::ApplicableFunction:
@@ -560,6 +563,10 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm, unsigned inden
     Out << " shape of ";
     break;
 
+  case ConstraintKind::ExplicitGenericArguments:
+    Out << " explicit generic argument binding ";
+    break;
+
   case ConstraintKind::Disjunction:
     llvm_unreachable("disjunction handled above");
   case ConstraintKind::Conjunction:
@@ -726,6 +733,7 @@ gatherReferencedTypeVars(Constraint *constraint,
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
+  case ConstraintKind::ExplicitGenericArguments:
     constraint->getFirstType()->getTypeVariables(typeVars);
     constraint->getSecondType()->getTypeVariables(typeVars);
     break;

lib/Sema/TypeCheckMacros.cpp

@@ -19,6 +19,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/CompilerPlugin.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/TypeCheckRequests.h"
@@ -100,7 +101,7 @@ getMacroSignatureContextBuffer(
 
 /// Compute the macro signature for a macro given the source code for its
 /// generic signature and type signature.
-static Optional<std::pair<GenericSignature, Type>>
+static Optional<std::tuple<GenericParamList *, GenericSignature, Type>>
 getMacroSignature(
     ModuleDecl *mod, Identifier macroName,
     Optional<StringRef> genericSignature,
@@ -161,8 +162,9 @@ getMacroSignature(
   performImportResolution(*macroSourceFile);
 
   auto typealias = cast<TypeAliasDecl>(decl);
-  return std::make_pair(
-      typealias->getGenericSignature(), typealias->getUnderlyingType());
+  return std::make_tuple(
+      typealias->getGenericParams(), typealias->getGenericSignature(),
+      typealias->getUnderlyingType());
 }
 
 /// Create a macro.
@@ -212,8 +214,9 @@ static MacroDecl *createMacro(
       opaqueHandle);
 
   // FIXME: Make these lazily computed.
-  macro->setGenericSignature(signature->first);
-  macro->setInterfaceType(signature->second);
+  GenericParamListRequest{macro}.cacheResult(std::get<0>(*signature));
+  macro->setGenericSignature(std::get<1>(*signature));
+  macro->setInterfaceType(std::get<2>(*signature));
 
   return macro;
 }

test/Macros/macros.swift

@@ -14,7 +14,7 @@ func test(a: Int, b: Int) {
   // CHECK: tuple_expr type='(() -> Bool, String)' location=Macro expansion of #stringify
 
   let (b2, s3) = #stringify<Double>(1 + 2)
-  // CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
+  // CHECK: macro_expansion_expr type='(Double, String)'{{.*}}name=stringify
   // CHECK-NEXT: argument_list
-  // CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
+  // CHECK: tuple_expr type='(Double, String)' location=Macro expansion of #stringify
 }