use crate::Result; use std::collections::BTreeMap; use std::fmt::Display; use std::ops::Deref; use super::activity::Activity; use super::activity::ActivityRoot; #[derive(Debug, Clone)] pub(crate) struct ActivityId { id: u64, } impl ActivityId { fn new(id: u64) -> ActivityId { ActivityId { id } } } impl Deref for ActivityId { type Target = u64; fn deref(&self) -> &Self::Target { &self.id } } impl Display for ActivityId { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { Display::fmt(&**self, f) } } pub(crate) struct ActivityTree { activities: BTreeMap, } impl ActivityTree { pub(crate) fn new() -> ActivityTree { let mut activities = BTreeMap::new(); activities.insert( 0, ActivityTreeEntry { id: ActivityId::new(0), parent_id: ActivityId::new(0), child_ids: Vec::new(), activity: Activity::Root(ActivityRoot {}), }, ); ActivityTree { activities } } pub(crate) fn add_activity( &mut self, activity_id: u64, parent_id: u64, activity: Activity, ) -> Result<&ActivityTreeEntry> { // The activity_id is not yet in the tree, so we can't use get_activity_id. let activity_id = ActivityId::new(activity_id); let parent_id = self.get_activity_id(parent_id)?; match self.activities.entry(activity_id.id) { std::collections::btree_map::Entry::Vacant(vacant_entry) => { vacant_entry.insert(ActivityTreeEntry { id: activity_id.clone(), parent_id: parent_id.clone(), child_ids: Vec::new(), activity, }); } std::collections::btree_map::Entry::Occupied(_occupied_entry) => { return Err(ActivityIdAlreadyInTreeError::new(*activity_id).into()); } }; self.get_mut(&parent_id).child_ids.push(activity_id.clone()); let tree_entry = self .activities .get(&*activity_id) .expect("We just created this entry so it must exist."); Ok(tree_entry) } pub(crate) fn get_activity_id(&self, activity_id: u64) -> Result { if activity_id == 0 || self.activities.contains_key(&activity_id) { Ok(ActivityId::new(activity_id)) } else { Err(ActivityIdNotInTreeError::new(activity_id).into()) } } pub(crate) fn get_root_id(&self) -> ActivityId { // The root always exists, so there is no need to check. ActivityId::new(0) } pub(crate) fn get(&self, activity_id: &ActivityId) -> &ActivityTreeEntry { self.activities .get(activity_id) .expect("You cannot create an ActivityId if the activity is not in the tree.") } pub(crate) fn get_mut(&mut self, activity_id: &ActivityId) -> &mut ActivityTreeEntry { self.activities .get_mut(activity_id) .expect("You cannot create an ActivityId if the activity is not in the tree.") } } pub(crate) struct ActivityTreeEntry { id: ActivityId, parent_id: ActivityId, child_ids: Vec, activity: Activity, } impl ActivityTreeEntry { pub(crate) fn get_activity_id(&self) -> &ActivityId { &self.id } pub(crate) fn get_parent_id(&self) -> &ActivityId { &self.parent_id } pub(crate) fn get_activity(&self) -> &Activity { &self.activity } pub(crate) fn get_mut_activity(&mut self) -> &mut Activity { &mut self.activity } pub(crate) fn get_child_ids(&self) -> &Vec { &self.child_ids } pub(crate) fn iter_children<'tree>( &'_ self, activity_tree: &'tree ActivityTree, ) -> impl Iterator { self.child_ids .iter() .map(|child_id| activity_tree.get(child_id)) } } #[derive(Debug, Clone)] pub(crate) struct ActivityIdNotInTreeError { id: u64, } impl ActivityIdNotInTreeError { fn new(id: u64) -> ActivityIdNotInTreeError { ActivityIdNotInTreeError { id } } } impl std::fmt::Display for ActivityIdNotInTreeError { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "Activity id {} not in the tree.", self.id) } } #[derive(Debug, Clone)] pub(crate) struct ActivityIdAlreadyInTreeError { id: u64, } impl ActivityIdAlreadyInTreeError { fn new(id: u64) -> ActivityIdAlreadyInTreeError { ActivityIdAlreadyInTreeError { id } } } impl std::fmt::Display for ActivityIdAlreadyInTreeError { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "Activity id {} already in the tree.", self.id) } }