Rust: PrintAst improvements

rust/ql/lib/codeql/rust/printast/PrintAst.qll

@@ -2,48 +2,51 @@
  * Provides queries to pretty-print a Rust AST as a graph.
  */
 
-import PrintAstNode
+import codeql.rust.printast.PrintAstNode
 
-cached
-private int getOrder(PrintAstNode node) {
-  node =
-    rank[result](PrintAstNode n, Location loc |
-      loc = n.getLocation()
-    |
-      n
-      order by
-        loc.getFile().getAbsolutePath(), loc.getStartLine(), loc.getStartColumn(), loc.getEndLine(),
-        loc.getEndColumn()
-    )
-}
+module PrintAst<shouldPrintSig/1 shouldPrint> {
+  import PrintAstNode<shouldPrint/1>
+
+  pragma[nomagic]
+  private predicate orderBy(
+    PrintAstNode n, string filepath, int startline, int startcolumn, int endline, int endcolumn
+  ) {
+    n.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn)
+  }
 
-/** 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) {
-  node.shouldBePrinted() and
-  (
+  private int getOrder(PrintAstNode node) {
+    node =
+      rank[result](PrintAstNode n, string filepath, int startline, int startcolumn, int endline,
+        int endcolumn |
+        orderBy(n, filepath, startline, startcolumn, endline, endcolumn)
+      |
+        n order by filepath, startline, startcolumn, endline, endcolumn
+      )
+  }
+
+  /** 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) {
     key = "semmle.label" and value = node.toString()
     or
     key = "semmle.order" and value = getOrder(node).toString()
     or
     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) {
-  source.shouldBePrinted() and
-  target.shouldBePrinted() and
-  exists(int index, string accessor | source.hasChild(target, index, accessor) |
-    key = "semmle.label" and value = accessor
-    or
-    key = "semmle.order" and value = index.toString()
-  )
-}
+  /**
+   * 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) {
+    exists(int index, string accessor | source.hasChild(target, index, accessor) |
+      key = "semmle.label" and value = accessor
+      or
+      key = "semmle.order" and value = index.toString()
+    )
+  }
 
-/** Holds if property `key` of the graph has the given `value`. */
-query predicate graphProperties(string key, string value) {
-  key = "semmle.graphKind" and value = "tree"
+  /** Holds if property `key` of the graph has the given `value`. */
+  query predicate graphProperties(string key, string value) {
+    key = "semmle.graphKind" and value = "tree"
+  }
 }

rust/ql/lib/codeql/rust/printast/PrintAstNode.qll

@@ -6,105 +6,81 @@
 import rust
 import codeql.rust.elements.internal.generated.ParentChild
 
-private newtype TPrintAstConfiguration = TMakePrintAstConfiguration()
+signature predicate shouldPrintSig(Locatable e);
 
-/**
- * The hook to customize the files and functions printed by this module.
- */
-class PrintAstConfiguration extends TPrintAstConfiguration {
+module PrintAstNode<shouldPrintSig/1 shouldPrint> {
   /**
-   * Gets the string representation of this singleton
+   * An AST node that should be printed.
    */
-  string toString() { result = "PrintAstConfiguration" }
+  private newtype TPrintAstNode = TPrintLocatable(Locatable ast) { shouldPrint(ast) }
 
   /**
-   * Holds if the AST for `e` should be printed. By default, holds for all.
+   * A node in the output tree.
    */
-  predicate shouldPrint(Locatable e) { any() }
-}
-
-private predicate shouldPrint(Locatable e) { any(PrintAstConfiguration config).shouldPrint(e) }
-
-/**
- * An AST node that should be printed.
- */
-private newtype TPrintAstNode = TPrintLocatable(Locatable ast)
-
-/**
- * A node in the output tree.
- */
-class PrintAstNode extends TPrintAstNode {
-  /**
-   * Gets a textual representation of this node.
-   */
-  abstract string toString();
-
-  /**
-   * Gets the child node at index `index`. Child indices must be unique,
-   * but need not be contiguous.
-   */
-  abstract predicate hasChild(PrintAstNode child, int index, string label);
-
-  /**
-   * Holds if this node should be printed in the output.
-   */
-  abstract predicate shouldBePrinted();
+  class PrintAstNode extends TPrintAstNode {
+    /**
+     * Gets a textual representation of this node.
+     */
+    abstract string toString();
+
+    /**
+     * Gets the child node at index `index`. Child indices must be unique,
+     * but need not be contiguous.
+     */
+    abstract predicate hasChild(PrintAstNode child, int index, string label);
+
+    /**
+     * Gets the location of this node in the source code.
+     */
+    abstract Location getLocation();
+
+    /**
+     * Gets the value of an additional property of this node, where the name of
+     * the property is `key`.
+     */
+    string getProperty(string key) { none() }
+
+    /**
+     * Gets the underlying AST node, if any.
+     */
+    abstract Locatable getAstNode();
+  }
 
-  /**
-   * Gets the location of this node in the source code.
-   */
-  abstract Location getLocation();
+  private string prettyPrint(Locatable e) { result = "[" + e.getPrimaryQlClasses() + "] " + e }
 
-  /**
-   * Gets the value of an additional property of this node, where the name of
-   * the property is `key`.
-   */
-  string getProperty(string key) { none() }
+  private class Unresolved extends Locatable {
+    Unresolved() { this != this.resolve() }
+  }
 
   /**
-   * Gets the underlying AST node, if any.
+   * A graph node representing a real Locatable node.
    */
-  abstract Locatable getAstNode();
-}
+  class PrintLocatable extends PrintAstNode, TPrintLocatable {
+    Locatable ast;
 
-private string prettyPrint(Locatable e) {
-  result = "[" + concat(e.getPrimaryQlClasses(), ", ") + "] " + e
-}
+    PrintLocatable() { this = TPrintLocatable(ast) }
 
-private class Unresolved extends Locatable {
-  Unresolved() { this != this.resolve() }
-}
+    override string toString() { result = prettyPrint(ast) }
 
-/**
- * A graph node representing a real Locatable node.
- */
-class PrintLocatable extends PrintAstNode, TPrintLocatable {
-  Locatable ast;
+    override predicate hasChild(PrintAstNode child, int index, string label) {
+      child = TPrintLocatable(any(Locatable c | c = getChildAndAccessor(ast, index, label)))
+    }
 
-  PrintLocatable() { this = TPrintLocatable(ast) }
+    final override Locatable getAstNode() { result = ast }
 
-  override string toString() { result = prettyPrint(ast) }
-
-  final override predicate shouldBePrinted() { shouldPrint(ast) }
-
-  override predicate hasChild(PrintAstNode child, int index, string label) {
-    child = TPrintLocatable(any(Locatable c | c = getChildAndAccessor(ast, index, label)))
+    final override Location getLocation() { result = ast.getLocation() }
   }
 
-  final override Locatable getAstNode() { result = ast }
-
-  final override Location getLocation() { result = ast.getLocation() }
-}
-
-/**
- * A specialization of graph node for "unresolved" children, that is nodes in
- * the parallel conversion AST.
- */
-class PrintUnresolved extends PrintLocatable {
-  override Unresolved ast;
+  /**
+   * A specialization of graph node for "unresolved" children, that is nodes in
+   * the parallel conversion AST.
+   */
+  class PrintUnresolved extends PrintLocatable {
+    override Unresolved ast;
 
-  override predicate hasChild(PrintAstNode child, int index, string label) {
-    // only print immediate unresolved children from the "parallel" AST
-    child = TPrintLocatable(getImmediateChildAndAccessor(ast, index, label).(Unresolved))
+    override predicate hasChild(PrintAstNode child, int index, string label) {
+      // only print immediate unresolved children from the "parallel" AST
+      child = TPrintLocatable(getImmediateChildAndAccessor(ast, index, label).(Unresolved))
+    }
   }
 }

rust/ql/src/queries/ide-contextual-queries/printAst.ql

@@ -17,17 +17,10 @@ import codeql.rust.elements.internal.generated.ParentChild
  */
 external string selectedSourceFile();
 
-class PrintAstConfigurationOverride extends PrintAstConfiguration {
-  /**
-   * Holds if the location matches the selected file in the VS Code extension and
-   * the element is `e`.
-   */
-  override predicate shouldPrint(Locatable e) {
-    super.shouldPrint(e) and
-    (
-      e.getFile() = getFileBySourceArchiveName(selectedSourceFile())
-      or
-      exists(Locatable parent | this.shouldPrint(parent) and parent = getImmediateParent(e))
-    )
-  }
+predicate shouldPrint(Locatable e) {
+  e.getFile() = getFileBySourceArchiveName(selectedSourceFile())
+  or
+  exists(Locatable parent | shouldPrint(parent) and parent = getImmediateParent(e))
 }
+
+import PrintAst<shouldPrint/1>