Opaque opaque types and compute substitutions in the constraint system

Opaque opaque types and record them within the "opened types" of the
constraint system, then use that information to compute the set of
substitutions needed for the opaque type declaration using the normal
mechanism of the constraint solver. Record these substitutions within
the underlying-to-opaque conversion.

Use the recorded substitutions in the underlying-to-opaque conversion
to set the underlying substitutions for the opaque type declaration
itself, rather than reconstructing the substitutions in an ad hoc manner
that does not account for structural opaque result types.

Author: DougGregor

include/swift/AST/Expr.h

@@ -3003,8 +3003,13 @@ class LinearToDifferentiableFunctionExpr : public ImplicitConversionExpr {
 /// (opaque type)" and "S<T> ---> S<(opaque type)>".
 class UnderlyingToOpaqueExpr : public ImplicitConversionExpr {
 public:
-  UnderlyingToOpaqueExpr(Expr *subExpr, Type ty)
-    : ImplicitConversionExpr(ExprKind::UnderlyingToOpaque, subExpr, ty) {}
+  /// The substitutions to be applied to the opaque type declaration to
+  /// produce the resulting type.
+  const SubstitutionMap substitutions;
+
+  UnderlyingToOpaqueExpr(Expr *subExpr, Type ty, SubstitutionMap substitutions)
+    : ImplicitConversionExpr(ExprKind::UnderlyingToOpaque, subExpr, ty),
+      substitutions(substitutions) {}
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::UnderlyingToOpaque;

include/swift/Sema/ConstraintSystem.h

@@ -3316,6 +3316,11 @@ class ConstraintSystem {
                        ArrayRef<ConstraintLocator::PathElement> path,
                        unsigned summaryFlags);
 
+  /// Retrieve a locator for opening the opaque archetype for the given
+  /// opaque type.
+  ConstraintLocator *getOpenOpaqueLocator(
+      ConstraintLocatorBuilder locator, OpaqueTypeDecl *opaqueDecl);
+
   /// Retrive the constraint locator for the given anchor and
   /// path, uniqued and automatically infer the summary flags
   ConstraintLocator *

lib/Sema/CSApply.cpp

@@ -7095,8 +7095,29 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
   // below is just an approximate check since the above would be expensive to
   // verify and still relies on the type checker ensuing `fromType` is
   // compatible with any opaque archetypes.
-  if (toType->hasOpaqueArchetype())
-    return cs.cacheType(new (ctx) UnderlyingToOpaqueExpr(expr, toType));
+  if (toType->hasOpaqueArchetype()) {
+    // Find the opaque type declaration. We need its generic signature.
+    OpaqueTypeDecl *opaqueDecl = nullptr;
+    bool found = toType.findIf([&](Type type) {
+      if (auto opaqueType = type->getAs<OpaqueTypeArchetypeType>()) {
+        opaqueDecl = opaqueType->getDecl();
+        return true;
+      }
+
+      return false;
+    });
+    (void)found;
+    assert(found && "No opaque type archetype?");
+
+    // Compute the substitutions for the opaque type declaration.
+    auto opaqueLocator = solution.getConstraintSystem().getOpenOpaqueLocator(
+        locator, opaqueDecl);
+    SubstitutionMap substitutions = solution.computeSubstitutions(
+        opaqueDecl->getOpaqueInterfaceGenericSignature(), opaqueLocator);
+    assert(!substitutions.empty() && "Missing substitutions for opaque type");
+    return cs.cacheType(
+        new (ctx) UnderlyingToOpaqueExpr(expr, toType, substitutions));
+  }
 
   llvm_unreachable("Unhandled coercion");
 }

lib/Sema/CSGen.cpp

@@ -1766,16 +1766,12 @@ namespace {
       auto contextualPurpose = CS.getContextualTypePurpose(expr);
 
       auto joinElementTypes = [&](Optional<Type> elementType) {
-        auto openedElementType = elementType.map([&](Type type) {
-          return CS.openOpaqueType(type, contextualPurpose, locator);
-        });
-
         const auto elements = expr->getElements();
         unsigned index = 0;
 
         using Iterator = decltype(elements)::iterator;
         CS.addJoinConstraint<Iterator>(
-            locator, elements.begin(), elements.end(), openedElementType,
+            locator, elements.begin(), elements.end(), elementType,
             [&](const auto it) {
               auto *locator = CS.getConstraintLocator(
                   expr, LocatorPathElt::TupleElement(index++));
@@ -1788,6 +1784,8 @@ namespace {
         // Now that we know we're actually going to use the type, get the
         // version for use in a constraint.
         contextualType = CS.getContextualType(expr, /*forConstraint=*/true);
+        contextualType = CS.openOpaqueType(
+            contextualType, contextualPurpose, locator);
         Optional<Type> arrayElementType =
             ConstraintSystem::isArrayType(contextualType);
         CS.addConstraint(ConstraintKind::LiteralConformsTo, contextualType,
@@ -1901,8 +1899,6 @@ namespace {
         contextualType = CS.getContextualType(expr, /*forConstraint=*/true);
         auto openedType =
             CS.openOpaqueType(contextualType, contextualPurpose, locator);
-        openedType = CS.replaceInferableTypesWithTypeVars(
-            openedType, CS.getConstraintLocator(expr));
         auto dictionaryKeyValue =
             ConstraintSystem::isDictionaryType(openedType);
         Type contextualDictionaryKeyType;

lib/Sema/ConstraintSystem.cpp

@@ -542,6 +542,14 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(
   return getConstraintLocator(anchor);
 }
 
+ConstraintLocator *ConstraintSystem::getOpenOpaqueLocator(
+    ConstraintLocatorBuilder locator, OpaqueTypeDecl *opaqueDecl) {
+  // Use only the opaque type declaration.
+  return getConstraintLocator(
+      ASTNode(opaqueDecl),
+      { LocatorPathElt::OpenedOpaqueArchetype(opaqueDecl) }, 0);
+}
+
 /// Extend the given depth map by adding depths for all of the subexpressions
 /// of the given expression.
 static void extendDepthMap(
@@ -829,25 +837,34 @@ Type ConstraintSystem::openType(Type type, OpenedTypeMap &replacements) {
 
 Type ConstraintSystem::openOpaqueType(OpaqueTypeArchetypeType *opaque,
                                       ConstraintLocatorBuilder locator) {
-  auto opaqueLocator = locator.withPathElement(
-      LocatorPathElt::OpenedOpaqueArchetype(opaque->getDecl()));
+  auto opaqueLocatorKey = getOpenOpaqueLocator(locator, opaque->getDecl());
+
+  // If we have already opened this opaque type, look in the known set of
+  // replacements.
+  auto knownReplacements = OpenedTypes.find(
+      getConstraintLocator(opaqueLocatorKey));
+  if (knownReplacements != OpenedTypes.end()) {
+    auto param = opaque->getInterfaceType()->castTo<GenericTypeParamType>();
+    for (const auto &replacement : knownReplacements->second) {
+      if (replacement.first == param)
+        return replacement.second;
+    }
+
+    llvm_unreachable("Missing opaque type replacement");
+  }
 
   // Open the generic signature of the opaque decl, and bind the "outer" generic
   // params to our context. The remaining axes of freedom on the type variable
   // corresponding to the underlying type should be the constraints on the
   // underlying return type.
+  auto opaqueLocator = locator.withPathElement(
+      LocatorPathElt::OpenedOpaqueArchetype(opaque->getDecl()));
   OpenedTypeMap replacements;
   openGeneric(DC, opaque->getBoundSignature(), opaqueLocator, replacements);
 
-  auto underlyingTyVar = openType(opaque->getInterfaceType(), replacements);
-  assert(underlyingTyVar);
-
-  for (auto param : DC->getGenericSignatureOfContext().getGenericParams()) {
-    addConstraint(ConstraintKind::Bind, openType(param, replacements),
-                  DC->mapTypeIntoContext(param), opaqueLocator);
-  }
+  recordOpenedTypes(opaqueLocatorKey, replacements);
 
-  return underlyingTyVar;
+  return openType(opaque->getInterfaceType(), replacements);
 }
 
 Type ConstraintSystem::openOpaqueType(Type type, ContextualTypePurpose context,

lib/Sema/MiscDiagnostics.cpp

@@ -2603,7 +2603,7 @@ class VarDeclUsageChecker : public ASTWalker {
 /// An AST walker that determines the underlying type of an opaque return decl
 /// from its associated function body.
 class OpaqueUnderlyingTypeChecker : public ASTWalker {
-  using Candidate = std::pair<Expr *, Type>;
+  using Candidate = std::pair<Expr *, SubstitutionMap>;
 
   ASTContext &Ctx;
   AbstractFunctionDecl *Implementation;
@@ -2642,20 +2642,21 @@ class OpaqueUnderlyingTypeChecker : public ASTWalker {
     }
 
     // Check whether all of the underlying type candidates match up.
-    // TODO [OPAQUE SUPPORT]: multiple opaque types
-    Type underlyingType = Candidates.front().second;
+    // TODO [OPAQUE SUPPORT]: diagnose multiple opaque types
+    SubstitutionMap underlyingSubs = Candidates.front().second;
+    SubstitutionMap canonicalUnderlyingSubs = underlyingSubs.getCanonical();
     bool mismatch =
         std::any_of(Candidates.begin() + 1, Candidates.end(),
                     [&](Candidate &otherCandidate) {
-                      return !underlyingType->isEqual(otherCandidate.second);
+                      return canonicalUnderlyingSubs != otherCandidate.second.getCanonical();
                     });
     if (mismatch) {
       Implementation->diagnose(
           diag::opaque_type_mismatched_underlying_type_candidates);
       for (auto candidate : Candidates) {
         Ctx.Diags.diagnose(candidate.first->getLoc(),
                            diag::opaque_type_underlying_type_candidate_here,
-                           candidate.second);
+                           candidate.second.getReplacementTypes()[0]);
       }
       return;
     }
@@ -2664,39 +2665,30 @@ class OpaqueUnderlyingTypeChecker : public ASTWalker {
     // in terms of the opaque type itself.
     auto opaqueTypeInContext = Implementation->mapTypeIntoContext(
         OpaqueDecl->getDeclaredInterfaceType());
-    auto isSelfReferencing = underlyingType.findIf([&](Type t) -> bool {
-      return t->isEqual(opaqueTypeInContext);
-    });
-    
-    if (isSelfReferencing) {
-      Ctx.Diags.diagnose(Candidates.front().first->getLoc(),
-                         diag::opaque_type_self_referential_underlying_type,
-                         underlyingType);
-      return;
+    for (auto genericParam : OpaqueDecl->getOpaqueGenericParams()) {
+      auto underlyingType = Type(genericParam).subst(underlyingSubs);
+      auto isSelfReferencing = underlyingType.findIf([&](Type t) -> bool {
+        return t->isEqual(opaqueTypeInContext);
+      });
+
+      if (isSelfReferencing) {
+        Ctx.Diags.diagnose(Candidates.front().first->getLoc(),
+                           diag::opaque_type_self_referential_underlying_type,
+                           underlyingSubs.getReplacementTypes()[0]);
+        return;
+      }
     }
-    
-    // If we have one successful candidate, then save it as the underlying type
-    // of the opaque decl.
-    // Form a substitution map that defines it in terms of the other context
-    // generic parameters.
-    underlyingType = underlyingType->mapTypeOutOfContext();
-    auto underlyingSubs = SubstitutionMap::get(
-      OpaqueDecl->getOpaqueInterfaceGenericSignature(),
-      [&](SubstitutableType *t) -> Type {
-        if (t->isEqual(OpaqueDecl->getUnderlyingInterfaceType())) {
-          return underlyingType;
-        }
-        return Type(t);
-      },
-      LookUpConformanceInModule(OpaqueDecl->getModuleContext()));
-    
-    OpaqueDecl->setUnderlyingTypeSubstitutions(underlyingSubs);
+
+    // If we have one successful candidate, then save it as the underlying
+    // substitutions of the opaque decl.
+    OpaqueDecl->setUnderlyingTypeSubstitutions(
+        underlyingSubs.mapReplacementTypesOutOfContext());
   }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
     if (auto underlyingToOpaque = dyn_cast<UnderlyingToOpaqueExpr>(E)) {
       Candidates.push_back(std::make_pair(underlyingToOpaque->getSubExpr(),
-                                  underlyingToOpaque->getSubExpr()->getType()));
+                                          underlyingToOpaque->substitutions));
       return {false, E};
     }
     return {true, E};

lib/Sema/TypeCheckConstraints.cpp

@@ -544,19 +544,8 @@ bool TypeChecker::typeCheckPatternBinding(PatternBindingDecl *PBD,
     if (auto opaque = var->getOpaqueResultTypeDecl()) {
       init->forEachChildExpr([&](Expr *expr) -> Expr * {
         if (auto coercionExpr = dyn_cast<UnderlyingToOpaqueExpr>(expr)) {
-          auto underlyingType =
-              coercionExpr->getSubExpr()->getType()->mapTypeOutOfContext();
-          auto underlyingSubs = SubstitutionMap::get(
-              opaque->getOpaqueInterfaceGenericSignature(),
-              [&](SubstitutableType *t) -> Type {
-                if (t->isEqual(opaque->getUnderlyingInterfaceType())) {
-                  return underlyingType;
-                }
-                return Type(t);
-              },
-              LookUpConformanceInModule(opaque->getModuleContext()));
-
-          opaque->setUnderlyingTypeSubstitutions(underlyingSubs);
+          opaque->setUnderlyingTypeSubstitutions(
+              coercionExpr->substitutions.mapReplacementTypesOutOfContext());
         }
         return expr;
       });