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

Rust: `PrintAst` improvements

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

@@ -2,31 +2,35 @@
  * Provides queries to pretty-print a Rust AST as a graph.
  */
 
-import PrintAstNode
+import codeql.rust.printast.PrintAstNode
+
+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)
+  }
 
-cached
   private int getOrder(PrintAstNode node) {
     node =
-    rank[result](PrintAstNode n, Location loc |
-      loc = n.getLocation()
+      rank[result](PrintAstNode n, string filepath, int startline, int startcolumn, int endline,
+        int endcolumn |
+        orderBy(n, filepath, startline, startcolumn, endline, endcolumn)
       |
-      n
-      order by
-        loc.getFile().getAbsolutePath(), loc.getStartLine(), loc.getStartColumn(), loc.getEndLine(),
-        loc.getEndColumn()
+        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) {
-  node.shouldBePrinted() and
-  (
     key = "semmle.label" and value = node.toString()
     or
     key = "semmle.order" and value = getOrder(node).toString()
     or
     value = node.getProperty(key)
-  )
   }
 
   /**
@@ -34,8 +38,6 @@ query predicate nodes(PrintAstNode node, string key, string value) {
    * 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
@@ -47,3 +49,4 @@ query predicate edges(PrintAstNode source, PrintAstNode target, string key, stri
   query predicate graphProperties(string key, string value) {
     key = "semmle.graphKind" and value = "tree"
   }
+}

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

@@ -6,29 +6,13 @@
 import rust
 import codeql.rust.elements.internal.generated.ParentChild
 
-private newtype TPrintAstConfiguration = TMakePrintAstConfiguration()
-
-/**
- * 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) }
+signature predicate shouldPrintSig(Locatable e);
 
+module PrintAstNode<shouldPrintSig/1 shouldPrint> {
   /**
    * 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.
@@ -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.
      */
@@ -67,9 +46,7 @@ class PrintAstNode extends TPrintAstNode {
     abstract Locatable getAstNode();
   }
 
-private string prettyPrint(Locatable e) {
-  result = "[" + concat(e.getPrimaryQlClasses(), ", ") + "] " + e
-}
+  private string prettyPrint(Locatable e) { result = "[" + e.getPrimaryQlClasses() + "] " + e }
 
   private class Unresolved extends Locatable {
     Unresolved() { this != this.resolve() }
@@ -85,8 +62,6 @@ class PrintLocatable extends PrintAstNode, TPrintLocatable {
 
     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)))
     }
@@ -108,3 +83,4 @@ class PrintUnresolved extends PrintLocatable {
       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
-    (
+predicate shouldPrint(Locatable e) {
   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>