1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080
#![allow(unsafe_code)] //! The core implementation of the concurrent trie data structure. //! //! This module contains the [`Raw`][crate::raw::Raw] type, which is the engine of all the data //! structures in this crate. This is exposed to allow wrapping it into further APIs, but is //! probably not the best thing for general use. // # The data structure // // The data structure is inspired by the [article] and [Wikipedia entry], however severely // simplified. Compared to the article, what we don't do (if you don't want to read the article, // that's fine, explanation is below): // // * We don't have variable-sized inner nodes. This wastes some more space, but also allows us to // keep the same node around instead of creating a new one every time we want to add or remove a // pointer. // * We don't do snapshots for iterations. // * We got rid of the I-nodes. This gets rid of half of the pointer loads on the way to the // element, so in theory it should make the data structure about twice faster. // // By this simplification, we lose the ability of having consistent iterations and we use somewhat // more memory, but get faster data structure. More importantly, the data structure is simpler to // implement and simpler to prove correct which makes it more likely to actually trust it with some // data. // // ## How it works // // The heart of the data structure is a trie where keys are prefixes of the 64bit hash of the key. // Each inner node has 16 pointer slots, indexed by the next 4 bits of the hash. When we reach a // level where the prefix is unique, we stop (we don't have all 16 levels of inner nodes if we // don't have to) and place a data node. // // A data nodes contain one or more elements (more in case we get a hash collision on the whole // hash ‒ in that case, we store all the colliding elements in an array and distinguish by equality // of the keys in a linear search through the array). // // On lookup, we either find the correct element or stop at the first null pointer encountered. // // On insertion, if we find a null pointer, we atomically replace that pointer to a new data node // containing (only) the new element, using the CaS operation. In case we reach a collision or // replace an existing element, we create a new data node and replace the pointer, again using the // CaS operation. If we find a non-matching data node in our way, we need to insert another level ‒ // we create a brand new inner node, link the old node there and again, replace the pointer (then // retry with our insertion on the next level). // // Deletion looks up the element and either replaces the pointer to the data node with null (if it // was the last one), or creates a new data node without the element. // // ## Pruning // // If implemented as above, the data structure would work. However, deletions would leave unneeded // dead branches or branches that don't branch (eg. linear ones) behind. That would make the data // structure perform worse than necessary, because the lookups would have to traverse the dead or // non-branching branches and it would need more memory. Therefore, after we remove an element, we // want to walk the path back towards the root and remove the nodes that are no longer needed // (either remove the branch completely or contract the end that doesn't branch). // // If we, however, started to simply delete the nodes and replace the pointers with nulls, we would // get race conditions: // // 1. We check that the current node is empty (or has only one pointer to data in it) and therefore // is unneeded. // 2. After we do the check but before we manage to update the pointer that points to the current // node, some other thread adds a new pointer into the node. This would make it ineligible for // deletion, but we've already done our check. // 3. We don't know about the addition from the other thread and go ahead with the deletion. This // loses and leaks the added element, or any update the other thread has done. // // The original article used the I-nodes and another kind of nodes to solve this problem. We are // going to get inspired by what they did, but will do it inline in the array of pointers. // // On any sane system, data structures like our inner or data nodes are aligned to multiples of // some number (assuming we have at least 32bit system, it's at least multiple of 4 bytes). This // leaves two bits that are always 0 in the pointers. We can abuse these bits to store additional // flags (we use the utilities of crossbeam_utils to manage the bits). One of these bits, when set // to 1, will mean that the pointer is no longer allowed to be updated. // // So, when removing, we first check once if the inner node may be removed. If not, we're done. If // it looks like it may be, we walk the pointers again, but this time we atomically read and mark // them with the flag. This'll make sure nobody is allowed to sneak an update to it past our second // check (the first one is just an optimisation). So, if we pass the second check, we can safely // proceed and remove the node. Hurray. We can move one level up towards the root and repeat. // // In case the second check failed, we have already marked all the pointers that they are never // ever allowed to be updated again. We can't leave a node like that in the data structure forever. // Therefore, we create a new copy of the node, with clean pointers and replace the node with that // (in that case we can stop the processing at this level). // // So, what the other thread that wants to update the pointer but can't because it is marked does? // It can't wait for the thread that did the marking to finish its job, because then the algorithm // would no longer be lock-free. But it can decide to do its work for it and also do the pruning // (only of this particular one inner node; it won't walk further towards the root). It proceeds to // the second check stage ‒ marks all the pointers (even if they are already marked) and deciding // if it can remove it completely or if it needs to create a brand new copy. One of the threads // competing for the prune operation will succeed, the other will fail during the CaS update of the // parent pointer, but both can proceed because the pruning already happened. // // As this collision on node being prune is likely to be rare in practice and it is already // relatively complex (and hard to test for) situation, the other thread simply completely restarts // the operation instead of trying to get all the corner cases right. The removal thread, however, // proceeds towards the root even in the collision situation ‒ it is responsible for pruning as far // as possible and when going up, the corner cases don't actually happen (well, with the exception // of its whole branch being already removed or contracted away by another removal, but in that // case it'll just waste a little bit of effort in trying to remove stuff in places that are going // to be thrown away anyway ‒ but thanks to crossbeam_epoch, it is still valid memory). // // ## Iteration // // Similar to lookups, iteration doesn't have to care about the flags about non-updateable // pointers ‒ even the old, marked, pointers form a valid representation of the map at a certain // point of time, though not necessarily optimally small. // // Therefore, the iterator simply keeps a stack of nodes it is in, with indices into either the // pointer array or the array of elements in a data node and does a DFS through the data structure. // // # Safety // // The current module contains a lot of unsafe code. In general, there are two kinds of things that // could go wrong. Well, in addition to coding bugs, of course. // // ## Lifetimes & invalid pointers // // First, we simply never insert pointers that would be invalid at that time into the data // structure ‒ whatever gets inserted is just brand new allocated thing. This boils down to just // being careful and, as this is relatively short code, this is possible to accomplish. // // So, we must make sure nothing gets destroyed too soon. To accomplish this, we use the mechanism // of crossbeam_epoch. When we remove something from the data structure, the destruction is // postponed to when all the threads leave the current epoch. We never hold onto the pointers or // references past the current method and we bind the lifetimes of the references to the lifetime // of passed pin and us. // // There are two exceptions to this: // // * The destructor deletes things right away. But as it holds a mutable reference, we can't be // destroying anything for any other thread ‒ the other thread no longer holds reference to us. // * The iterator binds the lifetimes to both itself and the map, but it holds a pin alive, so the // same would still apply. // // ## Inter-thread synchronization of data. // // In general, we use release ordering when putting data into the map and consume ordering when // reading it out of it. The claim is, this is enough. But as the elements are independent on each // other and only the inner nodes we traverse during the operation play any role to us (and we // access the further nodes through the loaded pointers) and there's exactly one path from the root // to each element (therefore anyone getting the element must have touched the same pointers), this // is basically the definition of how release/consume synchronization edges work. // // There are some other orderings through the code, though: // // * Relaxed in case we *fail* a CaS update. However, in such case nothing is modified and we just // throw the data we've created ourselves out, so there's nothing to synchronize. // * Relaxed in the first check for pruning. This one does not look at the actual data behind the // pointers, it simply counts how many pointers are non-null. The second pass does the actual // proper synchronization. // * Relaxed in the is_empty. As we don't care what data it points to (only that it's not null), // this again doesn't have to synchronize anything. // * Relaxed in the destructor. As argued above (and at the destructor itself), we have gained a // unique access to the whole map. Therefore, the whole map, containing the data in it must have // been properly synchronized into the thread already and we are in a single-threaded scenario. // * AcqRel in the pruning. This is because we need to acquire the data pointed to in the case // we'll be making a copy and we'll have to re-release it later on. We also modify the value of // the pointer (with the flag), therefore anyone reading it after us only synchronizes against // us, so we also need to re-release it right now onto that pointer. // // [article]: https://www.researchgate.net/publication/221643801_Concurrent_Tries_with_Efficient_Non-Blocking_Snapshots // [Wikipedia entry]: https://en.wikipedia.org/wiki/Ctrie use std::borrow::Borrow; use std::hash::{BuildHasher, Hash, Hasher}; use std::marker::PhantomData; use std::mem; use std::sync::atomic::Ordering; use arrayvec::ArrayVec; use bitflags::bitflags; use crossbeam_epoch::{Atomic, Guard, Owned, Shared}; use smallvec::SmallVec; pub mod config; pub mod debug; pub mod iterator; use self::config::Config; use crate::existing_or_new::ExistingOrNew; // All directly written, some things are not const fn yet :-(. But tested below. pub(crate) const LEVEL_BITS: usize = 4; pub(crate) const LEVEL_MASK: u64 = 0b1111; pub(crate) const LEVEL_CELLS: usize = 16; pub(crate) const MAX_LEVELS: usize = mem::size_of::<u64>() * 8 / LEVEL_BITS; bitflags! { /// Flags that can be put onto a pointer pointing to a node, specifying some interesting /// things. /// /// Note that this lives inside the unused bits of a pointer. All nodes align at least to a /// machine word and we assume it's at least 32bits, so we have at least 2 bits. struct NodeFlags: usize { /// The Inner containing this pointer is condemned to replacement/pruning. /// /// Changing this pointer is pointer is forbidden, and the containing Inner needs to be /// replaced first with a clean one. const CONDEMNED = 0b01; /// The pointer points not to an inner node, but to data node. /// /// # Rationale /// /// The [`Inner`] nodes are quite large. On the other hand, the values are usually just /// [`Arc`][std::sync::Arc] and there's usually just one at each leaf. That leaves a lot of /// wasted space. /// /// Therefore, instead of having an enum, we have nodes of two distinct types. We recognize /// them by this flag in the pointer pointing to them. If it is a leaf with data, this flag /// is set and anyone accessing it knows it needs to type cast the pointer before using. const DATA = 0b10; } } /// Extracts [`NodeFlags`] from a pointer. fn nf(node: Shared<Inner>) -> NodeFlags { NodeFlags::from_bits(node.tag()).expect("Invalid node flags") } /// Type-casts the pointer to a [`Data`] node. unsafe fn load_data<'a, C: Config>(node: Shared<'a, Inner>) -> &'a Data<C> { assert!( nf(node).contains(NodeFlags::DATA), "Tried to load data from inner node pointer" ); (node.as_raw() as usize as *const Data<C>) .as_ref() .expect("A null pointer with data flag found") } /// Moves a data node behind an [`Owned`] pointer, casts it and provides the correct flags. fn owned_data<C: Config>(data: Data<C>) -> Owned<Inner> { unsafe { Owned::<Inner>::from_raw(Box::into_raw(Box::new(data)) as usize as *mut _) .with_tag(NodeFlags::DATA.bits()) } } /// Type-casts and drops the node as data. unsafe fn drop_data<C: Config>(ptr: Shared<Inner>) { assert!( nf(ptr).contains(NodeFlags::DATA), "Tried to drop data from inner node pointer" ); drop(Owned::from_raw(ptr.as_raw() as usize as *mut Data<C>)); } /// An inner branching node of the trie. /// /// This is just a bunch of pointers to lower levels. #[derive(Default)] struct Inner([Atomic<Inner>; LEVEL_CELLS]); // Instead of distinguishing the very common case of single leaf and collision list in our code, we // just handle everything as a list, possibly with 1 element. // // However, as the case with 1 element is much more probable, we don't want the Vec indirection // there, so we let SmallVec to handle it by not spilling in that case. As the spilled Vec needs 2 // words in addition to the length (pointer and capacity), we have room for 2 Arcs in the not // spilled case too, so we as well might take advantage of it. // TODO: We want the union feature. // // Alternatively, we probably could use the raw allocator API and structure with len + [Arc<..>; 0]. // TODO: Compute the stack length based on the Payload size. type Data<C> = SmallVec<[<C as Config>::Payload; 2]>; enum TraverseState<C: Config, F> { Empty, // Invalid temporary state. Created(C::Payload), Future { key: C::Key, constructor: F }, } impl<C: Config, F: FnOnce(C::Key) -> C::Payload> TraverseState<C, F> { fn key(&self) -> &C::Key { match self { TraverseState::Empty => unreachable!("Not supposed to live in the empty state"), TraverseState::Created(payload) => payload.borrow(), TraverseState::Future { key, .. } => key, } } fn payload(&mut self) -> C::Payload { let (new_val, result) = match mem::replace(self, TraverseState::Empty) { TraverseState::Empty => unreachable!("Not supposed to live in the empty state"), TraverseState::Created(payload) => (TraverseState::Created(payload.clone()), payload), TraverseState::Future { key, constructor } => { let payload = constructor(key); let created = TraverseState::Created(payload.clone()); (created, payload) } }; *self = new_val; result } fn data_owned(&mut self) -> Owned<Inner> { let mut data = Data::<C>::new(); data.push(self.payload()); owned_data::<C>(data) } } #[derive(Copy, Clone, Eq, PartialEq)] enum TraverseMode { Overwrite, IfMissing, } /// How well pruning went. #[derive(Copy, Clone, Eq, PartialEq)] enum PruneResult { /// Removed the node completely, inserted NULL into the parent. Null, /// Contracted an edge, inserted a lone child. Singleton, /// Made a copy, as there were multiple pointers leading from the child. Copy, /// Failed to update the parent, some other thread updated it in the meantime. CasFail, } /// The raw hash trie data structure. /// /// This provides the low level data structure. It does provide the lock-free operations on some /// values. On the other hand, it does not provide user friendly interface. It is designed to /// separate the single implementation of the core algorithm and provide a way to wrap it into /// different interfaces for different use cases. /// /// It, however, can be used to fulfill some less common uses. /// /// The types stored inside and general behaviour is described by the [`Config`] type parameter and /// can be customized using that. /// /// As a general rule, this data structure takes the [`crossbeam_epoch`] [`Guard`] and returns /// borrowed data whenever appropriate. This allows cheaper manipulation if necessary or grouping /// multiple operations together. Note than even methods that would return owned values in /// single-threaded case (eg. [`insert`][Raw::insert] and [`remove`][Raw::remove] return borrowed /// values. This is because in concurrent situation some other thread might still be accessing /// them. They are scheduled for destruction once the epoch ends. /// /// For details of the internal implementation and correctness arguments, see the comments in /// source code (they probably don't belong into API documentation). pub struct Raw<C: Config, S> { hash_builder: S, root: Atomic<Inner>, _data: PhantomData<C::Payload>, } impl<C, S> Raw<C, S> where C: Config, S: BuildHasher, { /// Constructs an empty instance from the given hasher. pub fn with_hasher(hash_builder: S) -> Self { // Note: on any sane system, these assertions should actually never ever trigger no matter // what the user of the crate does. This is *internal* sanity check. If you ever find a // case where it *does* fail, open a bug report. assert!( mem::align_of::<Data<C>>().trailing_zeros() >= NodeFlags::all().bits().count_ones(), "BUG: Alignment of Data<Payload> is not large enough to store the internal flags", ); assert!( mem::align_of::<Inner>().trailing_zeros() >= NodeFlags::all().bits().count_ones(), "BUG: Alignment of Inner not large enough to store internal flags", ); Self { hash_builder, root: Atomic::null(), _data: PhantomData, } } /// Computes a hash (using the stored hasher) of a key. fn hash<Q>(&self, key: &Q) -> u64 where Q: ?Sized + Hash, { let mut hasher = self.hash_builder.build_hasher(); key.hash(&mut hasher); hasher.finish() } /// Inserts a new value, replacing and returning any previously held value. pub fn insert<'s, 'p, 'r>( &'s self, payload: C::Payload, pin: &'p Guard, ) -> Option<&'r C::Payload> where 's: 'r, 'p: 'r, { self.traverse( // Any way to do it without the type parameters here? Older rustc doesn't like them. TraverseState::<C, fn(C::Key) -> C::Payload>::Created(payload), TraverseMode::Overwrite, pin, ) // TODO: Should we sanity-check this is Existing because it returns the previous value? .map(ExistingOrNew::into_inner) } /// Prunes the given node. /// /// * The parent points to the child node. /// * The child must be valid pointer, of course. /// /// The parent is made to point to either: /// * NULL if child is empty. /// * child's only child. /// * A copy of child. /// /// Returns how the pruning went. unsafe fn prune(pin: &Guard, parent: &Atomic<Inner>, child: Shared<Inner>) -> PruneResult { assert!( !nf(child).contains(NodeFlags::DATA), "Child passed to prune must not be data" ); let inner = child.as_ref().expect("Null child node passed to prune"); let mut allow_contract = true; let mut child_cnt = 0; let mut last_leaf = None; let mut new_child = Inner::default(); // 1. Mark all the cells in this one as condemned. // 2. Look how many non-null branches are leading from there. // 3. Construct a copy of the child *without* the tags on the way. for (new, grandchild) in new_child.0.iter_mut().zip(&inner.0) { // Acquire ‒ Besides potentially looking at the child, we'll need to republish the // child in our swap of the pointer (this one and also the one below, in the CAS). To // do that we'll have to have acquired it first. // // Note that we don't need SeqCst here nor in the CaS below. We don't care about the // order ‒ the tagging is just making sure this particular slot never ever changes the // pointer. The CaS changes the trie in content-equivalent way, so observing either the // old or the new way is fine. let gc = grandchild.fetch_or(NodeFlags::CONDEMNED.bits(), Ordering::AcqRel, pin); // The flags we insert into the new one should not contain condemned flag even if it // was already present here. let flags = nf(gc) & !NodeFlags::CONDEMNED; let gc = gc.with_tag(flags.bits()); if gc.is_null() { // Do nothing, just skip } else if flags.contains(NodeFlags::DATA) { last_leaf.replace(gc); let gc = load_data::<C>(gc); child_cnt += gc.len(); } else { // If we have an inner node here, multiple leaves hang somewhere below there. More // importantly, we can't contrack the edge. allow_contract = false; child_cnt += 1; } *new = Atomic::from(gc); } // Now, decide what we want to put into the parent. let mut cleanup = None; let (insert, prune_result) = match (allow_contract, child_cnt, last_leaf) { // If there's exactly one leaf, we just contract the edge to lead there directly. Note // that we can't do that if this is not the leaf, because we would mess up the hash // matching on the way. But that's fine, we checked that above. (true, 1, Some(child)) => (child, PruneResult::Singleton), // If there's nothing, simply kill the node outright. (_, 0, None) => (Shared::null(), PruneResult::Null), // Many nodes (maybe somewhere below) ‒ someone must have inserted in between. But // we've already condemned this node, so create a new one and do the replacement. _ => { let new = Owned::new(new_child).into_shared(pin); // Note: we don't store Owned, because we may link it in. If we panicked before // disarming it, it would delete something linked in, which is bad. Instead, we // prefer deleting manually after the fact. cleanup = Some(new); (new, PruneResult::Copy) } }; assert_eq!( 0, child.tag(), "Attempt to replace condemned pointer or prune data node" ); // Orderings: We need to publish the new node. We don't need to acquire the previous value // to destroy, because we already have it in case of success and we don't care about it on // failure. let result = parent .compare_and_set(child, insert, (Ordering::Release, Ordering::Relaxed), pin) .is_ok(); if result { // We successfully unlinked the old child, so it's time to destroy it (as soon as // nobody is looking at it). pin.defer_destroy(child); prune_result } else { // We have failed to insert, so we need to clean up after ourselves. drop(cleanup.map(|c| Shared::into_owned(c))); PruneResult::CasFail } } /// Inner implementation of traversing the tree, creating missing branches and doing /// *something* at the leaf. fn traverse<'s, 'p, 'r, F>( &'s self, mut state: TraverseState<C, F>, mode: TraverseMode, pin: &'p Guard, ) -> Option<ExistingOrNew<&'r C::Payload>> where 's: 'r, 'p: 'r, F: FnOnce(C::Key) -> C::Payload, { let hash = self.hash(state.key()); let mut shift = 0; let mut current = &self.root; let mut parent = None; loop { let node = current.load_consume(&pin); let flags = nf(node); let replace = |with: Owned<Inner>, delete_previous| { // If we fail to set it, the `with` is dropped together with the Err case, freeing // whatever was inside it. let result = current.compare_and_set_weak( node, with, (Ordering::Release, Ordering::Relaxed), pin, ); match result { Ok(new) if !node.is_null() && delete_previous => { assert!(flags.contains(NodeFlags::DATA)); let node = Shared::from(node.as_raw() as usize as *const Data<C>); unsafe { pin.defer_destroy(node) }; Some(new) } Ok(new) => Some(new), Err(e) => { if NodeFlags::from_bits(e.new.tag()) .expect("Invalid flags") .contains(NodeFlags::DATA) { unsafe { drop_data::<C>(e.new.into_shared(&pin)) }; } // Else → just let e drop and destroy the owned in there None } } }; if flags.contains(NodeFlags::CONDEMNED) { // This one is going away. We are not allowed to modify the cell, we just have to // replace the inner node first. So, let's do some cleanup. // // TODO: In some cases we would not really *have* to do this (in particular, if we // just want to walk through and not modify it here at all, it's OK). unsafe { let (parent, child) = parent.expect("Condemned the root!"); Self::prune(&pin, parent, child); } // Either us or someone else modified the tree on our path. In many cases we // could just continue here, but some cases are complex. For now, we just restart // the whole traversal and try from the start, for simplicity. This should be rare // anyway, so complicating the code further probably is not worth it. shift = 0; current = &self.root; parent = None; } else if node.is_null() { // Not found, create it. if let Some(new) = replace(state.data_owned(), true) { if mode == TraverseMode::Overwrite { return None; } else { let new = unsafe { load_data::<C>(new) }; return Some(ExistingOrNew::New(&new[0])); } } // else -> retry } else if flags.contains(NodeFlags::DATA) { let data = unsafe { load_data::<C>(node) }; assert!(!data.is_empty(), "Empty data nodes must not be kept around"); if data[0].borrow() != state.key() && shift < mem::size_of_val(&hash) * 8 { assert!(data.len() == 1, "Collision node not deep enough"); // There's one data node at this pointer, but we want to place a different one // here too. So we create a new level, push the old one down. Note that we // check both that we are adding something else & that we still have some more // bits to distinguish by. // We need to add another level. Note: there *still* might be a collision. // Therefore, we just add the level and try again. let other_hash = self.hash(data[0].borrow()); let other_bits = (other_hash >> shift) & LEVEL_MASK; let mut inner = Inner::default(); inner.0[other_bits as usize] = Atomic::from(node); let split = Owned::new(inner); // No matter if it succeeds or fails, we try again. We'll either find the newly // inserted value here and continue with another level down, or it gets // destroyed and we try splitting again. replace(split, false); } else { // All the other cases: // * It has the same key // * There's already a collision on this level (because we've already run out of // bits previously). // * We've run out of the hash bits so there's nothing to split by any more. let mut result = data .iter() .find(|l| (*l).borrow().borrow() == state.key()) .map(ExistingOrNew::Existing); if result.is_none() || mode == TraverseMode::Overwrite { let mut new = Data::<C>::with_capacity(data.len() + 1); new.extend( data.iter() .filter(|l| (*l).borrow() != state.key()) .cloned(), ); new.push(state.payload()); new.shrink_to_fit(); let new = owned_data::<C>(new); if let Some(new) = replace(new, true) { if result.is_none() && mode == TraverseMode::IfMissing { let new = unsafe { load_data::<C>(new) }; result = Some(ExistingOrNew::New(new.last().unwrap())); } } else { continue; } } return result; } } else { // An inner node, go one level deeper. let inner = unsafe { node.as_ref().expect("We just checked this is not NULL") }; let bits = (hash >> shift) & LEVEL_MASK; shift += LEVEL_BITS; parent = Some((current, node)); current = &inner.0[bits as usize]; } } } /// Looks up a value. pub fn get<'r, 's, 'p, Q>(&'s self, key: &Q, pin: &'p Guard) -> Option<&'r C::Payload> where 's: 'r, 'p: 's, Q: ?Sized + Eq + Hash, C::Key: Borrow<Q>, { let mut current = &self.root; let mut hash = self.hash(key); loop { let node = current.load_consume(pin); let flags = nf(node); if node.is_null() { return None; } else if flags.contains(NodeFlags::DATA) { return unsafe { load_data::<C>(node) } .iter() .find(|l| (*l).borrow().borrow() == key); } else { let inner = unsafe { node.as_ref().expect("We just checked this is not NULL") }; let bits = hash & LEVEL_MASK; hash >>= LEVEL_BITS; current = &inner.0[bits as usize]; } } } /// Looks up a value or create (and insert) a new one. /// /// Either way, returns the value. pub fn get_or_insert_with<'s, 'p, 'r, F>( &'s self, key: C::Key, create: F, pin: &'p Guard, ) -> ExistingOrNew<&'r C::Payload> where 's: 'r, 'p: 'r, F: FnOnce(C::Key) -> C::Payload, { let state = TraverseState::Future { key, constructor: create, }; self.traverse(state, TraverseMode::IfMissing, pin) .expect("Should have created one for me") } /// Removes a value identified by the key from the trie, returning it if it was found. pub fn remove<'r, 's, 'p, Q>(&'s self, key: &Q, pin: &'p Guard) -> Option<&'r C::Payload> where 's: 'r, 'p: 'r, Q: ?Sized + Eq + Hash, C::Key: Borrow<Q>, { let mut current = &self.root; let hash = self.hash(key); let mut shift = 0; let mut levels: ArrayVec<[_; MAX_LEVELS]> = ArrayVec::new(); let deleted = loop { let node = current.load_consume(&pin); let flags = nf(node); let replace = |with: Shared<_>| { let result = current.compare_and_set_weak( node, with, (Ordering::Release, Ordering::Relaxed), &pin, ); match result { Ok(_) => { assert!(flags.contains(NodeFlags::DATA)); unsafe { let node = Shared::from(node.as_raw() as usize as *const Data<C>); pin.defer_destroy(node); } true } Err(ref e) if !e.new.is_null() => { assert!(nf(e.new).contains(NodeFlags::DATA)); unsafe { drop_data::<C>(e.new) }; false } Err(_) => false, } }; if node.is_null() { // Nothing to delete, so just give up (without pruning). return None; } else if flags.contains(NodeFlags::CONDEMNED) { unsafe { let (current, node) = levels.pop().expect("Condemned the root"); Self::prune(&pin, current, node); } // Retry by starting over from the top, for similar reasons to the one in // insert. levels.clear(); shift = 0; current = &self.root; } else if flags.contains(NodeFlags::DATA) { let data = unsafe { load_data::<C>(node) }; // Try deleting the thing. let mut deleted = None; let new = data .iter() .filter(|l| { if (*l).borrow().borrow() == key { deleted = Some(*l); false } else { true } }) .cloned() .collect::<Data<C>>(); if deleted.is_some() { let new = if new.is_empty() { Shared::null() } else { owned_data::<C>(new).into_shared(&pin) }; if !replace(new) { continue; } } break deleted; } else { let inner = unsafe { node.as_ref().expect("We just checked for NULL") }; levels.push((current, node)); let bits = (hash >> shift) & LEVEL_MASK; shift += LEVEL_BITS; current = &inner.0[bits as usize]; } }; // Go from the top and try to clean up. if deleted.is_some() { for (parent, child) in levels.into_iter().rev() { let inner = unsafe { child.as_ref().expect("We just checked for NULL") }; // This is an optimisation ‒ replacing the thing is expensive, so we want to check // first (which is cheaper). let non_null = inner .0 .iter() .filter(|ptr| !ptr.load(Ordering::Relaxed, &pin).is_null()) .count(); if non_null > 1 { // No reason to go into the upper levels. break; } // OK, we think we could remove this node. Try doing so. if let PruneResult::Copy = unsafe { Self::prune(&pin, parent, child) } { // Even though we tried to count how many pointers there are, someone must have // added some since. So there's no way we can prone anything higher up and we // give up. break; } // Else: // Just continue with higher levels. Even if someone made the contraction for // us, it should be safe to do so. } } deleted } } impl<C: Config, S> Raw<C, S> { /// Checks for emptiness. pub fn is_empty(&self) -> bool { // This relies on proper branch pruning. // We can use the unprotected here, because we are not actually interested in where the // pointer points to. Therefore we can also use the Relaxed ordering. unsafe { self.root .load(Ordering::Relaxed, &crossbeam_epoch::unprotected()) .is_null() } } /// Access to the hash builder. pub fn hash_builder(&self) -> &S { &self.hash_builder } } impl<C: Config, S> Drop for Raw<C, S> { fn drop(&mut self) { /* * Notes about unsafety here: * * We are in a destructor and that one is &mut self. There are no concurrent accesses to * this data structure any more, therefore we can safely assume we are the only ones * looking at the pointers inside. * * Therefore, using unprotected is also fine. * * Similarly, the Relaxed ordering here is fine too, as the whole data structure must * have been synchronized into our thread already by this time. * * The pointer inside this data structure is never dangling. */ unsafe fn drop_recursive<C: Config>(node: &Atomic<Inner>) { let pin = crossbeam_epoch::unprotected(); let extract = node.load(Ordering::Relaxed, &pin); let flags = nf(extract); if extract.is_null() { // Skip } else if flags.contains(NodeFlags::DATA) { drop_data::<C>(extract); } else { let owned = extract.into_owned(); for sub in &owned.0 { drop_recursive::<C>(sub); } drop(owned); } } unsafe { drop_recursive::<C>(&self.root) }; } } #[cfg(test)] pub(crate) mod tests { use std::ptr; use super::config::Trivial as TrivialConfig; use super::*; // A hasher to create collisions on purpose. Let's make the hash trie into a glorified array. // We allow tests in higher-level modules to reuse it for their tests. pub(crate) struct NoHasher; impl Hasher for NoHasher { fn finish(&self) -> u64 { 0 } fn write(&mut self, _: &[u8]) {} } impl BuildHasher for NoHasher { type Hasher = NoHasher; fn build_hasher(&self) -> NoHasher { NoHasher } } #[derive(Clone, Copy, Debug, Default)] pub(crate) struct SplatHasher(u64); impl Hasher for SplatHasher { fn finish(&self) -> u64 { self.0 } fn write(&mut self, value: &[u8]) { for val in value { for idx in 0..mem::size_of::<u64>() { self.0 ^= u64::from(*val) << (8 * idx); } } } } pub(crate) struct MakeSplatHasher; impl BuildHasher for MakeSplatHasher { type Hasher = SplatHasher; fn build_hasher(&self) -> SplatHasher { SplatHasher::default() } } /// Tests the test hasher. /// /// Because it was giving us some trouble ☹ #[test] fn splat_hasher() { let mut hasher = MakeSplatHasher.build_hasher(); hasher.write_u8(0); assert_eq!(0, hasher.finish()); hasher.write_u8(8); assert_eq!(0x0808_0808_0808_0808, hasher.finish()); } #[test] fn consts_consistent() { assert!(LEVEL_CELLS.is_power_of_two()); assert_eq!(LEVEL_BITS, LEVEL_MASK.count_ones() as usize); assert_eq!(LEVEL_BITS, (!LEVEL_MASK).trailing_zeros() as usize); assert_eq!(LEVEL_CELLS, 2usize.pow(LEVEL_BITS as u32)); } /// Pretend something left a condemned marker on one of the nodes when we insert. This will get /// cleaned up. /// /// And yes, the test abuses the fact that it knows how the specific hasher works and /// distributes the given values. #[test] fn prune_on_insert() { let mut map = Raw::<TrivialConfig<u8>, _>::with_hasher(MakeSplatHasher); let pin = crossbeam_epoch::pin(); for i in 0..LEVEL_CELLS as u8 { assert!(map.insert(i, &pin).is_none()); } eprintln!("{}", debug::PrintShape(&map)); // By now, we should have exactly one data node under each pointer under root. Sanity // check that (Relaxed is fine, we are in a single threaded test). let root = map.root.load(Ordering::Relaxed, &pin); let flags = nf(root); assert_eq!( NodeFlags::empty(), flags, "Root should be non-condemned inner node" ); assert!(!root.is_null()); let old_root = root.as_raw(); let root = unsafe { root.deref() }; for ptr in &root.0 { let ptr = ptr.load(Ordering::Relaxed, &pin); assert!(!ptr.is_null()); let flags = nf(ptr); assert_eq!( NodeFlags::DATA, flags, "Expected a data node, found {:?}", ptr ); } // Now, *start* condemning the node. Mark the first slot, the one we'll eventually use. root.0[0].fetch_or(NodeFlags::CONDEMNED.bits(), Ordering::Relaxed, &pin); // This touches the condemned slot, so it should trigger fixing stuff. let old = map.insert(0, &pin); assert_eq!(0, *old.unwrap()); // The condemned flag must have disappeared by now. map.assert_pruned(); // And the root should have changed for a brand new one. let new_root = map.root.load(Ordering::Relaxed, &pin).as_raw(); assert!(!ptr::eq(old_root, new_root), "Condemned node not replaced"); // But all the content is preserved for i in 0..LEVEL_CELLS as u8 { assert_eq!(i, *map.get(&i, &pin).unwrap()); } } /// Creates an effectively empty map with a leftover (unpruned) but condemned node. /// /// As the algorithm goes, almost everyone who finds it is responsible for cleaning it up. fn with_leftover() -> Raw<TrivialConfig<u8>, MakeSplatHasher> { let map = Raw::<TrivialConfig<u8>, _>::with_hasher(MakeSplatHasher); let pin = crossbeam_epoch::pin(); let i = Inner::default(); i.0[0].fetch_or(NodeFlags::CONDEMNED.bits(), Ordering::Relaxed, &pin); map.root.store(Owned::new(i), Ordering::Relaxed); // There's nothing in this map effectively, but it doesn't claim to be empty due to the // non-null pointer. assert!(iterator::Iter::new(&map).next().is_none()); assert!(!map.is_empty()); map } /// Similar as the above, but with empty condemned node. /// /// Here we put a fake node somewhere into the aether, make it condemned and see how it /// disappears on insertion. #[test] fn prune_on_insert_empty() { let mut map = with_leftover(); let pin = crossbeam_epoch::pin(); let old_root = map.root.load(Ordering::Relaxed, &pin).as_raw(); // Now, let's insert something so it meets the condemned mark assert!(map.insert(0, &pin).is_none()); map.assert_pruned(); let new_root = map.root.load(Ordering::Relaxed, &pin); // It got replaced and the root is directly the data node let new_flags = nf(new_root); assert_eq!(NodeFlags::DATA, new_flags); assert!( !ptr::eq(old_root, new_root.as_raw()), "Condemned node not replaced" ); } /// Test that if someone left a un-pruned node and remove finds it, it gets rid of it (even in /// cases it does not actually remove anything in particular). #[test] fn prune_on_remove() { let map = Raw::<TrivialConfig<u8>, _>::with_hasher(MakeSplatHasher); let pin = crossbeam_epoch::pin(); let i_inner = Inner::default(); let i_outer = Inner::default(); i_outer.0[0].store( Owned::new(i_inner).with_tag(NodeFlags::CONDEMNED.bits()), Ordering::Relaxed, ); map.root.store(Owned::new(i_outer), Ordering::Relaxed); // There's nothing in this map effectively, but it doesn't claim to be empty due to the // non-null pointer. assert!(iterator::Iter::new(&map).next().is_none()); assert!(!map.is_empty()); assert!(map.remove(&0, &pin).is_none()); eprintln!("{}", debug::PrintShape(&map)); assert_eq!(0, map.root.load(Ordering::Relaxed, &pin).tag()); // Note: it is still *not* properly pruned. The inner node should have a thread it'll clean // up later on. And we can't contract it as the one below is inner node, not data node. } }