refactor(lexer): mark AVX2 and Neon versions with cpu_version (#28669)

Part of effort to move to runtime CPU feature detection.

Mark the AVX2 and Neon versions of functions which `cpu_select!` picks between, and the helpers which only they use, with `#[apply(cpu_version(Avx2))]` or `#[apply(cpu_version(Neon))]`. 31 functions in all. The macro adds the `#[cfg(target_arch = ...)]` and `#[target_feature(...)]` attributes, so the hand-written ones are removed (14 `#[target_feature]`s and 7 `#[cfg]`s).

- `array_map` and `lic_first` had no CPU type parameter. They now take `_cpu: Avx2`, and their callers pass it. - `init_pair_luts_avx2` had no `#[target_feature]`, and now gets one from the macro. So its callers call it in an `unsafe` block, and it calls `PairLuts::new` without one. `PairLuts::new`'s `# SAFETY` section, which explained that `unsafe` block, is removed. - Remove the temporary `#[expect(unused_*)]`s which step G5 added.

`qcompact` and `bcompact` aren't marked. They're `const fn`s which only build `static` tables at compile time. Nor is `uninit_array_mut`, a function nested inside `PairLuts::new`. `get_pair_luts` keeps `cpu_versions` for now, because changing it changes codegen.

Codegen: no change in any build. Panic line numbers in `lanes.rs` change.

crates/oxc_lexer/src/cpu/mod.rs

@@ -60,9 +60,7 @@ mod versions_macro;
 mod tests;
 
 pub(crate) use select_macro::cpu_select;
-#[cfg_attr(not(test), expect(unused_imports))]
 pub use version_macro::__cpu_version;
-#[cfg_attr(not(test), expect(unused_imports))]
 pub(crate) use version_macro::cpu_version;
 pub(crate) use versions_macro::cpu_versions;
 pub use versions_macro::{__cpu_versions, apply};

crates/oxc_lexer/src/cpu/version_macro.rs

@@ -1,5 +1,3 @@
-#![cfg_attr(not(test), expect(unused_macros, unused_imports,))]
-
 /// Macro to mark a function as the version of some code for one CPU feature set.
 ///
 /// [`cpu_select!`] picks between the versions.

crates/oxc_lexer/src/pipeline/bitmap.rs

@@ -1,7 +1,7 @@
 #[cfg(target_arch = "x86_64")]
 use std::arch::x86_64::*;
 
-use crate::cpu::{Cpu, cpu_select};
+use crate::cpu::{Cpu, apply, cpu_select, cpu_version};
 
 #[cfg(target_arch = "x86_64")]
 use crate::cpu::Avx2;
@@ -157,9 +157,8 @@ pub(super) unsafe fn bm_any(bm: *const u64, nw: usize, cpu: impl Cpu) -> bool {
 /// # SAFETY
 ///
 /// Same as [`bm_any`].
-#[cfg(target_arch = "x86_64")]
+#[apply(cpu_version(Avx2))]
 #[inline]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
 unsafe fn bm_any_avx2(bm: *const u64, nw: usize, _cpu: Avx2) -> bool {
     let mut w = 0usize;
     while w + 4 <= nw {

crates/oxc_lexer/src/pipeline/carve/jsx/hyphens.rs

@@ -1,4 +1,4 @@
-use crate::cpu::{Cpu, cpu_select};
+use crate::cpu::{Cpu, apply, cpu_select, cpu_version};
 
 use crate::pipeline::bitmap::{bm_clear, bm_get};
 
@@ -35,7 +35,7 @@ pub(super) unsafe fn jsx_glue_hyphens(
 }
 
 /// [`jsx_glue_hyphens`] for AVX2.
-#[cfg(target_arch = "x86_64")]
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
 unsafe fn jsx_glue_hyphens_avx2(
     src: *const u8,

crates/oxc_lexer/src/pipeline/chunk.rs

@@ -2,7 +2,7 @@
 mod primitives_avx2 {
     use std::arch::x86_64::*;
 
-    use crate::cpu::Avx2;
+    use crate::cpu::{Avx2, apply, cpu_version};
 
     /// Load 256 bits starting at byte `i` as an `__m256i`.
     ///
@@ -11,6 +11,7 @@ mod primitives_avx2 {
     /// # SAFETY
     ///
     /// `src.add(i)` must be valid for reads of 32 bytes.
+    #[apply(cpu_version(Avx2))]
     #[inline(always)]
     pub unsafe fn load256(src: *const u8, i: usize, _cpu: Avx2) -> __m256i {
         _mm256_loadu_si256(src.add(i) as *const __m256i)
@@ -20,6 +21,7 @@ mod primitives_avx2 {
     ///
     /// Returns a `__m256i` where each of the 32 lanes is `0xFF`
     /// where the lane in `v` equals `c`, and `0x00` where it does not.
+    #[apply(cpu_version(Avx2))]
     #[inline(always)]
     pub fn veq(v: __m256i, c: u8, _cpu: Avx2) -> __m256i {
         // SAFETY: These intrinsics touch no memory. `Avx2` type guarantees CPU supports AVX2.
@@ -29,6 +31,7 @@ mod primitives_avx2 {
     /// Gather the top bit of each byte of chunk `v` into a 32-bit mask.
     ///
     /// Bit N of the result is the top bit of byte N in `v`.
+    #[apply(cpu_version(Avx2))]
     #[inline(always)]
     pub fn mm(v: __m256i, _cpu: Avx2) -> u32 {
         // SAFETY: This intrinsic touches no memory. `Avx2` type guarantees CPU supports AVX2.

crates/oxc_lexer/src/pipeline/classify/avx2.rs

@@ -1,6 +1,9 @@
 use std::arch::x86_64::*;
 
-use crate::{cpu::Avx2, token::tk};
+use crate::{
+    cpu::{Avx2, apply, cpu_version},
+    token::tk,
+};
 
 use crate::pipeline::chunk::{load256, mm, veq};
 
@@ -9,8 +12,8 @@ use super::{
     punct1_luts::{PH_A, PH_B, PH_T0, PH_T1},
 };
 
+#[apply(cpu_version(Avx2))]
 #[inline(never)]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
 pub(super) unsafe fn classify_avx2(
     ts: bool,
     src: *const u8,
@@ -142,6 +145,7 @@ pub(super) unsafe fn classify_avx2(
     }
 }
 
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
 fn vpunct1(
     v: __m256i,

crates/oxc_lexer/src/pipeline/classify/neon.rs

@@ -1,6 +1,9 @@
 use std::arch::aarch64::*;
 
-use crate::{cpu::Neon, token::tk};
+use crate::{
+    cpu::{Neon, apply, cpu_version},
+    token::tk,
+};
 
 use super::{
     luts::{MERGED_LUTS, WORD_LUTS},
@@ -29,8 +32,8 @@ struct Classified {
     kind: uint8x16_t,
 }
 
+#[apply(cpu_version(Neon))]
 #[inline(never)]
-#[target_feature(enable = "neon")]
 pub(super) unsafe fn classify_neon(
     ts: bool,
     src: *const u8,
@@ -106,6 +109,7 @@ pub(super) unsafe fn classify_neon(
     }
 }
 
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 fn chunk(v: uint8x16_t, l: &Luts, cpu: Neon) -> Classified {
     // SAFETY: `Neon` type guarantees required target features are present
@@ -140,6 +144,7 @@ fn chunk(v: uint8x16_t, l: &Luts, cpu: Neon) -> Classified {
     }
 }
 
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 fn punct1(v: uint8x16_t, lo: uint8x16_t, hn: uint8x16_t, l: &Luts, _cpu: Neon) -> uint8x16_t {
     // SAFETY: `Neon` type guarantees required target features are present
@@ -152,12 +157,14 @@ fn punct1(v: uint8x16_t, lo: uint8x16_t, hn: uint8x16_t, l: &Luts, _cpu: Neon) -
     }
 }
 
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 fn veq(v: uint8x16_t, c: u8, _cpu: Neon) -> uint8x16_t {
     // SAFETY: `Neon` type guarantees required target features are present
     unsafe { vceqq_u8(v, vdupq_n_u8(c)) }
 }
 
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 fn lut(t: &[u8; 16], _cpu: Neon) -> uint8x16_t {
     // SAFETY: `Neon` type guarantees required target features are present.
@@ -166,6 +173,7 @@ fn lut(t: &[u8; 16], _cpu: Neon) -> uint8x16_t {
 }
 
 /// Bit i is the top bit of byte i of a-d, gathered by summing the lane bit values.
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 fn mm64(a: uint8x16_t, b: uint8x16_t, c: uint8x16_t, d: uint8x16_t, cpu: Neon) -> u64 {
     const BITS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];

crates/oxc_lexer/src/pipeline/compress/avx2/mod.rs

@@ -3,7 +3,7 @@ use std::arch::x86_64::*;
 use oxc_span::Span;
 
 use crate::{
-    cpu::Avx2,
+    cpu::{Avx2, apply, cpu_version},
     token::{TRIVIA_MAX, TRIVIA_MIN, is_trivia_byte, tk},
 };
 
@@ -56,8 +56,8 @@ const fn bcompact() -> [[u8; 16]; 256] {
     t
 }
 
+#[apply(cpu_version(Avx2))]
 #[inline(never)]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
 pub(super) unsafe fn compress_blocks_avx2(
     st: *const u64,
     kind: *const u8,
@@ -116,7 +116,7 @@ pub(super) unsafe fn compress_blocks_avx2(
     m
 }
 
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 pub(super) unsafe fn build_spans_avx2(
     stage_kind: *const u8,
     stage_pos: *const u32,
@@ -181,7 +181,7 @@ pub(super) unsafe fn build_spans_avx2(
     w
 }
 
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 pub(super) unsafe fn lanes_post_avx2(
     src: &[u8],
     out_kinds: *const u8,

crates/oxc_lexer/src/pipeline/compress/avx2/pair_luts.rs

@@ -6,7 +6,7 @@ use std::{
     sync::OnceLock,
 };
 
-use crate::cpu::{Avx2, Cpu, apply, cpu_versions};
+use crate::cpu::{Avx2, Cpu, apply, cpu_version, cpu_versions};
 
 /// Single copy of [`PairLuts`], shared across all threads.
 ///
@@ -49,13 +49,7 @@ impl PairLuts {
     /// to zero both arrays, then all the data being copied twice - first from
     /// 2 stack temporaries for `lut0z` and `lutpad` arrays into another stack temporary
     /// for `PairLuts`, then a second copy of `PairLuts` from the stack to the heap.
-    ///
-    /// # SAFETY
-    ///
-    /// No safety preconditions.
-    /// Calling this function requires an `unsafe { ... }` block, because of `#[target_feature]`,
-    /// but it's always sound, because the `Avx2` parameter guarantees the CPU supports those features.
-    #[target_feature(enable = "avx2,bmi2,popcnt")]
+    #[apply(cpu_version(Avx2))]
     fn new(_cpu: Avx2) -> Box<Self> {
         // Allocate space for `PairLuts` on the heap, uninitialized
         let mut luts = Box::<PairLuts>::new_uninit();
@@ -105,11 +99,9 @@ impl PairLuts {
 /// Initialize the global [`PairLuts`] instance.
 ///
 /// This function must be called before calling [`get_pair_luts`].
+#[apply(cpu_version(Avx2))]
 pub fn init_pair_luts_avx2(cpu: Avx2) {
-    PAIR_LUTS.get_or_init(|| {
-        // SAFETY: `PairLuts::new` has no safety preconditions
-        unsafe { PairLuts::new(cpu) }
-    });
+    PAIR_LUTS.get_or_init(|| PairLuts::new(cpu));
 }
 
 /// Get reference to [`PairLuts`].

crates/oxc_lexer/src/pipeline/compress/avx2/tests.rs

@@ -10,7 +10,8 @@ fn compress_blocks_matches_scalar_reference() {
     let DynCpu::Avx2(cpu) = DynCpu::detect() else { return };
 
     // Required before calling `compress_blocks_avx2`
-    init_pair_luts_avx2(cpu);
+    // SAFETY: `Avx2` type guarantees required target features are present
+    unsafe { init_pair_luts_avx2(cpu) };
 
     let mut cases: Vec<Vec<u64>> = Vec::new();
     let all_pairs: Vec<u64> = (0..65536u64)

crates/oxc_lexer/src/pipeline/compress/mod.rs

@@ -74,7 +74,10 @@ pub unsafe fn write_sentinels(n: u32, spans: *mut Span, sig_kinds: *mut u8, _cpu
 /// `PairLuts` is only used by AVX2 implementation, so this function is a no-op on other CPUs.
 pub(super) fn init_pair_luts(cpu: DynCpu) {
     cpu_select!(cpu, {
-        Avx2(cpu) => init_pair_luts_avx2(cpu),
+        Avx2(cpu) => {
+            // SAFETY: `Avx2` type guarantees required target features are present
+            unsafe { init_pair_luts_avx2(cpu) }
+        }
         Neon(_) => {}
         Generic(_) => {}
     });

crates/oxc_lexer/src/pipeline/find/avx2.rs

@@ -3,7 +3,7 @@ use std::{
     mem::{self, MaybeUninit},
 };
 
-use crate::cpu::Avx2;
+use crate::cpu::{Avx2, apply, cpu_version};
 
 use crate::pipeline::chunk::{load256, mm, veq};
 
@@ -13,8 +13,8 @@ const BRACKET_BYTES: [u8; 6] = [b'(', b')', b'[', b'{', b']', b'}'];
 /// First byte at or after i that is one of needles, or n.
 ///
 /// Needle order matters for speed. See comment on [`any_of`] below.
+#[apply(cpu_version(Avx2))]
 #[inline]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
 pub unsafe fn find_avx2<const K: usize>(
     src: *const u8,
     n: usize,
@@ -39,8 +39,8 @@ pub unsafe fn find_avx2<const K: usize>(
 }
 
 /// Bits of the 64 bytes at `base` that are brackets (`(){}[]`): bit `i` for byte `base + i`.
+#[apply(cpu_version(Avx2))]
 #[inline]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
 pub fn bracket_bits_avx2(src: &[u8], base: usize, cpu: Avx2) -> u64 {
     let block = &src[base..base + 64];
     let mut out = 0u64;
@@ -70,9 +70,10 @@ pub fn bracket_bits_avx2(src: &[u8], base: usize, cpu: Avx2) -> u64 {
 // LLVM only does this when the 2 compares are OR-ed together directly,
 // i.e. when the bytes are at indexes `2 * j` and `2 * j + 1` in `needles`.
 // The byte sets passed to `find` are ordered to make such pairs.
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
 unsafe fn any_of<const N: usize>(v: __m256i, needles: [u8; N], cpu: Avx2) -> __m256i {
-    let mut m = array_map(needles, |c| veq(v, c, cpu));
+    let mut m = array_map(needles, |c| veq(v, c, cpu), cpu);
     let mut len = N;
     while len > 1 {
         let half = len / 2;
@@ -91,8 +92,9 @@ unsafe fn any_of<const N: usize>(v: __m256i, needles: [u8; N], cpu: Avx2) -> __m
 ///
 /// `map` is `std` code, which is compiled without AVX2.
 /// If LLVM doesn't inline it into the caller, it calls the intrinsics out of line.
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
-fn array_map<T, U, const N: usize>(array: [T; N], f: impl Fn(T) -> U) -> [U; N] {
+fn array_map<T, U, const N: usize>(array: [T; N], f: impl Fn(T) -> U, _cpu: Avx2) -> [U; N] {
     // If `f` panics, the values already in `out` aren't dropped. Ensure that can't leak memory.
     const { assert!(!mem::needs_drop::<U>()) };
 

crates/oxc_lexer/src/pipeline/keywords.rs

@@ -37,7 +37,7 @@ use constcat::concat_slices;
 use oxc_data_structures::{assert_unchecked, branch_hints::unlikely, str::const_str_eq};
 
 use crate::{
-    cpu::{Cpu, apply, cpu_select, cpu_versions},
+    cpu::{Cpu, apply, cpu_select, cpu_version, cpu_versions},
     token::{KW_KIND_BASE, TokenKind},
 };
 
@@ -511,7 +511,7 @@ unsafe fn text_matches(ptr: *const u8, len: usize, text: &[u64; 2], cpu: impl Cp
 /// # SAFETY
 ///
 /// Same as [`text_matches`].
-#[cfg(target_arch = "x86_64")]
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
 unsafe fn text_matches_avx2(ptr: *const u8, len: usize, text: &[u64; 2], _cpu: Avx2) -> bool {
     use std::arch::x86_64::{_mm_and_si128, _mm_cmpeq_epi8, _mm_loadu_si128, _mm_movemask_epi8};
@@ -529,7 +529,7 @@ unsafe fn text_matches_avx2(ptr: *const u8, len: usize, text: &[u64; 2], _cpu: A
 /// # SAFETY
 ///
 /// Same as [`text_matches`].
-#[cfg(target_arch = "aarch64")]
+#[apply(cpu_version(Neon))]
 #[inline(always)]
 unsafe fn text_matches_neon(ptr: *const u8, len: usize, text: &[u64; 2], _cpu: Neon) -> bool {
     use std::arch::aarch64::{vandq_u8, veorq_u8, vgetq_lane_u64, vld1q_u8, vreinterpretq_u64_u8};

crates/oxc_lexer/src/pipeline/lanes.rs

@@ -12,7 +12,7 @@ use oxc_ast::ast::{Comment, RegExpFlags};
 use oxc_syntax::identifier::{is_identifier_part, is_identifier_start};
 
 use crate::{
-    cpu::{Cpu, apply, cpu_select, cpu_versions},
+    cpu::{Cpu, apply, cpu_select, cpu_version, cpu_versions},
     error::{DiagCode, DiagSeverity, Diagnostic},
     token::StringSpan,
 };
@@ -466,8 +466,7 @@ fn cook_short<const EMIT: bool, const CRLF: bool>(
 }
 
 /// [`cook_short`] for AVX2.
-#[cfg(target_arch = "x86_64")]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 #[inline]
 fn cook_short_avx2<const EMIT: bool, const CRLF: bool>(
     src: &[u8],
@@ -540,8 +539,7 @@ fn span_has_bs(src: &[u8], bs: usize, be: usize, cpu: impl Cpu) -> bool {
 }
 
 /// [`span_has_bs`] for AVX2.
-#[cfg(target_arch = "x86_64")]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 #[inline]
 fn span_has_bs_avx2(src: &[u8], bs: usize, be: usize, _cpu: Avx2) -> bool {
     let mut i = bs;
@@ -582,8 +580,7 @@ fn span_has_bs_or_cr(src: &[u8], bs: usize, be: usize, cpu: impl Cpu) -> bool {
 }
 
 /// [`span_has_bs_or_cr`] for AVX2.
-#[cfg(target_arch = "x86_64")]
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 #[inline]
 fn span_has_bs_or_cr_avx2(src: &[u8], bs: usize, be: usize, _cpu: Avx2) -> bool {
     let mut i = bs;

crates/oxc_lexer/src/pipeline/scan/avx2.rs

@@ -1,12 +1,12 @@
 use std::arch::x86_64::*;
 
-use crate::cpu::Avx2;
+use crate::cpu::{Avx2, apply, cpu_version};
 
 use crate::pipeline::chunk::{load256, mm, veq};
 
 use super::common::lic_verify_at;
 
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 pub unsafe fn scan_block_comment_avx2(
     src: *const u8,
     n: usize,
@@ -32,14 +32,14 @@ pub unsafe fn scan_block_comment_avx2(
             if lic_q < 0 {
                 let am = at & bodymask;
                 if am != 0 {
-                    lic_q = lic_first(src, i, am);
+                    lic_q = lic_first(src, i, am, cpu);
                 }
             }
             return (i + j + 1, saw_nl, lic_q);
         }
         saw_nl |= nl != 0;
         if lic_q < 0 && at != 0 {
-            lic_q = lic_first(src, i, at);
+            lic_q = lic_first(src, i, at, cpu);
         }
         i += 32;
     }
@@ -64,7 +64,7 @@ pub unsafe fn scan_block_comment_avx2(
 /// metadata. A 0xE2 that isn't LS/PS (typographic punctuation in prose) is
 /// cleared from the mask and the scan continues; the LS/PS confirm can read
 /// the pad, which never matches 0x80.
-#[target_feature(enable = "avx2,bmi2,popcnt")]
+#[apply(cpu_version(Avx2))]
 pub unsafe fn scan_line_comment_avx2(
     src: *const u8,
     n: usize,
@@ -93,13 +93,13 @@ pub unsafe fn scan_line_comment_avx2(
             if lic_q < 0 {
                 let am = at & if t > 0 { (1u32 << t) - 1 } else { 0 };
                 if am != 0 {
-                    lic_q = lic_first(src, i, am);
+                    lic_q = lic_first(src, i, am, cpu);
                 }
             }
             return (i + t, lic_q);
         }
         if lic_q < 0 && at != 0 {
-            lic_q = lic_first(src, i, at);
+            lic_q = lic_first(src, i, at, cpu);
         }
         i += 32;
     }
@@ -122,8 +122,9 @@ pub unsafe fn scan_line_comment_avx2(
     (n, lic_q)
 }
 
+#[apply(cpu_version(Avx2))]
 #[inline(always)]
-unsafe fn lic_first(src: *const u8, base: usize, mut am: u32) -> i64 {
+unsafe fn lic_first(src: *const u8, base: usize, mut am: u32, _cpu: Avx2) -> i64 {
     while am != 0 {
         let q = base + am.trailing_zeros() as usize;
         am &= am - 1;