use alloc::borrow::ToOwned;
use alloc::string::String;
use hashbrown::HashMap;
use scale_info::{PortableRegistry, TypeDef, Variant, form::PortableForm};
#[derive(Debug, Clone)]
pub struct VariantIndex {
by_name: HashMap<String, usize>,
by_index: HashMap<u8, usize>,
}
impl VariantIndex {
pub fn build(variant_id: Option<u32>, types: &PortableRegistry) -> Self {
let Some(variants) = Self::get(variant_id, types) else {
return Self::empty();
};
let mut by_name = HashMap::new();
let mut by_index = HashMap::new();
for (pos, variant) in variants.iter().enumerate() {
by_name.insert(variant.name.to_owned(), pos);
by_index.insert(variant.index, pos);
}
Self { by_name, by_index }
}
pub fn empty() -> Self {
Self {
by_name: Default::default(),
by_index: Default::default(),
}
}
pub fn get(
variant_id: Option<u32>,
types: &PortableRegistry,
) -> Option<&[Variant<PortableForm>]> {
let variant_id = variant_id?;
let TypeDef::Variant(v) = &types.resolve(variant_id)?.type_def else {
return None;
};
Some(&v.variants)
}
pub fn lookup_by_name<'a, K>(
&self,
name: &K,
variant_id: Option<u32>,
types: &'a PortableRegistry,
) -> Option<&'a Variant<PortableForm>>
where
String: alloc::borrow::Borrow<K>,
K: core::hash::Hash + Eq + ?Sized,
{
let pos = *self.by_name.get(name)?;
let variants = Self::get(variant_id, types)?;
variants.get(pos)
}
pub fn lookup_by_index<'a>(
&self,
index: u8,
variant_id: Option<u32>,
types: &'a PortableRegistry,
) -> Option<&'a Variant<PortableForm>> {
let pos = *self.by_index.get(&index)?;
let variants = Self::get(variant_id, types)?;
variants.get(pos)
}
}