competitive/data_structure/splay_tree/
node.rs1use super::Allocator;
2use std::{
3 cmp::Ordering,
4 fmt::{self, Debug},
5 marker::PhantomData,
6 mem::{replace, swap},
7 ops::Bound,
8 ptr::NonNull,
9};
10
11pub trait SplaySpec: Sized {
12 type T;
13 fn has_bottom_up() -> bool {
14 false
15 }
16 fn top_down(_node: NodeRef<marker::DataMut<'_>, Self>) {}
17 fn bottom_up(_node: NodeRef<marker::DataMut<'_>, Self>) {}
18}
19
20pub trait SplaySeeker {
21 type S: SplaySpec;
22 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering;
23}
24
25pub struct SeekLeft<S> {
26 _marker: PhantomData<fn() -> S>,
27}
28impl<S> Default for SeekLeft<S> {
29 fn default() -> Self {
30 Self {
31 _marker: PhantomData,
32 }
33 }
34}
35impl<S> SeekLeft<S> {
36 pub fn new() -> Self {
37 Self::default()
38 }
39}
40impl<S> SplaySeeker for SeekLeft<S>
41where
42 S: SplaySpec,
43{
44 type S = S;
45 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
46 Ordering::Less
47 }
48}
49
50pub struct SeekRight<S> {
51 _marker: PhantomData<fn() -> S>,
52}
53impl<S> Default for SeekRight<S> {
54 fn default() -> Self {
55 Self {
56 _marker: PhantomData,
57 }
58 }
59}
60impl<S> SeekRight<S> {
61 pub fn new() -> Self {
62 Self::default()
63 }
64}
65impl<S> SplaySeeker for SeekRight<S>
66where
67 S: SplaySpec,
68{
69 type S = S;
70 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
71 Ordering::Greater
72 }
73}
74
75pub struct Node<T> {
76 data: T,
77 left: Option<NonNull<Node<T>>>,
78 right: Option<NonNull<Node<T>>>,
79}
80
81impl<T> Node<T> {
82 pub fn new(data: T) -> Self {
83 Self {
84 data,
85 left: None,
86 right: None,
87 }
88 }
89}
90
91pub struct NodeRef<B, S>
92where
93 S: SplaySpec,
94{
95 node: NonNull<Node<S::T>>,
96 _marker: PhantomData<B>,
97}
98
99pub struct Root<S>
100where
101 S: SplaySpec,
102{
103 root: Option<NodeRef<marker::Owned, S>>,
104}
105
106impl<S> Default for Root<S>
107where
108 S: SplaySpec,
109{
110 fn default() -> Self {
111 Self { root: None }
112 }
113}
114
115impl<'a, S> Copy for NodeRef<marker::Immut<'a>, S> where S: SplaySpec<T: 'a> {}
116impl<'a, S> Clone for NodeRef<marker::Immut<'a>, S>
117where
118 S: SplaySpec<T: 'a>,
119{
120 fn clone(&self) -> Self {
121 *self
122 }
123}
124
125impl<B, S> NodeRef<B, S>
126where
127 S: SplaySpec,
128{
129 unsafe fn new_unchecked(node: NonNull<Node<S::T>>) -> Self {
130 Self {
131 node,
132 _marker: PhantomData,
133 }
134 }
135 fn reborrow(&self) -> NodeRef<marker::Immut<'_>, S> {
136 unsafe { NodeRef::new_unchecked(self.node) }
137 }
138 fn as_ptr(&self) -> *mut Node<S::T> {
139 self.node.as_ptr()
140 }
141}
142
143impl<S> NodeRef<marker::Owned, S>
144where
145 S: SplaySpec,
146{
147 pub fn new(node: NonNull<Node<S::T>>) -> Self {
148 unsafe { NodeRef::new_unchecked(node) }
149 }
150 pub unsafe fn from_data<A>(data: S::T, allocator: &mut A) -> Self
151 where
152 A: Allocator<Node<S::T>>,
153 {
154 Self::new(allocator.allocate(Node::new(data)))
155 }
156 pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, S> {
157 unsafe { NodeRef::new_unchecked(self.node) }
158 }
159 pub fn borrow_datamut(&mut self) -> NodeRef<marker::DataMut<'_>, S> {
160 unsafe { NodeRef::new_unchecked(self.node) }
161 }
162 pub fn into_dying(self) -> NodeRef<marker::Dying, S> {
163 unsafe { NodeRef::new_unchecked(self.node) }
164 }
165}
166
167impl<'a, S> NodeRef<marker::Immut<'a>, S>
168where
169 S: SplaySpec<T: 'a>,
170{
171 pub fn data(&self) -> &'a S::T {
172 unsafe { &(*self.as_ptr()).data }
173 }
174 pub fn left(&self) -> Option<Self> {
175 unsafe {
176 (*self.as_ptr())
177 .left
178 .map(|node| NodeRef::new_unchecked(node))
179 }
180 }
181 pub fn right(&self) -> Option<Self> {
182 unsafe {
183 (*self.as_ptr())
184 .right
185 .map(|node| NodeRef::new_unchecked(node))
186 }
187 }
188 pub fn traverse<F>(self, f: &mut F)
189 where
190 S: SplaySpec<T: 'a>,
191 F: FnMut(Self),
192 {
193 if let Some(left) = self.clone().left() {
194 left.traverse(f);
195 }
196 f(self);
197 if let Some(right) = self.clone().right() {
198 right.traverse(f);
199 }
200 }
201}
202
203impl<'a, S> NodeRef<marker::DataMut<'a>, S>
204where
205 S: SplaySpec<T: 'a>,
206{
207 pub fn data(&self) -> &'a S::T {
208 unsafe { &(*self.as_ptr()).data }
209 }
210 pub fn data_mut(&self) -> &'a mut S::T {
211 unsafe { &mut (*self.as_ptr()).data }
212 }
213 pub fn left(&self) -> Option<Self> {
214 unsafe {
215 (*self.as_ptr())
216 .left
217 .map(|node| NodeRef::new_unchecked(node))
218 }
219 }
220 pub fn right(&self) -> Option<Self> {
221 unsafe {
222 (*self.as_ptr())
223 .right
224 .map(|node| NodeRef::new_unchecked(node))
225 }
226 }
227 pub fn reverse(&self) {
228 unsafe {
229 let node = &mut (*self.as_ptr());
230 swap(&mut node.left, &mut node.right);
231 }
232 }
233}
234
235impl<'a, S> NodeRef<marker::Mut<'a>, S>
236where
237 S: SplaySpec<T: 'a>,
238{
239 pub fn data(self) -> &'a S::T {
240 unsafe { &(*self.as_ptr()).data }
241 }
242 pub fn data_mut(self) -> &'a mut S::T {
243 unsafe { &mut (*self.as_ptr()).data }
244 }
245 pub fn take_left(&mut self) -> Option<NodeRef<marker::Owned, S>> {
246 Some(NodeRef::new(unsafe { (*self.as_ptr()).left.take()? }))
247 }
248 pub fn take_right(&mut self) -> Option<NodeRef<marker::Owned, S>> {
249 Some(NodeRef::new(unsafe { (*self.as_ptr()).right.take()? }))
250 }
251 pub fn set_left(&mut self, node: Option<NodeRef<marker::Owned, S>>) {
252 unsafe { (*self.as_ptr()).left = node.map(|node| node.node) }
253 }
254 pub fn set_right(&mut self, node: Option<NodeRef<marker::Owned, S>>) {
255 unsafe { (*self.as_ptr()).right = node.map(|node| node.node) }
256 }
257}
258
259impl<S> NodeRef<marker::Dying, S>
260where
261 S: SplaySpec,
262{
263 pub unsafe fn into_inner(self) -> NonNull<Node<S::T>> {
264 let node = self.node;
265 unsafe {
266 debug_assert!((*node.as_ptr()).left.is_none());
267 debug_assert!((*node.as_ptr()).right.is_none());
268 }
269 node
270 }
271 pub unsafe fn into_data<A>(self, allocator: &mut A) -> S::T
272 where
273 A: Allocator<Node<S::T>>,
274 {
275 let Node { data, left, right } = allocator.deallocate(self.node);
276 debug_assert!(left.is_none());
277 debug_assert!(right.is_none());
278 data
279 }
280}
281
282impl<S> NodeRef<marker::Owned, S>
283where
284 S: SplaySpec,
285{
286 pub fn splay_by<Seeker>(self, mut seeker: Seeker) -> (Ordering, Self)
288 where
289 Seeker: SplaySeeker<S = S>,
290 {
291 let mut x = self;
292 let mut left_subtree: Option<NonNull<Node<S::T>>> = None;
294 let mut right_subtree: Option<NonNull<Node<S::T>>> = None;
295 let mut left_entry = &mut left_subtree;
296 let mut right_entry = &mut right_subtree;
297 let mut stack = vec![];
298
299 macro_rules! add {
300 (@left Some($ptr:ident)) => { add!(@inner Some($ptr.node), left_entry $ptr right); };
301 (@right Some($ptr:ident)) => { add!(@inner Some($ptr.node), right_entry $ptr left); };
302 (@inner $node:expr, $entry:ident $ptr:ident $dir:ident) => {
303 *$entry = $node;
304 if S::has_bottom_up() {
305 stack.push($ptr.node);
306 }
307 $entry = unsafe { &mut (*$entry.as_mut().unwrap().as_ptr()).$dir };
308 };
309 }
310
311 let root_ord = loop {
312 S::top_down(x.borrow_datamut());
313 match seeker.splay_seek(x.reborrow()) {
314 Ordering::Less => {
315 if let Some(mut y) = x.borrow_mut().take_left() {
316 S::top_down(y.borrow_datamut());
317 match seeker.splay_seek(y.reborrow()) {
318 Ordering::Less => {
319 if let Some(mut z) = y.borrow_mut().take_left() {
320 S::top_down(z.borrow_datamut());
321 x.borrow_mut().set_left(y.borrow_mut().take_right());
322 S::bottom_up(x.borrow_datamut());
323 y.borrow_mut().set_right(Some(x));
324 add!(@right Some(y));
325 x = z;
326 } else {
327 add!(@right Some(x));
328 x = y;
329 break Ordering::Less;
330 }
331 }
332 Ordering::Equal => {
333 add!(@right Some(x));
334 x = y;
335 break Ordering::Equal;
336 }
337 Ordering::Greater => {
338 if let Some(mut z) = y.borrow_mut().take_right() {
339 S::top_down(z.borrow_datamut());
340 add!(@right Some(x));
341 add!(@left Some(y));
342 x = z;
343 } else {
344 add!(@right Some(x));
345 x = y;
346 break Ordering::Greater;
347 }
348 }
349 }
350 } else {
351 break Ordering::Less;
352 }
353 }
354 Ordering::Equal => break Ordering::Equal,
355 Ordering::Greater => {
356 if let Some(mut y) = x.borrow_mut().take_right() {
357 S::top_down(y.borrow_datamut());
358 match seeker.splay_seek(y.reborrow()) {
359 Ordering::Less => {
360 if let Some(mut z) = y.borrow_mut().take_left() {
361 S::top_down(z.borrow_datamut());
362 add!(@left Some(x));
363 add!(@right Some(y));
364 x = z;
365 } else {
366 add!(@left Some(x));
367 x = y;
368 break Ordering::Less;
369 }
370 }
371 Ordering::Equal => {
372 add!(@left Some(x));
373 x = y;
374 break Ordering::Equal;
375 }
376 Ordering::Greater => {
377 if let Some(mut z) = y.borrow_mut().take_right() {
378 S::top_down(z.borrow_datamut());
379 x.borrow_mut().set_right(y.borrow_mut().take_left());
380 S::bottom_up(x.borrow_datamut());
381 y.borrow_mut().set_left(Some(x));
382 add!(@left Some(y));
383 x = z;
384 } else {
385 add!(@left Some(x));
386 x = y;
387 break Ordering::Greater;
388 }
389 }
390 }
391 } else {
392 break Ordering::Greater;
393 }
394 }
395 }
396 };
397 *left_entry = x.borrow_mut().take_left().map(|node| node.node);
398 *right_entry = x.borrow_mut().take_right().map(|node| node.node);
399 unsafe {
400 x.borrow_mut()
401 .set_left(left_subtree.map(|node| NodeRef::new_unchecked(node)));
402 x.borrow_mut()
403 .set_right(right_subtree.map(|node| NodeRef::new_unchecked(node)));
404 if S::has_bottom_up() {
405 while let Some(node) = stack.pop() {
406 S::bottom_up(NodeRef::new_unchecked(node));
407 }
408 }
409 }
410 S::bottom_up(x.borrow_datamut());
411 (root_ord, x)
412 }
413 pub fn insert_left(mut self, mut node: Self) -> Self {
414 if let Some(left) = self.borrow_mut().take_left() {
415 node.borrow_mut().set_left(Some(left));
416 S::bottom_up(self.borrow_datamut());
417 };
418 node.borrow_mut().set_right(Some(self));
419 S::bottom_up(node.borrow_datamut());
420 node
421 }
422 pub fn insert_right(mut self, mut node: Self) -> Self {
423 if let Some(right) = self.borrow_mut().take_right() {
424 node.borrow_mut().set_right(Some(right));
425 S::bottom_up(self.borrow_datamut());
426 }
427 node.borrow_mut().set_left(Some(self));
428 S::bottom_up(node.borrow_datamut());
429 node
430 }
431 pub fn insert_first(self, mut node: Self) -> Self {
432 node.borrow_mut().set_right(Some(self));
433 S::bottom_up(node.borrow_datamut());
434 node
435 }
436 pub fn insert_last(self, mut node: Self) -> Self {
437 node.borrow_mut().set_left(Some(self));
438 S::bottom_up(node.borrow_datamut());
439 node
440 }
441 pub fn merge(mut self, mut other: Self) -> Self {
442 if other.reborrow().left().is_none() {
443 S::top_down(other.borrow_datamut());
444 other.borrow_mut().set_left(Some(self));
445 S::bottom_up(other.borrow_datamut());
446 other
447 } else {
448 self = self.splay_by(SeekRight::new()).1;
449 self.borrow_mut().set_right(Some(other));
450 S::bottom_up(self.borrow_datamut());
451 self
452 }
453 }
454}
455
456impl<S> Root<S>
457where
458 S: SplaySpec,
459{
460 pub fn new(root: Option<NodeRef<marker::Owned, S>>) -> Self {
461 Self { root }
462 }
463 pub unsafe fn from_single_nodes(nodes: Vec<NodeRef<marker::Owned, S>>) -> Self {
464 unsafe { Self::from_single_nodes_inner(&nodes) }
465 }
466 unsafe fn from_single_nodes_inner(nodes: &[NodeRef<marker::Owned, S>]) -> Self {
467 if nodes.is_empty() {
468 Self::new(None)
469 } else {
470 let m = nodes.len() / 2;
471 let left = unsafe { Self::from_single_nodes_inner(&nodes[..m]) };
472 let right = unsafe { Self::from_single_nodes_inner(&nodes[m + 1..]) };
473 let mut node = NodeRef::new(nodes[m].node);
474 node.borrow_mut().set_left(left.root);
475 node.borrow_mut().set_right(right.root);
476 S::bottom_up(node.borrow_datamut());
477 Self::new(Some(node))
478 }
479 }
480 pub fn is_empty(&self) -> bool {
481 self.root.is_none()
482 }
483 pub fn root(&self) -> Option<NodeRef<marker::Immut<'_>, S>> {
484 Some(self.root.as_ref()?.reborrow())
485 }
486 pub fn root_data_mut(&mut self) -> Option<NodeRef<marker::DataMut<'_>, S>> {
487 Some(self.root.as_mut()?.borrow_datamut())
488 }
489 pub fn splay_by<Seeker>(&mut self, seeker: Seeker) -> Option<Ordering>
490 where
491 Seeker: SplaySeeker<S = S>,
492 {
493 let (ord, root) = self.root.take()?.splay_by(seeker);
494 self.root = Some(root);
495 Some(ord)
496 }
497 pub fn insert_left(&mut self, mut node: NodeRef<marker::Owned, S>) {
498 self.root = Some(match self.root.take() {
499 Some(root) => root.insert_left(node),
500 None => {
501 S::bottom_up(node.borrow_datamut());
502 node
503 }
504 });
505 }
506 pub fn insert_right(&mut self, mut node: NodeRef<marker::Owned, S>) {
507 self.root = Some(match self.root.take() {
508 Some(root) => root.insert_right(node),
509 None => {
510 S::bottom_up(node.borrow_datamut());
511 node
512 }
513 });
514 }
515 pub fn insert_first(&mut self, mut node: NodeRef<marker::Owned, S>) {
516 self.root = Some(match self.root.take() {
517 Some(root) => root.insert_first(node),
518 None => {
519 S::bottom_up(node.borrow_datamut());
520 node
521 }
522 });
523 }
524 pub fn insert_last(&mut self, mut node: NodeRef<marker::Owned, S>) {
525 self.root = Some(match self.root.take() {
526 Some(root) => root.insert_last(node),
527 None => {
528 S::bottom_up(node.borrow_datamut());
529 node
530 }
531 });
532 }
533 pub fn take_root(&mut self) -> Option<NodeRef<marker::Owned, S>> {
534 let mut root = self.root.take()?;
535 let right = root.borrow_mut().take_right();
536 self.root = root.borrow_mut().take_left();
537 self.append(&mut Self::new(right));
538 Some(root)
539 }
540 pub fn take_first(&mut self) -> Option<NodeRef<marker::Owned, S>> {
541 let mut root = self.root.take()?.splay_by(SeekLeft::new()).1;
542 let right = root.borrow_mut().take_right();
543 self.root = right;
544 Some(root)
545 }
546 pub fn take_last(&mut self) -> Option<NodeRef<marker::Owned, S>> {
547 let mut root = self.root.take()?.splay_by(SeekRight::new()).1;
548 let left = root.borrow_mut().take_left();
549 self.root = left;
550 Some(root)
551 }
552 pub fn split_left(&mut self) -> Option<NodeRef<marker::Owned, S>> {
553 let root = self.root.as_mut()?;
554 let left = root.borrow_mut().take_left();
555 S::bottom_up(root.borrow_datamut());
556 left
557 }
558 pub fn split_right(&mut self) -> Option<NodeRef<marker::Owned, S>> {
559 let root = self.root.as_mut()?;
560 let right = root.borrow_mut().take_right();
561 S::bottom_up(root.borrow_datamut());
562 right
563 }
564 pub fn split_left_eq(&mut self) -> Option<NodeRef<marker::Owned, S>> {
565 let right = self.split_right();
566 replace(&mut self.root, right)
567 }
568 pub fn split_right_eq(&mut self) -> Option<NodeRef<marker::Owned, S>> {
569 let left = self.split_left();
570 replace(&mut self.root, left)
571 }
572 pub fn append(&mut self, other: &mut Self) {
573 self.root = match (self.root.take(), other.root.take()) {
574 (Some(node), None) | (None, Some(node)) => Some(node),
575 (Some(left), Some(right)) => Some(left.merge(right)),
576 (None, None) => None,
577 }
578 }
579 pub fn range<Seeker>(&mut self, start: Bound<Seeker>, end: Bound<Seeker>) -> NodeRange<'_, S>
580 where
581 Seeker: SplaySeeker<S = S>,
582 {
583 if self.is_empty() {
584 return NodeRange::new(self);
585 }
586 let right = match end {
587 Bound::Included(seeker) => match self.splay_by(seeker).unwrap() {
588 Ordering::Greater | Ordering::Equal => self.split_right(),
589 Ordering::Less => self.split_right_eq(),
590 },
591 Bound::Excluded(seeker) => match self.splay_by(seeker).unwrap() {
592 Ordering::Greater => self.split_right(),
593 Ordering::Less | Ordering::Equal => self.split_right_eq(),
594 },
595 Bound::Unbounded => None,
596 };
597 if self.is_empty() {
598 return NodeRange::three_way(None, self, right);
599 }
600 let left = match start {
601 Bound::Included(seeker) => match self.splay_by(seeker).unwrap() {
602 Ordering::Less | Ordering::Equal => self.split_left(),
603 Ordering::Greater => self.split_left_eq(),
604 },
605 Bound::Excluded(seeker) => match self.splay_by(seeker).unwrap() {
606 Ordering::Less => self.split_left(),
607 Ordering::Greater | Ordering::Equal => self.split_left_eq(),
608 },
609 Bound::Unbounded => None,
610 };
611 NodeRange::three_way(left, self, right)
612 }
613}
614
615impl<S> Debug for Root<S>
616where
617 S: SplaySpec<T: Debug>,
618{
619 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
620 fn fmt_recurse<'a, S>(
621 node: NodeRef<marker::Immut<'a>, S>,
622 f: &mut fmt::Formatter<'_>,
623 ) -> fmt::Result
624 where
625 S: SplaySpec<T: 'a + Debug>,
626 {
627 write!(f, "[")?;
628 if let Some(left) = node.left() {
629 fmt_recurse(left, f)?;
630 }
631 node.data().fmt(f)?;
632 if let Some(right) = node.right() {
633 fmt_recurse(right, f)?;
634 }
635 write!(f, "]")?;
636 Ok(())
637 }
638 if let Some(root) = self.root.as_ref() {
639 let root = root.reborrow();
640 fmt_recurse(root, f)?;
641 }
642 Ok(())
643 }
644}
645
646pub struct NodeRange<'a, S>
647where
648 S: SplaySpec<T: 'a>,
649{
650 front: Root<S>,
651 back: Root<S>,
652 root: &'a mut Root<S>,
653}
654
655impl<'a, S> Debug for NodeRange<'a, S>
656where
657 S: SplaySpec<T: 'a + Debug>,
658{
659 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
660 f.debug_struct("NodeRange")
661 .field("front", &self.front)
662 .field("back", &self.back)
663 .field("root", &self.root)
664 .finish()
665 }
666}
667impl<'a, S> NodeRange<'a, S>
668where
669 S: SplaySpec<T: 'a>,
670{
671 pub fn new(root: &'a mut Root<S>) -> Self {
672 Self {
673 front: Default::default(),
674 back: Default::default(),
675 root,
676 }
677 }
678 pub fn three_way(
679 front: Option<NodeRef<marker::Owned, S>>,
680 root: &'a mut Root<S>,
681 back: Option<NodeRef<marker::Owned, S>>,
682 ) -> Self {
683 Self {
684 front: Root::new(front),
685 back: Root::new(back),
686 root,
687 }
688 }
689 pub fn next_checked(&mut self) -> Option<NodeRef<marker::DataMut<'a>, S>> {
690 let first = self.root.take_first()?;
691 let noderef = unsafe { NodeRef::new_unchecked(first.node) };
692 self.front.insert_last(first);
693 Some(noderef)
694 }
695 pub fn next_back_checked(&mut self) -> Option<NodeRef<marker::DataMut<'a>, S>> {
696 let last = self.root.take_last()?;
697 let noderef = unsafe { NodeRef::new_unchecked(last.node) };
698 self.back.insert_first(last);
699 Some(noderef)
700 }
701 pub fn root(&self) -> &Root<S> {
702 self.root
703 }
704 pub fn root_mut(&mut self) -> &mut Root<S> {
705 self.root
706 }
707 pub fn front(&self) -> &Root<S> {
708 &self.front
709 }
710 pub fn drop_rotate_left(mut self) {
711 self.root.append(&mut self.back);
712 self.root.append(&mut self.front);
713 }
714}
715impl<'a, S> Drop for NodeRange<'a, S>
716where
717 S: SplaySpec<T: 'a>,
718{
719 fn drop(&mut self) {
720 swap(self.root, &mut self.front);
721 self.root.append(&mut self.front);
722 self.root.append(&mut self.back);
723 }
724}
725
726pub mod marker {
727 use std::marker::PhantomData;
728
729 pub enum Owned {}
730 pub enum Dying {}
731 pub struct Immut<'a>(PhantomData<&'a ()>);
732 pub struct Mut<'a>(PhantomData<&'a mut ()>);
733 pub struct DataMut<'a>(PhantomData<&'a mut ()>);
734}