refactor(rome_analyze): refactor how Visitors, Queryables and QueryMa…

…tches are related (#4063)

* refactor(rome_analyze): allow queryables to register their own visitor

* replace the `QueryMatch` enum with `Box<dyn Any>`

* store visitors in a BTreeMap to register them in a stable order

* fix CI checks

* move calculatiung the text range of query matches to lazy evaluation

* introduce a `Query` type to enforce the correct downcasting of query matches

* add documentation

crates/rome_analyze/CONTRIBUTING.md

@@ -387,6 +387,98 @@ If this the rule can retrieve its option with
 let options = ctx.options();
 ```
 
+#### Custom Visitors
+
+Some lint rules may need to deeply inspect the child nodes of a query match
+before deciding on whether they should emit a signal or not. These rules can be
+inefficient to implement using the query system, as they will lead to redundant
+traversal passes being executed over the same syntax tree. To make this more
+efficient, you can implement a custom `Queryable` type and and associated
+`Visitor` to emit it as part of the analyzer's main traversal pass. As an
+example, here's how this could be done to implement the `useYield` rule:
+
+```rust,ignore
+// First, create a visitor struct that holds a stack of function syntax nodes and booleans
+#[derive(Default)]
+struct MissingYieldVisitor {
+    stack: Vec<(AnyFunctionLike, bool)>,
+}
+
+// Implement the `Visitor` trait for this struct
+impl Visitor for MissingYieldVisitor {
+    fn visit(
+        &mut self,
+        event: &WalkEvent<SyntaxNode<Self::Language>>,
+        ctx: VisitorContext<Self::Language>,
+    ) {
+        match event {
+            WalkEvent::Enter(node) => {
+                // When the visitor enters a function node, push a new entry on the stack
+                if let Some(node) = AnyFunctionLike::cast_ref(node) {
+                    self.stack.push((node, false));
+                }
+                
+                if let Some((_, has_yield)) = self.stack.last_mut() {
+                    // When the visitor enters a `yield` expression, set the
+                    // `has_yield` flag for the top entry on the stack to `true`
+                    if JsYieldExpression::can_cast(node.kind()) {
+                        *has_yield = true;
+                    }
+                }
+            }
+            WalkEvent::Leave(node) => {
+                // When the visitor exits a function, if it matches the node of the top-most
+                // entry of the stack and the `has_yield` flag is `false`, emit a query match
+                if let Some(exit_node) = AnyFunctionLike::cast_ref(node) {
+                    if let Some((enter_node, has_yield)) = self.stack.pop() {
+                        assert_eq!(enter_node, exit_node);
+                        if !has_yield {
+                            ctx.match_query(MissingYield(enter_node));
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+// Declare a query match struct type containing a JavaScript function node
+struct MissingYield(AnyFunctionLike);
+
+// Implement the `Queryable` trait for this type
+impl Queryable for MissingYield {
+    // `Input` is the type that `ctx.match_query()` is called with in the visitor
+    type Input = Self;
+    // `Output` if the type that `ctx.query()` will return in the rule
+    type Output = AnyFunctionLike;
+    
+    fn build_visitor(
+        analyzer: &mut impl AddVisitor<Self::Language>,
+        _: &<Self::Language as Language>::Root,
+    ) {
+        // Register our custom visitor to run in the `Syntax` phase
+        analyzer.add_visitor(Phases::Syntax, MissingYieldVisitor::default());
+    }
+
+    // Extract the output object from the input type
+    fn unwrap_match(services: &ServiceBag, query: &Self::Input) -> Self::Output {
+        query.0.clone()
+    }
+}
+
+impl Rule for UseYield {
+    // Declare the custom `MissingYield` queryable as the rule's query
+    type Query = MissingYield;
+
+    fn run(ctx: &RuleContext<Self>) -> Self::Signals {
+        // Read the function's root node from the queryable output
+        let query: &AnyFunctionLike = ctx.query();
+
+        // ...
+    }
+}
+```
+
 # Using Just
 
 It is also possible to do all the steps above using our `Just` automation. For example, we can create

crates/rome_analyze/Cargo.toml

@@ -10,7 +10,6 @@ license = "MIT"
 
 [dependencies]
 rome_rowan = { path = "../rome_rowan" }
-rome_control_flow = { path = "../rome_control_flow" }
 rome_console = { path = "../rome_console" }
 rome_diagnostics = { path = "../rome_diagnostics" }
 bitflags = "1.3.2"

crates/rome_analyze/src/lib.rs

@@ -28,7 +28,7 @@ pub use crate::diagnostics::AnalyzerDiagnostic;
 use crate::diagnostics::SuppressionDiagnostic;
 pub use crate::matcher::{InspectMatcher, MatchQueryParams, QueryMatcher, RuleKey, SignalEntry};
 pub use crate::options::{AnalyzerConfiguration, AnalyzerOptions, AnalyzerRules};
-pub use crate::query::{Ast, QueryKey, QueryMatch, Queryable};
+pub use crate::query::{AddVisitor, QueryKey, QueryMatch, Queryable};
 pub use crate::registry::{
     LanguageRoot, MetadataRegistry, Phase, Phases, RegistryRuleMetadata, RegistryVisitor,
     RuleRegistry, RuleRegistryBuilder, RuleSuppressions,
@@ -40,7 +40,7 @@ pub use crate::rule::{
 pub use crate::services::{FromServices, MissingServicesDiagnostic, ServiceBag};
 use crate::signals::DiagnosticSignal;
 pub use crate::signals::{AnalyzerAction, AnalyzerSignal};
-pub use crate::syntax::SyntaxVisitor;
+pub use crate::syntax::{Ast, SyntaxVisitor};
 pub use crate::visitor::{NodeVisitor, Visitor, VisitorContext, VisitorFinishContext};
 pub use rule::DeserializableRuleOptions;
 
@@ -107,14 +107,13 @@ where
         }
     }
 
-    pub fn add_visitor<V>(&mut self, phase: Phases, visitor: V)
-    where
-        V: Visitor<Language = L> + 'analyzer,
-    {
-        self.phases
-            .entry(phase)
-            .or_default()
-            .push(Box::new(visitor));
+    /// Registers a [Visitor] to be executed as part of a given `phase` of the analyzer run
+    pub fn add_visitor(
+        &mut self,
+        phase: Phases,
+        visitor: Box<dyn Visitor<Language = L> + 'analyzer>,
+    ) {
+        self.phases.entry(phase).or_default().push(visitor);
     }
 
     pub fn run(self, mut ctx: AnalyzerContext<L>) -> Option<Break> {

crates/rome_analyze/src/matcher.rs

@@ -4,11 +4,16 @@ use crate::{
 };
 use rome_diagnostics::FileId;
 use rome_rowan::{Language, TextRange};
-use std::{cmp::Ordering, collections::BinaryHeap};
+use std::{
+    any::{Any, TypeId},
+    cmp::Ordering,
+    collections::BinaryHeap,
+};
 
 /// The [QueryMatcher] trait is responsible of running lint rules on
-/// [QueryMatch] instances emitted by the various [Visitor](crate::Visitor)
-/// and push signals wrapped in [SignalEntry] to the signal queue
+/// [QueryMatch](crate::QueryMatch) instances emitted by the various
+/// [Visitor](crate::Visitor) and push signals wrapped in [SignalEntry]
+/// to the signal queue
 pub trait QueryMatcher<L: Language> {
     /// Execute a single query match
     fn match_query(&mut self, params: MatchQueryParams<L>);
@@ -19,12 +24,52 @@ pub struct MatchQueryParams<'phase, 'query, L: Language> {
     pub phase: Phases,
     pub file_id: FileId,
     pub root: &'phase L::Root,
-    pub query: QueryMatch<L>,
+    pub query: Query,
     pub services: &'phase ServiceBag,
     pub signal_queue: &'query mut BinaryHeap<SignalEntry<'phase, L>>,
     pub apply_suppression_comment: SuppressionCommentEmitter<L>,
 }
 
+/// Wrapper type for a [QueryMatch]
+///
+/// This type is functionally equivalent to `Box<dyn QueryMatch + Any>`, it
+/// emulates dynamic dispatch for both traits and allows downcasting to a
+/// reference or owned type.
+pub struct Query {
+    data: Box<dyn Any>,
+    read_text_range: fn(&dyn Any) -> TextRange,
+}
+
+impl Query {
+    /// Construct a new [Query] instance from a [QueryMatch]
+    pub fn new<T: QueryMatch>(data: T) -> Self {
+        Self {
+            data: Box::new(data),
+            read_text_range: |query| query.downcast_ref::<T>().unwrap().text_range(),
+        }
+    }
+
+    /// Attempt to downcast the query to an owned type.
+    pub fn downcast<T: QueryMatch>(self) -> Option<T> {
+        Some(*self.data.downcast::<T>().ok()?)
+    }
+
+    /// Attempt to downcast the query to a reference type.
+    pub fn downcast_ref<T: QueryMatch>(&self) -> Option<&T> {
+        self.data.downcast_ref::<T>()
+    }
+
+    /// Returns the [TypeId] of this query, equivalent to calling [Any::type_id].
+    pub(crate) fn type_id(&self) -> TypeId {
+        self.data.as_ref().type_id()
+    }
+
+    /// Returns the [TextRange] of this query, equivalent to calling [QueryMatch::text_range].
+    pub(crate) fn text_range(&self) -> TextRange {
+        (self.read_text_range)(self.data.as_ref())
+    }
+}
+
 /// Opaque identifier for a group of rule
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct GroupKey {
@@ -158,14 +203,14 @@ mod tests {
     use rome_diagnostics::{Diagnostic, Severity};
     use rome_rowan::{
         raw_language::{RawLanguage, RawLanguageKind, RawLanguageRoot, RawSyntaxTreeBuilder},
-        AstNode, TextRange, TextSize, TriviaPiece, TriviaPieceKind,
+        AstNode, SyntaxNode, TextRange, TextSize, TriviaPiece, TriviaPieceKind,
     };
 
     use crate::SuppressionKind;
     use crate::{
         signals::DiagnosticSignal, Analyzer, AnalyzerContext, AnalyzerOptions, AnalyzerSignal,
-        ControlFlow, MetadataRegistry, Never, Phases, QueryMatch, QueryMatcher, RuleKey,
-        ServiceBag, SignalEntry, SyntaxVisitor,
+        ControlFlow, MetadataRegistry, Never, Phases, QueryMatcher, RuleKey, ServiceBag,
+        SignalEntry, SyntaxVisitor,
     };
 
     use super::MatchQueryParams;
@@ -184,10 +229,7 @@ mod tests {
     impl QueryMatcher<RawLanguage> for SuppressionMatcher {
         /// Emits a warning diagnostic for all literal expressions
         fn match_query(&mut self, params: MatchQueryParams<RawLanguage>) {
-            let node = match params.query {
-                QueryMatch::Syntax(node) => node,
-                _ => unreachable!(),
-            };
+            let node = params.query.downcast::<SyntaxNode<RawLanguage>>().unwrap();
 
             if node.kind() != RawLanguageKind::LITERAL_EXPRESSION {
                 return;
@@ -338,7 +380,7 @@ mod tests {
             &mut emit_signal,
         );
 
-        analyzer.add_visitor(Phases::Syntax, SyntaxVisitor::default());
+        analyzer.add_visitor(Phases::Syntax, Box::new(SyntaxVisitor::default()));
 
         let ctx: AnalyzerContext<RawLanguage> = AnalyzerContext {
             file_id: FileId::zero(),

crates/rome_analyze/src/query.rs

@@ -1,91 +1,58 @@
-use rome_control_flow::ControlFlowGraph;
-use rome_rowan::{AstNode, Language, SyntaxKindSet, SyntaxNode, TextRange};
+use std::any::TypeId;
 
-use crate::{
-    registry::{NodeLanguage, Phase},
-    services::FromServices,
-    ServiceBag,
-};
+use rome_rowan::{Language, SyntaxKindSet, TextRange};
 
-/// Trait implemented for all types, for example lint rules can query them to emit diagnostics or code actions.
+use crate::{registry::Phase, services::FromServices, Phases, ServiceBag, Visitor};
+
+/// Trait implemented for types that lint rules can query in order to emit diagnostics or code actions.
 pub trait Queryable: Sized {
+    type Input: QueryMatch;
     type Output;
+
     type Language: Language;
     type Services: FromServices + Phase;
 
-    /// Statically declares which [QueryMatch] variant is matched by this
-    /// [Queryable] type. For instance the [Ast] queryable matches on
-    /// [QueryMatch::Syntax], so its key is defined as [QueryKey::Syntax]
-    const KEY: QueryKey<Self::Language>;
+    /// Registers one or more [Visitor] that will emit `Self::Input` query
+    /// matches during the analyzer run
+    fn build_visitor(
+        analyzer: &mut impl AddVisitor<Self::Language>,
+        root: &<Self::Language as Language>::Root,
+    );
+
+    /// Returns the type of query matches this [Queryable] expects as inputs
+    ///
+    /// Unless your custom queryable needs to match on a specific
+    /// [SyntaxKindSet], you should not override the default implementation of
+    /// this method
+    fn key() -> QueryKey<Self::Language> {
+        QueryKey::TypeId(TypeId::of::<Self::Input>())
+    }
 
     /// Unwrap an instance of `Self` from a [QueryMatch].
     ///
     /// ## Panics
     ///
     /// If the [QueryMatch] variant of `query` doesn't match `Self::KEY`
-    fn unwrap_match(services: &ServiceBag, query: &QueryMatch<Self::Language>) -> Self::Output;
+    fn unwrap_match(services: &ServiceBag, query: &Self::Input) -> Self::Output;
 }
 
-/// Enumerate all the types of [Queryable] analyzer visitors may emit
-#[derive(Clone, Debug)]
-pub enum QueryMatch<L: Language> {
-    Syntax(SyntaxNode<L>),
-    SemanticModel(TextRange),
-    ControlFlowGraph(ControlFlowGraph<L>, TextRange),
+/// This trait is implemented on all types that supports the registration of [Visitor]
+pub trait AddVisitor<L: Language> {
+    /// Registers a [Visitor] for a given `phase`
+    fn add_visitor<F, V>(&mut self, phase: Phases, visitor: F)
+    where
+        F: FnOnce() -> V,
+        V: Visitor<Language = L> + 'static;
 }
 
-impl<L: Language> QueryMatch<L> {
-    pub fn text_range(&self) -> TextRange {
-        match self {
-            QueryMatch::Syntax(node) => node.text_trimmed_range(),
-            QueryMatch::SemanticModel(range) | QueryMatch::ControlFlowGraph(_, range) => *range,
-        }
-    }
+/// Marker trait implemented for all the types analyzer visitors may emit
+pub trait QueryMatch: 'static {
+    fn text_range(&self) -> TextRange;
 }
 
-/// Mirrors the variants of [QueryMatch] to statically compute which queries a
-/// given [Queryable] type can match
+/// Represents which type a given [Queryable] type can match, either a specific
+/// subset of syntax node kinds or any generic type
 pub enum QueryKey<L: Language> {
     Syntax(SyntaxKindSet<L>),
-    ControlFlowGraph,
-    SemanticModel,
-}
-
-/// Query type usable by lint rules to match on specific [AstNode] types
-#[derive(Clone)]
-pub struct Ast<N>(pub N);
-
-impl<N> Queryable for Ast<N>
-where
-    N: AstNode + 'static,
-{
-    type Output = N;
-    type Language = NodeLanguage<N>;
-    type Services = ();
-
-    /// Match on [QueryMatch::Syntax] if the kind of the syntax node matches
-    /// the kind set of `N`
-    const KEY: QueryKey<Self::Language> = QueryKey::Syntax(N::KIND_SET);
-
-    fn unwrap_match(_: &ServiceBag, query: &QueryMatch<Self::Language>) -> Self::Output {
-        match query {
-            QueryMatch::Syntax(node) => N::unwrap_cast(node.clone()),
-            _ => panic!("tried to unwrap unsupported QueryMatch kind, expected Syntax"),
-        }
-    }
-}
-
-impl<L: Language> Queryable for ControlFlowGraph<L> {
-    type Output = Self;
-    type Language = L;
-    type Services = ();
-
-    const KEY: QueryKey<Self::Language> = QueryKey::ControlFlowGraph;
-
-    fn unwrap_match(_: &ServiceBag, query: &QueryMatch<Self::Language>) -> Self::Output {
-        match query {
-            QueryMatch::ControlFlowGraph(cfg, _) => cfg.clone(),
-            _ => panic!("tried to unwrap unsupported QueryMatch kind, expected Syntax"),
-        }
-    }
+    TypeId(TypeId),
 }