[PATCH v2] rust: check type of `$ptr` in `container_of!`

[Tamir Duberstein <tamird@gmail.com>]

Add a compile-time check that `*$ptr` is of the type of `$type->$($f)*`.

Given the incorrect usage:

; diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
; index 8d978c896747..6a7089149878 100644
; --- a/rust/kernel/rbtree.rs
; +++ b/rust/kernel/rbtree.rs
; @@ -329,7 +329,7 @@ fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
;          while !(*child_field_of_parent).is_null() {
;              let curr = *child_field_of_parent;
;              // SAFETY: All links fields we create are in a `Node<K, V>`.
; -            let node = unsafe { container_of!(curr, Node<K, V>, links) };
; +            let node = unsafe { container_of!(curr, Node<K, V>, key) };
;
;              // SAFETY: `node` is a non-null node so it is valid by the type invariants.
;              match key.cmp(unsafe { &(*node).key }) {

this patch produces the compilation error:

; error[E0308]: mismatched types
;    --> rust/kernel/lib.rs:206:26
;     |
; 206 |             [$field_ptr, field_ptr]; // typeof(`$ptr_to_field`) == typeof(`$Container.$($fields)*`)
;     |                          ^^^^^^^^^ expected `*mut rb_node`, found `*mut K`
;     |
;    ::: rust/kernel/rbtree.rs:270:6
;     |
; 270 | impl<K, V> RBTree<K, V>
;     |      - found this type parameter
; ...
; 332 |             let node = unsafe { container_of!(curr, Node<K, V>, key) };
;     |                                 ------------------------------------ in this macro invocation
;     |
;     = note: expected raw pointer `*mut bindings::rb_node`
;                found raw pointer `*mut K`
;     = note: this error originates in the macro `container_of` (in Nightly builds, run with -Z macro-backtrace for more info)
;
; error: aborting due to 1 previous error

Suggested-by: Alice Ryhl <aliceryhl@google.com>
Link: https://lore.kernel.org/all/CAH5fLgh6gmqGBhPMi2SKn7mCmMWfOSiS0WP5wBuGPYh9ZTAiww@mail.gmail.com/
Signed-off-by: Tamir Duberstein <tamird@gmail.com>
---
I also considered an implementation using a function, but the resulting
compilation error was not as concise:

; error[E0308]: mismatched types
;    --> rust/kernel/lib.rs:207:42
;     |
; 207 |             assert_same_type($field_ptr, field_ptr);
;     |             ----------------             ^^^^^^^^^ expected `*const rb_node`, found `*mut K`
;     |             |
;     |             arguments to this function are incorrect
;     |
;    ::: rust/kernel/rbtree.rs:270:6
;     |
; 270 | impl<K, V> RBTree<K, V>
;     |      - found this type parameter
; ...
; 332 |             let node = unsafe { container_of!(curr, Node<K, V>, key) };
;     |                                 ------------------------------------ in this macro invocation
;     |
;     = note: expected raw pointer `*const bindings::rb_node`
;                found raw pointer `*mut K`
; note: function defined here
;    --> rust/kernel/lib.rs:205:16
;     |
; 205 |             fn assert_same_type<T>(_: *const T, _: *const T) {}
;     |                ^^^^^^^^^^^^^^^^                 -----------
;     |
;    ::: rust/kernel/rbtree.rs:332:33
;     |
; 332 |             let node = unsafe { container_of!(curr, Node<K, V>, key) };
;     |                                 ------------------------------------ in this macro invocation
;     = note: this error originates in the macro `container_of` (in Nightly builds, run with -Z macro-backtrace for more info)
; 
; error: aborting due to 1 previous error
---
Changes in v2:
- Wrap in `if false` to improve unoptimized codegen. (Alice Ryhl)
- Shrink implementation using an array literal instead of a function.
- Link to v1: https://lore.kernel.org/r/20250411-b4-container-of-type-check-v1-1-08262ef67c95@gmail.com
---
 rust/kernel/lib.rs | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 1df11156302a..6fbd4cc5afff 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -198,9 +198,14 @@ fn panic(info: &core::panic::PanicInfo<'_>) -> ! {
 /// ```
 #[macro_export]
 macro_rules! container_of {
-    ($ptr:expr, $type:ty, $($f:tt)*) => {{
-        let offset: usize = ::core::mem::offset_of!($type, $($f)*);
-        $ptr.byte_sub(offset).cast::<$type>()
+    ($field_ptr:expr, $Container:ty, $($fields:tt)*) => {{
+        let offset: usize = ::core::mem::offset_of!($Container, $($fields)*);
+        let container_ptr = $field_ptr.byte_sub(offset).cast::<$Container>();
+        if false {
+            let field_ptr = ::core::ptr::addr_of!((*container_ptr).$($fields)*).cast_mut();
+            [$field_ptr, field_ptr]; // typeof(`$ptr_to_field`) == typeof(`$Container.$($fields)*`)
+        }
+        container_ptr
     }}
 }
 

---
base-commit: 0af2f6be1b4281385b618cb86ad946eded089ac8
change-id: 20250411-b4-container-of-type-check-06af1c204f59
prerequisite-change-id: 20250409-container-of-mutness-b153dab4388d:v1
prerequisite-patch-id: 53d5889db599267f87642bb0ae3063c29bc24863

Best regards,
-- 
Tamir Duberstein <tamird@gmail.com>