Sema: Clean up resolveTypeInContext() and generalize typealiases in protocols

Previously a typealias in a protocol either had to be concrete,
or consist of a single path of member types from Self, eg
Self.A.B. Lift this restriction.

Fix a few other corner cases that came up in the validation
suite, and clean up the function in general.

Author: slavapestov

lib/AST/Type.cpp

@@ -2917,15 +2917,8 @@ TypeSubstitutionMap TypeBase::getMemberSubstitutions(const DeclContext *dc) {
   // If the member is part of a protocol or extension thereof, we need
   // to substitute in the type of Self.
   if (dc->getAsProtocolOrProtocolExtensionContext()) {
-    // We only substitute into archetypes for now for protocols.
-    // FIXME: This seems like an odd restriction. Whatever is depending on
-    // this, shouldn't.
-    if (!baseTy->is<ArchetypeType>() && isa<ProtocolDecl>(dc))
-      return substitutions;
-
     // FIXME: This feels painfully inefficient. We're creating a dense map
     // for a single substitution.
-    substitutions[dc->getProtocolSelf()->getArchetype()] = baseTy;
     substitutions[dc->getProtocolSelf()->getDeclaredType()
                     ->getCanonicalType()->castTo<GenericTypeParamType>()]
       = baseTy;

lib/Sema/TypeCheckType.cpp

@@ -193,23 +193,6 @@ void TypeChecker::forceExternalDeclMembers(NominalTypeDecl *nominalDecl) {
   }
 }
 
-static Optional<Type>
-resolveAssociatedTypeInContext(TypeChecker &TC, AssociatedTypeDecl *assocType,
-                               DeclContext *DC, GenericTypeResolver *resolver) {
-  auto protoSelf = DC->getProtocolSelf();
-  auto selfTy = protoSelf->getDeclaredType()->castTo<GenericTypeParamType>();
-  auto baseTy = resolver->resolveGenericTypeParamType(selfTy);
-
-  if (baseTy->isTypeParameter())
-    return resolver->resolveSelfAssociatedType(baseTy, DC, assocType);
-
-  if (assocType->getDeclContext() != DC)
-    return TC.substMemberTypeWithBase(DC->getParentModule(), assocType,
-                                      protoSelf->getArchetype(),
-                                      /*isTypeReference=*/true);
-  return None;
-}
-
 Type TypeChecker::resolveTypeInContext(
        TypeDecl *typeDecl,
        DeclContext *fromDC,
@@ -239,15 +222,12 @@ Type TypeChecker::resolveTypeInContext(
     nominalType = nullptr;
   }
 
-  // Walk up through the type scopes to find the context where the type
-  // declaration was found. When we find it, substitute the appropriate base
-  // type.
+  // Walk up through the type scopes to find the context containing the type
+  // being resolved.
   auto ownerDC = typeDecl->getDeclContext();
   bool nonTypeOwner = !ownerDC->isTypeContext();
   auto ownerNominal = ownerDC->getAsNominalTypeOrNominalTypeExtensionContext();
   auto assocType = dyn_cast<AssociatedTypeDecl>(typeDecl);
-  auto alias = dyn_cast<TypeAliasDecl>(typeDecl);
-  DeclContext *typeParent = nullptr;
   assert((ownerNominal || nonTypeOwner) &&
          "Owner must be a nominal type or a non type context");
 
@@ -273,140 +253,126 @@ Type TypeChecker::resolveTypeInContext(
       return typeDecl->getDeclaredType();
 
     // For the next steps we need our parentDC to be a type context
-    if (!parentDC->isTypeContext()) {
+    if (!parentDC->isTypeContext())
       continue;
-    } else if (!typeParent) {
-      // Remember the first type decl context in the hierarchy for later use
-      typeParent = parentDC;
-    }
-
-    // If we found an associated type in an inherited protocol, the base for our
-    // reference to this associated type is our own `Self`. If we can't resolve
-    // the associated type during this iteration, try again on the next.
-    if (assocType) {
-      if (auto proto = parentDC->getAsProtocolOrProtocolExtensionContext()) {
-        auto assocProto = assocType->getProtocol();
-        if (proto == assocProto || proto->inheritsFrom(assocProto)) {
-          // If the associated type is from our own protocol or we inherit from
-          // the associated type's protocol, resolve it
-          if (auto resolved = resolveAssociatedTypeInContext(
-                  *this, assocType, parentDC, resolver))
-            return *resolved;
-
-        } else if (auto ED = dyn_cast<ExtensionDecl>(parentDC)) {
-          // Otherwise, if we are in an extension there might be other
-          // associated types brought into the context through
-          // `extension ... where Self : SomeProtocol`
-          for (auto req : ED->getGenericParams()->getTrailingRequirements()) {
-            // Reject requirements other than constraints with an subject other
-            // than `Self`
-            if (req.getKind() != RequirementReprKind::TypeConstraint ||
-                !req.getSubject()->castTo<ArchetypeType>()->isSelfDerived())
+
+    // Search the type of this context and its supertypes (if its a
+    // class) or refined protocols (if its a protocol).
+    llvm::SmallPtrSet<const NominalTypeDecl *, 8> visited;
+    llvm::SmallVector<Type, 8> stack;
+
+    // Start with the type of the current context.
+    if (auto fromType = resolver->resolveTypeOfContext(parentDC))
+      stack.push_back(fromType);
+
+    // If we are in a protocol extension there might be other type aliases and
+    // nominal types brought into the context through requirements on Self,
+    // for example:
+    //
+    // extension MyProtocol where Self : YourProtocol { ... }
+    if (parentDC->getAsProtocolExtensionContext()) {
+      auto ED = cast<ExtensionDecl>(parentDC);
+      if (auto genericParams = ED->getGenericParams()) {
+        for (auto req : genericParams->getTrailingRequirements()) {
+          // We might be resolving 'req.getSubject()' itself.
+          // This whole case feels like a hack -- there should be a
+          // more principled way to represent extensions of protocol
+          // compositions.
+          if (req.getKind() == RequirementReprKind::TypeConstraint) {
+            if (!req.getSubject() ||
+                !req.getSubject()->is<ArchetypeType>() ||
+                !req.getSubject()->castTo<ArchetypeType>()->getSelfProtocol())
               continue;
 
-            // If the associated type is defined in the same protocol which is
-            // required for this extension, or if the required protocol inherits
-            // from the protocol the associated type is declared in, we can
-            // resolve the associated type with our `Self` as the reference
-            // point.
-            auto reqProto =
-                req.getConstraint()->castTo<ProtocolType>()->getDecl();
-            if (reqProto == assocProto || reqProto->inheritsFrom(assocProto)) {
-              if (auto resolved = resolveAssociatedTypeInContext(
-                      *this, assocType, parentDC, resolver))
-                return *resolved;
-              break;
-            }
+            stack.push_back(req.getConstraint());
           }
         }
       }
     }
-    
-    // If we found an alias type in an inherited protocol, resolve it based on our
-    // own `Self`.
-    if (alias && alias->hasInterfaceType()) {
-      auto metaType = alias->getInterfaceType()->getAs<MetatypeType>();
-      auto memberType = metaType ? metaType->getInstanceType()->getAs<DependentMemberType>() :
-                        nullptr;
-
-      if (memberType && parentDC->getAsProtocolOrProtocolExtensionContext()) {
-        auto protoSelf = parentDC->getProtocolSelf();
-        auto selfTy = protoSelf->getDeclaredType()->castTo<GenericTypeParamType>();
-        auto baseTy = resolver->resolveGenericTypeParamType(selfTy);
-
-        SmallVector<DependentMemberType *, 4> memberTypes;
-        do {
-          memberTypes.push_back(memberType);
-          memberType = memberType->getBase()->getAs<DependentMemberType>();
-        } while (memberType);
-
-        auto module = parentDC->getParentModule();
-        while (memberTypes.size()) {
-          baseTy = memberTypes.back()->substBaseType(module, baseTy, nullptr);
-          memberTypes.pop_back();
-        }
-        return baseTy;
-      }
-    }
 
-    // Search the type of this context and its supertypes.
-    llvm::SmallPtrSet<const NominalTypeDecl *, 8> visited;
-    for (auto fromType = resolver->resolveTypeOfContext(parentDC);
-         fromType;
-         fromType = getSuperClassOf(fromType)) {
+    while (!stack.empty()) {
+      auto fromType = stack.back();
+      auto *fromProto = parentDC->getAsProtocolOrProtocolExtensionContext();
+
+      stack.pop_back();
+
       // If we hit circularity, we will diagnose at some point in typeCheckDecl().
       // However we have to explicitly guard against that here because we get
       // called as part of validateDecl().
       if (!visited.insert(fromType->getAnyNominal()).second)
-        break;
+        continue;
+
+      // Handle this case:
+      // - Current context: concrete type
+      // - Nested type: associated type
+      // - Nested type's context: protocol or protocol extension
+      //
+      if (assocType && fromProto == nullptr) {
+        ProtocolConformance *conformance = nullptr;
+
+        // If the conformance check failed, the associated type is for a
+        // conformance of an outer context.
+        if (!options.contains(TR_InheritanceClause) &&
+            conformsToProtocol(fromType,
+                               cast<ProtocolDecl>(ownerNominal),
+                               parentDC, ConformanceCheckFlags::Used,
+                               &conformance) &&
+            conformance) {
+          return conformance->getTypeWitness(assocType, this).getReplacement();
+        }
+      }
 
-      // If the nominal type declaration of the context type we're looking at
-      // matches the owner's nominal type declaration, this is how we found
-      // the member type declaration. Substitute the type we're coming from as
-      // the base of the member type to produce the projected type result.
+      // Handle these cases:
+      // - Current context: concrete type
+      // - Nested type: concrete type or type alias
+      // - Nested type's context: concrete type
+      //
+      // - Current context: protocol or protocol extension
+      // - Nested type: type alias
+      // - Nested type's context: protocol or protocol extension
+      //
+      // Note: this is not supported yet, FIXME:
+      // - Current context: concrete type
+      // - Nested type: type alias
+      // - Nested type's context: protocol or protocol extension
+      //
       if (fromType->getAnyNominal() == ownerNominal) {
-        // If we are referring into a protocol or extension thereof,
-        // the base type is the 'Self'.
-        if (ownerDC->getAsProtocolOrProtocolExtensionContext()) {
-          auto selfTy = ownerDC->getProtocolSelf()->getDeclaredType()
-                          ->castTo<GenericTypeParamType>();
-          fromType = resolver->resolveGenericTypeParamType(selfTy);
+        if (fromProto &&
+            ownerNominal->getAsProtocolOrProtocolExtensionContext()) {
+          // If we are looking up an associated type or a protocol's type alias
+          // from a protocol or protocol extension, use the archetype for 'Self'
+          // instead of the existential type.
+          assert(fromType->is<ProtocolType>());
+
+          auto protoSelf = parentDC->getProtocolSelf();
+          auto selfType = protoSelf
+              ->getDeclaredType()
+              ->castTo<GenericTypeParamType>();
+          fromType = resolver->resolveGenericTypeParamType(selfType);
+
+          if (assocType) {
+            // Odd special case, ask Doug to explain it over pizza one day
+            if (fromType->isTypeParameter())
+              return resolver->resolveSelfAssociatedType(
+                  fromType, parentDC, assocType);
+          }
         }
 
-        // Perform the substitution.
         return substMemberTypeWithBase(parentDC->getParentModule(), typeDecl,
                                        fromType, /*isTypeReference=*/true);
       }
 
-      ProtocolConformance *conformance = nullptr;
-      if (assocType &&
-          !options.contains(TR_InheritanceClause) &&
-          conformsToProtocol(fromType,
-                             cast<ProtocolDecl>(assocType->getDeclContext()),
-                             parentDC, ConformanceCheckFlags::Used,
-                             &conformance) &&
-          conformance) {
-        return conformance->getTypeWitness(assocType, this).getReplacement();
+      if (auto superclassTy = getSuperClassOf(fromType))
+        stack.push_back(superclassTy);
+      else if (auto protoTy = fromType->getAs<ProtocolType>()) {
+        for (auto *proto : protoTy->getDecl()->getInheritedProtocols(this))
+          if (auto refinedTy = proto->getDeclaredTypeInContext())
+            stack.push_back(refinedTy);
       }
     }
   }
 
-  // At this point by iterating through the decl context hierarchy we should
-  // have encountered the first type context in the stack.
-  assert(typeParent && "incomplete iteration");
-  assert(!typeParent->isModuleContext());
-
-  // Substitute in the appropriate type for 'Self'.
-  // FIXME: We shouldn't have to guess here; the caller should tell us.
-  Type fromType;
-  if (typeParent->getAsProtocolOrProtocolExtensionContext())
-    fromType = typeParent->getProtocolSelf()->getArchetype();
-  else
-    fromType = resolver->resolveTypeOfContext(typeParent);
-
-  // Perform the substitution.
-  return substMemberTypeWithBase(typeParent->getParentModule(), typeDecl,
-                                 fromType, /*isTypeReference=*/true);
+  llvm_unreachable("Cannot resolve type");
 }
 
 Type TypeChecker::applyGenericArguments(Type type, SourceLoc loc,
@@ -537,7 +503,8 @@ Type TypeChecker::applyUnboundGenericArguments(
 
     // Check the generic arguments against the generic signature.
     auto genericSig = unbound->getDecl()->getGenericSignature();
-    if (unbound->getDecl()->isValidatingGenericSignature()) {
+    if (!unbound->getDecl()->hasType() ||
+        unbound->getDecl()->isValidatingGenericSignature()) {
       diagnose(loc, diag::recursive_requirement_reference);
       return nullptr;
     }
@@ -2474,7 +2441,6 @@ Type TypeChecker::substMemberTypeWithBase(Module *module,
   Type memberType = isTypeReference
                       ? cast<TypeDecl>(member)->getDeclaredInterfaceType()
                       : member->getInterfaceType();
-
   if (isTypeReference) {
     // The declared interface type for a generic type will have the type
     // arguments; strip them off.

test/decl/typealias/typealias.swift

@@ -158,8 +158,12 @@ protocol Col {
 protocol CB {
   associatedtype C : Col
   typealias E = C.Elem
+  typealias OE = C.Elem?
+  typealias FE = (C) -> (C.Elem) -> Self
 
   func setIt(_ element: E)
+  func setItMaybe(_ element: OE)
+  func teleport(_ fn: FE)
 }
 
 func go1<T : CB, U : Col where U.Elem == T.E>(_ col: U, builder: T) { // OK

validation-test/IDE/crashers/021-swift-typechecker-resolveidentifiertype.swift

@@ -0,0 +1,2 @@
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=A -source-filename=%s
+{class T{protocol b{#^A^#associatedtype b:B<T>struct B<b

validation-test/IDE/crashers/024-swift-archetypebuilder-potentialarchetype-getrepresentative.swift

@@ -0,0 +1,4 @@
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=A -source-filename=%s
+protocol c{func a:P
+protocol P{#^A^#func a:b
+associatedtype b:a

validation-test/IDE/crashers/044-swift-dependentgenerictyperesolver-resolveselfassociatedtype.swift

@@ -0,0 +1,5 @@
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=A -source-filename=%s
+protocol A{protocol A{
+associatedtype e:a
+func a<T:e
+enum b{func a{#^A^#

validation-test/IDE/crashers_fixed/021-swift-typechecker-resolveidentifiertype.swift

@@ -5,7 +5,7 @@
 // See http://swift.org/LICENSE.txt for license information
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 
-// RUN: not --crash %target-swift-frontend %s -parse
+// RUN: not %target-swift-frontend %s -parse
 protocol a {
 protocol A {
 }

validation-test/IDE/crashers_fixed/024-swift-archetypebuilder-potentialarchetype-getrepresentative.swift

@@ -6,5 +6,5 @@
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 
 // DUPLICATE-OF: 01766-swift-typechecker-validatedecl.swift
-// RUN: not --crash %target-swift-frontend %s -parse
+// RUN: not %target-swift-frontend %s -parse
 protocol e:A{protocol e:A class A}protocol A{enum A}protocol B:e

validation-test/IDE/crashers_fixed/044-swift-dependentgenerictyperesolver-resolveselfassociatedtype.swift

@@ -5,7 +5,7 @@
 // See http://swift.org/LICENSE.txt for license information
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 
-// RUN: not --crash %target-swift-frontend %s -parse
+// RUN: not %target-swift-frontend %s -parse
 struct B<T{protocol A{
 class B<T
 associatedtype e:B<T>

validation-test/compiler_crashers/01766-swift-typechecker-validatedecl.swift renamed to validation-test/compiler_crashers_fixed/01766-swift-typechecker-validatedecl.swift

@@ -5,6 +5,6 @@
 // See http://swift.org/LICENSE.txt for license information
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 
-// RUN: not --crash %target-swift-frontend %s -parse
+// RUN: not %target-swift-frontend %s -parse
 // REQUIRES: asserts
 var d:Collection{for c d in

validation-test/compiler_crashers/26083-std-function-func-swift-type-subst.swift renamed to validation-test/compiler_crashers_fixed/26083-std-function-func-swift-type-subst.swift

@@ -6,5 +6,5 @@
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 
 // DUPLICATE-OF: 01766-swift-typechecker-validatedecl.swift
-// RUN: not --crash %target-swift-frontend %s -parse
+// RUN: not %target-swift-frontend %s -parse
 protocol A{class A}protocol a:A{protocol P{associatedtype e:A}}a