Re: [PATCH v5 1/4] rust: clk: use the type-state pattern

["Alexandre Courbot" <acourbot@nvidia.com>]

On Mon Jul 6, 2026 at 11:37 PM JST, Daniel Almeida wrote:
> The current Clk abstraction can still be improved on the following issues:
>
> a) It only keeps track of a count to clk_get(), which means that users have
> to manually call disable() and unprepare(), or a variation of those, like
> disable_unprepare().
>
> b) It allows repeated calls to prepare() or enable(), but it keeps no track
> of how often these were called, i.e., it's currently legal to write the
> following:
>
> clk.prepare();
> clk.prepare();
> clk.enable();
> clk.enable();
>
> And nothing gets undone on drop().
>
> c) It adds a OptionalClk type that is probably not needed. There is no
> "struct optional_clk" in C and we should probably not add one.
>
> d) It does not let a user express the state of the clk through the
> type system. For example, there is currently no way to encode that a Clk is
> enabled via the type system alone.
>
> In light of the Regulator abstraction that was recently merged, switch this
> abstraction to use the type-state pattern instead. It solves both a) and b)
> by establishing a number of states and the valid ways to transition between
> them. It also automatically undoes any call to clk_get(), clk_prepare() and
> clk_enable() as applicable on drop(), so users do not have to do anything
> special before Clk goes out of scope.
>
> It solves c) by removing the OptionalClk type, which is now simply encoded
> as a Clk whose inner pointer is NULL.
>
> It solves d) by directly encoding the state of the Clk into the type, e.g.:
> Clk<Enabled> is now known to be a Clk that is enabled.
>
> The INVARIANTS section for Clk is expanded to highlight the relationship
> between the states and the respective reference counts that are owned by
> each of them.
>
> The examples are expanded to highlight how a user can transition between
> states, as well as highlight some of the shortcuts built into the API.
>
> The current implementation is also more flexible, in the sense that it
> allows for more states to be added in the future. This lets us implement
> different strategies for handling clocks, including one that mimics the
> current API, allowing for multiple calls to prepare() and enable().
>
> The users (cpufreq.rs/ rcpufreq_dt.rs) were updated by this patch (and not
> a separate one) to reflect the new changes. This is needed, because
> otherwise this patch would break the build.
>
> Link: https://crates.io/crates/sealed [1]
> Signed-off-by: Daniel Almeida <daniel.almeida@collabora.com>

The new API looks super nice; I really like it. A few nits/questions
inline, but regardless:

Reviewed-by: Alexandre Courbot <acourbot@nvidia.com>

I will try to go through the rest of the series shortly.

> ---
>  drivers/cpufreq/rcpufreq_dt.rs |   2 +-
>  drivers/gpu/drm/tyr/driver.rs  |  37 +--
>  drivers/pwm/pwm_th1520.rs      |  17 +-
>  rust/kernel/clk.rs             | 541 ++++++++++++++++++++++++++++++-----------
>  rust/kernel/cpufreq.rs         |   8 +-
>  5 files changed, 423 insertions(+), 182 deletions(-)
>
> diff --git a/drivers/cpufreq/rcpufreq_dt.rs b/drivers/cpufreq/rcpufreq_dt.rs
> index f17bf64c22e2..9d2ec7df4bac 100644
> --- a/drivers/cpufreq/rcpufreq_dt.rs
> +++ b/drivers/cpufreq/rcpufreq_dt.rs
> @@ -40,7 +40,7 @@ struct CPUFreqDTDevice {
>      freq_table: opp::FreqTable,
>      _mask: CpumaskVar,
>      _token: Option<opp::ConfigToken>,
> -    _clk: Clk,
> +    _clk: Clk<kernel::clk::Unprepared>,

Maybe import `kernel::clk` to shorten this a bit.

<...>
> +    /// An error that can occur when trying to convert a [`Clk`] between states.
> +    pub struct Error<State: ClkState> {
> +        /// The error that occurred.
> +        pub error: kernel::error::Error,
> +
> +        /// The [`Clk`] that caused the error, so that the operation may be
> +        /// retried.
> +        pub clk: Clk<State>,
> +    }

Can this have a `Debug` implementation? It can just forward to `error`.

> +
> +    impl<State: ClkState> From<Error<State>> for kernel::error::Error {
> +        /// Discards the [`Clk`] and keeps only the error code.
> +        ///
> +        /// This makes the fallible state transitions usable with the `?`
> +        /// operator when the caller does not need to retry the operation on the
> +        /// original [`Clk`], e.g.:
> +        ///
> +        /// ```
> +        /// use kernel::clk::{Clk, Enabled, Unprepared};
> +        /// use kernel::device::{Bound, Device};
> +        /// use kernel::error::Result;
> +        ///
> +        /// fn get_enabled(dev: &Device<Bound>) -> Result<Clk<Enabled>> {
> +        ///     let clk = Clk::<Unprepared>::get(dev, Some(c"apb_clk"))?
> +        ///         .prepare()?
> +        ///         .enable()?;
> +        ///     Ok(clk)
> +        /// }
> +        /// ```
> +        #[inline]
> +        fn from(err: Error<State>) -> Self {
> +            err.error
> +        }
> +    }
>  
>      /// A reference-counted clock.
>      ///
>      /// Rust abstraction for the C [`struct clk`].
>      ///
> +    /// A [`Clk`] instance represents a clock that can be in one of several
> +    /// states: [`Unprepared`], [`Prepared`], or [`Enabled`].
> +    ///
> +    /// No action needs to be taken when a [`Clk`] is dropped. The calls to
> +    /// `clk_unprepare()` and `clk_disable()` will be placed as applicable.

s/placed/made?

> +    ///
> +    /// An optional [`Clk`] is treated just like a regular [`Clk`], but its
> +    /// inner `struct clk` pointer is `NULL`. This interfaces correctly with the
> +    /// C API and also exposes all the methods of a regular [`Clk`] to users.
> +    ///
>      /// # Invariants
>      ///
>      /// A [`Clk`] instance holds either a pointer to a valid [`struct clk`] created by the C
> @@ -99,19 +185,36 @@ mod common_clk {
>      /// Instances of this type are reference-counted. Calling [`Clk::get`] ensures that the
>      /// allocation remains valid for the lifetime of the [`Clk`].
>      ///
> +    /// The [`Prepared`] state is associated with a single count of
> +    /// `clk_prepare()`, and the [`Enabled`] state is associated with a single
> +    /// count of both `clk_prepare()` and `clk_enable()`.
> +    ///
> +    /// All states are associated with a single count of `clk_get()`.
> +    ///
>      /// # Examples
>      ///
>      /// The following example demonstrates how to obtain and configure a clock for a device.
>      ///
>      /// ```
> -    /// use kernel::clk::{Clk, Hertz};
> -    /// use kernel::device::Device;
> +    /// use kernel::clk::{Clk, Enabled, Hertz, Unprepared, Prepared};
> +    /// use kernel::device::{Bound, Device};
>      /// use kernel::error::Result;
>      ///
> -    /// fn configure_clk(dev: &Device) -> Result {
> -    ///     let clk = Clk::get(dev, Some(c"apb_clk"))?;
> +    /// fn configure_clk(dev: &Device<Bound>) -> Result {
> +    ///     // The fastest way is to use a version of `Clk::get` for the desired
> +    ///     // state, i.e.:
> +    ///     let clk: Clk<Enabled> = Clk::<Enabled>::get(dev, Some(c"apb_clk"))?;
> +    ///
> +    ///     // Any other state is also possible, e.g.:
> +    ///     let clk: Clk<Prepared> = Clk::<Prepared>::get(dev, Some(c"apb_clk"))?;

nit: maybe use a different name as this is otherwise obtaining the same
clock.

>      ///
> -    ///     clk.prepare_enable()?;
> +    ///     // Later:
> +    ///     //
> +    ///     // `?` works directly thanks to `From<Error<State>>`; the failed
> +    ///     // `Clk` is dropped on error. Match on the returned `Error<State>`
> +    ///     // instead (its `clk` field is the original `Clk`) if you want to
> +    ///     // retry the operation.
> +    ///     let clk: Clk<Enabled> = clk.enable()?;
>      ///
>      ///     let expected_rate = Hertz::from_ghz(1);
>      ///
> @@ -119,122 +222,339 @@ mod common_clk {
>      ///         clk.set_rate(expected_rate)?;
>      ///     }
>      ///
> -    ///     clk.disable_unprepare();
> +    ///     // Nothing is needed here. The drop implementation will undo any
> +    ///     // operations as appropriate.
> +    ///     Ok(())
> +    /// }
> +    ///
> +    /// fn shutdown(clk: Clk<Enabled>) -> Result {
> +    ///     // The states can be traversed "in the reverse order" as well:
> +    ///     let clk: Clk<Prepared> = clk.disable();
> +    ///
> +    ///     // This is of type `Clk<Unprepared>`.
> +    ///     let clk = clk.unprepare();
> +    ///
>      ///     Ok(())
>      /// }
>      /// ```
>      ///
> +    /// Drivers that need to change a clock's state at runtime (for example to
> +    /// enable it on resume and disable it on suspend) can keep it in an enum
> +    /// and move between the variants:
> +    ///
> +    /// ```
> +    /// use kernel::clk::{Clk, Enabled, Prepared};

I know patch 4 eventually fixes the imports, but a more logical ordering
would be to fix them first, as it would avoid a bit of churn. Not a big
deal though.

<...>
> +        pub fn unprepare(self) -> Clk<Unprepared> {
> +            // We will be transferring the ownership of our `clk_get()` count to
> +            // `Clk<Unprepared>`.
> +            let clk = ManuallyDrop::new(self);
> +
> +            // SAFETY: By the type invariants, `clk.as_raw()` is a valid argument
> +            // for [`clk_unprepare`].
> +            unsafe { bindings::clk_unprepare(clk.as_raw()) }
> +
> +            // INVARIANT: The `clk_prepare()` count was released above, so the
> +            // returned `Clk<Unprepared>` owns only the `clk_get()` count.
> +            Clk {
> +                inner: clk.inner,
> +                _phantom: PhantomData,
> +            }
>          }
>  
> -        /// Prepare the clock.
> +        /// Attempts to convert the [`Clk`] to an [`Enabled`] state.
>          ///
> -        /// Equivalent to the kernel's [`clk_prepare`] API.
> +        /// Equivalent to the kernel's [`clk_enable`] API.
>          ///
> -        /// [`clk_prepare`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_prepare
> +        /// [`clk_enable`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_enable
>          #[inline]
> -        pub fn prepare(&self) -> Result {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_prepare`].
> -            to_result(unsafe { bindings::clk_prepare(self.as_raw()) })
> +        pub fn enable(self) -> Result<Clk<Enabled>, Error<Prepared>> {

Note that the `Regulator` API uses the `try_into_enabled` pattern for
state transitions that can fail. Now that these transitions are
consuming the clock, it might make sense to align?

> +            // We will be transferring the ownership of our `clk_get()` and
> +            // `clk_prepare()` counts to `Clk<Enabled>`.
> +            let clk = ManuallyDrop::new(self);
> +
> +            // SAFETY: By the type invariants, `clk.as_raw()` is a valid argument
> +            // for [`clk_enable`].
> +            to_result(unsafe { bindings::clk_enable(clk.as_raw()) })
> +                // INVARIANT: `clk_enable()` succeeded, so the returned
> +                // `Clk<Enabled>` owns a single count of it, which is released
> +                // when it leaves the [`Enabled`] state.
> +                .map(|()| Clk {
> +                    inner: clk.inner,
> +                    _phantom: PhantomData,
> +                })
> +                .map_err(|error| Error {
> +                    error,
> +                    clk: ManuallyDrop::into_inner(clk),
> +                })
>          }
>  
> -        /// Unprepare the clock.
> +        /// Runs `cb` with the clock temporarily enabled.
>          ///
> -        /// Equivalent to the kernel's [`clk_unprepare`] API.
> +        /// The clock is enabled before `cb` runs and disabled again afterwards,
> +        /// so the [`Enabled`] state is scoped to the closure and the [`Clk`]
> +        /// remains [`Prepared`]. This is convenient for drivers that only need
> +        /// the clock running for a short, well-defined section (e.g. while
> +        /// touching registers) without giving up ownership of the prepared
> +        /// clock or threading it through an intermediate state, e.g.:
>          ///
> -        /// [`clk_unprepare`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_unprepare
> +        /// ```
> +        /// use kernel::clk::{Clk, Enabled, Hertz, Prepared};
> +        /// use kernel::error::Result;
> +        ///
> +        /// fn read_rate(clk: &Clk<Prepared>) -> Result<Hertz> {
> +        ///     clk.with_enabled(|clk: &Clk<Enabled>| clk.rate())
> +        /// }
> +        /// ```
> +        ///
> +        /// Equivalent to a balanced [`clk_enable`]/[`clk_disable`] pair around
> +        /// `cb`.
> +        ///
> +        /// [`clk_enable`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_enable
> +        /// [`clk_disable`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_disable
>          #[inline]
> -        pub fn unprepare(&self) {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_unprepare`].
> -            unsafe { bindings::clk_unprepare(self.as_raw()) };
> +        pub fn with_enabled<R>(&self, cb: impl FnOnce(&Clk<Enabled>) -> R) -> Result<R> {
> +            // SAFETY: By the type invariants, `self.as_raw()` is a valid argument for
> +            // [`clk_enable`].
> +            to_result(unsafe { bindings::clk_enable(self.as_raw()) })?;
> +
> +            // Borrow the same clock as `Clk<Enabled>` for the duration of `cb`.
> +            // It must not be dropped, as that would run `clk_disable`/`clk_put`
> +            // against counts owned by `self`; the matching `clk_disable` below
> +            // balances the `clk_enable` above instead.
> +            //
> +            // INVARIANT: The clock is enabled for the lifetime of `enabled`.
> +            let enabled = ManuallyDrop::new(Clk::<Enabled> {
> +                inner: self.inner,
> +                _phantom: PhantomData,
> +            });
> +
> +            let ret = cb(&enabled);
> +
> +            // SAFETY: The `clk_enable` above succeeded, so this balances it.
> +            // `cb` only had a shared reference, so the enable count is unchanged.
> +            unsafe { bindings::clk_disable(self.as_raw()) };
> +
> +            Ok(ret)
>          }
> +    }
>  
> -        /// Prepare and enable the clock.
> +    impl Clk<Enabled> {
> +        /// Gets [`Clk`] corresponding to a bound [`Device`] and a connection id
> +        /// and then prepares and enables it.
>          ///
> -        /// Equivalent to calling [`Clk::prepare`] followed by [`Clk::enable`].
> +        /// Equivalent to calling [`Clk::get`], followed by [`Clk::prepare`],
> +        /// followed by [`Clk::enable`].
>          #[inline]
> -        pub fn prepare_enable(&self) -> Result {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_prepare_enable`].
> -            to_result(unsafe { bindings::clk_prepare_enable(self.as_raw()) })
> +        pub fn get(dev: &Device<Bound>, name: Option<&CStr>) -> Result<Clk<Enabled>> {
> +            Clk::<Prepared>::get(dev, name)?
> +                .enable()
> +                .map_err(|error| error.error)
> +        }
> +
> +        /// Behaves the same as [`Self::get`], except when there is no clock
> +        /// producer. In this case, instead of returning [`ENOENT`], it returns
> +        /// a dummy [`Clk`].
> +        #[inline]
> +        pub fn get_optional(dev: &Device<Bound>, name: Option<&CStr>) -> Result<Clk<Enabled>> {
> +            Clk::<Prepared>::get_optional(dev, name)?
> +                .enable()
> +                .map_err(|error| error.error)
>          }
>  
> -        /// Disable and unprepare the clock.
> +        /// Disables the [`Clk`] and converts it to the [`Prepared`] state.
>          ///
> -        /// Equivalent to calling [`Clk::disable`] followed by [`Clk::unprepare`].
> +        /// Equivalent to the kernel's [`clk_disable`] API.
> +        ///
> +        /// [`clk_disable`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_disable
>          #[inline]
> -        pub fn disable_unprepare(&self) {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_disable_unprepare`].
> -            unsafe { bindings::clk_disable_unprepare(self.as_raw()) };
> +        pub fn disable(self) -> Clk<Prepared> {
> +            // We will be transferring the ownership of our `clk_get()` and
> +            // `clk_prepare()` counts to `Clk<Prepared>`.
> +            let clk = ManuallyDrop::new(self);
> +
> +            // SAFETY: By the type invariants, `clk.as_raw()` is a valid argument
> +            // for [`clk_disable`].
> +            unsafe { bindings::clk_disable(clk.as_raw()) };
> +
> +            // INVARIANT: The `clk_enable()` count was released above, so the
> +            // returned `Clk<Prepared>` owns the `clk_get()` and `clk_prepare()`
> +            // counts.
> +            Clk {
> +                inner: clk.inner,
> +                _phantom: PhantomData,
> +            }
> +        }
> +    }
> +
> +    impl<T: ClkState> Clk<T> {
> +        /// Obtain the raw [`struct clk`] pointer.
> +        #[inline]
> +        pub fn as_raw(&self) -> *mut bindings::clk {
> +            self.inner
>          }
>  
>          /// Get clock's rate.
>          ///
>          /// Equivalent to the kernel's [`clk_get_rate`] API.
>          ///
> +        /// Note that the returned rate is only guaranteed to reflect what the
> +        /// hardware is doing once the clock is [`Enabled`].
> +        ///
>          /// [`clk_get_rate`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_get_rate
>          #[inline]
>          pub fn rate(&self) -> Hertz {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_get_rate`].
> +            // SAFETY: By the type invariants, `self.as_raw()` is a valid argument
> +            // for [`clk_get_rate`].

Ideally these cosmetic fixes would have been in their own patch to not
distract from the rest, but not a big deal.

>              Hertz(unsafe { bindings::clk_get_rate(self.as_raw()) })
>          }
>  
> @@ -245,88 +565,29 @@ pub fn rate(&self) -> Hertz {
>          /// [`clk_set_rate`]: https://docs.kernel.org/core-api/kernel-api.html#c.clk_set_rate
>          #[inline]
>          pub fn set_rate(&self, rate: Hertz) -> Result {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> -            // [`clk_set_rate`].
> +            // SAFETY: By the type invariants, `self.as_raw()` is a valid argument
> +            // for [`clk_set_rate`].
>              to_result(unsafe { bindings::clk_set_rate(self.as_raw(), rate.as_hz()) })
>          }
>      }
>  
> -    impl Drop for Clk {
> +    impl<T: ClkState> Drop for Clk<T> {
>          fn drop(&mut self) {
> -            // SAFETY: By the type invariants, self.as_raw() is a valid argument for [`clk_put`].
> -            unsafe { bindings::clk_put(self.as_raw()) };
> -        }
> -    }
> -
> -    /// A reference-counted optional clock.
> -    ///
> -    /// A lightweight wrapper around an optional [`Clk`]. An [`OptionalClk`] represents a [`Clk`]
> -    /// that a driver can function without but may improve performance or enable additional
> -    /// features when available.
> -    ///
> -    /// # Invariants
> -    ///
> -    /// An [`OptionalClk`] instance encapsulates a [`Clk`] with either a valid [`struct clk`] or
> -    /// `NULL` pointer.
> -    ///
> -    /// Instances of this type are reference-counted. Calling [`OptionalClk::get`] ensures that the
> -    /// allocation remains valid for the lifetime of the [`OptionalClk`].
> -    ///
> -    /// # Examples
> -    ///
> -    /// The following example demonstrates how to obtain and configure an optional clock for a
> -    /// device. The code functions correctly whether or not the clock is available.
> -    ///
> -    /// ```
> -    /// use kernel::clk::{OptionalClk, Hertz};
> -    /// use kernel::device::Device;
> -    /// use kernel::error::Result;
> -    ///
> -    /// fn configure_clk(dev: &Device) -> Result {
> -    ///     let clk = OptionalClk::get(dev, Some(c"apb_clk"))?;
> -    ///
> -    ///     clk.prepare_enable()?;
> -    ///
> -    ///     let expected_rate = Hertz::from_ghz(1);
> -    ///
> -    ///     if clk.rate() != expected_rate {
> -    ///         clk.set_rate(expected_rate)?;
> -    ///     }
> -    ///
> -    ///     clk.disable_unprepare();
> -    ///     Ok(())
> -    /// }
> -    /// ```
> -    ///
> -    /// [`struct clk`]: https://docs.kernel.org/driver-api/clk.html
> -    pub struct OptionalClk(Clk);
> -
> -    impl OptionalClk {
> -        /// Gets [`OptionalClk`] corresponding to a [`Device`] and a connection id.
> -        ///
> -        /// Equivalent to the kernel's [`clk_get_optional`] API.
> -        ///
> -        /// [`clk_get_optional`]:
> -        /// https://docs.kernel.org/core-api/kernel-api.html#c.clk_get_optional
> -        pub fn get(dev: &Device, name: Option<&CStr>) -> Result<Self> {
> -            let con_id = name.map_or(ptr::null(), |n| n.as_char_ptr());
> -
> -            // SAFETY: It is safe to call [`clk_get_optional`] for a valid device pointer.
> -            //
> -            // INVARIANT: The reference-count is decremented when [`OptionalClk`] goes out of
> -            // scope.
> -            Ok(Self(Clk(from_err_ptr(unsafe {
> -                bindings::clk_get_optional(dev.as_raw(), con_id)
> -            })?)))
> -        }
> -    }
> -
> -    // Make [`OptionalClk`] behave like [`Clk`].
> -    impl Deref for OptionalClk {
> -        type Target = Clk;
> +            if T::DISABLE_ON_DROP {
> +                // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> +                // [`clk_disable`].
> +                unsafe { bindings::clk_disable(self.as_raw()) };
> +            }
> +
> +            if T::UNPREPARE_ON_DROP {
> +                // SAFETY: By the type invariants, self.as_raw() is a valid argument for
> +                // [`clk_unprepare`].
> +                unsafe { bindings::clk_unprepare(self.as_raw()) };

With this `Drop` can sleep, this is probably worth mentioning in the
doccomment.