Possible optimization for mmap

Author: tibvdm

sa-index/src/array/mmap.rs

@@ -6,6 +6,89 @@ pub(super) fn read_u64_le(mmap: &Mmap, byte_offset: usize) -> u64 {
     u64::from_le_bytes(bytes)
 }
 
+/// Streaming sequential iterator over a contiguous range of a compressed or uncompressed
+/// mmap-backed suffix array.
+///
+/// Keeps `current_word` and `next_word` in local variables (register-allocated by the
+/// compiler) so that a new mmap read only occurs when crossing a 64-bit block boundary —
+/// roughly once per 1.6 entries for a 40-bit SA, vs 1–2 reads per entry with `get_mmap`.
+pub(crate) struct MmapSaRangeIter<'a> {
+    mmap: &'a Mmap,
+    data_offset: usize,
+    bits_per_value: usize,
+    mask: u64,
+    current_word: u64, // u64 block containing the next value to yield
+    next_word: u64,    // u64 block after current_word (pre-loaded)
+    block_idx: usize,  // index of current_word within the data section
+    bit_off: usize,    // bit offset of next value within current_word (0..64)
+    remaining: usize,  // entries left to yield
+}
+
+impl<'a> MmapSaRangeIter<'a> {
+    pub(crate) fn new(
+        mmap: &'a Mmap,
+        data_offset: usize,
+        bits_per_value: usize,
+        start: usize,
+        end: usize,
+    ) -> Self {
+        let remaining = end.saturating_sub(start);
+        if remaining == 0 {
+            return Self {
+                mmap, data_offset, bits_per_value,
+                mask: 0, current_word: 0, next_word: 0,
+                block_idx: 0, bit_off: 0, remaining: 0,
+            };
+        }
+
+        // (1u64 << 64) overflows; use u64::MAX for the 64-bit uncompressed case
+        let mask = if bits_per_value == 64 { u64::MAX } else { (1u64 << bits_per_value) - 1 };
+
+        let bit_pos   = start * bits_per_value;
+        let block_idx = bit_pos / 64;
+        let bit_off   = bit_pos % 64;
+
+        let current_word = read_u64_le(mmap, data_offset + block_idx * 8);
+        let next_off     = data_offset + (block_idx + 1) * 8;
+        let next_word    = if next_off + 8 <= mmap.len() { read_u64_le(mmap, next_off) } else { 0 };
+
+        Self { mmap, data_offset, bits_per_value, mask, current_word, next_word, block_idx, bit_off, remaining }
+    }
+}
+
+impl Iterator for MmapSaRangeIter<'_> {
+    type Item = i64;
+
+    #[inline]
+    fn next(&mut self) -> Option<i64> {
+        if self.remaining == 0 { return None; }
+        self.remaining -= 1;
+
+        let val = if self.bit_off + self.bits_per_value <= 64 {
+            // Value fits entirely within current_word
+            (self.current_word >> (64 - self.bit_off - self.bits_per_value)) & self.mask
+        } else {
+            // Value spans current_word and next_word
+            let end_off = (self.bit_off + self.bits_per_value) % 64;
+            ((self.current_word << end_off) | (self.next_word >> (64 - end_off))) & self.mask
+        };
+
+        // Advance bit cursor; load next word from mmap only on block-boundary crossing
+        self.bit_off += self.bits_per_value;
+        if self.bit_off >= 64 {
+            self.bit_off   -= 64;
+            self.block_idx += 1;
+            self.current_word = self.next_word;
+            let next_off = self.data_offset + (self.block_idx + 1) * 8;
+            self.next_word = if next_off + 8 <= self.mmap.len() {
+                read_u64_le(self.mmap, next_off)
+            } else { 0 };
+        }
+
+        Some(val as i64)
+    }
+}
+
 /// Returns the suffix array value at the given index from a memory-mapped file.
 pub(super) fn get_mmap(mmap: &Mmap, data_offset: usize, bits_per_value: usize, index: usize) -> i64 {
     if bits_per_value == 64 {

sa-index/src/array/mod.rs

@@ -105,6 +105,22 @@ impl SuffixArray {
         }
     }
 
+    /// Returns a streaming iterator over SA entries in `[start, end)`.
+    /// For `MmapBacked`, uses a streaming bit reader that keeps u64 words in CPU registers
+    /// and only accesses the mmap when crossing a 64-bit block boundary.
+    pub fn iter_range(&self, start: usize, end: usize) -> SuffixArrayRangeIter<'_> {
+        match self {
+            SuffixArray::Original(sa, _) =>
+                SuffixArrayRangeIter::Original(sa.get(start..end).unwrap_or(&[]).iter()),
+            SuffixArray::Compressed(ba, _) =>
+                SuffixArrayRangeIter::Compressed { ba, idx: start, end },
+            SuffixArray::MmapBacked { mmap, data_offset, bits_per_value, .. } =>
+                SuffixArrayRangeIter::Mmap(mmap::MmapSaRangeIter::new(
+                    mmap, *data_offset, *bits_per_value, start, end,
+                )),
+        }
+    }
+
     /// Issues an OS prefetch hint (`MADV_WILLNEED`) for the mmap pages covering SA indices
     /// `lo..hi_exclusive`.  No-op for in-memory variants and on non-Unix platforms.
     #[inline]
@@ -122,6 +138,35 @@ impl SuffixArray {
     }
 }
 
+/// Iterator over a contiguous range of SA entries.
+/// - `Original`: wraps a slice iterator — zero overhead.
+/// - `Compressed`: calls `BitArray::get()` per entry (heap-hot data).
+/// - `Mmap`: uses `MmapSaRangeIter` — keeps u64 words in registers, mmap only touched at block boundaries.
+pub enum SuffixArrayRangeIter<'a> {
+    Original(std::slice::Iter<'a, i64>),
+    Compressed { ba: &'a BitArray, idx: usize, end: usize },
+    #[doc(hidden)]
+    Mmap(mmap::MmapSaRangeIter<'a>),
+}
+
+impl Iterator for SuffixArrayRangeIter<'_> {
+    type Item = i64;
+
+    #[inline]
+    fn next(&mut self) -> Option<i64> {
+        match self {
+            Self::Original(iter) => iter.next().copied(),
+            Self::Compressed { ba, idx, end } => {
+                if *idx >= *end { return None; }
+                let val = ba.get(*idx) as i64;
+                *idx += 1;
+                Some(val)
+            }
+            Self::Mmap(iter) => iter.next(),
+        }
+    }
+}
+
 impl WriteBinary for SuffixArray {
     fn write_binary<W: Write>(self, writer: &mut W) -> Result<(), Box<dyn std::error::Error>> {
         match self {
@@ -411,4 +456,77 @@ mod tests {
             assert_eq!(loaded.get(i), i as i64 + 1);
         }
     }
+
+    /// Verifies that `iter_range(start, end)` yields the same values as repeated `get(i)`
+    /// calls for all three SA variants, including ranges that cross multiple 64-bit block
+    /// boundaries and a non-zero start offset.
+    #[test]
+    fn test_iter_range_matches_get() {
+        use tempdir::TempDir;
+
+        // 20 values — enough to cross multiple 64-bit blocks for a 40-bit SA (8 entries/cycle)
+        let values: Vec<i64> = (0..20).map(|i| i * 12345 + 7).collect();
+
+        // --- Original (Vec<i64>) ---
+        {
+            let sa = SuffixArray::Original(values.clone(), 1);
+            let collected: Vec<i64> = sa.iter_range(3, 17).collect();
+            let expected: Vec<i64> = (3..17).map(|i| sa.get(i)).collect();
+            assert_eq!(collected, expected, "Original iter_range mismatch");
+        }
+
+        // --- Compressed (BitArray, 40-bit) ---
+        {
+            let mut ba = BitArray::with_capacity(20, 40);
+            for (i, &v) in values.iter().enumerate() {
+                ba.set(i, v as u64);
+            }
+            let sa = SuffixArray::Compressed(ba, 1);
+            let collected: Vec<i64> = sa.iter_range(3, 17).collect();
+            let expected: Vec<i64> = (3..17).map(|i| sa.get(i)).collect();
+            assert_eq!(collected, expected, "Compressed iter_range mismatch");
+        }
+
+        // --- MmapBacked (40-bit compressed, via round-trip through file) ---
+        {
+            let tmp = TempDir::new("iter_range_mmap").unwrap();
+            let path = tmp.path().join("sa.bin");
+            let mut file = std::fs::File::create(&path).unwrap();
+            dump_compressed_suffix_array(values.clone(), 1, 40, &mut file).unwrap();
+            drop(file);
+
+            let sa = SuffixArray::read_binary_mmap(&path).unwrap();
+            let collected: Vec<i64> = sa.iter_range(3, 17).collect();
+            let expected: Vec<i64> = (3..17).map(|i| sa.get(i)).collect();
+            assert_eq!(collected, expected, "MmapBacked iter_range mismatch");
+        }
+
+        // --- MmapBacked (64-bit uncompressed) ---
+        {
+            let tmp = TempDir::new("iter_range_mmap64").unwrap();
+            let path = tmp.path().join("sa64.bin");
+            let mut file = std::fs::File::create(&path).unwrap();
+            dump_suffix_array(values.clone(), 1, &mut file).unwrap();
+            drop(file);
+
+            let sa = SuffixArray::read_binary_mmap(&path).unwrap();
+            let collected: Vec<i64> = sa.iter_range(3, 17).collect();
+            let expected: Vec<i64> = (3..17).map(|i| sa.get(i)).collect();
+            assert_eq!(collected, expected, "MmapBacked-64 iter_range mismatch");
+        }
+
+        // --- Edge case: empty range (start == end) ---
+        {
+            let sa = SuffixArray::Original(values.clone(), 1);
+            let collected: Vec<i64> = sa.iter_range(5, 5).collect();
+            assert!(collected.is_empty(), "Empty range should yield nothing");
+        }
+
+        // --- Edge case: inverted range (start > end) — must not panic ---
+        {
+            let sa = SuffixArray::Original(values.clone(), 1);
+            let collected: Vec<i64> = sa.iter_range(10, 3).collect();
+            assert!(collected.is_empty(), "Inverted range should yield nothing");
+        }
+    }
 }

sa-index/src/sa_searcher.rs

@@ -426,10 +426,9 @@ impl Searcher {
             if let BoundSearchResult::SearchResult((min_bound, max_bound)) = search_bound_result {
                 // try all the partially matched suffixes and store the matching suffixes in an
                 // array (stop when our max number of matches is reached)
-                let mut sa_index = min_bound;
                 let t_iter = Instant::now();
-                while sa_index < max_bound {
-                    let suffix = self.sa.get(sa_index) as usize;
+                for suffix in self.sa.iter_range(min_bound, max_bound) {
+                    let suffix = suffix as usize;
 
                     if suffix >= skip {
                         let match_start = suffix - skip;
@@ -465,8 +464,6 @@ impl Searcher {
                             }
                         }
                     }
-
-                    sa_index += 1;
                 }
                 self.match_iter_ns.fetch_add(t_iter.elapsed().as_nanos() as u64, Ordering::Relaxed);
             }