add bench for automaton sstable streaming

sstable/benches/stream_bench.rs

@@ -1,13 +1,68 @@
 use std::collections::BTreeSet;
+use std::hint::black_box;
 use std::io;
 
 use common::file_slice::FileSlice;
 use criterion::{Criterion, criterion_group, criterion_main};
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
+use tantivy_fst::Automaton;
 use tantivy_sstable::{Dictionary, MonotonicU64SSTable};
 
 const CHARSET: &[u8] = b"abcdefghij";
+const AUTOMATON_PREFIX: &[u8] = b"ab";
+const NUM_AUTOMATON_MATCHES: usize = 1_017;
+
+// Matches `prefix.*`, but only implement can_match/will_always_match if configured to
+//
+// this allow comparing effects of optimisations depending on these functions
+struct HintedPrefixAutomaton<'a> {
+    prefix: &'a [u8],
+    can_match_hint: bool,
+    always_match_hint: bool,
+}
+
+impl<'a> HintedPrefixAutomaton<'a> {
+    fn new(prefix: &'a [u8], can_match_hint: bool, always_match_hint: bool) -> Self {
+        Self {
+            prefix,
+            can_match_hint,
+            always_match_hint,
+        }
+    }
+}
+
+impl Automaton for HintedPrefixAutomaton<'_> {
+    type State = Option<usize>;
+
+    fn start(&self) -> Self::State {
+        Some(0)
+    }
+
+    fn is_match(&self, state: &Self::State) -> bool {
+        *state == Some(self.prefix.len())
+    }
+
+    fn can_match(&self, state: &Self::State) -> bool {
+        !self.can_match_hint || state.is_some()
+    }
+
+    fn will_always_match(&self, state: &Self::State) -> bool {
+        self.always_match_hint && self.is_match(state)
+    }
+
+    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
+        let Some(pos) = *state else { return None };
+        if pos == self.prefix.len() {
+            return Some(pos);
+        }
+        if self.prefix[pos] == byte {
+            Some(pos + 1)
+        } else {
+            None
+        }
+    }
+}
 
 fn generate_key(rng: &mut impl Rng) -> String {
     let len = rng.random_range(3..12);
@@ -56,6 +111,26 @@ fn stream_bench(
     count
 }
 
+fn automaton_bench(
+    dictionary: &Dictionary<MonotonicU64SSTable>,
+    can_match_hint: bool,
+    always_match_hint: bool,
+) -> usize {
+    let mut stream = dictionary
+        .search(HintedPrefixAutomaton::new(
+            AUTOMATON_PREFIX,
+            black_box(can_match_hint),
+            black_box(always_match_hint),
+        ))
+        .into_stream()
+        .unwrap();
+    let mut count = 0;
+    while stream.advance() {
+        count += 1;
+    }
+    count
+}
+
 pub fn criterion_benchmark(c: &mut Criterion) {
     let dict = prepare_sstable().unwrap();
     c.bench_function("short_scan_init", |b| {
@@ -63,7 +138,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     });
     c.bench_function("short_scan_init_and_scan", |b| {
         b.iter(|| {
-            assert_eq!(stream_bench(&dict, b"fa", b"faz", true), 971);
+            assert_eq!(stream_bench(&dict, b"fa", b"faz", true), 1051);
         })
     });
     c.bench_function("full_scan_init_and_scan_full_with_bound", |b| {
@@ -81,6 +156,18 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             count
         })
     });
+    c.bench_function("full_scan_prefix_automaton_no_hints", |b| {
+        b.iter(|| assert_eq!(automaton_bench(&dict, false, false), NUM_AUTOMATON_MATCHES))
+    });
+    c.bench_function("full_scan_prefix_automaton_can_match_hint_only", |b| {
+        b.iter(|| assert_eq!(automaton_bench(&dict, true, false), NUM_AUTOMATON_MATCHES))
+    });
+    c.bench_function("full_scan_prefix_automaton_always_match_hint_only", |b| {
+        b.iter(|| assert_eq!(automaton_bench(&dict, false, true), NUM_AUTOMATON_MATCHES))
+    });
+    c.bench_function("full_scan_prefix_automaton_both_hints", |b| {
+        b.iter(|| assert_eq!(automaton_bench(&dict, true, true), NUM_AUTOMATON_MATCHES))
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);