candding 0.1.0

//! minicoil's own word vocabulary and stem alias table, read from one JSON file.
//! Whole words resolve against it after `words::reconstruct`; it is not the tokenizer vocabulary.

use std::collections::HashMap;

use serde::Deserialize;

use crate::error::Result;
use crate::models::sparse::stemmer::stem;

/// `minicoil.triplet.model.vocab`'s schema: `vocab` is the ordered word list (id = position + 1, 0 means
/// unrecognized), `stem_mapping` maps other surface forms to its entries.
#[derive(Deserialize)]
struct Schema {
  vocab: Vec<String>,
  stem_mapping: HashMap<String, String>,
}

/// Word-to-id lookup plus the shipped stem alias table.
#[derive(Debug)]
pub struct Vocab {
  word_ids: HashMap<String, u32>,
  stem_mapping: HashMap<String, String>,
  len: usize,
}

impl Vocab {
  /// Parses the vocabulary JSON. Ids follow the rank of each distinct word's first appearance, so they stay in
  /// `1..=len()` even if the file repeats a word and no gather can reach past the matrix.
  pub fn parse(json: &str) -> Result<Self> {
    let schema: Schema = serde_json::from_str(json)?;
    let mut word_ids = HashMap::with_capacity(schema.vocab.len());
    let mut next_id = 1u32;
    for word in schema.vocab {
      word_ids.entry(word).or_insert_with(|| {
        let id = next_id;
        next_id += 1;
        id
      });
    }
    Ok(Self {
      len: word_ids.len(),
      word_ids,
      stem_mapping: schema.stem_mapping,
    })
  }

  /// Number of distinct words this vocabulary recognizes, not counting the reserved "unrecognized" id 0.
  pub fn len(&self) -> usize {
    self.len
  }

  pub fn is_empty(&self) -> bool {
    self.len == 0
  }

  /// Resolves `word` in the reference's order: vocabulary, then the alias table, then the alias table with the
  /// stem; an alias naming an absent word resolves to `None`.
  pub fn resolve(&self, word: &str) -> Option<u32> {
    if let Some(&id) = self.word_ids.get(word) {
      return Some(id);
    }
    if let Some(canonical) = self.stem_mapping.get(word) {
      return self.word_ids.get(canonical).copied();
    }
    let stemmed = stem(word);
    self
      .stem_mapping
      .get(&stemmed)
      .and_then(|canonical| self.word_ids.get(canonical).copied())
  }
}

#[cfg(test)]
mod tests;