diff --git a/afl/src/metamap.rs b/afl/src/metamap.rs index 0fbe44c1c3..979e641fcf 100644 --- a/afl/src/metamap.rs +++ b/afl/src/metamap.rs @@ -240,3 +240,103 @@ impl MetaInstanceMap { } } } + +pub struct NamedAnyMap { + map: HashMap>>, +} + +impl NamedAnyMap { + pub fn get(&self, name: &'static str) -> Option<&T> + where + T: Any, + { + match self.map.get(&TypeId::of::()) { + None => None, + Some(h) => h + .get(&name) + .map(|x| x.as_ref().downcast_ref::().unwrap()), + } + } + + pub fn get_mut(&mut self, name: &'static str) -> Option<&mut T> + where + T: Any, + { + match self.map.get_mut(&TypeId::of::()) { + None => None, + Some(h) => h + .get_mut(&name) + .map(|x| x.as_mut().downcast_mut::().unwrap()), + } + } + + pub fn get_all( + &self, + ) -> Option>, fn(&Box) -> &T>> + where + T: Any, + { + match self.map.get(&TypeId::of::()) { + None => None, + Some(h) => Some(h.values().map(|x| x.as_ref().downcast_ref::().unwrap())), + } + } + + pub fn get_all_mut( + &mut self, + ) -> Option< + core::iter::Map>, fn(&mut Box) -> &mut T>, + > + where + T: Any, + { + match self.map.get_mut(&TypeId::of::()) { + None => None, + Some(h) => Some( + h.values_mut() + .map(|x| x.as_mut().downcast_mut::().unwrap()), + ), + } + } + + pub fn insert(&mut self, t: T, name: &'static str) + where + T: Any, + { + let typeid = TypeId::of::(); + if !self.map.contains_key(&typeid) { + self.map.insert(typeid, HashMap::default()); + } + self.map + .get_mut(&typeid) + .unwrap() + .insert(name, Box::new(t)); + } + + pub fn len(&self) -> usize { + self.map.len() + } + + pub fn contains_type(&self) -> bool + where + T: Any, + { + self.map.contains_key(&TypeId::of::()) + } + + pub fn contains(&self, name: &'static str) -> bool + where + T: Any, + { + match self.map.get(&TypeId::of::()) { + None => false, + Some(h) => h.contains_key(&name), + } + } + + pub fn new() -> Self { + Self { + map: HashMap::default(), + } + } +}