treap_map/
map.rs

1// Copyright 2019 Conflux Foundation. All rights reserved.
2// Conflux is free software and distributed under GNU General Public License.
3// See http://www.gnu.org/licenses/
4
5use crate::{
6    config::{ConsoliableWeight, KeyMngTrait},
7    search::{accumulate_weight_search, SearchDirection},
8    update::{ApplyOp, ApplyOpOutcome, InsertOp, RemoveOp},
9    Direction, NoWeight, SearchResult,
10};
11
12use super::{config::TreapMapConfig, node::Node};
13use malloc_size_of::{MallocSizeOf, MallocSizeOfOps};
14use rand::{RngCore, SeedableRng};
15use rand_xorshift::XorShiftRng;
16
17/// A treap map data structure.
18///
19/// See [`TreapMapConfig`][crate::TreapMapConfig] for more details.
20pub struct TreapMap<C: TreapMapConfig> {
21    /// The root node of the treap.
22    #[cfg(test)]
23    pub(crate) root: Option<Box<Node<C>>>,
24    #[cfg(not(test))]
25    root: Option<Box<Node<C>>>,
26
27    /// A map for recovering the `sort_key` from the `search_key`.
28    /// This is useful when the `sort_key` is derived from `search_key` and
29    /// `value`.
30    ext_map: C::ExtMap,
31
32    /// A random number generator used for generating priority values for new
33    /// nodes.
34    rng: XorShiftRng,
35}
36
37impl<C: TreapMapConfig> MallocSizeOf for TreapMap<C>
38where
39    Node<C>: MallocSizeOf,
40    C::ExtMap: MallocSizeOf,
41{
42    fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
43        self.root.size_of(ops) + self.ext_map.size_of(ops)
44    }
45}
46
47impl<C: TreapMapConfig> Default for TreapMap<C> {
48    fn default() -> Self { Self::new() }
49}
50
51impl<C: TreapMapConfig> TreapMap<C> {
52    pub fn new() -> TreapMap<C> {
53        TreapMap {
54            root: None,
55            rng: XorShiftRng::from_os_rng(),
56            ext_map: Default::default(),
57        }
58    }
59
60    pub fn new_with_rng(rng: XorShiftRng) -> TreapMap<C> {
61        TreapMap {
62            root: None,
63            rng,
64            ext_map: Default::default(),
65        }
66    }
67
68    pub fn len(&self) -> usize { self.ext_map.len() }
69
70    pub fn is_empty(&self) -> bool { self.ext_map.len() == 0 }
71
72    pub fn contains_key(&self, key: &C::SearchKey) -> bool {
73        self.get(key).is_some()
74    }
75
76    pub fn insert(
77        &mut self, key: C::SearchKey, value: C::Value, weight: C::Weight,
78    ) -> Option<C::Value> {
79        let sort_key = self.ext_map.make_sort_key(&key, &value);
80
81        let node = Node::new(key, value, sort_key, weight, self.rng.next_u64());
82
83        let (result, _, _) = Node::update_inner(
84            &mut self.root,
85            InsertOp {
86                node: Box::new(node),
87                ext_map: &mut self.ext_map,
88            },
89        );
90
91        result
92    }
93
94    pub fn remove(&mut self, key: &C::SearchKey) -> Option<C::Value> {
95        let sort_key = self.ext_map.get_sort_key(key)?;
96
97        let (result, _, _) = Node::update_inner(
98            &mut self.root,
99            RemoveOp {
100                key: (&sort_key, key),
101                ext_map: &mut self.ext_map,
102            },
103        );
104
105        result
106    }
107
108    /// Updates the value of a node with the given key in the treap map.
109    ///
110    /// # Parameters
111    /// - `key`: The search key of the node to be updated.
112    /// - `update`: A function that is called if a node with the given key
113    ///   already exists. It takes a mutable reference to the node and returns
114    ///   an `ApplyOpOutcome<T>` or a custom error `E`. See
115    ///   [`ApplyOpOutcome`][crate::ApplyOpOutcome] for more details.
116    /// - `insert`: A function that is called if a node with the given key does
117    ///   not exist. It takes a mutable reference to a random number generator
118    ///   (for computing priority for a [`Node`][crate::Node]) and should return
119    ///   a tuple containing a new `Node<C>` and a value of type `T`, or an
120    ///   error of type `E`.
121    ///   - WARNING: The key of the new node must match the key provided to the
122    ///     function.
123    pub fn update<U, I, T, E>(
124        &mut self, key: &C::SearchKey, update: U, insert: I,
125    ) -> Result<T, E>
126    where
127        U: FnOnce(&mut Node<C>) -> Result<ApplyOpOutcome<T>, E>,
128        I: FnOnce(&mut dyn RngCore) -> Result<(Node<C>, T), E>,
129    {
130        let sort_key = if let Some(sort_key) = self.ext_map.get_sort_key(key) {
131            sort_key
132        } else {
133            return match insert(&mut self.rng) {
134                Ok((node, ret)) => {
135                    self.insert(node.key, node.value, node.weight);
136                    Ok(ret)
137                }
138                Err(err) => Err(err),
139            };
140        };
141        let rng = &mut self.rng;
142        let (res, _, _) = Node::update_inner(
143            &mut self.root,
144            ApplyOp {
145                key: (&sort_key, key),
146                update,
147                insert: || insert(rng),
148                ext_map: &mut self.ext_map,
149            },
150        );
151        let (ret, maybe_node) = res?;
152        if let Some(node) = maybe_node {
153            self.insert(node.key, node.value, node.weight);
154        }
155        Ok(ret)
156    }
157
158    pub fn sum_weight(&self) -> C::Weight {
159        match &self.root {
160            Some(node) => node.sum_weight(),
161            None => C::Weight::empty(),
162        }
163    }
164
165    pub fn get(&self, key: &C::SearchKey) -> Option<&C::Value> {
166        let sort_key = self.ext_map.get_sort_key(key)?;
167        self.root.as_ref().and_then(|x| x.get(&sort_key, key))
168    }
169
170    #[inline]
171    pub fn get_by_weight(&self, weight: C::Weight) -> Option<&C::Value>
172    where C::Weight: Ord {
173        use SearchDirection::*;
174        self.search(|base, mid| {
175            if &weight < base {
176                Left
177            } else {
178                let right_base = C::Weight::consolidate(base, &mid.weight);
179                if weight < right_base {
180                    Stop
181                } else {
182                    Right(right_base)
183                }
184            }
185        })?
186        .maybe_value()
187    }
188
189    /// See details in [`crate::accumulate_weight_search`]
190    pub fn search<F>(&self, f: F) -> Option<SearchResult<'_, C, C::Weight>>
191    where F: FnMut(&C::Weight, &Node<C>) -> SearchDirection<C::Weight> {
192        Some(accumulate_weight_search(self.root.as_ref()?, f, |weight| {
193            weight
194        }))
195    }
196
197    /// See details in [`crate::accumulate_weight_search`]
198    /// If the search process does not require accessing 'weight', this function
199    /// can outperform `search` by eliminating the maintenance of the 'weight'
200    /// dimension.
201    pub fn search_no_weight<F>(
202        &self, mut f: F,
203    ) -> Option<SearchResult<'_, C, NoWeight>>
204    where F: FnMut(&Node<C>) -> SearchDirection<()> {
205        static NW: NoWeight = NoWeight;
206        Some(accumulate_weight_search(
207            self.root.as_ref()?,
208            |_, node| f(node).map_into(|_| NoWeight),
209            |_| &NW,
210        ))
211    }
212
213    pub fn iter(&self) -> Iter<'_, C> {
214        let mut iter = Iter { nodes: vec![] };
215        if let Some(ref n) = self.root {
216            iter.nodes.push(&**n);
217            iter.extend_path();
218        }
219        iter
220    }
221
222    pub fn iter_range(&self, key: &C::SearchKey) -> Iter<'_, C>
223    where C: TreapMapConfig<SortKey = ()> {
224        let mut iter = Iter { nodes: vec![] };
225        if let Some(ref n) = self.root {
226            iter.nodes.push(&**n);
227            iter.extend_path_with_key((&(), key));
228        }
229        iter
230    }
231
232    pub fn values(&self) -> impl Iterator<Item = &C::Value> {
233        self.iter().map(|node| &node.value)
234    }
235
236    pub fn key_values(
237        &self,
238    ) -> impl Iterator<Item = (&C::SearchKey, &C::Value)> {
239        self.iter().map(|node| (&node.key, &node.value))
240    }
241
242    #[cfg(any(test, feature = "testonly_code"))]
243    pub fn assert_consistency(&self)
244    where C::Weight: std::fmt::Debug {
245        if let Some(node) = self.root.as_ref() {
246            node.assert_consistency()
247        }
248    }
249}
250
251pub struct Iter<'a, C: TreapMapConfig> {
252    nodes: Vec<&'a Node<C>>,
253}
254
255impl<'a, C: TreapMapConfig> Clone for Iter<'a, C> {
256    fn clone(&self) -> Self {
257        Self {
258            nodes: self.nodes.clone(),
259        }
260    }
261}
262
263impl<'a, C: TreapMapConfig> Iter<'a, C> {
264    fn extend_path(&mut self) {
265        loop {
266            let node = *self.nodes.last().unwrap();
267            match node.left {
268                None => return,
269                Some(ref n) => self.nodes.push(&**n),
270            }
271        }
272    }
273
274    fn extend_path_with_key(&mut self, key: (&C::SortKey, &C::SearchKey)) {
275        loop {
276            let node = *self.nodes.last().unwrap();
277            match C::next_node_dir(key, (&node.sort_key, &node.key)) {
278                Some(Direction::Left) => {
279                    if let Some(left) = &node.left {
280                        self.nodes.push(left);
281                    } else {
282                        return;
283                    }
284                }
285                None => {
286                    return;
287                }
288                Some(Direction::Right) => {
289                    let node = self.nodes.pop().unwrap();
290                    if let Some(right) = &node.right {
291                        self.nodes.push(right);
292                    } else {
293                        return;
294                    }
295                }
296            }
297        }
298    }
299}
300
301impl<'a, C: TreapMapConfig> Iterator for Iter<'a, C> {
302    type Item = &'a Node<C>;
303
304    fn next(&mut self) -> Option<Self::Item> {
305        match self.nodes.pop() {
306            None => None,
307            Some(node) => {
308                if let Some(ref n) = node.right {
309                    self.nodes.push(&**n);
310                    self.extend_path();
311                }
312                Some(node)
313            }
314        }
315    }
316}
treap_map/map.rs

treap_map/
map.rs
