perf: exclude nodes that don't need to be queried

Author: ematipico

crates/biome_analyze/src/registry.rs

@@ -1,20 +1,23 @@
 use crate::{
     AddVisitor, AnalysisFilter, GroupCategory, QueryMatcher, Rule, RuleGroup, RuleKey,
-    RuleMetadata, ServiceBag, SignalEntry, Visitor,
+    RuleMetadata, ServiceBag, SignalEntry, SyntaxVisitor, Visitor,
     context::RuleContext,
     matcher::{GroupKey, MatchQueryParams},
     query::{QueryKey, Queryable},
     signals::RuleSignal,
 };
 use biome_diagnostics::Error;
-use biome_rowan::{AstNode, Language, RawSyntaxKind, SyntaxKind, SyntaxNode};
+use biome_rowan::{AstNode, Language, RawSyntaxKind, SyntaxKind, SyntaxKindSet, SyntaxNode};
 use rustc_hash::{FxHashMap, FxHashSet};
 use std::{
     any::TypeId,
     borrow,
     collections::{BTreeMap, BTreeSet},
 };
 
+/// Number of phases supported by the [RuleRegistry].
+const PHASE_COUNT: usize = 2;
+
 /// Defines all the phases that the [RuleRegistry] supports.
 #[repr(usize)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -98,7 +101,7 @@ impl<L: Language> RegistryVisitor<L> for MetadataRegistry {
 ///   after the "SemanticModel" is ready, which demands a whole transverse of the parsed tree.
 pub struct RuleRegistry<L: Language> {
     /// Holds a collection of rules for each phase.
-    phase_rules: [PhaseRules<L>; 2],
+    phase_rules: [PhaseRules<L>; PHASE_COUNT],
 }
 
 impl<L: Language + Default> RuleRegistry<L> {
@@ -115,6 +118,7 @@ impl<L: Language + Default> RuleRegistry<L> {
             visitors: BTreeMap::default(),
             services: ServiceBag::default(),
             diagnostics: Vec::new(),
+            syntax_kind_sets: [SyntaxKindSet::empty(); PHASE_COUNT],
         }
     }
 }
@@ -143,6 +147,10 @@ pub struct RuleRegistryBuilder<'a, L: Language> {
     // Service Bag
     services: ServiceBag,
     diagnostics: Vec<Error>,
+    /// Per-phase union of all [SyntaxKindSet] values from registered rules.
+    /// Used to create a filtered [SyntaxVisitor] that only emits query matches
+    /// for node kinds that at least one rule cares about.
+    syntax_kind_sets: [SyntaxKindSet<L>; PHASE_COUNT],
 }
 
 impl<L: Language + Default + 'static> RegistryVisitor<L> for RuleRegistryBuilder<'_, L> {
@@ -167,13 +175,15 @@ impl<L: Language + Default + 'static> RegistryVisitor<L> for RuleRegistryBuilder
             return;
         }
 
-        let phase = R::phase() as usize;
-        let phase = &mut self.registry.phase_rules[phase];
+        let phase_index = R::phase() as usize;
+        let phase = &mut self.registry.phase_rules[phase_index];
 
         let rule = RegistryRule::new::<R>(phase.rule_states.len());
 
         match <R::Query as Queryable>::key() {
             QueryKey::Syntax(key) => {
+                self.syntax_kind_sets[phase_index] = self.syntax_kind_sets[phase_index].union(key);
+
                 let TypeRules::SyntaxRules { rules } = phase
                     .type_rules
                     .entry(TypeId::of::<SyntaxNode<L>>())
@@ -242,8 +252,25 @@ type BuilderResult<L> = (
     BTreeMap<(Phases, TypeId), Box<dyn Visitor<Language = L>>>,
 );
 
-impl<L: Language> RuleRegistryBuilder<'_, L> {
-    pub fn build(self) -> BuilderResult<L> {
+impl<L: Language + 'static> RuleRegistryBuilder<'_, L> {
+    pub fn build(mut self) -> BuilderResult<L> {
+        // Replace SyntaxVisitor instances with filtered versions that only
+        // emit query matches for node kinds that at least one rule queries.
+        // This avoids per-node Box::new() allocations for kinds no rule
+        // cares about.
+        let phases = [Phases::Syntax, Phases::Semantic];
+        for (phase_index, &phase) in phases.iter().enumerate() {
+            let key = (phase, TypeId::of::<SyntaxVisitor<L>>());
+            if self.visitors.contains_key(&key) {
+                self.visitors.insert(
+                    key,
+                    Box::new(SyntaxVisitor::with_kind_set(
+                        self.syntax_kind_sets[phase_index],
+                    )),
+                );
+            }
+        }
+
         (
             self.registry,
             self.services,

crates/biome_analyze/src/syntax.rs

@@ -1,4 +1,4 @@
-use biome_rowan::{AstNode, Language, SyntaxNode, WalkEvent};
+use biome_rowan::{AstNode, Language, SyntaxKindSet, SyntaxNode, WalkEvent};
 
 use crate::{
     AddVisitor, Phases, QueryKey, QueryMatch, Queryable, ServiceBag, Visitor, VisitorContext,
@@ -50,11 +50,31 @@ pub struct SyntaxVisitor<L: Language> {
     /// of that subtree. The visitor will then ignore all events until it
     /// receives a [WalkEvent::Leave] for the `skip_subtree` node
     skip_subtree: Option<SyntaxNode<L>>,
+    /// Optional filter to restrict which [SyntaxKind] values are emitted as
+    /// query matches. When [None], all nodes are emitted (unfiltered).
+    /// When [Some], only nodes whose kind is in the set are emitted,
+    /// avoiding unnecessary allocations for nodes no rule cares about.
+    kind_set: Option<SyntaxKindSet<L>>,
+}
+
+impl<L: Language> SyntaxVisitor<L> {
+    /// Create a [SyntaxVisitor] that only emits query matches for nodes whose
+    /// `SyntaxKind` is in the provided set. Nodes with kinds not in the set
+    /// are skipped entirely, avoiding per-node allocation overhead.
+    pub fn with_kind_set(kind_set: SyntaxKindSet<L>) -> Self {
+        Self {
+            skip_subtree: None,
+            kind_set: Some(kind_set),
+        }
+    }
 }
 
 impl<L: Language> Default for SyntaxVisitor<L> {
     fn default() -> Self {
-        Self { skip_subtree: None }
+        Self {
+            skip_subtree: None,
+            kind_set: None,
+        }
     }
 }
 
@@ -86,6 +106,12 @@ impl<L: Language + 'static> Visitor for SyntaxVisitor<L> {
             return;
         }
 
+        if let Some(kind_set) = self.kind_set
+            && !kind_set.matches(node.kind())
+        {
+            return;
+        }
+
         ctx.match_query(node.clone());
     }
 }

crates/biome_rowan/src/ast/mod.rs

@@ -24,19 +24,30 @@ use crate::{
 pub use batch::*;
 pub use mutation::{AstNodeExt, AstNodeListExt, AstSeparatedListExt};
 
+/// Number of `u128` buckets used by [SyntaxKindSet] to represent the bitfield.
+/// This determines the maximum [RawSyntaxKind] value that can be stored:
+/// `SYNTAX_KIND_SET_BUCKETS * 128 - 1`. All language syntax kind enums must
+/// have their `__LAST` variant within this limit.
+pub const SYNTAX_KIND_SET_BUCKETS: usize = 5;
+
 /// Represents a set of [SyntaxKind] as a bitfield, with each bit representing
 /// whether the corresponding [RawSyntaxKind] value is contained in the set
 ///
 /// This is similar to the `TokenSet` struct in `biome_js_parser`, with the
 /// bitfield here being twice as large as it needs to cover all nodes as well
 /// as all token kinds
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub struct SyntaxKindSet<L: Language>([u128; 5], PhantomData<L>);
+pub struct SyntaxKindSet<L: Language>([u128; SYNTAX_KIND_SET_BUCKETS], PhantomData<L>);
 
 impl<L> SyntaxKindSet<L>
 where
     L: Language,
 {
+    /// Create an empty [SyntaxKindSet] containing no [SyntaxKind] values
+    pub const fn empty() -> Self {
+        Self([0; SYNTAX_KIND_SET_BUCKETS], PhantomData)
+    }
+
     /// Create a new [SyntaxKindSet] containing only the provided [SyntaxKind]
     pub fn of(kind: L::Kind) -> Self {
         Self::from_raw(kind.to_raw())
@@ -60,7 +71,7 @@ where
         let shift = kind % u128::BITS as u16;
         let mask = 1 << shift;
 
-        let mut bits = [0; 5];
+        let mut bits = [0; SYNTAX_KIND_SET_BUCKETS];
         bits[index] = mask;
 
         Self(bits, PhantomData)