memmem: add new HeuristicFrequencyRank

bench/src/memmem/byterank.rs

@@ -0,0 +1,116 @@
+use criterion::Criterion;
+use memchr::memmem::HeuristicFrequencyRank;
+
+use crate::define;
+
+pub(crate) fn all(c: &mut Criterion) {
+    finder_construction(c);
+    byte_frequencies(c);
+}
+
+fn finder_construction(c: &mut Criterion) {
+    // This benchmark is purely for measuring the time taken to create a
+    // `Finder`. It is here to prevent regressions when adding new features
+    // to the `Finder`, such as the ability to construct with a custom
+    // `HeuristicFrequencyRank`.
+    const NEEDLES: [&str; 3] = ["a", "abcd", "abcdefgh12345678"];
+
+    for needle in NEEDLES {
+        define(
+            c,
+            &format!(
+                "memmem/krate/bytefreq/construct-finder/default(len={})",
+                needle.len()
+            ),
+            needle.as_bytes(),
+            Box::new(move |b| {
+                b.iter(|| {
+                    memchr::memmem::FinderBuilder::new()
+                        .build_forward(needle.as_bytes())
+                });
+            }),
+        );
+        define(
+            c,
+            &format!(
+                "memmem/krate/bytefreq/construct-finder/custom(len={})",
+                needle.len()
+            ),
+            needle.as_bytes(),
+            Box::new(move |b| {
+                b.iter(|| {
+                    memchr::memmem::FinderBuilder::new()
+                        .build_forward_with_ranker(Binary, needle.as_bytes())
+                });
+            }),
+        );
+    }
+}
+
+fn byte_frequencies(c: &mut Criterion) {
+    // This benchmark exists to demonstrate a common use case for
+    // customizing the byte frequency table used by a `Finder`
+    // and the relative performance gain from using an optimal table.
+    // This is essentially why `HeuristicFrequencyRank` was added.
+
+    // Bytes we want to scan for that are rare in strings but common in
+    // executables.
+    const NEEDLE: &[u8] = b"\x00\x00\xdd\xdd'";
+
+    // The input for the benchmark is the benchmark binary itself
+    let exe = std::env::args().next().unwrap();
+    let corpus = std::fs::read(exe).unwrap();
+
+    let bin = corpus.clone();
+    define(
+        c,
+        &format!("memmem/krate/bytefreq/default"),
+        &corpus,
+        Box::new(move |b| {
+            let finder =
+                memchr::memmem::FinderBuilder::new().build_forward(NEEDLE);
+            b.iter(|| {
+                assert_eq!(1, finder.find_iter(&bin).count());
+            });
+        }),
+    );
+
+    let bin = corpus.clone();
+    define(
+        c,
+        &format!("memmem/krate/bytefreq/custom"),
+        &corpus,
+        Box::new(move |b| {
+            let finder = memchr::memmem::FinderBuilder::new()
+                .build_forward_with_ranker(Binary, NEEDLE);
+            b.iter(|| {
+                assert_eq!(1, finder.find_iter(&bin).count());
+            });
+        }),
+    );
+}
+
+/// A byte-frequency table that is good for scanning binary executables.
+struct Binary;
+
+impl HeuristicFrequencyRank for Binary {
+    fn rank(&self, byte: u8) -> u8 {
+        const TABLE: [u8; 256] = [
+            255, 128, 61, 43, 50, 41, 27, 28, 57, 15, 21, 13, 24, 17, 17, 89,
+            58, 16, 11, 7, 14, 23, 7, 6, 24, 9, 6, 5, 9, 4, 7, 16, 68, 11, 9,
+            6, 88, 7, 4, 4, 23, 9, 4, 8, 8, 5, 10, 4, 30, 11, 9, 24, 11, 5, 5,
+            5, 19, 11, 6, 17, 9, 9, 6, 8, 48, 58, 11, 14, 53, 40, 9, 9, 254,
+            35, 3, 6, 52, 23, 6, 6, 27, 4, 7, 11, 14, 13, 10, 11, 11, 5, 2,
+            10, 16, 12, 6, 19, 19, 20, 5, 14, 16, 31, 19, 7, 14, 20, 4, 4, 19,
+            8, 18, 20, 24, 1, 25, 19, 58, 29, 10, 5, 15, 20, 2, 2, 9, 4, 3, 5,
+            51, 11, 4, 53, 23, 39, 6, 4, 13, 81, 4, 186, 5, 67, 3, 2, 15, 0,
+            0, 1, 3, 2, 0, 0, 5, 0, 0, 0, 2, 0, 0, 0, 12, 2, 1, 1, 3, 1, 1, 1,
+            6, 1, 2, 1, 3, 1, 1, 2, 9, 1, 1, 0, 2, 2, 4, 4, 11, 6, 7, 3, 6, 9,
+            4, 5, 46, 18, 8, 18, 17, 3, 8, 20, 16, 10, 3, 7, 175, 4, 6, 7, 13,
+            3, 7, 3, 3, 1, 3, 3, 10, 3, 1, 5, 2, 0, 1, 2, 16, 3, 5, 1, 6, 1,
+            1, 2, 58, 20, 3, 14, 12, 2, 1, 3, 16, 3, 5, 8, 3, 1, 8, 6, 17, 6,
+            5, 3, 8, 6, 13, 175,
+        ];
+        TABLE[byte as usize]
+    }
+}

bench/src/memmem/mod.rs

@@ -96,6 +96,7 @@ use criterion::Criterion;
 
 use crate::{define, memmem::inputs::INPUTS};
 
+mod byterank;
 mod imp;
 mod inputs;
 mod sliceslice;
@@ -106,6 +107,7 @@ pub fn all(c: &mut Criterion) {
     oneshot_iter(c);
     prebuilt_iter(c);
     sliceslice::all(c);
+    byterank::all(c);
 }
 
 fn oneshot(c: &mut Criterion) {

src/memmem/byte_frequencies.rs → src/memmem/byterank/default.rs

@@ -1,4 +1,4 @@
-pub const BYTE_FREQUENCIES: [u8; 256] = [
+pub const RANK: [u8; 256] = [
     55,  // '\x00'
     52,  // '\x01'
     51,  // '\x02'

src/memmem/byterank/mod.rs

@@ -0,0 +1,89 @@
+mod default;
+
+/// This trait allows the user to customize the heuristic used to determine the
+/// relative frequency of a given byte in the dataset being searched.
+///
+/// The use of this trait can have a dramatic impact on performance depending
+/// on the type of data being searched. The details of why are explained in the
+/// docs of [`prefilter::Prefilter`]. To summarize, the core algorithm uses a
+/// prefilter to quickly identify candidate matches that are later verified
+/// more slowly. This prefilter is implemented in terms of trying to find
+/// `rare` bytes at specific offsets that will occur less frequently in the
+/// dataset. While the concept of a `rare` byte is similar for most datasets,
+/// there are some specific datasets (like binary executables) that have
+/// dramatically different byte distributions. For these datasets customizing
+/// the byte frequency heuristic can have a massive impact on performance, and
+/// might even need to be done at runtime.
+///
+/// The default implementation of `HeuristicFrequencyRank` reads from the
+/// static frequency table defined in `src/memmem/byte_frequencies.rs`. This
+/// is optimal for most inputs, so if you are unsure of the impact of using a
+/// custom `HeuristicFrequencyRank` you should probably just use the default.
+///
+/// # Example
+///
+/// ```
+/// use memchr::memmem::{FinderBuilder, HeuristicFrequencyRank};
+///
+/// /// A byte-frequency table that is good for scanning binary executables.
+/// struct Binary;
+///
+/// impl HeuristicFrequencyRank for Binary {
+///     fn rank(&self, byte: u8) -> u8 {
+///         const TABLE: [u8; 256] = [
+///             255, 128, 61, 43, 50, 41, 27, 28, 57, 15, 21, 13, 24, 17, 17,
+///             89, 58, 16, 11, 7, 14, 23, 7, 6, 24, 9, 6, 5, 9, 4, 7, 16,
+///             68, 11, 9, 6, 88, 7, 4, 4, 23, 9, 4, 8, 8, 5, 10, 4, 30, 11,
+///             9, 24, 11, 5, 5, 5, 19, 11, 6, 17, 9, 9, 6, 8,
+///             48, 58, 11, 14, 53, 40, 9, 9, 254, 35, 3, 6, 52, 23, 6, 6, 27,
+///             4, 7, 11, 14, 13, 10, 11, 11, 5, 2, 10, 16, 12, 6, 19,
+///             19, 20, 5, 14, 16, 31, 19, 7, 14, 20, 4, 4, 19, 8, 18, 20, 24,
+///             1, 25, 19, 58, 29, 10, 5, 15, 20, 2, 2, 9, 4, 3, 5,
+///             51, 11, 4, 53, 23, 39, 6, 4, 13, 81, 4, 186, 5, 67, 3, 2, 15,
+///             0, 0, 1, 3, 2, 0, 0, 5, 0, 0, 0, 2, 0, 0, 0,
+///             12, 2, 1, 1, 3, 1, 1, 1, 6, 1, 2, 1, 3, 1, 1, 2, 9, 1, 1, 0,
+///             2, 2, 4, 4, 11, 6, 7, 3, 6, 9, 4, 5,
+///             46, 18, 8, 18, 17, 3, 8, 20, 16, 10, 3, 7, 175, 4, 6, 7, 13,
+///             3, 7, 3, 3, 1, 3, 3, 10, 3, 1, 5, 2, 0, 1, 2,
+///             16, 3, 5, 1, 6, 1, 1, 2, 58, 20, 3, 14, 12, 2, 1, 3, 16, 3, 5,
+///             8, 3, 1, 8, 6, 17, 6, 5, 3, 8, 6, 13, 175,
+///         ];
+///         TABLE[byte as usize]
+///     }
+/// }
+/// // Create a new finder with the custom heuristic.
+/// let finder = FinderBuilder::new()
+///     .build_forward_with_ranker(Binary, b"\x00\x00\xdd\xdd");
+/// // Find needle with custom heuristic.
+/// assert!(finder.find(b"\x00\x00\x00\xdd\xdd").is_some());
+/// ```
+pub trait HeuristicFrequencyRank {
+    /// Return the heuristic frequency rank of the given byte. A lower rank
+    /// means the byte is believed to occur less frequently in the haystack.
+    ///
+    /// Some uses of this heuristic may treat arbitrary absolute rank values as
+    /// significant. For example, an implementation detail in this crate may
+    /// determine that heuristic prefilters are inappropriate if every byte in
+    /// the needle has a "high" rank.
+    fn rank(&self, byte: u8) -> u8;
+}
+
+/// The default byte frequency heuristic that is good for most haystacks.
+pub(crate) struct DefaultFrequencyRank;
+
+impl HeuristicFrequencyRank for DefaultFrequencyRank {
+    fn rank(&self, byte: u8) -> u8 {
+        self::default::RANK[usize::from(byte)]
+    }
+}
+
+/// This permits passing any implementation of `HeuristicFrequencyRank` as a
+/// borrowed version of itself.
+impl<'a, R> HeuristicFrequencyRank for &'a R
+where
+    R: HeuristicFrequencyRank,
+{
+    fn rank(&self, byte: u8) -> u8 {
+        (**self).rank(byte)
+    }
+}

src/memmem/mod.rs

@@ -66,11 +66,12 @@ assert_eq!(None, finder.find(b"quux baz bar"));
 ```
 */
 
-pub use self::prefilter::Prefilter;
+pub use self::{byterank::HeuristicFrequencyRank, prefilter::Prefilter};
 
 use crate::{
     cow::CowBytes,
     memmem::{
+        byterank::DefaultFrequencyRank,
         prefilter::{Pre, PrefilterFn, PrefilterState},
         rabinkarp::NeedleHash,
         rarebytes::RareNeedleBytes,
@@ -145,7 +146,7 @@ macro_rules! define_memmem_simple_tests {
     };
 }
 
-mod byte_frequencies;
+mod byterank;
 #[cfg(memchr_runtime_simd)]
 mod genericsimd;
 mod prefilter;
@@ -712,7 +713,24 @@ impl FinderBuilder {
         &self,
         needle: &'n B,
     ) -> Finder<'n> {
-        Finder { searcher: Searcher::new(self.config, needle.as_ref()) }
+        self.build_forward_with_ranker(DefaultFrequencyRank, needle)
+    }
+
+    /// Build a forward finder using the given needle and a custom heuristic for
+    /// determining the frequency of a given byte in the dataset.
+    /// See [`HeuristicFrequencyRank`] for more details.
+    pub fn build_forward_with_ranker<
+        'n,
+        R: HeuristicFrequencyRank,
+        B: ?Sized + AsRef<[u8]>,
+    >(
+        &self,
+        ranker: R,
+        needle: &'n B,
+    ) -> Finder<'n> {
+        Finder {
+            searcher: Searcher::new(self.config, ranker, needle.as_ref()),
+        }
     }
 
     /// Build a reverse finder using the given needle from the current
@@ -817,14 +835,19 @@ enum SearcherKind {
 }
 
 impl<'n> Searcher<'n> {
-    fn new(config: SearcherConfig, needle: &'n [u8]) -> Searcher<'n> {
+    fn new<R: HeuristicFrequencyRank>(
+        config: SearcherConfig,
+        ranker: R,
+        needle: &'n [u8],
+    ) -> Searcher<'n> {
         use self::SearcherKind::*;
 
-        let ninfo = NeedleInfo::new(needle);
+        let ninfo = NeedleInfo::new(&ranker, needle);
         let mk = |kind: SearcherKind| {
             let prefn = prefilter::forward(
                 &config.prefilter,
                 &ninfo.rarebytes,
+                ranker,
                 needle,
             );
             Searcher { needle: CowBytes::new(needle), ninfo, prefn, kind }
@@ -1010,9 +1033,12 @@ impl<'n> Searcher<'n> {
 }
 
 impl NeedleInfo {
-    pub(crate) fn new(needle: &[u8]) -> NeedleInfo {
+    pub(crate) fn new<R: HeuristicFrequencyRank>(
+        ranker: &R,
+        needle: &[u8],
+    ) -> NeedleInfo {
         NeedleInfo {
-            rarebytes: RareNeedleBytes::forward(needle),
+            rarebytes: RareNeedleBytes::forward(ranker, needle),
             nhash: NeedleHash::forward(needle),
         }
     }

src/memmem/prefilter/fallback.rs

@@ -93,7 +93,10 @@ mod tests {
     use super::*;
 
     fn freqy_find(haystack: &[u8], needle: &[u8]) -> Option<usize> {
-        let ninfo = NeedleInfo::new(needle);
+        let ninfo = NeedleInfo::new(
+            &crate::memmem::byterank::DefaultFrequencyRank,
+            needle,
+        );
         let mut prestate = PrefilterState::new();
         find(&mut prestate, &ninfo, haystack, needle)
     }

src/memmem/prefilter/mod.rs

@@ -1,4 +1,6 @@
-use crate::memmem::{rarebytes::RareNeedleBytes, NeedleInfo};
+use crate::memmem::{
+    rarebytes::RareNeedleBytes, HeuristicFrequencyRank, NeedleInfo,
+};
 
 mod fallback;
 #[cfg(memchr_runtime_simd)]
@@ -287,9 +289,10 @@ impl PrefilterState {
 /// is the default). In general, we try to use an AVX prefilter, followed by
 /// SSE and then followed by a generic one based on memchr.
 #[inline(always)]
-pub(crate) fn forward(
+pub(crate) fn forward<R: HeuristicFrequencyRank>(
     config: &Prefilter,
     rare: &RareNeedleBytes,
+    ranker: R,
     needle: &[u8],
 ) -> Option<PrefilterFn> {
     if config.is_none() || needle.len() <= 1 {
@@ -327,7 +330,8 @@ pub(crate) fn forward(
     // Check that our rarest byte has a reasonably low rank. The main issue
     // here is that the fallback prefilter can perform pretty poorly if it's
     // given common bytes. So we try to avoid the worst cases here.
-    let (rare1_rank, _) = rare.as_ranks(needle);
+    let (rare1, _) = rare.as_rare_bytes(needle);
+    let rare1_rank = usize::from(ranker.rank(rare1));
     if rare1_rank <= MAX_FALLBACK_RANK {
         // SAFETY: fallback::find is safe to call in all environments.
         return unsafe { Some(PrefilterFn::new(fallback::find)) };