use core::iter::Enumerate;
use core::num::NonZeroU32;
use crate::charwise::CharwiseDoubleArrayAhoCorasick;
use crate::prefilter::{Prefilter, PrefilterGate};
use crate::{Match, ROOT_STATE_IDX};
#[inline(always)]
unsafe fn decode_char_unchecked(bytes: &[u8], pos: usize) -> (char, usize) {
let first = *bytes.get_unchecked(pos);
let (c, len) = if first < 0x80 {
(u32::from(first), 1)
} else {
let c = u32::from(*bytes.get_unchecked(pos + 1) & 0x3f);
if first < 0xe0 {
((u32::from(first & 0x1f) << 6) | c, 2)
} else {
let c = (c << 6) | u32::from(*bytes.get_unchecked(pos + 2) & 0x3f);
if first < 0xf0 {
((u32::from(first & 0x0f) << 12) | c, 3)
} else {
let c = (c << 6) | u32::from(*bytes.get_unchecked(pos + 3) & 0x3f);
((u32::from(first & 0x07) << 18) | c, 4)
}
}
};
(char::from_u32_unchecked(c), len)
}
#[doc(hidden)]
pub struct CharWithEndOffsetIterator<I> {
inner: Enumerate<I>,
}
impl<I> CharWithEndOffsetIterator<I>
where
I: Iterator<Item = u8>,
{
pub unsafe fn new(inner: I) -> Self {
Self {
inner: inner.enumerate(),
}
}
}
impl<I> Iterator for CharWithEndOffsetIterator<I>
where
I: Iterator<Item = u8>,
{
type Item = (usize, char);
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
let (i, first) = self.inner.next()?;
let (end_offset, c) = if first < 0x80 {
(i + 1, u32::from(first))
} else {
let (i, rest) = unsafe { self.inner.next().unwrap_unchecked() };
let c = u32::from(rest & 0x3f);
if first < 0xe0 {
(i + 1, (u32::from(first & 0x1f) << 6) | c)
} else {
let (i, rest) = unsafe { self.inner.next().unwrap_unchecked() };
let c = (c << 6) | u32::from(rest & 0x3f);
if first < 0xf0 {
(i + 1, (u32::from(first & 0x0f) << 12) | c)
} else {
let (i, rest) = unsafe { self.inner.next().unwrap_unchecked() };
let c = (c << 6) | u32::from(rest & 0x3f);
(i + 1, (u32::from(first & 0x07) << 18) | c)
}
}
};
Some((end_offset, unsafe { char::from_u32_unchecked(c) }))
}
}
#[inline(always)]
fn scan<V>(
pma: &CharwiseDoubleArrayAhoCorasick<V>,
prefilter: &mut Option<&Prefilter>,
gate: &mut PrefilterGate,
haystack: &str,
state_id: &mut u32,
pos: &mut usize,
) -> Option<NonZeroU32> {
let bytes = haystack.as_bytes();
if let Some(pf) = *prefilter {
loop {
if *state_id == ROOT_STATE_IDX {
let candidate_pos = pf.next_position_at_char_boundary(haystack, *pos);
gate.record(candidate_pos - *pos);
*pos = candidate_pos;
if !gate.is_enabled() {
*prefilter = None;
break;
}
}
if *pos >= bytes.len() {
return None;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, *pos) };
*state_id = unsafe { pma.next_state_id_unchecked(*state_id, c) };
*pos += char_len;
if let Some(output_pos) = unsafe { pma.output_pos_unchecked(*state_id) } {
return Some(output_pos);
}
}
}
loop {
if *pos >= bytes.len() {
return None;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, *pos) };
*state_id = unsafe { pma.next_state_id_unchecked(*state_id, c) };
*pos += char_len;
if let Some(output_pos) = unsafe { pma.output_pos_unchecked(*state_id) } {
return Some(output_pos);
}
}
}
pub struct FindIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: CharWithEndOffsetIterator<P>,
pub(crate) first_call: bool,
}
impl<P, V> Iterator for FindIterator<'_, P, V>
where
P: Iterator<Item = u8>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if let Some(value) = self.pma.root_output_value() {
return if self.first_call {
self.first_call = false;
Some(Match {
length: 0,
end: 0,
value,
})
} else {
self.haystack.next().map(|(pos, _)| Match {
length: 0,
end: pos,
value,
})
};
}
let mut state_id = ROOT_STATE_IDX;
for (pos, c) in self.haystack.by_ref() {
state_id = unsafe { self.pma.next_state_id_unchecked(state_id, c) };
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) } {
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(pos));
}
}
None
}
}
pub struct FindSliceIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: P,
pub(crate) pos: usize,
pub(crate) first_call: bool,
pub(crate) prefilter: Option<&'a Prefilter>,
pub(crate) gate: PrefilterGate,
}
impl<P, V> Iterator for FindSliceIterator<'_, P, V>
where
P: AsRef<str>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
let haystack = self.haystack.as_ref();
let bytes = haystack.as_bytes();
if let Some(value) = self.pma.root_output_value() {
if self.first_call {
self.first_call = false;
return Some(Match {
length: 0,
end: 0,
value,
});
}
if self.pos < bytes.len() {
let (_, char_len) = unsafe { decode_char_unchecked(bytes, self.pos) };
self.pos += char_len;
return Some(Match {
length: 0,
end: self.pos,
value,
});
}
return None;
}
let mut state_id = ROOT_STATE_IDX;
let mut pos = self.pos;
{
if pos >= bytes.len() {
return None;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, pos) };
state_id = unsafe { self.pma.next_state_id_unchecked(state_id, c) };
pos += char_len;
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) } {
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(pos));
}
}
let output_pos = scan(
self.pma,
&mut self.prefilter,
&mut self.gate,
haystack,
&mut state_id,
&mut pos,
);
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos?) };
Some(out.to_match(pos))
}
}
pub struct FindOverlappingIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: CharWithEndOffsetIterator<P>,
pub(crate) state_id: u32,
pub(crate) pos: usize,
pub(crate) output_pos: Option<NonZeroU32>,
}
impl<P, V> Iterator for FindOverlappingIterator<'_, P, V>
where
P: Iterator<Item = u8>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if let Some(output_pos) = self.output_pos {
let out = unsafe { self.pma.output_at(output_pos) };
self.output_pos = out.parent();
return Some(out.to_match(self.pos));
}
for (pos, c) in self.haystack.by_ref() {
self.pos = pos;
self.state_id = unsafe { self.pma.next_state_id_unchecked(self.state_id, c) };
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(self.state_id) } {
let out = unsafe { self.pma.output_at(output_pos) };
self.output_pos = out.parent();
return Some(out.to_match(pos));
}
}
None
}
}
pub struct FindOverlappingSliceIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: P,
pub(crate) state_id: u32,
pub(crate) pos: usize,
pub(crate) output_pos: Option<NonZeroU32>,
pub(crate) prefilter: Option<&'a Prefilter>,
pub(crate) gate: PrefilterGate,
}
impl<P, V> Iterator for FindOverlappingSliceIterator<'_, P, V>
where
P: AsRef<str>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if let Some(output_pos) = self.output_pos {
let out = unsafe { self.pma.output_at(output_pos) };
self.output_pos = out.parent();
return Some(out.to_match(self.pos));
}
let haystack = self.haystack.as_ref();
let bytes = haystack.as_bytes();
let mut state_id = self.state_id;
let mut pos = self.pos;
{
if pos >= bytes.len() {
return None;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, pos) };
state_id = unsafe { self.pma.next_state_id_unchecked(state_id, c) };
pos += char_len;
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) } {
self.state_id = state_id;
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos) };
self.output_pos = out.parent();
return Some(out.to_match(pos));
}
}
let output_pos = scan(
self.pma,
&mut self.prefilter,
&mut self.gate,
haystack,
&mut state_id,
&mut pos,
);
self.state_id = state_id;
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos?) };
self.output_pos = out.parent();
Some(out.to_match(pos))
}
}
pub struct FindOverlappingNoSuffixIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: CharWithEndOffsetIterator<P>,
pub(crate) state_id: u32,
pub(crate) first_call: bool,
}
impl<P, V> Iterator for FindOverlappingNoSuffixIterator<'_, P, V>
where
P: Iterator<Item = u8>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if self.first_call {
self.first_call = false;
if let Some(value) = self.pma.root_output_value() {
return Some(Match {
length: 0,
end: 0,
value,
});
}
}
for (pos, c) in self.haystack.by_ref() {
self.state_id = unsafe { self.pma.next_state_id_unchecked(self.state_id, c) };
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(self.state_id) } {
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(pos));
}
}
None
}
}
pub struct FindOverlappingNoSuffixSliceIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: P,
pub(crate) state_id: u32,
pub(crate) pos: usize,
pub(crate) first_call: bool,
pub(crate) prefilter: Option<&'a Prefilter>,
pub(crate) gate: PrefilterGate,
}
impl<P, V> Iterator for FindOverlappingNoSuffixSliceIterator<'_, P, V>
where
P: AsRef<str>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if self.first_call {
self.first_call = false;
if let Some(value) = self.pma.root_output_value() {
return Some(Match {
length: 0,
end: 0,
value,
});
}
}
let haystack = self.haystack.as_ref();
let bytes = haystack.as_bytes();
let mut state_id = self.state_id;
let mut pos = self.pos;
{
if pos >= bytes.len() {
return None;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, pos) };
state_id = unsafe { self.pma.next_state_id_unchecked(state_id, c) };
pos += char_len;
if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) } {
self.state_id = state_id;
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(pos));
}
}
let output_pos = scan(
self.pma,
&mut self.prefilter,
&mut self.gate,
haystack,
&mut state_id,
&mut pos,
);
self.state_id = state_id;
self.pos = pos;
let out = unsafe { self.pma.output_at(output_pos?) };
Some(out.to_match(pos))
}
}
pub struct LeftmostFindIterator<'a, P, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) haystack: P,
pub(crate) pos: usize,
pub(crate) init_output_pos: Option<NonZeroU32>,
pub(crate) skip_empty: bool,
pub(crate) gate: PrefilterGate,
}
impl<P, V> LeftmostFindIterator<'_, P, V>
where
P: AsRef<str>,
V: Copy,
{
#[inline(always)]
fn next_plain(&mut self) -> Option<Match<V>> {
let mut state_id = ROOT_STATE_IDX;
let mut last_output_pos: Option<NonZeroU32> = self.init_output_pos;
'a: loop {
let mut skips = 0;
for c in unsafe { self.haystack.as_ref().get_unchecked(self.pos..) }.chars() {
skips += c.len_utf8();
state_id = unsafe { self.pma.next_state_id_leftmost_unchecked(state_id, c) };
if state_id == ROOT_STATE_IDX {
if let Some(output_pos) = last_output_pos {
let end = self.pos;
if last_output_pos == self.init_output_pos {
self.pos += c.len_utf8();
if self.skip_empty {
self.skip_empty = false;
continue 'a;
}
} else {
self.skip_empty = true;
}
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(end));
}
} else if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) }
{
last_output_pos.replace(output_pos);
self.pos += skips;
skips = 0;
}
}
break;
}
if self.pos == self.haystack.as_ref().len() {
self.init_output_pos.take();
}
if let Some(output_pos) = last_output_pos {
let out = unsafe { self.pma.output_at(output_pos) };
Some(out.to_match(self.pos))
} else {
self.pos = self.haystack.as_ref().len();
None
}
}
#[inline(always)]
fn next_filtered(&mut self) -> Option<Match<V>> {
let mut state_id = ROOT_STATE_IDX;
let mut last_output_pos: Option<NonZeroU32> = self.init_output_pos;
let haystack = self.haystack.as_ref();
let bytes = haystack.as_bytes();
let mut pos = self.pos;
let mut prefilter = self.pma.prefilter.as_ref();
loop {
if let Some(pf) = prefilter {
if state_id == ROOT_STATE_IDX && last_output_pos.is_none() {
let candidate_pos = pf.next_position_at_char_boundary(haystack, pos);
self.gate.record(candidate_pos - pos);
pos = candidate_pos;
if !self.gate.is_enabled() {
prefilter = None;
}
}
}
if pos >= bytes.len() {
break;
}
let (c, char_len) = unsafe { decode_char_unchecked(bytes, pos) };
state_id = unsafe { self.pma.next_state_id_leftmost_unchecked(state_id, c) };
if state_id == ROOT_STATE_IDX {
if let Some(output_pos) = last_output_pos {
let end = self.pos;
if last_output_pos == self.init_output_pos {
self.pos += char_len;
if self.skip_empty {
self.skip_empty = false;
pos = self.pos;
continue;
}
} else {
self.skip_empty = true;
}
let out = unsafe { self.pma.output_at(output_pos) };
return Some(out.to_match(end));
}
} else if let Some(output_pos) = unsafe { self.pma.output_pos_unchecked(state_id) } {
last_output_pos.replace(output_pos);
self.pos = pos + char_len;
}
pos += char_len;
}
if self.pos == haystack.len() {
self.init_output_pos.take();
}
if let Some(output_pos) = last_output_pos {
let out = unsafe { self.pma.output_at(output_pos) };
Some(out.to_match(self.pos))
} else {
self.pos = haystack.len();
None
}
}
}
impl<P, V> Iterator for LeftmostFindIterator<'_, P, V>
where
P: AsRef<str>,
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if self.pma.prefilter.is_some() && self.gate.is_enabled() {
self.next_filtered()
} else {
self.next_plain()
}
}
}
pub struct FindStepper<'a, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) state_id: u32,
pub(crate) pos: usize,
pub(crate) output_pos: Option<NonZeroU32>,
}
impl<V> FindStepper<'_, V>
where
V: Copy,
{
#[inline(always)]
pub fn consume(&mut self, c: char) {
self.pos += c.len_utf8();
if self.pma.root_output_value().is_some() {
return;
}
unsafe {
self.state_id = self.pma.next_state_id_unchecked(self.state_id, c);
self.output_pos = self.pma.output_pos_unchecked(self.state_id);
}
if self.output_pos.is_some() {
self.state_id = ROOT_STATE_IDX;
}
}
#[must_use]
#[inline(always)]
pub fn matches(&self) -> Option<Match<V>> {
self.output_pos.map(|output_pos| unsafe {
let out = self.pma.output_at(output_pos);
out.to_match(self.pos)
})
}
}
pub struct FindOverlappingStepperIterator<'a, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) pos: usize,
pub(crate) output_pos: Option<NonZeroU32>,
}
impl<V> Iterator for FindOverlappingStepperIterator<'_, V>
where
V: Copy,
{
type Item = Match<V>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
if let Some(output_pos) = self.output_pos {
let out = unsafe { self.pma.output_at(output_pos) };
self.output_pos = out.parent();
return Some(out.to_match(self.pos));
}
None
}
}
pub struct FindOverlappingStepper<'a, V> {
pub(crate) pma: &'a CharwiseDoubleArrayAhoCorasick<V>,
pub(crate) state_id: u32,
pub(crate) pos: usize,
}
impl<'a, V> FindOverlappingStepper<'a, V>
where
V: Copy,
{
#[inline(always)]
pub fn consume(&mut self, c: char) {
self.state_id = unsafe { self.pma.next_state_id_unchecked(self.state_id, c) };
self.pos += c.len_utf8();
}
#[must_use]
#[inline(always)]
pub fn matches(&self) -> FindOverlappingStepperIterator<'a, V> {
let output_pos = unsafe {
self.pma.output_pos_unchecked(self.state_id)
};
FindOverlappingStepperIterator {
pma: self.pma,
pos: self.pos,
output_pos,
}
}
}
#[cfg(test)]
mod tests {
use alloc::vec::Vec;
use super::*;
#[test]
fn test_char_with_end_offset_iterator() {
let test_string =
"\u{0000}\u{0001}\u{0002}\u{0004}\u{0008}\u{0010}\u{001f}\u{0020}\u{0040}\
\u{007f}\u{0080}\u{0100}\u{01ff}\u{0200}\u{0400}\u{07ff}\u{0800}\u{1000}\
\u{1fff}\u{2000}\u{4000}\u{8000}\u{ffff}\u{10000}\
\u{1ffff}\u{20000}\u{40000}\u{80000}\u{100000}\u{10ffff}";
let mut it =
unsafe { CharWithEndOffsetIterator::new(test_string.as_bytes().iter().copied()) };
assert_eq!(Some((1, '\u{0000}')), it.next());
assert_eq!(Some((2, '\u{0001}')), it.next());
assert_eq!(Some((3, '\u{0002}')), it.next());
assert_eq!(Some((4, '\u{0004}')), it.next());
assert_eq!(Some((5, '\u{0008}')), it.next());
assert_eq!(Some((6, '\u{0010}')), it.next());
assert_eq!(Some((7, '\u{001f}')), it.next());
assert_eq!(Some((8, '\u{0020}')), it.next());
assert_eq!(Some((9, '\u{0040}')), it.next());
assert_eq!(Some((10, '\u{007f}')), it.next());
assert_eq!(Some((12, '\u{0080}')), it.next());
assert_eq!(Some((14, '\u{0100}')), it.next());
assert_eq!(Some((16, '\u{01ff}')), it.next());
assert_eq!(Some((18, '\u{0200}')), it.next());
assert_eq!(Some((20, '\u{0400}')), it.next());
assert_eq!(Some((22, '\u{07ff}')), it.next());
assert_eq!(Some((25, '\u{0800}')), it.next());
assert_eq!(Some((28, '\u{1000}')), it.next());
assert_eq!(Some((31, '\u{1fff}')), it.next());
assert_eq!(Some((34, '\u{2000}')), it.next());
assert_eq!(Some((37, '\u{4000}')), it.next());
assert_eq!(Some((40, '\u{8000}')), it.next());
assert_eq!(Some((43, '\u{ffff}')), it.next());
assert_eq!(Some((47, '\u{10000}')), it.next());
assert_eq!(Some((51, '\u{1ffff}')), it.next());
assert_eq!(Some((55, '\u{20000}')), it.next());
assert_eq!(Some((59, '\u{40000}')), it.next());
assert_eq!(Some((63, '\u{80000}')), it.next());
assert_eq!(Some((67, '\u{100000}')), it.next());
assert_eq!(Some((71, '\u{10ffff}')), it.next());
assert_eq!(None, it.next());
assert_eq!(None, it.next());
}
#[test]
fn test_overlapping_no_suffix_iter() {
let pma = CharwiseDoubleArrayAhoCorasick::<u32>::new(["a", "ab", ""]).unwrap();
let result = pma
.find_overlapping_no_suffix_iter("ab")
.collect::<Vec<_>>();
assert_eq!(
vec![
Match {
length: 0,
end: 0,
value: 2
},
Match {
length: 1,
end: 1,
value: 0
},
Match {
length: 2,
end: 2,
value: 1
},
],
result
);
}
#[test]
fn test_overlapping_stepper_lifetime() {
let pma = CharwiseDoubleArrayAhoCorasick::new(["a", "ab"]).unwrap();
let mut stepper = pma.find_overlapping_stepper();
stepper.consume('a');
let mut it1 = stepper.matches();
stepper.consume('b');
let mut it2 = stepper.matches();
assert_eq!(
Some(Match {
length: 1,
end: 1,
value: 0
}),
it1.next()
);
assert_eq!(
Some(Match {
length: 2,
end: 2,
value: 1
}),
it2.next()
);
}
}