use std::borrow::Cow;
#[cfg(feature = "dfa")]
use aho_corasick::{
AhoCorasick, AhoCorasickBuilder, AhoCorasickKind, MatchKind as AhoCorasickMatchKind,
};
use daachorse::{
DoubleArrayAhoCorasick, DoubleArrayAhoCorasickBuilder,
MatchKind as DoubleArrayAhoCorasickMatchKind,
charwise::{CharwiseDoubleArrayAhoCorasick, CharwiseDoubleArrayAhoCorasickBuilder},
};
use crate::MatcherError;
use super::rule::{PatternEntry, PatternIndex};
#[cfg(feature = "dfa")]
const AC_DFA_PATTERN_THRESHOLD: usize = 5_000;
#[derive(Clone)]
pub(super) struct ScanPlan {
ascii_matcher: Option<AsciiMatcher>,
non_ascii_matcher: Option<NonAsciiMatcher>,
patterns: PatternIndex,
}
#[derive(Clone)]
enum AsciiMatcher {
#[cfg(feature = "dfa")]
AcDfa {
matcher: AhoCorasick,
to_value: Vec<u32>,
},
DaacBytewise(DoubleArrayAhoCorasick<u32>),
}
#[derive(Clone)]
enum NonAsciiMatcher {
DaacCharwise(CharwiseDoubleArrayAhoCorasick<u32>),
}
impl ScanPlan {
pub(super) fn compile(
dedup_patterns: &[Cow<'_, str>],
dedup_entries: Vec<Vec<PatternEntry>>,
) -> Result<Self, MatcherError> {
let patterns = PatternIndex::new(dedup_entries);
let value_map = patterns.build_value_map();
let (ascii_matcher, non_ascii_matcher) = compile_automata(dedup_patterns, &value_map)?;
Ok(Self {
ascii_matcher,
non_ascii_matcher,
patterns,
})
}
#[inline(always)]
pub(super) fn patterns(&self) -> &PatternIndex {
&self.patterns
}
#[inline(always)]
pub(super) fn is_match(&self, text: &str) -> bool {
if self.non_ascii_matcher.is_none() {
return self
.ascii_matcher
.as_ref()
.is_some_and(|matcher| matcher.is_match(text));
}
if text.is_ascii() {
self.ascii_matcher
.as_ref()
.is_some_and(|matcher| matcher.is_match(text))
} else {
self.non_ascii_matcher
.as_ref()
.is_some_and(|matcher| matcher.is_match(text))
}
}
#[inline(always)]
pub(super) fn for_each_match_value(
&self,
text: &str,
is_ascii: bool,
on_value: impl FnMut(u32) -> bool,
) -> bool {
let use_ascii = self.non_ascii_matcher.is_none() || is_ascii;
if use_ascii {
if let Some(ref matcher) = self.ascii_matcher {
return matcher.for_each_match_value(text, on_value);
}
} else if let Some(ref matcher) = self.non_ascii_matcher {
return matcher.for_each_match_value(text, on_value);
}
false
}
}
impl AsciiMatcher {
#[inline(always)]
fn is_match(&self, text: &str) -> bool {
match self {
#[cfg(feature = "dfa")]
Self::AcDfa { matcher, .. } => matcher.is_match(text),
Self::DaacBytewise(matcher) => matcher.find_iter(text).next().is_some(),
}
}
#[inline(always)]
fn for_each_match_value(&self, text: &str, mut on_value: impl FnMut(u32) -> bool) -> bool {
match self {
#[cfg(feature = "dfa")]
Self::AcDfa { matcher, to_value } => {
for hit in matcher.find_overlapping_iter(text) {
let value = unsafe { *to_value.get_unchecked(hit.pattern().as_usize()) };
if on_value(value) {
return true;
}
}
false
}
Self::DaacBytewise(matcher) => {
for hit in matcher.find_overlapping_iter(text) {
if on_value(hit.value()) {
return true;
}
}
false
}
}
}
}
impl NonAsciiMatcher {
#[inline(always)]
fn is_match(&self, text: &str) -> bool {
match self {
Self::DaacCharwise(matcher) => matcher.find_iter(text).next().is_some(),
}
}
#[inline(always)]
fn for_each_match_value(&self, text: &str, mut on_value: impl FnMut(u32) -> bool) -> bool {
match self {
Self::DaacCharwise(matcher) => {
for hit in matcher.find_overlapping_iter(text) {
if on_value(hit.value()) {
return true;
}
}
false
}
}
}
}
fn compile_automata(
dedup_patterns: &[Cow<'_, str>],
value_map: &[u32],
) -> Result<(Option<AsciiMatcher>, Option<NonAsciiMatcher>), MatcherError> {
let cap = dedup_patterns.len();
let mut ascii_patvals: Vec<(&str, u32)> = Vec::with_capacity(cap);
let mut non_ascii_patvals: Vec<(&str, u32)> = Vec::with_capacity(cap);
#[cfg(feature = "dfa")]
let mut ascii_ac_to_value: Vec<u32> = Vec::with_capacity(cap);
for (dedup_idx, pattern) in dedup_patterns.iter().enumerate() {
let value = value_map[dedup_idx];
if pattern.as_ref().is_ascii() {
#[cfg(feature = "dfa")]
ascii_ac_to_value.push(value);
ascii_patvals.push((pattern.as_ref(), value));
} else {
non_ascii_patvals.push((pattern.as_ref(), value));
}
}
let has_ascii = !ascii_patvals.is_empty();
let has_non_ascii = !non_ascii_patvals.is_empty();
let full_charwise_patvals = if has_ascii && has_non_ascii {
Some(
dedup_patterns
.iter()
.enumerate()
.map(|(i, p)| (p.as_ref(), value_map[i]))
.collect::<Vec<_>>(),
)
} else {
None
};
let charwise_source = full_charwise_patvals
.as_deref()
.unwrap_or(non_ascii_patvals.as_slice());
let build_ascii = move || -> Result<AsciiMatcher, MatcherError> {
#[cfg(feature = "dfa")]
if ascii_patvals.len() <= AC_DFA_PATTERN_THRESHOLD {
return Ok(AsciiMatcher::AcDfa {
matcher: AhoCorasickBuilder::new()
.kind(Some(AhoCorasickKind::DFA))
.match_kind(AhoCorasickMatchKind::Standard)
.build(ascii_patvals.iter().map(|(p, _)| p))
.map_err(MatcherError::automaton_build)?,
to_value: ascii_ac_to_value,
});
}
Ok(AsciiMatcher::DaacBytewise(
DoubleArrayAhoCorasickBuilder::new()
.match_kind(DoubleArrayAhoCorasickMatchKind::Standard)
.build_with_values(ascii_patvals)
.map_err(MatcherError::automaton_build)?,
))
};
let build_charwise = || -> Result<NonAsciiMatcher, MatcherError> {
Ok(NonAsciiMatcher::DaacCharwise(
CharwiseDoubleArrayAhoCorasickBuilder::new()
.match_kind(DoubleArrayAhoCorasickMatchKind::Standard)
.build_with_values(charwise_source.iter().copied())
.map_err(MatcherError::automaton_build)?,
))
};
match (has_ascii, has_non_ascii) {
(true, true) => std::thread::scope(|s| {
let ascii_handle = s.spawn(build_ascii);
let charwise = build_charwise()?;
let ascii = ascii_handle
.join()
.expect("ASCII automaton build panicked")?;
Ok((Some(ascii), Some(charwise)))
}),
(true, false) => Ok((Some(build_ascii()?), None)),
(false, true) => Ok((None, Some(build_charwise()?))),
(false, false) => Ok((None, None)),
}
}