#![allow(clippy::module_name_repetitions)]
use crate::{Node, SetTrie};
pub struct EntryBuilder<'a, K, T, IK>
where
IK: Iterator<Item = K> + 'a,
K: Ord,
{
node: &'a mut Node<K, T>,
keys: IK,
}
impl<'a, K, T, IK> EntryBuilder<'a, K, T, IK>
where
IK: Iterator<Item = K> + 'a,
K: Ord,
{
pub(crate) fn new(trie: &'a mut SetTrie<K, T>, keys: IK) -> Self {
EntryBuilder {
node: &mut trie.0,
keys,
}
}
pub(crate) fn from_node(node: &'a mut Node<K, T>, keys: IK) -> Self {
EntryBuilder { node, keys }
}
}
pub enum Entry<'a, K, T>
where
K: Ord,
{
Created(CreatedEntry<'a, K, T>),
Existing(ExistingEntry<'a, K, T>),
}
pub struct CreatedEntry<'a, K, T>
where
K: Ord,
{
node: &'a mut Node<K, T>,
}
pub struct ExistingEntry<'a, K, T>
where
K: Ord,
{
node: &'a mut Node<K, T>,
}
impl<'a, K, T, IK> EntryBuilder<'a, K, T, IK>
where
IK: Iterator<Item = K> + 'a,
K: Ord,
{
pub fn and_extend(self, default: impl IntoIterator<Item = T>) -> Entry<'a, K, T> {
match self.or_create() {
Entry::Existing(e) => {
e.node.leaves.extend(default.into_iter());
Entry::Existing(e)
}
Entry::Created(e) => {
e.node.leaves.extend(default.into_iter());
Entry::Created(e)
}
}
}
pub fn and_insert(self, default: T) -> Entry<'a, K, T> {
match self.or_create() {
Entry::Existing(e) => {
e.node.leaves.push(default);
Entry::Existing(e)
}
Entry::Created(e) => {
e.node.leaves.push(default);
Entry::Created(e)
}
}
}
pub fn or_extend(self, default: impl IntoIterator<Item = T>) -> Entry<'a, K, T> {
match self.or_create() {
entry @ Entry::Existing(_) => entry,
Entry::Created(e) => {
e.node.leaves.extend(default.into_iter());
Entry::Created(e)
}
}
}
pub fn or_insert(self, default: T) -> Entry<'a, K, T> {
match self.or_create() {
entry @ Entry::Existing(_) => entry,
Entry::Created(e) => {
e.node.leaves.push(default);
Entry::Created(e)
}
}
}
pub fn or_create(self) -> Entry<'a, K, T> {
let mut node = self.node;
let mut created = false;
for key in self.keys {
node = match node.children.binary_search_by(|(k, _)| k.cmp(&key)) {
Ok(idx) => &mut (node.children[idx].1),
Err(idx) => {
created = true;
node.children.insert(idx, (key, Node::new()));
&mut (node.children[idx].1)
}
}
}
if created {
return Entry::Created(CreatedEntry { node });
}
Entry::Existing(ExistingEntry { node })
}
pub fn find(self) -> Option<ExistingEntry<'a, K, T>> {
let mut node = self.node;
for key in self.keys {
node = match node.children.binary_search_by(|(k, _)| k.cmp(&key)) {
Ok(idx) => &mut (node.children[idx].1),
Err(_) => return None,
}
}
Some(ExistingEntry { node })
}
pub fn items(self) -> Option<&'a Vec<T>> {
self.find().map(|node| &node.node.leaves)
}
pub fn items_mut(self) -> Option<&'a mut Vec<T>> {
self.find().map(|node| &mut node.node.leaves)
}
}
impl<'a, K, T> Entry<'a, K, T>
where
K: Ord,
{
fn node(&self) -> &Node<K, T> {
match self {
Entry::Existing(e) => e.node,
Entry::Created(e) => e.node,
}
}
fn node_mut(&mut self) -> &mut Node<K, T> {
match self {
Entry::Existing(e) => e.node,
Entry::Created(e) => e.node,
}
}
#[must_use]
pub fn items(&self) -> &Vec<T> {
&self.node().leaves
}
#[must_use]
pub fn items_mut(&mut self) -> &mut Vec<T> {
&mut self.node_mut().leaves
}
#[must_use]
pub fn entry<IK: IntoIterator<Item = K>>(
self,
keys: IK,
) -> EntryBuilder<'a, K, T, IK::IntoIter> {
match self {
Entry::Created(e) => EntryBuilder::from_node(e.node, keys.into_iter()),
Entry::Existing(e) => EntryBuilder::from_node(e.node, keys.into_iter()),
}
}
}