rust: driver: let probe() return impl PinInit<Self, Error>

The driver model defines the lifetime of the private data stored in (and
owned by) a bus device to be valid from when the driver is bound to a
device (i.e. from successful probe()) until the driver is unbound from
the device.

This is already taken care of by the Rust implementation of the driver
model. However, we still ask drivers to return a Result<Pin<KBox<Self>>>
from probe().

Unlike in C, where we do not have the concept of initializers, but
rather deal with uninitialized memory, drivers can just return an
impl PinInit<Self, Error> instead.

This contributes to more clarity to the fact that a driver returns it's
device private data in probe() and the Rust driver model owns the data,
manages the lifetime and - considering the lifetime - provides (safe)
accessors for the driver.

Hence, let probe() functions return an impl PinInit<Self, Error> instead
of Result<Pin<KBox<Self>>>.

Reviewed-by: Alice Ryhl <[e-mail]>
Acked-by: Viresh Kumar <[e-mail]>
Reviewed-by: Alexandre Courbot <[e-mail]>
Acked-by: Greg Kroah-Hartman <[e-mail]>
Signed-off-by: Danilo Krummrich <[e-mail]>

drivers/cpufreq/rcpufreq_dt.rs

···
     fn probe(
         pdev: &platform::Device<Core>,
         _id_info: Option<&Self::IdInfo>,
-    ) -> Result<Pin<KBox<Self>>> {
+    ) -> impl PinInit<Self, Error> {
         cpufreq::Registration::<CPUFreqDTDriver>::new_foreign_owned(pdev.as_ref())?;
-        Ok(KBox::new(Self {}, GFP_KERNEL)?.into())
+        Ok(Self {})
     }
 }
 

drivers/gpu/drm/nova/driver.rs

···
     type IdInfo = ();
     const ID_TABLE: auxiliary::IdTable<Self::IdInfo> = &AUX_TABLE;
 
-    fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
+    fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
         let data = try_pin_init!(NovaData { adev: adev.into() });
 
         let drm = drm::Device::<Self>::new(adev.as_ref(), data)?;
         drm::Registration::new_foreign_owned(&drm, adev.as_ref(), 0)?;
 
-        Ok(KBox::new(Self { drm }, GFP_KERNEL)?.into())
+        Ok(Self { drm })
     }
 }
 

drivers/gpu/drm/tyr/driver.rs

···
     fn probe(
         pdev: &platform::Device<Core>,
         _info: Option<&Self::IdInfo>,
-    ) -> Result<Pin<KBox<Self>>> {
+    ) -> impl PinInit<Self, Error> {
         let core_clk = Clk::get(pdev.as_ref(), Some(c_str!("core")))?;
         let stacks_clk = OptionalClk::get(pdev.as_ref(), Some(c_str!("stacks")))?;
         let coregroup_clk = OptionalClk::get(pdev.as_ref(), Some(c_str!("coregroup")))?;
···
         let tdev: ARef<TyrDevice> = drm::Device::new(pdev.as_ref(), data)?;
         drm::driver::Registration::new_foreign_owned(&tdev, pdev.as_ref(), 0)?;
 
-        let driver = KBox::pin_init(try_pin_init!(TyrDriver { device: tdev }), GFP_KERNEL)?;
+        let driver = TyrDriver { device: tdev };
 
         // We need this to be dev_info!() because dev_dbg!() does not work at
         // all in Rust for now, and we need to see whether probe succeeded.

drivers/gpu/nova-core/driver.rs

···
     type IdInfo = ();
     const ID_TABLE: pci::IdTable<Self::IdInfo> = &PCI_TABLE;
 
-    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
-        dev_dbg!(pdev.as_ref(), "Probe Nova Core GPU driver.\n");
+    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
+        pin_init::pin_init_scope(move || {
+            dev_dbg!(pdev.as_ref(), "Probe Nova Core GPU driver.\n");
 
-        pdev.enable_device_mem()?;
-        pdev.set_master();
+            pdev.enable_device_mem()?;
+            pdev.set_master();
 
-        let devres_bar = Arc::pin_init(
-            pdev.iomap_region_sized::<BAR0_SIZE>(0, c_str!("nova-core/bar0")),
-            GFP_KERNEL,
-        )?;
+            let bar = Arc::pin_init(
+                pdev.iomap_region_sized::<BAR0_SIZE>(0, c_str!("nova-core/bar0")),
+                GFP_KERNEL,
+            )?;
 
-        // Used to provided a `&Bar0` to `Gpu::new` without tying it to the lifetime of
-        // `devres_bar`.
-        let bar_clone = Arc::clone(&devres_bar);
-        let bar = bar_clone.access(pdev.as_ref())?;
-
-        let this = KBox::pin_init(
-            try_pin_init!(Self {
-                gpu <- Gpu::new(pdev, devres_bar, bar),
+            Ok(try_pin_init!(Self {
+                gpu <- Gpu::new(pdev, bar.clone(), bar.access(pdev.as_ref())?),
                 _reg: auxiliary::Registration::new(
                     pdev.as_ref(),
                     c_str!("nova-drm"),
                     0, // TODO[XARR]: Once it lands, use XArray; for now we don't use the ID.
                     crate::MODULE_NAME
                 )?,
-            }),
-            GFP_KERNEL,
-        )?;
-
-        Ok(this)
+            }))
+        })
     }
 
     fn unbind(pdev: &pci::Device<Core>, this: Pin<&Self>) {

rust/kernel/auxiliary.rs

···
         let info = T::ID_TABLE.info(id.index());
 
         from_result(|| {
-            let data = T::probe(adev, info)?;
+            let data = T::probe(adev, info);
 
-            adev.as_ref().set_drvdata(data);
+            adev.as_ref().set_drvdata(data)?;
             Ok(0)
         })
     }
···
     /// Auxiliary driver probe.
     ///
     /// Called when an auxiliary device is matches a corresponding driver.
-    fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> Result<Pin<KBox<Self>>>;
+    fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> impl PinInit<Self, Error>;
 }
 
 /// The auxiliary device representation.

rust/kernel/cpufreq.rs

···
 ///     fn probe(
 ///         pdev: &platform::Device<Core>,
 ///         _id_info: Option<&Self::IdInfo>,
-///     ) -> Result<Pin<KBox<Self>>> {
+///     ) -> impl PinInit<Self, Error> {
 ///         cpufreq::Registration::<SampleDriver>::new_foreign_owned(pdev.as_ref())?;
-///         Ok(KBox::new(Self {}, GFP_KERNEL)?.into())
+///         Ok(Self {})
 ///     }
 /// }
 /// ```

rust/kernel/device.rs

···
 
 use crate::{
     bindings, fmt,
+    prelude::*,
     sync::aref::ARef,
     types::{ForeignOwnable, Opaque},
 };
···
 
 impl Device<CoreInternal> {
     /// Store a pointer to the bound driver's private data.
-    pub fn set_drvdata(&self, data: impl ForeignOwnable) {
+    pub fn set_drvdata<T: 'static>(&self, data: impl PinInit<T, Error>) -> Result {
+        let data = KBox::pin_init(data, GFP_KERNEL)?;
+
         // SAFETY: By the type invariants, `self.as_raw()` is a valid pointer to a `struct device`.
-        unsafe { bindings::dev_set_drvdata(self.as_raw(), data.into_foreign().cast()) }
+        unsafe { bindings::dev_set_drvdata(self.as_raw(), data.into_foreign().cast()) };
+
+        Ok(())
     }
 
     /// Take ownership of the private data stored in this [`Device`].

rust/kernel/driver.rs

···
 //!     const ACPI_ID_TABLE: Option<acpi::IdTable<Self::IdInfo>> = None;
 //!
 //!     /// Driver probe.
-//!     fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> Result<Pin<KBox<Self>>>;
+//!     fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> impl PinInit<Self, Error>;
 //!
 //!     /// Driver unbind (optional).
 //!     fn unbind(dev: &Device<device::Core>, this: Pin<&Self>) {
···
 //!
 //! For specific examples see [`auxiliary::Driver`], [`pci::Driver`] and [`platform::Driver`].
 //!
-//! The `probe()` callback should return a `Result<Pin<KBox<Self>>>`, i.e. the driver's private
+//! The `probe()` callback should return a `impl PinInit<Self, Error>`, i.e. the driver's private
 //! data. The bus abstraction should store the pointer in the corresponding bus device. The generic
 //! [`Device`] infrastructure provides common helpers for this purpose on its
 //! [`Device<CoreInternal>`] implementation.

rust/kernel/io/mem.rs

···
     ///    fn probe(
     ///       pdev: &platform::Device<Core>,
     ///       info: Option<&Self::IdInfo>,
-    ///    ) -> Result<Pin<KBox<Self>>> {
+    ///    ) -> impl PinInit<Self, Error> {
     ///       let offset = 0; // Some offset.
     ///
     ///       // If the size is known at compile time, use [`Self::iomap_sized`].
···
     ///
     ///       io.write32_relaxed(data, offset);
     ///
-    ///       # Ok(KBox::new(SampleDriver, GFP_KERNEL)?.into())
+    ///       # Ok(SampleDriver)
     ///     }
     /// }
     /// ```
···
     ///    fn probe(
     ///       pdev: &platform::Device<Core>,
     ///       info: Option<&Self::IdInfo>,
-    ///    ) -> Result<Pin<KBox<Self>>> {
+    ///    ) -> impl PinInit<Self, Error> {
     ///       let offset = 0; // Some offset.
     ///
     ///       // Unlike [`Self::iomap_sized`], here the size of the memory region
···
     ///
     ///       io.try_write32_relaxed(data, offset)?;
     ///
-    ///       # Ok(KBox::new(SampleDriver, GFP_KERNEL)?.into())
+    ///       # Ok(SampleDriver)
     ///     }
     /// }
     /// ```

rust/kernel/pci.rs

···
         let info = T::ID_TABLE.info(id.index());
 
         from_result(|| {
-            let data = T::probe(pdev, info)?;
+            let data = T::probe(pdev, info);
 
-            pdev.as_ref().set_drvdata(data);
+            pdev.as_ref().set_drvdata(data)?;
             Ok(0)
         })
     }
···
 ///     fn probe(
 ///         _pdev: &pci::Device<Core>,
 ///         _id_info: &Self::IdInfo,
-///     ) -> Result<Pin<KBox<Self>>> {
+///     ) -> impl PinInit<Self, Error> {
 ///         Err(ENODEV)
 ///     }
 /// }
···
     ///
     /// Called when a new pci device is added or discovered. Implementers should
     /// attempt to initialize the device here.
-    fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> Result<Pin<KBox<Self>>>;
+    fn probe(dev: &Device<device::Core>, id_info: &Self::IdInfo) -> impl PinInit<Self, Error>;
 
     /// PCI driver unbind.
     ///

rust/kernel/platform.rs

···
         let info = <Self as driver::Adapter>::id_info(pdev.as_ref());
 
         from_result(|| {
-            let data = T::probe(pdev, info)?;
+            let data = T::probe(pdev, info);
 
-            pdev.as_ref().set_drvdata(data);
+            pdev.as_ref().set_drvdata(data)?;
             Ok(0)
         })
     }
···
 ///     fn probe(
 ///         _pdev: &platform::Device<Core>,
 ///         _id_info: Option<&Self::IdInfo>,
-///     ) -> Result<Pin<KBox<Self>>> {
+///     ) -> impl PinInit<Self, Error> {
 ///         Err(ENODEV)
 ///     }
 /// }
···
     ///
     /// Called when a new platform device is added or discovered.
     /// Implementers should attempt to initialize the device here.
-    fn probe(dev: &Device<device::Core>, id_info: Option<&Self::IdInfo>)
-        -> Result<Pin<KBox<Self>>>;
+    fn probe(
+        dev: &Device<device::Core>,
+        id_info: Option<&Self::IdInfo>,
+    ) -> impl PinInit<Self, Error>;
 
     /// Platform driver unbind.
     ///

rust/kernel/usb.rs

···
             let id = unsafe { &*id.cast::<DeviceId>() };
 
             let info = T::ID_TABLE.info(id.index());
-            let data = T::probe(intf, id, info)?;
+            let data = T::probe(intf, id, info);
 
             let dev: &device::Device<device::CoreInternal> = intf.as_ref();
-            dev.set_drvdata(data);
+            dev.set_drvdata(data)?;
             Ok(0)
         })
     }
···
 ///         _interface: &usb::Interface<Core>,
 ///         _id: &usb::DeviceId,
 ///         _info: &Self::IdInfo,
-///     ) -> Result<Pin<KBox<Self>>> {
+///     ) -> impl PinInit<Self, Error> {
 ///         Err(ENODEV)
 ///     }
 ///
···
         interface: &Interface<device::Core>,
         id: &DeviceId,
         id_info: &Self::IdInfo,
-    ) -> Result<Pin<KBox<Self>>>;
+    ) -> impl PinInit<Self, Error>;
 
     /// USB driver disconnect.
     ///

samples/rust/rust_debugfs.rs

···
     fn probe(
         pdev: &platform::Device<Core>,
         _info: Option<&Self::IdInfo>,
-    ) -> Result<Pin<KBox<Self>>> {
-        let result = KBox::try_pin_init(RustDebugFs::new(pdev), GFP_KERNEL)?;
-        // We can still mutate fields through the files which are atomic or mutexed:
-        result.counter.store(91, Ordering::Relaxed);
-        {
-            let mut guard = result.inner.lock();
-            guard.x = guard.y;
-            guard.y = 42;
-        }
-        Ok(result)
+    ) -> impl PinInit<Self, Error> {
+        RustDebugFs::new(pdev).pin_chain(|this| {
+            this.counter.store(91, Ordering::Relaxed);
+            {
+                let mut guard = this.inner.lock();
+                guard.x = guard.y;
+                guard.y = 42;
+            }
+
+            Ok(())
+        })
     }
 }
 

samples/rust/rust_dma.rs

···
     type IdInfo = ();
     const ID_TABLE: pci::IdTable<Self::IdInfo> = &PCI_TABLE;
 
-    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
-        dev_info!(pdev.as_ref(), "Probe DMA test driver.\n");
+    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
+        pin_init::pin_init_scope(move || {
+            dev_info!(pdev.as_ref(), "Probe DMA test driver.\n");
 
-        let mask = DmaMask::new::<64>();
+            let mask = DmaMask::new::<64>();
 
-        // SAFETY: There are no concurrent calls to DMA allocation and mapping primitives.
-        unsafe { pdev.dma_set_mask_and_coherent(mask)? };
+            // SAFETY: There are no concurrent calls to DMA allocation and mapping primitives.
+            unsafe { pdev.dma_set_mask_and_coherent(mask)? };
 
-        let ca: CoherentAllocation<MyStruct> =
-            CoherentAllocation::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
+            let ca: CoherentAllocation<MyStruct> =
+                CoherentAllocation::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
 
-        for (i, value) in TEST_VALUES.into_iter().enumerate() {
-            kernel::dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;
-        }
+            for (i, value) in TEST_VALUES.into_iter().enumerate() {
+                kernel::dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;
+            }
 
-        let size = 4 * page::PAGE_SIZE;
-        let pages = VVec::with_capacity(size, GFP_KERNEL)?;
+            let size = 4 * page::PAGE_SIZE;
+            let pages = VVec::with_capacity(size, GFP_KERNEL)?;
 
-        let sgt = SGTable::new(pdev.as_ref(), pages, DataDirection::ToDevice, GFP_KERNEL);
+            let sgt = SGTable::new(pdev.as_ref(), pages, DataDirection::ToDevice, GFP_KERNEL);
 
-        let drvdata = KBox::pin_init(
-            try_pin_init!(Self {
+            Ok(try_pin_init!(Self {
                 pdev: pdev.into(),
                 ca,
                 sgt <- sgt,
-            }),
-            GFP_KERNEL,
-        )?;
-
-        Ok(drvdata)
+            }))
+        })
     }
 }
 

samples/rust/rust_driver_auxiliary.rs

···
 
     const ID_TABLE: auxiliary::IdTable<Self::IdInfo> = &AUX_TABLE;
 
-    fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
+    fn probe(adev: &auxiliary::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
         dev_info!(
             adev.as_ref(),
             "Probing auxiliary driver for auxiliary device with id={}\n",
···
 
         ParentDriver::connect(adev)?;
 
-        let this = KBox::new(Self, GFP_KERNEL)?;
-
-        Ok(this.into())
+        Ok(Self)
     }
 }
 
···
 
     const ID_TABLE: pci::IdTable<Self::IdInfo> = &PCI_TABLE;
 
-    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
-        let this = KBox::new(
-            Self {
-                _reg: [
-                    auxiliary::Registration::new(pdev.as_ref(), AUXILIARY_NAME, 0, MODULE_NAME)?,
-                    auxiliary::Registration::new(pdev.as_ref(), AUXILIARY_NAME, 1, MODULE_NAME)?,
-                ],
-            },
-            GFP_KERNEL,
-        )?;
-
-        Ok(this.into())
+    fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, Error> {
+        Ok(Self {
+            _reg: [
+                auxiliary::Registration::new(pdev.as_ref(), AUXILIARY_NAME, 0, MODULE_NAME)?,
+                auxiliary::Registration::new(pdev.as_ref(), AUXILIARY_NAME, 1, MODULE_NAME)?,
+            ],
+        })
     }
 }
 

samples/rust/rust_driver_pci.rs

···
 
     const ID_TABLE: pci::IdTable<Self::IdInfo> = &PCI_TABLE;
 
-    fn probe(pdev: &pci::Device<Core>, info: &Self::IdInfo) -> Result<Pin<KBox<Self>>> {
-        let vendor = pdev.vendor_id();
-        dev_dbg!(
-            pdev.as_ref(),
-            "Probe Rust PCI driver sample (PCI ID: {}, 0x{:x}).\n",
-            vendor,
-            pdev.device_id()
-        );
+    fn probe(pdev: &pci::Device<Core>, info: &Self::IdInfo) -> impl PinInit<Self, Error> {
+        pin_init::pin_init_scope(move || {
+            let vendor = pdev.vendor_id();
+            dev_dbg!(
+                pdev.as_ref(),
+                "Probe Rust PCI driver sample (PCI ID: {}, 0x{:x}).\n",
+                vendor,
+                pdev.device_id()
+            );
 
-        pdev.enable_device_mem()?;
-        pdev.set_master();
+            pdev.enable_device_mem()?;
+            pdev.set_master();
 
-        let drvdata = KBox::pin_init(
-            try_pin_init!(Self {
+            Ok(try_pin_init!(Self {
                 bar <- pdev.iomap_region_sized::<{ Regs::END }>(0, c_str!("rust_driver_pci")),
-                pdev: pdev.into(),
                 index: *info,
-            }),
-            GFP_KERNEL,
-        )?;
+                _: {
+                    let bar = bar.access(pdev.as_ref())?;
 
-        let bar = drvdata.bar.access(pdev.as_ref())?;
-        dev_info!(
-            pdev.as_ref(),
-            "pci-testdev data-match count: {}\n",
-            Self::testdev(info, bar)?
-        );
-
-        Ok(drvdata)
+                    dev_info!(
+                        pdev.as_ref(),
+                        "pci-testdev data-match count: {}\n",
+                        Self::testdev(info, bar)?
+                    );
+                },
+                pdev: pdev.into(),
+            }))
+        })
     }
 
     fn unbind(pdev: &pci::Device<Core>, this: Pin<&Self>) {

samples/rust/rust_driver_platform.rs

···
     fn probe(
         pdev: &platform::Device<Core>,
         info: Option<&Self::IdInfo>,
-    ) -> Result<Pin<KBox<Self>>> {
+    ) -> impl PinInit<Self, Error> {
         let dev = pdev.as_ref();
 
         dev_dbg!(dev, "Probe Rust Platform driver sample.\n");
···
             Self::properties_parse(dev)?;
         }
 
-        let drvdata = KBox::new(Self { pdev: pdev.into() }, GFP_KERNEL)?;
-
-        Ok(drvdata.into())
+        Ok(Self { pdev: pdev.into() })
     }
 }
 

samples/rust/rust_driver_usb.rs

···
         intf: &usb::Interface<Core>,
         _id: &usb::DeviceId,
         _info: &Self::IdInfo,
-    ) -> Result<Pin<KBox<Self>>> {
+    ) -> impl PinInit<Self, Error> {
         let dev: &device::Device<Core> = intf.as_ref();
         dev_info!(dev, "Rust USB driver sample probed\n");
 
-        let drvdata = KBox::new(Self { _intf: intf.into() }, GFP_KERNEL)?;
-        Ok(drvdata.into())
+        Ok(Self { _intf: intf.into() })
     }
 
     fn disconnect(intf: &usb::Interface<Core>, _data: Pin<&Self>) {