From patchwork Mon Feb 19 11:48:13 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203036 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1239535dyc; Mon, 19 Feb 2024 04:10:58 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCUkKbijy+mlZFcSJBtlRFypYZCbC/QGrMEyYWpHQHj2oIQKneB8P2R8g4O7xEizkno9ESuP9Ht/aDpJxhh0NDkrPY6KQw== X-Google-Smtp-Source: AGHT+IE8t0jUWobkq6G7l1SIGkMq4cHG8IPfkpbdRzYkbe/fn007UHTBdyFQ3FA0k/obJ5a/3uRn X-Received: by 2002:a6b:e418:0:b0:7c4:808c:79cc with SMTP id u24-20020a6be418000000b007c4808c79ccmr14206887iog.16.1708344658688; Mon, 19 Feb 2024 04:10:58 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708344658; cv=pass; d=google.com; s=arc-20160816; b=Pk3dEapWWa4gMArQ5Lyrd36XRhEJcWtV4C3m5GEgAg5Ae9+tgtoSR4lijDtR8l67aZ fzQMlwcnpvspIWBDgrOoLHCRVCvmgh0HYjrlXM+QKCT7D1bUnDjsIBVWFEzmGknZQbxz kWpByPKHt11qpPR3DZdhpjj2ZKpPLr/U3zPb8i2E7dWq7gzrCM3Ifpx/8lVyvkYURAR8 QhHERrRMnQl3hx4GxAvxQN+oPbUGeKVtkwb0v4sp43dTsB2BddLdE3kbVE1guIvFiln8 FWJDwghYKPFMPOFjeGm7/SXzmTlNv9UBihdZvnhV539pyug+KJhKVcMEOBrAn8Vz/Zmd 8F4Q== ARC-Message-Signature: i=2; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=arc-20160816; h=cc:to:from:subject:message-id:references:mime-version :list-unsubscribe:list-subscribe:list-id:precedence:in-reply-to:date :dkim-signature; bh=2KqfkKTQJigXtQF4MQj5aQ7HbV15bzof55IeTcnSvrE=; fh=WbR/7G/9icIIgJuxncSxBxta0O6JFnnFU6hqY+ahi/w=; b=pSLGq/WiZWGuxJGM+1v+id7zZ1nyaO+lwUlAExY3dmAzeH1M+x7P8UnkebxFU8mOB/ rB4tSmtcpRDbxMpeHaQFFO8fLXEAaiOGB8em8ZMwYjdq5AIZyoXp7bTIFHlis8W7s3Uw OisBPDE7nzsdVbpc/reyU5FJ+vZcd5zwGcC4VK4/Kw6mbsEWGoziTCkMqXksInyLjAwI +M9v7/vS1pacyeyjPAlq7CoRHn8CHIrGFiJw5iCBRqCod00kNQWl6hk8iSJA2Nu59XPi IW4sp6/hwtI7btCfAmCdofq0lyCMC0CHPdzLv26T0g3xzZ5d7D3bMZxoDLAnuam/xYbe g5eQ==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b="3z/9me1T"; arc=pass (i=1 spf=pass spfdomain=flex--mattgilbride.bounces.google.com dkim=pass dkdomain=google.com dmarc=pass fromdomain=google.com); spf=pass (google.com: domain of linux-kernel+bounces-71246-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:40f1:3f00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71246-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from sy.mirrors.kernel.org (sy.mirrors.kernel.org. [2604:1380:40f1:3f00::1]) by mx.google.com with ESMTPS id l6-20020a656806000000b005dc4b41e2d0si4295497pgt.591.2024.02.19.04.10.58 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 04:10:58 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71246-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:40f1:3f00::1 as permitted sender) client-ip=2604:1380:40f1:3f00::1; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b="3z/9me1T"; arc=pass (i=1 spf=pass spfdomain=flex--mattgilbride.bounces.google.com dkim=pass dkdomain=google.com dmarc=pass fromdomain=google.com); spf=pass (google.com: domain of linux-kernel+bounces-71246-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:40f1:3f00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71246-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from smtp.subspace.kernel.org (wormhole.subspace.kernel.org [52.25.139.140]) (using TLSv1.2 with cipher ECDHE-RSA-AES256-GCM-SHA384 (256/256 bits)) (No client certificate requested) by sy.mirrors.kernel.org (Postfix) with ESMTPS id CE669B20DC7 for ; Mon, 19 Feb 2024 11:50:12 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id DA6F03418B; Mon, 19 Feb 2024 11:48:50 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="3z/9me1T" Received: from mail-yb1-f202.google.com (mail-yb1-f202.google.com [209.85.219.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 B413B2E84B for ; Mon, 19 Feb 2024 11:48:42 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.219.202 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343325; cv=none; b=n3UITNrKhFzc2i1YDPCgCYNYIVgmWf4hJnjIySMB4q742i47gDlExLl0Rh1egEUhZq67DJsI3wi0V39GLIJVJpCn39Dkvu3N7hfxGEnPXhynxhCmBdvL168sUvJJTLIpphvxwyBrDBXXkegxGG1UNe69Kn7vnqUgELrHAObodWQ= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343325; c=relaxed/simple; bh=XlC9T+CJckAuYVPNaoXXuIM2PiPS8NXGT3y0ZR8qbXs=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=m/wgrSKNX+WgU3PBda7iR3CngysDkixhdgNZFmsHrJwuRHvzzPhsXte0FLowIRH31MSrbPPXgtDi6I2uKK/B1pApWodEwPSKJEL+ONt6kvxtvzHpggYZ0G8yT2sSal8w5ijSWJtvb2+xRGvW80pNe5F6Znf8FUrja3UGYBcL7jI= 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=3z/9me1T; arc=none smtp.client-ip=209.85.219.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 Received: by mail-yb1-f202.google.com with SMTP id 3f1490d57ef6-dc74ac7d015so5326038276.0 for ; Mon, 19 Feb 2024 03:48:42 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343322; x=1708948122; 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=2KqfkKTQJigXtQF4MQj5aQ7HbV15bzof55IeTcnSvrE=; b=3z/9me1TVq5wY+GkDG8pdnIDtOrzeZJDOc62tQ4auaPgh3fRzzy3x5cZa/6r8N6HjZ v298bG6IuDIJEi5x2NsrdwvsClHHu0th2ybdXa/0WDwh8QtQhHT0/7rTBi87hdOTHaQO 3pxId/pnzaeBNBI1KGFy0KkFtZaf3k7V203NAIK7GjXNBA8bP1UjExmlez9jwJKtnO3r +1v+6m+kSBllL8HbTFe/dXCQmTjmp6Yfz/NBr/kU2h6VlXppCNKAtzizP1/nbvRzshEf YmrUMqkKKxj7EJoen2Tazo/Yoy164vL3v1dew/CZ1kvW+1in2d/4NGmr6TkEX8aNcbID o7xw== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343322; x=1708948122; 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=2KqfkKTQJigXtQF4MQj5aQ7HbV15bzof55IeTcnSvrE=; b=Q+EULdyrSf4pAlWRLqcJzBO1GWWilpB9t8t7NWVvlgcSc5eZmUY2lmm0WB7g1hAUzP FOStc4iRbU0NPPWnoCAgR5pxcV3vuJk//SYjwGdJVNulIosxoi/qmN4myLx05pT1RF06 Q4hbGz0qBQnWESS0e7e+SsWwLSVJzeXSUlCWuUXHqiZGijDVUIanYfmmcYBmZC6ZDkUZ rBKrnYJDti919UjITl7YMfQzoPGLEToVc0vmMs8SuIsNWmysemBzNKeAAojF1WlvaYA3 ovl3o3eJ9I9ahT6o4r/wY9b0ATJbdxH1VGBnBArZnTmHyjc9Yg+7Yr3xS2K77Vo7jjye jxeA== X-Forwarded-Encrypted: i=1; AJvYcCW+tu0wnC/jcNy93cdfgDC+bRUQ4wCraxL92LlKFSyx2ys3MwKyLkmPeOFwAc15phkxPiZEMBA+h+vVLRvoaPvSFqVSJCwRYU3pF9y9 X-Gm-Message-State: AOJu0YzDpcQc5o1UF7FzJQSnYEsEl6tzvc1kyiqp3fdIQvJdQIxqrgST UK3c/dPn/5OCxnIPK4rpjwYxBbVwHJiWuy/PvChssadqueL4E66C6NVvsfMFQd6h7wDVIIE6Zwz g1iikI2zb5Qb74Aw7A+8J63PQTg== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a25:9846:0:b0:dcd:c091:e86 with SMTP id k6-20020a259846000000b00dcdc0910e86mr478510ybo.13.1708343321754; Mon, 19 Feb 2024 03:48:41 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:13 +0000 In-Reply-To: <20240219-b4-rbtree-v2-0-0b113aab330d@google.com> Precedence: bulk X-Mailing-List: linux-kernel@vger.kernel.org List-Id: List-Subscribe: List-Unsubscribe: Mime-Version: 1.0 References: <20240219-b4-rbtree-v2-0-0b113aab330d@google.com> X-Mailer: b4 0.12.4 Message-ID: <20240219-b4-rbtree-v2-6-0b113aab330d@google.com> Subject: [PATCH v2 6/6] rust: rbtree: add `RBTree::entry` 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=B8n?= =?utf-8?q?nev=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 X-getmail-retrieved-from-mailbox: INBOX X-GMAIL-THRID: 1791329208849200405 X-GMAIL-MSGID: 1791329208849200405 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 Tested-by: Alice Ryhl Signed-off-by: Matt Gilbride --- 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 6c3b10e46517..b564e4eb733c 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -298,56 +298,64 @@ pub fn try_create_and_insert(&mut self, key: K, value: V) -> Result) -> Option> { - let RBTreeNode { node } = node; - let node = 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 = unsafe { addr_of_mut!((*node).links) }; + match self.raw_entry(&node.node.key) { + RawEntry::Occupied(entry) => Some(entry.replace(node)), + RawEntry::Vacant(entry) => { + entry.insert(node); + None + } + } + } + + fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { let mut new_link: &mut *mut bindings::rb_node = &mut self.root.rb_node; let mut parent = core::ptr::null_mut(); - while !new_link.is_null() { + while !(*new_link).is_null() { + let curr = *new_link; // SAFETY: All links fields we create are in a `Node`. - let this = unsafe { crate::container_of!(*new_link, Node, links) }; + let node = unsafe { crate::container_of!(curr, Node, links) }; - parent = *new_link; - - // SAFETY: `this` is a non-null node so it is valid by the type 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 the type invariants. - Ordering::Less => new_link = unsafe { &mut (*parent).rb_left }, - // SAFETY: `parent` is a non-null node so it is valid by the type invariants. - Ordering::Greater => new_link = unsafe { &mut (*parent).rb_right }, + // SAFETY: `node` is a non-null node so it is valid by the type invariants. + match key.cmp(unsafe { &(*node).key }) { + // SAFETY: `curr` is a non-null node so it is valid by the type invariants. + Ordering::Less => new_link = unsafe { &mut (*curr).rb_left }, + // SAFETY: `curr` is a non-null node so it is valid by the type invariants. + Ordering::Greater => new_link = unsafe { &mut (*curr).rb_right }, Ordering::Equal => { - // 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 caller may free it, however, - // it was removed from the tree. So the invariants still 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 = curr; } - // INVARIANT: We are linking in a new node, which is valid. It remains valid because we - // "forgot" it with `Box::into_raw`. - // SAFETY: All pointers are non-null and valid (`*new_link` is null, 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, + }) + } - // SAFETY: All pointers are valid. `node` has just been inserted into 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 manipulation. + pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { + match self.raw_entry(&key) { + RawEntry::Occupied(entry) => Entry::Occupied(entry), + RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }), + } } - /// 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) => Some(entry), + RawEntry::Vacant(_entry) => None, + } + } + + /// Returns a reference to the value corresponding to the key. + pub fn get(&self, key: &K) -> Option<&V> { let mut node = self.root.rb_node; while !node.is_null() { // SAFETY: All links fields we create are in a `Node`. @@ -358,54 +366,30 @@ fn find(&self, key: &K) -> Option>> { Ordering::Less => unsafe { (*node).rb_left }, // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Greater => unsafe { (*node).rb_right }, - Ordering::Equal => return NonNull::new(this as _), + // SAFETY: `node` is a non-null node so it is valid by the type invariants. + Ordering::Equal => return Some(unsafe { &(*this).value }), } } 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 }) + self.find_mut(key).map(|node| node.into_mut()) } /// Removes the node with the given key from the tree. /// /// It returns the node that was removed if one exists, or [`None`] otherwise. - fn remove_node(&mut self, key: &K) -> Option> { - let mut node = 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.root) }; - - // INVARIANT: The node is being returned and the caller may free it, 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) } /// Removes the node with the given key from the tree. /// /// It returns the value that was removed if one exists, or [`None`] otherwise. pub fn remove(&mut self, key: &K) -> Option { - let node = self.remove_node(key)?; - let RBTreeNode { node } = node; - let Node { - links: _, - key: _, - value, - } = *node; - Some(value) + self.find_mut(key).map(OccupiedEntry::remove) } /// Returns a cursor over the tree nodes based on the given key. @@ -1064,3 +1048,167 @@ unsafe impl Send for RBTreeNode {} // SAFETY: An [`RBTree`] allows the same kinds of access to its values that 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 different node (i.e., with + /// a different key and/or value). + /// + /// The existing key and value are dropped in-place as part of this operation, that is, memory + /// may be freed (but only for the key/value; memory for the node itself is kept for reuse). + pub fn into_reservation(self) -> RBTreeNodeReservation { + let raw = Box::into_raw(self.node); + let mut ret = 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 has 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 occupied. +/// +/// 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`. This 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 break the ordering of this + /// [`RBTree`]. + fn insert(self, node: RBTreeNode) -> &'a mut V { + let node = 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 = unsafe { addr_of_mut!((*node).links) }; + + // INVARIANT: We are linking in a new node, which is valid. It remains 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 into the tree. + unsafe { bindings::rb_insert_color(node_links, &mut self.rbtree.root) }; + + // 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 [`Entry`] 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 it, 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 = 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 = unsafe { addr_of_mut!((*node).links) }; + + // SAFETY: This updates the pointers so that `new_node_links` is in the tree where + // `self.node_links` used to be. + unsafe { + bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root) + }; + + // SAFETY: Now that we removed this entry from the tree, we can convert the node to a box. + let old_node = unsafe { Box::from_raw(self.node_ptr()) }; + + RBTreeNode { node: old_node } + } +}