feat(semantic): add ability to lookup if AST contains any node kinds

Author: camchenry

crates/oxc_ast/src/generated/ast_kind.rs

@@ -11,6 +11,9 @@ use oxc_span::{GetSpan, Span};
 
 use crate::ast::*;
 
+/// The largest integer value that can be mapped to an `AstType`/`AstKind` enum variant.
+pub const AST_TYPE_MAX: u8 = 186;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 #[repr(u8)]
 pub enum AstType {

crates/oxc_semantic/src/ast_types_bitset.rs

@@ -0,0 +1,116 @@
+use oxc_ast::{AstType, ast_kind::AST_TYPE_MAX};
+
+const USIZE_BITS: usize = usize::BITS as usize;
+
+/// Number of bytes required for bit set which can represent all [`AstType`]s.
+// Need to add plus one here because 0 is a possible value, but requires at least one bit to represent it.
+const NUM_USIZES: usize = (AST_TYPE_MAX as usize + 1).div_ceil(USIZE_BITS);
+
+/// Bit set with a bit for each [`AstType`].
+#[derive(Debug, Clone)]
+pub struct AstTypesBitset([usize; NUM_USIZES]);
+
+impl AstTypesBitset {
+    /// Create empty [`AstTypesBitset`] with no bits set.
+    pub const fn new() -> Self {
+        Self([0; NUM_USIZES])
+    }
+
+    /// Create a new [`AstTypesBitset`] from a slice of [`AstType`].
+    pub fn from_types(types: &[AstType]) -> Self {
+        let mut bitset = Self::new();
+        for &ty in types {
+            bitset.set(ty);
+        }
+        bitset
+    }
+
+    /// Returns `true` if bit is set for provided [`AstType`].
+    pub const fn has(&self, ty: AstType) -> bool {
+        let (index, mask) = Self::index_and_mask(ty);
+        (self.0[index] & mask) != 0
+    }
+
+    /// Set bit for provided [`AstType`].
+    pub const fn set(&mut self, ty: AstType) {
+        let (index, mask) = Self::index_and_mask(ty);
+        self.0[index] |= mask;
+    }
+
+    /// Returns `true` if any bit is set in both `self` and `other`.
+    pub fn intersects(&self, other: &Self) -> bool {
+        let mut intersection = 0;
+        for (&a, &b) in self.0.iter().zip(other.0.iter()) {
+            intersection |= a & b;
+        }
+        intersection != 0
+    }
+
+    /// Returns `true` if all bits in `other` are set in `self`.
+    pub fn contains(&self, other: &Self) -> bool {
+        let mut mismatches = 0;
+        for (&a, &b) in self.0.iter().zip(other.0.iter()) {
+            let set_in_both = a & b;
+            // 0 if `set_in_both == b`
+            let mismatch = set_in_both ^ b;
+            mismatches |= mismatch;
+        }
+        mismatches == 0
+    }
+
+    /// Get index and mask for an [`AstType`].
+    /// Returned `index` is guaranteed not to be out of bounds of the array.
+    const fn index_and_mask(ty: AstType) -> (usize, usize) {
+        let n = ty as usize;
+        let index = n / USIZE_BITS;
+        let mask = 1usize << (n % USIZE_BITS);
+        (index, mask)
+    }
+}
+
+impl Default for AstTypesBitset {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use oxc_ast::AstType;
+
+    #[test]
+    fn empty_bitset_has_no_bits_and_contains_empty() {
+        let bs = AstTypesBitset::new();
+        assert!(!bs.has(AstType::Program));
+        let other = AstTypesBitset::new();
+        assert!(bs.contains(&other), "Empty bitset should contain empty bitset");
+        assert!(!bs.intersects(&other));
+    }
+
+    #[test]
+    fn intersects_and_contains() {
+        let mut a = AstTypesBitset::from_types(&[AstType::Program, AstType::AssignmentPattern]);
+        let b = AstTypesBitset::from_types(&[AstType::TSTupleType]);
+        assert!(!a.intersects(&b));
+        a.set(AstType::TSTupleType);
+        assert!(a.intersects(&b));
+
+        let c = AstTypesBitset::from_types(&[AstType::Program]);
+        assert!(a.contains(&c));
+        assert!(!c.contains(&a));
+
+        // a should contain union of subset bits
+        let subset = AstTypesBitset::from_types(&[AstType::Program, AstType::AssignmentPattern]);
+        assert!(a.contains(&subset));
+        // subset does not contain a (missing TSTupleType)
+        assert!(!subset.contains(&a));
+    }
+
+    #[test]
+    fn contains_empty_is_true() {
+        let non_empty = AstTypesBitset::from_types(&[AstType::IdentifierName]);
+        let empty = AstTypesBitset::new();
+        assert!(non_empty.contains(&empty));
+    }
+}

crates/oxc_semantic/src/lib.rs

@@ -22,6 +22,7 @@ pub use oxc_syntax::{
 
 pub mod dot;
 
+mod ast_types_bitset;
 mod binder;
 mod builder;
 mod checker;

crates/oxc_semantic/src/node.rs

@@ -1,7 +1,7 @@
 use std::iter::FusedIterator;
 
 use oxc_allocator::{Address, GetAddress};
-use oxc_ast::{AstKind, ast::Program};
+use oxc_ast::{AstKind, AstType, ast::Program};
 use oxc_cfg::BlockNodeId;
 use oxc_index::{IndexSlice, IndexVec};
 use oxc_span::{GetSpan, Span};
@@ -10,6 +10,8 @@ use oxc_syntax::{
     scope::ScopeId,
 };
 
+use crate::ast_types_bitset::AstTypesBitset;
+
 /// Semantic node contains all the semantic information about an ast node.
 #[derive(Debug, Clone, Copy)]
 pub struct AstNode<'a> {
@@ -101,6 +103,10 @@ pub struct AstNodes<'a> {
     nodes: IndexVec<NodeId, AstNode<'a>>,
     /// `node` -> `parent`
     parent_ids: IndexVec<NodeId, NodeId>,
+    /// Stores a set of bits of a fixed size, where each bit represents a single [`AstKind`]. If the bit is set (1),
+    /// then the AST contains at least one node of that kind. If the bit is not set (0), then the AST does not contain
+    /// any nodes of that kind.
+    pub node_kinds_set: AstTypesBitset,
 }
 
 impl<'a> AstNodes<'a> {
@@ -212,6 +218,7 @@ impl<'a> AstNodes<'a> {
         let node_id = self.parent_ids.push(parent_node_id);
         let node = AstNode::new(kind, scope_id, cfg_id, flags, node_id);
         self.nodes.push(node);
+        self.node_kinds_set.set(kind.ty());
         node_id
     }
 
@@ -234,6 +241,7 @@ impl<'a> AstNodes<'a> {
         );
         self.parent_ids.push(NodeId::ROOT);
         self.nodes.push(AstNode::new(kind, scope_id, cfg_id, flags, NodeId::ROOT));
+        self.node_kinds_set.set(AstType::Program);
         NodeId::ROOT
     }
 
@@ -242,6 +250,51 @@ impl<'a> AstNodes<'a> {
         self.nodes.reserve(additional);
         self.parent_ids.reserve(additional);
     }
+
+    /// Checks if the AST contains any nodes of the given types.
+    ///
+    /// Example:
+    /// ```ignore
+    /// let for_stmt = AstTypesBitset::from_types(&[AstType::ForStatement]);
+    /// let import_export_decl = AstTypesBitset::from_types(&[AstType::ImportDeclaration, AstType::ExportDeclaration]);
+    ///
+    /// // returns true if there is a `for` loop anywhere in the AST.
+    /// nodes.contains_any(&for_stmt)
+    /// // returns true if there is at least one import OR one export in the AST.
+    /// nodes.contains_any(&import_export_decl)
+    /// ```
+    pub fn contains_any(&self, bitset: &AstTypesBitset) -> bool {
+        self.node_kinds_set.intersects(bitset)
+    }
+
+    /// Checks if the AST contains all of the given types.
+    ///
+    /// Example:
+    /// ```ignore
+    /// let for_stmt = AstTypesBitset::from_types(&[AstType::ForStatement]);
+    /// let import_export_decl = AstTypesBitset::from_types(&[AstType::ImportDeclaration, AstType::ExportDeclaration]);
+    ///
+    /// // returns true if there is a `for` loop anywhere in the AST.
+    /// nodes.contains_all(&for_stmt)
+    /// // returns true only if there is at least one import AND one export in the AST.
+    /// nodes.contains_all(&import_export_decl)
+    /// ```
+    pub fn contains_all(&self, bitset: &AstTypesBitset) -> bool {
+        self.node_kinds_set.contains(bitset)
+    }
+
+    /// Checks if the AST contains a node of the given type.
+    ///
+    /// Example:
+    /// ```ignore
+    /// // returns true if there is a `for` loop anywhere in the AST.
+    /// nodes.contains(AstType::ForStatement)
+    /// // returns true if there is an `ImportDeclaration` anywhere in the AST.
+    /// nodes.contains(AstType::ImportDeclaration)
+    /// ```
+    pub fn contains(&self, ty: AstType) -> bool {
+        self.node_kinds_set.has(ty)
+    }
 }
 
 impl<'a, 'n> IntoIterator for &'n AstNodes<'a> {

tasks/ast_tools/src/generators/ast_kind.rs

@@ -149,6 +149,8 @@ impl Generator for AstKindGenerator {
             next_index += 1;
         }
 
+        let ast_type_max = number_lit(next_index - 1);
+
         let output = quote! {
             #![expect(missing_docs)] ///@ FIXME (in ast_tools/src/generators/ast_kind.rs)
 
@@ -162,6 +164,10 @@ impl Generator for AstKindGenerator {
             ///@@line_break
             use crate::ast::*;
 
+            ///@@line_break
+            /// The largest integer value that can be mapped to an `AstType`/`AstKind` enum variant.
+            pub const AST_TYPE_MAX: u8 = #ast_type_max;
+
             ///@@line_break
             #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
             #[repr(u8)]