From 6a1d11bba2058924f03f8b131b176ac10e9f361c Mon Sep 17 00:00:00 2001
From: Enrico Granata <egranata@apple.com>
Date: Wed, 28 Sep 2016 11:27:30 -0700
Subject: [PATCH] Teach TypeReconstruction how to look through associated types
 of the Self protocol

---
 lib/IDE/TypeReconstruction.cpp | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/lib/IDE/TypeReconstruction.cpp b/lib/IDE/TypeReconstruction.cpp
index bcbb82d45d617..4a5626616de53 100644
--- a/lib/IDE/TypeReconstruction.cpp
+++ b/lib/IDE/TypeReconstruction.cpp
@@ -866,14 +866,27 @@ static void VisitNodeAssociatedTypeRef(
       ArchetypeType *archetype = type->getAs<ArchetypeType>();
       if (archetype) {
         Identifier identifier = ast->getIdentifier(ident->getText());
-        if (archetype->hasNestedType(identifier)) {
-          Type nested = archetype->getNestedTypeValue(identifier);
-          if (nested) {
-            result._types.push_back(nested);
-            result._module = type_result._module;
-            return;
+        Type nested;
+        if (archetype->hasNestedType(identifier))
+          nested = archetype->getNestedTypeValue(identifier);
+        else if (ProtocolDecl *self_protocol = archetype->getSelfProtocol()) {
+          for (auto self_member: self_protocol->getMembers()) {
+            if (AssociatedTypeDecl *associated_decl = llvm::dyn_cast_or_null<AssociatedTypeDecl>(self_member)) {
+              if (associated_decl->hasName()) {
+                llvm::StringRef decl_name = associated_decl->getNameStr();
+                if (decl_name == ident->getText()) {
+                  nested = associated_decl->getDeclaredType();
+                  break;
+                }
+              }
+            }
           }
         }
+        if (nested) {
+          result._types.push_back(nested);
+          result._module = type_result._module;
+          return;
+        }
       }
     }
   }