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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use alloc::vec::Vec;

/// Merkle selection.
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
pub enum IndexSelection {
    /// Choose left at current depth.
    Left,
    /// Choose right at current depth.
    Right,
}

/// Merkle route.
#[derive(Clone, Eq, PartialEq, Debug)]
pub enum IndexRoute {
    /// Root of the merkle tree.
    Root,
    /// Select items from the root.
    Select(Vec<IndexSelection>),
}

impl IndexRoute {
    /// Get selection at depth, where root is considered depth 0.
    pub fn at_depth(&self, depth: usize) -> Option<IndexSelection> {
        match self {
            IndexRoute::Root => None,
            IndexRoute::Select(selections) => {
                if depth == 0 || depth > selections.len() {
                    None
                } else {
                    Some(selections[depth - 1])
                }
            },
        }
    }
}

/// Raw merkle index.
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
pub struct Index(usize);

impl Index {
    /// Root merkle index.
    pub const fn root() -> Self {
        Self(1)
    }

    /// Get left child of current index.
    pub const fn left(&self) -> Self {
        Self(2 * self.0)
    }

    /// Get right child of current index.
    pub const fn right(&self) -> Self {
        Self(2 * self.0 + 1)
    }

    /// Get the parent of current merkle index.
    pub fn parent(&self) -> Option<Self> {
        if self.0 == 1 {
            None
        } else {
            Some(Self(self.0 / 2))
        }
    }

    /// Whether this index has given descendant.
    pub fn has_descendant(&self, other: &Index) -> bool {
        match other.parent() {
            Some(parent) => {
                if parent == *self {
                    true
                } else {
                    self.has_descendant(&parent)
                }
            }
            None => false,
        }
    }

    /// From one-based index.
    pub fn from_one(value: usize) -> Option<Self> {
        if value == 0 {
            None
        } else {
            Some(Self(value))
        }
    }

    /// From zero-based index.
    pub fn from_zero(value: usize) -> Self {
        Self(value + 1)
    }

    /// Get selections from current index.
    pub fn route(&self) -> IndexRoute {
        let mut value = self.0;
        let mut selections = Vec::<IndexSelection>::new();

        loop {
            if value >> 1 == 0 {
                debug_assert!(value == 1);

                if selections.is_empty() {
                    return IndexRoute::Root
                } else {
                    selections.reverse();
                    return IndexRoute::Select(selections)
                }
            }

            let sel = value & 0b1;
            if sel == 0 {
                selections.push(IndexSelection::Left)
            } else {
                selections.push(IndexSelection::Right)
            }

            value = value >> 1;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_descendant() {
        assert!(Index::root().left().has_descendant(&Index::root().left().right().left().right().right()));
        assert!(!Index::root().left().has_descendant(&Index::root().right().right().left().right().right()));
    }
}