[PREVIEW PATCH] rust: kernel: Support more jump_label api

[chenmiao]

The initial implementation of arch_static_branchwas achieved by accessing
the offset from the original type. However, this approach extended the
path and introduced redundant calculations when dealing with types like
`static_key_true/false`, as shown below:

```
static_brach_unlikely(tp, tracepoint, key)
=> tracepoint->key->key
=> &tracepoint->key(static_key_false) == &tracepoint->key.key(static_key)
=> off: tracepoint->key - tracepoint
```

In practice, the implementation of `arch_static_branch` overlooked many
detailed descriptions. To improve clarity, additional comments have been
added to the original logic. The approach has been modified to directly
locate the corresponding `static_key` instead of using offsets, thereby
reducing computational overhead.

Additionally, support for the `static_branch_enable/disable` APIs has been
introduced.

Signed-off-by: chenmiao <chen...@openatom.club>
---
.../generated_arch_static_branch_asm.rs.S | 2 +-
rust/kernel/jump_label.rs | 169 ++++++++++++++++--
2 files changed, 155 insertions(+), 16 deletions(-)

diff --git a/rust/kernel/generated_arch_static_branch_asm.rs.S b/rust/kernel/generated_arch_static_branch_asm.rs.S
index 2afb638708db3..08603dc2d61e7 100644
--- a/rust/kernel/generated_arch_static_branch_asm.rs.S
+++ b/rust/kernel/generated_arch_static_branch_asm.rs.S
@@ -4,4 +4,4 @@

// Cut here.

-::kernel::concat_literals!(ARCH_STATIC_BRANCH_ASM("{symb} + {off} + {branch}", "{l_yes}"))
+::kernel::concat_literals!(ARCH_STATIC_BRANCH_ASM("{real_key} + {branch}", "{l_yes}"))
diff --git a/rust/kernel/jump_label.rs b/rust/kernel/jump_label.rs
index 4e974c768dbd5..12ae78f95c761 100644
--- a/rust/kernel/jump_label.rs
+++ b/rust/kernel/jump_label.rs
@@ -6,22 +6,63 @@
//!
//! C header: [`include/linux/jump_label.h`](srctree/include/linux/jump_label.h).

-/// Branch based on a static key.
+/// The key used for the static_key_false/true.
///
-/// Takes three arguments:
+/// If the key just contains a static_key, like: `struct tracepoint`;
///
-/// * `key` - the path to the static variable containing the `static_key`.
-/// * `keytyp` - the type of `key`.
-/// * `field` - the name of the field of `key` that contains the `static_key`.
+/// ```
+/// pub struct tracepoint {
+/// ...,
+/// key: static_key,
+/// ...,
+/// }
///
-/// # Safety
+/// // When you use the tracepoint as the parameter.
+/// if static_branch_unlikely!(tp, tracepoint, key) {
+/// // Do something
+/// }
///
-/// The macro must be used with a real static key defined by C.
-#[macro_export]
-macro_rules! static_branch_unlikely {
- ($key:path, $keytyp:ty, $field:ident) => {{
+/// // It just like:
+/// let _key: *const crate::bindings::tracepoint = ::core::ptr::addr_of!($key);
+/// let _key: *const crate::bindings::static_key_false = ::core::ptr::addr_of!((*_key).key);
+/// let _key: *const crate::bindings::static_key = _key.cast();
+/// ```
+///
+/// If the key just contains a single static_key, like: `struct static_key_false`;
+///
+/// ```
+/// pub struct static_key_false {
+/// key: static_key,
+/// }
+///
+/// // When you use the static_key_false as the parameter.
+/// if static_branch_unlikely!(key) {
+/// // Do something
+/// }
+///
+/// // It just like:
+/// let _key: *const crate::bindings::static_key_false = ::core::ptr::addr_of!($key);
+/// let _key: *const crate::bindings::static_key = _key.cast();
+/// ```
+///
+macro_rules! __static_branch_base {
+ ($basety:ty, $branch:expr, $key:path) => {{
+ let _key: *const $basety = ::core::ptr::addr_of!($key);
+ let _key: *const $crate::bindings::static_key = _key.cast();
+
+ #[cfg(not(CONFIG_JUMP_LABEL))]
+ {
+ $crate::bindings::static_key_count(_key.cast_mut()) > 0
+ }
+
+ #[cfg(CONFIG_JUMP_LABEL)]
+ {
+ $crate::jump_label::arch_static_branch! { _key, $branch }
+ }
+ }};
+ ($basety:ty, $branch:expr, $key:path, $keytyp:ty, $field:ident) => {{
let _key: *const $keytyp = ::core::ptr::addr_of!($key);
- let _key: *const $crate::bindings::static_key_false = ::core::ptr::addr_of!((*_key).$field);
+ let _key: *const $basety = ::core::ptr::addr_of!((*_key).$field);
let _key: *const $crate::bindings::static_key = _key.cast();

#[cfg(not(CONFIG_JUMP_LABEL))]
@@ -30,7 +71,88 @@ macro_rules! static_branch_unlikely {
}

#[cfg(CONFIG_JUMP_LABEL)]
- $crate::jump_label::arch_static_branch! { $key, $keytyp, $field, false }
+ {
+ $crate::jump_label::arch_static_branch! { _key, $branch }
+ }
+ }};
+}
+
+/// Branch based on a static key.
+///
+/// Takes two type arguments:
+///
+/// First Type takes one argument:
+///
+/// * `key` - the static variable containing the `static_key`.
+///
+/// Second Type takes three arguments:
+///
+/// * `key` - the path to the static variable containing the `static_key`.
+/// * `keytyp` - the type of `key`.
+/// * `field` - the name of the field of `key` that contains the `static_key`.
+///
+/// # Safety
+///
+/// ```
+/// let tp: tracepoint = tracepoint::new();
+/// if static_key_likely!(tp, tracepoint, key) {
+/// // Do something
+/// }
+///
+/// let key: static_key_false = static_key_true::new();
+/// if static_key_likely!(key) {
+/// // Do something
+/// }
+/// ```
+///
+/// The macro must be used with a real static key defined by C.
+#[macro_export]
+macro_rules! static_branch_likely {
+ ($key:path) => {{
+ __static_branch_base! { $crate::bindings::static_key_true, true, $key }
+ }};
+ ($key:path, $keytyp:ty, $field:ident) => {{
+ __static_branch_base! { $crate::bindings::static_key_true, true, $key, $keytyp, $field }
+ }};
+}
+pub use static_branch_likely;
+
+/// Branch based on a static key.
+///
+/// Takes two type arguments:
+///
+/// First Type takes one argument:
+///
+/// * `key` - the static variable containing the `static_key`.
+///
+/// Second Type takes three arguments:
+///
+/// * `key` - the path to the static variable containing the `static_key`.
+/// * `keytyp` - the type of `key`.
+/// * `field` - the name of the field of `key` that contains the `static_key`.
+///
+/// # Safety
+///
+/// ```
+/// let tp: tracepoint = tracepoint::new();
+/// if static_key_unlikely!(tp, tracepoint, key) {
+/// // Do something
+/// }
+///
+/// let key: static_key_false = static_key_false::new();
+/// if static_key_unlikely!(key) {
+/// // Do something
+/// }
+/// ```
+///
+/// The macro must be used with a real static key defined by C.
+#[macro_export]
+macro_rules! static_branch_unlikely {
+ ($key:path) => {{
+ static_branch_base! { $crate::bindings::static_key_false, false, $key }
+ }};
+ ($key:path, $keytyp:ty, $field:ident) => {{
+ static_branch_base! { $crate::bindings::static_key_false, false, $key, $keytyp, $field }
}};
}
pub use static_branch_unlikely;
@@ -46,14 +168,13 @@ macro_rules! static_branch_unlikely {
#[doc(hidden)]
#[cfg(CONFIG_JUMP_LABEL)]
macro_rules! arch_static_branch {
- ($key:path, $keytyp:ty, $field:ident, $branch:expr) => {'my_label: {
+ ($key:path, $branch:expr) => {'my_label: {
$crate::asm!(
include!(concat!(env!("OBJTREE"), "/rust/kernel/generated_arch_static_branch_asm.rs"));
l_yes = label {
break 'my_label true;
},
- symb = sym $key,
- off = const ::core::mem::offset_of!($keytyp, $field),
+ real_key = sym $key,
branch = const $crate::jump_label::bool_to_int($branch),
);

@@ -72,3 +193,21 @@ macro_rules! arch_static_branch {
pub const fn bool_to_int(b: bool) -> i32 {
b as i32
}
+
+/// Enable a static branch.
+#[macro_export]
+macro_rules! static_branch_enable {
+ ($key:path) => {{
+ let _key: *const $crate::bindings::static_key = ::core::ptr::addr_of!($key);
+ $crate::bindings::static_key_enable(_key.cast_mut());
+ }};
+}
+
+/// Disable a static branch.
+#[macro_export]
+macro_rules! static_branch_disable {
+ ($key:path) => {{
+ let _key: *const $crate::bindings::static_key = ::core::ptr::addr_of!($key);
+ $crate::bindings::static_key_disable(_key.cast_mut());
+ }};
+}
--
2.43.0

[Chen Miao]

On 11/4/2025 10:07 PM, Alice Ryhl wrote:
> On Tue, Nov 04, 2025 at 02:04:17AM +0000, chenmiao wrote:
>> The initial implementation of arch_static_branch was achieved by accessing

>> the offset from the original type. However, this approach extended the
>> path and introduced redundant calculations when dealing with types like
>> `static_key_true/false`, as shown below:
>>
>> ```
>> static_brach_unlikely(tp, tracepoint, key)
>> => tracepoint->key->key
>> => &tracepoint->key(static_key_false) == &tracepoint->key.key(static_key)
>> => off: tracepoint->key - tracepoint
>> ```
>>
>> In practice, the implementation of `arch_static_branch` overlooked many
>> detailed descriptions. To improve clarity, additional comments have been
>> added to the original logic. The approach has been modified to directly
>> locate the corresponding `static_key` instead of using offsets, thereby
>> reducing computational overhead.
>>

>> If finding the offset from the primitive type is necessary for this
>> implementation, I will abandon this change.

>>
>> Additionally, support for the `static_branch_enable/disable` APIs has been
>> introduced.
>>
>> Signed-off-by: chenmiao <chen...@openatom.club>

>> + ($basety:ty, $branch:expr, $key:path, $keytyp:ty, $field:ident) => {{
>> let _key: *const $keytyp = ::core::ptr::addr_of!($key);
>> - let _key: *const $crate::bindings::static_key_false = ::core::ptr::addr_of!((*_key).$field);
>> + let _key: *const $basety = ::core::ptr::addr_of!((*_key).$field);
>> let _key: *const $crate::bindings::static_key = _key.cast();
>>
>> #[cfg(not(CONFIG_JUMP_LABEL))]
>> @@ -30,7 +71,88 @@ macro_rules! static_branch_unlikely {
>> }
>>
>> #[cfg(CONFIG_JUMP_LABEL)]
>> - $crate::jump_label::arch_static_branch! { $key, $keytyp, $field, false }
>> + {
>> + $crate::jump_label::arch_static_branch! { _key, $branch }
>> + }

> So ... this is changing from $key to _key. That's replacing the global
> variable with a local variable holding a pointer to the global variable.
> However, the arch_static_branch! macro uses the `sym` operand which
> requires you to pass it the global directly.
>
> Did you try this code? I don't believe it will work.
>
> Alice

I'm very sorry for making a fatal mistake. My misunderstanding of sym led to
this issue, so I shouldn't make any changes to that part. However, regarding
the other modifications, I believe it is necessary to support the direct
passing of structures similar to `static_key_false`, just as in C language.

Chen Miao

[Alice Ryhl]

It sounds like you are adding a new use-case for this macro. Can you
provide more information for this new feature? It is currently unclear
to me exactly how this will be used.

Alice

[Chen Miao]

If there's a binding-required driver implementation in the future where a key
function uses if (static_branch_unlikely(&zoned_enabled))— defined by
DEFINE_STATIC_KEY_FALSE(zoned_enabled);— then in Rust we can directly
implement it using if static_branch_unlikely!(zoned_enabled), without having
to call it via if static_branch_unlikely!(zoned_enabled,
bindings::static_key_false, key).

static_branch_unlikely!(zoned_enabled) instead of
static_branch_unlikely!(zoned_enabled, bindings::static_key_false, key).

Chen Miao

[Alice Ryhl]

In general, you would never use "static_key_false" as the second
argument to static_branch_unlikely!. The second argument is the name of
the struct *containing* a field of type static_key_false.

I guess your point is that there's no way to use the macro right now if
the global is a bare static_key_false that is not wrapped in a struct?

Alice

[Chen Miao]

Yes, you're right. But in fact, currently there are no other uses of
static_branch_unlikely in Rust for Linux (except for tracepoint), so this is
also a bit awkward.