Repository: firecracker-microvm/firecracker
File: src/vmm/src/device_manager/mod.rs (Rust)

Code before:
                    .device()
                    .clone(),
            )
        }
    }

    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
    pub fn try_with_virtio_device_with_id<T, F, R, E>(
        &self,
        virtio_type: u32,
        id: &str,
        f: F,
    ) -> Result<R, FindDeviceError>
    where
        T: VirtioDevice + 'static + Debug,
        E: std::error::Error + 'static + Send + Sync,
        F: FnOnce(&mut T) -> Result<R, E>,
    {
        if let Some(device) = self.get_virtio_device(virtio_type, id) {
            let mut dev = device.lock().expect("Poisoned lock");
            f(dev
                .as_mut_any()
                .downcast_mut::<T>()
                .ok_or(FindDeviceError::InvalidDeviceType)?)
            .map_err(|e| FindDeviceError::InternalDeviceError(e.into()))
        } else {
            Err(FindDeviceError::DeviceNotFound)
        }
    }

    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
    pub fn with_virtio_device_with_id<T, F, R>(
        &self,
        virtio_type: u32,
        id: &str,
        f: F,
    ) -> Result<R, FindDeviceError>
    where
        T: VirtioDevice + 'static + Debug,
        F: FnOnce(&mut T) -> R,
    {
        self.try_with_virtio_device_with_id(virtio_type, id, |dev: &mut T| {
            Ok::<R, FindDeviceError>(f(dev))
        })
    }
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
/// State of devices in the system
pub struct DevicesState {
    /// MMIO devices state

Review comment (suggestion):
could the device_type be an associated constant of the `VirtioDevice` trait, so that then we don't have to pass it here and instead get it from `T::DEVICE_ID`? Would give us static guarantees that `T` always matches the passing in ID 

Code after:
            Some(
                pci_device
                    .lock()
                    .expect("Poisoned lock")
                    .virtio_device()
                    .clone(),
            )
        } else {
            let mmio_device = self
                .mmio_devices
                .get_virtio_device(virtio_type, device_id)?;
            Some(
                mmio_device
                    .inner
                    .lock()
                    .expect("Poisoned lock")
                    .device()
                    .clone(),
            )
        }
    }

    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
    pub fn try_with_virtio_device_with_id<T, F, R, E>(
        &self,
        id: &str,
        f: F,
    ) -> Result<R, FindDeviceError>
    where
        T: VirtioDevice + 'static + Debug,
        E: std::error::Error + 'static + Send + Sync,
        F: FnOnce(&mut T) -> Result<R, E>,
    {
        if let Some(device) = self.get_virtio_device(T::const_device_type(), id) {
            let mut dev = device.lock().expect("Poisoned lock");
            f(dev
                .as_mut_any()
                .downcast_mut::<T>()
                .ok_or(FindDeviceError::InvalidDeviceType)?)
            .map_err(|e| FindDeviceError::InternalDeviceError(e.into()))
        } else {
            Err(FindDeviceError::DeviceNotFound)
        }
    }

    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
    pub fn with_virtio_device_with_id<T, F, R>(&self, id: &str, f: F) -> Result<R, FindDeviceError>
    where
        T: VirtioDevice + 'static + Debug,
        F: FnOnce(&mut T) -> R,
    {