1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100

///! Data structures to work with elimination trees (etree).
///! etrees arise when considering cholesky factorization, QR factorization, ...
use std::ops::{Deref, DerefMut};

pub type Parent = Option<usize>;

/// Store an etree as the parent information of each node.
/// This reflects the fact that etrees can in fact have multiple roots.
#[derive(Debug, Clone)]
pub struct Parents<S>
where
    S: Deref<Target = [Parent]>,
{
    parents: S,
}

pub type ParentsView<'a> = Parents<&'a [Parent]>;
pub type ParentsViewMut<'a> = Parents<&'a mut [Parent]>;
pub type ParentsOwned = Parents<Vec<Parent>>;

impl<S: Deref<Target = [Parent]>> Parents<S> {
    /// Get the parent of a node. Returns None if the node is a root.
    ///
    /// # Panics
    ///
    /// * if node is out of bounds
    pub fn get_parent(&self, node: usize) -> Option<usize> {
        self.parents[node]
    }

    /// Test whether a node is a root.
    ///
    /// # Panics
    ///
    /// * if node is out of bounds
    pub fn is_root(&self, node: usize) -> bool {
        self.parents[node].is_none()
    }

    /// The number of nodes in this tree.
    pub fn nb_nodes(&self) -> usize {
        self.parents.len()
    }

    /// Get a view of this object
    pub fn view(&self) -> ParentsView {
        ParentsView {
            parents: &self.parents[..],
        }
    }
}

impl<S: DerefMut<Target = [Parent]>> Parents<S> {
    /// Set the parent of a node.
    ///
    /// # Panics
    ///
    /// * if node is out of bounds
    /// * if parent is out of bounds
    pub fn set_parent(&mut self, node: usize, parent: usize) {
        assert!(parent < self.nb_nodes(), "parent is out of bounds");
        self.parents[node] = Some(parent)
    }

    /// Set a node as a root.
    ///
    /// # Panics
    ///
    /// * if node is out of bounds
    pub fn set_root(&mut self, node: usize) {
        self.parents[node] = None;
    }

    /// Give a parent to a root of the tree. No-op if the node was not a parent.
    ///
    /// # Panics
    ///
    /// if either node or parent is out of bounds
    pub fn uproot(&mut self, node: usize, parent: usize) {
        assert!(parent < self.nb_nodes(), "parent is out of bounds");
        if self.is_root(node) {
            self.set_parent(node, parent);
        }
    }

    pub fn view_mut(&mut self) -> ParentsViewMut {
        ParentsViewMut {
            parents: &mut self.parents[..],
        }
    }
}

impl Parents<Vec<Parent>> {
    /// Create a new tree with all nodes set as root
    pub fn new(nb_nodes: usize) -> ParentsOwned {
        ParentsOwned {
            parents: vec![None; nb_nodes],
        }
    }
}