From patchwork Mon Feb 19 11:48:08 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203023 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1231142dyc; Mon, 19 Feb 2024 03:56:10 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCVIGdDyQvR5cmuUIjSmHat4SPd9wzaTtEZEqKxg2h/+D3X14csE7uIJ1F4aq19SvFXga4v8XUlQBe8BQmy0gLG1OZvFtA== X-Google-Smtp-Source: AGHT+IFNzfQFksjySF2oAxBkoCvdxZJxrusFNa1oNdNUm3v0apivvZMYrNE3NXCac9rHwrQjkRHV X-Received: by 2002:a17:90a:df08:b0:299:7637:3dd5 with SMTP id gp8-20020a17090adf0800b0029976373dd5mr2251006pjb.36.1708343769946; Mon, 19 Feb 2024 03:56:09 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708343769; cv=pass; d=google.com; s=arc-20160816; b=wKXSow66Qpp3C8+KaLGep9fdGM6NbojlHDkWhm+9YXTF/ZQOHanmbT4vDVyi/rjtjW nIPYTgYIfrO6hwxg4t3vTday27a5KT5IN2kwkQJCsxO0VKeUb47pwcX+Lz58ekd4e95Q VYK+UtrFCL1xHPa0z34aCmtlAEV3MQ2uohaWAlk0krWEcTGZq1/FERG3YDip6vYAonpT jZWeX1jqXaPXljiARVJZrGWxfDVb8YjzvgWuPehBBOKdWZlB0eS1eGaXUVzIrCYKckOS PujN7nZkPVr9PRlmaGygFeQ8q0EdD9swh+UIywytt0smXXkD2+lDM7h47banJkZP1c5u PIaQ== 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=kQYC/9pHL/SpAUfLef4ziz+f1NtXPaV0zydD9EqqSTw=; fh=Qb0oVewIb5/OQu2TNeOJkeuAz7A6pMS/MDSQWrw8gUo=; b=nMVn1wqzCC/gCBEG8cP8rSr/a0a/Wgs1pehg4+KYFtS0Fv9SRmei7lXrePJqnwDnEq cdW1q5cBFu5PE8QxLn9FV9tlcXA2Dc6KkYEO6MQrwKi+sa3AXUUMuiBGz95YLGVeuZ7V 5Fxz6LTxFxpg2XkE51oMnev9GD3zRjErJIhK+izALkn708robFFj03pakYUqAaKCC7Fe BXk79E8gL0Ktc2N+617Qj8A/otqZaDXgq6KcSnKf74CuzytWPrBfWynW3uMXlYwXq0RT KA/aq3nwL+Wzlj4SsXzF4/nVB2jZb2TMeFsF1woT+l5jLdTObfjbudvRLbKwwyf3sdd9 Lfww==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=qtDFEN0+; 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-71241-ouuuleilei=gmail.com@vger.kernel.org designates 139.178.88.99 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71241-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from sv.mirrors.kernel.org (sv.mirrors.kernel.org. [139.178.88.99]) by mx.google.com with ESMTPS id oe11-20020a17090b394b00b002998c912f21si2202628pjb.14.2024.02.19.03.56.09 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 03:56:09 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71241-ouuuleilei=gmail.com@vger.kernel.org designates 139.178.88.99 as permitted sender) client-ip=139.178.88.99; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=qtDFEN0+; 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-71241-ouuuleilei=gmail.com@vger.kernel.org designates 139.178.88.99 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71241-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 sv.mirrors.kernel.org (Postfix) with ESMTPS id 89353285582 for ; Mon, 19 Feb 2024 11:49:06 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id 2A0F02E410; Mon, 19 Feb 2024 11:48:40 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="qtDFEN0+" 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 A2FF72C69E for ; Mon, 19 Feb 2024 11:48:35 +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=1708343317; cv=none; b=L5dfRy+r/QiSFNDJLY4vxTYAfpASaj0+p0g9utJ6iYbpKaQGpqKH2KykOmxa2fYPvN1pMnIfngedVdMlbFW8arXlWQMGZwvldfp5aVbbZT1jmuG4GcOrj11TDv+hIaFAYtHEduHg5gf5o0ruGgWQUkFdjFksb+p9JAsfzm+3jBM= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343317; c=relaxed/simple; bh=innSX49CwmZV7ZYVg2pc4oPmPaVu0uYAbVQ1SRQo7W0=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=G8TK2uOPRDcc37SaY8VfDAQKHQLwZmoM8zgsByKIJO3qTvLb1azs2XS99eIY0q5e0PgKXBsgi48DM8n8WUf1FjFZPxYMqurWu/CXwCu6Ll1/+HwcUdScYy6keGhUel6tBgKZiPEzupAFhEGslGlF3h2n8YkXzWwKzyNrpPMGJHE= 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=qtDFEN0+; 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-dc64e0fc7c8so5875207276.2 for ; Mon, 19 Feb 2024 03:48:35 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343314; x=1708948114; 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=kQYC/9pHL/SpAUfLef4ziz+f1NtXPaV0zydD9EqqSTw=; b=qtDFEN0+epX1qOO3GYygJ39pJdj5Ny8r6NzpJLVSlMIEHfkXz/4jPaPXOaFZ8isRE6 RUSEarDOJblEIExN0iwOxMpF0aBc2QSNiEPhGLHzhRTTfYRCCLX05dUrlcbM1OvlnEyw LVZyTftuKYD1XeKcqkoofr/ryhRHINwa84e0VVLED1T78gmu/og9XGYnP04Q9zTtAn8p Xb0Szzn6eLeeYlBkI1wVtpFZahMaZ0NcGpyp8vB7GP1XhXKp4PjGENa1uW0B3e1+T5J7 AdqXGgLS4YdOQhOs0n0FYWMSB90FZvu4/YGfYhclP34XUGDmC2JeROjcldlUBXGvziq8 StXA== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343314; x=1708948114; 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=kQYC/9pHL/SpAUfLef4ziz+f1NtXPaV0zydD9EqqSTw=; b=GBM9gHr2/gkfjUz0P1Jh9qh7zczmkqjP7hxH9Zm01naW0ie6oBIFUSNPsGPNaoCIIa 1VgDHANEfn9eMjspRrbxgSFK9W0rIkhWPoxnR3KgeVAb2JLy2ifIza4ggxQJygWHlgQB 1psy6nB5bwOUUgbE926rDxJ0puLjrJYyfjH92nW1jKzllE0PV/AvR7ltwGyr/M6raNuT VW9+dFTZid9wQzDxF4ewx4pkG/ENSt5ehgWerLnYxR/bkOmcOv4oXDXbrx3oWyEiH2T4 sO0Zdhxd+th1GcsyGY8+y9pTDFzqzwZ4Rh6NP5ViWXkZAYl6F7d+JkTAA4orTOh8PBsb koLw== X-Forwarded-Encrypted: i=1; AJvYcCXE9ST9vo2/9beBH7396dy9g6E9EgyMhMHO3AJBpfcPdVFyEZUiPud2TN7C8apYaIhPsCmbUzzTwoGUv7VEnxeaMvui+rMmcjgjPbBh X-Gm-Message-State: AOJu0Yx5zHF4ZASX0gFfY0Fx+x9O0hiG0stg4XUbj3BcCpJffdLnyQsX YKM3vycx9URhP6yTt7WgQLyLqeCsr8x7xJHX213iX7A/BqedKAOes9K0sfs7wgh4H7O04YxQs7O zO4j9ulYx88rzihUttqKc0HHeKw== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:1b85:b0:dc2:26f6:fbc8 with SMTP id ei5-20020a0569021b8500b00dc226f6fbc8mr479735ybb.7.1708343314722; Mon, 19 Feb 2024 03:48:34 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:08 +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-1-0b113aab330d@google.com> Subject: [PATCH v2 1/6] rust: add `container_of!` macro 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: 1791328276571707119 X-GMAIL-MSGID: 1791328276571707119 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 Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl Signed-off-by: Matt Gilbride --- 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 macro, must both be in +/// bounds of the same allocation. +/// +/// # Examples +/// +/// ``` +/// # use kernel::container_of; +/// struct Test { +/// a: u64, +/// b: u32, +/// } +/// +/// let test = Test { a: 10, b: 20 }; +/// let b_ptr = &test.b; +/// // SAFETY: The pointer points at the `b` field of a `Test`, so the resulting pointer will be +/// // in-bounds of the same allocation as `b_ptr`. +/// let test_alias = 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)*) => {{ + let ptr = $ptr as *const _ as *const u8; + let offset: usize = ::core::mem::offset_of!($type, $($f)*); + ptr.sub(offset) as *const $type + }} +} From patchwork Mon Feb 19 11:48:09 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203019 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1228956dyc; Mon, 19 Feb 2024 03:49:59 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCWXVyUz51XQYVM27pJPjI0T2GaRaqmm99F0tnSYeCYHd7QuH01arWJi+15wOgLx8v53ut6yaHAsqZvDyCwc9JJ+4tDVaA== X-Google-Smtp-Source: AGHT+IEwawFzxrseqH8YMNqeXMyuNnMvjqZXPiQO+sTwV5HqLPtjHC+61AytvvGDRbkIEcq8lO/l X-Received: by 2002:aa7:c896:0:b0:564:21fd:2710 with SMTP id p22-20020aa7c896000000b0056421fd2710mr3573516eds.37.1708343398932; Mon, 19 Feb 2024 03:49:58 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708343398; cv=pass; d=google.com; s=arc-20160816; b=howD5glCbQin8gsQKqf2LfyxwPiuolnW/uoB0xqLt31UIwO/sLbAJGzJb22aq/vDSw S6em8nkZ6sIfSpGuVKJSlzPg/qn1B8M+UV0+SfOg4tZ8+5w1ENUNzqUCyZStA7FQZbhG 8eOtu3Xo+IgVRdMQcB4DybuzmjTOd/EjkDnUUR5TBIxvRni2ocKYsVA1JRX6ODFxIKeb 0oEboq0Rkx7saHN4FPmoA5VoP+U175mozHmceUKSqqIzd6/Q6AVbiTJzA9xts/yx6bSQ Mu6PSEUjXk+fE4JJFaFQW0Dl1a0FVMDQNWbLu7BYihi9PnVdPDu4xyrpRdR/wyN+Vfpg Ut2Q== 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=KEgcT/UQTYpXxnEPnTMLLl8kAcJbkORmrTw0STIurYs=; fh=HYv6p8zUlkVDqtT8LmvsT0FQhIJwV5+oZEwcY6Kj3Js=; b=xtjRBQfvlTdHThlzIiQ4IQusjhzaDCE6D3IqgU9pZBYN2YJ1S8CJe2rBH4ZYlY7DW2 FIRlQXfPlwVZlQF+qnYoH3WhX8IwlldOLglr6eHw9hUxG3oZMNjMBmgb0RjR37ZrwR9O a8sbnif/F/+VWE2j0+7IwELhlBrOdPyV0nhKdXwV8Nkfs3jXZ1tqtt5iM/llJy6GPQIg nQ9uhS0xkUIsDuLufpClpgaCtC9ZmHJ73H3YuH0BQcnrIG6V/PrA0c3VH/cMCLVtByfw 8B5ITNcu5Rl4Desvy46a9fjUb2qkK6JNeQBehBFYG+RRuLkTX1nNzj0IUpKC7xHK+5QP gArQ==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=xQEZe4+k; 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-71242-ouuuleilei=gmail.com@vger.kernel.org designates 147.75.80.249 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71242-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from am.mirrors.kernel.org (am.mirrors.kernel.org. [147.75.80.249]) by mx.google.com with ESMTPS id co13-20020a0564020c0d00b0056464ecc409si1119463edb.124.2024.02.19.03.49.58 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 03:49:58 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71242-ouuuleilei=gmail.com@vger.kernel.org designates 147.75.80.249 as permitted sender) client-ip=147.75.80.249; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=xQEZe4+k; 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-71242-ouuuleilei=gmail.com@vger.kernel.org designates 147.75.80.249 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71242-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 am.mirrors.kernel.org (Postfix) with ESMTPS id 62B031F263B0 for ; Mon, 19 Feb 2024 11:49:58 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id EAD9B2C85D; Mon, 19 Feb 2024 11:48:46 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="xQEZe4+k" 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 1D58C2CCD6 for ; Mon, 19 Feb 2024 11:48:36 +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=1708343320; cv=none; b=f5WsLULiZZE+r/g4nR5kb5N8+w/X+k9fhyJ05v0PocM5kl2y1od9HcXLPYiveJrk3xoQzPAP4UeROU3qK5WahsDKZyTbIqYHRCn8XiruAGaPyEQP5UqVLYbBUtp28USlRWjWVW1U5KgnvwQE1efOE2qbeHF3ygayt4K9djQZlvM= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343320; c=relaxed/simple; bh=6SR52JAdmmCs4J7JswxAXBtZzhK6mj+MtYHKlaLG4L8=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=bvHeMjXrzQsjO/cTlLjyuUJEE+BqJHpt2kj0O/dyjg0TwIrt06Su43VXaYE3EliUBC0sNQHPQ2lvNlQY47Y7MvcuJxQumt0G0HOZbO7Pl86xEFYrIbxUuQSOpTTbhRy5TqN3/KVU6hW4O7rfSBZQMrolinYhOdCc4gh3AFYUHRw= 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=xQEZe4+k; 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-dc6b267bf11so3062801276.2 for ; Mon, 19 Feb 2024 03:48:36 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343316; x=1708948116; 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=KEgcT/UQTYpXxnEPnTMLLl8kAcJbkORmrTw0STIurYs=; b=xQEZe4+ktHQwcEcbs8D7xbad5eIMENaWWofuU0nn/fUB/o2yqFznAA0JuwspI+3Jxj huvGV37CDQ/5gxUZnttTNqL/921xbm59ZmJc7T2847OzY+wIhveShQavwORfl4l87CYy tia7PDHLCaSWXAMu8/Z0DGt/cuKqN3yOwkI4Xxc/SxJStKVwQrCtU4Y8/ySCYjQUrZj8 EL/DFJV2xK1sQCOzbwsWyhPl/WnSZUi/u771SgbjkHScLn1ionXUEtJ1iQ4PfzI2U244 giwLk3T4iYCxFWHLDaB6BEcGju4HfCw3D5ATJ2gBVxDKzkUMM1x/voLuH9R+sik63rFB zYQA== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343316; x=1708948116; 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=KEgcT/UQTYpXxnEPnTMLLl8kAcJbkORmrTw0STIurYs=; b=iojImM/x8ipSuco6oV2tf56FCnrwiYHRwEwjST/evh+EN/wKXeyhBHw2nluz7BDpQp tYX0QUe/2N+05PxE4+L6Fwcigah71YBDadH/p9t4ZGPUAg6HKYcShT8nSx3aA+rYyxIh UMvJubchAE9h2hvVMl+hgdCVjs4IwyaT67aHxKdxbDUSk017HCBDaNKAOfs5vqMHGRRw 2QsfVnR3O/aY9Y7XhSgolD1XHZCp/O2N6aumxGoQhOtt1l2q5Gqj92CM5VE5cIrWJb7t PXyZiDdiJTect8meBT3JyBEaMxnhTGVGrvoePfIiUTU5/I8RSTQm6cb6KFH1aVIoI4+n 7B9w== X-Forwarded-Encrypted: i=1; AJvYcCWoS0mpOiS9PTDFqHKjpix5AhHkZvCjKd+6hRo9MCRRsQ+1l6nQQjbPpXJVD0VK/Mfzxl2XbUz9V/UVsgM0LvxIfmKukKd6FVsVqlSq X-Gm-Message-State: AOJu0YyPpaM66OD6c5I3oXLGURHgJ8gbGdzwX84lDZLXOMzuSaU113aX VPZfKAUl4dLDNkDHEDv+DmOoSwHfwT7Lq/ByRdM1qUaroO/hgxdupv/KdDdsmHATykg/Kahf/KS BZBg8nwXakfWaQBwEdKOC/wR3MA== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a25:b18d:0:b0:dc7:82ba:ba6e with SMTP id h13-20020a25b18d000000b00dc782baba6emr447391ybj.7.1708343316077; Mon, 19 Feb 2024 03:48:36 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:09 +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-2-0b113aab330d@google.com> Subject: [PATCH v2 2/6] rust: rbtree: add red-black tree implementation backed by the C version 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: 1791327888074045303 X-GMAIL-MSGID: 1791327888074045303 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-08ba9197f637@google.com/ [1] Signed-off-by: Wedson Almeida Filho Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl Signed-off-by: Matt Gilbride --- rust/helpers.c | 7 + rust/kernel/lib.rs | 1 + rust/kernel/rbtree.rs | 404 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 412 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); +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) indices. 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..a72e9f57e660 --- /dev/null +++ b/rust/kernel/rbtree.rs @@ -0,0 +1,404 @@ +// SPDX-License-Identifier: GPL-2.0 + +//! Red-black trees. +//! +//! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h) +//! +//! Reference: + +use crate::{bindings, error::Result, prelude::*}; +use alloc::boxed::Box; +use core::{ + cmp::{Ord, Ordering}, + convert::Infallible, + 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 node. +/// +/// # 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 = 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() = 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, then insert the node into +/// the tree. This is useful when the insertion context does not allow sleeping, 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 = RBTree::try_allocate_node(10, 100)?; +/// +/// // Insert node while holding the lock. It is guaranteed to succeed with no allocation +/// // attempts. +/// let mut guard = tree.lock(); +/// guard.insert(node); +/// Ok(()) +/// } +/// ``` +/// +/// In the example below, we reuse an existing node allocation from an element we removed. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree = 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 = 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 different key/value. +/// let reservation = existing.into_reservation(); +/// +/// // Insert a new node into the tree, reusing the previous allocation. This 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 that 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 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 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 inserted into the tree via a + /// call to [`RBTree::insert`]. + pub fn try_reserve_node() -> Result> { + Ok(RBTreeNodeReservation { + node: Box::init::(crate::init::uninit())?, + }) + } + + /// Allocates and initialises a node that can be inserted into the tree 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 returns 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 returns 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 } = 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) }; + 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() { + // SAFETY: All links fields we create are in a `Node`. + let this = unsafe { crate::container_of!(*new_link, 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 }, + 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 _) }, + }); + } + } + } + + // 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) }; + + // SAFETY: All pointers are valid. `node` has just been inserted into 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 = self.root.rb_node; + while !node.is_null() { + // SAFETY: All links fields we create are in a `Node`. + let this = unsafe { crate::container_of!(node, Node, links) }; + // SAFETY: `this` is a non-null node so it is valid by the type invariants. + node = match key.cmp(unsafe { &(*this).key }) { + // SAFETY: `node` is a non-null node so it is valid by the type invariants. + 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 _), + } + } + 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`] 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()) }, + }) + } + + /// 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) + } +} + +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 = 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 = unsafe { crate::container_of!(next, Node, links) }; + + // Find out what the next node is before disposing of the current one. + // SAFETY: `next` and all nodes in postorder are still valid. + next = unsafe { bindings::rb_next_postorder(next) }; + + // INVARIANT: This is the destructor, so we break the type invariant during clean-up, + // but it is not observable. The loop invariant is still maintained. + // 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 that 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 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 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 = 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 = Box::into_raw(self.node); + RBTreeNode { + // SAFETY: The pointer came from a `MaybeUninit` whose fields 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 that 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 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 {} From patchwork Mon Feb 19 11:48:10 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203022 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1230758dyc; Mon, 19 Feb 2024 03:54:56 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCXaz88o8DiY5FhjA1mJjyGM0HPvTK1OAYKNwkGRQ3HonF/FAPPa6YGaUaYiljwgCGPoeTs+q2W60sKv2QUGad4o6H8AvA== X-Google-Smtp-Source: AGHT+IG0kj2nDG/UcqHCJvOoTEjeWIm4kJI5yDEbjtc4WPQPliDi7C4ePuzwZCpcVTn4sQo0jDOU X-Received: by 2002:a17:90b:315:b0:299:7d1a:7ad0 with SMTP id ay21-20020a17090b031500b002997d1a7ad0mr1996724pjb.32.1708343696807; Mon, 19 Feb 2024 03:54:56 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708343696; cv=pass; d=google.com; s=arc-20160816; b=bjhesbG5dYNqFY2j3roB3CsmUuxmhMbcZ1IrVfhEjLcjiqgNLdyAdVs9TZB6sIvGOP mj3465B69wlRkYeu9pwK4X3MeMm7KDcaR0koHJvbU0FOgZnZWfi7iWqeFUP0iRIwx1Zr mMXwgz2qv1WAljrprjjy9mUZGn9fFAlTEw4RlQJePiDmv/8EZcYs8M3GvzZ+oE67gpci umtnP7WHqYPPId/gkpUgwxqn3PaHrg0jHNQfGQ6Y5nkN1EhCJYVHHeC4y8Y3awjyMszT slD6Y/WS4GvoTN3qMkiOovh8pelpfxDxhqmtroBQvh5HfoRSycXPQXR62SdKb6geF6YP lk3w== 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=eQCD4AuKZId7gMhOchvG2XGZrBs0B56ew/r04xotUCQ=; fh=N10nev0sIqAT2OgN9GU0t6VGMsl4zMq0Pz5qHR1F8Ng=; b=rUtbeeqz+phvvoVO14WE4QPNFbiRK6PuR7e3WHelPwspxdg4IWKGIONiz4UGhOGNoR K+spEgjZiSbCtkTXf0dthexAC428xnVDjKAjUZsCvYr3J2ijdOvP7+84AMoMeIFvmRKC 2VbmxpRrqjwEh8yvCoAMJmRTkk8kP39EZpn8vUlCOfCeOYIDdf9jYAZcfTIYv0oOa1PG 38Yncj627/J++E47gSC2X237Uns5PgVAiwlvTtqGe00RhhyJRnqSXW1b8p+Yec6yucqo eHK7Flv73B8CiWOp9VcrLmPtQnYcIcV8cPpBheJc/SwGcJ50g1FodH3D2BCO+fD53lSc YJRg==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=DZsSTOZz; 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-71243-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45e3:2400::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71243-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from sv.mirrors.kernel.org (sv.mirrors.kernel.org. [2604:1380:45e3:2400::1]) by mx.google.com with ESMTPS id c17-20020a17090abf1100b002993718977fsi4394427pjs.91.2024.02.19.03.54.56 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 03:54:56 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71243-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45e3:2400::1 as permitted sender) client-ip=2604:1380:45e3:2400::1; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=DZsSTOZz; 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-71243-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45e3:2400::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71243-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 sv.mirrors.kernel.org (Postfix) with ESMTPS id 39E25285FE9 for ; Mon, 19 Feb 2024 11:49:40 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id 89B0B33080; Mon, 19 Feb 2024 11:48:44 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="DZsSTOZz" Received: from mail-qk1-f202.google.com (mail-qk1-f202.google.com [209.85.222.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 67E0B2D606 for ; Mon, 19 Feb 2024 11:48:38 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.222.202 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343320; cv=none; b=JerDzwu7JNLoOJDn+JKsQMfejUSuhPpGPAdGZApUD20eAbswvmVfsZ79fOoXTftnc9m8Kj08joVVYmtljRVKZ5b3+MatdgkZSDFIQ/+ea+UJQjxVvGFB+iBzOMCKzlMRMcZiGOXUOgnUK9GGETJLlHu9DB8fvnW+nLP46nqddI0= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343320; c=relaxed/simple; bh=Gsr8Np304/ZF/8MRD/n+jlVmItVNTviF52shcVw7ERw=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=EL3pnuxZplip7a9US11/a3BXwapNhzC9NkS5gP0VkzRlJZdyGQwn3kEgCnRgusHrXOv6FjuTtAToNT3Q89uQdVvTsY7TPMncsyvYt9+0+dHSH403k+UX+jiU2NCbzMkJn3ng6hXFO27J1mqlYmqjzbJwoMutlBQOK4VFfjq/wZE= 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=DZsSTOZz; arc=none smtp.client-ip=209.85.222.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-qk1-f202.google.com with SMTP id af79cd13be357-7816bea8d28so1015479985a.0 for ; Mon, 19 Feb 2024 03:48:38 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343317; x=1708948117; 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=eQCD4AuKZId7gMhOchvG2XGZrBs0B56ew/r04xotUCQ=; b=DZsSTOZzMjMot2AhkcMGcLko6acOMIDdSj8YEVYpLIsc3Rhy1M6LQhm0bSgYWyzmDI rd4BHMdalqe0XjGUsJetR2YLYWZZ5Cte3vfwxZbO9jqBeMjBMcGLmEh2Vr/zr2lFdTEL vhFaK4dNfY+GVz3fcfc0jEcPq1IiBBWtp/DutXBkTv3YEy4nqzb4cvKbqAwDwTYNCr01 iXW0qqSn/eZvM3gq38S3WFtADYkwuvs4DAGY2zkiOnP5g9bezQG6MgfQbTrGWgvLeUHm gYbEMUFNfh4dhkvPjjyXHDHFdlAL6AJv/Pr60Y1H9ANwaMnQCwxL906LtRwlujkkAF9R YI0w== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343317; x=1708948117; 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=eQCD4AuKZId7gMhOchvG2XGZrBs0B56ew/r04xotUCQ=; b=tuDBFZRUy9hY4PgTlAa+hq9/AxvBy62CblHHp/0/UlFT4Cr8XBhwhkcibQhwKg6rCC 2Q1IWzEZ5XLlEHkKq8arxYDo25CofGjd6sZai+KM9p0/bbQ0fa/DBLJ21lOip2utSeNR TQZ+IqPkCarT6ZwY3QZPmv49QgtOIdsg223StEFzD8Z/3xuIW+23xeB28Wytte0rGiI4 5tuUb/gpx5/N2BuVovBoeW3sg4gQrchem6WcBQoXbCSASnisw/CfjCRWswX0Wb6tsKy8 sAFBBJ8LZgMqrW8Qfr50EMnHSFd0ufOwRZvhiSHSWJH9kPGgx8+T0nhaY8GcB5jH3cWK 3hvA== X-Forwarded-Encrypted: i=1; AJvYcCV3Qz7BVsdpE6CAKjG/CjQeqL396d+HlPCo0pPPWllKtYPqSxGg0EE8rkv1r/eVQFW986Ea1/5hw3/SHWPRtaoGWFXYK3GCpbFIEtoF X-Gm-Message-State: AOJu0Yy9NMVqckH94lyU5aH3rHyoHdowPHADHHJEK7bcznLCinr8IRB2 +SfiEZTpi6KadcusNNX5MYwG89cOUoKJbY6DGc2bDXbKOv3R5jdgoTz4YTto/pdcyyCih6XC02u lB50Y7bdXUg6PD68wjIievmD1Mw== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:620a:40d0:b0:787:72d9:b7a0 with SMTP id g16-20020a05620a40d000b0078772d9b7a0mr3625qko.3.1708343317417; Mon, 19 Feb 2024 03:48:37 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:10 +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-3-0b113aab330d@google.com> Subject: [PATCH v2 3/6] rust: rbtree: add `RBTreeIterator` 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: 1791328199899502209 X-GMAIL-MSGID: 1791328199899502209 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-08ba9197f637@google.com/ [1] Signed-off-by: Wedson Almeida Filho Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl Signed-off-by: Matt Gilbride --- 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 a72e9f57e660..b1faac831cfc 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -54,14 +54,30 @@ struct Node { /// assert_eq!(tree.get(&30).unwrap(), &300); /// } /// +/// // Iterate over the nodes we just inserted. +/// { +/// let mut iter = 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!("{} = {}\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 = 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. @@ -69,9 +85,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 = 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. @@ -79,9 +97,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 = tree.iter(); +/// assert_eq!(iter.next().unwrap(), (&20, &200)); +/// assert_eq!(iter.next().unwrap(), (&30, &3000)); +/// assert!(iter.next().is_none()); /// } /// /// # Ok::<(), Error>(()) @@ -121,9 +140,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 = 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. @@ -131,9 +152,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 = 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 different key/value. @@ -145,9 +167,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 = 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>(()) @@ -188,6 +212,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 have a valid `self`. + next: unsafe { bindings::rb_first(&self.root) }, + } + } + + /// Returns an iterator over the keys of the nodes in the tree, in sorted order. + pub fn keys(&self) -> impl Iterator { + self.iter().map(|(k, _)| k) + } + + /// Returns an iterator over the values of the nodes in the tree, sorted by key. + pub fn values(&self) -> impl Iterator { + self.iter().map(|(_, v)| v) + } } impl RBTree @@ -350,6 +393,52 @@ fn drop(&mut self) { } } +impl<'a, K, V> IntoIterator for &'a RBTree { + type Item = (&'a K, &'a V); + type IntoIter = 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 that 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 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<'a, K: Sync, V: Sync> Sync for RBTreeIterator<'a, K, V> {} + +impl<'a, K, V> Iterator for RBTreeIterator<'a, K, V> { + type Item = (&'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 = unsafe { crate::container_of!(self.next, Node, links) }; + + // SAFETY: The reference to the tree used to create the iterator outlives the iterator, so + // the tree cannot change. By the tree invariant, all nodes are valid. + self.next = 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 From patchwork Mon Feb 19 11:48:11 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203018 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1228873dyc; Mon, 19 Feb 2024 03:49:43 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCU19hNr4Tka5ICJZdAptuIlizbjZqrE5eDf4wu+rgpAs1XwLMjrhhJOFqGIyLQB0WRSSe8gpYzK410jq9pY/mDZH9uMMg== X-Google-Smtp-Source: AGHT+IHKSoOJgH37qvYwOlLzxfYw+e7QvUzRmREMeJGveR72KR9+Al3K7HNyKJcpBB9Pj0YZmYh9 X-Received: by 2002:a05:620a:1a:b0:787:4082:56d0 with SMTP id j26-20020a05620a001a00b00787408256d0mr10189383qki.43.1708343383224; Mon, 19 Feb 2024 03:49:43 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708343383; cv=pass; d=google.com; s=arc-20160816; b=X8bpZMvNcZeOowcR4vBuyPJWcRXC+cSMvVnwpnz1SUGegY6BrZ9acVsrfIC37tFS/H /Y3sp5x6etiqL0R2wH4UcVtR7G1V0x/3ZpaifZ1C6J/IyAOrUZKwly3oMQG0vfxJUhdU xUCytkMiv1wtUnWY0A/3tH/48URIaaWixr6p/qdTQzgY5303C9FtRT014oVjlEd6awRi GhWuTW9FaH/yQRVq82we57r9fsN2cnuKsDZG/VN9Gpkda15xTIfbMa+gw7lBJRXTRQ5v V3bDYAft3J9KscnpiJbeLl/NPfIWoq+ynNNtTpOS1RdV8q9uuqZp9pbzwof8uP6SHENJ Y1qQ== 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=HB90XMvSMsoYLDGPsOGvbe6aP04mVYyG6I4g2jRdUgY=; fh=9XyANsR7bTFI6sPh/bVG6zClPTxSmOlZJTx/oggmvk0=; b=bD2PNaNLPFtwzqwqbfxWOh7jokzCYkx6/qiNtRvfU06AC7rmqI0kpite98numCdXfE DF4RSx6QODh4RAgKkYljR1ZwCEbEeouH6Cy3PMxcvbCmIJry6PnMw8YVlbVkb5BOkhBT K/ayOXgmD3hVAU44Cy3l6NXyTxC/P08dRuByrZmYSMsy5o1zMUEWCCxhEZrnPtYu21bK uoonEUqqUQR0ORTuZPnQYNLiUxtxB3acg2goPmyu2787NwQDLn1Be02Zvfu5a/yCo9BY vgFMgfIJtie2xsPrbPK9USNv1r58ycxwqtAzc4FjP/yo2SITZrfKhdACGBLGd/oJoEqj rY7g==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=cvsqMMEk; 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-71244-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71244-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from ny.mirrors.kernel.org (ny.mirrors.kernel.org. [2604:1380:45d1:ec00::1]) by mx.google.com with ESMTPS id z23-20020a05620a101700b0078553272301si5832485qkj.159.2024.02.19.03.49.43 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 03:49:43 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71244-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) client-ip=2604:1380:45d1:ec00::1; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b=cvsqMMEk; 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-71244-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71244-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 ny.mirrors.kernel.org (Postfix) with ESMTPS id EEF1B1C21FD3 for ; Mon, 19 Feb 2024 11:49:42 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id E492C339A4; Mon, 19 Feb 2024 11:48:44 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; dkim=pass (2048-bit key) header.d=google.com header.i=@google.com header.b="cvsqMMEk" 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 A651B2E400 for ; Mon, 19 Feb 2024 11:48:39 +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=1708343321; cv=none; b=EdKh5dJqgiQde4JbTXmtiuXPpl+oeFuFg77TfYRBaqH0VthO5wjnJgP5s4TvvZ9Aj5HsEVhcT1VPTJGF8Ql7LxV1mc0oCfpn8H41EuZLykcWNljE+6gdPZ1utKZuXRMT/ZQ0LZvRHQjd0K6a1dHI4FFIrX4NhercrB9JCXDUKZA= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343321; c=relaxed/simple; bh=qAEzaqvyHsbGb2Xo/aEXETVIQv22OgH/DTmjxPoN4uE=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=K9cIpZhNDcHE0JmqHIdIKeGx5wwdpqrHSPMYuLwBzdUuCyBGLk2PkXn+X1U3FIwTsMKcedb8RLp47Is867Euz/dbyBQminvvDJWIYB4HSx2jg2A976EvxqaIvSsQtM+idTmpRpSUE1AOAtgXWXF8R8NFo+51OsDSZdjyH0tjv2c= 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=cvsqMMEk; 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 Received: by mail-yb1-f201.google.com with SMTP id 3f1490d57ef6-dc64e0fc7c8so5875289276.2 for ; Mon, 19 Feb 2024 03:48:39 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343318; x=1708948118; 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=HB90XMvSMsoYLDGPsOGvbe6aP04mVYyG6I4g2jRdUgY=; b=cvsqMMEkSny5JvAufSXCKUJlYIq1T5Ws/TjioMMB2SznFA0czONrYEdwxesebqW4NS hAZsKCOBALhH37zQ/wYegZxpDmTWV+ybbF2WvyfrsOT4Z0AoVkNySwdeeaDUFtTTm5DL 31OlNLFshvgKUQ9pRZuTc+RGiJsb09zV+5Tq3NGqCZklN5yd6sk+9vCrKAKQ2+wCYsbB qwBQvLknQra+q/Q2ZqUXzRCQTuEYlcU5UAzS4NSitVkpTuEeMLkykwpLAk088a7yZ31t t5uUZmHJ2n8vrjmaEu12y6FZAWxEf6OYZjb1JT6vXTuBCrfDVm2jEb1oAm7gXTwQqt9W KJig== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343318; x=1708948118; 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=HB90XMvSMsoYLDGPsOGvbe6aP04mVYyG6I4g2jRdUgY=; b=Uq7khSdT/RW2hXPT8SiedUnWgfI/BfFqjykWtsSZ1upm2pDKEvSvCxdNGZ0jhzBxRN CivwICriG6P2fwtrGQ3m3tgjdQqjxRdzembnmp4GWdJX+fHbTQGzXr9DH0CjyHPIqhoB PZ8W+/w8Hwdwb3rXbYygSp9sVKjU5FrKUE3aVVRRXN9sOFCXusH12zVR4Jne0i9MUQz3 HvkT62S5WKksMi9+L7nigQTa2VQjkzB+Xmp1cRBOxfPS5FL+aG1IW/HnEcYIBeqUkdNw 6ITWljCX97Wv98uANEdhDw7gtEarI9V1rxR3kIHpvVRv4EI6HMUUrcEymZu0AK2zp3/F ZLqg== X-Forwarded-Encrypted: i=1; AJvYcCUOfg1lnbPQAjuPg7PmBOiQyE6FVxrp4HDCtM32HXYVf95U0LqEPwxwXQErQmRhAfpMB3pDPgtw9YiQoysNoLMmOrb2e8W6tMFgYXK/ X-Gm-Message-State: AOJu0Yy6IY+O3lnRvSpE1baPZpiAhVAXz3i0cGX5E3S8d57qo7OPcyLU wqpGIDIHiTJ5wHfDMk9pdpBW9wH/vygErS+kUnkBzYslP45EhjZ6RljHBdQx9uJJYxE7hVZQNx+ jix4cNQ7MtCEFrceoMMWCYWWy8g== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:6902:1001:b0:dcc:79ab:e522 with SMTP id w1-20020a056902100100b00dcc79abe522mr480849ybt.11.1708343318765; Mon, 19 Feb 2024 03:48:38 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:11 +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-4-0b113aab330d@google.com> Subject: [PATCH v2 4/6] rust: rbtree: add `RBTreeIteratorMut` 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: 1791327871656958590 X-GMAIL-MSGID: 1791327871656958590 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-08ba9197f637@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 b1faac831cfc..ccf74e0dc3ec 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -222,6 +222,15 @@ pub fn iter(&self) -> RBTreeIterator<'_, K, V> { } } + /// 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 have a valid `self`. + next: unsafe { bindings::rb_first(&self.root) }, + } + } + /// Returns an iterator over the keys of the nodes in the tree, in sorted order. pub fn keys(&self) -> impl Iterator { self.iter().map(|(k, _)| k) @@ -231,6 +240,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 tree, sorted by key. + pub fn values_mut(&mut self) -> impl Iterator { + self.iter_mut().map(|(_, v)| v) + } } impl RBTree @@ -439,6 +453,53 @@ fn next(&mut self) -> Option { } } +impl<'a, K, V> IntoIterator for &'a mut RBTree { + type Item = (&'a K, &'a mut V); + type IntoIter = 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 that 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 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<'a, K: Sync, V: Sync> Sync for RBTreeIteratorMut<'a, K, V> {} + +impl<'a, K, V> Iterator for RBTreeIteratorMut<'a, K, V> { + type Item = (&'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 = unsafe { crate::container_of!(self.next, Node, links) }.cast_mut(); + + // SAFETY: The reference to the tree used to create the iterator outlives 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 valid. + self.next = 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 From patchwork Mon Feb 19 11:48:12 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Matt Gilbride X-Patchwork-Id: 203020 Return-Path: Delivered-To: ouuuleilei@gmail.com Received: by 2002:a05:693c:2685:b0:108:e6aa:91d0 with SMTP id mn5csp1229114dyc; Mon, 19 Feb 2024 03:50:25 -0800 (PST) X-Forwarded-Encrypted: i=3; AJvYcCXY/bYCqVoPha/e+21m0Ig28DuV1asYUtJY006NkyMFi2caCY180JQyBc07SCu+/nMV/s0dXdoLqthokgg9GJQp06L7Hg== X-Google-Smtp-Source: AGHT+IFe9io6EqyiCjb2Sgh5F5wO+xoTze6dzig4bik39dxeeKyBklcA6vgRTF190/UprXgvpmDK X-Received: by 2002:a0c:8b1a:0:b0:68f:5901:c76a with SMTP id q26-20020a0c8b1a000000b0068f5901c76amr4290901qva.0.1708343425560; Mon, 19 Feb 2024 03:50:25 -0800 (PST) ARC-Seal: i=2; a=rsa-sha256; t=1708343425; cv=pass; d=google.com; s=arc-20160816; b=xO8+zZmJnfj7CXaUdh5DQaFhKYLNRaQh0rBL23Gw9Lt7qJ+t9vyzbyCxSjaTL2E5E2 HO5ULhRyJ6mtGsLTgBOrzLgqBGCbkLZGyuAzOFRVDr7/Q5zcXn9G15vf0fzXhky6zLH8 7ySk7Dg4Mo2ZDt0FeST08wewWN9GOkkcEHm3+HVb3OqVJTlGUXPSdZ3PfXsxW5NYTZQZ 35Q8lFTpc7h8TOjG3cKViuKO10b7ttPuQYTqOX7MFhBVqRi1+6FMyPkJMz76C6qapwTN tkMmfzwZqNR6ia2uCMjME6yd0tjqUyzfZx1uyFKves0TO7GSUD+vEdjGf+7trawiAmWw JImQ== 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=ZwZ2a7+2VGra5HwuOBKn0CmjD9GQoqV7rwcv+kyW2rs=; fh=It1QGjzp/mJOHfXHqtvGve/DGIeUNrbKokihhkKKAxk=; b=Cho7MvA1d8JJSsNcyekDjqzBaYmJLVP7dQJCM5WfwD+H9vmCYzLLFqJCzPYBZS3a5l OvSN6E0D5WHaJshacpxZ9PBA5My/7Eaw265j0TMLhNiTr0aUtikUz6cTdVdwXUZ0rLA9 ewCv6jfAhKDQDjEYlFii6R26vjpslp4LlpjtblyDv7hzGCtzCEwTOb4pWcc1/a2r5Z/9 unfsldPuYrFKhcNwqoeaLnpeU7RwPk0LsdUJclYpURFdbe/J2uM1EHEYgtcN/hf/ovz/ OfpAnH8KyY71yn7RVAmQVOooygSMuIfvCuL70Ry+J/DACcpVpQHVOaZp7/SC9CbErZKX kRlg==; dara=google.com ARC-Authentication-Results: i=2; mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b="PINxa/8u"; 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-71245-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71245-ouuuleilei=gmail.com@vger.kernel.org"; dmarc=pass (p=REJECT sp=REJECT dis=NONE) header.from=google.com Received: from ny.mirrors.kernel.org (ny.mirrors.kernel.org. [2604:1380:45d1:ec00::1]) by mx.google.com with ESMTPS id 13-20020a0562140ccd00b0068f29d0b05bsi6111935qvx.219.2024.02.19.03.50.25 for (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Mon, 19 Feb 2024 03:50:25 -0800 (PST) Received-SPF: pass (google.com: domain of linux-kernel+bounces-71245-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) client-ip=2604:1380:45d1:ec00::1; Authentication-Results: mx.google.com; dkim=pass header.i=@google.com header.s=20230601 header.b="PINxa/8u"; 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-71245-ouuuleilei=gmail.com@vger.kernel.org designates 2604:1380:45d1:ec00::1 as permitted sender) smtp.mailfrom="linux-kernel+bounces-71245-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 ny.mirrors.kernel.org (Postfix) with ESMTPS id 660141C210AA for ; Mon, 19 Feb 2024 11:50:10 +0000 (UTC) Received: from localhost.localdomain (localhost.localdomain [127.0.0.1]) by smtp.subspace.kernel.org (Postfix) with ESMTP id 7223733CF9; 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="PINxa/8u" Received: from mail-qk1-f202.google.com (mail-qk1-f202.google.com [209.85.222.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 3BF512E646 for ; Mon, 19 Feb 2024 11:48:41 +0000 (UTC) Authentication-Results: smtp.subspace.kernel.org; arc=none smtp.client-ip=209.85.222.202 ARC-Seal: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343323; cv=none; b=eQRuo834N1H7hdPDENdor/r37sMwlrVFUwXiR8CHoxeU+0aCea0JonKBmTzQhRO72w1qLGyJbHDl26AAiwrR3z5/hWd+qFXKQbqMAGq5rdo8Fm2oMmRo6hdNFjttWbWO5J+nexlTW4RHCya9sY5dG+UNwbR7FQPBM2Khpz6Wj1E= ARC-Message-Signature: i=1; a=rsa-sha256; d=subspace.kernel.org; s=arc-20240116; t=1708343323; c=relaxed/simple; bh=4NJvpz2RnE7LLpmvRQAsdRnn4Zgd+uKku+16slGylXE=; h=Date:In-Reply-To:Mime-Version:References:Message-ID:Subject:From: To:Cc:Content-Type; b=GkvPsCw3bREGTFfNxDlKG7VQs5yArNwYOpRTqejUiw7Mv4Xq/HbGDfbA80iXaOajevbxIZsOolfliKgjm7AHrUnvCbo6kOXmhAiCnif1unv/mrvpjvT++jFkVT312ekrulkXdpYnQB1jzOcyAvaJ++P86ISpnM0VtU4RpJlAH+o= 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=PINxa/8u; arc=none smtp.client-ip=209.85.222.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-qk1-f202.google.com with SMTP id af79cd13be357-7872570470eso573988685a.3 for ; Mon, 19 Feb 2024 03:48:41 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=20230601; t=1708343320; x=1708948120; 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=ZwZ2a7+2VGra5HwuOBKn0CmjD9GQoqV7rwcv+kyW2rs=; b=PINxa/8uYL6xxnUd1RyEcLVU3wLWM5VvWa5Ti3Ch3RL40ZWfo7dZ7CSdnlIOHgBDfw mpzqv4UgrG3gzlipciZjyDTks8YioBVFPGhkWM8lE9Wfi+DI30I8/N0epeQ+ZgOSzWng w+baQAXIkBG0liAeG9qVpZlO8d54/3xYKsYLMiMBzoS4MPvOdubOAtRTVvd+27+CQvyy 1Kku33i4DSRCCUQGUlRkEP22DsOE9QFjfyzf7vTPegnw/x4SfNp5vl2mEgSPqYENDlUJ elZPXLn9AiOUbj4CYRF/CwIzilyU8Q/+zSYH9VXbXYSN4doBCYwbqqALwEPNeblRlZif 2fyw== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1708343320; x=1708948120; 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=ZwZ2a7+2VGra5HwuOBKn0CmjD9GQoqV7rwcv+kyW2rs=; b=SSot9BmPPKCXitZhon5cMQjPhHzVgPgthdhNJZQpevX1S51mlX7Chn3E1oqKCFzcxB s1jhjZMW7zsvCuX/i8lKqFUkANIQMxRNYPNlX+oNQmH258MHnryUlppPfmOxgVWDhY1f VnnBpLtlApMiwTWqprkhsiiair17OfCTEKWbNWd+JHnKqUYf9LqcC9HIEXs/t9sWM6FO MhBW8C55Ywwu6weYTFG4hPAZH4WhysvfaV4agrYOYn0oUv/BRATEE3auoTjSn/IcZPfI dPupftUWofW7qmRw1pQFXtOka3SGFIRFoUiSt1NJJAi6wFoSvbrBYZUNJmRk3VqBqZGs TZdw== X-Forwarded-Encrypted: i=1; AJvYcCWlTnq0h5WfxhQ3crNcRk0EkT6lhq715e/wdj2wMnC2U6KuacUvRKY5q8tbZom6xCqFi7PuvJvoKEcSB/W+y4Q4CVZ0V3DbOMf4U4jZ X-Gm-Message-State: AOJu0YwHcV2IyFGkFM0qYXvL5XsgVIgDnN9CRcslUxMhCZxHsggAmwA6 WB4LvxWGrtaSt5d9t1Yd5+zLEJaEMfwvhg38VS19jdPtagaRybAWNdPmg68yLe9y006fIApShUc 2wHt9qlb23RmEJP8/Vnb/TpjpQw== X-Received: from mattgilbride.c.googlers.com ([fda3:e722:ac3:cc00:2b:7d90:c0a8:2ac5]) (user=mattgilbride job=sendgmr) by 2002:a05:620a:1909:b0:787:6f35:ba54 with SMTP id bj9-20020a05620a190900b007876f35ba54mr9547qkb.8.1708343320353; Mon, 19 Feb 2024 03:48:40 -0800 (PST) Date: Mon, 19 Feb 2024 11:48:12 +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-5-0b113aab330d@google.com> Subject: [PATCH v2 5/6] rust: rbtree: add `RBTreeCursor` 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: 1791327916308239702 X-GMAIL-MSGID: 1791327916308239702 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-08ba9197f637@google.com/ [1] Co-developed-by: Alice Ryhl Signed-off-by: Alice Ryhl Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl Signed-off-by: Matt Gilbride --- rust/kernel/rbtree.rs | 512 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 512 insertions(+) diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index ccf74e0dc3ec..6c3b10e46517 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -245,6 +245,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 key. + pub fn cursor_front(&mut self) -> Option> { + let root = addr_of_mut!(self.root); + // SAFETY: `self.root` is always a valid root node + let current = 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 key. + pub fn cursor_back(&mut self) -> Option> { + let root = addr_of_mut!(self.root); + // SAFETY: `self.root` is always a valid root node + let current = unsafe { bindings::rb_last(root) }; + if current.is_null() { + return None; + } + Some(RBTreeCursor { + _tree: PhantomData, + root, + current, + }) + } } impl RBTree @@ -377,6 +407,59 @@ pub fn remove(&mut self, key: &K) -> Option { } = *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 = self.root.rb_node; + let mut best_match: Option>> = None; + while !node.is_null() { + // SAFETY: All links fields we create are in a `Node`. + let this = unsafe { crate::container_of!(node, Node, links) }.cast_mut(); + // SAFETY: `this` is a non-null node so it is valid by the type invariants. + let this_key = unsafe { &(*this).key }; + // SAFETY: `node` is a non-null node so it is valid by the type invariants. + let left_child = unsafe { (*node).rb_left }; + // SAFETY: `node` is a non-null node so it is valid by the type invariants. + let right_child = unsafe { (*node).rb_right }; + if key == this_key { + return Some(RBTreeCursor { + _tree: PhantomData, + root: addr_of_mut!(self.root), + current: node, + }); + } else { + node = if key > this_key { + right_child + } else { + let is_better_match = match best_match { + None => true, + Some(best) => { + // SAFETY: `best` is a non-null node so it is valid by the type invariants. + let best_key = unsafe { &(*best.as_ptr()).key }; + best_key > this_key + } + }; + if is_better_match { + best_match = 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 type invariants. + current: unsafe { addr_of_mut!((*best.as_ptr()).links) }, + }) + } } impl Default for RBTree { @@ -407,6 +490,435 @@ fn drop(&mut self) { } } +/// 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 the tree. +/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree = 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 = tree.cursor_front().unwrap(); +/// let mut current = cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// +/// // Move the cursor, updating it to the 2nd element. +/// cursor = cursor.move_next().unwrap(); +/// current = cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Peek at the next element without impacting the cursor. +/// let next = cursor.peek_next().unwrap(); +/// assert_eq!(next, (&30, &300)); +/// current = cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Moving past the last element causes the cursor to return [`None`]. +/// cursor = cursor.move_next().unwrap(); +/// current = cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// let cursor = 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 = 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 = tree.cursor_back().unwrap(); +/// let current = 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 = RBTree::new(); +/// assert!(tree.cursor_front().is_none()); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree = 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 = tree.cursor_lower_bound(&20).unwrap(); +/// let current = cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned. +/// let cursor = tree.cursor_lower_bound(&25).unwrap(); +/// let current = cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// // If there is no larger key, [`None`] is returned. +/// let cursor = 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 = 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 = tree.cursor_front().unwrap(); +/// +/// // Get a mutable reference to the current value. +/// let (k, v) = cursor.current_mut(); +/// *v = 1000; +/// +/// // The updated value is reflected in the tree. +/// let updated = tree.get(&10).unwrap(); +/// assert_eq!(updated, &1000); +/// +/// # Ok::<(), Error>(()) +/// ``` +/// +/// It also allows node removal. The following examples demonstrate the behavior of removing the current node. +/// +/// ``` +/// use kernel::rbtree::RBTree; +/// +/// // Create a new tree. +/// let mut tree = 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 = tree.cursor_front().unwrap(); +/// let mut current = cursor.current(); +/// assert_eq!(current, (&10, &100)); +/// cursor = cursor.remove_current().unwrap(); +/// +/// // If a node exists after the current element, it is returned. +/// current = cursor.current(); +/// assert_eq!(current, (&20, &200)); +/// +/// // Get a cursor to the last element, and remove it. +/// cursor = tree.cursor_back().unwrap(); +/// current = cursor.current(); +/// assert_eq!(current, (&30, &300)); +/// +/// // Since there is no next node, the previous node is returned. +/// cursor = cursor.remove_current().unwrap(); +/// current = 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 = 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 = tree.cursor_front().unwrap(); +/// let mut current = 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 = tree.cursor_back().unwrap(); +/// current = 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 = cursor.move_prev().unwrap(); +/// current = 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 that 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 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<'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 = self.get_neighbor_raw(Direction::Prev); + let next = self.get_neighbor_raw(Direction::Next); + // SAFETY: All links fields we create are in a `Node`. + let this = unsafe { crate::container_of!(self.current, 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 are valid. + unsafe { bindings::rb_erase(&mut (*this).links, self.root) }; + + let current = match (prev, next) { + (_, Some(next)) => next, + (Some(prev), None) => prev, + (None, None) => { + 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) = self.get_neighbor_raw(direction) { + // SAFETY: All links fields we create are in a `Node`. + let this = 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 are 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 doesn'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 = unsafe { + match direction { + Direction::Prev => bindings::rb_prev(self.current), + Direction::Next => 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 = unsafe { + match direction { + Direction::Prev => bindings::rb_prev(self.current), + Direction::Next => 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 bindings::rb_node> { + // SAFETY: `self.current` is valid by the type invariants. + let neighbor = unsafe { + match direction { + Direction::Prev => bindings::rb_prev(self.current), + Direction::Next => bindings::rb_next(self.current), + } + }; + + if neighbor.is_null() { + return None; + } + + Some(neighbor) + } + + // This internal method should *only* be called with a valid pointer to 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 = unsafe { crate::container_of!(node, Node, links) }; + // SAFETY: The passed `node` is the current node or a non-null neighbor, + // thus `this` is valid by the type invariants. + let k = unsafe { &(*this).key }; + // SAFETY: The passed `node` is the current node or a non-null neighbor, + // thus `this` is valid by the type invariants. + let v = unsafe { &(*this).value }; + (k, v) + } + + // This internal method should *only* be called with a valid pointer to 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 = unsafe { crate::container_of!(node, Node, links) }.cast_mut(); + // SAFETY: The passed `node` is the current node or a non-null neighbor, + // thus `this` is valid by the type invariants. + let k = unsafe { &(*this).key }; + // SAFETY: The passed `node` is the current node or a non-null neighbor, + // thus `this` is valid by the type invariants. + let v = unsafe { &mut (*this).value }; + (k, v) + } + + // This internal method should *only* be called with a valid pointer to 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 = unsafe { crate::container_of!(node, Node, links) }.cast_mut(); + // SAFETY: The passed `node` is the current node or a non-null neighbor, + // thus `this` is valid by the type invariants. + let n = 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 = (&'a K, &'a V); type IntoIter = RBTreeIterator<'a, K, V>; 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 } + } +}