use super::config::ZiporaTrieConfig;
use super::trie::ZiporaTrie;
use crate::error::Result;
use crate::fsa::traits::{Trie, FiniteStateAutomaton};
use crate::succinct::RankSelectOps;
use crate::StateId;
#[derive(Debug)]
pub struct ZiporaTrieMap<V: Copy, R = crate::succinct::RankSelectInterleaved256>
where
R: RankSelectOps,
{
trie: ZiporaTrie<R>,
values: Vec<Option<V>>,
}
impl<V: Copy, R> ZiporaTrieMap<V, R>
where
R: RankSelectOps + Default,
{
pub fn new() -> Self {
Self {
trie: ZiporaTrie::new(),
values: Vec::new(),
}
}
pub fn with_config(config: ZiporaTrieConfig) -> Self {
Self {
trie: ZiporaTrie::with_config(config),
values: Vec::new(),
}
}
pub fn insert(&mut self, key: &[u8], value: V) -> Result<Option<V>> {
let state_id = <ZiporaTrie<R> as Trie>::insert(&mut self.trie, key)?;
let idx = state_id as usize;
if idx >= self.values.len() {
self.values.resize(idx + 1, None);
}
let prev = self.values[idx];
self.values[idx] = Some(value);
Ok(prev)
}
pub fn get(&self, key: &[u8]) -> Option<V> {
if !self.trie.contains(key) {
return None;
}
let state_id = self.find_state_for_key(key)?;
self.values.get(state_id as usize).and_then(|&v| v)
}
fn find_state_for_key(&self, key: &[u8]) -> Option<StateId> {
let mut state = self.trie.root();
for &symbol in key {
state = self.trie.transition(state, symbol)?;
}
Some(state)
}
pub fn contains(&self, key: &[u8]) -> bool {
self.trie.contains(key)
}
pub fn len(&self) -> usize {
self.trie.len()
}
pub fn is_empty(&self) -> bool {
self.trie.is_empty()
}
pub fn keys(&self) -> Vec<Vec<u8>> {
self.trie.keys()
}
}
impl<V: Copy, R> Default for ZiporaTrieMap<V, R>
where
R: RankSelectOps + Default,
{
fn default() -> Self {
Self::new()
}
}