avltriee/
search.rs

1use std::{cmp::Ordering, num::NonZeroU32, ops::Range};
2
3use crate::{Avltriee, AvltrieeAllocator, AvltrieeNode};
4
5pub(crate) type Edge = (Option<NonZeroU32>, Ordering);
6
7pub trait AvltrieeSearch<T, I: ?Sized, A: AvltrieeAllocator<T>>: AsRef<Avltriee<T, I, A>> {
8    fn cmp(left: &I, right: &I) -> Ordering;
9    fn value(&self, row: NonZeroU32) -> Option<&I>;
10    unsafe fn value_unchecked(&self, row: NonZeroU32) -> &I;
11    unsafe fn node_value_unchecked(&self, row: NonZeroU32) -> (&AvltrieeNode<T>, &I);
12
13    /// Search row of a value.
14    fn row(&self, value: &I) -> Option<NonZeroU32> {
15        let edge = self.edge(value);
16        (edge.1 == Ordering::Equal).then(|| edge.0).flatten()
17    }
18
19    /// Finds the edge of a node from the specified value.
20    fn edge(&self, value: &I) -> Edge {
21        let triee = self.as_ref();
22        let mut row: Option<NonZeroU32> = triee.root();
23        let mut ord = Ordering::Equal;
24        while let Some(row_inner) = row {
25            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
26            ord = Self::cmp(node_value, value);
27            match ord {
28                Ordering::Greater => {
29                    if node.left.is_some() {
30                        row = node.left;
31                    } else {
32                        break;
33                    }
34                }
35                Ordering::Equal => {
36                    break;
37                }
38                Ordering::Less => {
39                    if node.right.is_some() {
40                        row = node.right;
41                    } else {
42                        break;
43                    }
44                }
45            }
46        }
47        (row, ord)
48    }
49
50    /// Search >= value.
51    fn ge(&self, value: &I) -> Option<NonZeroU32> {
52        let triee = self.as_ref();
53        let mut row = triee.root();
54        let mut keep = None;
55        while let Some(row_inner) = row {
56            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
57            match Self::cmp(node_value, value) {
58                Ordering::Greater => {
59                    if node.left.is_some() {
60                        keep = row;
61                        row = node.left;
62                    } else {
63                        return row;
64                    }
65                }
66                Ordering::Equal => {
67                    return row;
68                }
69                Ordering::Less => {
70                    if node.right.is_some() {
71                        row = node.right;
72                    } else {
73                        break;
74                    }
75                }
76            }
77        }
78        keep
79    }
80
81    /// Search <= value.
82    fn le(&self, value: &I) -> Option<NonZeroU32> {
83        let triee = self.as_ref();
84        let mut row = triee.root();
85        let mut keep = None;
86        while let Some(row_inner) = row {
87            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
88            match Self::cmp(node_value, value) {
89                Ordering::Greater => {
90                    if node.left.is_some() {
91                        row = node.left;
92                    } else {
93                        break;
94                    }
95                }
96                Ordering::Equal => {
97                    return row;
98                }
99                Ordering::Less => {
100                    if node.right.is_some() {
101                        keep = row;
102                        row = node.right;
103                    } else {
104                        return row;
105                    }
106                }
107            }
108        }
109        keep
110    }
111
112    /// Search > value.
113    fn gt(&self, value: &I) -> Option<NonZeroU32> {
114        let triee = self.as_ref();
115        let mut row = triee.root();
116        let mut keep = None;
117        while let Some(row_inner) = row {
118            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
119            match Self::cmp(node_value, value) {
120                Ordering::Greater => {
121                    if node.left.is_some() {
122                        keep = row;
123                        row = node.left;
124                    } else {
125                        return row;
126                    }
127                }
128                Ordering::Equal => {
129                    if node.right.is_some() {
130                        return triee.min(node.right);
131                    }
132                    if let Some(parent) = node.parent {
133                        if unsafe { triee.node_unchecked(parent).left } == row {
134                            return node.parent;
135                        }
136                    }
137                    break;
138                }
139                Ordering::Less => {
140                    if node.right.is_some() {
141                        row = node.right;
142                    } else {
143                        break;
144                    }
145                }
146            }
147        }
148        keep
149    }
150
151    /// Search < value.
152    fn lt(&self, value: &I) -> Option<NonZeroU32> {
153        let triee = self.as_ref();
154        let mut row = triee.root();
155        let mut keep = None;
156        while let Some(row_inner) = row {
157            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
158            match Self::cmp(node_value, value) {
159                Ordering::Greater => {
160                    if node.left.is_some() {
161                        row = node.left;
162                    } else {
163                        break;
164                    }
165                }
166                Ordering::Equal => {
167                    if node.left.is_some() {
168                        return triee.max(node.left);
169                    }
170                    if let Some(parent) = node.parent {
171                        if unsafe { triee.node_unchecked(parent) }.right == row {
172                            return node.parent;
173                        }
174                    }
175                    break;
176                }
177                Ordering::Less => {
178                    if node.right.is_some() {
179                        keep = row;
180                        row = node.right;
181                    } else {
182                        return row;
183                    }
184                }
185            }
186        }
187        keep
188    }
189
190    /// Search with range value with custom ord.
191    fn range(&self, start_value: &I, end_value: &I) -> Option<Range<NonZeroU32>> {
192        let triee = self.as_ref();
193        let mut row = triee.root();
194        let mut start = None;
195        while let Some(row_inner) = row {
196            let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
197            match Self::cmp(node_value, start_value) {
198                Ordering::Greater => {
199                    start = row;
200                    if node.left.is_some() {
201                        row = node.left;
202                    } else {
203                        break;
204                    }
205                }
206                Ordering::Equal => {
207                    start = row;
208                    break;
209                }
210                Ordering::Less => {
211                    if node.right.is_some() {
212                        row = node.right;
213                    } else {
214                        break;
215                    }
216                }
217            }
218        }
219        if let Some(start) = start {
220            if Self::cmp(unsafe { self.value_unchecked(start) }, end_value) != Ordering::Greater {
221                row = triee.root();
222                let mut end = None;
223                while let Some(row_inner) = row {
224                    let (node, node_value) = unsafe { self.node_value_unchecked(row_inner) };
225                    match Self::cmp(node_value, end_value) {
226                        Ordering::Greater => {
227                            if node.left.is_some() {
228                                row = node.left;
229                            } else {
230                                break;
231                            }
232                        }
233                        Ordering::Equal => {
234                            end = row;
235                            break;
236                        }
237                        Ordering::Less => {
238                            end = row;
239                            if node.right.is_some() {
240                                row = node.right;
241                            } else {
242                                break;
243                            }
244                        }
245                    }
246                }
247                if let Some(end) = end {
248                    return Some(Range { start, end });
249                }
250            }
251        }
252        None
253    }
254}
avltriee/search.rs

avltriee/
search.rs
