Add 'print-ast' query

ql/lib/codeql/kaleidoscope/ideContextual/IDEContextual.qll

@@ -0,0 +1,19 @@
+private import codeql.files.FileSystem
+
+/**
+ * Returns an appropriately encoded version of a filename `name`
+ * passed by the VS Code extension in order to coincide with the
+ * output of `.getFile()` on locatable entities.
+ */
+cached
+File getFileBySourceArchiveName(string name) {
+  // The name provided for a file in the source archive by the VS Code extension
+  // has some differences from the absolute path in the database:
+  // 1. colons are replaced by underscores
+  // 2. there's a leading slash, even for Windows paths: "C:/foo/bar" ->
+  //    "/C_/foo/bar"
+  // 3. double slashes in UNC prefixes are replaced with a single slash
+  // We can handle 2 and 3 together by unconditionally adding a leading slash
+  // before replacing double slashes.
+  name = ("/" + result.getAbsolutePath().replaceAll(":", "_")).replaceAll("//", "/")
+}

ql/lib/codeql/kaleidoscope/ideContextual/printAstGenerated.qll

@@ -0,0 +1,109 @@
+/**
+ * Provides queries to pretty-print an Kaleidoscope abstract syntax tree as a graph.
+ *
+ * This representation is based on the TreeSitter auto-generated AST.
+ *
+ * By default, this will print the AST for all nodes in the database. To change
+ * this behavior, extend `PrintASTConfiguration` and override `shouldPrintNode`
+ * to hold for only the AST nodes you wish to view.
+ */
+
+import codeql.kaleidoscope.ast.internal.TreeSitter::Kaleidoscope
+private import codeql.Locations
+
+/**
+ * The query can extend this class to control which nodes are printed.
+ */
+class PrintAstConfiguration extends string {
+  PrintAstConfiguration() { this = "PrintAstConfiguration" }
+
+  /**
+   * Holds if the given node should be printed.
+   */
+  predicate shouldPrintNode(AstNode n) {
+    not n instanceof LineComment and
+    not n instanceof ReservedWord
+  }
+}
+
+/**
+ * Gets the `i`th child of parent.
+ * The ordering is location based and pretty arbitrary.
+ */
+AstNode getAstChild(PrintAstNode parent, int i) {
+  result =
+    rank[i](AstNode child, Location l |
+      child.getParent() = parent and
+      child.getLocation() = l
+    |
+      child
+      order by
+        l.getStartLine(), l.getStartColumn(), l.getEndColumn(), l.getEndLine(), child.toString()
+    )
+}
+
+/**
+ * A node in the output tree.
+ */
+class PrintAstNode extends AstNode {
+  PrintAstNode() { shouldPrintNode(this) }
+
+  string getProperty(string key) {
+    key = "semmle.label" and
+    result = "[" + concat(this.getAPrimaryQlClass(), ", ") + "] " + this.toString()
+    or
+    key = "semmle.order" and
+    result =
+      any(int i |
+        this =
+          rank[i](PrintAstNode p, Location l, File f |
+            l = p.getLocation() and
+            f = l.getFile()
+          |
+            p order by f.getBaseName(), f.getAbsolutePath(), l.getStartLine(), l.getStartColumn()
+          )
+      ).toString()
+  }
+
+  /**
+   * Gets the child node that is accessed using the predicate `edgeName`.
+   */
+  PrintAstNode getChild(string edgeName) {
+    exists(int i |
+      result = getAstChild(this, i) and
+      edgeName = i.toString()
+    )
+  }
+}
+
+private predicate shouldPrintNode(AstNode n) {
+  exists(PrintAstConfiguration config | config.shouldPrintNode(n))
+}
+
+/**
+ * Holds if `node` belongs to the output tree, and its property `key` has the
+ * given `value`.
+ */
+query predicate nodes(PrintAstNode node, string key, string value) { value = node.getProperty(key) }
+
+/**
+ * Holds if `target` is a child of `source` in the AST, and property `key` of
+ * the edge has the given `value`.
+ */
+query predicate edges(PrintAstNode source, PrintAstNode target, string key, string value) {
+  target = source.getChild(_) and
+  (
+    key = "semmle.label" and
+    value = strictconcat(string name | source.getChild(name) = target | name, "/")
+    or
+    key = "semmle.order" and
+    value = target.getProperty("semmle.order")
+  )
+}
+
+/**
+ * Holds if property `key` of the graph has the given `value`.
+ */
+query predicate graphProperties(string key, string value) {
+  key = "semmle.graphKind" and value = "tree"
+}

ql/lib/ide-contextual-queries/printAst.ql

@@ -0,0 +1,24 @@
+/**
+ * @name Print AST
+ * @description Produces a representation of a file's Abstract Syntax Tree.
+ *              This query is used by the VS Code extension.
+ * @id kaleidoscope/print-ast
+ * @kind graph
+ * @tags ide-contextual-queries/print-ast
+ */
+
+import codeql.kaleidoscope.ideContextual.printAstGenerated
+import codeql.kaleidoscope.ideContextual.IDEContextual
+
+/**
+ * Gets the source file to generate an AST from.
+ */
+external string selectedSourceFile();
+
+// Overrides the configuration to print only nodes in the selected source file.
+class Cfg extends PrintAstConfiguration {
+  override predicate shouldPrintNode(AstNode n) {
+    super.shouldPrintNode(n) and
+    n.getLocation().getFile() = getFileBySourceArchiveName(selectedSourceFile())
+  }
+}