word index for fast lookup
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::word::Word;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WordIndex {
|
||||
by_length: HashMap<usize, Vec<usize>>,
|
||||
letter_idx: HashMap<(usize, usize, u8), Vec<u64>>,
|
||||
chunk_counts: HashMap<usize, usize>,
|
||||
allowed: HashMap<(usize, usize), u32>,
|
||||
}
|
||||
|
||||
impl WordIndex {
|
||||
pub fn build(words: &[Word]) -> Self {
|
||||
let mut by_length: HashMap<usize, Vec<usize>> = HashMap::new();
|
||||
|
||||
for (i, w) in words.iter().enumerate() {
|
||||
by_length.entry(w.length).or_default().push(i);
|
||||
}
|
||||
|
||||
let mut letter_idx: HashMap<(usize, usize, u8), Vec<u64>> = HashMap::new();
|
||||
let mut chunk_counts: HashMap<usize, usize> = HashMap::new();
|
||||
|
||||
for (&len, indices) in &by_length {
|
||||
let n = indices.len();
|
||||
let chunks = (n + 63) / 64;
|
||||
chunk_counts.insert(len, chunks);
|
||||
|
||||
for (bit_pos, &word_idx) in indices.iter().enumerate() {
|
||||
let word = &words[word_idx];
|
||||
let chunk = bit_pos / 64;
|
||||
let bit = bit_pos % 64;
|
||||
|
||||
for pos in 0..word.length {
|
||||
let letter = word.char_at(pos);
|
||||
let key = (len, pos, letter);
|
||||
let entry = letter_idx.entry(key).or_insert_with(|| vec![0u64; chunks]);
|
||||
entry[chunk] |= 1u64 << bit;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut allowed: HashMap<(usize, usize), u32> = HashMap::new();
|
||||
for (&len, indices) in &by_length {
|
||||
for pos in 0..len {
|
||||
let mut mask = 0u32;
|
||||
for &wi in indices {
|
||||
let letter = words[wi].char_at(pos);
|
||||
mask |= 1u32 << (letter - b'A');
|
||||
}
|
||||
allowed.insert((len, pos), mask);
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
by_length,
|
||||
letter_idx,
|
||||
chunk_counts,
|
||||
allowed,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn allowed_letters(&self, len: usize, pos: usize) -> u32 {
|
||||
self.allowed.get(&(len, pos)).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn chunk_count(&self, len: usize) -> Option<usize> {
|
||||
self.chunk_counts.get(&len).copied()
|
||||
}
|
||||
|
||||
pub fn get_letter_bits(&self, len: usize, pos: usize, letter: u8) -> Option<&Vec<u64>> {
|
||||
self.letter_idx.get(&(len, pos, letter))
|
||||
}
|
||||
|
||||
pub fn words_of_length(&self, len: usize) -> &[usize] {
|
||||
self.by_length.get(&len).map(|v| v.as_slice()).unwrap_or(&[])
|
||||
}
|
||||
|
||||
pub fn count_of_length(&self, len: usize) -> usize {
|
||||
self.by_length.get(&len).map(|v| v.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Get candidates matching a pattern constraint.
|
||||
/// `constraints` is a list of (position, letter) pairs that must match.
|
||||
/// Returns indices into the master word list.
|
||||
pub fn candidates(&self, len: usize, constraints: &[(usize, u8)]) -> Vec<usize> {
|
||||
let indices = match self.by_length.get(&len) {
|
||||
Some(v) => v,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
|
||||
let chunks = match self.chunk_counts.get(&len) {
|
||||
Some(&c) => c,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
|
||||
if constraints.is_empty() {
|
||||
return indices.clone();
|
||||
}
|
||||
|
||||
// intersect bitsets for all constraints
|
||||
let mut result = vec![u64::MAX; chunks];
|
||||
// mask off unused bits in last chunk
|
||||
let n = indices.len();
|
||||
let last_bits = n % 64;
|
||||
if last_bits > 0 {
|
||||
result[chunks - 1] = (1u64 << last_bits) - 1;
|
||||
}
|
||||
|
||||
for &(pos, letter) in constraints {
|
||||
let key = (len, pos, letter);
|
||||
match self.letter_idx.get(&key) {
|
||||
Some(bits) => {
|
||||
for (i, chunk) in result.iter_mut().enumerate() {
|
||||
*chunk &= bits[i];
|
||||
}
|
||||
}
|
||||
None => return Vec::new(), // no words have this letter at this position
|
||||
}
|
||||
}
|
||||
|
||||
// extract set bits as word indices
|
||||
let mut out = Vec::new();
|
||||
for (chunk_i, &bits) in result.iter().enumerate() {
|
||||
let mut b = bits;
|
||||
while b != 0 {
|
||||
let bit = b.trailing_zeros() as usize;
|
||||
out.push(indices[chunk_i * 64 + bit]);
|
||||
b &= b - 1;
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Count matching candidates using popcount - no list allocation.
|
||||
pub fn candidate_count(&self, len: usize, constraints: &[(usize, u8)]) -> usize {
|
||||
let indices = match self.by_length.get(&len) {
|
||||
Some(v) => v,
|
||||
None => return 0,
|
||||
};
|
||||
let chunks = match self.chunk_counts.get(&len) {
|
||||
Some(&c) => c,
|
||||
None => return 0,
|
||||
};
|
||||
if constraints.is_empty() {
|
||||
return indices.len();
|
||||
}
|
||||
|
||||
let n = indices.len();
|
||||
let mut count = 0usize;
|
||||
for chunk_i in 0..chunks {
|
||||
let mut bits = if chunk_i == chunks - 1 && n % 64 != 0 {
|
||||
(1u64 << (n % 64)) - 1
|
||||
} else {
|
||||
u64::MAX
|
||||
};
|
||||
for &(pos, letter) in constraints {
|
||||
let key = (len, pos, letter);
|
||||
match self.letter_idx.get(&key) {
|
||||
Some(b) => bits &= b[chunk_i],
|
||||
None => return 0,
|
||||
}
|
||||
}
|
||||
count += bits.count_ones() as usize;
|
||||
}
|
||||
count
|
||||
}
|
||||
|
||||
/// Check if at least one candidate exists - early exit.
|
||||
pub fn has_any_candidate(&self, len: usize, constraints: &[(usize, u8)]) -> bool {
|
||||
let indices = match self.by_length.get(&len) {
|
||||
Some(v) => v,
|
||||
None => return false,
|
||||
};
|
||||
let chunks = match self.chunk_counts.get(&len) {
|
||||
Some(&c) => c,
|
||||
None => return false,
|
||||
};
|
||||
if constraints.is_empty() {
|
||||
return !indices.is_empty();
|
||||
}
|
||||
|
||||
let n = indices.len();
|
||||
for chunk_i in 0..chunks {
|
||||
let mut bits = if chunk_i == chunks - 1 && n % 64 != 0 {
|
||||
(1u64 << (n % 64)) - 1
|
||||
} else {
|
||||
u64::MAX
|
||||
};
|
||||
for &(pos, letter) in constraints {
|
||||
let key = (len, pos, letter);
|
||||
match self.letter_idx.get(&key) {
|
||||
Some(b) => bits &= b[chunk_i],
|
||||
None => { bits = 0; break; }
|
||||
}
|
||||
}
|
||||
if bits != 0 {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_words() -> Vec<Word> {
|
||||
["CAT", "CAR", "COP", "DOG", "DIG", "BAT", "BAR", "BIG"]
|
||||
.iter()
|
||||
.map(|&w| Word::new(w.to_string(), 0.5))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_no_constraints() {
|
||||
let words = test_words();
|
||||
let idx = WordIndex::build(&words);
|
||||
let c = idx.candidates(3, &[]);
|
||||
assert_eq!(c.len(), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_single_constraint() {
|
||||
let words = test_words();
|
||||
let idx = WordIndex::build(&words);
|
||||
// words starting with 'C': CAT, CAR, COP
|
||||
let c = idx.candidates(3, &[(0, b'C')]);
|
||||
assert_eq!(c.len(), 3);
|
||||
for &i in &c {
|
||||
assert!(words[i].text.starts_with('C'));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_multiple_constraints() {
|
||||
let words = test_words();
|
||||
let idx = WordIndex::build(&words);
|
||||
// words starting with 'C' and ending with 'T': CAT
|
||||
let c = idx.candidates(3, &[(0, b'C'), (2, b'T')]);
|
||||
assert_eq!(c.len(), 1);
|
||||
assert_eq!(words[c[0]].text, "CAT");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_no_match() {
|
||||
let words = test_words();
|
||||
let idx = WordIndex::build(&words);
|
||||
let c = idx.candidates(3, &[(0, b'Z')]);
|
||||
assert!(c.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidates_wrong_length() {
|
||||
let words = test_words();
|
||||
let idx = WordIndex::build(&words);
|
||||
let c = idx.candidates(5, &[]);
|
||||
assert!(c.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user