mod builder;
pub mod iter;
use core::mem;
use core::num::NonZeroU32;
use alloc::vec::Vec;
use crate::build_helper::BuildHelper;
pub use crate::bytewise::builder::DoubleArrayAhoCorasickBuilder;
use crate::bytewise::iter::{
FindIterator, FindOverlappingIterator, FindOverlappingNoSuffixIterator, FindOverlappingStepper,
FindStepper, LeftmostFindIterator, U8SliceIterator,
};
use crate::errors::{DaachorseError, Result};
use crate::intpack::{U24nU8, U24};
use crate::serializer::{Serializable, SerializableVec};
use crate::utils::FromU32;
use crate::{MatchKind, Output};
const ROOT_STATE_IDX: u32 = 0;
const DEAD_STATE_IDX: u32 = 1;
#[derive(Clone, Eq, Hash, PartialEq)]
pub struct DoubleArrayAhoCorasick<V> {
states: Vec<State>,
outputs: Vec<Output<V>>,
match_kind: MatchKind,
num_states: u32,
}
impl<V> DoubleArrayAhoCorasick<V> {
pub fn new<I, P>(patterns: I) -> Result<Self>
where
I: IntoIterator<Item = P>,
P: AsRef<[u8]>,
V: Copy + TryFrom<usize>,
{
DoubleArrayAhoCorasickBuilder::new().build(patterns)
}
pub fn with_values<I, P>(patvals: I) -> Result<Self>
where
I: IntoIterator<Item = (P, V)>,
P: AsRef<[u8]>,
V: Copy,
{
DoubleArrayAhoCorasickBuilder::new().build_with_values(patvals)
}
pub fn find_iter<P>(&self, haystack: P) -> FindIterator<'_, U8SliceIterator<P>, V>
where
P: AsRef<[u8]>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindIterator {
pma: self,
haystack: U8SliceIterator::new(haystack).enumerate(),
}
}
pub fn find_iter_from_iter<P>(&self, haystack: P) -> FindIterator<'_, P, V>
where
P: Iterator<Item = u8>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindIterator {
pma: self,
haystack: haystack.enumerate(),
}
}
pub fn find_overlapping_iter<P>(
&self,
haystack: P,
) -> FindOverlappingIterator<'_, U8SliceIterator<P>, V>
where
P: AsRef<[u8]>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindOverlappingIterator {
pma: self,
haystack: U8SliceIterator::new(haystack).enumerate(),
state_id: ROOT_STATE_IDX,
output_pos: None,
pos: 0,
}
}
pub fn find_overlapping_iter_from_iter<P>(
&self,
haystack: P,
) -> FindOverlappingIterator<'_, P, V>
where
P: Iterator<Item = u8>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindOverlappingIterator {
pma: self,
haystack: haystack.enumerate(),
state_id: ROOT_STATE_IDX,
output_pos: None,
pos: 0,
}
}
pub fn find_overlapping_no_suffix_iter<P>(
&self,
haystack: P,
) -> FindOverlappingNoSuffixIterator<'_, U8SliceIterator<P>, V>
where
P: AsRef<[u8]>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindOverlappingNoSuffixIterator {
pma: self,
haystack: U8SliceIterator::new(haystack).enumerate(),
state_id: ROOT_STATE_IDX,
}
}
pub fn find_overlapping_no_suffix_iter_from_iter<P>(
&self,
haystack: P,
) -> FindOverlappingNoSuffixIterator<'_, P, V>
where
P: Iterator<Item = u8>,
{
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindOverlappingNoSuffixIterator {
pma: self,
haystack: haystack.enumerate(),
state_id: ROOT_STATE_IDX,
}
}
pub fn leftmost_find_iter<P>(&self, haystack: P) -> LeftmostFindIterator<'_, P, V>
where
P: AsRef<[u8]>,
{
assert!(
self.match_kind.is_leftmost(),
"Error: match_kind must be leftmost."
);
LeftmostFindIterator {
pma: self,
haystack,
pos: 0,
}
}
#[must_use]
pub fn find_stepper(&self) -> FindStepper<'_, V> {
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindStepper {
pma: self,
state_id: ROOT_STATE_IDX,
pos: 0,
}
}
#[must_use]
pub fn find_overlapping_stepper(&self) -> FindOverlappingStepper<'_, V> {
assert!(
self.match_kind.is_standard(),
"Error: match_kind must be standard."
);
FindOverlappingStepper {
pma: self,
state_id: ROOT_STATE_IDX,
pos: 0,
}
}
#[must_use]
pub fn heap_bytes(&self) -> usize {
self.states.len() * mem::size_of::<State>()
+ self.outputs.len() * mem::size_of::<Output<V>>()
}
#[must_use]
pub fn num_states(&self) -> usize {
usize::from_u32(self.num_states)
}
#[must_use]
pub fn serialize(&self) -> Vec<u8>
where
V: Serializable,
{
let mut result = Vec::with_capacity(
self.states.serialized_bytes()
+ self.outputs.serialized_bytes()
+ MatchKind::serialized_bytes()
+ u32::serialized_bytes(),
);
self.states.serialize_to_vec(&mut result);
self.outputs.serialize_to_vec(&mut result);
self.match_kind.serialize_to_vec(&mut result);
self.num_states.serialize_to_vec(&mut result);
result
}
pub fn deserialize(source: &[u8]) -> Result<(Self, &[u8])>
where
V: Serializable,
{
let (states, source) = Vec::<State>::deserialize_from_slice(source)?;
let (outputs, source) = Vec::<Output<V>>::deserialize_from_slice(source)?;
let (match_kind, source) = MatchKind::deserialize_from_slice(source)?;
let (num_states, source) = u32::deserialize_from_slice(source)?;
let pma = Self {
states,
outputs,
match_kind,
num_states,
};
let block_len = 256;
if pma.states.len() % block_len != 0 {
return Err(DaachorseError::invalid_automaton());
}
let states_len = pma.states.len();
let outputs_len = pma.outputs.len();
for state in &pma.states {
if let Some(base) = state.base() {
if usize::from_u32(base.get()) >= states_len {
return Err(DaachorseError::invalid_automaton());
}
};
if usize::from_u32(state.fail()) >= states_len {
return Err(DaachorseError::invalid_automaton());
}
if let Some(output_pos) = state.output_pos() {
if usize::from_u32(output_pos.get() - 1) >= outputs_len {
return Err(DaachorseError::invalid_automaton());
}
};
}
for (i, output) in pma.outputs.iter().enumerate() {
if let Some(parent) = output.parent {
if usize::from_u32(parent.get() - 1) >= i {
return Err(DaachorseError::invalid_automaton());
}
};
}
Ok((pma, source))
}
#[must_use]
pub unsafe fn deserialize_unchecked(source: &[u8]) -> (Self, &[u8])
where
V: Serializable,
{
let (states, source) = Vec::<State>::deserialize_from_slice(source).unwrap_unchecked();
let (outputs, source) = Vec::<Output<V>>::deserialize_from_slice(source).unwrap_unchecked();
let (match_kind, source) = MatchKind::deserialize_from_slice(source).unwrap_unchecked();
let (num_states, source) = u32::deserialize_from_slice(source).unwrap_unchecked();
(
Self {
states,
outputs,
match_kind,
num_states,
},
source,
)
}
#[inline(always)]
unsafe fn child_index_unchecked(&self, state_id: u32, c: u8) -> Option<u32> {
self.states
.get_unchecked(usize::from_u32(state_id))
.base()
.and_then(|base| {
let child_idx = base.get() ^ u32::from(c);
Some(child_idx)
.filter(|&x| self.states.get_unchecked(usize::from_u32(x)).check() == c)
})
}
#[inline(always)]
unsafe fn next_state_id_unchecked(&self, mut state_id: u32, c: u8) -> u32 {
loop {
if let Some(state_id) = self.child_index_unchecked(state_id, c) {
return state_id;
}
if state_id == ROOT_STATE_IDX {
return ROOT_STATE_IDX;
}
state_id = self.states.get_unchecked(usize::from_u32(state_id)).fail();
}
}
#[inline(always)]
unsafe fn next_state_id_leftmost_unchecked(&self, mut state_id: u32, c: u8) -> u32 {
loop {
if let Some(state_id) = self.child_index_unchecked(state_id, c) {
return state_id;
}
if state_id == ROOT_STATE_IDX {
return ROOT_STATE_IDX;
}
let fail_id = self.states.get_unchecked(usize::from_u32(state_id)).fail();
if fail_id == DEAD_STATE_IDX {
return ROOT_STATE_IDX;
}
state_id = fail_id;
}
}
}
#[derive(Clone, Copy, Default, Eq, Hash, PartialEq)]
struct State {
base: Option<NonZeroU32>,
fail: u32,
opos_ch: U24nU8,
}
impl State {
#[inline(always)]
pub const fn base(&self) -> Option<NonZeroU32> {
self.base
}
#[inline(always)]
pub fn check(&self) -> u8 {
self.opos_ch.b()
}
#[inline(always)]
pub const fn fail(&self) -> u32 {
self.fail
}
#[inline(always)]
pub const fn output_pos(&self) -> Option<NonZeroU32> {
NonZeroU32::new(self.opos_ch.a().get())
}
#[inline(always)]
pub fn set_base(&mut self, x: NonZeroU32) {
self.base = Some(x);
}
#[inline(always)]
pub fn set_check(&mut self, x: u8) {
self.opos_ch.set_b(x);
}
#[inline(always)]
pub fn set_fail(&mut self, x: u32) {
self.fail = x;
}
#[inline(always)]
pub fn set_output_pos(&mut self, x: Option<NonZeroU32>) -> Result<()> {
let x = x.map_or(0, NonZeroU32::get);
if let Ok(x) = U24::try_from(x) {
self.opos_ch.set_a(x);
Ok(())
} else {
Err(DaachorseError::automaton_scale("output_pos", U24::MAX))
}
}
}
impl Serializable for State {
#[inline(always)]
fn serialize_to_vec(&self, dst: &mut Vec<u8>) {
self.base.serialize_to_vec(dst);
self.fail.serialize_to_vec(dst);
self.opos_ch.serialize_to_vec(dst);
}
#[inline(always)]
fn deserialize_from_slice(src: &[u8]) -> Result<(Self, &[u8])> {
let (base, src) = Option::<NonZeroU32>::deserialize_from_slice(src)?;
let (fail, src) = u32::deserialize_from_slice(src)?;
let (opos_ch, src) = U24nU8::deserialize_from_slice(src)?;
Ok((
Self {
base,
fail,
opos_ch,
},
src,
))
}
#[inline(always)]
fn serialized_bytes() -> usize {
Option::<NonZeroU32>::serialized_bytes()
+ u32::serialized_bytes()
+ U24nU8::serialized_bytes()
}
}
impl core::fmt::Debug for State {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("State")
.field("base", &self.base())
.field("check", &self.check())
.field("fail", &self.fail())
.field("output_pos", &self.output_pos())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_double_array() {
let patterns = vec![vec![0, 0], vec![0, 2], vec![1, 2], vec![2]];
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
let base_expected = vec![
NonZeroU32::new(4), None, None, None, NonZeroU32::new(8), NonZeroU32::new(1), None, None, None, None, None, ];
let check_expected = vec![
0, 1, 2, 2, 0, 1, 2, 7, 0, 9, 2, ];
let fail_expected = vec![
ROOT_STATE_IDX, ROOT_STATE_IDX, ROOT_STATE_IDX, 6, ROOT_STATE_IDX, ROOT_STATE_IDX, ROOT_STATE_IDX, ROOT_STATE_IDX, 4, ROOT_STATE_IDX, 6, ];
let pma_base: Vec<_> = pma.states[0..11].iter().map(|state| state.base()).collect();
let pma_check: Vec<_> = pma.states[0..11]
.iter()
.map(|state| state.check())
.collect();
let pma_fail: Vec<_> = pma.states[0..11].iter().map(|state| state.fail()).collect();
assert_eq!(base_expected, pma_base);
assert_eq!(check_expected, pma_check);
assert_eq!(fail_expected, pma_fail);
}
#[test]
fn test_num_states() {
let patterns = vec!["abba", "baaba", "ababa"];
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
assert_eq!(13, pma.num_states());
}
#[test]
fn test_input_order() {
let patvals_sorted = vec![("ababa", 0), ("abba", 1), ("baaba", 2)];
let patvals_unsorted = vec![("abba", 1), ("baaba", 2), ("ababa", 0)];
let pma_sorted = DoubleArrayAhoCorasick::with_values(patvals_sorted).unwrap();
let pma_unsorted = DoubleArrayAhoCorasick::with_values(patvals_unsorted).unwrap();
assert_eq!(pma_sorted.states, pma_unsorted.states);
assert_eq!(pma_sorted.outputs, pma_unsorted.outputs);
}
#[test]
fn test_n_blocks_1_1() {
let mut patterns = vec![];
for i in 0x00..=0xfd {
let pattern = vec![i];
patterns.push(pattern);
}
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
assert_eq!(255, pma.num_states());
assert_eq!(256, pma.states.len());
assert_eq!(0xfe, pma.states[0].base().unwrap().get());
}
#[test]
fn test_n_blocks_1_2() {
let mut patterns = vec![];
patterns.push(vec![0x00]);
patterns.push(vec![0x02]);
for i in 0x04..=0xff {
patterns.push(vec![i]);
}
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
assert_eq!(255, pma.num_states());
assert_eq!(512, pma.states.len());
assert_eq!(0x100, pma.states[0].base().unwrap().get());
}
#[test]
fn test_n_blocks_2_1() {
let mut patterns = vec![];
for i in 0x00..=0x7f {
let pattern = vec![i];
patterns.push(pattern);
}
for i in 0x00..=0x7d {
let pattern = vec![0x00, i];
patterns.push(pattern);
}
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
assert_eq!(255, pma.num_states());
assert_eq!(256, pma.states.len());
assert_eq!(0x80, pma.states[0].base().unwrap().get());
assert_eq!(0x7e, pma.states[0x80].base().unwrap().get());
}
#[test]
fn test_n_blocks_2_2() {
let mut patterns = vec![];
for i in 0..=0x7f {
let pattern = vec![i];
patterns.push(pattern);
}
patterns.push(vec![0, 0x00]);
patterns.push(vec![0, 0x02]);
for i in 0x04..=0x7f {
let pattern = vec![0x00, i];
patterns.push(pattern);
}
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
assert_eq!(255, pma.num_states());
assert_eq!(512, pma.states.len());
assert_eq!(0x80, pma.states[0].base().unwrap().get());
assert_eq!(0x100, pma.states[0x80].base().unwrap().get());
}
#[test]
fn test_serialize_state() {
let mut opos_ch = U24nU8::default();
opos_ch.set_a(U24::try_from(57).unwrap());
opos_ch.set_b(77);
let x = State {
base: NonZeroU32::new(42),
fail: 13,
opos_ch,
};
let mut data = vec![];
x.serialize_to_vec(&mut data);
assert_eq!(data.len(), State::serialized_bytes());
let (y, rest) = State::deserialize_from_slice(&data).unwrap();
assert!(rest.is_empty());
assert_eq!(x, y);
}
#[test]
fn test_serialize_pma() {
let patterns = vec!["abba", "baaba", "ababa"];
let pma = DoubleArrayAhoCorasick::<u32>::new(patterns).unwrap();
let bytes = pma.serialize();
let (other, rest) = unsafe { DoubleArrayAhoCorasick::deserialize_unchecked(&bytes) };
assert!(rest.is_empty());
assert_eq!(pma.states, other.states);
assert_eq!(pma.outputs, other.outputs);
assert_eq!(pma.match_kind, other.match_kind);
assert_eq!(pma.num_states, other.num_states);
}
}