feat: allow the matcher to use user defined stack implementation and use by default a static stack

Author: WolverinDEV

bmatcher-core/benches/matcher.rs

@@ -23,7 +23,37 @@ pub fn simple_binary_pattern(instance: &mut Criterion) {
         let pattern = compiler::parse_pattern("FF FF 00 00 FF DE AD BE EF").unwrap();
         let pattern = compiler::optimize_pattern(&pattern);
 
-        let mut group = instance.benchmark_group("binary match throughput");
+        let mut group = instance.benchmark_group("binary");
+        for target_size in [128usize, 1024, 1024 * 4, 1024 * 1024] {
+            group.throughput(Throughput::Bytes(target_size as u64));
+            group.bench_with_input(
+                BenchmarkId::from_parameter(target_size),
+                &target_size,
+                |bencher, size| {
+                    let mut buffer = Vec::with_capacity(*size);
+
+                    let mut rng = StdRng::from_seed([0xEEu8; 32]);
+                    buffer.resize_with(*size, || rng.next_u32() as u8);
+
+                    let buffer = buffer.as_slice();
+                    bencher.iter(|| {
+                        let mut matcher =
+                            BinaryMatcher::new(black_box(&pattern), black_box(&buffer));
+                        while let Some(_match) = matcher.next_match() {
+                            panic!("pattern should not match");
+                        }
+                    });
+                },
+            );
+        }
+        group.finish();
+    }
+
+    {
+        let pattern = compiler::parse_pattern("FF FF 00 00 ? ? ? ? FF DE AD BE EF").unwrap();
+        let pattern = compiler::optimize_pattern(&pattern);
+
+        let mut group = instance.benchmark_group("binary with wildcard");
         for target_size in [128usize, 1024, 1024 * 4, 1024 * 1024] {
             group.throughput(Throughput::Bytes(target_size as u64));
             group.bench_with_input(

bmatcher-core/src/lib.rs

@@ -16,6 +16,13 @@ pub use target::MatchTarget;
 mod matcher;
 pub use matcher::BinaryMatcher;
 
+mod stack;
+pub use stack::{
+    HeapStack,
+    Stack,
+    StaticStack,
+};
+
 mod pattern;
 pub use pattern::{
     BinaryPattern,

bmatcher-core/src/matcher.rs

@@ -1,41 +1,78 @@
-use alloc::{
-    vec,
-    vec::Vec,
-};
-
 use crate::{
     Atom,
     BinaryPattern,
+    HeapStack,
     JumpType,
     MatchTarget,
     ReadWidth,
+    Stack,
+    StaticStack,
 };
 
 /// The `BinaryMatcher` is responsible for searching a [BinaryPattern] within a [MatchTarget].
 ///
 /// Use [`BinaryMatcher::next_match`] to iterate through matches of the specified pattern.
-pub struct BinaryMatcher<'a> {
+pub struct BinaryMatcher<
+    'a,
+    S: Stack<u32> = StaticStack<0x10, u32>,
+    C: Stack<usize> = StaticStack<0x10, usize>,
+> {
     pattern_atoms: &'a [Atom],
     pattern_byte_sequence: &'a [u8],
 
     target: &'a dyn MatchTarget,
 
     match_offset: usize,
 
-    save_stack: Vec<u32>,
-    cursor_stack: Vec<usize>,
+    save_stack: S,
+    cursor_stack: C,
 }
 
 impl<'a> BinaryMatcher<'a> {
+    /// Create a new BinaryMatcher instance with a statically allocated stack.
+    /// The default stack size is 0x10 for the save and cursor stack.
     pub fn new(pattern: &'a dyn BinaryPattern, target: &'a dyn MatchTarget) -> Self {
+        Self::new_with_stack(
+            pattern,
+            target,
+            StaticStack::<0x10, u32>::new(),
+            StaticStack::<0x10, usize>::new(),
+        )
+    }
+
+    /// Create a new BinaryMatcher instance with a heap allocated save and cursor stack
+    pub fn new_heap_stack(
+        pattern: &'a dyn BinaryPattern,
+        target: &'a dyn MatchTarget,
+    ) -> BinaryMatcher<'a, HeapStack<u32>, HeapStack<usize>> {
+        BinaryMatcher::new_with_stack(
+            pattern,
+            target,
+            HeapStack::<u32>::new(),
+            HeapStack::<usize>::new(),
+        )
+    }
+}
+
+impl<'a, S: Stack<u32>, C: Stack<usize>> BinaryMatcher<'a, S, C> {
+    /// Create a new binary matcher and supply the save and cursor stacks on your own
+    pub fn new_with_stack(
+        pattern: &'a dyn BinaryPattern,
+        target: &'a dyn MatchTarget,
+        mut save_stack: S,
+        cursor_stack: C,
+    ) -> Self {
+        save_stack.truncate(0);
+        save_stack.push_value(0x00);
+
         Self {
             pattern_atoms: pattern.atoms(),
             pattern_byte_sequence: pattern.byte_sequence(),
 
             target,
 
-            save_stack: vec![0; 8],
-            cursor_stack: vec![0; 8],
+            save_stack,
+            cursor_stack,
 
             match_offset: 0,
         }
@@ -86,11 +123,20 @@ impl<'a> BinaryMatcher<'a> {
                 }
 
                 Atom::CursorPush => {
-                    self.cursor_stack.push(data_cursor);
+                    if !self.cursor_stack.push_value(data_cursor) {
+                        /* TODO: Return error instead of abort search */
+                        return None;
+                    }
+
                     atom_cursor += 1;
                 }
                 Atom::CursorPop { advance } => {
-                    data_cursor = self.cursor_stack.pop().unwrap() + advance as usize;
+                    let Some(value) = self.cursor_stack.pop_value() else {
+                        /* TODO: Return error instead of abort search */
+                        return None;
+                    };
+
+                    data_cursor = value + advance as usize;
                     atom_cursor += 1;
                 }
 
@@ -157,18 +203,27 @@ impl<'a> BinaryMatcher<'a> {
                             (u32::from_le_bytes(value.try_into().unwrap()), 4)
                         }
                     };
-                    self.save_stack.push(value);
+                    if !self.save_stack.push_value(value) {
+                        /* TODO: Return error instead of abort search */
+                        return None;
+                    }
 
                     atom_cursor += 1;
                     data_cursor += width;
                 }
 
                 Atom::SaveCursor => {
-                    self.save_stack.push(data_cursor as u32);
+                    if !self.save_stack.push_value(data_cursor as u32) {
+                        /* TODO: Return error instead of abort search */
+                        return None;
+                    }
                     atom_cursor += 1;
                 }
                 Atom::SaveConstant(value) => {
-                    self.save_stack.push(value);
+                    if !self.save_stack.push_value(value) {
+                        /* TODO: Return error instead of abort search */
+                        return None;
+                    }
                     atom_cursor += 1;
                 }
             }
@@ -194,8 +249,10 @@ impl<'a> BinaryMatcher<'a> {
             }
 
             self.match_offset = match_offset + 1;
-            self.save_stack[0] = match_offset as u32;
-            return Some(&self.save_stack);
+
+            let save_stack = self.save_stack.stack_mut();
+            save_stack[0] = match_offset as u32;
+            return Some(save_stack);
         }
 
         None

bmatcher-core/src/stack.rs

@@ -0,0 +1,96 @@
+use alloc::vec::Vec;
+
+pub trait Stack<T> {
+    fn new() -> Self;
+
+    fn len(&self) -> usize;
+    fn reserve(&mut self, _size: usize) {}
+    fn truncate(&mut self, size: usize);
+
+    fn push_value(&mut self, value: T) -> bool;
+    fn pop_value(&mut self) -> Option<T>;
+
+    fn stack_mut(&mut self) -> &mut [T];
+}
+
+pub struct StaticStack<const N: usize, T> {
+    stack: [T; N],
+    length: usize,
+}
+
+impl<const N: usize, T: Default + Copy> Stack<T> for StaticStack<N, T> {
+    fn new() -> Self {
+        Self {
+            stack: [Default::default(); N],
+            length: 0,
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.length
+    }
+
+    fn stack_mut(&mut self) -> &mut [T] {
+        &mut self.stack[0..self.length]
+    }
+
+    fn truncate(&mut self, size: usize) {
+        self.length = size;
+        assert!(size <= self.stack.len());
+    }
+
+    fn push_value(&mut self, value: T) -> bool {
+        if self.length == self.stack.len() {
+            return false;
+        }
+
+        self.stack[self.length] = value;
+        self.length += 1;
+        true
+    }
+
+    fn pop_value(&mut self) -> Option<T> {
+        if self.length > 0 {
+            let value = self.stack[self.length];
+            self.length -= 1;
+            Some(value)
+        } else {
+            None
+        }
+    }
+}
+
+pub struct HeapStack<T> {
+    stack: Vec<T>,
+}
+
+impl<T> Stack<T> for HeapStack<T> {
+    fn new() -> Self {
+        Self { stack: Vec::new() }
+    }
+
+    fn len(&self) -> usize {
+        self.stack.len()
+    }
+
+    fn stack_mut(&mut self) -> &mut [T] {
+        &mut self.stack
+    }
+
+    fn reserve(&mut self, size: usize) {
+        self.stack.reserve(size);
+    }
+
+    fn truncate(&mut self, size: usize) {
+        self.stack.truncate(size);
+    }
+
+    fn push_value(&mut self, value: T) -> bool {
+        self.stack.push(value);
+        true
+    }
+
+    fn pop_value(&mut self) -> Option<T> {
+        self.stack.pop()
+    }
+}