[IDE] Teach type checker about conditional conformance extensions.

Before conditional conformances, the archetypes in conformance
extensions (i.e. extension Foo: SomeProtocol) were equivalent to those
in the type decl, with the same protocol bounds and so on. The code for
printing "synthesized" members relied on this fact. This commit teaches
that code to deal with archetypes in the conditional conformance
extension when required.

Fixes rdar://problem/36553066 and SR-6930.

include/swift/AST/ASTContext.h

@@ -174,10 +174,6 @@ public:
 };
 
 class SILLayout; // From SIL
-/// \brief Describes either a nominal type declaration or an extension
-/// declaration.
-typedef llvm::PointerUnion<NominalTypeDecl *, ExtensionDecl *>
-  TypeOrExtensionDecl;
 
 /// ASTContext - This object creates and owns the AST objects.
 /// However, this class does more than just maintain context within an AST.

include/swift/AST/ASTPrinter.h

@@ -84,7 +84,7 @@ enum class PrintStructureKind {
 class ASTPrinter {
   unsigned CurrentIndentation = 0;
   unsigned PendingNewlines = 0;
-  const NominalTypeDecl *SynthesizeTarget = nullptr;
+  TypeOrExtensionDecl SynthesizeTarget;
 
   void printTextImpl(StringRef Text);
 
@@ -136,13 +136,14 @@ public:
 
   /// Called before printing a synthesized extension.
   virtual void printSynthesizedExtensionPre(const ExtensionDecl *ED,
-                                            const NominalTypeDecl *NTD,
+                                            TypeOrExtensionDecl NTD,
                                             Optional<BracketOptions> Bracket) {}
 
   /// Called after printing a synthesized extension.
   virtual void printSynthesizedExtensionPost(const ExtensionDecl *ED,
-                                             const NominalTypeDecl *NTD,
-                                             Optional<BracketOptions> Bracket) {}
+                                             TypeOrExtensionDecl TargetDecl,
+                                             Optional<BracketOptions> Bracket) {
+  }
 
   /// Called before printing a structured entity.
   ///
@@ -213,7 +214,7 @@ public:
     CurrentIndentation = NumSpaces;
   }
 
-  void setSynthesizedTarget(NominalTypeDecl *Target) {
+  void setSynthesizedTarget(TypeOrExtensionDecl Target) {
     assert((!SynthesizeTarget || !Target || Target == SynthesizeTarget) &&
            "unexpected change of setSynthesizedTarget");
     // FIXME: this can overwrite the original target with nullptr.

include/swift/AST/PrintOptions.h

@@ -36,16 +36,45 @@ enum DeclAttrKind : unsigned;
 class SynthesizedExtensionAnalyzer;
 struct PrintOptions;
 
+/// \brief Describes either a nominal type declaration or an extension
+/// declaration.
+struct TypeOrExtensionDecl {
+  llvm::PointerUnion<NominalTypeDecl *, ExtensionDecl *> Decl;
+
+  TypeOrExtensionDecl() = default;
+
+  TypeOrExtensionDecl(NominalTypeDecl *D);
+  TypeOrExtensionDecl(ExtensionDecl *D);
+
+  /// \brief Return the contained *Decl as the Decl superclass.
+  class Decl *getAsDecl() const;
+  /// \brief Return the contained *Decl as the DeclContext superclass.
+  DeclContext *getAsDeclContext() const;
+  /// \brief Return the contained NominalTypeDecl or that of the extended type
+  /// in the ExtensionDecl.
+  NominalTypeDecl *getBaseNominal() const;
+
+  /// \brief Is the contained pointer null?
+  bool isNull() const;
+  explicit operator bool() const { return !isNull(); }
+
+  bool operator==(TypeOrExtensionDecl rhs) { return Decl == rhs.Decl; }
+  bool operator!=(TypeOrExtensionDecl rhs) { return Decl != rhs.Decl; }
+  bool operator<(TypeOrExtensionDecl rhs) { return Decl < rhs.Decl; }
+};
+
 /// Necessary information for archetype transformation during printing.
 struct TypeTransformContext {
   TypeBase *BaseType;
-  NominalTypeDecl *Nominal = nullptr;
+  TypeOrExtensionDecl Decl;
 
   explicit TypeTransformContext(Type T);
-  explicit TypeTransformContext(NominalTypeDecl* NTD);
+  explicit TypeTransformContext(TypeOrExtensionDecl D);
 
   Type getBaseType() const;
-  NominalTypeDecl *getNominal() const;
+  TypeOrExtensionDecl getDecl() const;
+
+  DeclContext *getDeclContext() const;
 
   bool isPrintingSynthesizedExtension() const;
 };
@@ -421,7 +450,7 @@ struct PrintOptions {
 
   void setBaseType(Type T);
 
-  void initForSynthesizedExtension(NominalTypeDecl *D);
+  void initForSynthesizedExtension(TypeOrExtensionDecl D);
 
   void clearSynthesizedExtension();
 

include/swift/Sema/IDETypeChecking.h

@@ -137,9 +137,17 @@ namespace swift {
   /// Creates a lazy type resolver for use in lookups.
   OwnedResolver createLazyResolver(ASTContext &Ctx);
 
-  typedef std::pair<ExtensionDecl*, bool> ExtensionAndIsSynthesized;
+  struct ExtensionInfo {
+    // The extension with the declarations to apply.
+    ExtensionDecl *Ext;
+    // The extension that enables the former to apply, if any (i.e. a
+    // conditional
+    // conformance to Foo enables 'extension Foo').
+    ExtensionDecl *EnablingExt;
+    bool IsSynthesized;
+  };
 
-  typedef llvm::function_ref<void(ArrayRef<ExtensionAndIsSynthesized>)>
+  typedef llvm::function_ref<void(ArrayRef<ExtensionInfo>)>
       ExtensionGroupOperation;
 
   class SynthesizedExtensionAnalyzer {

lib/AST/ASTPrinter.cpp

@@ -55,7 +55,7 @@ void PrintOptions::setBaseType(Type T) {
   TransformContext = TypeTransformContext(T);
 }
 
-void PrintOptions::initForSynthesizedExtension(NominalTypeDecl *D) {
+void PrintOptions::initForSynthesizedExtension(TypeOrExtensionDecl D) {
   TransformContext = TypeTransformContext(D);
 }
 
@@ -63,16 +63,40 @@ void PrintOptions::clearSynthesizedExtension() {
   TransformContext.reset();
 }
 
+TypeOrExtensionDecl::TypeOrExtensionDecl(NominalTypeDecl *D) : Decl(D) {}
+TypeOrExtensionDecl::TypeOrExtensionDecl(ExtensionDecl *D) : Decl(D) {}
+
+Decl *TypeOrExtensionDecl::getAsDecl() const {
+  if (auto NTD = Decl.dyn_cast<NominalTypeDecl *>())
+    return NTD;
+
+  return Decl.get<ExtensionDecl *>();
+}
+DeclContext *TypeOrExtensionDecl::getAsDeclContext() const {
+  return getAsDecl()->getInnermostDeclContext();
+}
+NominalTypeDecl *TypeOrExtensionDecl::getBaseNominal() const {
+  return getAsDeclContext()->getAsNominalTypeOrNominalTypeExtensionContext();
+}
+bool TypeOrExtensionDecl::isNull() const { return Decl.isNull(); }
+
 TypeTransformContext::TypeTransformContext(Type T)
     : BaseType(T.getPointer()) {
   assert(T->mayHaveMembers());
 }
 
-TypeTransformContext::TypeTransformContext(NominalTypeDecl *NTD)
-    : BaseType(NTD->getDeclaredTypeInContext().getPointer()), Nominal(NTD) {}
+TypeTransformContext::TypeTransformContext(TypeOrExtensionDecl D)
+    : BaseType(nullptr), Decl(D) {
+  if (auto NTD = Decl.Decl.dyn_cast<NominalTypeDecl *>())
+    BaseType = NTD->getDeclaredTypeInContext().getPointer();
+  else
+    BaseType = Decl.Decl.get<ExtensionDecl *>()->getExtendedType().getPointer();
+}
 
-NominalTypeDecl *TypeTransformContext::getNominal() const {
-  return Nominal;
+TypeOrExtensionDecl TypeTransformContext::getDecl() const { return Decl; }
+
+DeclContext *TypeTransformContext::getDeclContext() const {
+  return Decl.getAsDecl()->getDeclContext();
 }
 
 Type TypeTransformContext::getBaseType() const {
@@ -80,7 +104,7 @@ Type TypeTransformContext::getBaseType() const {
 }
 
 bool TypeTransformContext::isPrintingSynthesizedExtension() const {
-  return Nominal != nullptr;
+  return !Decl.isNull();
 }
 
 std::string ASTPrinter::sanitizeUtf8(StringRef Text) {
@@ -631,8 +655,7 @@ private:
 #define STMT(Name, Parent) void visit##Name##Stmt(Name##Stmt *stmt);
 #include "swift/AST/StmtNodes.def"
 
-  void printSynthesizedExtension(NominalTypeDecl* Decl,
-                                 ExtensionDecl* ExtDecl);
+  void printSynthesizedExtension(Type ExtendedType, ExtensionDecl *ExtDecl);
 
   void printExtension(ExtensionDecl* ExtDecl);
 
@@ -668,8 +691,9 @@ public:
         Options.TransformContext &&
         Options.TransformContext->isPrintingSynthesizedExtension() &&
         isa<ExtensionDecl>(D);
-    if (Synthesize)
-      Printer.setSynthesizedTarget(Options.TransformContext->getNominal());
+    if (Synthesize) {
+      Printer.setSynthesizedTarget(Options.TransformContext->getDecl());
+    }
 
     // We want to print a newline before doc comments.  Swift code already
     // handles this, but we need to insert it for clang doc comments when not
@@ -692,9 +716,9 @@ public:
     ASTVisitor::visit(D);
 
     if (Synthesize) {
-      Printer.setSynthesizedTarget(nullptr);
-      Printer.printSynthesizedExtensionPost(
-          cast<ExtensionDecl>(D), Options.TransformContext->getNominal(),
+      Printer.setSynthesizedTarget({});
+      Printer.printSynthesizedExtensionPost(cast<ExtensionDecl>(D),
+                                            Options.TransformContext->getDecl(),
                                             Options.BracketOptions);
     } else {
       Printer.callPrintDeclPost(D, Options.BracketOptions);
@@ -1775,14 +1799,14 @@ static void printExtendedTypeName(Type ExtendedType, ASTPrinter &Printer,
   Printer.printTypeRef(ExtendedType, Nominal, Nominal->getName());
 }
 
-void PrintAST::
-printSynthesizedExtension(NominalTypeDecl* Decl, ExtensionDecl *ExtDecl) {
+void PrintAST::printSynthesizedExtension(Type ExtendedType,
+                                         ExtensionDecl *ExtDecl) {
   if (Options.BracketOptions.shouldOpenExtension(ExtDecl)) {
     printDocumentationComment(ExtDecl);
     printAttributes(ExtDecl);
     Printer << tok::kw_extension << " ";
 
-    printExtendedTypeName(Decl->getDeclaredType(), Printer, Options);
+    printExtendedTypeName(ExtendedType, Printer, Options);
     printInherited(ExtDecl);
 
     if (ExtDecl->getGenericParams())
@@ -1832,9 +1856,11 @@ void PrintAST::printExtension(ExtensionDecl *decl) {
 
 void PrintAST::visitExtensionDecl(ExtensionDecl *decl) {
   if (Options.TransformContext &&
-      Options.TransformContext->isPrintingSynthesizedExtension())
-    printSynthesizedExtension(Options.TransformContext->getNominal(), decl);
-  else
+      Options.TransformContext->isPrintingSynthesizedExtension()) {
+    auto extendedType =
+        Options.TransformContext->getBaseType()->mapTypeOutOfContext();
+    printSynthesizedExtension(extendedType, decl);
+  } else
     printExtension(decl);
 }
 

lib/IDE/IDETypeChecking.cpp

@@ -31,14 +31,13 @@ static bool shouldPrintAsFavorable(const Decl *D, PrintOptions &Options) {
       !D->getDeclContext()->isExtensionContext() ||
       !Options.TransformContext->isPrintingSynthesizedExtension())
     return true;
-  NominalTypeDecl *Target = Options.TransformContext->getNominal();
-  Type BaseTy = Target->getDeclaredTypeInContext();
+  auto DC = Options.TransformContext->getDeclContext();
+  auto BaseTy = Options.TransformContext->getBaseType();
   const auto *FD = dyn_cast<FuncDecl>(D);
   if (!FD)
     return true;
   ResolvedMemberResult Result =
-  resolveValueMember(*Target->getDeclContext(), BaseTy,
-                     FD->getEffectiveFullName());
+      resolveValueMember(*DC, BaseTy, FD->getEffectiveFullName());
   return !(Result.hasBestOverload() && Result.getBestOverload() != D);
 }
 
@@ -122,9 +121,11 @@ struct SynthesizedExtensionAnalyzer::Implementation {
   struct SynthesizedExtensionInfo {
     ExtensionDecl *Ext = nullptr;
     bool IsSynthesized;
+    ExtensionDecl *EnablingExt = nullptr;
     operator bool() const { return Ext; }
-    SynthesizedExtensionInfo(bool IsSynthesized = true) :
-    IsSynthesized(IsSynthesized) {}
+    SynthesizedExtensionInfo(bool IsSynthesized = false,
+                             ExtensionDecl *EnablingExt = nullptr)
+        : IsSynthesized(IsSynthesized), EnablingExt(EnablingExt) {}
     bool operator< (const SynthesizedExtensionInfo& Rhs) const {
 
       // Synthesized are always after actual ones.
@@ -146,6 +147,7 @@ struct SynthesizedExtensionAnalyzer::Implementation {
           return LeftOrder.getValue() < RightOrder.getValue();
         }
       }
+
       return false;
     }
   };
@@ -269,28 +271,31 @@ struct SynthesizedExtensionAnalyzer::Implementation {
   }
 
   std::pair<SynthesizedExtensionInfo, ExtensionMergeInfo>
-  isApplicable(ExtensionDecl *Ext, bool IsSynthesized) {
-    SynthesizedExtensionInfo Result(IsSynthesized);
+  isApplicable(ExtensionDecl *Ext, bool IsSynthesized,
+               ExtensionDecl *EnablingExt, NormalProtocolConformance *Conf) {
+    SynthesizedExtensionInfo Result(IsSynthesized, EnablingExt);
     ExtensionMergeInfo MergeInfo;
     MergeInfo.HasDocComment = !Ext->getRawComment().isEmpty();
     MergeInfo.InheritsCount = countInherits(Ext);
-    if (!Ext->isConstrainedExtension()) {
+
+    // There's (up to) two extensions here: the extension with the items that we
+    // might be merging, plus the "enabling extension", which is the route
+    // through which \c Ext itself applies, e.g. extension SomeProtocol {}
+    // extension SomeType: SomeProtocol where T: SomeProtocol {}. The former is
+    // Ext and the latter is EnablingExt/Conf. Either of these can be
+    // conditional in ways that need to be considered when merging.
+    auto conformanceIsConditional =
+        Conf && !Conf->getConditionalRequirements().empty();
+    if (!Ext->isConstrainedExtension() && !conformanceIsConditional) {
       if (IncludeUnconditional)
         Result.Ext = Ext;
       return {Result, MergeInfo};
     }
 
-    // Get the substitutions from the generic signature of
-    // the extension to the interface types of the base type's
-    // declaration.
-    auto *M = DC->getParentModule();
-    SubstitutionMap subMap;
-    if (!BaseType->isExistentialType())
-      subMap = BaseType->getContextSubstitutionMap(M, Ext);
-
-    assert(Ext->getGenericSignature() && "No generic signature.");
-    auto GenericSig = Ext->getGenericSignature();
-    for (auto Req : GenericSig->getRequirements()) {
+    auto handleRequirements = [&](SubstitutionMap subMap,
+                                  GenericSignature *GenericSig,
+                                  ArrayRef<Requirement> Reqs) {
+      for (auto Req : Reqs) {
         auto Kind = Req.getKind();
 
         // FIXME: Could do something here
@@ -307,7 +312,7 @@ struct SynthesizedExtensionAnalyzer::Implementation {
             // Substitution with interface type bases can only fail
             // if a concrete type fails to conform to a protocol.
             // In this case, just give up on the extension altogether.
-          return {Result, MergeInfo};
+            return true;
           }
         }
 
@@ -317,14 +322,14 @@ struct SynthesizedExtensionAnalyzer::Implementation {
           // FIXME: This could be more accurate; check
           // conformance instead of subtyping
           if (!canPossiblyConvertTo(First, Second, *DC))
-            return {Result, MergeInfo};
+            return true;
           else if (!isConvertibleTo(First, Second, *DC))
             MergeInfo.addRequirement(GenericSig, First, Second, Kind);
           break;
 
         case RequirementKind::SameType:
           if (!canPossiblyEqual(First, Second, *DC)) {
-            return {Result, MergeInfo};
+            return true;
           } else if (!First->isEqual(Second)) {
             MergeInfo.addRequirement(GenericSig, First, Second, Kind);
           }
@@ -334,6 +339,29 @@ struct SynthesizedExtensionAnalyzer::Implementation {
           llvm_unreachable("Handled above");
         }
       }
+      return false;
+    };
+
+    auto *M = DC->getParentModule();
+    if (Ext->isConstrainedExtension()) {
+      // Get the substitutions from the generic signature of
+      // the extension to the interface types of the base type's
+      // declaration.
+      SubstitutionMap subMap;
+      if (!BaseType->isExistentialType())
+        subMap = BaseType->getContextSubstitutionMap(M, Ext);
+
+      assert(Ext->getGenericSignature() && "No generic signature.");
+      auto GenericSig = Ext->getGenericSignature();
+      if (handleRequirements(subMap, GenericSig, GenericSig->getRequirements()))
+        return {Result, MergeInfo};
+    }
+
+    if (Conf && handleRequirements(Conf->getSubstitutions(M),
+                                   Conf->getGenericSignature(),
+                                   Conf->getConditionalRequirements()))
+      return {Result, MergeInfo};
+
     Result.Ext = Ext;
     return {Result, MergeInfo};
   }
@@ -368,7 +396,9 @@ struct SynthesizedExtensionAnalyzer::Implementation {
     for (auto *E : Target->getExtensions()) {
       if (!Options.shouldPrint(E))
         continue;
-      auto Pair = isApplicable(E, /*Synthesized*/false);
+      auto Pair = isApplicable(E, /*Synthesized*/ false,
+                               /*EnablingExt*/ nullptr,
+                               /*Conf*/ nullptr);
       if (Pair.first) {
         InfoMap->insert({E, Pair.first});
         MergeInfoMap.insert({E, Pair.second});
@@ -400,9 +430,11 @@ struct SynthesizedExtensionAnalyzer::Implementation {
     ExtensionMergeInfoMap MergeInfoMap;
     std::vector<NominalTypeDecl*> Unhandled;
 
-    auto handleExtension = [&](ExtensionDecl *E, bool Synthesized) {
+    auto handleExtension = [&](ExtensionDecl *E, bool Synthesized,
+                               ExtensionDecl *EnablingE,
+                               NormalProtocolConformance *Conf) {
       if (Options.shouldPrint(E)) {
-        auto Pair = isApplicable(E, Synthesized);
+        auto Pair = isApplicable(E, Synthesized, EnablingE, Conf);
         if (Pair.first) {
           InfoMap->insert({E, Pair.first});
           MergeInfoMap.insert({E, Pair.second});
@@ -421,7 +453,7 @@ struct SynthesizedExtensionAnalyzer::Implementation {
       NominalTypeDecl* Back = Unhandled.back();
       Unhandled.pop_back();
       for (ExtensionDecl *E : Back->getExtensions()) {
-        handleExtension(E, true);
+        handleExtension(E, true, nullptr, nullptr);
       }
       for (auto *Conf : Back->getLocalConformances()) {
         Unhandled.push_back(Conf->getProtocol());
@@ -433,11 +465,11 @@ struct SynthesizedExtensionAnalyzer::Implementation {
     }
 
     // Merge with actual extensions.
-    for (auto *E : Target->getExtensions()) {
-      handleExtension(E, false);
-      for (auto *Conf : E->getLocalConformances()) {
+    for (auto *EnablingE : Target->getExtensions()) {
+      handleExtension(EnablingE, false, nullptr, nullptr);
+      for (auto *Conf : EnablingE->getLocalConformances()) {
         for (auto E : Conf->getProtocol()->getExtensions())
-          handleExtension(E, true);
+          handleExtension(E, true, EnablingE, Conf->getRootNormalConformance());
       }
     }
 
@@ -482,9 +514,10 @@ forEachExtensionMergeGroup(MergeGroupKind Kind, ExtensionGroupOperation Fn) {
       if (Kind != Group.Kind)
         continue;
     }
-    std::vector<ExtensionAndIsSynthesized> GroupContent;
+    std::vector<ExtensionInfo> GroupContent;
     for (auto &Member : Group.Members) {
-      GroupContent.push_back({Member->Ext, Member->IsSynthesized});
+      GroupContent.push_back(
+          {Member->Ext, Member->EnablingExt, Member->IsSynthesized});
     }
     Fn(llvm::makeArrayRef(GroupContent));
   }

lib/IDE/ModuleInterfacePrinting.cpp

@@ -81,15 +81,16 @@ private:
     return OtherPrinter.printModuleRef(Mod, Name);
   }
   void printSynthesizedExtensionPre(const ExtensionDecl *ED,
-                                    const NominalTypeDecl *NTD,
+                                    TypeOrExtensionDecl Target,
                                     Optional<BracketOptions> Bracket) override {
-    return OtherPrinter.printSynthesizedExtensionPre(ED, NTD, Bracket);
+    return OtherPrinter.printSynthesizedExtensionPre(ED, Target, Bracket);
   }
 
-  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
-                                     const NominalTypeDecl *NTD,
+  void
+  printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                TypeOrExtensionDecl Target,
                                 Optional<BracketOptions> Bracket) override {
-    return OtherPrinter.printSynthesizedExtensionPost(ED, NTD, Bracket);
+    return OtherPrinter.printSynthesizedExtensionPost(ED, Target, Bracket);
   }
 
   void printStructurePre(PrintStructureKind Kind, const Decl *D) override {
@@ -589,17 +590,19 @@ void swift::ide::printSubmoduleInterface(
           if (IsTopLevelDecl) {
           // Print the part that should be merged with the type decl.
           pAnalyzer->forEachExtensionMergeGroup(
-            SynthesizedExtensionAnalyzer::MergeGroupKind::MergeableWithTypeDef,
-            [&](ArrayRef<ExtensionAndIsSynthesized> Decls){
+              SynthesizedExtensionAnalyzer::MergeGroupKind::
+                  MergeableWithTypeDef,
+              [&](ArrayRef<ExtensionInfo> Decls) {
                 for (auto ET : Decls) {
-                AdjustedOptions.BracketOptions = {ET.first, false,
-                                        Decls.back().first == ET.first, true};
-                if (ET.second)
+                  AdjustedOptions.BracketOptions = {
+                      ET.Ext, false, Decls.back().Ext == ET.Ext, true};
+                  if (ET.IsSynthesized)
                     AdjustedOptions.initForSynthesizedExtension(NTD);
-                ET.first->print(Printer, AdjustedOptions);
-                if (ET.second)
+                  ET.Ext->print(Printer, AdjustedOptions);
+                  if (ET.IsSynthesized)
                     AdjustedOptions.clearSynthesizedExtension();
-                if (AdjustedOptions.BracketOptions.shouldCloseExtension(ET.first))
+                  if (AdjustedOptions.BracketOptions.shouldCloseExtension(
+                          ET.Ext))
                     Printer << "\n";
                 }
               });
@@ -611,25 +614,33 @@ void swift::ide::printSubmoduleInterface(
 
           // Print the rest as synthesized extensions.
           pAnalyzer->forEachExtensionMergeGroup(
-            // For top-level decls, only constraint extensions are to print;
-            // Since the rest are merged into the main body.
-            IsTopLevelDecl ?
-              SynthesizedExtensionAnalyzer::MergeGroupKind::UnmergeableWithTypeDef :
+              // For top-level decls, only constraint extensions need to be
+              // printed, since the rest are merged into the main body.
+              IsTopLevelDecl ? SynthesizedExtensionAnalyzer::MergeGroupKind::
+                                   UnmergeableWithTypeDef
+                             :
                              // For sub-decls, all extensions should be printed.
                   SynthesizedExtensionAnalyzer::MergeGroupKind::All,
-            [&](ArrayRef<ExtensionAndIsSynthesized> Decls){
+              [&](ArrayRef<ExtensionInfo> Decls) {
                 for (auto ET : Decls) {
-                AdjustedOptions.BracketOptions = {ET.first,
-                  Decls.front().first == ET.first,
-                  Decls.back().first == ET.first, true};
-                if (AdjustedOptions.BracketOptions.shouldOpenExtension(ET.first))
+                  AdjustedOptions.BracketOptions = {
+                      ET.Ext, Decls.front().Ext == ET.Ext,
+                      Decls.back().Ext == ET.Ext, true};
+                  if (AdjustedOptions.BracketOptions.shouldOpenExtension(
+                          ET.Ext))
                     Printer << "\n";
-                if (ET.second)
+                  if (ET.IsSynthesized) {
+                    if (ET.EnablingExt)
+                      AdjustedOptions.initForSynthesizedExtension(
+                          ET.EnablingExt);
+                    else
                       AdjustedOptions.initForSynthesizedExtension(NTD);
-                ET.first->print(Printer, AdjustedOptions);
-                if (ET.second)
+                  }
+                  ET.Ext->print(Printer, AdjustedOptions);
+                  if (ET.IsSynthesized)
                     AdjustedOptions.clearSynthesizedExtension();
-                if (AdjustedOptions.BracketOptions.shouldCloseExtension(ET.first))
+                  if (AdjustedOptions.BracketOptions.shouldCloseExtension(
+                          ET.Ext))
                     Printer << "\n";
                 }
               });

test/SourceKit/DocSupport/Inputs/cake.swift

@@ -83,3 +83,17 @@ extension C1 : P4 {
     public func C1S1foo(a : P4) {}
   }
 }
+
+// rdar://problem/36553066
+
+public protocol P5 {
+    associatedtype Element
+}
+public protocol P6: P5 {}
+extension P6 {
+    public var null: Element? { return nil }
+}
+public struct S3<Wrapped: P5>: P5 {
+    public typealias Element = Wrapped.Element
+}
+extension S3: P6 where Wrapped: P6 {}

tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp

@@ -62,7 +62,7 @@ struct TextRange {
 
 struct TextEntity {
   const Decl *Dcl = nullptr;
-  const Decl *SynthesizeTarget = nullptr;
+  TypeOrExtensionDecl SynthesizeTarget;
   const Decl *DefaultImplementationOf = nullptr;
   StringRef Argument;
   TextRange Range;
@@ -70,26 +70,26 @@ struct TextEntity {
   std::vector<TextEntity> SubEntities;
   const bool IsSynthesizedExtension;
 
-  TextEntity(const Decl *D, const Decl *SynthesizeTarget,
-             const Decl* DefaultImplementationOf,
-             unsigned StartOffset, bool IsSynthesizedExtension)
+  TextEntity(const Decl *D, TypeOrExtensionDecl SynthesizeTarget,
+             const Decl *DefaultImplementationOf, unsigned StartOffset,
+             bool IsSynthesizedExtension)
       : Dcl(D), SynthesizeTarget(SynthesizeTarget),
         DefaultImplementationOf(DefaultImplementationOf), Range{StartOffset, 0},
         IsSynthesizedExtension(IsSynthesizedExtension) {}
 
-  TextEntity(const Decl *D, const Decl *SynthesizeTarget,
+  TextEntity(const Decl *D, TypeOrExtensionDecl SynthesizeTarget,
              const Decl *DefaultImplementationOf, TextRange TR,
-             unsigned LocOffset, bool IsSynthesizedExtension) : Dcl(D),
-             DefaultImplementationOf(DefaultImplementationOf), Range(TR),
-              LocOffset(LocOffset),
-              IsSynthesizedExtension(IsSynthesizedExtension) {}
-
-  TextEntity(const Decl *D, const Decl *SynthesizeTarget,
-             const Decl* DefaultImplementationOf, StringRef Arg,
-             TextRange TR, unsigned LocOffset, bool IsSynthesizedExtension)
-    : Dcl(D), SynthesizeTarget(SynthesizeTarget),
-      DefaultImplementationOf(DefaultImplementationOf), Argument(Arg), Range(TR),
+             unsigned LocOffset, bool IsSynthesizedExtension)
+      : Dcl(D), DefaultImplementationOf(DefaultImplementationOf), Range(TR),
         LocOffset(LocOffset), IsSynthesizedExtension(IsSynthesizedExtension) {}
+
+  TextEntity(const Decl *D, TypeOrExtensionDecl SynthesizeTarget,
+             const Decl *DefaultImplementationOf, StringRef Arg, TextRange TR,
+             unsigned LocOffset, bool IsSynthesizedExtension)
+      : Dcl(D), SynthesizeTarget(SynthesizeTarget),
+        DefaultImplementationOf(DefaultImplementationOf), Argument(Arg),
+        Range(TR), LocOffset(LocOffset),
+        IsSynthesizedExtension(IsSynthesizedExtension) {}
 };
 
 struct TextReference {
@@ -103,8 +103,8 @@ struct TextReference {
 
 class AnnotatingPrinter : public StreamPrinter {
 
-  std::pair<const ExtensionDecl *, const NominalTypeDecl *>
-    SynthesizedExtensionInfo = {nullptr, nullptr};
+  std::pair<const ExtensionDecl *, TypeOrExtensionDecl>
+      SynthesizedExtensionInfo = {nullptr, {}};
 
   typedef llvm::SmallDenseMap<ValueDecl*, ValueDecl*> DefaultImplementMap;
   llvm::SmallDenseMap<ProtocolDecl*, DefaultImplementMap> AllDefaultMaps;
@@ -175,22 +175,22 @@ public:
   }
 
   void printSynthesizedExtensionPre(const ExtensionDecl *ED,
-                                    const NominalTypeDecl *NTD,
+                                    TypeOrExtensionDecl Target,
                                     Optional<BracketOptions> Bracket) override {
     assert(!SynthesizedExtensionInfo.first);
-    SynthesizedExtensionInfo = {ED, NTD};
+    SynthesizedExtensionInfo = {ED, Target};
     if (!shouldContinuePre(ED, Bracket))
       return;
     unsigned StartOffset = OS.tell();
-    EntitiesStack.emplace_back(ED, SynthesizedExtensionInfo.second, nullptr,
-                               StartOffset, true);
+    EntitiesStack.emplace_back(ED, Target, nullptr, StartOffset, true);
   }
 
-  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
-                                     const NominalTypeDecl *NTD,
+  void
+  printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                TypeOrExtensionDecl Target,
                                 Optional<BracketOptions> Bracket) override {
     assert(SynthesizedExtensionInfo.first);
-    SynthesizedExtensionInfo = {nullptr, nullptr};
+    SynthesizedExtensionInfo = {nullptr, {}};
     if (!shouldContinuePost(ED, Bracket))
       return;
     TextEntity Entity = std::move(EntitiesStack.back());
@@ -207,8 +207,8 @@ public:
       return;
     unsigned StartOffset = OS.tell();
     initDefaultMapToUse(D);
-    const NominalTypeDecl *SynthesizedTarget = nullptr;
     // If D is declared in the extension, then the synthesized target is valid.
+    TypeOrExtensionDecl SynthesizedTarget;
     if (D->getDeclContext() == SynthesizedExtensionInfo.first)
       SynthesizedTarget = SynthesizedExtensionInfo.second;
     EntitiesStack.emplace_back(D, SynthesizedTarget,
@@ -297,10 +297,10 @@ static void initDocGenericParams(const Decl *D, DocEntityInfo &Info) {
   }
 }
 
-static bool initDocEntityInfo(const Decl *D, const Decl *SynthesizedTarget,
-                              const Decl *DefaultImplementationOf,
-                              bool IsRef, bool IsSynthesizedExtension,
-                              DocEntityInfo &Info,
+static bool initDocEntityInfo(const Decl *D,
+                              TypeOrExtensionDecl SynthesizedTarget,
+                              const Decl *DefaultImplementationOf, bool IsRef,
+                              bool IsSynthesizedExtension, DocEntityInfo &Info,
                               StringRef Arg = StringRef()) {
   if (!IsRef && D->isImplicit())
     return true;
@@ -320,9 +320,13 @@ static bool initDocEntityInfo(const Decl *D, const Decl *SynthesizedTarget,
     return false;
   }
 
-  if (IsSynthesizedExtension)
-    Info.Kind = SwiftLangSupport::getUIDForExtensionOfDecl(SynthesizedTarget);
-  else
+  auto SynthesizedTargetNTD =
+      SynthesizedTarget ? SynthesizedTarget.getBaseNominal() : nullptr;
+
+  if (IsSynthesizedExtension) {
+    Info.Kind =
+        SwiftLangSupport::getUIDForExtensionOfDecl(SynthesizedTargetNTD);
+  } else
     Info.Kind = SwiftLangSupport::getUIDForDecl(D, IsRef);
 
   if (Info.Kind.isInvalid())
@@ -335,7 +339,7 @@ static bool initDocEntityInfo(const Decl *D, const Decl *SynthesizedTarget,
       SwiftLangSupport::printUSR(VD, OS);
       if (SynthesizedTarget) {
         OS << SwiftLangSupport::SynthesizedUSRSeparator;
-        SwiftLangSupport::printUSR(dyn_cast<ValueDecl>(SynthesizedTarget), OS);
+        SwiftLangSupport::printUSR(SynthesizedTargetNTD, OS);
         {
           llvm::raw_svector_ostream OS(Info.OriginalUSR);
           SwiftLangSupport::printUSR(VD, OS);
@@ -375,8 +379,7 @@ static bool initDocEntityInfo(const Decl *D, const Decl *SynthesizedTarget,
         SecondPart = SecondPart.substr(ExtendedName.size());
         llvm::SmallString<128> UpdatedDocBuffer;
         UpdatedDocBuffer.append(FirstPart);
-        UpdatedDocBuffer.append(((const NominalTypeDecl*)SynthesizedTarget)->getName().
-                                str());
+        UpdatedDocBuffer.append(SynthesizedTargetNTD->getName().str());
         UpdatedDocBuffer.append(SecondPart);
         OS << UpdatedDocBuffer;
       } else
@@ -392,9 +395,7 @@ static bool initDocEntityInfo(const Decl *D, const Decl *SynthesizedTarget,
       llvm::raw_svector_ostream OS(Info.FullyAnnotatedDecl);
       if (SynthesizedTarget)
         SwiftLangSupport::printFullyAnnotatedSynthesizedDeclaration(
-            VD, const_cast<NominalTypeDecl *>(
-                    static_cast<const NominalTypeDecl *>(SynthesizedTarget)),
-            OS);
+            VD, SynthesizedTarget, OS);
       else
         SwiftLangSupport::printFullyAnnotatedDeclaration(VD, Type(), OS);
     }
@@ -461,13 +462,13 @@ static const TypeDecl *getTypeDeclFromType(Type Ty) {
 
 static void passInherits(const ValueDecl *D, DocInfoConsumer &Consumer) {
   DocEntityInfo EntInfo;
-  if (initDocEntityInfo(D, nullptr, nullptr, /*IsRef=*/true, false, EntInfo))
+  if (initDocEntityInfo(D, {}, nullptr, /*IsRef=*/true, false, EntInfo))
     return;
   Consumer.handleInheritsEntity(EntInfo);
 }
 static void passConforms(const ValueDecl *D, DocInfoConsumer &Consumer) {
   DocEntityInfo EntInfo;
-  if (initDocEntityInfo(D, nullptr, nullptr, /*IsRef=*/true, false, EntInfo))
+  if (initDocEntityInfo(D, {}, nullptr, /*IsRef=*/true, false, EntInfo))
     return;
   Consumer.handleConformsToEntity(EntInfo);
 }
@@ -503,7 +504,7 @@ static void passConforms(ArrayRef<ValueDecl *> Dcls,
 }
 static void passExtends(const ValueDecl *D, DocInfoConsumer &Consumer) {
   DocEntityInfo EntInfo;
-  if (initDocEntityInfo(D, nullptr, nullptr, /*IsRef=*/true, false, EntInfo))
+  if (initDocEntityInfo(D, {}, nullptr, /*IsRef=*/true, false, EntInfo))
     return;
   Consumer.handleExtendsEntity(EntInfo);
 }
@@ -516,15 +517,15 @@ static void passInheritsAndConformancesForValueDecl(const ValueDecl *VD,
                Consumer);
 }
 
-static void reportRelated(ASTContext &Ctx,
-                          const Decl *D,
-                          const Decl* SynthesizedTarget,
+static void reportRelated(ASTContext &Ctx, const Decl *D,
+                          TypeOrExtensionDecl SynthesizedTarget,
                           DocInfoConsumer &Consumer) {
   if (!D || isa<ParamDecl>(D))
     return;
   if (const auto *ED = dyn_cast<ExtensionDecl>(D)) {
     if (SynthesizedTarget) {
-      passExtends((const ValueDecl*)SynthesizedTarget, Consumer);
+      auto Extends = SynthesizedTarget.getBaseNominal();
+      passExtends(Extends, Consumer);
     } else if (Type T = ED->getExtendedType()) {
       if (auto TD = getTypeDeclFromType(T))
         passExtends(TD, Consumer);
@@ -645,7 +646,8 @@ static void reportDocEntities(ASTContext &Ctx,
       continue;
     Consumer.startSourceEntity(EntInfo);
     reportRelated(Ctx, Entity.Dcl,
-          Entity.IsSynthesizedExtension ? Entity.SynthesizeTarget : nullptr,
+                  Entity.IsSynthesizedExtension ? Entity.SynthesizeTarget
+                                                : TypeOrExtensionDecl(),
                   Consumer);
     reportDocEntities(Ctx, Entity.SubEntities, Consumer);
     reportAttributes(Ctx, Entity.Dcl, Consumer);
@@ -759,8 +761,7 @@ private:
       const TextReference &Ref = References.front();
       References = References.slice(1);
       DocEntityInfo Info;
-      if (initDocEntityInfo(Ref.Dcl, nullptr, nullptr, /*IsRef=*/true, false,
-                            Info))
+      if (initDocEntityInfo(Ref.Dcl, {}, nullptr, /*IsRef=*/true, false, Info))
         continue;
       Info.Offset = Ref.Range.Offset;
       Info.Length = Ref.Range.Length;
@@ -826,7 +827,7 @@ static void addParameters(ArrayRef<Identifier> &ArgNames,
         SM.getLocOffsetInBuffer(Lexer::getLocForEndOfToken(SM, TypeRange.End),
                                 BufferID);
       TextRange TR{ StartOffs, EndOffs-StartOffs };
-      TextEntity Param(param, nullptr, nullptr, Arg, TR, StartOffs, false);
+      TextEntity Param(param, {}, nullptr, Arg, TR, StartOffs, false);
       Ent.SubEntities.push_back(std::move(Param));
     }
   }
@@ -1020,7 +1021,8 @@ public:
       return true;
     TextRange TR = getTextRange(D->getSourceRange());
     unsigned LocOffset = getOffset(Range.getStart());
-    EntitiesStack.emplace_back(D, nullptr, nullptr, TR, LocOffset, false);
+    EntitiesStack.emplace_back(D, TypeOrExtensionDecl(), nullptr, TR, LocOffset,
+                               false);
     return true;
   }
 

tools/SourceKit/lib/SwiftLang/SwiftEditorInterfaceGen.cpp

@@ -129,20 +129,21 @@ public:
     assert(DeclUSRs.empty() && "unmatched printDeclLoc call ?");
   }
 
-
   void printSynthesizedExtensionPre(const ExtensionDecl *ED,
-                                    const NominalTypeDecl *Target,
+                                    TypeOrExtensionDecl Target,
                                     Optional<BracketOptions> Bracket) override {
     // When we start print a synthesized extension, record the target's USR.
     llvm::SmallString<64> Buf;
     llvm::raw_svector_ostream OS(Buf);
-    if (!SwiftLangSupport::printUSR(Target, OS)) {
+    auto TargetNTD = Target.getBaseNominal();
+    if (!SwiftLangSupport::printUSR(TargetNTD, OS)) {
       TargetUSR = OS.str();
     }
   }
 
-  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
-                                     const NominalTypeDecl *Target,
+  void
+  printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                TypeOrExtensionDecl Target,
                                 Optional<BracketOptions> Bracket) override {
     // When we leave a synthesized extension, clear target's USR.
     TargetUSR = "";

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -354,9 +354,9 @@ public:
                                              swift::Type BaseTy,
                                              llvm::raw_ostream &OS);
 
-  static void printFullyAnnotatedSynthesizedDeclaration(
-                                            const swift::ValueDecl *VD,
-                                            swift::NominalTypeDecl *Target,
+  static void
+  printFullyAnnotatedSynthesizedDeclaration(const swift::ValueDecl *VD,
+                                            swift::TypeOrExtensionDecl Target,
                                             llvm::raw_ostream &OS);
 
   /// Tries to resolve the path to the real file-system path. If it fails it

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -433,19 +433,20 @@ void SwiftLangSupport::printFullyAnnotatedDeclaration(const ValueDecl *VD,
   VD->print(Printer, PO);
 }
 
-void SwiftLangSupport::
-printFullyAnnotatedSynthesizedDeclaration(const swift::ValueDecl *VD,
-                                          swift::NominalTypeDecl *Target,
+void SwiftLangSupport::printFullyAnnotatedSynthesizedDeclaration(
+    const swift::ValueDecl *VD, TypeOrExtensionDecl Target,
     llvm::raw_ostream &OS) {
   // FIXME: Mutable global variable - gross!
   static llvm::SmallDenseMap<swift::ValueDecl*,
     std::unique_ptr<swift::SynthesizedExtensionAnalyzer>> TargetToAnalyzerMap;
   FullyAnnotatedDeclarationPrinter Printer(OS);
   PrintOptions PO = PrintOptions::printQuickHelpDeclaration();
-  if (TargetToAnalyzerMap.count(Target) == 0) {
+  NominalTypeDecl *TargetNTD = Target.getBaseNominal();
+
+  if (TargetToAnalyzerMap.count(TargetNTD) == 0) {
     std::unique_ptr<SynthesizedExtensionAnalyzer> Analyzer(
-      new SynthesizedExtensionAnalyzer(Target, PO));
-    TargetToAnalyzerMap.insert({Target, std::move(Analyzer)});
+        new SynthesizedExtensionAnalyzer(TargetNTD, PO));
+    TargetToAnalyzerMap.insert({TargetNTD, std::move(Analyzer)});
   }
   PO.initForSynthesizedExtension(Target);
   PO.PrintAsMember = true;

tools/swift-ide-test/swift-ide-test.cpp

@@ -1700,17 +1700,17 @@ public:
       printDeclPost(D, None);
   }
 
-
   void printSynthesizedExtensionPre(const ExtensionDecl *ED,
-                                    const NominalTypeDecl *NTD,
+                                    TypeOrExtensionDecl Target,
                                     Optional<BracketOptions> Bracket) override {
     if (Bracket.hasValue() && !Bracket.getValue().shouldOpenExtension(ED))
       return;
     OS << "<synthesized>";
   }
 
-  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
-                                     const NominalTypeDecl *NTD,
+  void
+  printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                TypeOrExtensionDecl Target,
                                 Optional<BracketOptions> Bracket) override {
     if (Bracket.hasValue() && !Bracket.getValue().shouldCloseExtension(ED))
       return;