Switch to a new ActivityTree implementation.

This new implementation explicitly tracks the children of each node, eliminating many iterations through the full tree.

src/nix_util/activity_tree.rs

@@ -0,0 +1,180 @@
+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<u64, ActivityTreeEntry>,
+}
+
+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<ActivityId> {
+        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<ActivityId>,
+    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<ActivityId> {
+        &self.child_ids
+    }
+
+    pub(crate) fn iter_children<'tree>(
+        &'_ self,
+        activity_tree: &'tree ActivityTree,
+    ) -> impl Iterator<Item = &'tree ActivityTreeEntry> {
+        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)
+    }
+}