From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yw1-f201.google.com (mail-yw1-f201.google.com [209.85.128.201]) (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 6BB8E3FB21 for ; Mon, 5 Feb 2024 15:50:06 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.128.201 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148208; cv=none; b=CFTY0AUb2R924kFv5XCQgvlwCAd/mZsvwdumeIcxcpyLQFPz2alzYAoPmJ0QvVOTSASC5riJ6/pzUe5MpfLyR8R00lV5yG7poIOaDZYlPvOupY7Jm1erRLCyUWz9Hfy05XMPDdawyCUW/hkz2HKWFi1vuFe4o72ZUTEizudBz0M= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148208; c=relaxed/simple; bh=FR4T3Qjs2jKRsHMGwYv7pMd/5Tm8HyvFjH7AWhAgulg=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=iKCg1bv/ONpfwCYJp3k6BvipMQS/RlWey77rlBqYdeLVLM5w/XUpHvJ3CqtJY6Vx0sS9fLBYzDtjN5ByZDvjGVdGzWNXcqDIY0tyRvRIdrPi8Em6t3UgA5brDB2TQTIRti89zstw7s9IRWp91dQwebaxRza4+zsCrIGibyIi/es= 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=wPZ8etjv; arc=none smtp.client-ip=209.85.128.201 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="wPZ8etjv" Received: by mail-yw1-f201.google.com with SMTP id 00721157ae682-60404b12af2so72290047b3.2 for ; Mon, 05 Feb 2024 07:50:06 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148205; x=1707753005; 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=KixoytXQDGjoXAfvxRuyy14dtLsh1IkvCMs1kLMTxgU=; b=wPZ8etjvu/a0We6zybD5oAoPX7szcBceYVUEtMiWa3ZZrix76d58rDCkj7FojQPqdH ZmXPORLgWNqQWFWzdBV107HrPBpy1c6A1q5DQ1u9JNkVQwCuQlpScUulMKP6VyCuleJ0 6dHcmeSwZEVGPfGT1D7jS+IB6RZrt6f2h2cLrGku+leR467FVVGe3vT213AA6N5XcWS0 pa8RGs4idtdY2QEOyZyWoZilowaWZL+zqoh0ulBe0R8jHu1HKxVCoH2NOEQvKFHhkPDw aAPL5DKsWpmAz2R3dU4gHLZrV25PVdlMnaztzFDVtmXUN4/Ns7TRSNikrsJ0LIWWPr90 3TWA== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148205; x=1707753005; 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=KixoytXQDGjoXAfvxRuyy14dtLsh1IkvCMs1kLMTxgU=; b=fOOPwbgeGvGWL4kLY8doMMtNFB2Bi8mO5sVD6zyjEvdL28Hqs3RjrtUmbM2ahe4Owj 6WAdEWShO+Gkr9jElqWUb8cwBz6Xln3kVbIve9RQGnRc/uA/EUxLuyd8aeZbVk3WT14m rCSGR+XDyQT1Vy7Zlp6J31XiuEk2dyfkO0YbsCtVH6dyXO0YhpmYScDcfurxhMzSMj8W S6gCS1rzhnOnxNGFXpahmw1s7JBWE9rtc8OvBV7jmNkoy7nEVfQ0zR8IoXw4e2kDFezM PPjQVFw0NtF5lG1uaq2LbS9czH/iiAQUqiavDfrf0YLbKC+S49Vp05T82Koo07yVXCOV GYIg== X-Gm-Message-State: AOJu0Yw1FCp9vF4QTpuA982GtKa15cOndTtlgxe48DOo3rAreUp4yrqf smaLGZ8gKIfktT9r8cYrVOf83yaX/jc7m2VMYaYjhfzfbLHI3zNENtU0TcaBK3c6hP8u64Je3Wl tJ/+w+622p4Qn0p9FUIkpPmvEKA== X-Google-Smtp-Source: AGHT+IEU6z6eWwsI2wrzuaRC8jm+CHnZ6wBunkVHRxWVayHsmL35Oo/zIT20NC39+nbywkp3VaDfFSfpxPdFXrAwlcw= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:138f:b0:dc6:dd74:de68 with SMTP id x15-20020a056902138f00b00dc6dd74de68mr440165ybu.12.1707148205361; Mon, 05 Feb 2024 07:50:05 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:01 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-1-995e3eee38c0@google.com> Subject: [PATCH 1/6] rust: add `container_of!` macro From: mattgilbride@google.com 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 From: Wedson Almeida Filho This macro is used to obtain a pointer to an entire struct when given a pointer to a field in that struct. Signed-off-by: Wedson Almeida Filho Signed-off-by: Matt Gilbride Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl --- rust/kernel/lib.rs | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs index 7ac39874aeac..c7963efd1318 100644 --- a/rust/kernel/lib.rs +++ b/rust/kernel/lib.rs @@ -102,3 +102,35 @@ fn panic(info: &core::panic::PanicInfo<'_>) -> ! { // SAFETY: FFI call. unsafe { bindings::BUG() }; } + +/// Produces a pointer to an object from a pointer to one of its fields. +/// +/// # Safety +/// +/// The pointer passed to this macro, and the pointer returned by this mac= ro, must both be in +/// bounds of the same allocation. +/// +/// # Examples +/// +/// ``` +/// # use kernel::container_of; +/// struct Test { +/// a: u64, +/// b: u32, +/// } +/// +/// let test =3D Test { a: 10, b: 20 }; +/// let b_ptr =3D &test.b; +/// // SAFETY: The pointer points at the `b` field of a `Test`, so the res= ulting pointer will be +/// // in-bounds of the same allocation as `b_ptr`. +/// let test_alias =3D unsafe { container_of!(b_ptr, Test, b) }; +/// assert!(core::ptr::eq(&test, test_alias)); +/// ``` +#[macro_export] +macro_rules! container_of { + ($ptr:expr, $type:ty, $($f:tt)*) =3D> {{ + let ptr =3D $ptr as *const _ as *const u8; + let offset: usize =3D ::core::mem::offset_of!($type, $($f)*); + ptr.sub(offset) as *const $type + }} +} --=20 2.43.0.594.gd9cf4e227d-goog From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yb1-f201.google.com (mail-yb1-f201.google.com [209.85.219.201]) (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 B057B40BE4 for ; Mon, 5 Feb 2024 15:50:07 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.219.201 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148212; cv=none; b=UlbKgPSHKn0bVb1KR2FFpdKB9qtT0lb2HcJgUL9JXtnRzPWP/STLEPeH6iN1bn1jQeN+h5j5VUgsxfVfZF/w+xCNl2laqgUHgVWvG64txsFV1i+RmyaahLY07YZQEybIknndqs3co/usqhzYr0aNIUbhstB1khy896TfUv4hppc= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148212; c=relaxed/simple; bh=uvN4HYBhK3WEe9AZbrcSp8gNrWklicLvg2W1iLCEJOQ=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=lkLxPe+mhblV3EHdgo2aKGuMX93v4XgusM2CD0DnDu3F1oFoaQ4li038VTlRqLPNsU1XHoZ4X3TEQ4Lt3hH6rYPaNtyEGDiP7DdezWcp/JCdudU8d4U+Ak476sKqR3p0CfRNL4mZLngExpnRM0B4PPw23efrcP669ssdp8zYkns= 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=Ehtup+4w; arc=none smtp.client-ip=209.85.219.201 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="Ehtup+4w" Received: by mail-yb1-f201.google.com with SMTP id 3f1490d57ef6-dc6cd10fd4aso5913536276.1 for ; Mon, 05 Feb 2024 07:50:07 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148206; x=1707753006; 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=BSUqxiDZoHSJvKfsybLXdkjXBa+halpHRGLUo35nIOk=; b=Ehtup+4wegNChfVqQCiCZ0G2EiJXeIvUO6/gm5eixoukZ0+DVKTKTL/PZt9JTPVU+h f9HDkJJDHwCpKZmq7uC9mccXMkY+CSliGhArCjS06MGhHEn32NHohbI9tJUVNVu+0155 VNKJ4Engo0Wngg9LzURxtZxtbDLLbXxJWaWPbAYk65Jqz5kHmArGGsgCT86cQcR5lOh1 tq8HmKWed/hBdsuVjuypv+t6MxwrxBg10+yCXhCpivdFVUsGJD0psqadTj5NW5i008SW X0RR+r+Tf5uPdoqqDB/SuniI7cAZ6sj0waoOYt0JfgUef0qRE2Ivt62VIIBJSIiJmpC+ MV9Q== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148206; x=1707753006; 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=BSUqxiDZoHSJvKfsybLXdkjXBa+halpHRGLUo35nIOk=; b=o4B0jvtR9BmN9m6I4S/IyRmaI9L1RiPzIolTqR+ifemnFPQ+EU8PZLIO2f1yKtfTUE DkEholJgItsiY+f6ZOOn7ul3Ss0uvUyM7bXQSNcG+8yYNGY/XZJ9rexnQRkjg59dNgyu nE4lphTfFoLC4WDYKN7UIAOiCzQ1D/A4fGIKPWM9XYsnnkAz/+idYjQfhsY3+r4EZjZX ficsPJT2tPjgepqKnKWlRUx+IxblB9OXE/EmkjwaSb5Xe/scnCsg1BHedzTAq4e5HVNB rDqCOyTs55e1sJ6eJydOzzSO0mBuntYA8wnvWOk1vU7q2TFO1RyRHD/s7t+A+ABSj4v0 6UsQ== X-Gm-Message-State: AOJu0Yx+ojndmHgfY/y6gZ6gns8ehL32N5Gmbfe/8UUXYg2SfE+iBoFL Hr7ClOK5VCA64k5GGuJfNOVsE6ZRgmWGu6IVvBxXlmGMnFSgMTMCGuZdl66FaGZx4vfl8ZSQxIF Al1TERfdlRTSkUa4rZ0r7t1XGaA== X-Google-Smtp-Source: AGHT+IHx/++5HZZOfNcEEytS2+DQyLLaGUfJVofe6ZrtLiFHJiWrqNHshAkbTSlhSIHnSI8H4hdXpkyMEQESNw4cxGI= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:160d:b0:dc2:421d:ee30 with SMTP id bw13-20020a056902160d00b00dc2421dee30mr454909ybb.6.1707148206511; Mon, 05 Feb 2024 07:50:06 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:02 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-2-995e3eee38c0@google.com> Subject: [PATCH 2/6] rust: rbtree: add red-black tree implementation backed by the C version From: mattgilbride@google.com 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 From: Wedson Almeida Filho The rust rbtree exposes a map-like interface over keys and values, backed by the kernel red-black tree implementation. Values can be inserted, deleted, and retrieved from a `RBTree` by key. This base abstraction is used by binder to store key/value pairs and perform lookups, for example the patch "[PATCH RFC 03/20] rust_binder: add threading support" in the binder RFC [1]. Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-3-08ba= 9197f637@google.com/ [1] Signed-off-by: Wedson Almeida Filho Signed-off-by: Matt Gilbride Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl --- rust/helpers.c | 7 + rust/kernel/lib.rs | 1 + rust/kernel/rbtree.rs | 403 ++++++++++++++++++++++++++++++++++++++++++++++= ++++ 3 files changed, 411 insertions(+) diff --git a/rust/helpers.c b/rust/helpers.c index 70e59efd92bc..56ec79e823df 100644 --- a/rust/helpers.c +++ b/rust/helpers.c @@ -157,6 +157,13 @@ void rust_helper_init_work_with_key(struct work_struct= *work, work_func_t func, } EXPORT_SYMBOL_GPL(rust_helper_init_work_with_key); =20 +void rust_helper_rb_link_node(struct rb_node *node, struct rb_node *parent, + struct rb_node **rb_link) +{ + rb_link_node(node, parent, rb_link); +} +EXPORT_SYMBOL_GPL(rust_helper_rb_link_node); + /* * `bindgen` binds the C `size_t` type as the Rust `usize` type, so we can * use it in contexts where Rust expects a `usize` like slice (array) indi= ces. diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs index c7963efd1318..eb84ffba1831 100644 --- a/rust/kernel/lib.rs +++ b/rust/kernel/lib.rs @@ -43,6 +43,7 @@ pub mod net; pub mod prelude; pub mod print; +pub mod rbtree; mod static_assert; #[doc(hidden)] pub mod std_vendor; diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs new file mode 100644 index 000000000000..f33650258743 --- /dev/null +++ b/rust/kernel/rbtree.rs @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: GPL-2.0 + +//! Red-black trees. +//! +//! C header: [`include/linux/rbtree.h`](../../../../include/linux/rbtree.= h) +//! +//! Reference: + +use crate::{bindings, error::Result}; +use alloc::boxed::Box; +use core::{ + cmp::{Ord, Ordering}, + marker::PhantomData, + mem::MaybeUninit, + ptr::{addr_of_mut, NonNull}, +}; + +struct Node { + links: bindings::rb_node, + key: K, + value: V, +} + +/// A red-black tree with owned nodes. +/// +/// It is backed by the kernel C red-black trees. +/// +/// # Invariants +/// +/// Non-null parent/children pointers stored in instances of the `rb_node`= C struct are always +/// valid, and pointing to a field of our internal representation of a nod= e. +/// +/// # Examples +/// +/// In the example below we do several operations on a tree. We note that = insertions may fail if +/// the system is out of memory. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // Check the nodes we just inserted. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &100); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30).unwrap(), &300); +/// } +/// +/// // Replace one of the elements. +/// tree.try_create_and_insert(10, 1000)?; +/// +/// // Check that the tree reflects the replacement. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &1000); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30).unwrap(), &300); +/// } +/// +/// // Change the value of one of the elements. +/// *tree.get_mut(&30).unwrap() =3D 3000; +/// +/// // Check that the tree reflects the update. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &1000); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30).unwrap(), &3000); +/// } +/// +/// // Remove an element. +/// tree.remove(&10); +/// +/// // Check that the tree reflects the removal. +/// { +/// assert_eq!(tree.get(&10), None); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30).unwrap(), &3000); +/// } +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// In the example below, we first allocate a node, acquire a spinlock, th= en insert the node into +/// the tree. This is useful when the insertion context does not allow sle= eping, for example, when +/// holding a spinlock. +/// +/// ``` +/// use kernel::{rbtree::RBTree, sync::SpinLock}; +/// +/// fn insert_test(tree: &SpinLock>) -> Result { +/// // Pre-allocate node. This may fail (as it allocates memory). +/// let node =3D RBTree::try_allocate_node(10, 100)?; +/// +/// // Insert node while holding the lock. It is guaranteed to succeed= with no allocation +/// // attempts. +/// let mut guard =3D tree.lock(); +/// guard.insert(node); +/// Ok(()) +/// } +/// ``` +/// +/// In the example below, we reuse an existing node allocation from an ele= ment we removed. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // Check the nodes we just inserted. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &100); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30).unwrap(), &300); +/// } +/// +/// // Remove a node, getting back ownership of it. +/// let existing =3D tree.remove_node(&30).unwrap(); +/// +/// // Check that the tree reflects the removal. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &100); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// assert_eq!(tree.get(&30), None); +/// } +/// +/// // Turn the node into a reservation so that we can reuse it with a dif= ferent key/value. +/// let reservation =3D existing.into_reservation(); +/// +/// // Insert a new node into the tree, reusing the previous allocation. T= his is guaranteed to +/// // succeed (no memory allocations). +/// tree.insert(reservation.into_node(15, 150)); +/// +/// // Check that the tree reflect the new insertion. +/// { +/// assert_eq!(tree.get(&10).unwrap(), &100); +/// assert_eq!(tree.get(&15).unwrap(), &150); +/// assert_eq!(tree.get(&20).unwrap(), &200); +/// } +/// +/// # Ok::<(), Error>(()) +/// ``` +pub struct RBTree { + root: bindings::rb_root, + _p: PhantomData>, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl Send for RBTree {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl Sync for RBTree {} + +impl RBTree { + /// Creates a new and empty tree. + pub fn new() -> Self { + Self { + // INVARIANT: There are no nodes in the tree, so the invariant= holds vacuously. + root: bindings::rb_root::default(), + _p: PhantomData, + } + } + + /// Allocates memory for a node to be eventually initialised and inser= ted into the tree via a + /// call to [`RBTree::insert`]. + pub fn try_reserve_node() -> Result> { + Ok(RBTreeNodeReservation { + node: Box::try_new(MaybeUninit::uninit())?, + }) + } + + /// Allocates and initialises a node that can be inserted into the tre= e via + /// [`RBTree::insert`]. + pub fn try_allocate_node(key: K, value: V) -> Result>= { + Ok(Self::try_reserve_node()?.into_node(key, value)) + } +} + +impl RBTree +where + K: Ord, +{ + /// Tries to insert a new value into the tree. + /// + /// It overwrites a node if one already exists with the same key and r= eturns it (containing the + /// key/value pair). Returns [`None`] if a node with the same key didn= 't already exist. + /// + /// Returns an error if it cannot allocate memory for the new node. + pub fn try_create_and_insert(&mut self, key: K, value: V) -> Result>> { + Ok(self.insert(Self::try_allocate_node(key, value)?)) + } + + /// Inserts a new node into the tree. + /// + /// It overwrites a node if one already exists with the same key and r= eturns it (containing the + /// key/value pair). Returns [`None`] if a node with the same key didn= 't already exist. + /// + /// This function always succeeds. + pub fn insert(&mut self, node: RBTreeNode) -> Option> { + let RBTreeNode { node } =3D node; + let node =3D Box::into_raw(node); + // SAFETY: `node` is valid at least until we call `Box::from_raw`,= which only happens when + // the node is removed or replaced. + let node_links =3D unsafe { addr_of_mut!((*node).links) }; + let mut new_link: &mut *mut bindings::rb_node =3D &mut self.root.r= b_node; + let mut parent =3D core::ptr::null_mut(); + while !new_link.is_null() { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(*new_link, Node, links) }; + + parent =3D *new_link; + + // SAFETY: `this` is a non-null node so it is valid by the typ= e invariants. `node` is + // valid until the node is removed. + match unsafe { (*node).key.cmp(&(*this).key) } { + // SAFETY: `parent` is a non-null node so it is valid by t= he type invariants. + Ordering::Less =3D> new_link =3D unsafe { &mut (*parent).r= b_left }, + // SAFETY: `parent` is a non-null node so it is valid by t= he type invariants. + Ordering::Greater =3D> new_link =3D unsafe { &mut (*parent= ).rb_right }, + Ordering::Equal =3D> { + // INVARIANT: We are replacing an existing node with a= new one, which is valid. + // It remains valid because we "forgot" it with `Box::= into_raw`. + // SAFETY: All pointers are non-null and valid (parent= , despite the name, really + // is the node we're replacing). + unsafe { bindings::rb_replace_node(parent, node_links,= &mut self.root) }; + + // INVARIANT: The node is being returned and the calle= r may free it, however, + // it was removed from the tree. So the invariants sti= ll hold. + return Some(RBTreeNode { + // SAFETY: `this` was a node in the tree, so it is= valid. + node: unsafe { Box::from_raw(this as _) }, + }); + } + } + } + + // INVARIANT: We are linking in a new node, which is valid. It rem= ains valid because we + // "forgot" it with `Box::into_raw`. + // SAFETY: All pointers are non-null and valid (`*new_link` is nul= l, but `new_link` is a + // mutable reference). + unsafe { bindings::rb_link_node(node_links, parent, new_link) }; + + // SAFETY: All pointers are valid. `node` has just been inserted i= nto the tree. + unsafe { bindings::rb_insert_color(node_links, &mut self.root) }; + None + } + + /// Returns a node with the given key, if one exists. + fn find(&self, key: &K) -> Option>> { + let mut node =3D self.root.rb_node; + while !node.is_null() { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(node, Node, l= inks) }; + // SAFETY: `this` is a non-null node so it is valid by the typ= e invariants. + node =3D match key.cmp(unsafe { &(*this).key }) { + // SAFETY: `node` is a non-null node so it is valid by the= type invariants. + Ordering::Less =3D> unsafe { (*node).rb_left }, + // SAFETY: `node` is a non-null node so it is valid by the= type invariants. + Ordering::Greater =3D> unsafe { (*node).rb_right }, + Ordering::Equal =3D> return NonNull::new(this as _), + } + } + None + } + + /// Returns a reference to the value corresponding to the key. + pub fn get(&self, key: &K) -> Option<&V> { + // SAFETY: The `find` return value is a node in the tree, so it is= valid. + self.find(key).map(|node| unsafe { &node.as_ref().value }) + } + + /// Returns a mutable reference to the value corresponding to the key. + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + // SAFETY: The `find` return value is a node in the tree, so it is= valid. + self.find(key) + .map(|mut node| unsafe { &mut node.as_mut().value }) + } + + /// Removes the node with the given key from the tree. + /// + /// It returns the node that was removed if one exists, or [`None`] ot= herwise. + fn remove_node(&mut self, key: &K) -> Option> { + let mut node =3D self.find(key)?; + + // SAFETY: The `find` return value is a node in the tree, so it is= valid. + unsafe { bindings::rb_erase(&mut node.as_mut().links, &mut self.ro= ot) }; + + // INVARIANT: The node is being returned and the caller may free i= t, however, it was + // removed from the tree. So the invariants still hold. + Some(RBTreeNode { + // SAFETY: The `find` return value was a node in the tree, so = it is valid. + node: unsafe { Box::from_raw(node.as_ptr()) }, + }) + } + + /// Removes the node with the given key from the tree. + /// + /// It returns the value that was removed if one exists, or [`None`] o= therwise. + pub fn remove(&mut self, key: &K) -> Option { + let node =3D self.remove_node(key)?; + let RBTreeNode { node } =3D node; + let Node { + links: _, + key: _, + value, + } =3D *node; + Some(value) + } +} + +impl Default for RBTree { + fn default() -> Self { + Self::new() + } +} + +impl Drop for RBTree { + fn drop(&mut self) { + // SAFETY: `root` is valid as it's embedded in `self` and we have = a valid `self`. + let mut next =3D unsafe { bindings::rb_first_postorder(&self.root)= }; + + // INVARIANT: The loop invariant is that all tree nodes from `next= ` in postorder are valid. + while !next.is_null() { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(next, Node, l= inks) }; + + // Find out what the next node is before disposing of the curr= ent one. + // SAFETY: `next` and all nodes in postorder are still valid. + next =3D unsafe { bindings::rb_next_postorder(next) }; + + // INVARIANT: This is the destructor, so we break the type inv= ariant during clean-up, + // but it is not observable. The loop invariant is still maint= ained. + // SAFETY: `this` is valid per the loop invariant. + unsafe { drop(Box::from_raw(this as *mut Node)) }; + } + } +} + +/// A memory reservation for a red-black tree node. +/// +/// It contains the memory needed to hold a node that can be inserted into= a red-black tree. One +/// can be obtained by directly allocating it ([`RBTree::try_reserve_node`= ]). +pub struct RBTreeNodeReservation { + node: Box>>, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl Send for RBTreeNodeReservation {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl Sync for RBTreeNodeReservation {} + +impl RBTreeNodeReservation { + /// Initialises a node reservation. + /// + /// It then becomes an [`RBTreeNode`] that can be inserted into a tree. + pub fn into_node(mut self, key: K, value: V) -> RBTreeNode { + let node_ptr =3D self.node.as_mut_ptr(); + // SAFETY: `node_ptr` is valid, and so are its fields. + unsafe { addr_of_mut!((*node_ptr).links).write(bindings::rb_node::= default()) }; + // SAFETY: `node_ptr` is valid, and so are its fields. + unsafe { addr_of_mut!((*node_ptr).key).write(key) }; + // SAFETY: `node_ptr` is valid, and so are its fields. + unsafe { addr_of_mut!((*node_ptr).value).write(value) }; + let raw =3D Box::into_raw(self.node); + RBTreeNode { + // SAFETY: The pointer came from a `MaybeUninit` whose f= ields have all been + // initialised. Additionally, it has the same layout as `Node`. + node: unsafe { Box::from_raw(raw as _) }, + } + } +} + +/// A red-black tree node. +/// +/// The node is fully initialised (with key and value) and can be inserted= into a tree without any +/// extra allocations or failure paths. +pub struct RBTreeNode { + node: Box>, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl Send for RBTreeNode {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl Sync for RBTreeNode {} --=20 2.43.0.594.gd9cf4e227d-goog From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yb1-f201.google.com (mail-yb1-f201.google.com [209.85.219.201]) (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 C96554176E for ; Mon, 5 Feb 2024 15:50:08 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.219.201 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148210; cv=none; b=pTxpcC9fQ6Wc7jPf09Tqc+FKF5BE9ghtrwykXqGqaOW+qj7Q7bt9SZuN8xTxY14iBluROEoOsWp9IY/s2L1TfoeV/7esp89V3dW9e/PkU6dFiKzWmLw0dIH+F+1kt2Imn6c5jNWSpYYO/Fq1ARrkTbJu33z1Slf2qPm5Yb4xsG8= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148210; c=relaxed/simple; bh=bjenK6W2mteZ4lawE7+mh5kopqmKFi1EJSc5e4+d2Jc=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=GSAkjaC//M6Yxy+uIVvujJaPkwvUx9KHPSaiZbR10xo88yigI9pkd+xosaAdiMxt6brw4SqgBoqwZ+O0/lwnmyKksNNMyx6n9GOCEMdx7mHvvxDlJy+Exo8F4eURWqjQWekPT0mAanm8PnytkTpmfuVVwGjTu92K2HGT+FRy84Q= 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=L6cwn6Vh; arc=none smtp.client-ip=209.85.219.201 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="L6cwn6Vh" Received: by mail-yb1-f201.google.com with SMTP id 3f1490d57ef6-dc6ba69e803so7673004276.2 for ; Mon, 05 Feb 2024 07:50:08 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148208; x=1707753008; 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=arEHk0ctRB/sznZOfkcY6q8hHQuRWxqyzdzyyUckFyo=; b=L6cwn6Vh1U5fLZ5yAFbkudnqYFiZ6g2UMmu/55AH5fv9PlIC83lXqcr2lu4yyp3o0m vobfwf9n5cNiZSsw5eRaUyR6j4wJ4zFZ3VINFr3ZYaZibQYtJ+jDB+EKCSQzRux66fB9 I2c+pZ/MypqSdlqi6Oufnl2/omuVXwx9NRKYcLZRlIhTlTtufXjVnnI4k3BBrpd3AJeo F7T7tOPp7MQvGYEJtSkqjEtGhLJ6zlJJ+nx10BuaUXug38lgBue0xZyOi5HBhfUflmMI 5B2mHWnrmLOBoqAwmEkqXis9xcuaRrDvuDNow8txIaLIumKy57Bd4KbtNGAAiE5NCi+S KLnw== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148208; x=1707753008; 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=arEHk0ctRB/sznZOfkcY6q8hHQuRWxqyzdzyyUckFyo=; b=Rbh2I0NYIgmQWhj7rgQK6Sj8hd9i/zMT0guDOXdujyBdVG1LxISJhfCVwJnq9wfm4E cw4OJlq7KTZnEchhVTDIGR79kkQuXmkK3P4VJPexVDRFwAjl6eLO3wkv8Bm0hxZ3MZAX auSYw6KLU0UpAh2HhNspWqn3GwxV+A1E/Z9ghcnOFm6qX3w9iH1R233UbXypIYw43Tsf gKMZJkUp5lXvtFw6S7P/Ord1RMQz9YsAOyKBGZAxBOsBV7a/1e5EV7kUuGWE13/iaSh8 bKnGfdLyE1c8YqSaz+wVH3cA1oZYio9er72dC8h6M1yVbSBpiWt7UwKl3qVAne3jHQqo BeIw== X-Gm-Message-State: AOJu0YxW2faGHfaNlM/UE/Rbyivbaw94g61b/Cv+qLGF75fiGTwhxcu0 nFnbiYdcNqaBsY2m7Sh9vJ7wHxtvvl9Uhajtr41BbrcI8iYScTMsosqVKi7nlzV/eNVNTqzU/di oDbHsHjpdF7FiTmGy5p+R+YgdhQ== X-Google-Smtp-Source: AGHT+IFrW68xCg2j39/5TZqCiA+L40nU0XhZR6jp/fIklvNF0+R996XTcQnLLtdTcbbz7LXffGN1aBPZpJp5h+lA1ZQ= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:230d:b0:dc6:519b:5425 with SMTP id do13-20020a056902230d00b00dc6519b5425mr2873212ybb.11.1707148207870; Mon, 05 Feb 2024 07:50:07 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:03 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-3-995e3eee38c0@google.com> Subject: [PATCH 3/6] rust: rbtree: add `RBTreeIterator` From: mattgilbride@google.com 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 From: Wedson Almeida Filho - Add Iterator implementation (`RBTreeIterator`) for `RBTree`, allowing iteration over (key, value) pairs in key order. - Add individual `keys()` and `values()` functions to iterate over keys or values alone. - Update doctests to use iteration instead of explicitly getting items. Iteration is needed by the binder driver to enumerate all values in a tree for oneway spam detection [1]. Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-17-08b= a9197f637@google.com/ [1] Signed-off-by: Wedson Almeida Filho Signed-off-by: Matt Gilbride Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl --- rust/kernel/rbtree.rs | 125 ++++++++++++++++++++++++++++++++++++++++++----= ---- 1 file changed, 107 insertions(+), 18 deletions(-) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index f33650258743..29d8c2f6bd7b 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -53,14 +53,30 @@ struct Node { /// assert_eq!(tree.get(&30).unwrap(), &300); /// } /// +/// // Iterate over the nodes we just inserted. +/// { +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &100)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &300)); +/// assert!(iter.next().is_none()); +/// } +/// +/// // Print all elements. +/// for (key, value) in &tree { +/// pr_info!("{} =3D {}\n", key, value); +/// } +/// /// // Replace one of the elements. /// tree.try_create_and_insert(10, 1000)?; /// /// // Check that the tree reflects the replacement. /// { -/// assert_eq!(tree.get(&10).unwrap(), &1000); -/// assert_eq!(tree.get(&20).unwrap(), &200); -/// assert_eq!(tree.get(&30).unwrap(), &300); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &1000)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &300)); +/// assert!(iter.next().is_none()); /// } /// /// // Change the value of one of the elements. @@ -68,9 +84,11 @@ struct Node { /// /// // Check that the tree reflects the update. /// { -/// assert_eq!(tree.get(&10).unwrap(), &1000); -/// assert_eq!(tree.get(&20).unwrap(), &200); -/// assert_eq!(tree.get(&30).unwrap(), &3000); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &1000)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &3000)); +/// assert!(iter.next().is_none()); /// } /// /// // Remove an element. @@ -78,9 +96,10 @@ struct Node { /// /// // Check that the tree reflects the removal. /// { -/// assert_eq!(tree.get(&10), None); -/// assert_eq!(tree.get(&20).unwrap(), &200); -/// assert_eq!(tree.get(&30).unwrap(), &3000); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &3000)); +/// assert!(iter.next().is_none()); /// } /// /// # Ok::<(), Error>(()) @@ -120,9 +139,11 @@ struct Node { /// /// // Check the nodes we just inserted. /// { -/// assert_eq!(tree.get(&10).unwrap(), &100); -/// assert_eq!(tree.get(&20).unwrap(), &200); -/// assert_eq!(tree.get(&30).unwrap(), &300); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &100)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &300)); +/// assert!(iter.next().is_none()); /// } /// /// // Remove a node, getting back ownership of it. @@ -130,9 +151,10 @@ struct Node { /// /// // Check that the tree reflects the removal. /// { -/// assert_eq!(tree.get(&10).unwrap(), &100); -/// assert_eq!(tree.get(&20).unwrap(), &200); -/// assert_eq!(tree.get(&30), None); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &100)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert!(iter.next().is_none()); /// } /// /// // Turn the node into a reservation so that we can reuse it with a dif= ferent key/value. @@ -144,9 +166,11 @@ struct Node { /// /// // Check that the tree reflect the new insertion. /// { -/// assert_eq!(tree.get(&10).unwrap(), &100); -/// assert_eq!(tree.get(&15).unwrap(), &150); -/// assert_eq!(tree.get(&20).unwrap(), &200); +/// let mut iter =3D tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&10, &100)); +/// assert_eq!(iter.next().unwrap(), (&15, &150)); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert!(iter.next().is_none()); /// } /// /// # Ok::<(), Error>(()) @@ -187,6 +211,25 @@ pub fn try_reserve_node() -> Result> { pub fn try_allocate_node(key: K, value: V) -> Result>= { Ok(Self::try_reserve_node()?.into_node(key, value)) } + + /// Returns an iterator over the tree nodes, sorted by key. + pub fn iter(&self) -> RBTreeIterator<'_, K, V> { + RBTreeIterator { + _tree: PhantomData, + // SAFETY: `root` is valid as it's embedded in `self` and we h= ave a valid `self`. + next: unsafe { bindings::rb_first(&self.root) }, + } + } + + /// Returns an iterator over the keys of the nodes in the tree, in sor= ted order. + pub fn keys(&self) -> impl Iterator { + self.iter().map(|(k, _)| k) + } + + /// Returns an iterator over the values of the nodes in the tree, sort= ed by key. + pub fn values(&self) -> impl Iterator { + self.iter().map(|(_, v)| v) + } } =20 impl RBTree @@ -349,6 +392,52 @@ fn drop(&mut self) { } } =20 +impl<'a, K, V> IntoIterator for &'a RBTree { + type Item =3D (&'a K, &'a V); + type IntoIter =3D RBTreeIterator<'a, K, V>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +/// An iterator over the nodes of a [`RBTree`]. +/// +/// Instances are created by calling [`RBTree::iter`]. +pub struct RBTreeIterator<'a, K, V> { + _tree: PhantomData<&'a RBTree>, + next: *mut bindings::rb_node, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Send, V: Send> Send for RBTreeIterator<'a, K, V> {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeIterator<'a, K, V> {} + +impl<'a, K, V> Iterator for RBTreeIterator<'a, K, V> { + type Item =3D (&'a K, &'a V); + + fn next(&mut self) -> Option { + if self.next.is_null() { + return None; + } + + // SAFETY: All links fields we create are in a `Node`. + let cur =3D unsafe { crate::container_of!(self.next, Node, l= inks) }; + + // SAFETY: The reference to the tree used to create the iterator o= utlives the iterator, so + // the tree cannot change. By the tree invariant, all nodes are va= lid. + self.next =3D unsafe { bindings::rb_next(self.next) }; + + // SAFETY: By the same reasoning above, it is safe to dereference = the node. Additionally, + // it is ok to return a reference to members because the iterator = must outlive it. + Some(unsafe { (&(*cur).key, &(*cur).value) }) + } +} + /// A memory reservation for a red-black tree node. /// /// It contains the memory needed to hold a node that can be inserted into= a red-black tree. One --=20 2.43.0.594.gd9cf4e227d-goog From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yb1-f201.google.com (mail-yb1-f201.google.com [209.85.219.201]) (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 F10C044C6B for ; Mon, 5 Feb 2024 15:50:09 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.219.201 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148211; cv=none; b=jj8Hsx09FDlq77YL58cajkGpVt9lzd33o6tTEY09K1Vsis3TtFreClMPqGQWDBsSp/mKFHg9VFeaRvSfh30o5HsjhLjqZy9VIerbfLJ8aBKwWOp84hdX2HN0MfLiAIioSjETiE5V1oaT6J6379ksLFM037o+r1jtdABW/AnSO+M= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148211; c=relaxed/simple; bh=DVzO7sH9br7pDvZZ99ERESu2r5MxqGc7Jk4Dqa4lhEA=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=LswUCLemHka+WRglfhceeRidqg+vUSHzW3HIGUjZnvwm/K0zesE2misp11hfh3Wr3js0erOyCu2BiQVkEpcpRiY+TcT5ar3F8yUnd8ZSLtTD8QlGQrvPQDfK88/lZXx70NlBvslOY38WgO1yC/Io55yiqeZS0Azuss+1T0aeRXE= 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=rmJ96hhe; arc=none smtp.client-ip=209.85.219.201 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="rmJ96hhe" Received: by mail-yb1-f201.google.com with SMTP id 3f1490d57ef6-dc6b26783b4so5014396276.0 for ; Mon, 05 Feb 2024 07:50:09 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148209; x=1707753009; 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=LqF9zLVi16rKyrWpAJ0bhub0WiHD3paBZjfcYyqlQjA=; b=rmJ96hhe7ME3DH4E5j4xu9kXfo5ePa2Y076KjrfzoAHjREOfp3/qEVTRSnLNG6d4XF My7XQLdVTYMhcgmLzM+sOXi9gqik4KadaB/40HH0/qmtLN5nckfk6N8jDjzuyvny/zLa WPqpTpGpsNClkHtmagWqRn8OYYvaZ49heOc41gZZFKboODWCpEWnYWoD28ARGv7En7mF q9wAL6Ab6zB5GZDr7sMbI/0zUBnywe2Ra7mz3Lcjn5b2k9/Z2JhZaVsmCbTrIs+8CoaI QHmdFVL/fP5K4m2Ft8IsiXi53TNvumHWP/kvnAAInAKs7jfmOlEXlpdZNcIAj1EQBTK9 hTvQ== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148209; x=1707753009; 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=LqF9zLVi16rKyrWpAJ0bhub0WiHD3paBZjfcYyqlQjA=; b=RcnptAOdm3ATa5HepPCDWUKz9i8OAWxUfMtn5JX8jPNbAf3QHlK+iveMs/YvBuxNLZ GrDPkDYTXYSBw6Vk5jsXpkBHNKXsWxcxzWdtbBfYuFjAH8ud++plEbslYYJ0VhG1mpTk dgvEMMVY25c5wJBt0At0ZMRkVO8/HJ6UVZyXuv0U+z1jVvHhFpQoxLhgLNkUi0m1RXtM tDA3uWJFBREsMjw8K0vnjGBjDfgEjQ2EBqoH8FEzMujPDbWfDyWu8RfIPLSTWb1+cSfC zrjFbuhWuxm/HSYk4nSp+5U+ItW2DGg2MNbqoznHqykJ2JB2znXNlE4SX3zanBl3FXOa LAMw== X-Gm-Message-State: AOJu0YwoNfLXFf84X6JWcDg/ypoblPZ2IUgRABXcXgUkImxGosPCfVE9 PuOlmkMXkJw4bQZg4j4QmnF2rPV8vedhOEmqj9Ek3CUh/g212QZgtPIdqvAWxu+BLfowYJu4M74 Vkuypx+pqJefrTlxnFcrDYT5SsA== X-Google-Smtp-Source: AGHT+IF6Fis7YrbjmpE2NhJ3s+XZIGQdOXc4JxZ0aS+q/QDNDI/ezUPodv7GcBiAVYfBKU0ROCQDSTmv9hF+DkEZOzA= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:993:b0:dc2:25fd:eff1 with SMTP id bv19-20020a056902099300b00dc225fdeff1mr454327ybb.4.1707148209024; Mon, 05 Feb 2024 07:50:09 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:04 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-4-995e3eee38c0@google.com> Subject: [PATCH 4/6] rust: rbtree: add `RBTreeIteratorMut` From: mattgilbride@google.com 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 , Matt Gilbride Content-Type: text/plain; charset="utf-8" Content-Transfer-Encoding: quoted-printable From: Wedson Almeida Filho Add mutable Iterator implementation (`RBTreeIteratorMut`) for `RBTree`, allowing iteration over (key, value) pairs in key order. Only values are mutable, as mutating keys implies modifying a node's position in the tree. Mutable iteration is used by the binder driver during shutdown to clean up the tree maintained by the "range allocator" [1]. Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba= 9197f637@google.com/ [1] Signed-off-by: Wedson Almeida Filho Signed-off-by: Matt Gilbride Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl --- rust/kernel/rbtree.rs | 61 +++++++++++++++++++++++++++++++++++++++++++++++= ++++ 1 file changed, 61 insertions(+) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index 29d8c2f6bd7b..db17734b3fa1 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -221,6 +221,15 @@ pub fn iter(&self) -> RBTreeIterator<'_, K, V> { } } =20 + /// Returns a mutable iterator over the tree nodes, sorted by key. + pub fn iter_mut(&mut self) -> RBTreeIteratorMut<'_, K, V> { + RBTreeIteratorMut { + _tree: PhantomData, + // SAFETY: `root` is valid as it's embedded in `self` and we h= ave a valid `self`. + next: unsafe { bindings::rb_first(&self.root) }, + } + } + /// Returns an iterator over the keys of the nodes in the tree, in sor= ted order. pub fn keys(&self) -> impl Iterator { self.iter().map(|(k, _)| k) @@ -230,6 +239,11 @@ pub fn keys(&self) -> impl Iterator { pub fn values(&self) -> impl Iterator { self.iter().map(|(_, v)| v) } + + /// Returns a mutable iterator over the values of the nodes in the tre= e, sorted by key. + pub fn values_mut(&mut self) -> impl Iterator { + self.iter_mut().map(|(_, v)| v) + } } =20 impl RBTree @@ -438,6 +452,53 @@ fn next(&mut self) -> Option { } } =20 +impl<'a, K, V> IntoIterator for &'a mut RBTree { + type Item =3D (&'a K, &'a mut V); + type IntoIter =3D RBTreeIteratorMut<'a, K, V>; + + fn into_iter(self) -> Self::IntoIter { + self.iter_mut() + } +} + +/// A mutable iterator over the nodes of a [`RBTree`]. +/// +/// Instances are created by calling [`RBTree::iter_mut`]. +pub struct RBTreeIteratorMut<'a, K, V> { + _tree: PhantomData<&'a RBTree>, + next: *mut bindings::rb_node, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Send, V: Send> Send for RBTreeIteratorMut<'a, K, V> {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeIteratorMut<'a, K, V> {} + +impl<'a, K, V> Iterator for RBTreeIteratorMut<'a, K, V> { + type Item =3D (&'a K, &'a mut V); + + fn next(&mut self) -> Option { + if self.next.is_null() { + return None; + } + + // SAFETY: All links fields we create are in a `Node`. + let cur =3D unsafe { crate::container_of!(self.next, Node, l= inks) }.cast_mut(); + + // SAFETY: The reference to the tree used to create the iterator o= utlives the iterator, so + // the tree cannot change (except for the value of previous nodes,= but those don't affect + // the iteration process). By the tree invariant, all nodes are va= lid. + self.next =3D unsafe { bindings::rb_next(self.next) }; + + // SAFETY: By the same reasoning above, it is safe to dereference = the node. Additionally, + // it is ok to return a reference to members because the iterator = must outlive it. + Some(unsafe { (&(*cur).key, &mut (*cur).value) }) + } +} + /// A memory reservation for a red-black tree node. /// /// It contains the memory needed to hold a node that can be inserted into= a red-black tree. One --=20 2.43.0.594.gd9cf4e227d-goog From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yw1-f201.google.com (mail-yw1-f201.google.com [209.85.128.201]) (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 5794B4595A for ; Mon, 5 Feb 2024 15:50:11 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.128.201 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148214; cv=none; b=n/zPCrs0uMqNvipqSR7pZeUc0APT3VtOBYVi8ZyfmRtI4ZT+jo9kAX4UWxcdXj2Xx+/g8fgfv8p4vRqst9nDAqqj1KFYUKW2xh9J4+gtnzqiv30Bo7VrK1fEr6PDf4Q/J25Yw2XiRXXRQOw3c39MUfpGsiFgLAA1qvobSyM/SMs= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148214; c=relaxed/simple; bh=fEdkpQRAKxcK+Vq+Ej5tRRpaejTCMjdLRVmNroqCOfw=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=gm68aH6pQrwxRzEC0OY3p+FytD2fhvGLJNnoJzMW+et0sMRiST/n0wKikUl3ekXrMWo8eyPrF3JUITUNjaHsDnkYOuIyx4f78kCHdnwyDPU351rk3iqLQFUNo2mg4NZvzX3GDyK5UUAxKyU9qglYhGyFQWGC3wImQraLXnGSeiw= 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=ECodOzTp; arc=none smtp.client-ip=209.85.128.201 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="ECodOzTp" Received: by mail-yw1-f201.google.com with SMTP id 00721157ae682-6041bb56dbfso62798797b3.2 for ; Mon, 05 Feb 2024 07:50:11 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148210; x=1707753010; 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=JtORM6z3n6Nx7zIdzVBRR85gUfsLFA4+3fHETUbUVdg=; b=ECodOzTpDLuiK1B/JxpKMoBWzzAohfjIZoTaPXbUaFquLe4BQ2TSTgBw4puALcHz6P OhUgsUxh3ApkCDy5hngMt117MZld7tUKkbsoIVV8HB/UkCDws2rifg9SjeEMyovtqJkU Xow/+San3+ysBIYScMFujk7zV5o539QJiLFf8zsI3GLg0JUC4rq36MJDnWr5O7ZZwC5N Kz/pLltL2Qn59v8K3zg9tyn6E3V1K2YF0QUjZobPOXecbGLVAYMYhsT2mOafh3gvqb19 mM4KkF8Dhc+NGxSVjezDq/v1UIGPaD8vZTQ8G7uw1mi9A7a8bXzuikUmJ/lwG/96tvr6 ajEQ== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148210; x=1707753010; 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=JtORM6z3n6Nx7zIdzVBRR85gUfsLFA4+3fHETUbUVdg=; b=jHGb5oq2q3TUGVr0VIZBP3FBJy+9u5F70V2McPIIsjvFFJZEVJMQOZeAn/adiZAnWP meYgLvJQNOJTXB05sKhELaKz9iid+q7xIyDzLxmI1lUSp0/BgdgPGZE6E8YsD/dnLBb7 MlNcAYiRZKlGZLeyb75WLX54tCreMZfOUVXuUI2AlG4L0XLDIWXrv82hSyNLUsvNojpr 3geRMXq52qx15R24nltjZ8F1XThGIDjwVQ04PCuzDuhtKopplVHqqEfkZpknlx5bKa7Q yB0I1/umOyXXHm9S3IpU+CChHbvceHZR1RiPTMSbmWc6+Qm+VoCgEfTWbCtXTOajx7L7 mTjg== X-Gm-Message-State: AOJu0YyXIOX5PLfBSt2NlyJhYi9CtydzK7xiRsTx6kVrI/Y7QOueqaEg x+638pjWyX37JicaPA0X/ysuVcoc8q0LEniQGuMzA2SQcucOMNt53dzP5aMhFwtBKlvuH33aAQ5 Wv/i8AVwHL3JqIcrV3UmKyye/zA== X-Google-Smtp-Source: AGHT+IEnEzg36o2wODfPotcfBiZl9rChq+AHl88nPPHI6AibfHOM9IDUYQiwQyQs1syRA3Qg9H3P3oAC99qmsmsUrB0= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a25:b183:0:b0:dc6:ebd4:cca2 with SMTP id h3-20020a25b183000000b00dc6ebd4cca2mr327456ybj.11.1707148210290; Mon, 05 Feb 2024 07:50:10 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:05 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-5-995e3eee38c0@google.com> Subject: [PATCH 5/6] rust: rbtree: add `RBTreeCursor` From: mattgilbride@google.com 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 Signed-off-by: Matt Gilbride Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl --- rust/kernel/rbtree.rs | 512 ++++++++++++++++++++++++++++++++++++++++++++++= ++++ 1 file changed, 512 insertions(+) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index db17734b3fa1..0db6a584a9fc 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -244,6 +244,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) }; + if current.is_null() { + return None; + } + Some(RBTreeCursor { + _tree: PhantomData, + root, + current, + }) + } + + /// 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) }; + if current.is_null() { + return None; + } + Some(RBTreeCursor { + _tree: PhantomData, + root, + current, + }) + } } =20 impl RBTree @@ -376,6 +406,59 @@ pub fn remove(&mut self, key: &K) -> Option { } =3D *node; Some(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: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(node, Node, l= inks) }.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 }; + if key =3D=3D this_key { + return Some(RBTreeCursor { + _tree: PhantomData, + root: addr_of_mut!(self.root), + current: node, + }); + } else { + node =3D if key > this_key { + right_child + } else { + 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); + } + left_child + } + }; + } + best_match.map(|best| RBTreeCursor { + _tree: PhantomData, + root: addr_of_mut!(self.root), + // SAFETY: `best` is a non-null node so it is valid by the typ= e invariants. + current: unsafe { addr_of_mut!((*best.as_ptr()).links) }, + }) + } } =20 impl Default for RBTree { @@ -406,6 +489,435 @@ fn drop(&mut self) { } } =20 +/// A bidirectional cursor over the tree nodes, sorted by key. +/// +/// # Invariants +/// +/// In instance of `RBTreeCursor` is only acquired from [`RBTree`]. +/// A reference to the tree used to create the cursor outlives the cursor,= so +/// the tree cannot change. By the tree invariant, all nodes are valid. +/// +/// # 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::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // 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::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// 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::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert five elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// tree.try_create_and_insert(40, 400)?; +/// tree.try_create_and_insert(50, 500)?; +/// +/// // 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::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // 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::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // 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().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().unwrap(); +/// current =3D cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Removing the last element in the tree returns [`None`]. +/// assert!(cursor.remove_current().is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// Nodes adjacent to the current node can also be removed. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree =3D RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(10, 100)?; +/// tree.try_create_and_insert(20, 200)?; +/// tree.try_create_and_insert(30, 300)?; +/// +/// // 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(), (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(), (30, 300)); +/// +/// # Ok::<(), Error>(()) +/// ``` +pub struct RBTreeCursor<'a, K, V> { + _tree: PhantomData<&'a RBTree>, + root: *mut bindings::rb_root, + current: *mut bindings::rb_node, +} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Send condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {} + +// SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its +// fields, so we use the same Sync condition as would be used for a struct= with K and V fields. +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeCursor<'a, K, V> {} + +impl<'a, K, V> RBTreeCursor<'a, K, V> { + /// The current node + pub fn current(&self) -> (&K, &V) { + Self::to_key_value(self.current) + } + + /// The current node, with a mutable value + pub fn current_mut(&mut self) -> (&K, &mut V) { + Self::to_key_value_mut(self.current) + } + + /// Remove the current node from the tree. + /// + /// Returns a cursor to the next node, if it exists, + /// else the previous node. Returns [`None`] if the tree + /// becomes empty. + pub fn remove_current(mut self) -> Option { + let prev =3D self.get_neighbor_raw(Direction::Prev); + let next =3D self.get_neighbor_raw(Direction::Next); + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(self.current, Node, links) }.cast_mut(); + // 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, self.root) }; + + let current =3D match (prev, next) { + (_, Some(next)) =3D> next, + (Some(prev), None) =3D> prev, + (None, None) =3D> { + return None; + } + }; + + Some(Self { + current, + _tree: self._tree, + root: self.root, + }) + } + + /// Remove the previous node, returning it if it exists. + pub fn remove_prev(&mut self) -> Option<(K, V)> { + self.remove_neighbor(Direction::Prev) + } + + /// Remove the next node, returning it if it exists. + pub fn remove_next(&mut self) -> Option<(K, V)> { + self.remove_neighbor(Direction::Next) + } + + fn remove_neighbor(&mut self, direction: Direction) -> Option<(K, V)> { + if let Some(neighbor) =3D self.get_neighbor_raw(direction) { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(neighbor, Node, links) }.cast_mut(); + // 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(&mut (*this).links, self.root) }; + return Some(Self::to_key_value_owned(neighbor)); + } + 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(mut self, direction: Direction) -> Option { + self.get_neighbor_raw(direction).map(|neighbor| Self { + _tree: self._tree, + root: self.root, + 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)> { + // SAFETY: `self.current` is valid by the type invariants. + let neighbor =3D unsafe { + match direction { + Direction::Prev =3D> bindings::rb_prev(self.current), + Direction::Next =3D> bindings::rb_next(self.current), + } + }; + + if neighbor.is_null() { + return None; + } + + Some(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)> { + // SAFETY: `self.current` is valid by the type invariants. + let neighbor =3D unsafe { + match direction { + Direction::Prev =3D> bindings::rb_prev(self.current), + Direction::Next =3D> bindings::rb_next(self.current), + } + }; + + if neighbor.is_null() { + return None; + } + + Some(Self::to_key_value_mut(neighbor)) + } + + fn get_neighbor_raw(&mut self, direction: Direction) -> Option<*mut bi= ndings::rb_node> { + // SAFETY: `self.current` is valid by the type invariants. + let neighbor =3D unsafe { + match direction { + Direction::Prev =3D> bindings::rb_prev(self.current), + Direction::Next =3D> bindings::rb_next(self.current), + } + }; + + if neighbor.is_null() { + return None; + } + + Some(neighbor) + } + + // This internal method should *only* be called with a valid pointer t= o a node. + fn to_key_value(node: *mut bindings::rb_node) -> (&'a K, &'a V) { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(node, Node, links= ) }; + // 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 { &(*this).value }; + (k, v) + } + + // This internal method should *only* be called with a valid pointer t= o a node. + fn to_key_value_mut(node: *mut bindings::rb_node) -> (&'a K, &'a mut V= ) { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(node, Node, links= ) }.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 { &mut (*this).value }; + (k, v) + } + + // This internal method should *only* be called with a valid pointer t= o a node *that is being removed*. + fn to_key_value_owned(node: *mut bindings::rb_node) -> (K, V) { + // SAFETY: All links fields we create are in a `Node`. + let this =3D unsafe { crate::container_of!(node, Node, links= ) }.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 n =3D unsafe { Box::from_raw(this) }; + + (n.key, n.value) + } +} + +/// Direction for [`RBTreeCursor`] 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 RBTreeIterator<'a, K, V>; --=20 2.43.0.594.gd9cf4e227d-goog From nobody Sun Feb 8 03:26:57 2026 Received: from mail-yw1-f202.google.com (mail-yw1-f202.google.com [209.85.128.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 520D645950 for ; Mon, 5 Feb 2024 15:50:12 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.128.202 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148214; cv=none; b=C0GFoGHuzJYI4+WkUbieQ8nXia+jA7MXmVg0tSbwbTEpkcw93bqFxWfI/FRcIZGgx93STkA6iiLgxyB49WqvpEZMziWzeYCHv5F5wb4JKAalcRERMNY7UlQ+6tb7V/zfQlrNYwk1FapO1/98ACZBvV0uS1vk/ZsqCua2ZPlMHb4= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1707148214; c=relaxed/simple; bh=22UG/FoayaE5gos56FE9J5zNHaONu8WJgvgbCJ9Ju5I=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=atHhdRr1GWILNuGDzZda7MUsNku9Qg8i+SfR/CkAKj+6K1FRrbmkxdFk8cob5LlrItW2dKXH23CSYds5gkFy/DQN5FbC6jG9JhI5cy//ksRgzHafOHi19RasQhb8+3dS62iPpcbB2+uu3br6kTpgzS5WDlzK0tYQcPVVjNh2A30= 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=PJrGdy5x; arc=none smtp.client-ip=209.85.128.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="PJrGdy5x" Received: by mail-yw1-f202.google.com with SMTP id 00721157ae682-5ee22efe5eeso70388267b3.3 for ; Mon, 05 Feb 2024 07:50:12 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1707148211; x=1707753011; 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=yvwbZDgWN0lh70oVs2Q46rrbtwbVSHs2JN97boJgs94=; b=PJrGdy5xE1H5o1QnDlFkdEmsS6OqYnjQUGGJuv9EfxoIFn8bg3LmvxN2fpaRmQqVCR 2a+huzEaH9BCRUgzrY5napM+oTqQEedsaQT1FjAJ6JWL2igFb26/GKAq9B3x9wVnmgA6 MUEEACIay0VgQeTQk+PIuWzsDgiJeYM1cwrBTAU3v9KpOsFHNxTfRd9YYi/n7IvJIDX8 wv1nRKwZB94FkhJLn44oSXKe2jWhwWfTuxTZQkszgcA8fyyIVecI6WzPA5/lAGPlXwjC 8VX7HW/oTNDGuRlh/B3qAw8MPyuL5X3O8U/IVonr5pv/eoWLwcsnr92tdqaRcDa2jeCj kIog== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1707148211; x=1707753011; 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=yvwbZDgWN0lh70oVs2Q46rrbtwbVSHs2JN97boJgs94=; b=jGwhqDwuFRTE/uphuNYzuwLQJ5y5n0J/Gp9DlvHnvm4b10EWN056B78tqW0dSnXeUL jfZzyp2FwRMONB8n+2JJHUW0G5Kj7aphfT2gwIi4fwXGSgTKvISK90rqZGphpqqNz0XZ Q0lIqPYNLEDVG5mNRcaBpl4s1X1X8A3Y5oQ7TTi+T6Ly0jF6Jld7S3oj27XGGlOgnCcp WeZLDTt3YXfUaQh7++M03aPqDu2ES8RQ5Y6hlmMPkA2nVAR3PExY9+eSSVmIHoTlNlum DS8+DEGpqL/fuAyeOGFknEzwNLq9mxuiwPGc2B4wTVDPdMEbhLQ+/fe6xC7CLrlqANVw GcPA== X-Gm-Message-State: AOJu0Yxf6Lf5Kpbhfzi6wHMwZEL/kDP22W/JSt7pwaJCcx3NFxaHW6Rf +l3TdVeqY4ZtV5yEHZr9lkE2xjnvXdVk2sJY/+XCro30DXJaVfT7eWt37QECAIEWPb44VZmtUZ8 Vbknu9hmHpleAhdu5fG1lUa5wdA== X-Google-Smtp-Source: AGHT+IEMft5d2wJebWkpwaL5gkrPglrd6gRNOeZLonZHEJQD2i80Ih8PdQSInGrpCo2h1yhXn22srVVwbWt57OQAkKY= X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:2206:b0:dc2:4921:cc0 with SMTP id dm6-20020a056902220600b00dc249210cc0mr455698ybb.5.1707148211433; Mon, 05 Feb 2024 07:50:11 -0800 (PST) Date: Mon, 05 Feb 2024 15:50:06 +0000 In-Reply-To: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240205-b4-rbtree-v1-0-995e3eee38c0@google.com> X-Mailer: b4 0.12.3 Message-ID: <20240205-b4-rbtree-v1-6-995e3eee38c0@google.com> Subject: [PATCH 6/6] rust: rbtree: add `RBTree::entry` From: mattgilbride@google.com 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 From: Alice Ryhl This mirrors the entry API [1] from the Rust standard library on `RBTree`. This API can be used to access the entry at a specific key and make modifications depending on whether the key is vacant or occupied. This API is useful because it can often be used to avoid traversing the tree multiple times. This is used by binder to look up and conditionally access or insert a value, depending on whether it is there or not [2]. Link: https://doc.rust-lang.org/stable/std/collections/btree_map/enum.Entry= .html [1] Link: https://android-review.googlesource.com/c/kernel/common/+/2849906 [2] Signed-off-by: Alice Ryhl Signed-off-by: Matt Gilbride Tested-by: Alice Ryhl --- rust/kernel/rbtree.rs | 284 ++++++++++++++++++++++++++++++++++++++--------= ---- 1 file changed, 216 insertions(+), 68 deletions(-) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index 0db6a584a9fc..9ee953df2347 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -297,56 +297,64 @@ pub fn try_create_and_insert(&mut self, key: K, value= : V) -> Result) -> Option> { - let RBTreeNode { node } =3D node; - let node =3D Box::into_raw(node); - // SAFETY: `node` is valid at least until we call `Box::from_raw`,= which only happens when - // the node is removed or replaced. - let node_links =3D unsafe { addr_of_mut!((*node).links) }; + match self.raw_entry(&node.node.key) { + RawEntry::Occupied(entry) =3D> Some(entry.replace(node)), + RawEntry::Vacant(entry) =3D> { + entry.insert(node); + None + } + } + } + + fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { let mut new_link: &mut *mut bindings::rb_node =3D &mut self.root.r= b_node; let mut parent =3D core::ptr::null_mut(); - while !new_link.is_null() { + while !(*new_link).is_null() { + let curr =3D *new_link; // SAFETY: All links fields we create are in a `Node`. - let this =3D unsafe { crate::container_of!(*new_link, Node, links) }; + let node =3D unsafe { crate::container_of!(curr, Node, l= inks) }; =20 - parent =3D *new_link; - - // SAFETY: `this` is a non-null node so it is valid by the typ= e invariants. `node` is - // valid until the node is removed. - match unsafe { (*node).key.cmp(&(*this).key) } { - // SAFETY: `parent` is a non-null node so it is valid by t= he type invariants. - Ordering::Less =3D> new_link =3D unsafe { &mut (*parent).r= b_left }, - // SAFETY: `parent` is a non-null node so it is valid by t= he type invariants. - Ordering::Greater =3D> new_link =3D unsafe { &mut (*parent= ).rb_right }, + // SAFETY: `node` is a non-null node so it is valid by the typ= e invariants. + match key.cmp(unsafe { &(*node).key }) { + // SAFETY: `curr` is a non-null node so it is valid by the= type invariants. + Ordering::Less =3D> new_link =3D unsafe { &mut (*curr).rb_= left }, + // SAFETY: `curr` is a non-null node so it is valid by the= type invariants. + Ordering::Greater =3D> new_link =3D unsafe { &mut (*curr).= rb_right }, Ordering::Equal =3D> { - // INVARIANT: We are replacing an existing node with a= new one, which is valid. - // It remains valid because we "forgot" it with `Box::= into_raw`. - // SAFETY: All pointers are non-null and valid (parent= , despite the name, really - // is the node we're replacing). - unsafe { bindings::rb_replace_node(parent, node_links,= &mut self.root) }; - - // INVARIANT: The node is being returned and the calle= r may free it, however, - // it was removed from the tree. So the invariants sti= ll hold. - return Some(RBTreeNode { - // SAFETY: `this` was a node in the tree, so it is= valid. - node: unsafe { Box::from_raw(this as _) }, - }); + return RawEntry::Occupied(OccupiedEntry { + rbtree: self, + node_links: curr, + }) } } + parent =3D curr; } =20 - // INVARIANT: We are linking in a new node, which is valid. It rem= ains valid because we - // "forgot" it with `Box::into_raw`. - // SAFETY: All pointers are non-null and valid (`*new_link` is nul= l, but `new_link` is a - // mutable reference). - unsafe { bindings::rb_link_node(node_links, parent, new_link) }; + RawEntry::Vacant(RawVacantEntry { + parent, + new_link, + rbtree: self, + }) + } =20 - // SAFETY: All pointers are valid. `node` has just been inserted i= nto the tree. - unsafe { bindings::rb_insert_color(node_links, &mut self.root) }; - None + /// Gets the given key's corresponding entry in the map for in-place m= anipulation. + pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { + match self.raw_entry(&key) { + RawEntry::Occupied(entry) =3D> Entry::Occupied(entry), + RawEntry::Vacant(entry) =3D> Entry::Vacant(VacantEntry { raw: = entry, key }), + } } =20 - /// Returns a node with the given key, if one exists. - fn find(&self, key: &K) -> Option>> { + /// Used for accessing the given node, if it exists. + pub fn find_mut(&mut self, key: &K) -> Option>= { + match self.raw_entry(key) { + RawEntry::Occupied(entry) =3D> Some(entry), + RawEntry::Vacant(_entry) =3D> None, + } + } + + /// Returns a reference to the value corresponding to the key. + pub fn get(&self, key: &K) -> Option<&V> { let mut node =3D self.root.rb_node; while !node.is_null() { // SAFETY: All links fields we create are in a `Node`. @@ -357,54 +365,30 @@ fn find(&self, key: &K) -> Option>= > { Ordering::Less =3D> unsafe { (*node).rb_left }, // SAFETY: `node` is a non-null node so it is valid by the= type invariants. Ordering::Greater =3D> unsafe { (*node).rb_right }, - Ordering::Equal =3D> return NonNull::new(this as _), + // SAFETY: `node` is a non-null node so it is valid by the= type invariants. + Ordering::Equal =3D> return Some(unsafe { &(*this).value }= ), } } None } =20 - /// Returns a reference to the value corresponding to the key. - pub fn get(&self, key: &K) -> Option<&V> { - // SAFETY: The `find` return value is a node in the tree, so it is= valid. - self.find(key).map(|node| unsafe { &node.as_ref().value }) - } - /// Returns a mutable reference to the value corresponding to the key. pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { - // SAFETY: The `find` return value is a node in the tree, so it is= valid. - self.find(key) - .map(|mut node| unsafe { &mut node.as_mut().value }) + self.find_mut(key).map(|node| node.into_mut()) } =20 /// Removes the node with the given key from the tree. /// /// It returns the node that was removed if one exists, or [`None`] ot= herwise. - fn remove_node(&mut self, key: &K) -> Option> { - let mut node =3D self.find(key)?; - - // SAFETY: The `find` return value is a node in the tree, so it is= valid. - unsafe { bindings::rb_erase(&mut node.as_mut().links, &mut self.ro= ot) }; - - // INVARIANT: The node is being returned and the caller may free i= t, however, it was - // removed from the tree. So the invariants still hold. - Some(RBTreeNode { - // SAFETY: The `find` return value was a node in the tree, so = it is valid. - node: unsafe { Box::from_raw(node.as_ptr()) }, - }) + pub fn remove_node(&mut self, key: &K) -> Option> { + self.find_mut(key).map(OccupiedEntry::remove_node) } =20 /// Removes the node with the given key from the tree. /// /// It returns the value that was removed if one exists, or [`None`] o= therwise. pub fn remove(&mut self, key: &K) -> Option { - let node =3D self.remove_node(key)?; - let RBTreeNode { node } =3D node; - let Node { - links: _, - key: _, - value, - } =3D *node; - Some(value) + self.find_mut(key).map(OccupiedEntry::remove) } =20 /// Returns a cursor over the tree nodes based on the given key. @@ -1063,3 +1047,167 @@ unsafe impl Send for RBTreeNode {} // SAFETY: An [`RBTree`] allows the same kinds of access to its values tha= t a struct allows to its // fields, so we use the same Sync condition as would be used for a struct= with K and V fields. unsafe impl Sync for RBTreeNode {} + +impl RBTreeNode { + /// "Uninitialises" a node. + /// + /// It then becomes a reservation that can be re-initialised into a di= fferent node (i.e., with + /// a different key and/or value). + /// + /// The existing key and value are dropped in-place as part of this op= eration, that is, memory + /// may be freed (but only for the key/value; memory for the node itse= lf is kept for reuse). + pub fn into_reservation(self) -> RBTreeNodeReservation { + let raw =3D Box::into_raw(self.node); + let mut ret =3D RBTreeNodeReservation { + // SAFETY: The pointer came from a valid `Node`, which has the= same layout as + // `MaybeUninit`. + node: unsafe { Box::from_raw(raw as _) }, + }; + // SAFETY: Although the type is `MaybeUninit`, we know it ha= s been initialised + // because it came from a `Node`. So it is safe to drop it. + unsafe { core::ptr::drop_in_place(ret.node.as_mut_ptr()) }; + ret + } +} + +/// A view into a single entry in a map, which may either be vacant or occ= upied. +/// +/// This enum is constructed from the [`entry`] method on [`RBTree`]. +/// +/// [`entry`]: fn@RBTree::entry +pub enum Entry<'a, K, V> { + /// This [`RBTree`] does not have a node with this key. + Vacant(VacantEntry<'a, K, V>), + /// This [`RBTree`] already has a node with this key. + Occupied(OccupiedEntry<'a, K, V>), +} + +/// Like [`Entry`], except that it doesn't have ownership of the key. +enum RawEntry<'a, K, V> { + Vacant(RawVacantEntry<'a, K, V>), + Occupied(OccupiedEntry<'a, K, V>), +} + +/// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`= ] enum. +pub struct VacantEntry<'a, K, V> { + key: K, + raw: RawVacantEntry<'a, K, V>, +} + +/// Like [`VacantEntry`], but doesn't hold on to the key. +struct RawVacantEntry<'a, K, V> { + rbtree: &'a mut RBTree, + /// The node that will become the parent of the new node if we insert = one. + /// + /// This pointer may be null if the new node becomes the root. + parent: *mut bindings::rb_node, + /// This points to the left-child or right-child field of `parent`. Th= is controls whether the + /// new node will become the left or right child of `parent`. + /// + /// If `parent` is null, then this points at `rbtree.root`. + new_link: *mut *mut bindings::rb_node, +} + +impl<'a, K, V> RawVacantEntry<'a, K, V> { + /// Inserts the given node into the [`RBTree`] at this entry. + /// + /// The `node` must have a key such that inserting it here does not br= eak the ordering of this + /// [`RBTree`]. + fn insert(self, node: RBTreeNode) -> &'a mut V { + let node =3D Box::into_raw(node.node); + + // SAFETY: `node` is valid at least until we call `Box::from_raw`,= which only happens when + // the node is removed or replaced. + let node_links =3D unsafe { addr_of_mut!((*node).links) }; + + // INVARIANT: We are linking in a new node, which is valid. It rem= ains valid because we + // "forgot" it with `Box::into_raw`. + // SAFETY: All pointers are null or valid in an appropriate way. + unsafe { bindings::rb_link_node(node_links, self.parent, self.new_= link) }; + + // SAFETY: All pointers are valid. `node` has just been inserted i= nto the tree. + unsafe { bindings::rb_insert_color(node_links, &mut self.rbtree.ro= ot) }; + + // SAFETY: The node is valid until we remove it from the tree. + unsafe { &mut (*node).value } + } +} + +impl<'a, K, V> VacantEntry<'a, K, V> { + /// Inserts the given node into the [`RBTree`] at this entry. + pub fn insert(self, value: V, reservation: RBTreeNodeReservation= ) -> &'a mut V { + self.raw.insert(reservation.into_node(self.key, value)) + } +} + +/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Ent= ry`] enum. +pub struct OccupiedEntry<'a, K, V> { + rbtree: &'a mut RBTree, + /// The node that this entry corresponds to. Non null. + node_links: *mut bindings::rb_node, +} + +impl<'a, K, V> OccupiedEntry<'a, K, V> { + fn node_ptr(&self) -> *mut Node { + // SAFETY: All links fields we create are in a `Node`. + unsafe { crate::container_of!(self.node_links, Node, links) = }.cast_mut() + } + + /// Gets a reference to the value in the entry. + pub fn get(&self) -> &V { + unsafe { &(*self.node_ptr()).value } + } + + /// Gets a mutable reference to the value in the entry. + pub fn get_mut(&mut self) -> &mut V { + unsafe { &mut (*self.node_ptr()).value } + } + + /// Converts the entry into a mutable reference to its value. + /// + /// If you need multiple references to the `OccupiedEntry`, see [`self= #get_mut`]. + pub fn into_mut(self) -> &'a mut V { + unsafe { &mut (*self.node_ptr()).value } + } + + /// Remove this entry from the [`RBTree`]. + pub fn remove_node(self) -> RBTreeNode { + // SAFETY: The node is a node in the tree, so it is valid. + unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.root= ) }; + + // INVARIANT: The node is being returned and the caller may free i= t, however, it was + // removed from the tree. So the invariants still hold. + RBTreeNode { + // SAFETY: The node was a node in the tree, but we removed it,= so we can convert it + // back into a box. + node: unsafe { Box::from_raw(self.node_ptr()) }, + } + } + + /// Takes the value of the entry out of the map, and returns it. + pub fn remove(self) -> V { + self.remove_node().node.value + } + + /// Swap the current node for the provided node. + /// + /// The key of both nodes must be equal. + fn replace(self, node: RBTreeNode) -> RBTreeNode { + let node =3D Box::into_raw(node.node); + + // SAFETY: `node` is valid at least until we call `Box::from_raw`,= which only happens when + // the node is removed or replaced. + let new_node_links =3D unsafe { addr_of_mut!((*node).links) }; + + // SAFETY: This updates the pointers so that `new_node_links` is i= n the tree where + // `self.node_links` used to be. + unsafe { + bindings::rb_replace_node(self.node_links, new_node_links, &mu= t self.rbtree.root) + }; + + // SAFETY: Now that we removed this entry from the tree, we can co= nvert the node to a box. + let old_node =3D unsafe { Box::from_raw(self.node_ptr()) }; + + RBTreeNode { node: old_node } + } +} --=20 2.43.0.594.gd9cf4e227d-goog