competitive/data_structure/binary_search_tree/
node.rs

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

competitive/data_structure/binary_search_tree/
node.rs
