competitive/data_structure/binary_search_tree/
node.rs

1use super::{Allocator, BstSeeker};
2use std::{cmp::Ordering, marker::PhantomData, ptr::NonNull};
3
4#[derive(Clone, Copy, Debug, Eq, PartialEq)]
5pub enum EqualSide {
6    Left,
7    Right,
8}
9
10impl EqualSide {
11    pub fn goes_left(self, ordering: Ordering) -> bool {
12        match ordering {
13            Ordering::Less => true,
14            Ordering::Equal => matches!(self, Self::Left),
15            Ordering::Greater => false,
16        }
17    }
18}
19
20pub trait BstSpec: Sized {
21    type Parent: ParentStrategy<Data = Self::Data>;
22    type Data;
23
24    fn top_down(_node: BstDataMutRef<'_, Self>) {}
25
26    fn bottom_up(_node: BstDataMutRef<'_, Self>) {}
27
28    fn merge(left: Option<BstRoot<Self>>, right: Option<BstRoot<Self>>) -> Option<BstRoot<Self>>;
29
30    fn split<Seeker>(
31        node: Option<BstRoot<Self>>,
32        seeker: Seeker,
33        equal_side: EqualSide,
34    ) -> (Option<BstRoot<Self>>, Option<BstRoot<Self>>)
35    where
36        Seeker: BstSeeker<Spec = Self>;
37}
38
39pub struct BstNode<Data, Parent = WithNoParent<Data>> {
40    pub data: Data,
41    pub parent: Parent,
42    pub child: [Option<NonNull<BstNode<Data, Parent>>>; 2],
43}
44
45impl<Data, Parent> BstNode<Data, Parent>
46where
47    Parent: Default,
48{
49    pub fn new(data: Data) -> Self {
50        Self {
51            data,
52            parent: Parent::default(),
53            child: [None, None],
54        }
55    }
56}
57
58pub trait ParentStrategy: Sized + Default {
59    type Data;
60
61    fn take_parent<Spec>(_node: BstNodeRef<marker::Mut<'_>, Spec>)
62    where
63        Spec: BstSpec<Data = Self::Data, Parent = Self>,
64    {
65    }
66
67    fn set_parent<Spec>(
68        _node: BstNodeRef<marker::Mut<'_>, Spec>,
69        _parent: Option<NonNull<BstNode<Spec::Data, Self>>>,
70    ) where
71        Spec: BstSpec<Data = Self::Data, Parent = Self>,
72    {
73    }
74}
75
76pub struct WithNoParent<Data> {
77    _marker: PhantomData<fn() -> Data>,
78}
79
80impl<Data> Default for WithNoParent<Data> {
81    fn default() -> Self {
82        Self {
83            _marker: PhantomData,
84        }
85    }
86}
87
88impl<Data> ParentStrategy for WithNoParent<Data> {
89    type Data = Data;
90
91    fn take_parent<Spec>(_node: BstNodeRef<marker::Mut<'_>, Spec>)
92    where
93        Spec: BstSpec<Data = Self::Data, Parent = Self>,
94    {
95    }
96
97    fn set_parent<Spec>(
98        _node: BstNodeRef<marker::Mut<'_>, Spec>,
99        _parent: Option<NonNull<BstNode<Spec::Data, Self>>>,
100    ) where
101        Spec: BstSpec<Data = Self::Data, Parent = Self>,
102    {
103    }
104}
105
106pub struct WithParent<Data> {
107    pub parent: Option<NonNull<BstNode<Data, Self>>>,
108}
109
110impl<Data> Default for WithParent<Data> {
111    fn default() -> Self {
112        Self {
113            parent: Default::default(),
114        }
115    }
116}
117
118impl<Data> ParentStrategy for WithParent<Data> {
119    type Data = Data;
120
121    fn take_parent<Spec>(mut node: BstNodeRef<marker::Mut<'_>, Spec>)
122    where
123        Spec: BstSpec<Data = Self::Data, Parent = Self>,
124    {
125        unsafe { node.node.as_mut().parent = Default::default() };
126    }
127
128    fn set_parent<Spec>(
129        mut node: BstNodeRef<marker::Mut<'_>, Spec>,
130        parent: Option<NonNull<BstNode<Spec::Data, Self>>>,
131    ) where
132        Spec: BstSpec<Data = Self::Data, Parent = Self>,
133    {
134        unsafe { node.node.as_mut().parent.parent = parent };
135    }
136}
137
138impl<Data> WithParent<Data> {
139    pub fn resolve_top_down<Spec>(node: BstNodeRef<marker::DataMut<'_>, Spec>)
140    where
141        Spec: BstSpec<Data = Data, Parent = Self>,
142    {
143        unsafe {
144            let (mut node, mut stack) = node.root_path();
145            while let Some(is_left) = stack.pop() {
146                Spec::top_down(node.reborrow_datamut());
147                if is_left {
148                    node = node.left().descend().unwrap_unchecked();
149                } else {
150                    node = node.right().descend().unwrap_unchecked();
151                }
152            }
153            Spec::top_down(node.reborrow_datamut());
154        }
155    }
156
157    pub fn resolve_bottom_up<Spec>(mut node: BstNodeRef<marker::DataMut<'_>, Spec>)
158    where
159        Spec: BstSpec<Data = Data, Parent = Self>,
160    {
161        loop {
162            Spec::bottom_up(node.reborrow_datamut());
163            match node.ascend() {
164                Ok(parent) => node = parent,
165                Err(_) => break,
166            }
167        }
168    }
169
170    pub fn is_root<Spec>(node: BstNodeRef<marker::Immut<'_>, Spec>) -> bool
171    where
172        Spec: BstSpec<Data = Data, Parent = Self>,
173    {
174        unsafe { node.node.as_ref().parent.parent.is_none() }
175    }
176
177    pub unsafe fn remove_root<Spec>(
178        root: &mut Option<BstRoot<Spec>>,
179    ) -> Option<BstNodeRef<marker::Owned, Spec>>
180    where
181        Spec: BstSpec<Data = Data, Parent = Self>,
182    {
183        let mut node = root.take()?;
184        unsafe {
185            let left = node.borrow_mut().left_mut().take();
186            let right = node.borrow_mut().right_mut().take();
187            *root = Spec::merge(left, right);
188            Spec::bottom_up(node.borrow_datamut());
189            Some(node)
190        }
191    }
192
193    pub unsafe fn remove_not_root<Spec>(
194        mut node: BstNodeRef<marker::Mut<'_>, Spec>,
195    ) -> BstNodeRef<marker::Owned, Spec>
196    where
197        Spec: BstSpec<Data = Data, Parent = Self>,
198    {
199        assert!(!Self::is_root(node.reborrow()));
200        unsafe {
201            let left = node.left_mut().take();
202            let right = node.right_mut().take();
203            let merged = Spec::merge(left, right);
204            let node_inner = node.node;
205            let mut parent = node.ascend().unwrap_unchecked();
206            let mut node = if let Some(merged) = merged {
207                let node = if parent
208                    .reborrow()
209                    .left()
210                    .descend()
211                    .is_ok_and(|n| n.node == node_inner)
212                {
213                    parent.left_mut().replace(merged)
214                } else {
215                    parent.right_mut().replace(merged)
216                };
217                Self::resolve_bottom_up(parent.reborrow_datamut());
218                node.unwrap_unchecked()
219            } else {
220                let node = if parent
221                    .reborrow()
222                    .left()
223                    .descend()
224                    .is_ok_and(|n| n.node == node_inner)
225                {
226                    parent.left_mut().take()
227                } else {
228                    parent.right_mut().take()
229                };
230                Self::resolve_bottom_up(parent.reborrow_datamut());
231                node.unwrap_unchecked()
232            };
233            Spec::bottom_up(node.borrow_datamut());
234            node
235        }
236    }
237}
238
239pub struct BstNodeRef<BorrowType, Spec>
240where
241    Spec: BstSpec,
242{
243    pub node: NonNull<BstNode<Spec::Data, Spec::Parent>>,
244    _marker: PhantomData<BorrowType>,
245}
246
247impl<'a, Spec> Copy for BstNodeRef<marker::Immut<'a>, Spec> where Spec: BstSpec<Data: 'a> {}
248impl<'a, Spec> Clone for BstNodeRef<marker::Immut<'a>, Spec>
249where
250    Spec: BstSpec<Data: 'a>,
251{
252    fn clone(&self) -> Self {
253        *self
254    }
255}
256
257impl<BorrowType, Spec> BstNodeRef<BorrowType, Spec>
258where
259    Spec: BstSpec,
260    BorrowType: marker::BorrowType,
261{
262    pub unsafe fn new_unchecked(node: NonNull<BstNode<Spec::Data, Spec::Parent>>) -> Self {
263        Self {
264            node,
265            _marker: PhantomData,
266        }
267    }
268    pub fn reborrow(&self) -> BstNodeRef<marker::Immut<'_>, Spec> {
269        BstNodeRef {
270            node: self.node,
271            _marker: PhantomData,
272        }
273    }
274    pub fn left(self) -> BstEdgeHandle<Self, marker::Left> {
275        BstEdgeHandle {
276            node: self,
277            _marker: PhantomData,
278        }
279    }
280    pub fn right(self) -> BstEdgeHandle<Self, marker::Right> {
281        BstEdgeHandle {
282            node: self,
283            _marker: PhantomData,
284        }
285    }
286}
287
288impl<BorrowType, Spec, Data> BstNodeRef<BorrowType, Spec>
289where
290    Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
291    BorrowType: marker::BorrowType,
292{
293    pub fn ascend(self) -> Result<BstNodeRef<BorrowType, Spec>, Self> {
294        const {
295            assert!(BorrowType::TRAVERSAL_PERMIT);
296        };
297        let parent = unsafe { self.node.as_ref().parent.parent };
298        parent
299            .map(|node| BstNodeRef {
300                node,
301                _marker: PhantomData,
302            })
303            .ok_or(self)
304    }
305    pub fn root_path(self) -> (Self, Vec<bool>) {
306        let mut node = self;
307        let mut nn = node.node;
308        let mut stack = vec![];
309        let root = loop {
310            match node.ascend() {
311                Ok(parent) => {
312                    node = parent;
313                    stack.push(
314                        node.reborrow()
315                            .left()
316                            .descend()
317                            .is_ok_and(|node| node.node == nn),
318                    );
319                    nn = node.node;
320                }
321                Err(node) => {
322                    break node;
323                }
324            }
325        };
326        (root, stack)
327    }
328}
329
330impl<Spec> BstNodeRef<marker::Owned, Spec>
331where
332    Spec: BstSpec,
333{
334    pub fn new(node: NonNull<BstNode<Spec::Data, Spec::Parent>>) -> Self {
335        Self {
336            node,
337            _marker: PhantomData,
338        }
339    }
340    pub fn from_data<A>(data: Spec::Data, allocator: &mut A) -> Self
341    where
342        A: Allocator<BstNode<Spec::Data, Spec::Parent>>,
343    {
344        Self::new(allocator.allocate(BstNode::new(data)))
345    }
346    pub fn borrow_mut(&mut self) -> BstNodeRef<marker::Mut<'_>, Spec> {
347        BstNodeRef {
348            node: self.node,
349            _marker: PhantomData,
350        }
351    }
352    pub fn borrow_datamut(&mut self) -> BstNodeRef<marker::DataMut<'_>, Spec> {
353        BstNodeRef {
354            node: self.node,
355            _marker: PhantomData,
356        }
357    }
358    pub fn into_dying(self) -> BstNodeRef<marker::Dying, Spec> {
359        BstNodeRef {
360            node: self.node,
361            _marker: PhantomData,
362        }
363    }
364}
365
366impl<'a, Spec> BstNodeRef<marker::Immut<'a>, Spec>
367where
368    Spec: BstSpec<Parent: 'a, Data: 'a>,
369{
370    pub fn into_data(self) -> &'a Spec::Data {
371        unsafe { &self.node.as_ref().data }
372    }
373
374    pub fn traverse<F>(self, f: &mut F)
375    where
376        F: FnMut(Self),
377    {
378        if let Ok(left) = self.left().descend() {
379            left.traverse(f);
380        }
381        f(self);
382        if let Ok(right) = self.right().descend() {
383            right.traverse(f);
384        }
385    }
386
387    pub fn leftmost(self) -> Self {
388        let mut node = self;
389        while let Ok(left) = node.left().descend() {
390            node = left;
391        }
392        node
393    }
394
395    pub fn rightmost(self) -> Self {
396        let mut node = self;
397        while let Ok(right) = node.right().descend() {
398            node = right;
399        }
400        node
401    }
402}
403
404impl<'a, Spec> BstNodeRef<marker::DataMut<'a>, Spec>
405where
406    Spec: BstSpec,
407{
408    pub fn reborrow_datamut(&mut self) -> BstNodeRef<marker::DataMut<'_>, Spec> {
409        BstNodeRef {
410            node: self.node,
411            _marker: PhantomData,
412        }
413    }
414    pub fn data_mut(&mut self) -> &mut Spec::Data {
415        unsafe { &mut self.node.as_mut().data }
416    }
417
418    pub fn swap_children(&mut self) {
419        unsafe { self.node.as_mut().child.swap(0, 1) };
420    }
421}
422
423impl<'a, Spec> BstNodeRef<marker::DataMut<'a>, Spec>
424where
425    Spec: BstSpec<Parent: 'a, Data: 'a>,
426{
427    pub fn into_data_mut(mut self) -> &'a mut Spec::Data {
428        unsafe { &mut self.node.as_mut().data }
429    }
430}
431
432impl<'a, Spec> BstNodeRef<marker::Mut<'a>, Spec>
433where
434    Spec: BstSpec,
435{
436    pub fn reborrow_datamut(&mut self) -> BstNodeRef<marker::DataMut<'_>, Spec> {
437        BstNodeRef {
438            node: self.node,
439            _marker: PhantomData,
440        }
441    }
442
443    pub fn left_mut(&mut self) -> BstEdgeHandle<BstNodeRef<marker::Mut<'_>, Spec>, marker::Left> {
444        BstEdgeHandle {
445            node: BstNodeRef {
446                node: self.node,
447                _marker: PhantomData,
448            },
449            _marker: PhantomData,
450        }
451    }
452
453    pub fn right_mut(&mut self) -> BstEdgeHandle<BstNodeRef<marker::Mut<'_>, Spec>, marker::Right> {
454        BstEdgeHandle {
455            node: BstNodeRef {
456                node: self.node,
457                _marker: PhantomData,
458            },
459            _marker: PhantomData,
460        }
461    }
462}
463
464impl<'a, Spec> BstNodeRef<marker::Mut<'a>, Spec>
465where
466    Spec: BstSpec<Data: 'a>,
467{
468    pub fn dormant(self) -> BstNodeRef<marker::DormantMut, Spec> {
469        BstNodeRef {
470            node: self.node,
471            _marker: PhantomData,
472        }
473    }
474}
475
476impl<Spec> BstNodeRef<marker::DormantMut, Spec>
477where
478    Spec: BstSpec,
479{
480    pub unsafe fn awaken<'a>(self) -> BstNodeRef<marker::Mut<'a>, Spec> {
481        BstNodeRef {
482            node: self.node,
483            _marker: PhantomData,
484        }
485    }
486}
487
488impl<Spec> BstNodeRef<marker::Dying, Spec>
489where
490    Spec: BstSpec,
491{
492    pub unsafe fn into_data<A>(self, allocator: &mut A) -> Spec::Data
493    where
494        A: Allocator<BstNode<Spec::Data, Spec::Parent>>,
495    {
496        debug_assert!(self.reborrow().left().descend().is_err());
497        debug_assert!(self.reborrow().right().descend().is_err());
498        allocator.deallocate(self.node).data
499    }
500
501    pub unsafe fn drop_all<A>(self, allocator: &mut A)
502    where
503        A: Allocator<BstNode<Spec::Data, Spec::Parent>>,
504    {
505        let mut stack = vec![self.node];
506        while let Some(node) = stack.pop() {
507            stack.extend(allocator.deallocate(node).child.into_iter().flatten());
508        }
509    }
510}
511
512pub struct BstEdgeHandle<Node, Dir> {
513    node: Node,
514    _marker: PhantomData<Dir>,
515}
516
517impl<BorrowType, Spec, Dir> BstEdgeHandle<BstNodeRef<BorrowType, Spec>, Dir>
518where
519    Spec: BstSpec,
520    BorrowType: marker::BorrowType,
521    Dir: marker::BstDirection,
522{
523    pub fn descend(self) -> Result<BstNodeRef<BorrowType, Spec>, Self> {
524        const {
525            assert!(BorrowType::TRAVERSAL_PERMIT);
526        };
527        let child = unsafe { self.node.node.as_ref().child.get_unchecked(Dir::IDX) };
528        child
529            .map(|node| BstNodeRef {
530                node,
531                _marker: PhantomData,
532            })
533            .ok_or(self)
534    }
535}
536
537impl<'a, Spec, Dir> BstEdgeHandle<BstNodeRef<marker::Mut<'a>, Spec>, Dir>
538where
539    Spec: BstSpec,
540    Dir: marker::BstDirection,
541{
542    pub unsafe fn take(&mut self) -> Option<BstNodeRef<marker::Owned, Spec>> {
543        let child = unsafe { self.node.node.as_mut().child.get_unchecked_mut(Dir::IDX) };
544        child.take().map(|node| {
545            let mut node = BstNodeRef {
546                node,
547                _marker: PhantomData,
548            };
549            Spec::Parent::take_parent(node.borrow_mut());
550            node
551        })
552    }
553    pub unsafe fn replace(
554        &mut self,
555        mut other: BstNodeRef<marker::Owned, Spec>,
556    ) -> Option<BstNodeRef<marker::Owned, Spec>> {
557        let child = unsafe { self.node.node.as_mut().child.get_unchecked_mut(Dir::IDX) };
558        Spec::Parent::set_parent(other.borrow_mut(), Some(self.node.node));
559        child.replace(other.node).map(|node| {
560            let mut node = BstNodeRef {
561                node,
562                _marker: PhantomData,
563            };
564            Spec::Parent::take_parent(node.borrow_mut());
565            node
566        })
567    }
568    pub unsafe fn set(&mut self, mut other: BstNodeRef<marker::Owned, Spec>) {
569        let child = unsafe { self.node.node.as_mut().child.get_unchecked_mut(Dir::IDX) };
570        Spec::Parent::set_parent(other.borrow_mut(), Some(self.node.node));
571        *child = Some(other.node);
572    }
573}
574
575pub type BstRoot<Spec> = BstNodeRef<marker::Owned, Spec>;
576pub type BstDataMutRef<'a, Spec> = BstNodeRef<marker::DataMut<'a>, Spec>;
577pub type BstImmutRef<'a, Spec> = BstNodeRef<marker::Immut<'a>, Spec>;
578pub type BstNodePtr<Data, Parent> = NonNull<BstNode<Data, Parent>>;
579
580pub mod marker {
581    use std::marker::PhantomData;
582
583    pub enum Left {}
584    pub enum Right {}
585    pub trait BstDirection {
586        const IDX: usize;
587    }
588    impl BstDirection for Left {
589        const IDX: usize = 0;
590    }
591    impl BstDirection for Right {
592        const IDX: usize = 1;
593    }
594
595    pub enum Owned {}
596    pub enum Dying {}
597    pub enum DormantMut {}
598    pub struct Immut<'a>(PhantomData<&'a ()>);
599    pub struct Mut<'a>(PhantomData<&'a mut ()>);
600    pub struct DataMut<'a>(PhantomData<&'a mut ()>);
601
602    pub trait BorrowType {
603        const TRAVERSAL_PERMIT: bool = true;
604    }
605    impl BorrowType for Owned {
606        const TRAVERSAL_PERMIT: bool = false;
607    }
608    impl BorrowType for Dying {}
609    impl BorrowType for DormantMut {}
610    impl<'a> BorrowType for Immut<'a> {}
611    impl<'a> BorrowType for Mut<'a> {}
612    impl<'a> BorrowType for DataMut<'a> {}
613}
competitive/data_structure/binary_search_tree/node.rs

competitive/data_structure/binary_search_tree/
node.rs
