From 65af58a57c35797037821c75f399ced986cccba2 Mon Sep 17 00:00:00 2001 From: Igor Malovitsa Date: Thu, 17 Sep 2026 05:49:03 +0000 Subject: [PATCH 1/4] Fix use-after-free between ZipperHead readers and writers A head reader owned whatever node the walk down its path reached, which for a path that stops inside a node -- or is not in the trie at all -- is a node the head's writers descend through. The next exclusive writer copied it, and the writers already made kept pointing into the copy that was then dropped; the reproducer segfaults. A value at the reader's root was borrowed from the node above it, which writers copy the same way. Keep owning the node when the path lands on one and it carries no value of its own: no writer may be at or below the reader's path, so none is inside it, and that is the common case. Otherwise copy just this entry -- its value, its subtrie, or its dangling path -- into a private node, and root the reader one byte above it. --- src/zipper.rs | 99 ++++++++++++++++++++++++++++++++++++++++------ src/zipper_head.rs | 30 ++++++++++++-- 2 files changed, 114 insertions(+), 15 deletions(-) diff --git a/src/zipper.rs b/src/zipper.rs index 15fd560d..a99b3486 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1493,10 +1493,13 @@ impl<'a, V: Clone + Send + Sync + Unpin + 'a, A: Allocator + 'a> ZipperReadOnlyP } impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ReadZipperTracked<'a, 'path, V, A> { - /// See [ReadZipperCore::new_with_node_and_path] - pub(crate) fn new_with_node_and_path_in(root_node: &'a TrieNodeODRc, owned_root: bool, path: &'path [u8], root_prefix_len: usize, root_key_start: usize, root_val: Option<&'a V>, alloc: A, tracker: Option>) -> Self { - let core = ReadZipperCore::new_with_node_and_path_in(root_node, owned_root, path, root_prefix_len, root_key_start, root_val, alloc); - Self { z: core, tracker } + /// See [ReadZipperCore::new_isolated_in] + pub(crate) fn new_isolated_in(root_node: &'a TrieNodeODRc, path: &'path [u8], root_val: Option<&'a V>, alloc: A, tracker: Option>) -> Self { + Self { z: ReadZipperCore::new_isolated_in(root_node, path, root_val, alloc), tracker } + } + /// See [ReadZipperCore::new_isolated_cloned_path_in] + pub(crate) fn new_isolated_cloned_path_in(root_node: &'a TrieNodeODRc, path: &[u8], root_val: Option<&'a V>, alloc: A, tracker: Option>) -> Self { + Self { z: ReadZipperCore::new_isolated_cloned_path_in(root_node, path, root_val, alloc), tracker } } /// See [ReadZipperCore::new_with_node_and_cloned_path] pub(crate) fn new_with_node_and_cloned_path_in(root_node: &'a TrieNodeODRc, owned_root: bool, path: &[u8], root_prefix_len: usize, root_key_start: usize, root_val: Option<&'a V>, alloc: A, tracker: Option>) -> Self { @@ -2630,11 +2633,11 @@ pub(crate) mod read_zipper_core { let (_key_len, focus_node) = parent.node_get_child(self.parent_key()).unwrap(); !focus_node.is_empty() && focus_node.refcount() > 1 } else { - match &self.root_node { - OwnedOrBorrowed::Owned(root) => !root.is_empty() && root.refcount() > 1, - OwnedOrBorrowed::Borrowed(root) => !root.is_empty() && root.refcount() > 1, - OwnedOrBorrowed::None => false, - } + let focus = match &self.root_node { + OwnedOrBorrowed::None => return false, + _ => self.focus_parent(), + }; + !focus.is_empty() && focus.refcount() > 1 } } } @@ -2806,6 +2809,67 @@ pub(crate) mod read_zipper_core { new_zipper.make_static_path() } + /// Like [Self::new_with_node_and_path_in] with an owned root, but never holding a node that a + /// live `ZipperHead` writer can reach. Sharing one lets the next exclusive writer copy it, + /// leaving the writers already made pointing into the copy that is about to be dropped + pub(crate) fn new_isolated_in(root_node: &'a TrieNodeODRc, path: &'path [u8], root_val: Option<&'a V>, alloc: A) -> Self { + //A reader at the head's own root excludes every writer, so it may hold the root node + let Some(&last) = path.last() else { + return Self::new_with_node_and_path_in(root_node, true, path, 0, 0, root_val, alloc) + }; + let (node, key, val) = node_along_path(root_node, path, root_val, false); + + //The focus is a whole node and carries no value of its own, so the zipper can own that + // node outright: no writer may be at or below the reader's path, so none is inside it + if key.is_empty() && val.is_none() { + return Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Owned(node.clone()), path, path.len(), None, alloc) + } + + Self::new_isolated_entry_in(node, key, val, path, last, alloc) + } + /// The uncommon half of [Self::new_isolated_in]: the entry at the focus lives in a node that + /// `ZipperHead` writers reach as well, because its value sits beside their branches or the + /// node is one they descend through. Copy just this entry into a private node. + #[inline(never)] + #[cold] + fn new_isolated_entry_in(node: &'a TrieNodeODRc, key: &[u8], val: Option<&'a V>, path: &'path [u8], last: u8, alloc: A) -> Self { + let (val, child, dangling) = if key.is_empty() { + //A value at the focus lives in the node above it, which is not ours to hold + let child = (!node.is_empty()).then(|| node.clone()); + (val.cloned(), child, false) + } else { + let node = node.as_tagged(); + let val = node.node_get_val(key).cloned(); + let child = node.get_node_at_key(key).into_option(); + let dangling = val.is_none() && child.is_none() && node.node_contains_partial_key(key); + (val, child, dangling) + }; + #[cfg(not(feature = "all_dense_nodes"))] + let mut root = TrieNodeODRc::new_in(crate::line_list_node::LineListNode::new_in(alloc.clone()), alloc.clone()); + #[cfg(feature = "all_dense_nodes")] + let mut root = TrieNodeODRc::new_in(crate::dense_byte_node::DenseByteNode::new_in(alloc.clone()), alloc.clone()); + if let Some(val) = val { + if let Err(n) = root.make_mut().node_set_val(&[last], val) { root = n } + } + if let Some(child) = child { + if let Err(n) = root.make_mut().node_set_branch(&[last], child) { root = n } + } + if dangling { + if let Err(n) = root.make_mut().node_create_dangling(&[last]) { root = n } + } + //The root value is read from `root`, via `root_parent_key_start` + Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Owned(root), path, path.len() - 1, None, alloc) + } + /// Same as [Self::new_isolated_in], but with a `'static` path + pub(crate) fn new_isolated_cloned_path_in(root_node: &'a TrieNodeODRc, path: &[u8], root_val: Option<&'a V>, alloc: A) -> ReadZipperCore<'a, 'static, V, A> { + let mut new_zipper = ReadZipperCore::<'a, '_, V, A>::new_isolated_in(root_node, path, root_val, alloc); + new_zipper.prefix_buf = Vec::with_capacity(EXPECTED_PATH_LEN); + new_zipper.prefix_buf.extend(path); + new_zipper.origin_path = SliceOrLen::new_owned(path.len()); + new_zipper.ancestors = Vec::with_capacity(EXPECTED_DEPTH); + new_zipper.make_static_path() + } + /// Makes a version of `self` that has an allocated path buffer and a `'static`` path lifetime #[inline] pub(crate) fn make_static_path(mut self) -> ReadZipperCore<'a, 'static, V, A> { @@ -2888,7 +2952,12 @@ pub(crate) mod read_zipper_core { // we currently share the same implementation between `val()` and `get_val()` because the only difference is the return // lifetime, and the current ZipperHead implementation is actually ok with referencing the value in the root of the ZipperHead. // debug_assert!(self.root_node.is_borrowed()); - self.root_val + if self.root_val.is_some() || self.root_parent_key_start == usize::MAX || !self.root_node.is_owned() { + self.root_val + } else { + //SAFETY: see the note on this method + self.root_node.as_ref().as_tagged().node_get_val(self.root_node_key()).map(|v| unsafe{ &*(v as *const V) }) + } } } } @@ -3017,6 +3086,10 @@ pub(crate) mod read_zipper_core { if parent_key.len() == 0 { return self.root_node.as_ref() } + if self.ancestors.is_empty() { + //At the root, with the focus on a child of the root node + return self.root_node.as_ref().as_tagged().node_get_child(parent_key).unwrap().1 + } self.focus_parent_borrowed() } @@ -3239,8 +3312,10 @@ pub(crate) mod read_zipper_core { } else { if let Some((parent, _iter_tok, _prefix_offset)) = self.ancestors.last() { parent.node_contains_val(self.parent_key()) - } else { + } else if self.root_val.is_some() || self.root_parent_key_start == usize::MAX || !self.root_node.is_owned() { self.root_val.is_some() + } else { + self.root_node.as_ref().as_tagged().node_contains_val(self.root_node_key()) } } } @@ -3327,6 +3402,8 @@ pub(crate) mod read_zipper_core { if self.prefix_buf.len() > 0 { let key_start = if self.ancestors.len() > 1 { unsafe{ self.ancestors.get_unchecked(self.ancestors.len()-2) }.2 + } else if self.ancestors.is_empty() && self.root_parent_key_start != usize::MAX { + self.root_parent_key_start } else { self.root_key_start }; diff --git a/src/zipper_head.rs b/src/zipper_head.rs index 20357a74..df405619 100644 --- a/src/zipper_head.rs +++ b/src/zipper_head.rs @@ -209,7 +209,7 @@ impl<'trie, Z, V: 'trie + Clone + Send + Sync + Unpin, A: Allocator + 'trie> Zip // logic makes sure conflicting paths aren't permitted, so we should not get aliased &mut borrows let root_node: &'trie TrieNodeODRc = unsafe{ core::mem::transmute(root_node) }; let root_val: Option<&'trie V> = root_val.map(|v| unsafe{ &*v.as_ptr() } ); - let new_zipper = ReadZipperTracked::new_with_node_and_path_in(root_node, true, path.as_ref(), path.len(), 0, root_val, z.alloc.clone(), zipper_tracker); + let new_zipper = ReadZipperTracked::new_isolated_in(root_node, path, root_val, z.alloc.clone(), zipper_tracker); Ok(new_zipper) }) } @@ -229,7 +229,7 @@ impl<'trie, Z, V: 'trie + Clone + Send + Sync + Unpin, A: Allocator + 'trie> Zip #[cfg(not(debug_assertions))] let zipper_tracker = None; - ReadZipperTracked::new_with_node_and_path_in(root_node, true, path.as_ref(), path.len(), 0, root_val, z.alloc.clone(), zipper_tracker) + ReadZipperTracked::new_isolated_in(root_node, path, root_val, z.alloc.clone(), zipper_tracker) }) } fn read_zipper_at_path<'a, K: AsRef<[u8]>>(&'a self, path: K) -> Result, Conflict> where 'trie: 'a { @@ -242,7 +242,7 @@ impl<'trie, Z, V: 'trie + Clone + Send + Sync + Unpin, A: Allocator + 'trie> Zip let root_node: &'trie TrieNodeODRc = unsafe{ core::mem::transmute(root_node) }; let root_val: Option<&'trie V> = root_val.map(|v| unsafe{ &*v.as_ptr() } ); - let new_zipper = ReadZipperTracked::new_with_node_and_cloned_path_in(root_node, true, path.as_ref(), path.len(), 0, root_val, z.alloc.clone(), Some(zipper_tracker)); + let new_zipper = ReadZipperTracked::new_isolated_cloned_path_in(root_node, path, root_val, z.alloc.clone(), Some(zipper_tracker)); Ok(new_zipper) }) } @@ -263,7 +263,7 @@ impl<'trie, Z, V: 'trie + Clone + Send + Sync + Unpin, A: Allocator + 'trie> Zip #[cfg(not(debug_assertions))] let zipper_tracker = None; - ReadZipperTracked::new_with_node_and_cloned_path_in(root_node, true, path.as_ref(), path.len(), 0, root_val, z.alloc.clone(), zipper_tracker) + ReadZipperTracked::new_isolated_cloned_path_in(root_node, path, root_val, z.alloc.clone(), zipper_tracker) }) } fn write_zipper_at_exclusive_path<'a, K: AsRef<[u8]>>(&'a self, path: K) -> Result, Conflict> where 'trie: 'a { @@ -1524,6 +1524,28 @@ mod tests { assert_eq!(paths, vec![b"ax".to_vec(), b"bx".to_vec(), b"c".to_vec(), b"dx".to_vec()]); } + /// A reader must not hold a node that live writers point into + #[test] + fn head_reader_beside_live_writers() { + let mut map = PathMap::::new(); + map.set_val_at(&[0u8, 0], 1); + let zh = map.into_zipper_head(&[]); + { + let mut w1 = zh.write_zipper_at_exclusive_path(&[0x11u8]).unwrap(); + //Not in the trie, so it used to hold the root node, which the next writer then copied + let r0 = zh.read_zipper_at_path(&[0x22u8, 0, 0]).unwrap(); + let w0 = zh.write_zipper_at_exclusive_path(&[0u8]).unwrap(); + assert!(!r0.path_exists()); + drop(r0); + w1.set_val(5); + drop(w0); + drop(w1); + } + let map = zh.into_map(); + assert_eq!(map.get_val_at(&[0x11u8]), Some(&5)); + assert_eq!(map.get_val_at(&[0u8, 0]), Some(&1)); + } + /// `get_trie_ref`, `get_focus` and forks from a head's read zipper, which owns its root node #[test] fn head_read_zipper_trie_refs() { From b3beb65ce436c18bb947820d3858c376183d40ad Mon Sep 17 00:00:00 2001 From: Luke Peterson Date: Tue, 29 Sep 2026 20:58:00 -0600 Subject: [PATCH 2/4] Updating is_shared docs, adding more zipper_head focussed benchmarks, squishing warnings in test --- benches/zipper_head_owned.rs | 128 +++++++++++++++++++++++++++++++++++ src/zipper.rs | 4 ++ src/zipper_head.rs | 4 +- 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/benches/zipper_head_owned.rs b/benches/zipper_head_owned.rs index 7e8534ad..4508d526 100644 --- a/benches/zipper_head_owned.rs +++ b/benches/zipper_head_owned.rs @@ -33,6 +33,22 @@ where }); } +fn bench_head_read_at_path<'trie, H>(bencher: Bencher, head: &H, path: &[u8]) +where + H: ZipperCreation<'trie, usize>, +{ + bencher.bench_local(|| { + let mut observed = 0usize; + for _ in 0..REPEATS { + let reader = head.read_zipper_at_borrowed_path(black_box(path)).unwrap(); + observed += reader.val().copied().unwrap_or_default(); + observed += reader.child_count(); + drop(reader); + } + black_box(observed); + }); +} + fn bench_head_write_creation_cleanup<'trie, H, const CHECKED: bool>(bencher: Bencher, head: &H) where H: ZipperCreation<'trie, usize>, @@ -72,6 +88,118 @@ fn owned_head_read_creation(bencher: Bencher) { bench_head_read_creation(bencher, &head); } +#[divan::bench] +fn borrowed_head_read_value_in_shared_parent(bencher: Bencher) { + let mut map = PathMap::::new(); + map.set_val_at([0x22u8], 22); + map.set_val_at([0x22u8, 0x01], 1); + let head = black_box(&mut map).zipper_head(); + let _writer = head.write_zipper_at_exclusive_path([0x11u8]).unwrap(); + bench_head_read_at_path(bencher, &head, &[0x22]); +} + +#[divan::bench] +fn borrowed_head_read_missing_path(bencher: Bencher) { + let mut map = zipper_head_fixture(); + let head = black_box(&mut map).zipper_head(); + let _writer = head.write_zipper_at_exclusive_path([0x11u8]).unwrap(); + bench_head_read_at_path(bencher, &head, &[0xee, 0x00]); +} + +const READ_ACCESS_REPEATS: usize = 1000; + +fn bench_shared_parent_reader_access(bencher: Bencher) { + let mut map = PathMap::::new(); + map.set_val_at([0x22u8], 22); + map.set_val_at([0x22u8, 0x01], 1); + let head = black_box(&mut map).zipper_head(); + let _writer = head.write_zipper_at_exclusive_path([0x11u8]).unwrap(); + let reader = head.read_zipper_at_borrowed_path(&[0x22u8]).unwrap(); + assert_eq!(reader.val(), Some(&22)); + assert_eq!(reader.child_count(), 1); + bencher.bench_local(|| { + let mut observed = 0usize; + for _ in 0..READ_ACCESS_REPEATS { + let reader = black_box(&reader); + if OP == 0 { + observed += reader.val().copied().unwrap_or_default(); + } else if OP == 1 { + observed += reader.is_val() as usize; + } else if OP == 2 { + observed += reader.is_shared() as usize; + } else { + black_box(reader.get_focus()); + } + } + black_box(observed); + }); +} + +#[divan::bench] +fn shared_parent_reader_val(bencher: Bencher) { + bench_shared_parent_reader_access::<0>(bencher); +} + +#[divan::bench] +fn shared_parent_reader_is_val(bencher: Bencher) { + bench_shared_parent_reader_access::<1>(bencher); +} + +#[divan::bench] +fn shared_parent_reader_is_shared(bencher: Bencher) { + bench_shared_parent_reader_access::<2>(bencher); +} + +#[divan::bench] +fn shared_parent_reader_get_focus(bencher: Bencher) { + bench_shared_parent_reader_access::<3>(bencher); +} + +fn bench_map_reader_access(bencher: Bencher) { + let mut map = PathMap::::new(); + map.set_val_at([0x22u8], 22); + map.set_val_at([0x22u8, 0x01], 1); + let reader = map.read_zipper_at_path([0x22u8]); + assert_eq!(reader.val(), Some(&22)); + assert_eq!(reader.child_count(), 1); + bencher.bench_local(|| { + let mut observed = 0usize; + for _ in 0..READ_ACCESS_REPEATS { + let reader = black_box(&reader); + if OP == 0 { + observed += reader.val().copied().unwrap_or_default(); + } else if OP == 1 { + observed += reader.is_val() as usize; + } else if OP == 2 { + observed += reader.is_shared() as usize; + } else { + black_box(reader.get_focus()); + } + } + black_box(observed); + }); +} + +#[divan::bench] +fn map_reader_val(bencher: Bencher) { + bench_map_reader_access::<0>(bencher); +} + +#[divan::bench] +fn map_reader_is_val(bencher: Bencher) { + bench_map_reader_access::<1>(bencher); +} + +#[divan::bench] +fn map_reader_is_shared(bencher: Bencher) { + bench_map_reader_access::<2>(bencher); +} + +#[divan::bench] +fn map_reader_get_focus(bencher: Bencher) { + bench_map_reader_access::<3>(bencher); +} + #[divan::bench] fn borrowed_head_write_creation_cleanup(bencher: Bencher) { let mut map = zipper_head_fixture(); diff --git a/src/zipper.rs b/src/zipper.rs index a99b3486..0f35c6b8 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1223,6 +1223,10 @@ pub trait ZipperConcrete { /// subtrie may be copied for thread isolation, or the internal trie representation might otherwise /// change, and alter the shared property. /// + /// NOTE: A `ZipperHead` may produce read zippers that own private references to subtries. Therefore + /// `is_shared` does not always indicate genuine structural sharing (i.e. sharing between upstream + /// parents) when it is called at the root of a read zipper created from a zipper head. + /// /// GOAT: Make a graphic diagram to illustrate the `shared` property. The graphic should have /// multiple shared subtries accessible via distinct paths, and highlight which locations will be /// considered `shared` from the perspective of this method. diff --git a/src/zipper_head.rs b/src/zipper_head.rs index df405619..c7abc1c0 100644 --- a/src/zipper_head.rs +++ b/src/zipper_head.rs @@ -1542,8 +1542,8 @@ mod tests { drop(w1); } let map = zh.into_map(); - assert_eq!(map.get_val_at(&[0x11u8]), Some(&5)); - assert_eq!(map.get_val_at(&[0u8, 0]), Some(&1)); + assert_eq!(map.val_at(&[0x11u8]), Some(&5)); + assert_eq!(map.val_at(&[0u8, 0]), Some(&1)); } /// `get_trie_ref`, `get_focus` and forks from a head's read zipper, which owns its root node From eb8f7c59f76007c36370d37468c3571d2b14120e Mon Sep 17 00:00:00 2001 From: Luke Peterson Date: Thu, 1 Oct 2026 01:46:06 -0600 Subject: [PATCH 3/4] Adding some new tests to cover potential unsoundness with witness calls --- src/zipper_head.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/src/zipper_head.rs b/src/zipper_head.rs index c7abc1c0..69333943 100644 --- a/src/zipper_head.rs +++ b/src/zipper_head.rs @@ -1546,6 +1546,71 @@ mod tests { assert_eq!(map.val_at(&[0u8, 0]), Some(&1)); } + #[test] + fn head_reader_private_root_value_and_witness() { + let mut map = PathMap::::new(); + map.set_val_at([0x22], 22); + map.set_val_at([0x22, 0x01], 1); + map.set_val_at([0x44, 0x55], 55); + let zh = map.into_zipper_head([]); + let mut other_writer = zh.write_zipper_at_exclusive_path([0x11]).unwrap(); + + let reader = zh.read_zipper_at_borrowed_path(&[0x22]).unwrap(); + assert!(reader.path_exists()); + assert!(reader.is_val()); + assert_eq!(reader.val(), Some(&22)); + assert_eq!(reader.val_at([0x01]), Some(&1)); + let cloned_reader = reader.clone(); + let witness = reader.witness(); + let held_value = reader.get_val_with_witness(&witness).unwrap(); + drop(reader); + assert_eq!(cloned_reader.val(), Some(&22)); + drop(cloned_reader); + + let mut writer = zh.write_zipper_at_exclusive_path([0x22]).unwrap(); + writer.set_val(99); + drop(writer); + assert_eq!(*held_value, 22); + + let prefix = zh.read_zipper_at_borrowed_path(&[0x44]).unwrap(); + assert!(prefix.path_exists()); + assert!(!prefix.is_val()); + assert_eq!(prefix.val_at([0x55]), Some(&55)); + drop(prefix); + + let missing = zh.read_zipper_at_borrowed_path(&[0xee, 0x00]).unwrap(); + assert!(!missing.path_exists()); + assert!(!missing.is_val()); + assert_eq!(missing.child_count(), 0); + drop(missing); + + other_writer.set_val(11); + } + + /// A witness must keep the head's root value alive after its reader releases the path lock. + #[test] + fn head_root_value_witness_survives_reader_and_writer() { + use std::sync::Arc; + + let original = Arc::new(String::from("old")); + let mut map = PathMap::>::new(); + map.set_val_at([], original.clone()); + let head = map.zipper_head(); + + let reader = head.read_zipper_at_path([]).unwrap(); + let witness = reader.witness(); + assert_eq!(reader.get_val_with_witness(&witness).map(|v| v.as_str()), Some("old")); + drop(reader); + + let mut writer = head.write_zipper_at_exclusive_path([]).unwrap(); + drop(writer.remove_val(false)); + drop(writer); + + // The map no longer owns the old value. A valid witness must still own it. + assert!(Arc::strong_count(&original) > 1); + drop(witness); + } + /// `get_trie_ref`, `get_focus` and forks from a head's read zipper, which owns its root node #[test] fn head_read_zipper_trie_refs() { From b9ef9810eea10bce6fcf15e31e068a8a0aff834b Mon Sep 17 00:00:00 2001 From: Luke Peterson Date: Thu, 1 Oct 2026 18:02:08 -0600 Subject: [PATCH 4/4] Clawing back perf and fixing issue with incomplete witnesses --- src/zipper.rs | 203 ++++++++++++++++++-------------------------------- 1 file changed, 71 insertions(+), 132 deletions(-) diff --git a/src/zipper.rs b/src/zipper.rs index 0f35c6b8..675f41ca 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -764,7 +764,7 @@ pub trait ZipperReadOnlyValues<'a, V>: ZipperValues { } /// A [`witness`](ZipperReadOnlyConditionalValues::witness) type used by [`ReadZipperTracked`] and [`ReadZipperOwned`] -pub struct ReadZipperWitness(pub(crate) Option>); +pub struct ReadZipperWitness(pub(crate) Option>, Option); crate::impl_name_only_debug!( impl core::fmt::Debug for ReadZipperWitness @@ -1798,12 +1798,10 @@ pub(crate) mod read_zipper_core { /// descended from a [ZipperHead], and thus the `root_node` field is `Owned` /// In the cases where this field is not used, it will be set to [usize::MAX] root_parent_key_start: usize, - /// A special-case to access a value at the root node, because that value would be otherwise inaccessible - /// NOTE: `root_val` will be `None`, except in situations where `ReadZipper` at the root of a map or `ZipperHead` + /// A value at the zipper root, including a value held by `root_node`. root_val: Option<&'a V>, - /// The [TrieNodeODRc] that contains the root node from which the zipper is descended. If the zipper - /// is descended from a `ZipperHead`, this field will be `Owned`, otherwise it will be `Borrowed` - root_node: OwnedOrBorrowed<'a, TrieNodeODRc>, + /// The node from which the zipper is descended. + root_node: OwnedOrBorrowed<'a, TrieNodeODRc, V>, /// A reference to the focus node focus_node: MiriWrapper>, /// An iter token corresponding to the location of the `node_key` within the `focus_node`, or NODE_ITER_INVALID @@ -1817,77 +1815,46 @@ pub(crate) mod read_zipper_core { pub(crate) alloc: A, } - //GOAT-TODO, we should unify this `OwnedOrBorrowed` type with [`AbstractNodeRef`], and it should be able to - // be packed into a single 64-bit word, and do the right thing when it is dropped. #[derive(Clone, Debug)] - pub enum OwnedOrBorrowed<'a, T> { + pub(crate) enum OwnedOrBorrowed<'a, T, V> { Owned(T), Borrowed(&'a T), None, + /// A private node and root value with stable addresses for the zipper's references. + Private(Box<(T, V)>), } - impl<'a, T> From> for OwnedOrBorrowed<'a, T> { - fn from(opt: Option<&'a T>) -> Self { - match opt { - Some(t) => Self::Borrowed(t), - None => Self::None - } - } - } - - impl<'a, T> OwnedOrBorrowed<'a, T> { - /// Returns a reference to the content, regardless of whether it is owned or borrowed + impl<'a, T, V> OwnedOrBorrowed<'a, T, V> { #[inline] - pub fn as_ref(&self) -> &T { + fn as_option(&self) -> Option<&T> { match self { - Self::Owned(t) => &t, - Self::Borrowed(t) => t, - Self::None => panic!(), + Self::Owned(value) => Some(value), + Self::Private(value) => Some(&value.0), + Self::Borrowed(value) => Some(value), + Self::None => None, } } - //GOAT, may be unneeded - // /// Returns a reference to the - // #[inline] - // pub fn as_option(&self) -> Option<&T> { - // match self { - // Self::Owned(t) => Some(&t), - // Self::Borrowed(t) => Some(t), - // Self::None => None, - // } - // } - /// Returns a reference to the content in the reference lifetime, if it's borrowed. Panics if the content is owned - pub fn as_borrowed_ref(&self) -> &'a T { + + #[inline] + fn as_ref(&self) -> &T { match self { - Self::Borrowed(t) => t, - Self::Owned(_) => panic!(), + Self::Borrowed(value) => value, + Self::Owned(value) => value, + Self::Private(value) => &value.0, Self::None => panic!(), } } - /// Returns a reference to the owned content, or `None` if the content is borrowed - pub fn get_owned_ref(&self) -> Option<&T> { - match self { - Self::Owned(t) => Some(&t), - Self::Borrowed(_) => None, - Self::None => None, - } - } - /// Returns `true` if the content is owned, or `false` otherwise - pub fn is_owned(&self) -> bool { + + #[inline] + fn as_borrowed_ref(&self) -> &'a T { match self { - Self::Owned(_) => true, - Self::Borrowed(_) => false, - Self::None => false, + Self::Borrowed(value) => value, + _ => panic!("root node is not borrowed"), } } - //GOAT, maybe unneeded - // /// Returns `true` if the content is borrowed, or `false` otherwise - // pub fn is_borrowed(&self) -> bool { - // match self { - // Self::Borrowed(_) => true, - // Self::Owned(_) => false, - // Self::None => false, - // } - // } + + #[inline] + fn is_owned(&self) -> bool { matches!(self, Self::Owned(_) | Self::Private(_)) } } #[cfg(miri)] @@ -1928,12 +1895,18 @@ pub(crate) mod read_zipper_core { impl Clone for ReadZipperCore<'_, '_, V, A> where V: Clone { fn clone(&self) -> Self { + let root_node = self.root_node.clone(); + // The cloned private root has a new address. Its reference is valid while the clone lives. + let root_val = match &root_node { + OwnedOrBorrowed::Private(private) => Some(unsafe { &*(&private.1 as *const V) }), + _ => self.root_val, + }; Self { origin_path: self.origin_path.clone(), root_key_start: self.root_key_start, root_parent_key_start: self.root_parent_key_start, - root_val: self.root_val, - root_node: self.root_node.clone(), + root_val, + root_node, focus_node: self.focus_node.clone(), focus_iter_token: NODE_ITER_INVALID, prefix_buf: self.prefix_buf.clone(), @@ -2581,13 +2554,14 @@ pub(crate) mod read_zipper_core { type WitnessT = ReadZipperWitness; fn witness<'w>(&self) -> Self::WitnessT { if self.root_node.is_owned() { - ReadZipperWitness(Some(self.root_node.as_ref().clone())) + ReadZipperWitness(Some(self.root_node.as_ref().clone()), self.root_val.cloned()) } else { - ReadZipperWitness(None) + ReadZipperWitness(None, None) } } fn get_val_with_witness<'w>(&self, witness: &'w Self::WitnessT) -> Option<&'w V> where 'trie: 'w { - assert_eq!(witness.0.as_ref(), self.root_node.get_owned_ref()); + assert_eq!(witness.0.as_ref(), self.root_node.is_owned().then(|| self.root_node.as_ref())); + assert_eq!(witness.1.is_some(), self.root_node.is_owned() && self.root_val.is_some()); let key = self.node_key(); if key.len() > 0 { self.focus_node.node_get_val(key) @@ -2597,7 +2571,7 @@ pub(crate) mod read_zipper_core { } else { if self.root_val.is_some() || self.root_parent_key_start == usize::MAX { //No parent key: the zipper root is the root node itself, and its value is `root_val` - self.root_val + if self.root_node.is_owned() { witness.1.as_ref() } else { self.root_val } } else { //We know the node in the witness and the node in self.root_node are the same, // but we borrow it from the witness here because that has the lifetime we need @@ -2637,11 +2611,7 @@ pub(crate) mod read_zipper_core { let (_key_len, focus_node) = parent.node_get_child(self.parent_key()).unwrap(); !focus_node.is_empty() && focus_node.refcount() > 1 } else { - let focus = match &self.root_node { - OwnedOrBorrowed::None => return false, - _ => self.focus_parent(), - }; - !focus.is_empty() && focus.refcount() > 1 + self.root_node.as_option().is_some_and(|root| !root.is_empty() && root.refcount() > 1) } } } @@ -2754,17 +2724,21 @@ pub(crate) mod read_zipper_core { /// /// NOTE: This method currently doesn't descend subnodes. Use [Self::new_with_node_and_path_in] if you can't /// guarantee the path is within the supplied node. - pub(crate) fn new_with_node_and_path_internal_in(root_node: OwnedOrBorrowed<'a, TrieNodeODRc>, path: &'path [u8], mut root_key_start: usize, root_val: Option<&'a V>, alloc: A) -> Self { + pub(crate) fn new_with_node_and_path_internal_in(root_node: OwnedOrBorrowed<'a, TrieNodeODRc, V>, path: &'path [u8], mut root_key_start: usize, root_val: Option<&'a V>, alloc: A) -> Self { let mut focus: TaggedNodeRef<'a, V, A> = match &root_node { - OwnedOrBorrowed::Owned(root_node) => { + OwnedOrBorrowed::Owned(node) => { // SAFETY: The root_node makes the ReadZipper essentially a self-referential type. As long // as the ReadZipperCore is alive, this will remain valid, and the "witness" mechanism in // `ZipperReadOnlyConditionalValues` makes sure the references returned from the `get_val` // method remain valid - unsafe{ core::mem::transmute(root_node.as_tagged()) } + unsafe{ core::mem::transmute(node.as_tagged()) } }, - OwnedOrBorrowed::Borrowed(root_node) => { - root_node.as_tagged() + OwnedOrBorrowed::Private(private) => { + // SAFETY: The boxed node keeps the address stable for the zipper's lifetime. + unsafe { core::mem::transmute(private.0.as_tagged()) } + }, + OwnedOrBorrowed::Borrowed(node) => { + node.as_tagged() }, OwnedOrBorrowed::None => { TaggedNodeRef::empty_node() @@ -2784,7 +2758,7 @@ pub(crate) mod read_zipper_core { origin_path: SliceOrLen::from(path), root_key_start, root_parent_key_start, - root_val: root_val.map(|v| v.into()), + root_val, focus_node: MiriWrapper::new(focus), root_node, focus_iter_token: NODE_ITER_INVALID, @@ -2813,56 +2787,34 @@ pub(crate) mod read_zipper_core { new_zipper.make_static_path() } - /// Like [Self::new_with_node_and_path_in] with an owned root, but never holding a node that a - /// live `ZipperHead` writer can reach. Sharing one lets the next exclusive writer copy it, - /// leaving the writers already made pointing into the copy that is about to be dropped + /// Own only the subtree at the reader's origin, so no live writer can reach its root node. pub(crate) fn new_isolated_in(root_node: &'a TrieNodeODRc, path: &'path [u8], root_val: Option<&'a V>, alloc: A) -> Self { //A reader at the head's own root excludes every writer, so it may hold the root node - let Some(&last) = path.last() else { + if path.is_empty() { return Self::new_with_node_and_path_in(root_node, true, path, 0, 0, root_val, alloc) - }; + } let (node, key, val) = node_along_path(root_node, path, root_val, false); - - //The focus is a whole node and carries no value of its own, so the zipper can own that - // node outright: no writer may be at or below the reader's path, so none is inside it + //A whole value-free node needs no private value or split. if key.is_empty() && val.is_none() { return Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Owned(node.clone()), path, path.len(), None, alloc) } - - Self::new_isolated_entry_in(node, key, val, path, last, alloc) - } - /// The uncommon half of [Self::new_isolated_in]: the entry at the focus lives in a node that - /// `ZipperHead` writers reach as well, because its value sits beside their branches or the - /// node is one they descend through. Copy just this entry into a private node. - #[inline(never)] - #[cold] - fn new_isolated_entry_in(node: &'a TrieNodeODRc, key: &[u8], val: Option<&'a V>, path: &'path [u8], last: u8, alloc: A) -> Self { - let (val, child, dangling) = if key.is_empty() { - //A value at the focus lives in the node above it, which is not ours to hold - let child = (!node.is_empty()).then(|| node.clone()); - (val.cloned(), child, false) + let (subtree, val, exists) = if key.is_empty() { + (Some(node.clone()), val.cloned(), true) } else { - let node = node.as_tagged(); - let val = node.node_get_val(key).cloned(); - let child = node.get_node_at_key(key).into_option(); - let dangling = val.is_none() && child.is_none() && node.node_contains_partial_key(key); - (val, child, dangling) + let tagged = node.as_tagged(); + (tagged.get_node_at_key(key).into_option(), tagged.node_get_val(key).cloned(), tagged.node_contains_partial_key(key)) }; - #[cfg(not(feature = "all_dense_nodes"))] - let mut root = TrieNodeODRc::new_in(crate::line_list_node::LineListNode::new_in(alloc.clone()), alloc.clone()); - #[cfg(feature = "all_dense_nodes")] - let mut root = TrieNodeODRc::new_in(crate::dense_byte_node::DenseByteNode::new_in(alloc.clone()), alloc.clone()); + //A missing path must retain a nonempty node key so `path_exists` remains false. + let root_key_start = if exists { path.len() } else { path.len() - 1 }; + let node = subtree.unwrap_or_else(TrieNodeODRc::new_empty); if let Some(val) = val { - if let Err(n) = root.make_mut().node_set_val(&[last], val) { root = n } - } - if let Some(child) = child { - if let Err(n) = root.make_mut().node_set_branch(&[last], child) { root = n } - } - if dangling { - if let Err(n) = root.make_mut().node_create_dangling(&[last]) { root = n } + let private = Box::new((node, val)); + // Box storage remains at the same address when the zipper moves. + let root_val = Some(unsafe { &*(&private.1 as *const V) }); + Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Private(private), path, root_key_start, root_val, alloc) + } else { + Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Owned(node), path, root_key_start, None, alloc) } - //The root value is read from `root`, via `root_parent_key_start` - Self::new_with_node_and_path_internal_in(OwnedOrBorrowed::Owned(root), path, path.len() - 1, None, alloc) } /// Same as [Self::new_isolated_in], but with a `'static` path pub(crate) fn new_isolated_cloned_path_in(root_node: &'a TrieNodeODRc, path: &[u8], root_val: Option<&'a V>, alloc: A) -> ReadZipperCore<'a, 'static, V, A> { @@ -2956,12 +2908,7 @@ pub(crate) mod read_zipper_core { // we currently share the same implementation between `val()` and `get_val()` because the only difference is the return // lifetime, and the current ZipperHead implementation is actually ok with referencing the value in the root of the ZipperHead. // debug_assert!(self.root_node.is_borrowed()); - if self.root_val.is_some() || self.root_parent_key_start == usize::MAX || !self.root_node.is_owned() { - self.root_val - } else { - //SAFETY: see the note on this method - self.root_node.as_ref().as_tagged().node_get_val(self.root_node_key()).map(|v| unsafe{ &*(v as *const V) }) - } + self.root_val } } } @@ -3090,10 +3037,6 @@ pub(crate) mod read_zipper_core { if parent_key.len() == 0 { return self.root_node.as_ref() } - if self.ancestors.is_empty() { - //At the root, with the focus on a child of the root node - return self.root_node.as_ref().as_tagged().node_get_child(parent_key).unwrap().1 - } self.focus_parent_borrowed() } @@ -3316,10 +3259,8 @@ pub(crate) mod read_zipper_core { } else { if let Some((parent, _iter_tok, _prefix_offset)) = self.ancestors.last() { parent.node_contains_val(self.parent_key()) - } else if self.root_val.is_some() || self.root_parent_key_start == usize::MAX || !self.root_node.is_owned() { - self.root_val.is_some() } else { - self.root_node.as_ref().as_tagged().node_contains_val(self.root_node_key()) + self.root_val.is_some() } } } @@ -3392,7 +3333,7 @@ pub(crate) mod read_zipper_core { fn root_node_key(&self) -> &[u8] { //This method should only be called when we have an owned `root_node`, which should go with // a valid value for `root_parent_key_start` - debug_assert!(matches!(self.root_node, OwnedOrBorrowed::Owned(_))); + debug_assert!(self.root_node.is_owned()); debug_assert!(self.root_parent_key_start < usize::MAX); if self.prefix_buf.capacity() == 0 && self.origin_path.len() > 0 { unsafe{ &self.origin_path.as_slice_unchecked()[self.root_parent_key_start..] } @@ -3406,8 +3347,6 @@ pub(crate) mod read_zipper_core { if self.prefix_buf.len() > 0 { let key_start = if self.ancestors.len() > 1 { unsafe{ self.ancestors.get_unchecked(self.ancestors.len()-2) }.2 - } else if self.ancestors.is_empty() && self.root_parent_key_start != usize::MAX { - self.root_parent_key_start } else { self.root_key_start };