QL: Add convenience methods for aggregates

Cargo.lock

@@ -589,7 +589,7 @@ dependencies = [
 [[package]]
 name = "tree-sitter-ql"
 version = "0.19.0"
-source = "git+https://github.com/tausbn/tree-sitter-ql.git?rev=a58a56d7ee47bfd75728916a49b576af29238a1d#a58a56d7ee47bfd75728916a49b576af29238a1d"
+source = "git+https://github.com/tausbn/tree-sitter-ql.git?rev=5f3e790557ad6d6612e269808168e7a8f6413dc0#5f3e790557ad6d6612e269808168e7a8f6413dc0"
 dependencies = [
  "cc",
  "tree-sitter",

extractor/Cargo.toml

@@ -10,7 +10,7 @@ edition = "2018"
 flate2 = "1.0"
 node-types = { path = "../node-types" }
 tree-sitter = "0.19"
-tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "a58a56d7ee47bfd75728916a49b576af29238a1d" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "5f3e790557ad6d6612e269808168e7a8f6413dc0" }
 clap = "2.33"
 tracing = "0.1"
 tracing-subscriber = { version = "0.2", features = ["env-filter"] }

generator/Cargo.toml

@@ -10,4 +10,4 @@ edition = "2018"
 node-types = { path = "../node-types" }
 tracing = "0.1"
 tracing-subscriber = { version = "0.2", features = ["env-filter"] }
-tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "a58a56d7ee47bfd75728916a49b576af29238a1d" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "5f3e790557ad6d6612e269808168e7a8f6413dc0" }

ql/src/codeql_ql/ast/Ast.qll

@@ -383,7 +383,7 @@ class Quantifier extends TQuantifier, Formula {
   Formula getFormula() { toGenerated(result) = quant.getFormula() }
 
   /** Gets the expression of this quantifier, if it is of the expression only form of an exists. */
-  Expr getExpr() { toGenerated(result) = quant.getChild(1) }
+  Expr getExpr() { toGenerated(result) = quant.getExpr() }
 
   /** Holds if this is the expression only form of an exists quantifier. */
   predicate hasExpr() { exists(getExpr()) }

ql/src/codeql_ql/ast/internal/TreeSitter.qll

@@ -344,11 +344,15 @@ module Generated {
   class ExprAggregateBody extends @expr_aggregate_body, AstNode {
     override string getAPrimaryQlClass() { result = "ExprAggregateBody" }
 
-    override Location getLocation() { expr_aggregate_body_def(this, result) }
+    override Location getLocation() { expr_aggregate_body_def(this, _, result) }
 
-    AstNode getChild(int i) { expr_aggregate_body_child(this, i, result) }
+    AsExprs getAsExprs() { expr_aggregate_body_def(this, result, _) }
 
-    override AstNode getAFieldOrChild() { expr_aggregate_body_child(this, _, result) }
+    OrderBys getOrderBys() { expr_aggregate_body_order_bys(this, result) }
+
+    override AstNode getAFieldOrChild() {
+      expr_aggregate_body_def(this, result, _) or expr_aggregate_body_order_bys(this, result)
+    }
   }
 
   class ExprAnnotation extends @expr_annotation, AstNode {
@@ -392,9 +396,20 @@ module Generated {
 
     override Location getLocation() { full_aggregate_body_def(this, result) }
 
-    AstNode getChild(int i) { full_aggregate_body_child(this, i, result) }
+    AsExprs getAsExprs() { full_aggregate_body_as_exprs(this, result) }
 
-    override AstNode getAFieldOrChild() { full_aggregate_body_child(this, _, result) }
+    AstNode getGuard() { full_aggregate_body_guard(this, result) }
+
+    OrderBys getOrderBys() { full_aggregate_body_order_bys(this, result) }
+
+    VarDecl getChild(int i) { full_aggregate_body_child(this, i, result) }
+
+    override AstNode getAFieldOrChild() {
+      full_aggregate_body_as_exprs(this, result) or
+      full_aggregate_body_guard(this, result) or
+      full_aggregate_body_order_bys(this, result) or
+      full_aggregate_body_child(this, _, result)
+    }
   }
 
   class HigherOrderTerm extends @higher_order_term, AstNode {
@@ -746,6 +761,8 @@ module Generated {
 
     override Location getLocation() { quantified_def(this, result) }
 
+    AstNode getExpr() { quantified_expr(this, result) }
+
     AstNode getFormula() { quantified_formula(this, result) }
 
     AstNode getRange() { quantified_range(this, result) }
@@ -753,6 +770,7 @@ module Generated {
     AstNode getChild(int i) { quantified_child(this, i, result) }
 
     override AstNode getAFieldOrChild() {
+      quantified_expr(this, result) or
       quantified_formula(this, result) or
       quantified_range(this, result) or
       quantified_child(this, _, result)
@@ -902,9 +920,17 @@ module Generated {
 
     override Location getLocation() { unqual_agg_body_def(this, result) }
 
-    AstNode getChild(int i) { unqual_agg_body_child(this, i, result) }
+    AstNode getAsExprs(int i) { unqual_agg_body_as_exprs(this, i, result) }
 
-    override AstNode getAFieldOrChild() { unqual_agg_body_child(this, _, result) }
+    AstNode getGuard() { unqual_agg_body_guard(this, result) }
+
+    VarDecl getChild(int i) { unqual_agg_body_child(this, i, result) }
+
+    override AstNode getAFieldOrChild() {
+      unqual_agg_body_as_exprs(this, _, result) or
+      unqual_agg_body_guard(this, result) or
+      unqual_agg_body_child(this, _, result)
+    }
   }
 
   class VarDecl extends @var_decl, AstNode {

ql/src/ql.dbscheme

@@ -300,17 +300,14 @@ disjunction_def(
   int loc: @location ref
 );
 
-@expr_aggregate_body_child_type = @as_exprs | @order_bys
-
-#keyset[expr_aggregate_body, index]
-expr_aggregate_body_child(
-  int expr_aggregate_body: @expr_aggregate_body ref,
-  int index: int ref,
-  unique int child: @expr_aggregate_body_child_type ref
+expr_aggregate_body_order_bys(
+  unique int expr_aggregate_body: @expr_aggregate_body ref,
+  unique int orderBys: @order_bys ref
 );
 
 expr_aggregate_body_def(
   unique int id: @expr_aggregate_body,
+  int as_exprs: @as_exprs ref,
   int loc: @location ref
 );
 
@@ -330,13 +327,28 @@ field_def(
   int loc: @location ref
 );
 
-@full_aggregate_body_child_type = @add_expr | @aggregate | @as_exprs | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @order_bys | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @var_decl | @variable
+full_aggregate_body_as_exprs(
+  unique int full_aggregate_body: @full_aggregate_body ref,
+  unique int asExprs: @as_exprs ref
+);
+
+@full_aggregate_body_guard_type = @add_expr | @aggregate | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @variable
+
+full_aggregate_body_guard(
+  unique int full_aggregate_body: @full_aggregate_body ref,
+  unique int guard: @full_aggregate_body_guard_type ref
+);
+
+full_aggregate_body_order_bys(
+  unique int full_aggregate_body: @full_aggregate_body ref,
+  unique int orderBys: @order_bys ref
+);
 
 #keyset[full_aggregate_body, index]
 full_aggregate_body_child(
   int full_aggregate_body: @full_aggregate_body ref,
   int index: int ref,
-  unique int child: @full_aggregate_body_child_type ref
+  unique int child: @var_decl ref
 );
 
 full_aggregate_body_def(
@@ -661,6 +673,13 @@ qualified_expr_def(
   int loc: @location ref
 );
 
+@quantified_expr_type = @add_expr | @aggregate | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @variable
+
+quantified_expr(
+  unique int quantified: @quantified ref,
+  unique int expr: @quantified_expr_type ref
+);
+
 @quantified_formula_type = @add_expr | @aggregate | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @variable
 
 quantified_formula(
@@ -675,7 +694,7 @@ quantified_range(
   unique int range: @quantified_range_type ref
 );
 
-@quantified_child_type = @add_expr | @aggregate | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @token_quantifier | @unary_expr | @var_decl | @variable
+@quantified_child_type = @token_quantifier | @var_decl
 
 #keyset[quantified, index]
 quantified_child(
@@ -797,13 +816,27 @@ unary_expr_def(
   int loc: @location ref
 );
 
-@unqual_agg_body_child_type = @add_expr | @aggregate | @as_exprs | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @var_decl | @variable
+@unqual_agg_body_asExprs_type = @as_exprs | @reserved_word
+
+#keyset[unqual_agg_body, index]
+unqual_agg_body_as_exprs(
+  int unqual_agg_body: @unqual_agg_body ref,
+  int index: int ref,
+  unique int asExprs: @unqual_agg_body_asExprs_type ref
+);
+
+@unqual_agg_body_guard_type = @add_expr | @aggregate | @call_or_unqual_agg_expr | @comp_term | @conjunction | @disjunction | @expr_annotation | @if_term | @implication | @in_expr | @instance_of | @literal | @mul_expr | @negation | @par_expr | @prefix_cast | @qualified_expr | @quantified | @range | @set_literal | @special_call | @super_ref | @unary_expr | @variable
+
+unqual_agg_body_guard(
+  unique int unqual_agg_body: @unqual_agg_body ref,
+  unique int guard: @unqual_agg_body_guard_type ref
+);
 
 #keyset[unqual_agg_body, index]
 unqual_agg_body_child(
   int unqual_agg_body: @unqual_agg_body ref,
   int index: int ref,
-  unique int child: @unqual_agg_body_child_type ref
+  unique int child: @var_decl ref
 );
 
 unqual_agg_body_def(