nix_builder/src/nix_util/activity_tree.rs

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)
    }
}