Merge pull request #264 from smowton/smowton/feature/printast-restrict-files

PrintAst: improve support for restricting subsets of the AST to print

Author: smowton

ql/src/printAst.ql

@@ -20,7 +20,9 @@ external string selectedSourceFile();
  * Hook to customize the functions printed by this query.
  */
 class Cfg extends PrintAstConfiguration {
-  override predicate shouldPrintFunction(FuncDef func) {
-    func.getFile() = getEncodedFile(selectedSourceFile())
-  }
+  override predicate shouldPrintFunction(FuncDecl func) { shouldPrintFile(func.getFile()) }
+
+  override predicate shouldPrintFile(File file) { file = getEncodedFile(selectedSourceFile()) }
+
+  override predicate shouldPrintComments(File file) { none() }
 }

ql/src/semmle/go/PrintAst.ql

@@ -12,5 +12,9 @@ import PrintAst
  * Hook to customize the functions printed by this query.
  */
 class Cfg extends PrintAstConfiguration {
-  override predicate shouldPrintFunction(FuncDef func) { any() }
+  override predicate shouldPrintFunction(FuncDecl func) { any() }
+
+  override predicate shouldPrintFile(File file) { any() }
+
+  override predicate shouldPrintComments(File file) { any() }
 }

ql/src/semmle/go/PrintAst.qll

@@ -5,7 +5,11 @@
 import go
 
 /**
- * Hook to customize the functions printed by this module.
+ * Hook to customize the files and functions printed by this module.
+ *
+ * For an AstNode to be printed, it always requires `shouldPrintFile(f)` to hold
+ * for its containing file `f`, and additionally requires `shouldPrintFunction(fun)`
+ * to hold if it is, or is a child of, function `fun`.
  */
 class PrintAstConfiguration extends string {
   /**
@@ -17,20 +21,48 @@ class PrintAstConfiguration extends string {
    * Holds if the AST for `func` should be printed. By default, holds for all
    * functions.
    */
-  predicate shouldPrintFunction(FuncDef func) { any() }
+  predicate shouldPrintFunction(FuncDecl func) { any() }
+
+  /**
+   * Holds if the AST for `file` should be printed. By default, holds for all
+   * files.
+   */
+  predicate shouldPrintFile(File file) { any() }
+
+  /**
+   * Holds if the AST for `file` should include comments. By default, holds for all
+   * files.
+   */
+  predicate shouldPrintComments(File file) { any() }
 }
 
 private predicate shouldPrintFunction(FuncDef func) {
   exists(PrintAstConfiguration config | config.shouldPrintFunction(func))
 }
 
+private predicate shouldPrintFile(File file) {
+  exists(PrintAstConfiguration config | config.shouldPrintFile(file))
+}
+
+private predicate shouldPrintComments(File file) {
+  exists(PrintAstConfiguration config | config.shouldPrintComments(file))
+}
+
+private FuncDecl getEnclosingFunctionDecl(AstNode n) { result = n.getParent*() }
+
 /**
  * An AST node that should be printed.
  */
 private newtype TPrintAstNode =
   TAstNode(AstNode ast) {
-    // Do print ast nodes without an enclosing function, e.g. file headers
-    forall(FuncDef f | f = ast.getEnclosingFunction() | shouldPrintFunction(f))
+    shouldPrintFile(ast.getFile()) and
+    // Do print ast nodes without an enclosing function, e.g. file headers, that are not otherwise excluded
+    forall(FuncDecl f | f = getEnclosingFunctionDecl(ast) | shouldPrintFunction(f)) and
+    (
+      shouldPrintComments(ast.getFile())
+      or
+      not ast instanceof Comment and not ast instanceof CommentGroup
+    )
   }
 
 /**
@@ -149,6 +181,19 @@ class ExprNode extends BaseAstNode {
 class FileNode extends BaseAstNode {
   override File ast;
 
+  private string getRelativePath() { result = ast.getRelativePath() }
+
+  private int getSortOrder() {
+    rank[result](FileNode fn | any() | fn order by fn.getRelativePath()) = this
+  }
+
+  override string getProperty(string key) {
+    result = super.getProperty(key)
+    or
+    key = "semmle.order" and
+    result = getSortOrder().toString()
+  }
+
   /**
    * Gets the string representation of this File. Note explicitly using a relative path
    * like this rather than absolute as per default for the File class is a workaround for

ql/test/library-tests/semmle/go/PrintAst/PrintAst.expected

@@ -1,3 +1,7 @@
+go.mod:
+#    0| [GoModFile] library-tests/semmle/go/PrintAst/go.mod
+#    1|   0: [GoModModuleLine] go.mod module line
+#    3|   1: [GoModGoLine] go.mod go line
 input.go:
 #    0| [File] library-tests/semmle/go/PrintAst/input.go
 #    5|   0: [CommentGroup] comment group
@@ -592,3 +596,55 @@ input.go:
 #  145|             Type = []int
 #  145|         1: [BlockStmt] block statement
 #    1|   18: [Ident] main
+other.go:
+#    0| [File] library-tests/semmle/go/PrintAst/other.go
+#    3|   0: [FuncDecl] function declaration
+#    3|     0: [FunctionName, Ident] main
+#    3|         Type = func() 
+#    3|     1: [FuncTypeExpr] function type
+#    3|     2: [BlockStmt] block statement
+#    5|   1: [FuncDecl] function declaration
+#    5|     0: [FunctionName, Ident] f
+#    5|         Type = func() 
+#    5|     1: [FuncTypeExpr] function type
+#    5|     2: [BlockStmt] block statement
+#    6|   2: [FuncDecl] function declaration
+#    6|     0: [FunctionName, Ident] g
+#    6|         Type = func() 
+#    6|     1: [FuncTypeExpr] function type
+#    6|     2: [BlockStmt] block statement
+#    8|   3: [FuncDecl] function declaration
+#    8|     0: [FunctionName, Ident] hasNested
+#    8|         Type = func() 
+#    8|     1: [FuncTypeExpr] function type
+#    8|     2: [BlockStmt] block statement
+#   10|       0: [DefineStmt] ... := ...
+#   10|         0: [Ident, VariableName] myNested
+#   10|             Type = func() int
+#   10|         1: [FuncLit] function literal
+#   10|             Type = func() int
+#   10|           0: [FuncTypeExpr] function type
+#   10|               Type = func() int
+#   10|             0: [ResultVariableDecl] result variable declaration
+#   10|               0: [Ident, TypeName] int
+#   10|                   Type = int
+#   10|           1: [BlockStmt] block statement
+#   10|             0: [ReturnStmt] return statement
+#   10|               0: [IntLit] 1
+#   10|                   Type = int
+#   10|                   Value = [IntLit] 1
+#   11|       1: [ExprStmt] expression statement
+#   11|         0: [CallExpr] call to myNested
+#   11|             Type = int
+#   11|           0: [Ident, VariableName] myNested
+#   11|               Type = func() int
+#   15|   4: [VarDecl] variable declaration
+#   15|     0: [ValueSpec] value declaration specifier
+#   15|       0: [Ident, VariableName] x
+#   15|           Type = int
+#   15|       1: [Ident, TypeName] int
+#   15|           Type = int
+#   15|       2: [IntLit] 0
+#   15|           Type = int
+#   15|           Value = [IntLit] 0
+#    1|   5: [Ident] main