From b80a70bb2d48d1918373c2e49ac7fd50f22f5896 Mon Sep 17 00:00:00 2001 From: Zixuan Chen Date: Thu, 27 Oct 2022 13:05:52 +0800 Subject: [PATCH] fix: reduce unsafe code --- .../src/container/text/tracker/cursor_map.rs | 18 +- .../src/container/text/tracker/yata_impl.rs | 2 +- crates/rle/src/rle_tree.rs | 173 ++++++++---------- crates/rle/src/rle_tree/iter.rs | 4 +- 4 files changed, 88 insertions(+), 109 deletions(-) diff --git a/crates/loro-core/src/container/text/tracker/cursor_map.rs b/crates/loro-core/src/container/text/tracker/cursor_map.rs index c948ffd4..6e298f8e 100644 --- a/crates/loro-core/src/container/text/tracker/cursor_map.rs +++ b/crates/loro-core/src/container/text/tracker/cursor_map.rs @@ -41,7 +41,10 @@ impl ZeroElement for Marker { } impl Marker { - pub fn as_cursor(&self, id: ID) -> Option> { + pub(super) fn as_cursor<'a, 'b>( + &'a self, + id: ID, + ) -> Option> { match self { Marker::Insert { ptr, len: _ } => { // SAFETY: tree data is always valid @@ -50,16 +53,17 @@ impl Marker { let child = &node.children()[position]; let start_counter = child.id.counter; let offset = id.counter - start_counter; - // SAFETY: we just checked it is valid - Some(unsafe { - SafeCursor::new( - *ptr, + // SAFETY: we transform lifetime from SafeCursor<'static> to SafeCursor<'b> to suit the need. + // Its safety is guaranteed by the caller, who has access to the underlying tree + unsafe { + std::mem::transmute(Some(SafeCursor::from_leaf( + node, position, offset as usize, Position::from_offset(offset as isize, child.atom_len()), 0, - ) - }) + ))) + } } Marker::Delete(_) => None, } diff --git a/crates/loro-core/src/container/text/tracker/yata_impl.rs b/crates/loro-core/src/container/text/tracker/yata_impl.rs index 308b2d2c..b890aa30 100644 --- a/crates/loro-core/src/container/text/tracker/yata_impl.rs +++ b/crates/loro-core/src/container/text/tracker/yata_impl.rs @@ -79,7 +79,7 @@ impl ListCrdt for YataImpl { // dbg!(&container.content); // SAFETY: loosen lifetime requirement here. It's safe because the function // signature can limit the lifetime of the returned iterator - unsafe { std::mem::transmute(container.content.iter_mut_in(from, to)) } + container.content.iter_mut_in(from, to) } fn id(op: &Self::OpUnit) -> Self::OpId { diff --git a/crates/rle/src/rle_tree.rs b/crates/rle/src/rle_tree.rs index db2d390d..63cbcbbc 100644 --- a/crates/rle/src/rle_tree.rs +++ b/crates/rle/src/rle_tree.rs @@ -38,6 +38,16 @@ impl + 'static> Default for RleTree { } impl> RleTree { + fn root(&self) -> &Node { + // SAFETY: self can be shared ref so the root node must be valid and can be shared ref + self.with_node(|node| unsafe { std::mem::transmute::<_, &Node>(&**node) }) + } + + fn root_mut(&mut self) -> &mut Node { + // SAFETY: self can be exclusively ref so the root node must be valid and can be exclusively ref + self.with_node_mut(|node| unsafe { std::mem::transmute::<_, &mut Node>(&mut **node) }) + } + pub fn insert_at_first(&mut self, value: T, notify: &mut F) where F: FnMut(&T, *mut LeafNode<'_, T, A>), @@ -84,76 +94,68 @@ impl> RleTree { /// return a cursor at the given index #[inline] pub fn get(&self, mut index: A::Int) -> Option> { - self.with_node(|mut node| { - loop { - match node { - Node::Internal(internal_node) => { - let result = A::find_pos_internal(internal_node, index); - if !result.found { - return None; - } - - node = &internal_node.children[result.child_index]; - index = result.offset; + let mut node = self.root(); + loop { + match node { + Node::Internal(internal_node) => { + let result = A::find_pos_internal(internal_node, index); + if !result.found { + return None; } - Node::Leaf(leaf) => { - let result = A::find_pos_leaf(leaf, index); - if !result.found { - return None; - } - // SAFETY: result is valid - return Some(unsafe { - std::mem::transmute(SafeCursor::from_leaf( - leaf, - result.child_index, - result.offset, - result.pos, - 0, - )) - }); + node = internal_node.children[result.child_index]; + index = result.offset; + } + Node::Leaf(leaf) => { + let result = A::find_pos_leaf(leaf, index); + if !result.found { + return None; } + + return Some(SafeCursor::from_leaf( + leaf, + result.child_index, + result.offset, + result.pos, + 0, + )); } } - }) + } } /// return the first valid cursor after the given index /// reviewed by @Leeeon233 #[inline] fn get_cursor_ge(&self, mut index: A::Int) -> Option> { - self.with_node(|mut node| { - loop { - match node { - Node::Internal(internal_node) => { - let result = A::find_pos_internal(internal_node, index); - if result.child_index >= internal_node.children.len() { - return None; - } - - node = &internal_node.children[result.child_index]; - index = result.offset; + let mut node = self.root(); + loop { + match node { + Node::Internal(internal_node) => { + let result = A::find_pos_internal(internal_node, index); + if result.child_index >= internal_node.children.len() { + return None; } - Node::Leaf(leaf) => { - let result = A::find_pos_leaf(leaf, index); - if result.child_index >= leaf.children.len() { - return None; - } - // SAFETY: result is valid - return Some(unsafe { - std::mem::transmute(SafeCursor::new( - leaf.into(), - result.child_index, - result.offset, - result.pos, - 0, - )) - }); + node = internal_node.children[result.child_index]; + index = result.offset; + } + Node::Leaf(leaf) => { + let result = A::find_pos_leaf(leaf, index); + if result.child_index >= leaf.children.len() { + return None; } + + return Some(SafeCursor::from_leaf( + leaf, + result.child_index, + result.offset, + result.pos, + 0, + )); } } - }) + } } #[inline] @@ -165,18 +167,13 @@ impl> RleTree { #[inline] pub fn iter(&self) -> iter::Iter<'_, T, A> { - // SAFETY: the cursor and iter cannot outlive self - self.with_node(|node| unsafe { - iter::Iter::new(std::mem::transmute(node.get_first_leaf())) - }) + iter::Iter::new(self.root().get_first_leaf()) } #[inline] pub fn iter_mut(&mut self) -> iter::IterMut<'_, T, A> { // SAFETY: the cursor and iter cannot outlive self - self.with_node_mut(|node| unsafe { - iter::IterMut::new(std::mem::transmute(node.get_first_leaf_mut())) - }) + iter::IterMut::new(self.root_mut().get_first_leaf_mut()) } #[inline] @@ -184,27 +181,25 @@ impl> RleTree { self.len() == A::Int::from_usize(0).unwrap() } - pub fn iter_mut_in( - &mut self, - start: Option>, - end: Option>, - ) -> iter::IterMut<'_, T, A> { + pub fn iter_mut_in<'a>( + &'a mut self, + start: Option>, + end: Option>, + ) -> iter::IterMut<'a, T, A> { if start.is_none() && end.is_none() { self.iter_mut() } else { - // SAFETY: the cursor cannot outlive self, so we are safe here - self.with_node_mut(|node| unsafe { - let leaf = node.get_first_leaf().unwrap(); - // SAFETY: this is safe because we know there are at least one element in the tree - let start = start.unwrap_or_else(|| { - std::mem::transmute(SafeCursor::from_leaf(leaf, 0, 0, Position::Start, 0)) - }); - let start: SafeCursorMut<'_, T, A> = SafeCursorMut::from(start.0); - std::mem::transmute::<_, iter::IterMut<'_, T, A>>(iter::IterMut::from_cursor( - std::mem::transmute::<_, SafeCursorMut<'_, T, A>>(start), - end, - )) - }) + let leaf = self.root().get_first_leaf().unwrap(); + // SAFETY: this is safe because we know there are at least one element in the tree + let start = + start.unwrap_or_else(|| SafeCursor::from_leaf(leaf, 0, 0, Position::Start, 0)); + + // SAFETY: we have exclusive ref to the tree, so it's safe to have an exclusive ref to its elements + let start: SafeCursorMut<'a, T, A> = unsafe { SafeCursorMut::from(start.0) }; + iter::IterMut::from_cursor( + start, + end.map(|x| UnsafeCursor::new(x.0.leaf, x.0.index, x.0.offset, x.0.pos, 0)), + ) } } @@ -373,26 +368,6 @@ impl> RleTree { } } - pub fn update_range( - &mut self, - start: A::Int, - end: Option, - update_fn: &mut U, - notify: &mut F, - ) where - U: FnMut(&mut T), - F: FnMut(&T, *mut LeafNode<'_, T, A>), - { - let mut cursors = Vec::new(); - for cursor in self.iter_range(start, end) { - cursors.push(cursor.0); - } - - // SAFETY: it's perfectly safe here because we know what we are doing in the update_at_cursors - let mut cursors: Vec<_> = unsafe { std::mem::transmute(cursors) }; - self.update_at_cursors(&mut cursors, update_fn, notify); - } - pub fn debug_check(&mut self) { self.with_node_mut(|node| { node.as_internal_mut().unwrap().check(); diff --git a/crates/rle/src/rle_tree/iter.rs b/crates/rle/src/rle_tree/iter.rs index d8af3166..e5119ba9 100644 --- a/crates/rle/src/rle_tree/iter.rs +++ b/crates/rle/src/rle_tree/iter.rs @@ -55,7 +55,7 @@ impl<'tree, T: Rle, A: RleTreeTrait> IterMut<'tree, T, A> { #[inline] pub fn from_cursor( mut start: SafeCursorMut<'tree, T, A>, - end: Option>, + end: Option>, ) -> Self { if start.0.pos == Position::After { match start.next_elem_start() { @@ -77,7 +77,7 @@ impl<'tree, T: Rle, A: RleTreeTrait> IterMut<'tree, T, A> { start.0.pos, 0, )), - end_cursor: end.map(|x| UnsafeCursor::new(x.0.leaf, x.0.index, x.0.offset, x.0.pos, 0)), + end_cursor: end, } } }