From nobody Sat Nov 9 16:56:51 2024 Received: from mail-qt1-f202.google.com (mail-qt1-f202.google.com [209.85.160.202]) (using TLSv1.2 with cipher ECDHE-RSA-AES128-GCM-SHA256 (128/128 bits)) (No client certificate requested) by smtp.subspace.kernel.org (Postfix) with ESMTPS id 406811CDFA9 for ; Thu, 22 Aug 2024 16:38:05 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.160.202 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1724344687; cv=none; b=MNYK0ZPRMPN/xXp6DvzS1PVcoUb2mR0c+BZFrZefVmlDgpu9OL0hiE+1JRxT7a9ZJxoCMDkWWgRE2KMmmpAhkYBwvXL1gdXxUStzj8UBA30XnT2opvEDU5ndcthdu/glNc5SP0VLNjbS5B67+q2vtWgKk7o/5ayOuZtMja7xv1U= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1724344687; c=relaxed/simple; bh=TRJjiULNayI0Bn/88ItSNVr5LOGlugJqIZvCXlzvGYA=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=ASobQbhb9aXN+XNYU3ZrhPg/9XftCB+Q4dgFX6miI23WQw1U+ECcaMau8SwSx3G7UDX6cjIzuHjzIr6ovNk6LX4vZkjCVk+G2EiPCz/oK86QsrAHoXrTyAuK3sIhCpcx7jRcVjzQTJPLvTfXP20+DEykp1F0NHPWNAdd8H2RRLY= ARC-Authentication-Results: i=1; smtp.subspace.kernel.org; dmarc=pass (p=reject dis=none) header.from=google.com; spf=pass smtp.mailfrom=flex--mattgilbride.bounces.google.com; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b=T/PxTLjR; arc=none smtp.client-ip=209.85.160.202 Authentication-Results: smtp.subspace.kernel.org; dmarc=pass (p=reject dis=none) header.from=google.com Authentication-Results: smtp.subspace.kernel.org; spf=pass smtp.mailfrom=flex--mattgilbride.bounces.google.com Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="T/PxTLjR" Received: by mail-qt1-f202.google.com with SMTP id d75a77b69052e-44ffa2bf13fso12689001cf.3 for ; Thu, 22 Aug 2024 09:38:05 -0700 (PDT) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1724344684; x=1724949484; darn=vger.kernel.org; h=cc:to:from:subject:message-id:references:mime-version:in-reply-to :date:from:to:cc:subject:date:message-id:reply-to; bh=qzLOCRmg4omxdJQNiqf0cVUIyGdYk0mbYUu6qYq7nkM=; b=T/PxTLjRV0pofIWG6e2ruQft8Ol727ehRg5PvwPQnWWlIsVFsRBvFfPvZHRICivVpZ uS/6bzhvkwjXhYbtIFLSLYuJ8uv8rRoGTyWkR9OnuTsYimSHvFGwSV9l801wKSOZ2o0d UFZW8QU2qzvN36ivhgiN654s35ZjTcd7XE6jkz/m4TZhiZ2fLruidJR+GF5c+mRUAbj3 pTKZ5cYg+f2a+k8y/g9gIMhP0GTDaSJL7XyXZ/UB3yoDrukjrsxZVaDoHAOGKsfUNjmU RnFpbNuRT6dklo3kE3vRpSKZA0MEC+fAMyGj5x0srExBgjvaLQCuKrP/8hj1pr9BWCG0 4flg== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1724344684; x=1724949484; h=cc:to:from:subject:message-id:references:mime-version:in-reply-to :date:x-gm-message-state:from:to:cc:subject:date:message-id:reply-to; bh=qzLOCRmg4omxdJQNiqf0cVUIyGdYk0mbYUu6qYq7nkM=; b=P60suW5AycrTDdIoLdgm5uUK0lQdKJP1YjkXWulgpSBV8MLKUJi/rWsPfgubvClUdC Geve8gCll8ajSD9GZb91B1mPDOg9HqvyW5gZUkOrrSqisOiuLWpyoVFEebXlkI58sb5R +3brRZYHX2f3Ak5+IgknYz2t0tTBmQm0Df+wfBJqJWD9Pof2ankUYPtzVon2M415cTyM RcsAIdcPGE5yN8VDZsZv9CtdNCTW7juxRNcfVRfGxJmpn6WaZh+m5ic00e/ERy/1SelB LCSnvXxnT2DuNYIEytezw64hR8qPvV9UWlFZcnEwUidmecCxDk8U5BAhF7Iz3txqJkD5 muPw== X-Forwarded-Encrypted: i=1; AJvYcCXjUCICHxfVTQeZmfcoI5WTf2nhQYbLO16mFdpRBryvmGxvQ+MBBXjakgtt9ba5OLavELJuPWkYxr3fzSw=@vger.kernel.org X-Gm-Message-State: AOJu0YwMmsuP2H60ZWIiek0OAqd199i8WUKQuDPckM9/ZIWNaMKeiucq jMadrvTHuMttTUVFpGLvNQNpxupymnltnt28TIJSICR0yaTk6hDV52sXgSajSKwwAjcFnLeAFQo 10bZKNZxmSGg70nekGwMy0HBrww== X-Google-Smtp-Source: AGHT+IHbgYpoW9oiyBHrrQSpK+91uMQ8QlYAO9qc6eXT0moqWV71r/XIQiCq3B7ZC2BVEaIYJYKikQ96xY1DaJuy0xU= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:ac8:109a:0:b0:44f:e134:3430 with SMTP id d75a77b69052e-454f21f27e2mr310071cf.4.1724344684179; Thu, 22 Aug 2024 09:38:04 -0700 (PDT) Date: Thu, 22 Aug 2024 16:37:56 +0000 In-Reply-To: <20240822-b4-rbtree-v12-0-014561758a57@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240822-b4-rbtree-v12-0-014561758a57@google.com> X-Mailer: b4 0.13.0 Message-ID: <20240822-b4-rbtree-v12-4-014561758a57@google.com> Subject: [PATCH v12 4/5] rust: rbtree: add cursor From: Matt Gilbride To: Miguel Ojeda , Alex Gaynor , Wedson Almeida Filho , Boqun Feng , Gary Guo , "=?utf-8?q?Bj=C3=B6rn_Roy_Baron?=" , Benno Lossin , Andreas Hindborg , Alice Ryhl , Greg Kroah-Hartman , "=?utf-8?q?Arve_Hj=C3=B8nnev=C3=A5g?=" , Todd Kjos , Martijn Coenen , Joel Fernandes , Carlos Llamas , Suren Baghdasaryan , Christian Brauner Cc: Rob Landley , Davidlohr Bueso , Michel Lespinasse , rust-for-linux@vger.kernel.org, linux-kernel@vger.kernel.org, Matt Gilbride Content-Type: text/plain; charset="utf-8" Content-Transfer-Encoding: quoted-printable Add a cursor interface to `RBTree`, supporting the following use cases: - Inspect the current node pointed to by the cursor, inspect/move to it's neighbors in sort order (bidirectionally). - Mutate the tree itself by removing the current node pointed to by the cursor, or one of its neighbors. Add functions to obtain a cursor to the tree by key: - The node with the smallest key - The node with the largest key - The node matching the given key, or the one with the next larger key The cursor abstraction is needed by the binder driver to efficiently search for nodes and (conditionally) modify them, as well as their neighbors [1]. Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba= 9197f637@google.com/ [1] Co-developed-by: Alice Ryhl Signed-off-by: Alice Ryhl Tested-by: Alice Ryhl Reviewed-by: Boqun Feng Reviewed-by: Benno Lossin Signed-off-by: Matt Gilbride --- rust/kernel/rbtree.rs | 523 ++++++++++++++++++++++++++++++++++++++++++++++= ++++ 1 file changed, 523 insertions(+) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index 399ace4e1949..fcf0d756c5be 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -236,6 +236,36 @@ pub fn values(&self) -> impl Iterator { pub fn values_mut(&mut self) -> impl Iterator { self.iter_mut().map(|(_, v)| v) } + + /// Returns a cursor over the tree nodes, starting with the smallest k= ey. + pub fn cursor_front(&mut self) -> Option> { + let root =3D addr_of_mut!(self.root); + // SAFETY: `self.root` is always a valid root node + let current =3D unsafe { bindings::rb_first(root) }; + NonNull::new(current).map(|current| { + // INVARIANT: + // - `current` is a valid node in the [`RBTree`] pointed to by= `self`. + Cursor { + current, + tree: self, + } + }) + } + + /// Returns a cursor over the tree nodes, starting with the largest ke= y. + pub fn cursor_back(&mut self) -> Option> { + let root =3D addr_of_mut!(self.root); + // SAFETY: `self.root` is always a valid root node + let current =3D unsafe { bindings::rb_last(root) }; + NonNull::new(current).map(|current| { + // INVARIANT: + // - `current` is a valid node in the [`RBTree`] pointed to by= `self`. + Cursor { + current, + tree: self, + } + }) + } } =20 impl RBTree @@ -396,6 +426,67 @@ fn remove_node(&mut self, key: &K) -> Option> { pub fn remove(&mut self, key: &K) -> Option { self.remove_node(key).map(|node| node.node.value) } + + /// Returns a cursor over the tree nodes based on the given key. + /// + /// If the given key exists, the cursor starts there. + /// Otherwise it starts with the first larger key in sort order. + /// If there is no larger key, it returns [`None`]. + pub fn cursor_lower_bound(&mut self, key: &K) -> Option> + where + K: Ord, + { + let mut node =3D self.root.rb_node; + let mut best_match: Option>> =3D None; + while !node.is_null() { + // SAFETY: By the type invariant of `Self`, all non-null `rb_n= ode` pointers stored in `self` + // point to the links field of `Node` objects. + let this =3D unsafe { container_of!(node, Node, links) }= .cast_mut(); + // SAFETY: `this` is a non-null node so it is valid by the typ= e invariants. + let this_key =3D unsafe { &(*this).key }; + // SAFETY: `node` is a non-null node so it is valid by the typ= e invariants. + let left_child =3D unsafe { (*node).rb_left }; + // SAFETY: `node` is a non-null node so it is valid by the typ= e invariants. + let right_child =3D unsafe { (*node).rb_right }; + match key.cmp(this_key) { + Ordering::Equal =3D> { + best_match =3D NonNull::new(this); + break; + } + Ordering::Greater =3D> { + node =3D right_child; + } + Ordering::Less =3D> { + let is_better_match =3D match best_match { + None =3D> true, + Some(best) =3D> { + // SAFETY: `best` is a non-null node so it is = valid by the type invariants. + let best_key =3D unsafe { &(*best.as_ptr()).ke= y }; + best_key > this_key + } + }; + if is_better_match { + best_match =3D NonNull::new(this); + } + node =3D left_child; + } + }; + } + + let best =3D best_match?; + + // SAFETY: `best` is a non-null node so it is valid by the type in= variants. + let links =3D unsafe { addr_of_mut!((*best.as_ptr()).links) }; + + NonNull::new(links).map(|current| { + // INVARIANT: + // - `current` is a valid node in the [`RBTree`] pointed to by= `self`. + Cursor { + current, + tree: self, + } + }) + } } =20 impl Default for RBTree { @@ -427,6 +518,433 @@ fn drop(&mut self) { } } =20 +/// A bidirectional cursor over the tree nodes, sorted by key. +/// +/// # Examples +/// +/// In the following example, we obtain a cursor to the first element in t= he tree. +/// The cursor allows us to iterate bidirectionally over key/value pairs i= n the tree. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // Get a cursor to the first element. +/// let mut cursor =3D tree.cursor_front().unwrap(); +/// let mut current =3D cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// +/// // Move the cursor, updating it to the 2nd element. +/// cursor =3D cursor.move_next().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Peek at the next element without impacting the cursor. +/// let next =3D cursor.peek_next().unwrap(); +/// assert_eq!(next, (&30, &300)); +/// current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Moving past the last element causes the cursor to return [`None`]. +/// cursor =3D cursor.move_next().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// let cursor =3D cursor.move_next(); +/// assert!(cursor.is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// A cursor can also be obtained at the last element in the tree. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// let mut cursor =3D tree.cursor_back().unwrap(); +/// let current =3D cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// Obtaining a cursor returns [`None`] if the tree is empty. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// let mut tree: RBTree =3D RBTree::new(); +/// assert!(tree.cursor_front().is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary no= de in the tree. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert five elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(40, 400, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(50, 500, flags::GFP_KERNEL)?; +/// +/// // If the provided key exists, a cursor to that key is returned. +/// let cursor =3D tree.cursor_lower_bound(&20).unwrap(); +/// let current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // If the provided key doesn't exist, a cursor to the first larger ele= ment in sort order is returned. +/// let cursor =3D tree.cursor_lower_bound(&25).unwrap(); +/// let current =3D cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// // If there is no larger key, [`None`] is returned. +/// let cursor =3D tree.cursor_lower_bound(&55); +/// assert!(cursor.is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// The cursor allows mutation of values in the tree. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // Retrieve a cursor. +/// let mut cursor =3D tree.cursor_front().unwrap(); +/// +/// // Get a mutable reference to the current value. +/// let (k, v) =3D cursor.current_mut(); +/// *v =3D 1000; +/// +/// // The updated value is reflected in the tree. +/// let updated =3D tree.get(&10).unwrap(); +/// assert_eq!(updated, &1000); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// It also allows node removal. The following examples demonstrate the be= havior of removing the current node. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // Remove the first element. +/// let mut cursor =3D tree.cursor_front().unwrap(); +/// let mut current =3D cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// cursor =3D cursor.remove_current().0.unwrap(); +/// +/// // If a node exists after the current element, it is returned. +/// current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Get a cursor to the last element, and remove it. +/// cursor =3D tree.cursor_back().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// // Since there is no next node, the previous node is returned. +/// cursor =3D cursor.remove_current().0.unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Removing the last element in the tree returns [`None`]. +/// assert!(cursor.remove_current().0.is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// Nodes adjacent to the current node can also be removed. +/// +/// ``` +/// use kernel::{alloc::flags, rbtree::RBTree}; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // Get a cursor to the first element. +/// let mut cursor =3D tree.cursor_front().unwrap(); +/// let mut current =3D cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// +/// // Calling `remove_prev` from the first element returns [`None`]. +/// assert!(cursor.remove_prev().is_none()); +/// +/// // Get a cursor to the last element. +/// cursor =3D tree.cursor_back().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// // Calling `remove_prev` removes and returns the middle element. +/// assert_eq!(cursor.remove_prev().unwrap().to_key_value(), (20, 200)); +/// +/// // Calling `remove_next` from the last element returns [`None`]. +/// assert!(cursor.remove_next().is_none()); +/// +/// // Move to the first element +/// cursor =3D cursor.move_prev().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// +/// // Calling `remove_next` removes and returns the last element. +/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300)); +/// +/// # Ok::<(), Error>(()) +/// +/// ``` +/// +/// # Invariants +/// - `current` points to a node that is in the same [`RBTree`] as `tree`. +pub struct Cursor<'a, K, V> { + tree: &'a mut RBTree, + current: NonNull, +} + +// SAFETY: The [`Cursor`] has exclusive access to both `K` and `V`, so it = is sufficient to require them to be `Send`. +// The cursor only gives out immutable references to the keys, but since i= t has excusive access to those same +// keys, `Send` is sufficient. `Sync` would be okay, but it is more restri= ctive to the user. +unsafe impl<'a, K: Send, V: Send> Send for Cursor<'a, K, V> {} + +// SAFETY: The [`Cursor`] gives out immutable references to K and mutable = references to V, +// so it has the same thread safety requirements as mutable references. +unsafe impl<'a, K: Sync, V: Sync> Sync for Cursor<'a, K, V> {} + +impl<'a, K, V> Cursor<'a, K, V> { + /// The current node + pub fn current(&self) -> (&K, &V) { + // SAFETY: + // - `self.current` is a valid node by the type invariants. + // - We have an immutable reference by the function signature. + unsafe { Self::to_key_value(self.current) } + } + + /// The current node, with a mutable value + pub fn current_mut(&mut self) -> (&K, &mut V) { + // SAFETY: + // - `self.current` is a valid node by the type invariants. + // - We have an mutable reference by the function signature. + unsafe { Self::to_key_value_mut(self.current) } + } + + /// Remove the current node from the tree. + /// + /// Returns a tuple where the first element is a cursor to the next no= de, if it exists, + /// else the previous node, else [`None`] (if the tree becomes empty).= The second element + /// is the removed node. + pub fn remove_current(self) -> (Option, RBTreeNode) { + let prev =3D self.get_neighbor_raw(Direction::Prev); + let next =3D self.get_neighbor_raw(Direction::Next); + // SAFETY: By the type invariant of `Self`, all non-null `rb_node`= pointers stored in `self` + // point to the links field of `Node` objects. + let this =3D unsafe { container_of!(self.current.as_ptr(), Node, links) }.cast_mut(); + // SAFETY: `this` is valid by the type invariants as described abo= ve. + let node =3D unsafe { Box::from_raw(this) }; + let node =3D RBTreeNode { node }; + // SAFETY: The reference to the tree used to create the cursor out= lives the cursor, so + // the tree cannot change. By the tree invariant, all nodes are va= lid. + unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.= tree.root)) }; + + let current =3D match (prev, next) { + (_, Some(next)) =3D> next, + (Some(prev), None) =3D> prev, + (None, None) =3D> { + return (None, node); + } + }; + + ( + // INVARIANT: + // - `current` is a valid node in the [`RBTree`] pointed to by= `self.tree`. + Some(Self { + current, + tree: self.tree, + }), + node, + ) + } + + /// Remove the previous node, returning it if it exists. + pub fn remove_prev(&mut self) -> Option> { + self.remove_neighbor(Direction::Prev) + } + + /// Remove the next node, returning it if it exists. + pub fn remove_next(&mut self) -> Option> { + self.remove_neighbor(Direction::Next) + } + + fn remove_neighbor(&mut self, direction: Direction) -> Option> { + if let Some(neighbor) =3D self.get_neighbor_raw(direction) { + let neighbor =3D neighbor.as_ptr(); + // SAFETY: The reference to the tree used to create the cursor= outlives the cursor, so + // the tree cannot change. By the tree invariant, all nodes ar= e valid. + unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.r= oot)) }; + // SAFETY: By the type invariant of `Self`, all non-null `rb_n= ode` pointers stored in `self` + // point to the links field of `Node` objects. + let this =3D unsafe { container_of!(neighbor, Node, link= s) }.cast_mut(); + // SAFETY: `this` is valid by the type invariants as described= above. + let node =3D unsafe { Box::from_raw(this) }; + return Some(RBTreeNode { node }); + } + None + } + + /// Move the cursor to the previous node, returning [`None`] if it doe= sn't exist. + pub fn move_prev(self) -> Option { + self.mv(Direction::Prev) + } + + /// Move the cursor to the next node, returning [`None`] if it doesn't= exist. + pub fn move_next(self) -> Option { + self.mv(Direction::Next) + } + + fn mv(self, direction: Direction) -> Option { + // INVARIANT: + // - `neighbor` is a valid node in the [`RBTree`] pointed to by `s= elf.tree`. + self.get_neighbor_raw(direction).map(|neighbor| Self { + tree: self.tree, + current: neighbor, + }) + } + + /// Access the previous node without moving the cursor. + pub fn peek_prev(&self) -> Option<(&K, &V)> { + self.peek(Direction::Prev) + } + + /// Access the previous node without moving the cursor. + pub fn peek_next(&self) -> Option<(&K, &V)> { + self.peek(Direction::Next) + } + + fn peek(&self, direction: Direction) -> Option<(&K, &V)> { + self.get_neighbor_raw(direction).map(|neighbor| { + // SAFETY: + // - `neighbor` is a valid tree node. + // - By the function signature, we have an immutable reference= to `self`. + unsafe { Self::to_key_value(neighbor) } + }) + } + + /// Access the previous node mutably without moving the cursor. + pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> { + self.peek_mut(Direction::Prev) + } + + /// Access the next node mutably without moving the cursor. + pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> { + self.peek_mut(Direction::Next) + } + + fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> { + self.get_neighbor_raw(direction).map(|neighbor| { + // SAFETY: + // - `neighbor` is a valid tree node. + // - By the function signature, we have a mutable reference to= `self`. + unsafe { Self::to_key_value_mut(neighbor) } + }) + } + + fn get_neighbor_raw(&self, direction: Direction) -> Option> { + // SAFETY: `self.current` is valid by the type invariants. + let neighbor =3D unsafe { + match direction { + Direction::Prev =3D> bindings::rb_prev(self.current.as_ptr= ()), + Direction::Next =3D> bindings::rb_next(self.current.as_ptr= ()), + } + }; + + NonNull::new(neighbor) + } + + /// SAFETY: + /// - `node` must be a valid pointer to a node in an [`RBTree`]. + /// - The caller has immutable access to `node` for the duration of 'b. + unsafe fn to_key_value<'b>(node: NonNull) -> (&'b K= , &'b V) { + // SAFETY: the caller guarantees that `node` is a valid pointer in= an `RBTree`. + let (k, v) =3D unsafe { Self::to_key_value_raw(node) }; + // SAFETY: the caller guarantees immutable access to `node`. + (k, unsafe { &*v }) + } + + /// SAFETY: + /// - `node` must be a valid pointer to a node in an [`RBTree`]. + /// - The caller has mutable access to `node` for the duration of 'b. + unsafe fn to_key_value_mut<'b>(node: NonNull) -> (&= 'b K, &'b mut V) { + // SAFETY: the caller guarantees that `node` is a valid pointer in= an `RBTree`. + let (k, v) =3D unsafe { Self::to_key_value_raw(node) }; + // SAFETY: the caller guarantees mutable access to `node`. + (k, unsafe { &mut *v }) + } + + /// SAFETY: + /// - `node` must be a valid pointer to a node in an [`RBTree`]. + /// - The caller has immutable access to the key for the duration of '= b. + unsafe fn to_key_value_raw<'b>(node: NonNull) -> (&= 'b K, *mut V) { + // SAFETY: By the type invariant of `Self`, all non-null `rb_node`= pointers stored in `self` + // point to the links field of `Node` objects. + let this =3D unsafe { container_of!(node.as_ptr(), Node, lin= ks) }.cast_mut(); + // SAFETY: The passed `node` is the current node or a non-null nei= ghbor, + // thus `this` is valid by the type invariants. + let k =3D unsafe { &(*this).key }; + // SAFETY: The passed `node` is the current node or a non-null nei= ghbor, + // thus `this` is valid by the type invariants. + let v =3D unsafe { addr_of_mut!((*this).value) }; + (k, v) + } +} + +/// Direction for [`Cursor`] operations. +enum Direction { + /// the node immediately before, in sort order + Prev, + /// the node immediately after, in sort order + Next, +} + impl<'a, K, V> IntoIterator for &'a RBTree { type Item =3D (&'a K, &'a V); type IntoIter =3D Iter<'a, K, V>; @@ -584,6 +1102,11 @@ impl RBTreeNode { pub fn new(key: K, value: V, flags: Flags) -> Result>= { Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value)) } + + /// Get the key and value from inside the node. + pub fn to_key_value(self) -> (K, V) { + (self.node.key, self.node.value) + } } =20 // SAFETY: If K and V can be sent across threads, then it's also okay to s= end [`RBTreeNode`] across --=20 2.46.0.184.g6999bdac58-goog