From cb49c1ec24e925499732303ce7714fc6e0683afd Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 14:04:23 +0000
Subject: [PATCH 01/13] switch printAst query to use new pretty AST layer

---
 ql/src/codeql_ql/ast/Ast.qll                  |   6 +
 ql/src/codeql_ql/printAstAst.qll              | 118 ++++++++++++++++++
 .../{printAst.qll => printAstGenerated.qll}   |   0
 ql/src/ide-contextual-queries/printAst.ql     |   4 +-
 4 files changed, 127 insertions(+), 1 deletion(-)
 create mode 100644 ql/src/codeql_ql/printAstAst.qll
 rename ql/src/codeql_ql/{printAst.qll => printAstGenerated.qll} (100%)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 75ca484c730..21a750bfc61 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -6,6 +6,10 @@ class AstNode extends TAstNode {
   string toString() { result = "ASTNode" }
 
   Location getLocation() { result = toGenerated(this).getLocation() }
+
+  AstNode getParent() { toGenerated(result) = toGenerated(this).getParent() }
+
+  string getAPrimaryQlClass() { result = "???" }
 }
 
 /**
@@ -20,4 +24,6 @@ class ClasslessPredicate extends TClasslessPredicate, AstNode {
   predicate isPrivate() {
     member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
   }
+
+  override string getAPrimaryQlClass() { result = "ClasslessPredicate" }
 }
diff --git a/ql/src/codeql_ql/printAstAst.qll b/ql/src/codeql_ql/printAstAst.qll
new file mode 100644
index 00000000000..db26b2a629d
--- /dev/null
+++ b/ql/src/codeql_ql/printAstAst.qll
@@ -0,0 +1,118 @@
+/**
+ * Provides queries to pretty-print a Ruby abstract syntax tree as a graph.
+ *
+ * By default, this will print the AST for all nodes in the database. To change
+ * this behavior, extend `PrintASTConfiguration` and override `shouldPrintNode`
+ * to hold for only the AST nodes you wish to view.
+ */
+
+import ast.Ast
+private import codeql.Locations
+
+/**
+ * The query can extend this class to control which nodes are printed.
+ */
+class PrintAstConfiguration extends string {
+  PrintAstConfiguration() { this = "PrintAstConfiguration" }
+
+  /**
+   * Holds if the given node should be printed.
+   */
+  predicate shouldPrintNode(AstNode n) { any() }
+}
+
+/**
+ * Gets the `i`th child of parent.
+ * The ordering is location based and pretty arbitary.
+ */
+AstNode getAstChild(PrintAstNode parent, int i) {
+  parent.shouldPrint() and
+  result =
+    rank[i](AstNode child, Location l |
+      child.getParent() = parent and
+      child.getLocation() = l
+    |
+      child
+      order by
+        l.getStartLine(), l.getStartColumn(), l.getEndColumn(), l.getEndLine(), child.toString()
+    )
+}
+
+/**
+ * A node in the output tree.
+ */
+class PrintAstNode extends AstNode {
+  string getProperty(string key) {
+    this.shouldPrint() and
+    (
+      key = "semmle.label" and
+      result = "[" + concat(this.getAPrimaryQlClass(), ", ") + "] " + this.toString()
+      or
+      key = "semmle.order" and
+      result =
+        any(int i |
+          this =
+            rank[i](AstNode p, Location l, File f |
+              l = p.getLocation() and
+              f = l.getFile()
+            |
+              p order by f.getBaseName(), f.getAbsolutePath(), l.getStartLine(), l.getStartColumn()
+            )
+        ).toString()
+    )
+  }
+
+  /**
+   * Holds if this node should be printed in the output. By default, all nodes
+   * are printed, but the query can override
+   * `PrintAstConfiguration.shouldPrintNode` to filter the output.
+   */
+  predicate shouldPrint() { shouldPrintNode(this) }
+
+  /**
+   * Gets the child node that is accessed using the predicate `edgeName`.
+   */
+  PrintAstNode getChild(string edgeName) {
+    exists(int i |
+      result = getAstChild(this, i) and
+      edgeName = i.toString()
+    )
+  }
+}
+
+private predicate shouldPrintNode(AstNode n) {
+  exists(PrintAstConfiguration config | config.shouldPrintNode(n))
+}
+
+/**
+ * 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.shouldPrint() and
+  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.shouldPrint() and
+  target.shouldPrint() and
+  target = source.getChild(_) and
+  (
+    key = "semmle.label" and
+    value = strictconcat(string name | source.getChild(name) = target | name, "/")
+    or
+    key = "semmle.order" and
+    value = target.getProperty("semmle.order")
+  )
+}
+
+/**
+ * Holds if property `key` of the graph has the given `value`.
+ */
+query predicate graphProperties(string key, string value) {
+  key = "semmle.graphKind" and value = "tree"
+}
diff --git a/ql/src/codeql_ql/printAst.qll b/ql/src/codeql_ql/printAstGenerated.qll
similarity index 100%
rename from ql/src/codeql_ql/printAst.qll
rename to ql/src/codeql_ql/printAstGenerated.qll
diff --git a/ql/src/ide-contextual-queries/printAst.ql b/ql/src/ide-contextual-queries/printAst.ql
index da0dd6ecd87..1fd07db49d5 100644
--- a/ql/src/ide-contextual-queries/printAst.ql
+++ b/ql/src/ide-contextual-queries/printAst.ql
@@ -7,7 +7,9 @@
  * @tags ide-contextual-queries/print-ast
  */
 
-import codeql_ql.printAst
+// Switch between the below two to switch between generated and pretty AST.
+// import codeql_ql.printAstGenerated
+import codeql_ql.printAstAst
 import codeql.IDEContextual
 
 /**

From 8fa3a425a5dbca1970920c61eb21c7ad52ba07b7 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 14:18:25 +0000
Subject: [PATCH 02/13] slightly more AST stuff

---
 ql/src/codeql_ql/ast/Ast.qll               | 33 ++++++++++++++++++++++
 ql/src/codeql_ql/ast/internal/AstNodes.qll | 12 ++++++--
 ql/src/ide-contextual-queries/printAst.ql  |  4 +--
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 21a750bfc61..e161366b7f7 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -26,4 +26,37 @@ class ClasslessPredicate extends TClasslessPredicate, AstNode {
   }
 
   override string getAPrimaryQlClass() { result = "ClasslessPredicate" }
+
+  /**
+   * Gets the `i`th parameter of the predicate.
+   */
+  VarDecl getParameter(int i) {
+    toGenerated(result) =
+      rank[i](Generated::VarDecl decl, int index | decl = pred.getChild(index) | decl order by index)
+  }
+
+  /**
+   * Gets the body of the predicate.
+   */
+  Body getBody() { toGenerated(result) = pred.getChild(_) }
+}
+
+/**
+ * A variable declaration, with a type and a name.
+ */
+class VarDecl extends TVarDecl, AstNode {
+  Generated::VarDecl var;
+
+  VarDecl() { this = TVarDecl(var) }
+  // TODO: Type and name getters.
+}
+
+/**
+ * The body of a predicate.
+ */
+class Body extends TBody, AstNode {
+  Generated::Body body;
+
+  Body() { this = TBody(body) }
+  // TODO: Children.
 }
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 0213b2d6fad..5199d0649ea 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -5,9 +5,17 @@ cached
 newtype TAstNode =
   TClasslessPredicate(Generated::ModuleMember member, Generated::ClasslessPredicate pred) {
     pred.getParent() = member
-  }
+  } or
+  TVarDecl(Generated::VarDecl decl) or
+  TBody(Generated::Body body)
 
 /**
  * Gets the underlying TreeSitter entity for a given AST node.
  */
-Generated::AstNode toGenerated(AST::AstNode n) { n = TClasslessPredicate(_, result) }
+Generated::AstNode toGenerated(AST::AstNode n) {
+  n = TClasslessPredicate(_, result)
+  or
+  n = TVarDecl(result)
+  or
+  n = TBody(result)
+}
diff --git a/ql/src/ide-contextual-queries/printAst.ql b/ql/src/ide-contextual-queries/printAst.ql
index 1fd07db49d5..ddffb099796 100644
--- a/ql/src/ide-contextual-queries/printAst.ql
+++ b/ql/src/ide-contextual-queries/printAst.ql
@@ -8,8 +8,8 @@
  */
 
 // Switch between the below two to switch between generated and pretty AST.
-// import codeql_ql.printAstGenerated
-import codeql_ql.printAstAst
+import codeql_ql.printAstGenerated
+// import codeql_ql.printAstAst
 import codeql.IDEContextual
 
 /**

From 7a4a4e9cdd9114d9aa5f0cc0e4413023b96314ba Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 15:01:55 +0000
Subject: [PATCH 03/13] implement more AST cases

---
 ql/src/codeql_ql/ast/Ast.qll               | 113 ++++++++++++++++++++-
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  11 +-
 2 files changed, 122 insertions(+), 2 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index e161366b7f7..1202af4dffd 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -16,6 +16,7 @@ class AstNode extends TAstNode {
  * A classless predicate.
  */
 class ClasslessPredicate extends TClasslessPredicate, AstNode {
+  // TODO: Make super class for predicate things. (classless, class predicate, charpred)
   Generated::ModuleMember member;
   Generated::ClasslessPredicate pred;
 
@@ -39,6 +40,54 @@ class ClasslessPredicate extends TClasslessPredicate, AstNode {
    * Gets the body of the predicate.
    */
   Body getBody() { toGenerated(result) = pred.getChild(_) }
+
+  /**
+   * Gets the name of the predicate
+   */
+  string getName() { result = pred.getName().getValue() }
+}
+
+/**
+ * A predicate in a class.
+ */
+class ClassPredicate extends TClassPredicate, AstNode {
+  Generated::MemberPredicate pred;
+
+  ClassPredicate() { this = TClassPredicate(pred) }
+
+  /**
+   * Gets the name of the predicate.
+   */
+  string getName() { result = pred.getName().getValue() }
+
+  /**
+   * Gets the body containing the implementation of the predicate.
+   */
+  Body getBody() { toGenerated(result) = pred.getChild(_) }
+
+  // TODO: ReturnType.
+  override string getAPrimaryQlClass() { result = "ClassPredicate" }
+
+  override Class getParent() { result.getAClassPredicate() = this }
+}
+
+/**
+ * A characteristic predicate of a class.
+ */
+class CharPred extends TCharPred, AstNode {
+  Generated::Charpred pred;
+
+  CharPred() { this = TCharPred(pred) }
+
+  override string getAPrimaryQlClass() { result = "CharPred" }
+
+  /*
+   * Gets the body of the predicate.
+   * TODO: The body is just directly an Expr in the generated AST. Wait for Expr type to appear.
+   * Body getBody() { toGenerated(result) = pred.getChild(_) }MemberPredicate
+   */
+
+  override Class getParent() { result.getCharPred() = this }
 }
 
 /**
@@ -48,7 +97,67 @@ class VarDecl extends TVarDecl, AstNode {
   Generated::VarDecl var;
 
   VarDecl() { this = TVarDecl(var) }
-  // TODO: Type and name getters.
+
+  /**
+   * Gets the name for this variable declaration.
+   */
+  string getName() { result = var.getChild(_).(Generated::VarName).getChild().getValue() }
+
+  override string getAPrimaryQlClass() { result = "VarDecl" }
+
+  override AstNode getParent() {
+    result = super.getParent()
+    or
+    result.(Class).getAField() = this
+  }
+  // TODO: Getter for the Type.
+}
+
+/**
+ * A QL class.
+ */
+class Class extends TClass, AstNode {
+  Generated::Dataclass cls;
+
+  Class() { this = TClass(cls) }
+
+  override string getAPrimaryQlClass() { result = "Class" }
+
+  /**
+   * Gets the name of the class.
+   */
+  string getName() { result = cls.getName().getValue() }
+
+  /**
+   * Gets the charateristic predicate for this class.
+   */
+  CharPred getCharPred() {
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(0).(Generated::Charpred)
+  }
+
+  /**
+   * Gets a predicate in this class.
+   */
+  ClassPredicate getAClassPredicate() {
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(0)
+  }
+
+  /**
+   * Gets predicate `name` implemented in this class.
+   */
+  ClassPredicate getClassPredicate(string name) {
+    result = getAClassPredicate() and
+    result.getName() = name
+  }
+
+  /**
+   * Gets a field in this class.
+   */
+  VarDecl getAField() {
+    toGenerated(result) =
+      cls.getChild(_).(Generated::ClassMember).getChild(0).(Generated::Field).getChild()
+  }
+  // TODO: extends, modifiers.
 }
 
 /**
@@ -58,5 +167,7 @@ class Body extends TBody, AstNode {
   Generated::Body body;
 
   Body() { this = TBody(body) }
+
+  override string getAPrimaryQlClass() { result = "Body" }
   // TODO: Children.
 }
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 5199d0649ea..49775112e94 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -7,7 +7,10 @@ newtype TAstNode =
     pred.getParent() = member
   } or
   TVarDecl(Generated::VarDecl decl) or
-  TBody(Generated::Body body)
+  TBody(Generated::Body body) or
+  TClass(Generated::Dataclass dc) or
+  TCharPred(Generated::Charpred pred) or
+  TClassPredicate(Generated::MemberPredicate pred)
 
 /**
  * Gets the underlying TreeSitter entity for a given AST node.
@@ -18,4 +21,10 @@ Generated::AstNode toGenerated(AST::AstNode n) {
   n = TVarDecl(result)
   or
   n = TBody(result)
+  or
+  n = TClass(result)
+  or
+  n = TCharPred(result)
+  or
+  n = TClassPredicate(result)
 }

From 94c1321e4336fd586c5446f05e7a25e42eb10c42 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Wed, 26 May 2021 16:21:01 +0000
Subject: [PATCH 04/13] WIP formulas and expressions

Joint work with shati-patel.
---
 ql/src/codeql_ql/ast/Ast.qll               | 65 ++++++++++++++++++++++
 ql/src/codeql_ql/ast/internal/AstNodes.qll | 22 +++++++-
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 1202af4dffd..91528697ab1 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -171,3 +171,68 @@ class Body extends TBody, AstNode {
   override string getAPrimaryQlClass() { result = "Body" }
   // TODO: Children.
 }
+
+/** A formula, such as `x = 6 and y < 5`. */
+abstract class Formula extends AstNode { }
+
+/** An `and` formula, with 2 or more operands. */
+class Conjunction extends TConjunction, AstNode, Formula {
+  Generated::Conjunction conj;
+
+  Conjunction() { this = TConjunction(conj) }
+
+  override string getAPrimaryQlClass() { result = "Conjunction" }
+
+  /** Gets an operand to this formula. */
+  Formula getAnOperand() { toGenerated(result) in [conj.getLeft(), conj.getRight()] }
+}
+
+/** An `or` formula, with 2 or more operands. */
+class Disjunction extends TDisjunction, AstNode {
+  Generated::Disjunction disj;
+
+  Disjunction() { this = TDisjunction(disj) }
+
+  override string getAPrimaryQlClass() { result = "Disjunction" }
+
+  /** Gets an operand to this formula. */
+  Formula getAnOperand() { toGenerated(result) in [disj.getLeft(), disj.getRight()] }
+}
+
+class ComparisonOp extends TComparisonOp, AstNode {
+  Generated::Compop op;
+
+  ComparisonOp() {this = TComparisonOp(op)}
+
+
+}
+
+class ComparisonFormula extends TComparisonFormula, Formula {
+  Expr getLeftOperand() {none()}
+  Expr getRightOperand() {none()}
+  Expr getAnOperand() {none()}
+  ComparisonOp getOperator() {none()}
+  //ComparisonSymbol getSymbol() {none()}
+}
+
+/** An expression, such as `x+4`. */
+abstract class Expr extends AstNode {}
+
+/** A function symbol, such as `+` or `*`. */
+class FunctionSymbol extends string {
+  FunctionSymbol() { this = "+" or this = "-" or this = "*" or this = "/" or this = "%" }
+}
+
+/** A binary operation, such as `x+3` or `y/2` */
+abstract class BinOpExpr extends Expr {}
+
+class AddExpr extends TAddExpr, BinOpExpr {
+  Generated::AddExpr addexpr;
+
+  AddExpr() {this = TAddExpr(addexpr)}
+  Expr getLeftOperand() { toGenerated(result) = addexpr.getLeft() }
+  Expr getRightOperand() { toGenerated(result) = addexpr.getRight() }
+  Expr getAnOperand() { result = getLeftOperand() or result = getRightOperand() }
+  
+  FunctionSymbol getOperator() { result = addexpr.getChild().getValue() }
+}
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 49775112e94..27115f7bcff 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -10,12 +10,32 @@ newtype TAstNode =
   TBody(Generated::Body body) or
   TClass(Generated::Dataclass dc) or
   TCharPred(Generated::Charpred pred) or
-  TClassPredicate(Generated::MemberPredicate pred)
+  TClassPredicate(Generated::MemberPredicate pred) or
+  TDisjunction(Generated::Disjunction disj) or
+  TConjunction(Generated::Conjunction conj) or
+  TComparisonFormula(Generated::CompTerm comp) or
+  TComparisonOp(Generated::Compop op) or
+  TAddExpr(Generated::AddExpr addexp)
+
+
+class TBinOpExpr = TAddExpr; // Can this work?
+
+Generated::AstNode toGeneratedFormula(AST::AstNode n) {
+  n = TConjunction(result) or
+  n = TDisjunction(result) or
+  n = TComparisonFormula(result)
+}
+
+Generated::AstNode toGeneratedExpr(AST::AstNode n) { n = TAddExpr(result) }
 
 /**
  * Gets the underlying TreeSitter entity for a given AST node.
  */
 Generated::AstNode toGenerated(AST::AstNode n) {
+  result = toGeneratedExpr(n)
+  or
+  result = toGeneratedFormula(n)
+  or
   n = TClasslessPredicate(_, result)
   or
   n = TVarDecl(result)

From 2addbfabd08c4927ed632481f4c4ba940abaa583 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Wed, 26 May 2021 18:12:06 +0000
Subject: [PATCH 05/13] More work on `ComparisonFormula`

---
 ql/src/codeql_ql/ast/Ast.qll               | 49 +++++++++++++++++-----
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  8 +++-
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 91528697ab1..748101aff41 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -173,7 +173,7 @@ class Body extends TBody, AstNode {
 }
 
 /** A formula, such as `x = 6 and y < 5`. */
-abstract class Formula extends AstNode { }
+class Formula extends TFormula, AstNode { }
 
 /** An `and` formula, with 2 or more operands. */
 class Conjunction extends TConjunction, AstNode, Formula {
@@ -202,21 +202,45 @@ class Disjunction extends TDisjunction, AstNode {
 class ComparisonOp extends TComparisonOp, AstNode {
   Generated::Compop op;
 
-  ComparisonOp() {this = TComparisonOp(op)}
+  ComparisonOp() { this = TComparisonOp(op) }
 
+  ComparisonSymbol getSymbol() { result = op.getValue() }
 
+  override string getAPrimaryQlClass() { result = "ComparisonOp" }
+}
+
+/** A comparison symbol, such as `<` or `=`. */
+class ComparisonSymbol extends string {
+  ComparisonSymbol() {
+    this = "=" or
+    this = "!=" or
+    this = "<" or
+    this = ">" or
+    this = "<=" or
+    this = ">="
+  }
 }
 
 class ComparisonFormula extends TComparisonFormula, Formula {
-  Expr getLeftOperand() {none()}
-  Expr getRightOperand() {none()}
-  Expr getAnOperand() {none()}
-  ComparisonOp getOperator() {none()}
-  //ComparisonSymbol getSymbol() {none()}
+  Generated::CompTerm comp;
+
+  ComparisonFormula() { this = TComparisonFormula(comp) }
+
+  Expr getLeftOperand() { toGenerated(result) = comp.getLeft() }
+
+  Expr getRightOperand() { toGenerated(result) = comp.getRight() }
+
+  Expr getAnOperand() { result in [getLeftOperand(), getRightOperand()] }
+
+  ComparisonOp getOperator() { toGenerated(result) = comp.getChild() }
+
+  ComparisonSymbol getSymbol() { result = this.getOperator().getSymbol() }
+
+  override string getAPrimaryQlClass() { result = "ComparisonFormula" }
 }
 
 /** An expression, such as `x+4`. */
-abstract class Expr extends AstNode {}
+class Expr extends TExpr, AstNode { }
 
 /** A function symbol, such as `+` or `*`. */
 class FunctionSymbol extends string {
@@ -224,15 +248,18 @@ class FunctionSymbol extends string {
 }
 
 /** A binary operation, such as `x+3` or `y/2` */
-abstract class BinOpExpr extends Expr {}
+class BinOpExpr extends TBinOpExpr, Expr { }
 
 class AddExpr extends TAddExpr, BinOpExpr {
   Generated::AddExpr addexpr;
 
-  AddExpr() {this = TAddExpr(addexpr)}
+  AddExpr() { this = TAddExpr(addexpr) }
+
   Expr getLeftOperand() { toGenerated(result) = addexpr.getLeft() }
+
   Expr getRightOperand() { toGenerated(result) = addexpr.getRight() }
+
   Expr getAnOperand() { result = getLeftOperand() or result = getRightOperand() }
-  
+
   FunctionSymbol getOperator() { result = addexpr.getChild().getValue() }
 }
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 27115f7bcff..5290b75835f 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -17,13 +17,17 @@ newtype TAstNode =
   TComparisonOp(Generated::Compop op) or
   TAddExpr(Generated::AddExpr addexp)
 
+class TFormula = TDisjunction or TConjunction or TComparisonFormula;
 
-class TBinOpExpr = TAddExpr; // Can this work?
+class TBinOpExpr = TAddExpr;
+
+class TExpr = TBinOpExpr;
 
 Generated::AstNode toGeneratedFormula(AST::AstNode n) {
   n = TConjunction(result) or
   n = TDisjunction(result) or
-  n = TComparisonFormula(result)
+  n = TComparisonFormula(result) or
+  n = TComparisonOp(result)
 }
 
 Generated::AstNode toGeneratedExpr(AST::AstNode n) { n = TAddExpr(result) }

From 30577900714aa3486c61f887ada446197068ed58 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 19:39:54 +0000
Subject: [PATCH 06/13] ast for imports, modules, and various fixes

---
 ql/src/codeql_ql/ast/Ast.qll               | 204 ++++++++++++++++-----
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  22 ++-
 2 files changed, 179 insertions(+), 47 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 748101aff41..194c0f1b0cd 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -3,7 +3,7 @@ private import codeql_ql.ast.internal.AstNodes
 
 /** An AST node of a QL program */
 class AstNode extends TAstNode {
-  string toString() { result = "ASTNode" }
+  string toString() { result = getAPrimaryQlClass() }
 
   Location getLocation() { result = toGenerated(this).getLocation() }
 
@@ -12,11 +12,40 @@ class AstNode extends TAstNode {
   string getAPrimaryQlClass() { result = "???" }
 }
 
+/**
+ * The `from, where, select` part of a QL query.
+ */
+class Select extends TSelect, AstNode {
+  Generated::Select sel;
+
+  Select() { this = TSelect(sel) }
+
+  override string getAPrimaryQlClass() { result = "Select" }
+  // TODO: Getters for VarDecls, Where-clause, selects.
+}
+
+class Predicate extends TPredicate, AstNode {
+  /**
+   * Gets the body of the predicate.
+   */
+  Formula getBody() { none() }
+
+  /**
+   * Gets the name of the predicate
+   */
+  string getName() { none() }
+
+  /**
+   * Gets the `i`th parameter of the predicate.
+   */
+  VarDecl getParameter(int i) { none() }
+  // TODO: ReturnType.
+}
+
 /**
  * A classless predicate.
  */
-class ClasslessPredicate extends TClasslessPredicate, AstNode {
-  // TODO: Make super class for predicate things. (classless, class predicate, charpred)
+class ClasslessPredicate extends TClasslessPredicate, Predicate, ModuleMember {
   Generated::ModuleMember member;
   Generated::ClasslessPredicate pred;
 
@@ -28,64 +57,51 @@ class ClasslessPredicate extends TClasslessPredicate, AstNode {
 
   override string getAPrimaryQlClass() { result = "ClasslessPredicate" }
 
-  /**
-   * Gets the `i`th parameter of the predicate.
-   */
-  VarDecl getParameter(int i) {
+  override Formula getBody() { toGenerated(result) = pred.getChild(_).(Generated::Body).getChild() }
+
+  override string getName() { result = pred.getName().getValue() }
+
+  override VarDecl getParameter(int i) {
     toGenerated(result) =
       rank[i](Generated::VarDecl decl, int index | decl = pred.getChild(index) | decl order by index)
   }
-
-  /**
-   * Gets the body of the predicate.
-   */
-  Body getBody() { toGenerated(result) = pred.getChild(_) }
-
-  /**
-   * Gets the name of the predicate
-   */
-  string getName() { result = pred.getName().getValue() }
 }
 
 /**
  * A predicate in a class.
  */
-class ClassPredicate extends TClassPredicate, AstNode {
+class ClassPredicate extends TClassPredicate, Predicate {
   Generated::MemberPredicate pred;
 
   ClassPredicate() { this = TClassPredicate(pred) }
 
-  /**
-   * Gets the name of the predicate.
-   */
-  string getName() { result = pred.getName().getValue() }
+  override string getName() { result = pred.getName().getValue() }
 
-  /**
-   * Gets the body containing the implementation of the predicate.
-   */
-  Body getBody() { toGenerated(result) = pred.getChild(_) }
+  override Formula getBody() { toGenerated(result) = pred.getChild(_).(Generated::Body).getChild() }
 
-  // TODO: ReturnType.
   override string getAPrimaryQlClass() { result = "ClassPredicate" }
 
   override Class getParent() { result.getAClassPredicate() = this }
+
+  override VarDecl getParameter(int i) {
+    toGenerated(result) =
+      rank[i](Generated::VarDecl decl, int index | decl = pred.getChild(index) | decl order by index)
+  }
 }
 
 /**
  * A characteristic predicate of a class.
  */
-class CharPred extends TCharPred, AstNode {
+class CharPred extends TCharPred, Predicate {
   Generated::Charpred pred;
 
   CharPred() { this = TCharPred(pred) }
 
   override string getAPrimaryQlClass() { result = "CharPred" }
 
-  /*
-   * Gets the body of the predicate.
-   * TODO: The body is just directly an Expr in the generated AST. Wait for Expr type to appear.
-   * Body getBody() { toGenerated(result) = pred.getChild(_) }MemberPredicate
-   */
+  override Formula getBody() { toGenerated(result) = pred.getBody() }
+
+  override string getName() { result = getParent().getName() }
 
   override Class getParent() { result.getCharPred() = this }
 }
@@ -113,10 +129,40 @@ class VarDecl extends TVarDecl, AstNode {
   // TODO: Getter for the Type.
 }
 
+/**
+ * A QL module.
+ */
+class Module extends TModule, AstNode, ModuleMember {
+  Generated::Module mod;
+
+  Module() { this = TModule(mod) }
+
+  override string getAPrimaryQlClass() { result = "Module" }
+
+  /**
+   * Gets the name of the module.
+   */
+  string getName() { result = mod.getName().(Generated::ModuleName).getChild().getValue() }
+
+  /**
+   * Gets a member of the module.
+   */
+  AstNode getAMember() {
+    toGenerated(result) = mod.getChild(_).(Generated::ModuleMember).getChild(_)
+  }
+}
+
+/**
+ * Something that can be member of a module.
+ */
+class ModuleMember extends TModuleMember, AstNode {
+  override AstNode getParent() { result.(Module).getAMember() = this }
+}
+
 /**
  * A QL class.
  */
-class Class extends TClass, AstNode {
+class Class extends TClass, AstNode, ModuleMember {
   Generated::Dataclass cls;
 
   Class() { this = TClass(cls) }
@@ -132,14 +178,14 @@ class Class extends TClass, AstNode {
    * Gets the charateristic predicate for this class.
    */
   CharPred getCharPred() {
-    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(0).(Generated::Charpred)
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(_)
   }
 
   /**
    * Gets a predicate in this class.
    */
   ClassPredicate getAClassPredicate() {
-    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(0)
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(_)
   }
 
   /**
@@ -155,25 +201,95 @@ class Class extends TClass, AstNode {
    */
   VarDecl getAField() {
     toGenerated(result) =
-      cls.getChild(_).(Generated::ClassMember).getChild(0).(Generated::Field).getChild()
+      cls.getChild(_).(Generated::ClassMember).getChild(_).(Generated::Field).getChild()
   }
   // TODO: extends, modifiers.
 }
 
 /**
- * The body of a predicate.
+ * A `newtype Foo` declaration.
  */
-class Body extends TBody, AstNode {
-  Generated::Body body;
+class NewType extends TNewType, ModuleMember {
+  Generated::Datatype type;
 
-  Body() { this = TBody(body) }
+  NewType() { this = TNewType(type) }
 
-  override string getAPrimaryQlClass() { result = "Body" }
-  // TODO: Children.
+  string getName() { result = type.getName().getValue() }
+
+  override string getAPrimaryQlClass() { result = "DataType" }
+
+  NewTypeBranch getABranch() { toGenerated(result) = type.getChild().getChild(_) }
+}
+
+/**
+ * A branch in a `newtype`.
+ */
+class NewTypeBranch extends TNewTypeBranch, AstNode {
+  Generated::DatatypeBranch branch;
+
+  NewTypeBranch() { this = TNewTypeBranch(branch) }
+
+  override string getAPrimaryQlClass() { result = "NewTypeBranch" }
+
+  string getName() { result = branch.getName().getValue() }
+
+  VarDecl getField(int i) {
+    toGenerated(result) =
+      rank[i](Generated::VarDecl var | var = branch.getChild(i) | var order by i)
+  }
+
+  Formula getBody() { toGenerated(result) = branch.getChild(_).(Generated::Body).getChild() }
+}
+
+/**
+ * An import statement.
+ */
+class Import extends TImport, ModuleMember {
+  Generated::ImportDirective imp;
+
+  Import() { this = TImport(imp) }
+
+  override string getAPrimaryQlClass() { result = "Import" }
+
+  /**
+   * Gets the name under which this import is imported, if such a name exists.
+   * E.g. the `Flow` in:
+   * ```
+   * import semmle.javascript.dataflow.Configuration as Flow
+   * ```
+   */
+  string importedAs() { result = imp.getChild(1).(Generated::ModuleName).getChild().getValue() }
+
+  /**
+   * Gets the `i`th selected name from the imported module.
+   * E.g. for
+   * `import foo.bar::Baz::Qux`
+   * It is true that `getSelectionName(0) = "Baz"` and `getSelectionName(1) = "Qux"`.
+   */
+  string getSelectionName(int i) {
+    result = imp.getChild(0).(Generated::ImportModuleExpr).getName(i).getValue()
+  }
+
+  /**
+   * Gets the `i`th imported module.
+   * E.g. for
+   * `import foo.bar::Baz::Qux`
+   * It is true that `getQualifiedName(0) = "foo"` and `getQualifiedName(1) = "bar"`.
+   */
+  string getQualifiedName(int i) {
+    result = imp.getChild(0).(Generated::ImportModuleExpr).getChild().getName(i).getValue()
+  }
+  // TODO: private modifier.
 }
 
 /** A formula, such as `x = 6 and y < 5`. */
-class Formula extends TFormula, AstNode { }
+class Formula extends TFormula, AstNode {
+  override AstNode getParent() {
+    result = super.getParent()
+    or
+    result.(Predicate).getBody() = this
+  }
+}
 
 /** An `and` formula, with 2 or more operands. */
 class Conjunction extends TConjunction, AstNode, Formula {
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 5290b75835f..54c1f61ac82 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -7,10 +7,14 @@ newtype TAstNode =
     pred.getParent() = member
   } or
   TVarDecl(Generated::VarDecl decl) or
-  TBody(Generated::Body body) or
   TClass(Generated::Dataclass dc) or
   TCharPred(Generated::Charpred pred) or
   TClassPredicate(Generated::MemberPredicate pred) or
+  TSelect(Generated::Select sel) or
+  TModule(Generated::Module mod) or
+  TNewType(Generated::Datatype dt) or
+  TNewTypeBranch(Generated::DatatypeBranch branch) or
+  TImport(Generated::ImportDirective imp) or
   TDisjunction(Generated::Disjunction disj) or
   TConjunction(Generated::Conjunction conj) or
   TComparisonFormula(Generated::CompTerm comp) or
@@ -44,11 +48,23 @@ Generated::AstNode toGenerated(AST::AstNode n) {
   or
   n = TVarDecl(result)
   or
-  n = TBody(result)
-  or
   n = TClass(result)
   or
   n = TCharPred(result)
   or
   n = TClassPredicate(result)
+  or
+  n = TSelect(result)
+  or
+  n = TModule(result)
+  or
+  n = TNewType(result)
+  or
+  n = TNewTypeBranch(result)
+  or
+  n = TImport(result)
 }
+
+class TPredicate = TCharPred or TClasslessPredicate or TClassPredicate;
+
+class TModuleMember = TClasslessPredicate or TClass or TModule or TNewType or TImport;

From 4eb836ca9b7c66153b7a7d469e1b8006ac27f793 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Wed, 26 May 2021 21:06:04 +0000
Subject: [PATCH 07/13] AST for quantifiers, negation

---
 Cargo.lock                                   |  2 +-
 extractor/Cargo.toml                         |  2 +-
 generator/Cargo.toml                         |  2 +-
 ql/src/codeql_ql/ast/Ast.qll                 | 64 ++++++++++++++++++++
 ql/src/codeql_ql/ast/internal/AstNodes.qll   |  8 ++-
 ql/src/codeql_ql/ast/internal/TreeSitter.qll | 10 ++-
 ql/src/ql.dbscheme                           | 14 +++++
 7 files changed, 96 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d70d80f9439..bb7f945c686 100644
--- a/Cargo.lock
+++ b/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=4125f8b919f80889dec4b74842419d287ccf1782#4125f8b919f80889dec4b74842419d287ccf1782"
+source = "git+https://github.com/tausbn/tree-sitter-ql.git?rev=a58a56d7ee47bfd75728916a49b576af29238a1d#a58a56d7ee47bfd75728916a49b576af29238a1d"
 dependencies = [
  "cc",
  "tree-sitter",
diff --git a/extractor/Cargo.toml b/extractor/Cargo.toml
index 44fb96c3cbb..6ac0c3e3e72 100644
--- a/extractor/Cargo.toml
+++ b/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 = "4125f8b919f80889dec4b74842419d287ccf1782" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "a58a56d7ee47bfd75728916a49b576af29238a1d" }
 clap = "2.33"
 tracing = "0.1"
 tracing-subscriber = { version = "0.2", features = ["env-filter"] }
diff --git a/generator/Cargo.toml b/generator/Cargo.toml
index 755353b7011..1ca391c42f3 100644
--- a/generator/Cargo.toml
+++ b/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 = "4125f8b919f80889dec4b74842419d287ccf1782" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "a58a56d7ee47bfd75728916a49b576af29238a1d" }
diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 194c0f1b0cd..5ea9c6de2e0 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -355,6 +355,70 @@ class ComparisonFormula extends TComparisonFormula, Formula {
   override string getAPrimaryQlClass() { result = "ComparisonFormula" }
 }
 
+class Quantifier extends TQuantifier, Formula {
+  Generated::Quantified quant;
+  string kind;
+
+  Quantifier() {
+    this = TQuantifier(quant) and kind = quant.getChild(0).(Generated::Quantifier).getValue()
+  }
+
+  /** Gets the ith declared argument of this quantifier. */
+  VarDecl getArgument(int i) {
+    i >= 1 and
+    toGenerated(result) = quant.getChild(i - 1) and
+    exists()
+  }
+
+  /** Gets an argument of this quantifier. */
+  VarDecl getAnArgument() { result = this.getArgument(_) }
+
+  /** Gets the formula restricting the range of this quantifier, if any. */
+  Formula getRange() { toGenerated(result) = quant.getRange() }
+
+  /** Holds if this quantifier has a range formula. */
+  predicate hasRange() { exists(this.getRange()) }
+
+  /** Gets the main body of the quantifier. */
+  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) }
+
+  /** Holds if this is the expression only form of an exists quantifier. */
+  predicate hasExpr() { exists(getExpr()) }
+
+  override string getAPrimaryQlClass() { result = "Quantifier" }
+}
+
+class Exists extends Quantifier {
+  Exists() { kind = "exists" }
+
+  override string getAPrimaryQlClass() { result = "Exists" }
+}
+
+class Forall extends Quantifier {
+  Forall() { kind = "forall" }
+
+  override string getAPrimaryQlClass() { result = "Forall" }
+}
+
+class Forex extends Quantifier {
+  Forex() { kind = "forex" }
+
+  override string getAPrimaryQlClass() { result = "Forex" }
+}
+
+class Negation extends TNegation, Formula {
+  Generated::Negation neg;
+
+  Negation() { this = TNegation(neg) }
+
+  Formula getFormula() { toGenerated(result) = neg.getChild() }
+
+  override string getAPrimaryQlClass() { result = "Negation" }
+}
+
 /** An expression, such as `x+4`. */
 class Expr extends TExpr, AstNode { }
 
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 54c1f61ac82..66edaea97a8 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -19,9 +19,11 @@ newtype TAstNode =
   TConjunction(Generated::Conjunction conj) or
   TComparisonFormula(Generated::CompTerm comp) or
   TComparisonOp(Generated::Compop op) or
+  TQuantifier(Generated::Quantified quant) or
+  TNegation(Generated::Negation neg) or
   TAddExpr(Generated::AddExpr addexp)
 
-class TFormula = TDisjunction or TConjunction or TComparisonFormula;
+class TFormula = TDisjunction or TConjunction or TComparisonFormula or TQuantifier or TNegation;
 
 class TBinOpExpr = TAddExpr;
 
@@ -31,7 +33,9 @@ Generated::AstNode toGeneratedFormula(AST::AstNode n) {
   n = TConjunction(result) or
   n = TDisjunction(result) or
   n = TComparisonFormula(result) or
-  n = TComparisonOp(result)
+  n = TComparisonOp(result) or
+  n = TQuantifier(result) or
+  n = TNegation(result)
 }
 
 Generated::AstNode toGeneratedExpr(AST::AstNode n) { n = TAddExpr(result) }
diff --git a/ql/src/codeql_ql/ast/internal/TreeSitter.qll b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
index a585096e7ea..b4533ee9735 100644
--- a/ql/src/codeql_ql/ast/internal/TreeSitter.qll
+++ b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
@@ -746,9 +746,17 @@ module Generated {
 
     override Location getLocation() { quantified_def(this, result) }
 
+    AstNode getFormula() { quantified_formula(this, result) }
+
+    AstNode getRange() { quantified_range(this, result) }
+
     AstNode getChild(int i) { quantified_child(this, i, result) }
 
-    override AstNode getAFieldOrChild() { quantified_child(this, _, result) }
+    override AstNode getAFieldOrChild() {
+      quantified_formula(this, result) or
+      quantified_range(this, result) or
+      quantified_child(this, _, result)
+    }
   }
 
   class Quantifier extends @token_quantifier, Token {
diff --git a/ql/src/ql.dbscheme b/ql/src/ql.dbscheme
index 9041b7801fc..dcb57473c3a 100644
--- a/ql/src/ql.dbscheme
+++ b/ql/src/ql.dbscheme
@@ -661,6 +661,20 @@ qualified_expr_def(
   int loc: @location 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(
+  unique int quantified: @quantified ref,
+  unique int formula: @quantified_formula_type ref
+);
+
+@quantified_range_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_range(
+  unique int quantified: @quantified ref,
+  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
 
 #keyset[quantified, index]

From ec98e8a82da6b55b97686c341985375d2093a2d4 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Wed, 26 May 2021 21:25:42 +0000
Subject: [PATCH 08/13] Add convenience methods for aggregates

---
 Cargo.lock                                   |  2 +-
 extractor/Cargo.toml                         |  2 +-
 generator/Cargo.toml                         |  2 +-
 ql/src/codeql_ql/ast/Ast.qll                 |  2 +-
 ql/src/codeql_ql/ast/internal/TreeSitter.qll | 40 +++++++++++---
 ql/src/ql.dbscheme                           | 57 +++++++++++++++-----
 6 files changed, 82 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index bb7f945c686..d37b84d6dda 100644
--- a/Cargo.lock
+++ b/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",
diff --git a/extractor/Cargo.toml b/extractor/Cargo.toml
index 6ac0c3e3e72..7b238a75ccb 100644
--- a/extractor/Cargo.toml
+++ b/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"] }
diff --git a/generator/Cargo.toml b/generator/Cargo.toml
index 1ca391c42f3..a26291bcb55 100644
--- a/generator/Cargo.toml
+++ b/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" }
diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 5ea9c6de2e0..9237f9e38f1 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/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()) }
diff --git a/ql/src/codeql_ql/ast/internal/TreeSitter.qll b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
index b4533ee9735..c9684cb2e6e 100644
--- a/ql/src/codeql_ql/ast/internal/TreeSitter.qll
+++ b/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 {
diff --git a/ql/src/ql.dbscheme b/ql/src/ql.dbscheme
index dcb57473c3a..a0e8d866cf6 100644
--- a/ql/src/ql.dbscheme
+++ b/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(

From a692794178acae0eafd6d4ebeff768d4ac8bc29d Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 21:10:37 +0000
Subject: [PATCH 09/13] add types

---
 ql/src/codeql_ql/ast/Ast.qll               | 38 ++++++++++++++++++++--
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  3 ++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 9237f9e38f1..e5a60b8b14a 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -117,7 +117,7 @@ class VarDecl extends TVarDecl, AstNode {
   /**
    * Gets the name for this variable declaration.
    */
-  string getName() { result = var.getChild(_).(Generated::VarName).getChild().getValue() }
+  string getName() { result = var.getChild(1).(Generated::VarName).getChild().getValue() }
 
   override string getAPrimaryQlClass() { result = "VarDecl" }
 
@@ -126,7 +126,41 @@ class VarDecl extends TVarDecl, AstNode {
     or
     result.(Class).getAField() = this
   }
-  // TODO: Getter for the Type.
+
+  Type getType() { toGenerated(result) = var.getChild(0) }
+}
+
+/**
+ * A type, such as `DataFlow::Node`.
+ */
+class Type extends TType, AstNode {
+  Generated::TypeExpr type;
+
+  Type() { this = TType(type) }
+
+  override string getAPrimaryQlClass() { result = "Type" }
+
+  /**
+   * Gets the class name for the type.
+   * E.g. `Node` in `DataFlow::Node`.
+   * Also gets the name for primitive types such as `string` or `int`.
+   */
+  string getClassName() {
+    result = type.getName().getValue()
+    or
+    result = type.getChild().(Generated::PrimitiveType).getValue()
+  }
+
+  /**
+   * Holds if this type is a primitive such as `string` or `int`.
+   */
+  predicate isPrimitive() { type.getChild() instanceof Generated::PrimitiveType }
+
+  /**
+   * Gets the module name of the type, if it exists.
+   * E.g. `DataFlow` in `DataFlow::Node`.
+   */
+  string getModuleName() { result = type.getChild().(Generated::ModuleExpr).getName().getValue() }
 }
 
 /**
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 66edaea97a8..e482d673655 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -15,6 +15,7 @@ newtype TAstNode =
   TNewType(Generated::Datatype dt) or
   TNewTypeBranch(Generated::DatatypeBranch branch) or
   TImport(Generated::ImportDirective imp) or
+  TType(Generated::TypeExpr type) or
   TDisjunction(Generated::Disjunction disj) or
   TConjunction(Generated::Conjunction conj) or
   TComparisonFormula(Generated::CompTerm comp) or
@@ -67,6 +68,8 @@ Generated::AstNode toGenerated(AST::AstNode n) {
   n = TNewTypeBranch(result)
   or
   n = TImport(result)
+  or
+  n = TType(result)
 }
 
 class TPredicate = TCharPred or TClasslessPredicate or TClassPredicate;

From ee9e1914b0fd465780b7367710e188d5e2b93362 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Wed, 26 May 2021 21:34:33 +0000
Subject: [PATCH 10/13] literals and fixes

---
 ql/src/codeql_ql/ast/Ast.qll               | 41 ++++++++++++++++++++--
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  7 ++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index e5a60b8b14a..6a88b7cd3bb 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -143,12 +143,15 @@ class Type extends TType, AstNode {
   /**
    * Gets the class name for the type.
    * E.g. `Node` in `DataFlow::Node`.
-   * Also gets the name for primitive types such as `string` or `int`.
+   * Also gets the name for primitive types such as `string` or `int`
+   * or db-types such as `@locateable`.
    */
   string getClassName() {
     result = type.getName().getValue()
     or
     result = type.getChild().(Generated::PrimitiveType).getValue()
+    or
+    result = type.getChild().(Generated::Dbtype).getValue()
   }
 
   /**
@@ -156,6 +159,11 @@ class Type extends TType, AstNode {
    */
   predicate isPrimitive() { type.getChild() instanceof Generated::PrimitiveType }
 
+  /**
+   * Holds if this type is a db-type.
+   */
+  predicate isDBType() { type.getChild() instanceof Generated::Dbtype }
+
   /**
    * Gets the module name of the type, if it exists.
    * E.g. `DataFlow` in `DataFlow::Node`.
@@ -237,7 +245,12 @@ class Class extends TClass, AstNode, ModuleMember {
     toGenerated(result) =
       cls.getChild(_).(Generated::ClassMember).getChild(_).(Generated::Field).getChild()
   }
-  // TODO: extends, modifiers.
+
+  /**
+   * Gets a super-type for this class.
+   * That is: a type after the `extends` keyword.
+   */
+  Type getASuperType() { toGenerated(result) = cls.getChild(_) }
 }
 
 /**
@@ -359,6 +372,30 @@ class ComparisonOp extends TComparisonOp, AstNode {
   override string getAPrimaryQlClass() { result = "ComparisonOp" }
 }
 
+class Literal extends TLiteral, AstNode {
+  Generated::Literal lit;
+
+  Literal() { this = TLiteral(lit) }
+
+  override string getAPrimaryQlClass() { result = "??Literal??" }
+}
+
+class String extends Literal {
+  String() { lit.getChild() instanceof Generated::String }
+
+  override string getAPrimaryQlClass() { result = "String" }
+
+  string getValue() { result = lit.getChild().(Generated::String).getValue() }
+}
+
+class Integer extends Literal {
+  Integer() { lit.getChild() instanceof Generated::Integer }
+
+  override string getAPrimaryQlClass() { result = "Integer" }
+
+  int getValue() { result = lit.getChild().(Generated::Integer).getValue().toInt() }
+}
+
 /** A comparison symbol, such as `<` or `=`. */
 class ComparisonSymbol extends string {
   ComparisonSymbol() {
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index e482d673655..ee0a5be93b1 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -22,13 +22,14 @@ newtype TAstNode =
   TComparisonOp(Generated::Compop op) or
   TQuantifier(Generated::Quantified quant) or
   TNegation(Generated::Negation neg) or
-  TAddExpr(Generated::AddExpr addexp)
+  TAddExpr(Generated::AddExpr addexp) or
+  TLiteral(Generated::Literal lit)
 
 class TFormula = TDisjunction or TConjunction or TComparisonFormula or TQuantifier or TNegation;
 
 class TBinOpExpr = TAddExpr;
 
-class TExpr = TBinOpExpr;
+class TExpr = TBinOpExpr or TLiteral;
 
 Generated::AstNode toGeneratedFormula(AST::AstNode n) {
   n = TConjunction(result) or
@@ -70,6 +71,8 @@ Generated::AstNode toGenerated(AST::AstNode n) {
   n = TImport(result)
   or
   n = TType(result)
+  or
+  n = TLiteral(result)
 }
 
 class TPredicate = TCharPred or TClasslessPredicate or TClassPredicate;

From fec3d745cac9adc39115e97ebc21f28bd0b01e68 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Thu, 27 May 2021 07:31:44 +0000
Subject: [PATCH 11/13] make sure Literal is an Expr

---
 ql/src/codeql_ql/ast/Ast.qll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 6a88b7cd3bb..0f59c07f092 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -372,7 +372,7 @@ class ComparisonOp extends TComparisonOp, AstNode {
   override string getAPrimaryQlClass() { result = "ComparisonOp" }
 }
 
-class Literal extends TLiteral, AstNode {
+class Literal extends TLiteral, Expr {
   Generated::Literal lit;
 
   Literal() { this = TLiteral(lit) }

From ab20f8f74e15fc8ed5f8911717451623dd7b6ed2 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Thu, 27 May 2021 08:23:15 +0000
Subject: [PATCH 12/13] add support for aggregates

---
 ql/src/codeql_ql/ast/Ast.qll                 | 84 +++++++++++++++++++-
 ql/src/codeql_ql/ast/internal/AstNodes.qll   |  9 ++-
 ql/src/codeql_ql/ast/internal/TreeSitter.qll |  5 ++
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 0f59c07f092..60646043ba3 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -125,6 +125,10 @@ class VarDecl extends TVarDecl, AstNode {
     result = super.getParent()
     or
     result.(Class).getAField() = this
+    or
+    result.(Aggregate).getAnArgument() = this
+    or
+    result.(Quantifier).getAnArgument() = this
   }
 
   Type getType() { toGenerated(result) = var.getChild(0) }
@@ -335,6 +339,8 @@ class Formula extends TFormula, AstNode {
     result = super.getParent()
     or
     result.(Predicate).getBody() = this
+    or
+    result.(Aggregate).getGuard() = this
   }
 }
 
@@ -437,8 +443,7 @@ class Quantifier extends TQuantifier, Formula {
   /** Gets the ith declared argument of this quantifier. */
   VarDecl getArgument(int i) {
     i >= 1 and
-    toGenerated(result) = quant.getChild(i - 1) and
-    exists()
+    toGenerated(result) = quant.getChild(i - 1)
   }
 
   /** Gets an argument of this quantifier. */
@@ -480,6 +485,81 @@ class Forex extends Quantifier {
   override string getAPrimaryQlClass() { result = "Forex" }
 }
 
+class Aggregate extends TAggregate, Expr {
+  Generated::Aggregate agg;
+  Generated::FullAggregateBody body;
+  string kind;
+
+  Aggregate() {
+    this = TAggregate(agg) and
+    kind = agg.getChild(0).(Generated::AggId).getValue() and
+    body = agg.getChild(_)
+  }
+
+  string getKind() { result = kind }
+
+  /** Gets the ith declared argument of this quantifier. */
+  VarDecl getArgument(int i) { toGenerated(result) = body.getChild(i) }
+
+  /** Gets an argument of this quantifier. */
+  VarDecl getAnArgument() { result = this.getArgument(_) }
+
+  Formula getGuard() { toGenerated(result) = body.getGuard() }
+
+  AsExpr getAsExpr(int i) { toGenerated(result) = body.getAsExprs().getChild(i) }
+
+  Expr getOrderBy(int i) { toGenerated(result) = body.getOrderBys().getChild(i).getChild(0) }
+
+  string getOrderbyDirection(int i) {
+    result = body.getOrderBys().getChild(i).getChild(1).(Generated::Direction).getValue()
+  }
+
+  override string getAPrimaryQlClass() { result = "Aggregate[" + kind + "]" }
+}
+
+class Rank extends Aggregate {
+  Rank() { kind = "rank" }
+
+  override string getAPrimaryQlClass() { result = "Rank" }
+
+  /**
+   * The `i` in `rank[i]( | | )`.
+   */
+  Expr getRankExpr() { toGenerated(result) = agg.getChild(1) }
+}
+
+class AsExpr extends TAsExpr, AstNode {
+  Generated::AsExpr asExpr;
+
+  AsExpr() { this = TAsExpr(asExpr) }
+
+  override string getAPrimaryQlClass() { result = "AsExpr" }
+
+  /**
+   * Gets the name the inner expression gets "saved" under.
+   * If such a name exists.
+   */
+  string getAsName() { result = asExpr.getChild(1).(Generated::VarName).getChild().getValue() }
+
+  Expr getInnerExpr() { toGenerated(result) = asExpr.getChild(0) }
+
+  override AstNode getParent() {
+    result = super.getParent()
+    or
+    result.(Aggregate).getAsExpr(_) = this
+  }
+}
+
+class Identifier extends TIdentifier, Expr {
+  Generated::Variable id;
+
+  Identifier() { this = TIdentifier(id) }
+
+  string getName() { result = id.getChild().(Generated::VarName).getChild().getValue() }
+
+  override string getAPrimaryQlClass() { result = "Identifier" }
+}
+
 class Negation extends TNegation, Formula {
   Generated::Negation neg;
 
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index ee0a5be93b1..cc9d0ae624c 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -21,6 +21,9 @@ newtype TAstNode =
   TComparisonFormula(Generated::CompTerm comp) or
   TComparisonOp(Generated::Compop op) or
   TQuantifier(Generated::Quantified quant) or
+  TAggregate(Generated::Aggregate agg) or
+  TIdentifier(Generated::Variable var) or
+  TAsExpr(Generated::AsExpr asExpr) or
   TNegation(Generated::Negation neg) or
   TAddExpr(Generated::AddExpr addexp) or
   TLiteral(Generated::Literal lit)
@@ -29,7 +32,7 @@ class TFormula = TDisjunction or TConjunction or TComparisonFormula or TQuantifi
 
 class TBinOpExpr = TAddExpr;
 
-class TExpr = TBinOpExpr or TLiteral;
+class TExpr = TBinOpExpr or TLiteral or TAggregate or TIdentifier;
 
 Generated::AstNode toGeneratedFormula(AST::AstNode n) {
   n = TConjunction(result) or
@@ -37,6 +40,8 @@ Generated::AstNode toGeneratedFormula(AST::AstNode n) {
   n = TComparisonFormula(result) or
   n = TComparisonOp(result) or
   n = TQuantifier(result) or
+  n = TAggregate(result) or
+  n = TIdentifier(result) or
   n = TNegation(result)
 }
 
@@ -73,6 +78,8 @@ Generated::AstNode toGenerated(AST::AstNode n) {
   n = TType(result)
   or
   n = TLiteral(result)
+  or
+  n = TAsExpr(result)
 }
 
 class TPredicate = TCharPred or TClasslessPredicate or TClassPredicate;
diff --git a/ql/src/codeql_ql/ast/internal/TreeSitter.qll b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
index c9684cb2e6e..e353a0678f3 100644
--- a/ql/src/codeql_ql/ast/internal/TreeSitter.qll
+++ b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
@@ -123,6 +123,11 @@ module Generated {
     override AstNode getAFieldOrChild() { as_expr_child(this, _, result) }
   }
 
+  AsExpr childThing(int i, AstNode child) {
+    result.getChild(i) = child and
+    i != 0
+  }
+
   class AsExprs extends @as_exprs, AstNode {
     override string getAPrimaryQlClass() { result = "AsExprs" }
 

From 09d0cdbaf8c4697f04cb05e59e2b9da305501889 Mon Sep 17 00:00:00 2001
From: Tom Hvitved <hvitved@github.com>
Date: Thu, 27 May 2021 10:26:33 +0200
Subject: [PATCH 13/13] Add `ModuleMember::isPrivate()`

---
 ql/src/codeql_ql/ast/Ast.qll | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index 60646043ba3..f47f950c600 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -51,7 +51,7 @@ class ClasslessPredicate extends TClasslessPredicate, Predicate, ModuleMember {
 
   ClasslessPredicate() { this = TClasslessPredicate(member, pred) }
 
-  predicate isPrivate() {
+  final override predicate isPrivate() {
     member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
   }
 
@@ -185,6 +185,13 @@ class Module extends TModule, AstNode, ModuleMember {
 
   override string getAPrimaryQlClass() { result = "Module" }
 
+  final override predicate isPrivate() {
+    exists(Generated::ModuleMember member |
+      mod = member.getChild(_) and
+      member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
+    )
+  }
+
   /**
    * Gets the name of the module.
    */
@@ -203,6 +210,9 @@ class Module extends TModule, AstNode, ModuleMember {
  */
 class ModuleMember extends TModuleMember, AstNode {
   override AstNode getParent() { result.(Module).getAMember() = this }
+
+  /** Holds if this member is declared as `private`. */
+  predicate isPrivate() { none() }
 }
 
 /**
@@ -215,6 +225,13 @@ class Class extends TClass, AstNode, ModuleMember {
 
   override string getAPrimaryQlClass() { result = "Class" }
 
+  final override predicate isPrivate() {
+    exists(Generated::ModuleMember member |
+      cls = member.getChild(_) and
+      member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
+    )
+  }
+
   /**
    * Gets the name of the class.
    */
@@ -269,6 +286,13 @@ class NewType extends TNewType, ModuleMember {
 
   override string getAPrimaryQlClass() { result = "DataType" }
 
+  final override predicate isPrivate() {
+    exists(Generated::ModuleMember member |
+      type = member.getChild(_) and
+      member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
+    )
+  }
+
   NewTypeBranch getABranch() { toGenerated(result) = type.getChild().getChild(_) }
 }
 
@@ -330,7 +354,13 @@ class Import extends TImport, ModuleMember {
   string getQualifiedName(int i) {
     result = imp.getChild(0).(Generated::ImportModuleExpr).getChild().getName(i).getValue()
   }
-  // TODO: private modifier.
+
+  final override predicate isPrivate() {
+    exists(Generated::ModuleMember member |
+      imp = member.getChild(_) and
+      member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
+    )
+  }
 }
 
 /** A formula, such as `x = 6 and y < 5`. */