From 6ae75e5920687b3877993c319ab79eb04dd5cfcf Mon Sep 17 00:00:00 2001
From: Asger F <asgerf@github.com>
Date: Thu, 25 Jun 2026 16:55:40 +0200
Subject: [PATCH] unified: Add PrintAst query

---
 unified/ql/lib/codeql/IDEContextual.qll       | 16 ++++
 unified/ql/lib/codeql/unified/printAst.qll    | 96 +++++++++++++++++++
 .../ql/lib/ide-contextual-queries/printAst.ql | 27 ++++++
 3 files changed, 139 insertions(+)
 create mode 100644 unified/ql/lib/codeql/IDEContextual.qll
 create mode 100644 unified/ql/lib/codeql/unified/printAst.qll
 create mode 100644 unified/ql/lib/ide-contextual-queries/printAst.ql

diff --git a/unified/ql/lib/codeql/IDEContextual.qll b/unified/ql/lib/codeql/IDEContextual.qll
new file mode 100644
index 00000000000..3b8486b4526
--- /dev/null
+++ b/unified/ql/lib/codeql/IDEContextual.qll
@@ -0,0 +1,16 @@
+/**
+ * Provides shared predicates related to contextual queries in the code viewer.
+ */
+
+private import codeql.files.FileSystem
+private import codeql.util.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) {
+  result = IdeContextual<File>::getFileBySourceArchiveName(name)
+}
diff --git a/unified/ql/lib/codeql/unified/printAst.qll b/unified/ql/lib/codeql/unified/printAst.qll
new file mode 100644
index 00000000000..93ff11f5c8b
--- /dev/null
+++ b/unified/ql/lib/codeql/unified/printAst.qll
@@ -0,0 +1,96 @@
+/** Provides a configurable query for printing AST nodes */
+
+private import unified
+
+/**
+ * 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 TriviaToken }
+
+  /**
+   * Holds if the given edge should be printed.
+   */
+  predicate shouldPrintAstEdge(AstNode parent, string edgeName, AstNode child) {
+    exists(string name, int i |
+      child = PrintAst::getChild(parent, name, i) and
+      (if i = -1 then edgeName = name else edgeName = name + "(" + i + ")")
+    )
+  }
+}
+
+private predicate shouldPrintNode(AstNode n) {
+  any(PrintAstConfiguration config).shouldPrintNode(n)
+}
+
+private predicate shouldPrintAstEdge(AstNode parent, string edgeName, AstNode child) {
+  any(PrintAstConfiguration config).shouldPrintAstEdge(parent, edgeName, child) and
+  shouldPrintNode(parent) and
+  shouldPrintNode(child)
+}
+
+/**
+ * Get an alias for the predicate `name` to use for ordering purposes, to control where
+ * in the list of children it should appear.
+ */
+private string reorderName1(string name) { name = "getModifier" and result = "00_getModifier" }
+
+bindingset[name]
+private string reorderName(string name) {
+  result = reorderName1(name)
+  or
+  not exists(reorderName1(name)) and
+  result = name
+}
+
+class PrintAstNode extends AstNode {
+  final int getOrder() {
+    this =
+      rank[result](AstNode parent, AstNode child, string name, int i |
+        child = PrintAst::getChild(parent, name, i)
+      |
+        child order by reorderName(name), i
+      )
+  }
+
+  final string getProperty(string key) {
+    key = "semmle.label" and
+    result = this.toString()
+    or
+    key = "semmle.order" and result = this.getOrder().toString()
+  }
+}
+
+/**
+ * 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) {
+  shouldPrintNode(node) and
+  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) {
+  key = "semmle.label" and
+  shouldPrintAstEdge(source, value, target)
+  or
+  key = "semmle.order" and
+  shouldPrintAstEdge(source, _, target) 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"
+}
diff --git a/unified/ql/lib/ide-contextual-queries/printAst.ql b/unified/ql/lib/ide-contextual-queries/printAst.ql
new file mode 100644
index 00000000000..3622babe07f
--- /dev/null
+++ b/unified/ql/lib/ide-contextual-queries/printAst.ql
@@ -0,0 +1,27 @@
+/**
+ * @name Print AST
+ * @description Produces a representation of a file's Abstract Syntax Tree.
+ *              This query is used by the VS Code extension.
+ * @id unified/print-ast
+ * @kind graph
+ * @tags ide-contextual-queries/print-ast
+ */
+
+private import codeql.IDEContextual
+private import unified
+private import codeql.unified.printAst
+
+/**
+ * The source file to generate an AST from.
+ */
+external string selectedSourceFile();
+
+/**
+ * A configuration that only prints nodes in the selected source file.
+ */
+class Cfg extends PrintAstConfiguration {
+  override predicate shouldPrintNode(AstNode n) {
+    super.shouldPrintNode(n) and
+    n.getLocation().getFile() = getFileBySourceArchiveName(selectedSourceFile())
+  }
+}