perf(linter): jest/prefer-hooks-in-order: rewrite rule to allocate less and iterate fewer times (#6030)

Profiling this in a single-threaded run showed this rule to be relatively slow compared to other rules:

<img width="907" alt="image" src="https://github.com/user-attachments/assets/789cdc1f-d1ce-4d4a-bfa3-5882bd92cfc4">

The profile shows that a lot of time is spent doing pushes on `Vec`:
<img width="896" alt="image" src="https://github.com/user-attachments/assets/cde361d7-091d-42db-abc2-76a0c8a7fa44">

I believe this is because we were essentially duplicating the entirety of `ctx.nodes()`, and then iterating over it again for the actual rule.

I addressed this by:

- We no longer store any nodes (we only store the previous hook span + order), which saves many allocations.
- We only iterate over the nodes once and store the previous hook function in a per-scope hash map. This should be faster since we now only do `n` iterations, instead of `2n` iterations.

Benchmarking on the `microsoft/vscode` repository shows that this is probably faster (by up to ~3%), and probably not any slower:

<img width="992" alt="Screenshot 2024-09-24 at 12 06 32 PM" src="https://github.com/user-attachments/assets/ab916464-1b50-48e1-a65d-f9b3a4b4607d">

The snapshot change in this PR is due to slightly different ordering, I think, but it appears to be the same diagnostics overall.

crates/oxc_linter/src/rules/jest/prefer_hooks_in_order.rs

@@ -1,7 +1,7 @@
-use oxc_ast::{ast::CallExpression, AstKind};
+use oxc_ast::AstKind;
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
-use oxc_semantic::{AstNode, ScopeId};
+use oxc_semantic::ScopeId;
 use oxc_span::Span;
 use rustc_hash::FxHashMap;
 
@@ -13,10 +13,10 @@ use crate::{
     },
 };
 
-fn reorder_hooks(x1: &str, x2: &str, span: Span) -> OxcDiagnostic {
+fn reorder_hooks(hook: (&str, Span), previous_hook: (&str, Span)) -> OxcDiagnostic {
     OxcDiagnostic::warn("Prefer having hooks in a consistent order.")
-        .with_help(format!("{x1:?} hooks should be before any {x2:?} hooks"))
+        .with_help(format!("{:?} hooks should be before any {:?} hooks", hook.0, previous_hook.0))
-        .with_label(span)
+        .with_label(hook.1)
 }
 
 #[derive(Debug, Default, Clone)]
@@ -140,60 +140,66 @@ declare_oxc_lint!(
 
 impl Rule for PreferHooksInOrder {
     fn run_once(&self, ctx: &LintContext) {
-        let mut hook_groups: FxHashMap<ScopeId, Vec<AstNode>> = FxHashMap::default();
+        let mut previous_hook_orders: FxHashMap<ScopeId, (u8, Span)> = FxHashMap::default();
 
         for node in ctx.nodes() {
-            hook_groups.entry(node.scope_id()).or_default().push(*node);
-        }
-
-        for (_, nodes) in hook_groups {
-            let mut previous_hook_index = -1;
-
-            for node in nodes {
             if let AstKind::CallExpression(call_expr) = node.kind() {
-                    let possible_jest_node = &PossibleJestNode { node: &node, original: None };
+                let possible_jest_node = &PossibleJestNode { node, original: None };
-                    Self::check(&mut previous_hook_index, possible_jest_node, call_expr, ctx);
+                let Some(ParsedJestFnCallNew::GeneralJest(jest_fn_call)) =
+                    parse_jest_fn_call(call_expr, possible_jest_node, ctx)
+                else {
+                    previous_hook_orders.remove(&node.scope_id());
+                    continue;
                 };
+
+                if !matches!(jest_fn_call.kind, JestFnKind::General(JestGeneralFnKind::Hook)) {
+                    previous_hook_orders.remove(&node.scope_id());
+                    continue;
                 }
+
+                let previous_hook_order = previous_hook_orders.get(&node.scope_id());
+
+                let hook_name = jest_fn_call.name.as_ref();
+                let Some(hook_order) = get_hook_order(hook_name) else {
+                    continue;
+                };
+
+                if let Some((previous_hook_order, previous_hook_span)) = previous_hook_order {
+                    if hook_order < *previous_hook_order {
+                        let Some(previous_hook_name) = get_hook_name(*previous_hook_order) else {
+                            continue;
+                        };
+
+                        ctx.diagnostic(reorder_hooks(
+                            (hook_name, call_expr.span),
+                            (previous_hook_name, *previous_hook_span),
+                        ));
+                        continue;
+                    }
+                }
+                previous_hook_orders.insert(node.scope_id(), (hook_order, call_expr.span));
+            };
         }
     }
 }
 
-impl PreferHooksInOrder {
+fn get_hook_order(hook_name: &str) -> Option<u8> {
-    fn check<'a>(
+    match hook_name {
-        previous_hook_index: &mut i32,
+        "beforeAll" => Some(0),
-        possible_jest_node: &PossibleJestNode<'a, '_>,
+        "beforeEach" => Some(1),
-        call_expr: &'a CallExpression<'_>,
+        "afterEach" => Some(2),
-        ctx: &LintContext<'a>,
+        "afterAll" => Some(3),
-    ) {
+        _ => None,
-        let Some(ParsedJestFnCallNew::GeneralJest(jest_fn_call)) =
-            parse_jest_fn_call(call_expr, possible_jest_node, ctx)
-        else {
-            *previous_hook_index = -1;
-            return;
-        };
-
-        if !matches!(jest_fn_call.kind, JestFnKind::General(JestGeneralFnKind::Hook)) {
-            *previous_hook_index = -1;
-            return;
     }
+}
 
-        let hook_orders = ["beforeAll", "beforeEach", "afterEach", "afterAll"];
+fn get_hook_name(hook_order: u8) -> Option<&'static str> {
-        let hook_name = jest_fn_call.name;
+    match hook_order {
-        let hook_pos =
+        0 => Some("beforeAll"),
-            hook_orders.iter().position(|h| h.eq_ignore_ascii_case(&hook_name)).unwrap_or_default();
+        1 => Some("beforeEach"),
-        let previous_hook_pos = usize::try_from(*previous_hook_index).unwrap_or(0);
+        2 => Some("afterEach"),
-
+        3 => Some("afterAll"),
-        if hook_pos < previous_hook_pos {
+        _ => None,
-            let Some(previous_hook_name) = hook_orders.get(previous_hook_pos) else {
-                return;
-            };
-
-            ctx.diagnostic(reorder_hooks(&hook_name, previous_hook_name, call_expr.span));
-            return;
-        }
-
-        *previous_hook_index = i32::try_from(hook_pos).unwrap_or(-1);
     }
 }
 

crates/oxc_linter/src/snapshots/prefer_hooks_in_order.snap

@@ -1,6 +1,5 @@
 ---
 source: crates/oxc_linter/src/tester.rs
-assertion_line: 216
 ---
   ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
    ╭─[prefer_hooks_in_order.tsx:6:21]
@@ -185,16 +184,6 @@ assertion_line: 216
     ╰────
   help: "afterEach" hooks should be before any "afterAll" hooks
 
-  ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
-    ╭─[prefer_hooks_in_order.tsx:38:25]
- 37 │     
- 38 │ ╭─▶                         beforeEach(() => {
- 39 │ │                               mockLogger();
- 40 │ ╰─▶                         });
- 41 │     
-    ╰────
-  help: "beforeEach" hooks should be before any "afterEach" hooks
-
   ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
     ╭─[prefer_hooks_in_order.tsx:7:21]
   6 │     
@@ -205,6 +194,16 @@ assertion_line: 216
     ╰────
   help: "beforeAll" hooks should be before any "beforeEach" hooks
 
+  ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
+    ╭─[prefer_hooks_in_order.tsx:38:25]
+ 37 │     
+ 38 │ ╭─▶                         beforeEach(() => {
+ 39 │ │                               mockLogger();
+ 40 │ ╰─▶                         });
+ 41 │     
+    ╰────
+  help: "beforeEach" hooks should be before any "afterEach" hooks
+
   ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
    ╭─[prefer_hooks_in_order.tsx:6:17]
  5 │                     });
@@ -388,16 +387,6 @@ assertion_line: 216
     ╰────
   help: "afterEach" hooks should be before any "afterAll" hooks
 
-  ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
-    ╭─[prefer_hooks_in_order.tsx:38:21]
- 37 │     
- 38 │ ╭─▶                     beforeEach(() => {
- 39 │ │                           mockLogger();
- 40 │ ╰─▶                     });
- 41 │                     
-    ╰────
-  help: "beforeEach" hooks should be before any "afterEach" hooks
-
   ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
     ╭─[prefer_hooks_in_order.tsx:7:17]
   6 │     
@@ -407,3 +396,13 @@ assertion_line: 216
  10 │     
     ╰────
   help: "beforeAll" hooks should be before any "beforeEach" hooks
+
+  ⚠ eslint-plugin-vitest(prefer-hooks-in-order): Prefer having hooks in a consistent order.
+    ╭─[prefer_hooks_in_order.tsx:38:21]
+ 37 │     
+ 38 │ ╭─▶                     beforeEach(() => {
+ 39 │ │                           mockLogger();
+ 40 │ ╰─▶                     });
+ 41 │                     
+    ╰────
+  help: "beforeEach" hooks should be before any "afterEach" hooks