use crate::{
component::{metadata::MetaData, params::Param},
runtime_services::StandardRuntimeServices,
};
use crate::{BinaryGuid, boot_services::StandardBootServices};
use alloc::{borrow::Cow, boxed::Box, collections::BTreeMap, vec::Vec};
use core::{
any::{Any, TypeId},
cell::{Ref, RefCell, RefMut, UnsafeCell},
fmt::Debug,
marker::PhantomData,
ops::{Deref, DerefMut},
ptr,
};
use super::{
hob::{FromHob, Hob},
service::{IntoService, Service},
};
type HobParsers = BTreeMap<BinaryGuid, BTreeMap<TypeId, fn(&[u8], &mut Storage)>>;
#[derive(Debug)]
pub(crate) struct SparseVec<V> {
values: Vec<Option<V>>,
}
impl<V> SparseVec<V> {
pub const fn new() -> Self {
Self { values: Vec::new() }
}
#[inline]
pub fn contains(&self, index: usize) -> bool {
self.values.get(index).map(|v| v.is_some()).unwrap_or(false)
}
#[inline]
pub fn get(&self, index: usize) -> Option<&V> {
self.values.get(index).map(|v| v.as_ref())?
}
#[inline]
pub fn get_mut(&mut self, index: usize) -> Option<&mut V> {
self.values.get_mut(index).map(|v| v.as_mut())?
}
#[inline]
pub fn insert(&mut self, index: usize, value: V) {
if index >= self.values.len() {
self.values.resize_with(index + 1, || None);
}
self.values[index] = Some(value);
}
}
impl<'a, V> IntoIterator for &'a SparseVec<V> {
type Item = &'a Option<V>;
type IntoIter = core::slice::Iter<'a, Option<V>>;
fn into_iter(self) -> Self::IntoIter {
self.values.iter()
}
}
impl<'a, V> IntoIterator for &'a mut SparseVec<V> {
type Item = &'a mut Option<V>;
type IntoIter = core::slice::IterMut<'a, Option<V>>;
fn into_iter(self) -> Self::IntoIter {
self.values.iter_mut()
}
}
pub struct ConfigRaw(bool, Box<dyn Any>);
impl Debug for ConfigRaw {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ConfigRaw").field("is_locked", &self.0).field("raw_config", &self.1).finish()
}
}
impl ConfigRaw {
pub fn new(locked: bool, config: Box<dyn Any>) -> Self {
Self(locked, config)
}
pub fn lock(&mut self) {
self.0 = true;
}
pub(crate) fn unlock(&mut self) {
self.0 = false;
}
pub fn is_locked(&self) -> bool {
self.0
}
}
impl Deref for ConfigRaw {
type Target = dyn Any;
fn deref(&self) -> &Self::Target {
self.1.as_ref()
}
}
impl DerefMut for ConfigRaw {
fn deref_mut(&mut self) -> &mut Self::Target {
self.1.as_mut()
}
}
#[derive(Default)]
#[allow(clippy::type_complexity)]
pub(crate) struct Deferred {
queue: Vec<Box<dyn FnOnce(&mut Storage)>>,
}
impl Deferred {
pub(crate) fn add_command<F: FnOnce(&mut Storage) + 'static>(&mut self, command: F) {
self.queue.push(Box::new(command));
}
fn apply(&mut self, storage: &mut Storage) {
for command in self.queue.drain(..) {
command(storage);
}
}
#[cfg(test)]
pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
}
impl Debug for Deferred {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Deferred").field("queue", &self.queue.len()).finish()
}
}
#[derive(Debug)]
pub struct Storage {
deferred: Option<Deferred>,
configs: SparseVec<RefCell<ConfigRaw>>,
config_indices: BTreeMap<TypeId, usize>,
services: SparseVec<&'static dyn Any>,
service_indices: BTreeMap<TypeId, usize>,
hob_parsers: HobParsers,
hobs: SparseVec<Vec<Box<dyn Any>>>,
hob_indices: BTreeMap<TypeId, usize>,
boot_services: StandardBootServices,
runtime_services: StandardRuntimeServices,
image_handle: Option<r_efi::efi::Handle>,
}
impl Default for Storage {
fn default() -> Self {
Self::new()
}
}
impl Storage {
pub const fn new() -> Self {
Self {
deferred: None,
configs: SparseVec::new(),
config_indices: BTreeMap::new(),
services: SparseVec::new(),
service_indices: BTreeMap::new(),
hob_parsers: BTreeMap::new(),
hobs: SparseVec::new(),
hob_indices: BTreeMap::new(),
boot_services: StandardBootServices::new_uninit(),
runtime_services: StandardRuntimeServices::new_uninit(),
image_handle: None,
}
}
pub(crate) fn apply_deferred(&mut self) {
if let Some(mut deferred) = self.deferred.take() {
deferred.apply(self);
}
}
pub(crate) fn deferred(&mut self) -> &mut Deferred {
if self.deferred.is_none() {
self.deferred = Some(Deferred::default());
}
self.deferred.as_mut().unwrap()
}
pub fn set_boot_services(&mut self, bs: StandardBootServices) {
self.boot_services = bs;
}
pub fn set_runtime_services(&mut self, rs: StandardRuntimeServices) {
self.runtime_services = rs;
}
pub fn boot_services(&self) -> &StandardBootServices {
&self.boot_services
}
pub fn runtime_services(&self) -> &StandardRuntimeServices {
&self.runtime_services
}
pub fn set_image_handle(&mut self, handle: r_efi::efi::Handle) {
self.image_handle = Some(handle);
}
pub fn image_handle(&self) -> Option<r_efi::efi::Handle> {
self.image_handle
}
pub(crate) fn register_config<C: Default + 'static>(&mut self) -> usize {
let idx = self.config_indices.len();
*self.config_indices.entry(TypeId::of::<C>()).or_insert(idx)
}
pub(crate) fn add_config_default_if_not_present<C: Default + 'static>(&mut self) -> usize {
let idx = self.register_config::<C>();
if !self.configs.contains(idx) {
self.configs.insert(idx, RefCell::new(ConfigRaw::new(true, Box::<C>::default())));
}
idx
}
pub fn add_config<C: Default + 'static>(&mut self, config: C) {
let id = self.register_config::<C>();
self.configs.insert(id, RefCell::new(ConfigRaw::new(true, Box::new(config))));
}
pub fn get_config<C: Default + 'static>(&self) -> Option<crate::component::params::Config<'_, C>> {
let id = self.config_indices.get(&TypeId::of::<C>())?;
let untyped = self.get_raw_config(*id);
Some(crate::component::params::Config::from(untyped))
}
pub fn get_config_mut<C: Default + 'static>(&mut self) -> Option<crate::component::params::ConfigMut<'_, C>> {
let id = self.config_indices.get(&TypeId::of::<C>())?;
let untyped = self.get_raw_config_mut(*id);
Some(crate::component::params::ConfigMut::from(untyped))
}
pub(crate) fn get_raw_config(&self, id: usize) -> Ref<'_, ConfigRaw> {
self.configs
.get(id)
.unwrap_or_else(|| panic!("Could not find Config value when with id [{id}] it should always exist."))
.borrow()
}
pub(crate) fn get_raw_config_mut(&self, id: usize) -> RefMut<'_, ConfigRaw> {
self.configs
.get(id)
.unwrap_or_else(|| panic!("Could not find Config value when with id [{id}] it should always exist."))
.borrow_mut()
}
pub(crate) fn unlock_config(&self, id: usize) {
if let Some(config) = self.configs.get(id) {
config.borrow_mut().unlock();
}
}
pub fn lock_configs(&self) {
(&self.configs).into_iter().flatten().for_each(|config| config.borrow_mut().lock());
}
pub(crate) fn register_service<C: ?Sized + 'static>(&mut self) -> usize {
self.get_or_register_service(TypeId::of::<C>())
}
pub(crate) fn get_or_register_service(&mut self, id: TypeId) -> usize {
let idx = self.service_indices.len();
*self.service_indices.entry(id).or_insert(idx)
}
pub(crate) fn insert_service(&mut self, id: usize, service: &'static dyn Any) {
self.services.insert(id, service);
}
pub fn add_service<S: IntoService + 'static>(&mut self, service: S) {
service.register(self);
}
pub(crate) fn get_raw_service(&self, id: usize) -> Option<&'static dyn Any> {
self.services.get(id).copied()
}
pub fn get_service<S: ?Sized + 'static>(&self) -> Option<Service<S>> {
let idx = *self.service_indices.get(&TypeId::of::<S>())?;
Some(Service::from(self.get_raw_service(idx)?))
}
pub(crate) fn add_hob_parser<T: FromHob>(&mut self) {
self.hob_parsers.entry(T::HOB_GUID).or_default().insert(TypeId::of::<T>(), T::register);
}
pub(crate) fn register_hob<T: FromHob>(&mut self) -> usize {
self.get_or_register_hob(TypeId::of::<T>())
}
pub(crate) fn get_or_register_hob(&mut self, id: TypeId) -> usize {
let idx = self.hob_indices.len();
let idx = self.hob_indices.entry(id).or_insert(idx);
if self.hobs.get(*idx).is_none() {
self.hobs.insert(*idx, Vec::new());
}
*idx
}
pub(crate) fn add_hob<H: FromHob>(&mut self, hob: H) {
let id = self.register_hob::<H>(); self.hobs.get_mut(id).expect("Hob Index should always exist.").push(Box::new(hob));
}
pub(crate) fn get_raw_hob(&self, id: usize) -> &[Box<dyn Any>] {
self.hobs
.get(id)
.unwrap_or_else(|| panic!("Could not find Hob value when with id [{id}] it should always exist."))
}
pub fn get_hob<T: FromHob>(&self) -> Option<Hob<'_, T>> {
let id = self.hob_indices.get(&TypeId::of::<T>())?;
self.hobs.get(*id).and_then(|hob| {
if hob.is_empty() {
return None;
}
Some(Hob::from(hob.as_slice()))
})
}
pub fn get_hob_parsers(&self, guid: &BinaryGuid) -> Vec<fn(&[u8], &mut Storage)> {
self.hob_parsers.get(guid).map(|type_map| type_map.values().copied().collect()).unwrap_or_default()
}
}
#[derive(Copy, Clone)]
pub struct UnsafeStorageCell<'s>(*mut Storage, PhantomData<(&'s Storage, &'s UnsafeCell<Storage>)>);
unsafe impl Send for UnsafeStorageCell<'_> {}
unsafe impl Sync for UnsafeStorageCell<'_> {}
impl<'s> From<&'s mut Storage> for UnsafeStorageCell<'s> {
fn from(storage: &'s mut Storage) -> Self {
UnsafeStorageCell::new_mutable(storage)
}
}
impl<'s> From<&'s Storage> for UnsafeStorageCell<'s> {
fn from(storage: &'s Storage) -> Self {
UnsafeStorageCell::new_readonly(storage)
}
}
impl<'s> UnsafeStorageCell<'s> {
#[inline]
pub fn new_readonly(storage: &'s Storage) -> Self {
Self(ptr::from_ref(storage).cast_mut(), PhantomData)
}
#[inline]
pub fn new_mutable(storage: &'s mut Storage) -> Self {
Self(ptr::from_mut(storage), PhantomData)
}
#[inline]
pub unsafe fn storage_mut(self) -> &'s mut Storage {
unsafe { &mut *self.0 }
}
#[inline]
pub unsafe fn storage(self) -> &'s Storage {
unsafe { &*self.0 }
}
}
unsafe impl Param for &mut Storage {
type State = ();
type Item<'storage, 'state> = &'storage mut Storage;
unsafe fn get_param<'storage, 'state>(
_state: &'state Self::State,
storage: UnsafeStorageCell<'storage>,
) -> Self::Item<'storage, 'state> {
unsafe { storage.storage_mut() }
}
fn validate(_state: &Self::State, _storage: UnsafeStorageCell) -> bool {
true
}
fn init_state(_storage: &mut Storage, meta: &mut MetaData) -> Result<Self::State, Cow<'static, str>> {
if meta.access().has_writes_all_configs() {
return Err(Cow::from("&mut Storage conflicts with a previous &mut Storage access."));
}
if meta.access().has_reads_all_configs() {
return Err(Cow::from("&mut Storage conflicts with a previous &Storage access."));
}
if meta.access().has_any_config_write() {
return Err(Cow::from("&mut Storage conflicts with a previous ConfigMut<T> access."));
}
if meta.access().has_any_config_read() {
return Err(Cow::from("&mut Storage conflicts with a previous Config<T> access."));
}
meta.access_mut().writes_all_configs();
Ok(())
}
}
unsafe impl Param for &Storage {
type State = ();
type Item<'storage, 'state> = &'storage Storage;
unsafe fn get_param<'storage, 'state>(
_state: &'state Self::State,
storage: UnsafeStorageCell<'storage>,
) -> Self::Item<'storage, 'state> {
unsafe { storage.storage() }
}
fn validate(_state: &Self::State, _storage: UnsafeStorageCell) -> bool {
true
}
fn init_state(_storage: &mut Storage, meta: &mut MetaData) -> Result<Self::State, Cow<'static, str>> {
if meta.access().has_writes_all_configs() {
return Err(Cow::from("&Storage conflicts with a previous &mut Storage access."));
}
if meta.access().has_any_config_write() {
return Err(Cow::from("&Storage conflicts with a previous ConfigMut<T> access."));
}
meta.access_mut().reads_all_configs();
Ok(())
}
}
#[cfg(test)]
#[coverage(off)]
mod tests {
use super::*;
#[test]
fn validate_iter_works() {
let mut v: SparseVec<u32> = SparseVec::new();
v.insert(4, 100);
v.insert(9, 200);
v.insert(75, 300);
assert_eq!(v.into_iter().filter(|o| o.is_some()).count(), 3);
for v in &mut v {
if let Some(v) = v.as_mut() {
*v += 1;
}
}
assert_eq!(v.get(4), Some(&101));
assert_eq!(v.get(9), Some(&201));
assert_eq!(v.get(75), Some(&301));
for v in &mut v {
*v = None;
}
assert_eq!(v.into_iter().filter(|o| o.is_some()).count(), 0);
}
#[test]
fn test_hob_functionality() {
use crate as patina;
#[derive(FromHob, zerocopy_derive::FromBytes)]
#[repr(C)]
#[hob = "12345678-1234-1234-1234-123456789012"]
struct MyStruct;
let mut storage = Storage::new();
storage.register_hob::<MyStruct>();
assert!(storage.get_hob::<MyStruct>().is_none());
assert!(storage.get_hob_parsers(&MyStruct::HOB_GUID).is_empty());
storage.add_hob(MyStruct);
assert!(storage.get_hob::<MyStruct>().is_some());
assert_eq!(storage.get_hob::<MyStruct>().unwrap().iter().count(), 1);
storage.add_hob(MyStruct);
assert_eq!(storage.get_hob::<MyStruct>().unwrap().iter().count(), 2);
storage.add_hob_parser::<MyStruct>();
assert!(!storage.get_hob_parsers(&MyStruct::HOB_GUID).is_empty());
}
#[test]
fn test_services_still_work_if_storage_requires_re_alloc() {
use crate as patina;
trait TestService {
fn test(&self) -> usize;
}
#[derive(IntoService)]
#[service(dyn TestService)]
struct TestServiceImpl {
id: usize,
}
impl TestService for TestServiceImpl {
fn test(&self) -> usize {
self.id
}
}
let mut storage = Storage::new();
storage.add_service(TestServiceImpl { id: 42 });
let service = storage.get_service::<dyn TestService>().unwrap();
assert_eq!(service.test(), 42);
drop(storage);
assert_eq!(service.test(), 42);
}
#[test]
fn test_apply_deferred_storage() {
use crate as patina;
use patina::component::params::Commands;
trait TestService {
fn test(&self) -> usize;
}
#[derive(IntoService)]
#[service(dyn TestService)]
struct TestServiceImpl {
id: usize,
}
impl TestService for TestServiceImpl {
fn test(&self) -> usize {
self.id
}
}
let mut storage = Storage::new();
assert!(storage.get_service::<dyn TestService>().is_none());
{
let mut commands = unsafe { <Commands as Param>::get_param(&(), UnsafeStorageCell::from(&mut storage)) };
commands.add_service(TestServiceImpl { id: 42 });
}
assert!(storage.get_service::<dyn TestService>().is_none());
storage.apply_deferred();
assert!(storage.get_service::<dyn TestService>().is_some());
let service = storage.get_service::<dyn TestService>().unwrap();
assert_eq!(service.test(), 42);
}
#[test]
fn test_mutable_storage_param_conflict_scenarios() {
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<&mut Storage as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&mut Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&mut Storage conflicts with a previous &mut Storage access."
);
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<&Storage as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&mut Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&mut Storage conflicts with a previous &Storage access."
);
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<crate::component::params::ConfigMut<u32> as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&mut Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&mut Storage conflicts with a previous ConfigMut<T> access."
);
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<crate::component::params::Config<u32> as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&mut Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&mut Storage conflicts with a previous Config<T> access."
);
}
#[test]
fn test_immutable_storage_param_conflict_scenarios() {
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<&mut Storage as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&Storage conflicts with a previous &mut Storage access."
);
let mut storage = Storage::new();
let mut metadata = MetaData::new::<bool>();
assert!(<crate::component::params::ConfigMut<u32> as Param>::init_state(&mut storage, &mut metadata).is_ok());
assert_eq!(
<&Storage as Param>::init_state(&mut storage, &mut metadata).unwrap_err(),
"&Storage conflicts with a previous ConfigMut<T> access."
);
}
}