Merge pull request #17689 from hvitved/rust/print-ast

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
+module PrintAst<shouldPrintSig/1 shouldPrint> {
-private int getOrder(PrintAstNode node) {
+  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)
+  }
+
+  private int getOrder(PrintAstNode node) {
     node =
-    rank[result](PrintAstNode n, Location loc |
+      rank[result](PrintAstNode n, string filepath, int startline, int startcolumn, int endline,
-      loc = n.getLocation()
+        int endcolumn |
+        orderBy(n, filepath, startline, startcolumn, endline, endcolumn)
       |
-      n
+        n order by filepath, startline, startcolumn, endline, endcolumn
-      order by
-        loc.getFile().getAbsolutePath(), loc.getStartLine(), loc.getStartColumn(), loc.getEndLine(),
-        loc.getEndColumn()
       )
-}
+  }
 
-/** Holds if `node` belongs to the output tree, and its property `key` has the given `value`. */
+  /** 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) {
+  query predicate nodes(PrintAstNode node, string key, string value) {
-  node.shouldBePrinted() and
-  (
     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) {
+  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 property `key` of the graph has the given `value`. */
+  /** Holds if property `key` of the graph has the given `value`. */
-query predicate graphProperties(string key, string value) {
+  query predicate graphProperties(string key, string value) {
     key = "semmle.graphKind" and value = "tree"
+  }
 }

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

@@ -6,34 +6,18 @@
 import rust
 import codeql.rust.elements.internal.generated.ParentChild
 
-private newtype TPrintAstConfiguration = TMakePrintAstConfiguration()
+signature predicate shouldPrintSig(Locatable e);
 
-/**
+module PrintAstNode<shouldPrintSig/1 shouldPrint> {
- * The hook to customize the files and functions printed by this module.
- */
-class PrintAstConfiguration extends TPrintAstConfiguration {
   /**
-   * Gets the string representation of this singleton
-   */
-  string toString() { result = "PrintAstConfiguration" }
-
-  /**
-   * Holds if the AST for `e` should be printed. By default, holds for all.
-   */
-  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)
+  private newtype TPrintAstNode = TPrintLocatable(Locatable ast) { shouldPrint(ast) }
 
-/**
+  /**
    * A node in the output tree.
    */
-class PrintAstNode extends TPrintAstNode {
+  class PrintAstNode extends TPrintAstNode {
     /**
      * Gets a textual representation of this node.
      */
@@ -45,11 +29,6 @@ class PrintAstNode extends TPrintAstNode {
      */
     abstract predicate hasChild(PrintAstNode child, int index, string label);
 
-  /**
-   * Holds if this node should be printed in the output.
-   */
-  abstract predicate shouldBePrinted();
-
     /**
      * Gets the location of this node in the source code.
      */
@@ -65,28 +44,24 @@ class PrintAstNode extends TPrintAstNode {
      * Gets the underlying AST node, if any.
      */
     abstract Locatable getAstNode();
-}
+  }
 
-private string prettyPrint(Locatable e) {
+  private string prettyPrint(Locatable e) { result = "[" + e.getPrimaryQlClasses() + "] " + e }
-  result = "[" + concat(e.getPrimaryQlClasses(), ", ") + "] " + e
-}
 
-private class Unresolved extends Locatable {
+  private class Unresolved extends Locatable {
     Unresolved() { this != this.resolve() }
-}
+  }
 
-/**
+  /**
    * A graph node representing a real Locatable node.
    */
-class PrintLocatable extends PrintAstNode, TPrintLocatable {
+  class PrintLocatable extends PrintAstNode, TPrintLocatable {
     Locatable ast;
 
     PrintLocatable() { this = TPrintLocatable(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)))
     }
@@ -94,17 +69,18 @@ class PrintLocatable extends PrintAstNode, TPrintLocatable {
     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 {
+  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))
     }
+  }
 }

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 {
+predicate shouldPrint(Locatable e) {
-  /**
-   * 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))
+  exists(Locatable parent | shouldPrint(parent) and parent = getImmediateParent(e))
-    )
-  }
 }
+
+import PrintAst<shouldPrint/1>