[AST] Allow tentative type witnesses to be plumbed through SubstOptions.

Extend SubstOptions, which controls how substitution is performed, to
allow the caller to subst() to provide a callback function that may
provide a type witness for a normal protocol conformance that is
undergoing type witness inference. In effect, it's allowing us to
provide tentative bindings for type witnesses so we can see the
effects of substitution.

Author: DougGregor

include/swift/AST/ProtocolConformance.h

@@ -145,13 +145,15 @@ class alignas(1 << DeclAlignInBits) ProtocolConformance {
 
   /// Retrieve the type witness for the given associated type.
   Type getTypeWitness(AssociatedTypeDecl *assocType,
-                      LazyResolver *resolver) const;
+                      LazyResolver *resolver,
+                      SubstOptions options = None) const;
 
   /// Retrieve the type witness and type decl (if one exists)
   /// for the given associated type.
   std::pair<Type, TypeDecl *>
   getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                        LazyResolver *resolver) const;
+                        LazyResolver *resolver,
+                        SubstOptions options = None) const;
 
   /// Apply the given function object to each type witness within this
   /// protocol conformance.
@@ -399,7 +401,8 @@ class NormalProtocolConformance : public ProtocolConformance,
   /// for the given associated type.
   std::pair<Type, TypeDecl *>
   getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                        LazyResolver *resolver) const;
+                        LazyResolver *resolver,
+                        SubstOptions options = None) const;
 
   /// Determine whether the protocol conformance has a type witness for the
   /// given associated type.
@@ -568,7 +571,8 @@ class SpecializedProtocolConformance : public ProtocolConformance,
   /// for the given associated type.
   std::pair<Type, TypeDecl *>
   getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                        LazyResolver *resolver) const;
+                        LazyResolver *resolver,
+                        SubstOptions options = None) const;
 
   /// Given that the requirement signature of the protocol directly states
   /// that the given dependent type must conform to the given protocol,
@@ -662,8 +666,10 @@ class InheritedProtocolConformance : public ProtocolConformance,
   /// for the given associated type.
   std::pair<Type, TypeDecl *>
   getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                        LazyResolver *resolver) const {
-    return InheritedConformance->getTypeWitnessAndDecl(assocType, resolver);
+                        LazyResolver *resolver,
+                        SubstOptions options = None) const {
+    return InheritedConformance->getTypeWitnessAndDecl(assocType, resolver,
+                                                       options);
   }
 
   /// Given that the requirement signature of the protocol directly states

include/swift/AST/Type.h

@@ -165,7 +165,26 @@ enum class SubstFlags {
 };
 
 /// Options for performing substitutions into a type.
-typedef OptionSet<SubstFlags> SubstOptions;
+struct SubstOptions : public OptionSet<SubstFlags> {
+  typedef std::function<Type(const NormalProtocolConformance *,
+                             AssociatedTypeDecl *)>
+    GetTentativeTypeWitness;
+
+  /// Function that retrieves a tentative type witness for a protocol
+  /// conformance with the state \c CheckingTypeWitnesses.
+  GetTentativeTypeWitness getTentativeTypeWitness;
+
+  SubstOptions(llvm::NoneType) : OptionSet(None) { }
+
+  SubstOptions(SubstFlags flags) : OptionSet(flags) { }
+
+  SubstOptions(OptionSet<SubstFlags> options) : OptionSet(options) { }
+
+  SubstOptions(OptionSet<SubstFlags> options,
+              GetTentativeTypeWitness getTentativeTypeWitness)
+    : OptionSet(options),
+      getTentativeTypeWitness(std::move(getTentativeTypeWitness)) { }
+};
 
 inline SubstOptions operator|(SubstFlags lhs, SubstFlags rhs) {
   return SubstOptions(lhs) | rhs;

lib/AST/ProtocolConformance.cpp

@@ -215,13 +215,16 @@ ProtocolConformance::hasTypeWitness(AssociatedTypeDecl *assocType,
 
 std::pair<Type, TypeDecl *>
 ProtocolConformance::getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                                                LazyResolver *resolver) const {
-  CONFORMANCE_SUBCLASS_DISPATCH(getTypeWitnessAndDecl, (assocType, resolver))
+                                           LazyResolver *resolver,
+                                           SubstOptions options) const {
+  CONFORMANCE_SUBCLASS_DISPATCH(getTypeWitnessAndDecl,
+                                (assocType, resolver, options))
 }
 
 Type ProtocolConformance::getTypeWitness(AssociatedTypeDecl *assocType,
-                                         LazyResolver *resolver) const {
-  return getTypeWitnessAndDecl(assocType, resolver).first;
+                                         LazyResolver *resolver,
+                                         SubstOptions options) const {
+  return getTypeWitnessAndDecl(assocType, resolver, options).first;
 }
 
 ValueDecl *ProtocolConformance::getWitnessDecl(ValueDecl *requirement,
@@ -472,22 +475,39 @@ static bool resolveKnownTypeWitness(NormalProtocolConformance *conformance,
 
 std::pair<Type, TypeDecl *>
 NormalProtocolConformance::getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
-                                                 LazyResolver *resolver) const {
+                                                 LazyResolver *resolver,
+                                                 SubstOptions options) const {
   if (Resolver)
     resolveLazyInfo();
 
+  // Check whether we already have a type witness.
   auto known = TypeWitnesses.find(assocType);
-  if (known == TypeWitnesses.end()) {
-    PrettyStackTraceRequirement trace("resolving", this, assocType);
-    if (!resolveKnownTypeWitness(const_cast<NormalProtocolConformance *>(this),
-                                 assocType)) {
-      assert(resolver && "Unable to resolve type witness");
-      resolver->resolveTypeWitness(this, assocType);
+  if (known != TypeWitnesses.end())
+    return known->second;
+
+  // If this conformance is in a state where it is inferring type witnesses,
+  // check tentative witnesses.
+  if (getState() == ProtocolConformanceState::CheckingTypeWitnesses) {
+    // If there is a tentative-type-witness function, use it.
+    if (options.getTentativeTypeWitness) {
+      if (Type witnessType = options.getTentativeTypeWitness(this, assocType))
+        return { witnessType, nullptr };
     }
-    known = TypeWitnesses.find(assocType);
-    assert(known != TypeWitnesses.end() && "Didn't resolve witness?");
+
+    // Otherwise, we fail; this is the only case in which we can retturn a
+    // null type.
+    return { Type(), nullptr };
   }
 
+  // Otherwise, resolve the type witness.
+  PrettyStackTraceRequirement trace("resolving", this, assocType);
+  if (!resolveKnownTypeWitness(const_cast<NormalProtocolConformance *>(this),
+                               assocType)) {
+    assert(resolver && "Unable to resolve type witness");
+    resolver->resolveTypeWitness(this, assocType);
+  }
+  known = TypeWitnesses.find(assocType);
+  assert(known != TypeWitnesses.end() && "Didn't resolve witness?");
   return known->second;
 }
 
@@ -656,7 +676,8 @@ bool SpecializedProtocolConformance::hasTypeWitness(
 std::pair<Type, TypeDecl *>
 SpecializedProtocolConformance::getTypeWitnessAndDecl(
                       AssociatedTypeDecl *assocType, 
-                      LazyResolver *resolver) const {
+                      LazyResolver *resolver,
+                      SubstOptions options) const {
   // If we've already created this type witness, return it.
   auto known = TypeWitnesses.find(assocType);
   if (known != TypeWitnesses.end()) {
@@ -669,26 +690,41 @@ SpecializedProtocolConformance::getTypeWitnessAndDecl(
   auto substitutionMap =
       genericSig->getSubstitutionMap(GenericSubstitutions);
 
+  // Local function to determine whether we will end up
+  auto normal = GenericConformance->getRootNormalConformance();
+  auto isTentativeWitness = [&] {
+    if (normal->getState() != ProtocolConformanceState::CheckingTypeWitnesses)
+      return false;
+
+    return !normal->hasTypeWitness(assocType, nullptr);
+  };
+
   auto genericWitnessAndDecl
-    = GenericConformance->getTypeWitnessAndDecl(assocType, resolver);
+    = GenericConformance->getTypeWitnessAndDecl(assocType, resolver, options);
+
+  auto genericWitness = genericWitnessAndDecl.first;
+  if (!genericWitness)
+    return { Type(), nullptr };
 
-  auto &genericWitness = genericWitnessAndDecl.first;
   auto *typeDecl = genericWitnessAndDecl.second;
 
   // Apply the substitution we computed above
   auto specializedType
-    = genericWitness.subst(substitutionMap);
-  if (!specializedType)
-    specializedType = ErrorType::get(genericWitness);
+    = genericWitness.subst(substitutionMap, options);
+  if (!specializedType) {
+    if (isTentativeWitness())
+      return { Type(), nullptr };
 
-  // If the type witness was unchanged, just copy it directly.
-  if (specializedType.getPointer() == genericWitness.getPointer()) {
-    TypeWitnesses[assocType] = genericWitnessAndDecl;
-    return TypeWitnesses[assocType];
+    specializedType = ErrorType::get(genericWitness);
   }
 
-  TypeWitnesses[assocType] = std::make_pair(specializedType, typeDecl);
-  return TypeWitnesses[assocType];
+  // If we aren't in a case where we used the tentative type witness
+  // information, cache the result.
+  auto specializedWitnessAndDecl = std::make_pair(specializedType, typeDecl);
+  if (!isTentativeWitness() && !specializedType->hasError())
+    TypeWitnesses[assocType] = specializedWitnessAndDecl;
+
+  return specializedWitnessAndDecl;
 }
 
 ProtocolConformanceRef

lib/AST/Type.cpp

@@ -2909,17 +2909,11 @@ static Type getMemberForBaseType(LookupConformanceFn lookupConformances,
     if (!conformance) return failed();
     if (!conformance->isConcrete()) return failed();
 
-    // If we have an unsatisfied type witness while we're checking the
-    // conformances we're supposed to skip this conformance's unsatisfied type
-    // witnesses, and we have an unsatisfied type witness, return
-    // "missing".
-    if (conformance->getConcrete()->getRootNormalConformance()->getState()
-          == ProtocolConformanceState::CheckingTypeWitnesses &&
-        !conformance->getConcrete()->hasTypeWitness(assocType, nullptr))
-      return failed();
-
+    // Retrieve the type witness.
     auto witness =
-      conformance->getConcrete()->getTypeWitness(assocType, resolver);
+      conformance->getConcrete()->getTypeWitness(assocType, resolver, options);
+    if (!witness)
+      return failed();
 
     // This is a hacky feature allowing code completion to migrate to
     // using Type::subst() without changing output.