cruciverb/cruciverb-core/src/dict/index.rs

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());
    }
}