From 20e72d33044a73f0dd2520407f85cf8f7f76fc33 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Thu, 27 May 2021 12:52:37 +0000
Subject: [PATCH] Add AST node for expression annotations

---
 ql/src/codeql_ql/ast/Ast.qll               | 14 ++++++++++++++
 ql/src/codeql_ql/ast/internal/AstNodes.qll |  4 +++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index a9400e3b07b..52bb7eb420f 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -857,6 +857,20 @@ class Expr extends TExpr, AstNode {
   }
 }
 
+class ExprAnnotation extends TExprAnnotation, Expr {
+  Generated::ExprAnnotation expr_anno;
+
+  ExprAnnotation() { this = TExprAnnotation(expr_anno) }
+
+  string getName() { result = expr_anno.getName().getValue() }
+
+  string getAnnotationArgument() { result = expr_anno.getAnnotArg().getValue() }
+
+  Expr getExpression() { toGenerated(result) = expr_anno.getChild() }
+
+  override string getAPrimaryQlClass() { result = "ExprAnnotation" }
+}
+
 /** A function symbol, such as `+` or `*`. */
 class FunctionSymbol extends string {
   FunctionSymbol() { this = "+" or this = "-" or this = "*" or this = "/" or this = "%" }
diff --git a/ql/src/codeql_ql/ast/internal/AstNodes.qll b/ql/src/codeql_ql/ast/internal/AstNodes.qll
index 2a22d8af2fb..ff5e50a34ec 100644
--- a/ql/src/codeql_ql/ast/internal/AstNodes.qll
+++ b/ql/src/codeql_ql/ast/internal/AstNodes.qll
@@ -40,6 +40,7 @@ newtype TAstNode =
   TIfFormula(Generated::IfTerm ifterm) or
   TImplication(Generated::Implication impl) or
   TInstanceOf(Generated::InstanceOf inst) or
+  TExprAnnotation(Generated::ExprAnnotation expr_anno) or
   TAddExpr(Generated::AddExpr addexp) or
   TLiteral(Generated::Literal lit) or
   TUnaryExpr(Generated::UnaryExpr unaryexpr) or
@@ -55,7 +56,7 @@ class TBinOpExpr = TAddExpr;
 
 class TExpr =
   TBinOpExpr or TLiteral or TAggregate or TIdentifier or TInlineCast or TCall or TUnaryExpr or
-      TDontCare;
+      TExprAnnotation or TDontCare;
 
 class TCall = TPredicateCall or TMemberCall or TNoneCall or TAnyCall;
 
@@ -77,6 +78,7 @@ Generated::AstNode toGeneratedFormula(AST::AstNode n) {
 
 Generated::AstNode toGeneratedExpr(AST::AstNode n) {
   n = TAddExpr(result) or
+  n = TExprAnnotation(result) or
   n = TLiteral(result) or
   n = TAggregate(result) or
   n = TIdentifier(result) or