[AST] Add minimal JSON support to ASTDumper.

This only takes the existing AST information and writes it as JSON
instead of S-expressions. Since many of these fields are stringified,
they're not ideal for the kind of analysis clients of the JSON format
would want to do. A future commit will update these values to use a
more structured representation.

Author: allevato

include/swift/AST/DiagnosticsFrontend.def

@@ -589,5 +589,8 @@ ERROR(experimental_not_supported_in_production,none,
       "experimental feature '%0' cannot be enabled in production compiler",
       (StringRef))
 
+ERROR(zlib_not_supported,none,
+       "this compiler was not built with zlib compression support enabled", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/SourceFile.h

@@ -686,6 +686,9 @@ class SourceFile final : public FileUnit {
   SWIFT_DEBUG_DUMP;
   void dump(raw_ostream &os, bool parseIfNeeded = false) const;
 
+  /// Dumps this source file's AST in JSON format to the given output stream.
+  void dumpJSON(raw_ostream &os) const;
+
   /// Pretty-print the contents of this source file.
   ///
   /// \param Printer The AST printer used for printing the contents.

include/swift/Frontend/FrontendOptions.h

@@ -433,6 +433,16 @@ class FrontendOptions {
   /// -dump-scope-maps.
   SmallVector<std::pair<unsigned, unsigned>, 2> DumpScopeMapLocations;
 
+  /// The possible output formats supported for dumping ASTs.
+  enum class ASTFormat {
+    Default,  ///< S-expressions for debugging
+    JSON,     ///< Structured JSON for analysis
+    JSONZlib, ///< Like JSON, but zlib-compressed
+  };
+
+  /// The output format generated by the `-dump-ast` flag.
+  ASTFormat DumpASTFormat = ASTFormat::Default;
+
   /// Determines whether the static or shared resource folder is used.
   /// When set to `true`, the default resource folder will be set to
   /// '.../lib/swift', otherwise '.../lib/swift_static'.

include/swift/Option/Options.td

@@ -1339,6 +1339,10 @@ def emit_parse : Flag<["-"], "emit-parse">, Alias<dump_parse>,
 def dump_ast : Flag<["-"], "dump-ast">,
   HelpText<"Parse and type-check input file(s) and dump AST(s)">, ModeOpt,
   Flags<[FrontendOption, NoInteractiveOption, DoesNotAffectIncrementalBuild]>;
+def dump_ast_format : Separate<["-"], "dump-ast-format">,
+  HelpText<"Desired format for -dump-ast output ('default', 'json', or 'json-zlib')">,
+  MetaVarName<"<format>">,
+  Flags<[FrontendOption, NoInteractiveOption, DoesNotAffectIncrementalBuild]>;
 def emit_ast : Flag<["-"], "emit-ast">, Alias<dump_ast>,
   Flags<[FrontendOption, NoInteractiveOption, DoesNotAffectIncrementalBuild]>;
 def dump_scope_maps : Separate<["-"], "dump-scope-maps">,

lib/AST/ASTDumper.cpp

@@ -32,13 +32,16 @@
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/QuotedString.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/Basic/SourceLoc.h"
+#include "swift/Basic/SourceManager.h"
 #include "swift/Basic/StringExtras.h"
 #include "clang/AST/Type.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/JSON.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/raw_ostream.h"
@@ -760,6 +763,111 @@ namespace {
     bool isParsable() const override { return false; }
   };
 
+  /// Implements JSON formatted output for `-ast-dump`.
+  class JSONWriter : public PrintWriterBase {
+    llvm::json::OStream OS;
+    std::vector<bool> InObjectStack;
+
+  public:
+    JSONWriter(raw_ostream &os, unsigned indent = 0) : OS(os, indent) {}
+
+    void printRecArbitrary(std::function<void(Label)> body,
+                           Label label) override {
+      // The label is ignored if we're not printing inside an object (meaning
+      // we must be in an array).
+      if (InObjectStack.back()) {
+        OS.attributeBegin(label.text());
+        body(Label::optional(""));
+        OS.attributeEnd();
+      } else {
+        body(label);
+      }
+    }
+
+    void printRecRange(std::function<void(Label)> body, Label label) override {
+      printListArbitrary([&]{ body(label); }, label);
+    }
+
+    void printListArbitrary(std::function<void()> body, Label label) override {
+      OS.attributeBegin(label.text());
+      OS.arrayBegin();
+      InObjectStack.push_back(false);
+      body();
+      InObjectStack.pop_back();
+      OS.arrayEnd();
+      OS.attributeEnd();
+    }
+
+    void printHead(StringRef name, TerminalColor color, Label label) override {
+      OS.objectBegin();
+      InObjectStack.push_back(true);
+      OS.attribute(label.empty() ? "_kind" : label.text(), name);
+    }
+
+    void printFoot() override {
+      InObjectStack.pop_back();
+      OS.objectEnd();
+    }
+
+    void printFieldRaw(std::function<void(llvm::raw_ostream &)> body,
+                       Label label, TerminalColor color) override {
+      std::string value;
+      llvm::raw_string_ostream SOS(value);
+      body(SOS);
+      // The label is ignored if we're not printing inside an object (meaning
+      // we must be in an array).
+      if (InObjectStack.back()) {
+        OS.attribute(label.text(), value);
+      } else {
+        OS.value(value);
+      }
+    }
+
+    void printFieldQuotedRaw(std::function<void(llvm::raw_ostream &)> body,
+                             Label label, TerminalColor color) override {
+      // No need to do special quoting for complex values; the JSON output
+      // stream will do this for us.
+      printFieldRaw(body, label, color);
+    }
+
+    void printFlagRaw(std::function<void(llvm::raw_ostream &)> body,
+                      TerminalColor color) override {
+      std::string flag;
+      llvm::raw_string_ostream SOS(flag);
+      body(SOS);
+      OS.attribute(flag, true);
+    }
+
+    void printSourceLoc(const SourceLoc L, const ASTContext *Ctx,
+                        Label label) override {
+      // For compactness, we only print source ranges in JSON, since they
+      // provide a superset of this information.
+    }
+
+    void printSourceRange(const SourceRange R,
+                          const ASTContext *Ctx) override {
+      OS.attributeBegin("range");
+      OS.objectBegin();
+
+      SourceManager &srcMgr = Ctx->SourceMgr;
+      unsigned startBufferID = srcMgr.findBufferContainingLoc(R.Start);
+      unsigned startOffset = srcMgr.getLocOffsetInBuffer(R.Start,
+                                                         startBufferID);
+      OS.attribute("start", startOffset);
+
+      unsigned endBufferID = srcMgr.findBufferContainingLoc(R.End);
+      unsigned endOffset = srcMgr.getLocOffsetInBuffer(R.End, endBufferID);
+      OS.attribute("end", endOffset);
+
+      OS.objectEnd();
+      OS.attributeEnd();
+    }
+
+    bool hasNonStandardOutput() const override { return true; }
+
+    bool isParsable() const override { return true; }
+  };
+
   /// PrintBase - Base type for recursive structured dumps of AST nodes.
   ///
   /// Please keep direct I/O, especially of structural elements like
@@ -2228,6 +2336,11 @@ void SourceFile::dump(llvm::raw_ostream &OS, bool parseIfNeeded) const {
   llvm::errs() << '\n';
 }
 
+void SourceFile::dumpJSON(llvm::raw_ostream &OS) const {
+  JSONWriter writer(OS, /*indent*/ 2);
+  PrintDecl(writer, /*parseIfNeeded*/ true).visitSourceFile(*this);
+}
+
 void Pattern::dump() const {
   dump(llvm::errs());
 }

lib/Driver/ToolChains.cpp

@@ -350,6 +350,7 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_enable_builtin_module);
   inputArgs.AddLastArg(arguments, options::OPT_compiler_assertions);
   inputArgs.AddLastArg(arguments, options::OPT_load_pass_plugin_EQ);
+  inputArgs.AddLastArg(arguments, options::OPT_dump_ast_format);
 
   // Pass on any build config options
   inputArgs.AddAllArgs(arguments, options::OPT_D);

lib/Frontend/ArgsToFrontendOptionsConverter.cpp

@@ -30,6 +30,7 @@
 #include "llvm/Option/Arg.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Option/Option.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LineIterator.h"
@@ -204,6 +205,21 @@ bool ArgsToFrontendOptionsConverter::convert(
 
   computeDumpScopeMapLocations();
 
+  // Ensure that the compiler was built with zlib support if it was the
+  // requested AST format.
+  if (const Arg *A = Args.getLastArg(OPT_dump_ast_format)) {
+    Opts.DumpASTFormat =
+        llvm::StringSwitch<FrontendOptions::ASTFormat>(A->getValue())
+            .Case("json", FrontendOptions::ASTFormat::JSON)
+            .Case("json-zlib", FrontendOptions::ASTFormat::JSONZlib)
+            .Default(FrontendOptions::ASTFormat::Default);
+    if (Opts.DumpASTFormat == FrontendOptions::ASTFormat::JSONZlib &&
+        !llvm::compression::zlib::isAvailable()) {
+      Diags.diagnose(SourceLoc(), diag::zlib_not_supported);
+      return true;
+    }
+  }
+
   std::optional<FrontendInputsAndOutputs> inputsAndOutputs =
       ArgsToFrontendInputsConverter(Diags, Args).convert(buffers);
 

lib/FrontendTool/FrontendTool.cpp

@@ -82,6 +82,7 @@
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Option/OptTable.h"
 #include "llvm/Option/Option.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
@@ -463,6 +464,28 @@ getPrimaryOrMainSourceFile(const CompilerInstance &Instance) {
 /// Dumps the AST of all available primary source files. If corresponding output
 /// files were specified, use them; otherwise, dump the AST to stdout.
 static bool dumpAST(CompilerInstance &Instance) {
+  const FrontendOptions &opts = Instance.getInvocation().getFrontendOptions();
+  auto dumpAST = [&](SourceFile *SF, raw_ostream &out) {
+    switch (opts.DumpASTFormat) {
+    case FrontendOptions::ASTFormat::Default:
+      SF->dump(out, /*parseIfNeeded*/ true);
+      break;
+    case FrontendOptions::ASTFormat::JSON:
+      SF->dumpJSON(out);
+      break;
+    case FrontendOptions::ASTFormat::JSONZlib:
+      std::string jsonText;
+      llvm::raw_string_ostream jsonTextStream(jsonText);
+      SF->dumpJSON(jsonTextStream);
+
+      SmallVector<uint8_t, 0> compressed;
+      llvm::compression::zlib::compress(llvm::arrayRefFromStringRef(jsonText),
+                                        compressed);
+      out << llvm::toStringRef(compressed);
+      break;
+    }
+  };
+
   auto primaryFiles = Instance.getPrimarySourceFiles();
   if (!primaryFiles.empty()) {
     for (SourceFile *sourceFile: primaryFiles) {
@@ -471,7 +494,7 @@ static bool dumpAST(CompilerInstance &Instance) {
       if (withOutputPath(Instance.getASTContext().Diags,
                          Instance.getOutputBackend(), OutputFilename,
                          [&](raw_ostream &out) -> bool {
-                           sourceFile->dump(out, /*parseIfNeeded*/ true);
+                           dumpAST(sourceFile, out);
                            return false;
                          }))
         return true;
@@ -480,7 +503,7 @@ static bool dumpAST(CompilerInstance &Instance) {
     // Some invocations don't have primary files. In that case, we default to
     // looking for the main file and dumping it to `stdout`.
     auto &SF = getPrimaryOrMainSourceFile(Instance);
-    SF.dump(llvm::outs(), /*parseIfNeeded*/ true);
+    dumpAST(&SF, llvm::outs());
   }
   return Instance.getASTContext().hadError();
 }