ASTPrinter: Add pre and post callbacks for printing synthesized extensions and call them.

This is necessary for jump to synthesized extensions in IDE.

Author: nkcsgexi

include/swift/AST/ASTPrinter.h

@@ -25,6 +25,8 @@ namespace swift {
   class TypeDecl;
   class Type;
   class Pattern;
+  class ExtensionDecl;
+  class NominalTypeDecl;
   struct PrintOptions;
 
 /// Describes the context in which a name is being printed, which
@@ -45,6 +47,7 @@ class ASTPrinter {
   unsigned PendingNewlines = 0;
   const Decl *PendingDeclPreCallback = nullptr;
   const Decl *PendingDeclLocCallback = nullptr;
+  const NominalTypeDecl *SynthesizeTarget = nullptr;
 
   void printTextImpl(StringRef Text);
 
@@ -74,6 +77,14 @@ class ASTPrinter {
   /// Called when printing the referenced name of a module.
   virtual void printModuleRef(ModuleEntity Mod, Identifier Name);
 
+  /// Called before printing a synthesized extesion.
+  virtual void printSynthesizedExtensionPre(const ExtensionDecl *ED,
+                                            const NominalTypeDecl *NTD) {}
+
+  /// Called after printing a synthesized extension.
+  virtual void printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                             const NominalTypeDecl *NTD) {}
+
   // Helper functions.
 
   ASTPrinter &operator<<(StringRef Text) {
@@ -93,6 +104,10 @@ class ASTPrinter {
     CurrentIndentation = NumSpaces;
   }
 
+  void setSynthesizedTarget(NominalTypeDecl *Target) {
+    SynthesizeTarget = Target;
+  }
+
   void printNewline() {
     PendingNewlines++;
   }

include/swift/AST/PrintOptions.h

@@ -257,6 +257,8 @@ struct PrintOptions {
 
   void initArchetypeTransformerForSynthesizedExtensions(NominalTypeDecl *D);
 
+  bool isPrintingSynthesizedExtension();
+
   void clearArchetypeTransformerForSynthesizedExtensions();
 
   /// Retrieve the print options that are suitable to print the testable interface.

lib/AST/ASTPrinter.cpp

@@ -217,6 +217,10 @@ void PrintOptions::initArchetypeTransformerForSynthesizedExtensions(NominalTypeD
   SynthesizedTarget = D;
 }
 
+bool PrintOptions::isPrintingSynthesizedExtension() {
+  return pTransformer && SynthesizedTarget;
+}
+
 void PrintOptions::clearArchetypeTransformerForSynthesizedExtensions() {
   pTransformer = nullptr;
   SynthesizedTarget = nullptr;
@@ -301,7 +305,10 @@ void ASTPrinter::printTextImpl(StringRef Text) {
   PendingDeclLocCallback = nullptr;
 
   if (PreD) {
-    printDeclPre(PreD);
+    if (SynthesizeTarget && PreD->getKind() == DeclKind::Extension)
+      printSynthesizedExtensionPre(cast<ExtensionDecl>(PreD), SynthesizeTarget);
+    else
+      printDeclPre(PreD);
   }
   if (LocD) {
     printDeclLoc(LocD);
@@ -636,9 +643,19 @@ class PrintAST : public ASTVisitor<PrintAST> {
     if (!shouldPrint(D, true))
       return false;
 
+    bool Synthesize = Options.isPrintingSynthesizedExtension() &&
+                      D->getKind() == DeclKind::Extension;
+    if (Synthesize)
+      Printer.setSynthesizedTarget(Options.SynthesizedTarget);
     Printer.callPrintDeclPre(D);
     ASTVisitor::visit(D);
-    Printer.printDeclPost(D);
+    if (Synthesize) {
+      Printer.setSynthesizedTarget(nullptr);
+      Printer.printSynthesizedExtensionPost(cast<ExtensionDecl>(D),
+                                            Options.SynthesizedTarget);
+    } else {
+      Printer.printDeclPost(D);
+    }
     return true;
   }
 
@@ -1415,6 +1432,8 @@ void PrintAST::visitImportDecl(ImportDecl *decl) {
 }
 
 void PrintAST::printSynthesizedExtension(NominalTypeDecl* Decl, ExtensionDecl* ExtDecl) {
+  Printer << "/// Synthesized extension from " <<
+    ExtDecl->getExtendedType()->getAnyNominal()->getName().str() << "\n";
   printDocumentationComment(ExtDecl);
   printAttributes(ExtDecl);
   Printer << "extension ";

lib/IDE/ModuleInterfacePrinting.cpp

@@ -68,6 +68,15 @@ class ClangCommentPrinter : public ASTPrinter {
   void printModuleRef(ModuleEntity Mod, Identifier Name) override {
     return OtherPrinter.printModuleRef(Mod, Name);
   }
+  void printSynthesizedExtensionPre(const ExtensionDecl *ED,
+                                    const NominalTypeDecl *NTD) override {
+    return OtherPrinter.printSynthesizedExtensionPre(ED, NTD);
+  }
+
+  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                     const NominalTypeDecl *NTD) override {
+    return OtherPrinter.printSynthesizedExtensionPost(ED, NTD);
+  }
 
   // Prints regular comments of the header the clang node comes from, until
   // the location of the node. Keeps track of the comments that were printed
@@ -435,9 +444,6 @@ void swift::ide::printSubmoduleInterface(
             if (!shouldPrint(ET, AdjustedOptions))
               continue;
             Printer << "\n";
-            Printer << "/// Synthesized extension from ";
-            ET->getExtendedTypeLoc().getType().print(Printer, AdjustedOptions);
-            Printer << "\n";
             ET->print(Printer, AdjustedOptions);
             Printer << "\n";
           }

test/IDE/print_synthesized_extensions.swift

@@ -1,6 +1,6 @@
 // RUN: rm -rf %t && mkdir %t
 // RUN: %target-swift-frontend -emit-module-path %t/print_synthesized_extensions.swiftmodule %s
-// RUN: %target-swift-ide-test -print-module -synthesize-extension -print-interface -module-to-print=print_synthesized_extensions -I %t -source-filename=%s | FileCheck %s
+// RUN: %target-swift-ide-test -print-module -annotate-print -synthesize-extension -print-interface -module-to-print=print_synthesized_extensions -I %t -source-filename=%s | FileCheck %s
 
 public protocol P1{
   associatedtype T1
@@ -47,24 +47,24 @@ public struct S1<T> : P1, P2 {
   }
 }
 
-// CHECK: /// Synthesized extension from P2
+// CHECK: <synthesized>/// Synthesized extension from P2
 // CHECK: extension S1 where T : P2 {
-// CHECK:     public func p2member()
-// CHECK: }
+// CHECK:     <decl:Func>public func <loc>p2member()</loc></decl>
+// CHECK: }</synthesized>
 
-// CHECK: /// Synthesized extension from P1
+// CHECK: <synthesized>/// Synthesized extension from P1
 // CHECK: extension S1 where T : P2 {
-// CHECK:     public func ef1(t: T)
-// CHECK:     public func ef2(t: S2)
-// CHECK: }
+// CHECK:     <decl:Func>public func <loc>ef1(t: T)</loc></decl>
+// CHECK:     <decl:Func>public func <loc>ef2(t: <ref:Struct>S2</ref>)</loc></decl>
+// CHECK: }</synthesized>
 
-// CHECK: /// Synthesized extension from P1
+// CHECK: <synthesized>/// Synthesized extension from P1
 // CHECK: extension S1 where T == P2, S2 : P3 {
-// CHECK:     public func ef3(t: P2)
-// CHECK:     public func ef4(t: P2)
-// CHECK: }
+// CHECK:     <decl:Func>public func <loc>ef3(t: <ref:Protocol>P2</ref>)</loc></decl>
+// CHECK:     <decl:Func>public func <loc>ef4(t: <ref:Protocol>P2</ref>)</loc></decl>
+// CHECK: }</synthesized>
 
-// CHECK: /// Synthesized extension from P1
+// CHECK: <synthesized>/// Synthesized extension from P1
 // CHECK: extension S1 where S2 : P3 {
-// CHECK:    public func ef5(t: S2)
-// CHECK: }
+// CHECK:     <decl:Func>public func <loc>ef5(t: <ref:Struct>S2</ref>)</loc></decl>
+// CHECK: }</synthesized>

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

@@ -1571,6 +1571,16 @@ class AnnotatingPrinter : public StreamPrinter {
     OS << "</decl>";
   }
 
+  void printSynthesizedExtensionPre(const ExtensionDecl *ED,
+                                    const NominalTypeDecl *NTD) override {
+    OS << "<synthesized>";
+  }
+
+  void printSynthesizedExtensionPost(const ExtensionDecl *ED,
+                                     const NominalTypeDecl *NTD) override {
+    OS << "</synthesized>";
+  }
+
   void printTypeRef(const TypeDecl *TD, Identifier Name) override {
     OS << "<ref:" << Decl::getKindName(TD->getKind()) << '>';
     StreamPrinter::printTypeRef(TD, Name);