From f2edd7724c2fc632b845bdc1b9e909917c7f9480 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Wed, 26 May 2021 13:14:52 +0000
Subject: [PATCH] Record annotation names in the AST

---
 Cargo.lock                                   |  2 +-
 extractor/Cargo.toml                         |  2 +-
 generator/Cargo.toml                         |  2 +-
 ql/src/codeql_ql/ast/internal/TreeSitter.qll | 14 +++++++++++---
 ql/src/ql.dbscheme                           |  2 ++
 5 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 39742029999..62f1da345fe 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=064e438e92a89d99a2b061eca39e0ea225e32341#064e438e92a89d99a2b061eca39e0ea225e32341"
+source = "git+https://github.com/tausbn/tree-sitter-ql.git?rev=dee83b7a1c159f2c71d8d3d3bce3ee1203ed653a#dee83b7a1c159f2c71d8d3d3bce3ee1203ed653a"
 dependencies = [
  "cc",
  "tree-sitter",
diff --git a/extractor/Cargo.toml b/extractor/Cargo.toml
index 6fd657be513..723b5b95ce1 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 = "064e438e92a89d99a2b061eca39e0ea225e32341" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "dee83b7a1c159f2c71d8d3d3bce3ee1203ed653a" }
 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 28b6097dd59..3eab976ab87 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 = "064e438e92a89d99a2b061eca39e0ea225e32341" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "dee83b7a1c159f2c71d8d3d3bce3ee1203ed653a" }
diff --git a/ql/src/codeql_ql/ast/internal/TreeSitter.qll b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
index a77bd9c552a..768fe85f81e 100644
--- a/ql/src/codeql_ql/ast/internal/TreeSitter.qll
+++ b/ql/src/codeql_ql/ast/internal/TreeSitter.qll
@@ -350,11 +350,19 @@ module Generated {
   class ExprAnnotation extends @expr_annotation, AstNode {
     override string getAPrimaryQlClass() { result = "ExprAnnotation" }
 
-    override Location getLocation() { expr_annotation_def(this, _, result) }
+    override Location getLocation() { expr_annotation_def(this, _, _, _, result) }
 
-    AstNode getChild() { expr_annotation_def(this, result, _) }
+    AnnotName getAnnotArg() { expr_annotation_def(this, result, _, _, _) }
 
-    override AstNode getAFieldOrChild() { expr_annotation_def(this, result, _) }
+    AnnotName getName() { expr_annotation_def(this, _, result, _, _) }
+
+    AstNode getChild() { expr_annotation_def(this, _, _, result, _) }
+
+    override AstNode getAFieldOrChild() {
+      expr_annotation_def(this, result, _, _, _) or
+      expr_annotation_def(this, _, result, _, _) or
+      expr_annotation_def(this, _, _, result, _)
+    }
   }
 
   class False extends @token_false, Token {
diff --git a/ql/src/ql.dbscheme b/ql/src/ql.dbscheme
index 48bc35f70a2..e334a7a2fc4 100644
--- a/ql/src/ql.dbscheme
+++ b/ql/src/ql.dbscheme
@@ -323,6 +323,8 @@ expr_aggregate_body_def(
 
 expr_annotation_def(
   unique int id: @expr_annotation,
+  int annot_arg: @token_annot_name ref,
+  int name: @token_annot_name ref,
   int child: @expr_annotation_child_type ref,
   int loc: @location ref
 );