Ruby: generate overlay discard predicates

ruby/ql/lib/codeql/ruby/ast/internal/TreeSitter.qll

@@ -5,6 +5,11 @@
 
 import codeql.Locations as L
 
+/** Holds if and only if the database is an overlay. */
+overlay[local]
+pragma[nomagic]
+private predicate isOverlay() { databaseMetadata("isOverlay", "true") }
+
 module Ruby {
   /** The base class for all AST nodes */
   class AstNode extends @ruby_ast_node {
@@ -48,6 +53,36 @@ module Ruby {
     final override string getAPrimaryQlClass() { result = "ReservedWord" }
   }
 
+  /** Gets the file containing the given `node`. */
+  overlay[local]
+  pragma[nomagic]
+  private @file getNodeFile(@ruby_ast_node node) {
+    exists(@location_default loc | ruby_ast_node_location(node, loc) |
+      locations_default(loc, result, _, _, _, _)
+    )
+  }
+
+  /** Holds if `file` was extracted as part of the overlay database. */
+  overlay[local]
+  pragma[nomagic]
+  private predicate discardFile(@file file) {
+    isOverlay() and exists(@ruby_ast_node node | file = getNodeFile(node))
+  }
+
+  /** Holds if `node` is in the `file` and is part of the overlay base database. */
+  overlay[local]
+  pragma[nomagic]
+  private predicate discardableAstNode(@file file, @ruby_ast_node node) {
+    not isOverlay() and file = getNodeFile(node)
+  }
+
+  /** Holds if `node` should be discarded, because it is part of the overlay base and is in a file that was also extracted as part of the overlay database. */
+  overlay[discard_entity]
+  pragma[nomagic]
+  private predicate discardAstNode(@ruby_ast_node node) {
+    exists(@file file | discardableAstNode(file, node) and discardFile(file))
+  }
+
   class UnderscoreArg extends @ruby_underscore_arg, AstNode { }
 
   class UnderscoreCallOperator extends @ruby_underscore_call_operator, AstNode { }
@@ -1970,6 +2005,36 @@ module Erb {
     final override string getAPrimaryQlClass() { result = "ReservedWord" }
   }
 
+  /** Gets the file containing the given `node`. */
+  overlay[local]
+  pragma[nomagic]
+  private @file getNodeFile(@erb_ast_node node) {
+    exists(@location_default loc | erb_ast_node_location(node, loc) |
+      locations_default(loc, result, _, _, _, _)
+    )
+  }
+
+  /** Holds if `file` was extracted as part of the overlay database. */
+  overlay[local]
+  pragma[nomagic]
+  private predicate discardFile(@file file) {
+    isOverlay() and exists(@erb_ast_node node | file = getNodeFile(node))
+  }
+
+  /** Holds if `node` is in the `file` and is part of the overlay base database. */
+  overlay[local]
+  pragma[nomagic]
+  private predicate discardableAstNode(@file file, @erb_ast_node node) {
+    not isOverlay() and file = getNodeFile(node)
+  }
+
+  /** Holds if `node` should be discarded, because it is part of the overlay base and is in a file that was also extracted as part of the overlay database. */
+  overlay[discard_entity]
+  pragma[nomagic]
+  private predicate discardAstNode(@erb_ast_node node) {
+    exists(@file file | discardableAstNode(file, node) and discardFile(file))
+  }
+
   /** A class representing `code` tokens. */
   class Code extends @erb_token_code, Token {
     /** Gets the name of the primary QL class for this element. */

shared/tree-sitter-extractor/src/generator/mod.rs

@@ -17,7 +17,7 @@ pub fn generate(
     languages: Vec<language::Language>,
     dbscheme_path: PathBuf,
     ql_library_path: PathBuf,
-    add_metadata_relation: bool,
+    overlay_support: bool,
 ) -> std::io::Result<()> {
     let dbscheme_file = File::create(dbscheme_path).map_err(|e| {
         tracing::error!("Failed to create dbscheme file: {}", e);
@@ -35,7 +35,7 @@ pub fn generate(
 
     // Eventually all languages will have the metadata relation (for overlay support), at which
     // point this could be moved to prefix.dbscheme.
-    if add_metadata_relation {
+    if overlay_support {
         writeln!(dbscheme_writer, "/*- Database metadata -*/",)?;
         dbscheme::write(
             &mut dbscheme_writer,
@@ -60,6 +60,15 @@ pub fn generate(
         })],
     )?;
 
+    if overlay_support {
+        ql::write(
+            &mut ql_writer,
+            &[ql::TopLevel::Predicate(
+                ql_gen::create_is_overlay_predicate(),
+            )],
+        )?;
+    }
+
     for language in languages {
         let prefix = node_types::to_snake_case(&language.name);
         let ast_node_name = format!("{}_ast_node", &prefix);
@@ -103,6 +112,22 @@ pub fn generate(
             ql::TopLevel::Class(ql_gen::create_token_class(&token_name, &tokeninfo_name)),
             ql::TopLevel::Class(ql_gen::create_reserved_word_class(&reserved_word_name)),
         ];
+
+        if overlay_support {
+            body.push(ql::TopLevel::Predicate(
+                ql_gen::create_get_node_file_predicate(&ast_node_name, &node_location_table_name),
+            ));
+            body.push(ql::TopLevel::Predicate(
+                ql_gen::create_discard_file_predicate(&ast_node_name),
+            ));
+            body.push(ql::TopLevel::Predicate(
+                ql_gen::create_discardable_ast_node_predicate(&ast_node_name),
+            ));
+            body.push(ql::TopLevel::Predicate(
+                ql_gen::create_discard_ast_node_predicate(&ast_node_name),
+            ));
+        }
+
         body.append(&mut ql_gen::convert_nodes(&nodes));
         ql::write(
             &mut ql_writer,

shared/tree-sitter-extractor/src/generator/ql.rs

@@ -6,6 +6,7 @@ pub enum TopLevel<'a> {
     Class(Class<'a>),
     Import(Import<'a>),
     Module(Module<'a>),
+    Predicate(Predicate<'a>),
 }
 
 impl fmt::Display for TopLevel<'_> {
@@ -14,6 +15,7 @@ impl fmt::Display for TopLevel<'_> {
             TopLevel::Import(imp) => write!(f, "{}", imp),
             TopLevel::Class(cls) => write!(f, "{}", cls),
             TopLevel::Module(m) => write!(f, "{}", m),
+            TopLevel::Predicate(pred) => write!(f, "{}", pred),
         }
     }
 }
@@ -68,10 +70,13 @@ impl fmt::Display for Class<'_> {
                     qldoc: None,
                     name: self.name,
                     overridden: false,
+                    is_private: false,
                     is_final: false,
                     return_type: None,
                     formal_parameters: vec![],
                     body: charpred.clone(),
+                    pragma: None,
+                    overlay: None,
                 }
             )?;
         }
@@ -150,6 +155,7 @@ pub enum Expression<'a> {
         expr: Box<Expression<'a>>,
         second_expr: Option<Box<Expression<'a>>>,
     },
+    Negation(Box<Expression<'a>>),
 }
 
 impl fmt::Display for Expression<'_> {
@@ -231,26 +237,59 @@ impl fmt::Display for Expression<'_> {
                 }
                 write!(f, ")")
             }
+            Expression::Negation(e) => write!(f, "not ({})", e),
         }
     }
 }
 
+#[derive(Clone, Eq, PartialEq, Hash)]
+pub enum PragmaAnnotation {
+    Nomagic,
+}
+
+#[derive(Clone, Eq, PartialEq, Hash)]
+pub enum OverlayAnnotation {
+    Local,
+    DiscardEntity,
+}
+
 #[derive(Clone, Eq, PartialEq, Hash)]
 pub struct Predicate<'a> {
     pub qldoc: Option<String>,
     pub name: &'a str,
     pub overridden: bool,
+    pub is_private: bool,
     pub is_final: bool,
     pub return_type: Option<Type<'a>>,
     pub formal_parameters: Vec<FormalParameter<'a>>,
     pub body: Expression<'a>,
+    pub pragma: Option<PragmaAnnotation>,
+    pub overlay: Option<OverlayAnnotation>,
 }
 
 impl fmt::Display for Predicate<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if let Some(qldoc) = &self.qldoc {
             write!(f, "/** {} */", qldoc)?;
         }
+        if let Some(overlay_annotation) = &self.overlay {
+            write!(f, "overlay[")?;
+            match overlay_annotation {
+                OverlayAnnotation::Local => write!(f, "local")?,
+                OverlayAnnotation::DiscardEntity => write!(f, "discard_entity")?,
+            }
+            write!(f, "] ")?;
+        }
+        if let Some(pragma) = &self.pragma {
+            write!(f, "pragma[")?;
+            match pragma {
+                PragmaAnnotation::Nomagic => write!(f, "nomagic")?,
+            }
+            write!(f, "] ")?;
+        }
+        if self.is_private {
+            write!(f, "private ")?;
+        }
         if self.is_final {
             write!(f, "final ")?;
         }