[ASTPrinter] Enforce pairing of pre/post decl callbacks

We can have multiple printDeclPre callbacks pending (e.g top-level-code
decls), so use a vector to ensure we don't lose the earlier callbacks.

We also may end up calling printDeclPost without forcing the
corresponding printDeclPre first if the decl doesn't actually print
anything (e.g. an if-config statement when skipping those). So add a
wrapper callPrintDeclPost that can check for this and skip both
callbacks.  In theory, we could handle this case by instead making all
ast nodes go through something like shouldPrint() and making an
invariant that something will be printed if and only if shouldPrint
returns true.  However, that is not an obvious win, because it forces us
to walk all the first-level statements and decls inside a top-level-code
decl to determine if anything will be printed, and it also means we can
never make local decisions about whether something will be printed.  For
now, I've chosen to maintain flexibility by recovering from unprinted
decls.

Finally, add a bunch of assertions to try to keep callbacks sane.

Author: benlangmuir

include/swift/AST/ASTPrinter.h

@@ -47,7 +47,9 @@ enum class PrintNameContext {
 class ASTPrinter {
   unsigned CurrentIndentation = 0;
   unsigned PendingNewlines = 0;
-  const Decl *PendingDeclPreCallback = nullptr;
+  /// The queue of pending printDeclPre callbacks that will be run when we print
+  /// a non-whitespace character.
+  SmallVector<const Decl *, 4> PendingDeclPreCallbacks;
   const Decl *PendingDeclLocCallback = nullptr;
   Optional<PrintNameContext> PendingNamePreCallback;
   const NominalTypeDecl *SynthesizeTarget = nullptr;
@@ -59,19 +61,29 @@ class ASTPrinter {
 
   virtual void printText(StringRef Text) = 0;
 
+  // MARK: Callback interface.
+
   /// Called after the printer decides not to print D.
+  ///
+  /// Callers should use callAvoidPrintDeclPost().
   virtual void avoidPrintDeclPost(const Decl *D) {};
   /// Called before printing of a declaration.
+  ///
+  /// Callers should use callPrintDeclPre().
   virtual void printDeclPre(const Decl *D) {}
   /// Called before printing at the point which would be considered the location
   /// of the declaration (normally the name of the declaration).
+  ///
+  /// Callers should use callPrintDeclLoc().
   virtual void printDeclLoc(const Decl *D) {}
   /// Called after printing the name of the declaration.
   virtual void printDeclNameEndLoc(const Decl *D) {}
   /// Called after printing the name of a declaration, or in the case of
   /// functions its signature.
   virtual void printDeclNameOrSignatureEndLoc(const Decl *D) {}
   /// Called after finishing printing of a declaration.
+  ///
+  /// Callers should use callPrintDeclPost().
   virtual void printDeclPost(const Decl *D) {}
 
   /// Called before printing a type.
@@ -127,6 +139,9 @@ class ASTPrinter {
   }
 
   void setSynthesizedTarget(NominalTypeDecl *Target) {
+    assert((!SynthesizeTarget || !Target || Target == SynthesizeTarget) &&
+           "unexpected change of setSynthesizedTarget");
+    // FIXME: this can overwrite the original target with nullptr.
     SynthesizeTarget = Target;
   }
 
@@ -136,21 +151,48 @@ class ASTPrinter {
 
   virtual void printIndent();
 
+  // MARK: Callback interface wrappers that perform ASTPrinter bookkeeping.
+
   /// Schedule a \c printDeclPre callback to be called as soon as a
   /// non-whitespace character is printed.
   void callPrintDeclPre(const Decl *D) {
-    PendingDeclPreCallback = D;
+    PendingDeclPreCallbacks.emplace_back(D);
+  }
+
+  /// Make a callback to printDeclPost(), performing any necessary bookeeping.
+  void callPrintDeclPost(const Decl *D) {
+    if (!PendingDeclPreCallbacks.empty() &&
+        PendingDeclPreCallbacks.back() == D) {
+      // Nothing printed for D; skip both pre and post callbacks.
+      // Ideally we wouldn't get as far as setting up the callback if we aren't
+      // going to print anything, but currently that would mean walking the
+      // children of top-level code decls to determine.
+      PendingDeclPreCallbacks.pop_back();
+      return;
+    }
+    printDeclPost(D);
+  }
+
+  /// Make a callback to avoidPrintDeclPost(), performing any necessary
+  /// bookkeeping.
+  void callAvoidPrintDeclPost(const Decl *D) {
+    assert((PendingDeclPreCallbacks.empty() ||
+            PendingDeclPreCallbacks.back() != D) &&
+           "printDeclPre should not be called on avoided decl");
+    avoidPrintDeclPost(D);
   }
 
   /// Schedule a \c printDeclLoc callback to be called as soon as a
   /// non-whitespace character is printed.
   void callPrintDeclLoc(const Decl *D) {
+    assert(!PendingDeclLocCallback && "unexpected nested callPrintDeclLoc");
     PendingDeclLocCallback = D;
   }
 
   /// Schedule a \c printNamePre callback to be called as soon as a
   /// non-whitespace character is printed.
   void callPrintNamePre(PrintNameContext Context) {
+    assert(!PendingNamePreCallback && "unexpected nested callPrintNamePre");
     PendingNamePreCallback = Context;
   }
 

lib/AST/ASTPrinter.cpp

@@ -330,26 +330,27 @@ void ASTPrinter::printTextImpl(StringRef Text) {
     printText(Str);
     printIndent();
   }
-  
-  const Decl *PreD = PendingDeclPreCallback;
-  const Decl *LocD = PendingDeclLocCallback;
-  auto NameContext = PendingNamePreCallback;
-  PendingDeclPreCallback = nullptr;
+
+  // Get the pending callbacks and remove them from the printer. They must all
+  // be removed before calling any of them to ensure correct ordering.
+  auto pendingDeclPre = PendingDeclPreCallbacks;
+  PendingDeclPreCallbacks.clear();
+  auto *LocD = PendingDeclLocCallback;
   PendingDeclLocCallback = nullptr;
+  auto NameContext = PendingNamePreCallback;
   PendingNamePreCallback.reset();
 
-  if (PreD) {
+  // Perform pending callbacks.
+  for (const Decl *PreD : pendingDeclPre) {
     if (SynthesizeTarget && PreD->getKind() == DeclKind::Extension)
       printSynthesizedExtensionPre(cast<ExtensionDecl>(PreD), SynthesizeTarget);
     else
       printDeclPre(PreD);
   }
-  if (LocD) {
+  if (LocD)
     printDeclLoc(LocD);
-  }
-  if (NameContext) {
+  if (NameContext)
     printNamePre(*NameContext);
-  }
 
   printText(Text);
 }
@@ -702,7 +703,7 @@ class PrintAST : public ASTVisitor<PrintAST> {
       Printer.printSynthesizedExtensionPost(cast<ExtensionDecl>(D),
                                             Options.TransformContext->getNominal());
     } else {
-      Printer.printDeclPost(D);
+      Printer.callPrintDeclPost(D);
     }
     return true;
   }
@@ -1039,7 +1040,7 @@ bool swift::shouldPrint(const Decl *D, PrintOptions &Options) {
 bool PrintAST::shouldPrint(const Decl *D, bool Notify) {
   auto Result = swift::shouldPrint(D, Options);
   if (!Result && Notify)
-    Printer.avoidPrintDeclPost(D);
+    Printer.callAvoidPrintDeclPost(D);
   return Result;
 }
 
@@ -1826,9 +1827,7 @@ void PrintAST::visitParamDecl(ParamDecl *decl) {
 void PrintAST::printOneParameter(const ParamDecl *param, bool Curried,
                                  bool ArgNameIsAPIByDefault) {
   Printer.callPrintDeclPre(param);
-  defer {
-    Printer.printDeclPost(param);
-  };
+  defer { Printer.callPrintDeclPost(param); };
 
   auto printArgName = [&]() {
     // Print argument name.

lib/IDE/ModuleInterfacePrinting.cpp

@@ -420,7 +420,7 @@ void swift::ide::printSubmoduleInterface(
   auto PrintDecl = [&](Decl *D) -> bool {
     ASTPrinter &Printer = *PrinterToUse;
     if (!shouldPrint(D, AdjustedOptions)) {
-      Printer.avoidPrintDeclPost(D);
+      Printer.callAvoidPrintDeclPost(D);
       return false;
     }
     if (auto Ext = dyn_cast<ExtensionDecl>(D)) {
@@ -451,7 +451,7 @@ void swift::ide::printSubmoduleInterface(
           // Print Ext and add sub-types of Ext.
           for (auto Ext : NTD->getExtensions()) {
             if (!shouldPrint(Ext, AdjustedOptions)) {
-              Printer.avoidPrintDeclPost(Ext);
+              Printer.callAvoidPrintDeclPost(Ext);
               continue;
             }
             if (Ext->hasClangNode())
@@ -618,7 +618,7 @@ void swift::ide::printHeaderInterface(
   for (auto *D : ClangDecls) {
     ASTPrinter &Printer = *PrinterToUse;
     if (!shouldPrint(D, AdjustedOptions)) {
-      Printer.avoidPrintDeclPost(D);
+      Printer.callAvoidPrintDeclPost(D);
       continue;
     }
     if (D->print(Printer, AdjustedOptions))