feat(lexer): add SIMD optimization to the lexer

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+ "rust-analyzer.showUnlinkedFileNotification": false
+}

Cargo.toml

@@ -137,6 +137,7 @@ tsify = { version = "0.4.5" }
 wasm-bindgen = { version = "0.2" }
 serde-wasm-bindgen = { version = "0.6.3" }
 
+
 [profile.release.package.oxc_wasm]
 opt-level = 'z'
 

crates/oxc_parser/src/lexer/identifier.rs

@@ -1,29 +1,36 @@
 use super::{
  cold_branch,
- search::{byte_search, safe_byte_match_table, SafeByteMatchTable},
+ search::{byte_search, simd_byte_match_table, SimdByteMatchTable, SEARCH_BATCH_SIZE},
  Kind, Lexer, SourcePosition,
 };
 use crate::diagnostics;
 
-use std::cmp::max;
-
 use oxc_allocator::String;
 use oxc_span::Span;
 use oxc_syntax::identifier::{
  is_identifier_part, is_identifier_part_unicode, is_identifier_start_unicode,
 };
+use std::{borrow::Cow, cmp::max};
 
 const MIN_ESCAPED_STR_LEN: usize = 16;
 
-static ASCII_ID_START_TABLE: SafeByteMatchTable =
- safe_byte_match_table!(|b| b.is_ascii_alphabetic() || b == b'_' || b == b'$');
+static ASCII_ID_START_TABLE: SimdByteMatchTable =
+ simd_byte_match_table!(|b| b.is_ascii_alphabetic() || b == b'_' || b == b'$', false);
 
-static NOT_ASCII_ID_CONTINUE_TABLE: SafeByteMatchTable =
- safe_byte_match_table!(|b| !(b.is_ascii_alphanumeric() || b == b'_' || b == b'$'));
+static NOT_ASCII_ID_CONTINUE_TABLE: SimdByteMatchTable =
+ simd_byte_match_table!(|b| !(b.is_ascii_alphanumeric() || b == b'_' || b == b'$'), true);
 
 #[inline]
-fn is_identifier_start_ascii_byte(byte: u8) -> bool {
- ASCII_ID_START_TABLE.matches(byte)
+fn is_identifier_start_ascii_byte(data: Option<(Cow<[u8; SEARCH_BATCH_SIZE]>, usize)>) -> bool {
+ let data = match data {
+ Some(data) => data,
+ None => return false,
+ };
+ let mut iter = ASCII_ID_START_TABLE.matches(data.0.as_ref(), data.1);
+ match iter.next() {
+ Some((offset, _)) => offset == 0,
+ None => false,
+ }
 }
 
 impl<'a> Lexer<'a> {
@@ -224,10 +231,11 @@ impl<'a> Lexer<'a> {
  });
  }
 
+ let pos = self.source.position();
  // Handle if not an ASCII identifier byte.
  // SAFETY: Not at EOF, so safe to read a byte.
- let b = unsafe { start_pos.read() };
- if !is_identifier_start_ascii_byte(b) {
+ let data = unsafe { pos.peek_n_with_padding::<SEARCH_BATCH_SIZE>(self.source.end_addr()) };
+ if !is_identifier_start_ascii_byte(data) {
  return self.private_identifier_not_ascii_id();
  }
 

crates/oxc_parser/src/lexer/mod.rs

@@ -17,6 +17,7 @@ mod numeric;
 mod punctuation;
 mod regex;
 mod search;
+mod simd;
 mod source;
 mod string;
 mod string_builder;

crates/oxc_parser/src/lexer/search.rs

@@ -4,8 +4,52 @@
 //! * `byte_match_table!` and `safe_byte_match_table!` macros create those tables at compile time.
 //! * `byte_search!` macro searches source text for first byte matching a byte table.
 
+use super::simd;
+
 /// Batch size for searching
-pub const SEARCH_BATCH_SIZE: usize = 32;
+pub const SEARCH_BATCH_SIZE: usize = simd::ALIGNMENT;
+
+pub struct SimdByteMatchTable(simd::MatchTable);
+
+#[allow(dead_code)]
+impl SimdByteMatchTable {
+ // Create new `SimdByteMatchTable`.
+ pub const fn new(bytes: [bool; 256], reverse: bool) -> Self {
+ Self(simd::MatchTable::new(bytes, reverse))
+ }
+
+ /// Declare that using this table for searching.
+ /// An unsafe function here, whereas for `SafeByteMatchTable` it's safe.
+ /// `byte_search!` macro calls `.use_table()` on whatever table it's provided, which makes
+ /// using the macro unsafe for `ByteMatchTable`, but safe for `SafeByteMatchTable`.
+ #[allow(clippy::unused_self)]
+ #[inline]
+ pub const fn use_table(&self) {}
+
+ /// Test a value against this `ByteMatchTable`.
+ #[inline]
+ pub fn matches<'a>(
+ &'a self,
+ data: &'a [u8; SEARCH_BATCH_SIZE],
+ actual_len: usize,
+ ) -> impl Iterator<Item = (usize, u8)> + 'a {
+ self.0.matches(data, actual_len)
+ }
+}
+
+macro_rules! simd_byte_match_table {
+ (|$byte:ident| $res:expr, $reverse:expr) => {{
+ use crate::lexer::search::SimdByteMatchTable;
+ // Clippy creates warnings because e.g. `byte_match_table!(|b| b == 0)`
+ // is expanded to `SimdByteMatchTable([(0 == 0), ... ])`
+ #[allow(clippy::eq_op)]
+ const TABLE: SimdByteMatchTable = seq_macro::seq!($byte in 0u8..=255 {
+ SimdByteMatchTable::new([ #($res,)* ], $reverse)
+ });
+ TABLE
+ }};
+}
+pub(crate) use simd_byte_match_table;
 
 /// Byte matcher lookup table.
 ///
@@ -158,6 +202,7 @@ pub(crate) use byte_match_table;
 /// }
 /// }
 /// ```
+#[derive(Debug)]
 #[repr(C, align(64))]
 pub struct SafeByteMatchTable([bool; 256]);
 
@@ -207,10 +252,96 @@ impl SafeByteMatchTable {
  #[inline]
  pub const fn use_table(&self) {}
 
- /// Test a value against this `SafeByteMatchTable`.
+ /// Returns the position of matched first delimiter and the matched first byte.
  #[inline]
- pub const fn matches(&self, b: u8) -> bool {
- self.0[b as usize]
+ pub fn matches<'a>(
+ &'a self,
+ data: &'a [u8; SEARCH_BATCH_SIZE],
+ actual_len: usize,
+ ) -> impl Iterator<Item = (usize, u8)> + 'a {
+ SafeByteMatchTableIter { table: self, data, actual_len, offset: 0 }
+ }
+}
+
+struct SafeByteMatchTableIter<'a> {
+ table: &'a SafeByteMatchTable,
+ data: &'a [u8; SEARCH_BATCH_SIZE],
+ actual_len: usize,
+ offset: usize,
+}
+
+impl Iterator for SafeByteMatchTableIter<'_> {
+ type Item = (usize, u8);
+
+ fn next(&mut self) -> Option<Self::Item> {
+ for (i, &b) in self.data[self.offset..self.actual_len].iter().enumerate() {
+ self.offset += 1;
+ if self.table.0[b as usize] {
+ return Some((i, b));
+ }
+ }
+ None
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::SafeByteMatchTable;
+ use crate::lexer::{source::Source, UniquePromise};
+
+ const SEARCH_BATCH_SIZE: usize = 16;
+ #[test]
+ fn neon_find_non_ascii() {
+ let table = seq_macro::seq!(b in 0u8..=255 {
+ SafeByteMatchTable::new([#(!(b.is_ascii_alphanumeric() || b == b'_' || b == b'$'),)*])
+ });
+ let data = [
+ "AAAAAAAA\"\rAAAAAA",
+ "AAAAAAAAAAAAAAA\"",
+ "AAAAAAAAAAAAAAAA",
+ "AAAAAAAA",
+ "AAAAAAAA\r",
+ "AAAAAAAAAAAAAAA\r",
+ ]
+ .map(|x| Source::new(x, UniquePromise::new_for_tests()));
+ let expected = [
+ (vec![Some((8, b'"')), Some((0, b'\r')), None], SEARCH_BATCH_SIZE),
+ (vec![Some((15, b'"')), None], SEARCH_BATCH_SIZE),
+ (vec![None], SEARCH_BATCH_SIZE),
+ (vec![None], 8),
+ (vec![Some((8, b'\r')), None], 9),
+ (vec![Some((15, b'\r')), None], SEARCH_BATCH_SIZE),
+ ];
+
+ for (idx, d) in data.into_iter().enumerate() {
+ let pos = d.position();
+ let (data, actual_len) =
+ unsafe { pos.peek_n_with_padding::<SEARCH_BATCH_SIZE>(d.end_addr()) }.unwrap();
+ let mut result = table.matches(&data, actual_len);
+ for val in &expected[idx].0 {
+ assert_eq!(result.next(), *val);
+ }
+ assert_eq!(actual_len, expected[idx].1);
+ }
+ }
+
+ #[test]
+ fn neon_find_single_quote_string() {
+ let table = seq_macro::seq!(b in 0u8..=255 {
+ // find non ascii
+ SafeByteMatchTable::new([#(matches!(b, b'\'' | b'\r' | b'\n' | b'\\'),)*])
+ });
+ let s1 = String::from(138u8 as char);
+ let data = [&s1].map(|x| Source::new(x, UniquePromise::new_for_tests()));
+ let expected = [(None, 2)];
+
+ for (idx, d) in data.into_iter().enumerate() {
+ let pos = d.position();
+ let (data, actual_len) =
+ unsafe { pos.peek_n_with_padding::<SEARCH_BATCH_SIZE>(d.end_addr()) }.unwrap();
+ let mut result = table.matches(&data, actual_len);
+ assert_eq!((result.next(), actual_len), expected[idx]);
+ }
  }
 }
 
@@ -495,93 +626,54 @@ macro_rules! byte_search {
 
  let mut $pos = $start;
  #[allow(unused_unsafe)] // Silence warnings if macro called in unsafe code
- loop {
- if $pos.addr() <= $lexer.source.end_for_batch_search_addr() {
- // Search a batch of `SEARCH_BATCH_SIZE` bytes.
- // The compiler unrolls this loop.
- // SAFETY:
- // `$pos.addr() > lexer.source.end_for_batch_search_addr()` check above ensures there are
- // at least `SEARCH_BATCH_SIZE` bytes remaining in `lexer.source`.
- // So calls to `$pos.read()` and `$pos.add(1)` in this loop cannot go out of bounds.
- for _i in 0..crate::lexer::search::SEARCH_BATCH_SIZE {
- // SAFETY: `$pos` cannot go out of bounds in this loop (see above).
- let $match_byte = unsafe { $pos.read() };
- if $table.matches($match_byte) {
- // Found match.
- // Check if should continue.
- {
- let $continue_byte = $match_byte;
- if $should_continue {
- // Not a match after all - continue searching.
- // SAFETY: `pos` is not at end of source, so safe to advance 1 byte.
- // See above about UTF-8 character boundaries invariant.
- $pos = unsafe { $pos.add(1) };
- continue;
- }
- }
-
- // Advance `lexer.source`'s position up to `$pos`, consuming unmatched bytes.
- // SAFETY: See above about UTF-8 character boundaries invariant.
- $lexer.source.set_position($pos);
-
- let $match_start = $start;
- return $match_handler;
- }
-
- // No match - continue searching
- // SAFETY: `$pos` cannot go out of bounds in this loop (see above).
- // Also see above about UTF-8 character boundaries invariant.
- $pos = unsafe { $pos.add(1) };
- }
- // No match in batch - loop round and searching next batch
- } else {
- // Not enough bytes remaining to process as a batch.
- // This branch marked `#[cold]` as should be very uncommon in normal-length JS files.
- // Very short JS files will be penalized, but they'll be very fast to parse anyway.
- // TODO: Could extend very short files with padding during parser initialization
- // to remove that problem.
- return crate::lexer::cold_branch(|| {
- let end_addr = $lexer.source.end_addr();
- while $pos.addr() < end_addr {
- // SAFETY: `pos` is not at end of source, so safe to read a byte
- let $match_byte = unsafe { $pos.read() };
- if $table.matches($match_byte) {
- // Found match.
- // Check if should continue.
- {
- let $continue_byte = $match_byte;
- if $should_continue {
- // Not a match after all - continue searching.
- // SAFETY: `pos` is not at end of source, so safe to advance 1 byte.
- // See above about UTF-8 character boundaries invariant.
- $pos = unsafe { $pos.add(1) };
- continue;
- }
- }
-
- // Advance `lexer.source`'s position up to `pos`, consuming unmatched bytes.
- // SAFETY: See above about UTF-8 character boundaries invariant.
- $lexer.source.set_position($pos);
-
- let $match_start = $start;
- return $match_handler;
- }
-
- // No match - continue searching
+ while let Some((data, actual_len)) = unsafe {
+ $pos.peek_n_with_padding::<{ crate::lexer::search::SEARCH_BATCH_SIZE }>(
+ $lexer.source.end_addr(),
+ )
+ } {
+ let mut iter = $table.matches(&data, actual_len);
+ let mut remaining = actual_len;
+ while let Some((offset, b)) = iter.next() {
+ // Advance the $pos with the batch matched pos
+ // SAFETY: `pos` is not at end of source, so safe to advance `pos` bytes.
+ // See above about UTF-8 character boundaries invariant.
+ $pos = unsafe { $pos.add(offset) };
+ remaining -= offset;
+ // SAFETY: `$pos` cannot go out of bounds in this loop (see above).
+ let $match_byte = b;
+ // Found match.
+ // Check if should continue.
+ {
+ let $continue_byte = $match_byte;
+ if $should_continue {
+ // Not a match after all - continue searching.
  // SAFETY: `pos` is not at end of source, so safe to advance 1 byte.
  // See above about UTF-8 character boundaries invariant.
  $pos = unsafe { $pos.add(1) };
+ remaining -= 1;
+ continue;
  }
+ }
+ // Advance `lexer.source`'s position up to `$pos`, consuming unmatched bytes.
+ // SAFETY: See above about UTF-8 character boundaries invariant.
+ $lexer.source.set_position($pos);
 
- // EOF.
- // Advance `lexer.source`'s position to end of file.
- $lexer.source.set_position($pos);
-
- let $eof_start = $start;
- $eof_handler
- });
+ let $match_start = $start;
+ return $match_handler;
  }
+ // No match in batch - loop round and searching next batch
+
+ // No match - continue searching
+ // SAFETY: `$pos` cannot go out of bounds in this loop (see above).
+ // Also see above about UTF-8 character boundaries invariant.
+ $pos = unsafe { $pos.add(remaining) };
  }
+
+ // EOF.
+ // Advance `lexer.source`'s position to end of file.
+ $lexer.source.set_position($pos);
+ let $eof_start = $start;
+ return $eof_handler;
  }};
 }
 pub(crate) use byte_search;