Merge pull request #7 from github/ast

AST layer

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=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 = "4125f8b919f80889dec4b74842419d287ccf1782" }
+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 = "4125f8b919f80889dec4b74842419d287ccf1782" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "5f3e790557ad6d6612e269808168e7a8f6413dc0" }

ql/src/codeql_ql/ast/Ast.qll

@@ -3,21 +3,624 @@ 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() }
+
+  AstNode getParent() { toGenerated(result) = toGenerated(this).getParent() }
+
+  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 {
+class ClasslessPredicate extends TClasslessPredicate, Predicate, ModuleMember {
   Generated::ModuleMember member;
   Generated::ClasslessPredicate pred;
 
   ClasslessPredicate() { this = TClasslessPredicate(member, pred) }
 
-  predicate isPrivate() {
+  final override predicate isPrivate() {
     member.getAFieldOrChild().(Generated::Annotation).getName().getValue() = "private"
   }
+
+  override string getAPrimaryQlClass() { result = "ClasslessPredicate" }
+
+  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)
+  }
+}
+
+/**
+ * A predicate in a class.
+ */
+class ClassPredicate extends TClassPredicate, Predicate {
+  Generated::MemberPredicate pred;
+
+  ClassPredicate() { this = TClassPredicate(pred) }
+
+  override string getName() { result = pred.getName().getValue() }
+
+  override Formula getBody() { toGenerated(result) = pred.getChild(_).(Generated::Body).getChild() }
+
+  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, Predicate {
+  Generated::Charpred pred;
+
+  CharPred() { this = TCharPred(pred) }
+
+  override string getAPrimaryQlClass() { result = "CharPred" }
+
+  override Formula getBody() { toGenerated(result) = pred.getBody() }
+
+  override string getName() { result = getParent().getName() }
+
+  override Class getParent() { result.getCharPred() = this }
+}
+
+/**
+ * A variable declaration, with a type and a name.
+ */
+class VarDecl extends TVarDecl, AstNode {
+  Generated::VarDecl var;
+
+  VarDecl() { this = TVarDecl(var) }
+
+  /**
+   * Gets the name for this variable declaration.
+   */
+  string getName() { result = var.getChild(1).(Generated::VarName).getChild().getValue() }
+
+  override string getAPrimaryQlClass() { result = "VarDecl" }
+
+  override AstNode getParent() {
+    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) }
+}
+
+/**
+ * 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`
+   * 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()
+  }
+
+  /**
+   * Holds if this type is a primitive such as `string` or `int`.
+   */
+  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`.
+   */
+  string getModuleName() { result = type.getChild().(Generated::ModuleExpr).getName().getValue() }
+}
+
+/**
+ * A QL module.
+ */
+class Module extends TModule, AstNode, ModuleMember {
+  Generated::Module mod;
+
+  Module() { this = TModule(mod) }
+
+  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.
+   */
+  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 }
+
+  /** Holds if this member is declared as `private`. */
+  predicate isPrivate() { none() }
+}
+
+/**
+ * A QL class.
+ */
+class Class extends TClass, AstNode, ModuleMember {
+  Generated::Dataclass cls;
+
+  Class() { this = TClass(cls) }
+
+  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.
+   */
+  string getName() { result = cls.getName().getValue() }
+
+  /**
+   * Gets the charateristic predicate for this class.
+   */
+  CharPred getCharPred() {
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(_)
+  }
+
+  /**
+   * Gets a predicate in this class.
+   */
+  ClassPredicate getAClassPredicate() {
+    toGenerated(result) = cls.getChild(_).(Generated::ClassMember).getChild(_)
+  }
+
+  /**
+   * 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(_).(Generated::Field).getChild()
+  }
+
+  /**
+   * Gets a super-type for this class.
+   * That is: a type after the `extends` keyword.
+   */
+  Type getASuperType() { toGenerated(result) = cls.getChild(_) }
+}
+
+/**
+ * A `newtype Foo` declaration.
+ */
+class NewType extends TNewType, ModuleMember {
+  Generated::Datatype type;
+
+  NewType() { this = TNewType(type) }
+
+  string getName() { result = type.getName().getValue() }
+
+  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(_) }
+}
+
+/**
+ * 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()
+  }
+
+  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`. */
+class Formula extends TFormula, AstNode {
+  override AstNode getParent() {
+    result = super.getParent()
+    or
+    result.(Predicate).getBody() = this
+    or
+    result.(Aggregate).getGuard() = this
+  }
+}
+
+/** 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) }
+
+  ComparisonSymbol getSymbol() { result = op.getValue() }
+
+  override string getAPrimaryQlClass() { result = "ComparisonOp" }
+}
+
+class Literal extends TLiteral, Expr {
+  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() {
+    this = "=" or
+    this = "!=" or
+    this = "<" or
+    this = ">" or
+    this = "<=" or
+    this = ">="
+  }
+}
+
+class ComparisonFormula extends TComparisonFormula, Formula {
+  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" }
+}
+
+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)
+  }
+
+  /** 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.getExpr() }
+
+  /** 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 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;
+
+  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 { }
+
+/** 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` */
+class BinOpExpr extends TBinOpExpr, 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() }
 }

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

@@ -5,9 +5,83 @@ cached
 newtype TAstNode =
   TClasslessPredicate(Generated::ModuleMember member, Generated::ClasslessPredicate pred) {
     pred.getParent() = member
-  }
+  } or
+  TVarDecl(Generated::VarDecl decl) 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
+  TType(Generated::TypeExpr type) or
+  TDisjunction(Generated::Disjunction disj) or
+  TConjunction(Generated::Conjunction conj) or
+  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)
+
+class TFormula = TDisjunction or TConjunction or TComparisonFormula or TQuantifier or TNegation;
+
+class TBinOpExpr = TAddExpr;
+
+class TExpr = TBinOpExpr or TLiteral or TAggregate or TIdentifier;
+
+Generated::AstNode toGeneratedFormula(AST::AstNode n) {
+  n = TConjunction(result) or
+  n = TDisjunction(result) or
+  n = TComparisonFormula(result) or
+  n = TComparisonOp(result) or
+  n = TQuantifier(result) or
+  n = TAggregate(result) or
+  n = TIdentifier(result) or
+  n = TNegation(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) { n = TClasslessPredicate(_, result) }
+Generated::AstNode toGenerated(AST::AstNode n) {
+  result = toGeneratedExpr(n)
+  or
+  result = toGeneratedFormula(n)
+  or
+  n = TClasslessPredicate(_, result)
+  or
+  n = TVarDecl(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)
+  or
+  n = TType(result)
+  or
+  n = TLiteral(result)
+  or
+  n = TAsExpr(result)
+}
+
+class TPredicate = TCharPred or TClasslessPredicate or TClassPredicate;
+
+class TModuleMember = TClasslessPredicate or TClass or TModule or TNewType or TImport;

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" }
 
@@ -344,11 +349,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 +401,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,9 +766,20 @@ 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) }
+
     AstNode getChild(int i) { quantified_child(this, i, result) }
 
-    override AstNode getAFieldOrChild() { quantified_child(this, _, result) }
+    override AstNode getAFieldOrChild() {
+      quantified_expr(this, result) or
+      quantified_formula(this, result) or
+      quantified_range(this, result) or
+      quantified_child(this, _, result)
+    }
   }
 
   class Quantifier extends @token_quantifier, Token {
@@ -894,9 +925,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/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"
+}

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
 
 /**

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,7 +673,28 @@ qualified_expr_def(
   int loc: @location 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_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(
+  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 = @token_quantifier | @var_decl
 
 #keyset[quantified, index]
 quantified_child(
@@ -783,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(