set_trie/
entry.rs

1#![allow(clippy::module_name_repetitions)]
2
3use crate::{Node, SetTrie};
4
5/// `EntryBuilder` for the [entry](SetTrie::entry) method. Entries are lazily evaluated, thus the builder
6/// is used to provide the configuration, while the [entry](Entry) is already evaluated.
7pub struct EntryBuilder<'a, K, T, IK>
8where
9    IK: Iterator<Item = K> + 'a,
10    K: Ord,
11{
12    node: &'a mut Node<K, T>,
13    keys: IK,
14}
15
16impl<'a, K, T, IK> EntryBuilder<'a, K, T, IK>
17where
18    IK: Iterator<Item = K> + 'a,
19    K: Ord,
20{
21    pub(crate) fn new(trie: &'a mut SetTrie<K, T>, keys: IK) -> Self {
22        EntryBuilder {
23            node: &mut trie.0,
24            keys,
25        }
26    }
27
28    pub(crate) fn from_node(node: &'a mut Node<K, T>, keys: IK) -> Self {
29        EntryBuilder { node, keys }
30    }
31}
32
33/// A view into a node of a [`SetTrie`](SetTrie), either created or already existing.
34pub enum Entry<'a, K, T>
35where
36    K: Ord,
37{
38    /// Indicates that the entry was created.
39    Created(CreatedEntry<'a, K, T>),
40
41    /// Indicates that the entry already existed.
42    Existing(ExistingEntry<'a, K, T>),
43}
44
45/// Indicates that the entry was created.
46pub struct CreatedEntry<'a, K, T>
47where
48    K: Ord,
49{
50    node: &'a mut Node<K, T>,
51}
52
53/// Indicates that the entry already exists.
54pub struct ExistingEntry<'a, K, T>
55where
56    K: Ord,
57{
58    node: &'a mut Node<K, T>,
59}
60
61impl<'a, K, T, IK> EntryBuilder<'a, K, T, IK>
62where
63    IK: Iterator<Item = K> + 'a,
64    K: Ord,
65{
66    /// Extends the entry, creating it if needed
67    pub fn and_extend(self, default: impl IntoIterator<Item = T>) -> Entry<'a, K, T> {
68        match self.or_create() {
69            Entry::Existing(e) => {
70                e.node.leaves.extend(default.into_iter());
71                Entry::Existing(e)
72            }
73            Entry::Created(e) => {
74                e.node.leaves.extend(default.into_iter());
75                Entry::Created(e)
76            }
77        }
78    }
79
80    /// Inserts into the entry, creating it if needed
81    pub fn and_insert(self, default: T) -> Entry<'a, K, T> {
82        match self.or_create() {
83            Entry::Existing(e) => {
84                e.node.leaves.push(default);
85                Entry::Existing(e)
86            }
87            Entry::Created(e) => {
88                e.node.leaves.push(default);
89                Entry::Created(e)
90            }
91        }
92    }
93
94    /// Finds the entry, and if it does not exist, extends with the provided value.
95    pub fn or_extend(self, default: impl IntoIterator<Item = T>) -> Entry<'a, K, T> {
96        match self.or_create() {
97            entry @ Entry::Existing(_) => entry,
98            Entry::Created(e) => {
99                e.node.leaves.extend(default.into_iter());
100                Entry::Created(e)
101            }
102        }
103    }
104
105    /// Finds the entry, and if it does not exist, inserts the value.
106    pub fn or_insert(self, default: T) -> Entry<'a, K, T> {
107        match self.or_create() {
108            entry @ Entry::Existing(_) => entry,
109            Entry::Created(e) => {
110                e.node.leaves.push(default);
111                Entry::Created(e)
112            }
113        }
114    }
115
116    /// Finds the entry, and if it does not exist, creates it.
117    pub fn or_create(self) -> Entry<'a, K, T> {
118        let mut node = self.node;
119        let mut created = false;
120
121        for key in self.keys {
122            node = match node.children.binary_search_by(|(k, _)| k.cmp(&key)) {
123                Ok(idx) => &mut (node.children[idx].1),
124                Err(idx) => {
125                    created = true;
126                    node.children.insert(idx, (key, Node::new()));
127                    &mut (node.children[idx].1)
128                }
129            }
130        }
131
132        if created {
133            return Entry::Created(CreatedEntry { node });
134        }
135        Entry::Existing(ExistingEntry { node })
136    }
137
138    /// Finds the entry, but does not create one. This method short circuits on the first missing
139    /// key.
140    pub fn find(self) -> Option<ExistingEntry<'a, K, T>> {
141        let mut node = self.node;
142
143        for key in self.keys {
144            node = match node.children.binary_search_by(|(k, _)| k.cmp(&key)) {
145                Ok(idx) => &mut (node.children[idx].1),
146                Err(_) => return None,
147            }
148        }
149        Some(ExistingEntry { node })
150    }
151
152    /// Returns all associated items of an entry.
153    pub fn items(self) -> Option<&'a Vec<T>> {
154        self.find().map(|node| &node.node.leaves)
155    }
156
157    /// Mutably returns all associated items of an entry.
158    pub fn items_mut(self) -> Option<&'a mut Vec<T>> {
159        self.find().map(|node| &mut node.node.leaves)
160    }
161}
162
163impl<'a, K, T> Entry<'a, K, T>
164where
165    K: Ord,
166{
167    fn node(&self) -> &Node<K, T> {
168        match self {
169            Entry::Existing(e) => e.node,
170            Entry::Created(e) => e.node,
171        }
172    }
173
174    fn node_mut(&mut self) -> &mut Node<K, T> {
175        match self {
176            Entry::Existing(e) => e.node,
177            Entry::Created(e) => e.node,
178        }
179    }
180
181    /// Returns all associated items of an entry.
182    #[must_use]
183    pub fn items(&self) -> &Vec<T> {
184        &self.node().leaves
185    }
186
187    /// Mutably returns all associated items of an entry.
188    #[must_use]
189    pub fn items_mut(&mut self) -> &mut Vec<T> {
190        &mut self.node_mut().leaves
191    }
192
193    /// Provides a view into a child of the entry. If you are sequentially inserting longer keys,
194    /// reusing the entry is more efficient than starting from the root.
195    ///
196    /// ```rust
197    /// // Equivalent to first inserting 0..1, then 0..2 etc.
198    ///
199    /// let mut trie = set_trie::SetTrie::new();
200    /// let mut current = trie.entry(0..1).or_insert(0);
201    ///     for i in 1..5000 {
202    ///         current = {
203    ///             current.entry(i-1..i).or_insert(i)
204    ///         }
205    ///     }
206    /// ```
207    #[must_use]
208    pub fn entry<IK: IntoIterator<Item = K>>(
209        self,
210        keys: IK,
211    ) -> EntryBuilder<'a, K, T, IK::IntoIter> {
212        match self {
213            Entry::Created(e) => EntryBuilder::from_node(e.node, keys.into_iter()),
214            Entry::Existing(e) => EntryBuilder::from_node(e.node, keys.into_iter()),
215        }
216    }
217}
set_trie/entry.rs

set_trie/
entry.rs
