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

use crate::{prelude::*, utils};
use alloc::collections::BTreeMap;

pub fn is_left_index(index: usize) -> bool {
    index % 2 == 0
}

pub fn get_sibling_index(index: usize) -> usize {
    if is_left_index(index) {
        // Right sibling index
        return index + 1;
    }
    // Left sibling index
    index - 1
}

pub fn sibling_indices(indices: &[usize]) -> Vec<usize> {
    indices.iter().cloned().map(get_sibling_index).collect()
}

pub fn parent_index(index: usize) -> usize {
    if is_left_index(index) {
        return index / 2;
    }
    get_sibling_index(index) / 2
}

pub fn parent_indices(indices: &[usize]) -> Vec<usize> {
    let mut parents: Vec<usize> = indices.iter().cloned().map(parent_index).collect();
    parents.dedup();
    parents
}

pub fn tree_depth(leaves_count: usize) -> usize {
    if leaves_count == 1 {
        1
    } else {
        let val = micromath::F32(leaves_count as f32);
        val.log2().ceil().0 as usize
    }
}

pub fn uneven_layers(tree_leaves_count: usize) -> BTreeMap<usize, usize> {
    let mut leaves_count = tree_leaves_count;
    let depth = tree_depth(tree_leaves_count);

    let mut uneven_layers = BTreeMap::new();

    for index in 0..depth {
        let uneven_layer = leaves_count % 2 != 0;
        if uneven_layer {
            uneven_layers.insert(index, leaves_count);
        }
        leaves_count = div_ceil(leaves_count, 2);
    }

    uneven_layers
}

/// Returns layered proof indices
pub fn proof_indices_by_layers(
    sorted_leaf_indices: &[usize],
    leaves_count: usize,
) -> Vec<Vec<usize>> {
    let depth = tree_depth(leaves_count);
    let uneven_layers = uneven_layers(leaves_count);

    let mut layer_nodes = sorted_leaf_indices.to_vec();
    let mut proof_indices: Vec<Vec<usize>> = Vec::new();

    for layer_index in 0..depth {
        let mut sibling_indices = sibling_indices(&layer_nodes);
        // The last node of that layer doesn't have another hash to the right, so no need to include
        // that index
        if let Some(leaves_count) = uneven_layers.get(&layer_index) {
            if let Some(layer_last_node_index) = layer_nodes.last() {
                if *layer_last_node_index == leaves_count - 1 {
                    sibling_indices.pop();
                }
            }
        }

        // Figuring out indices that are already siblings and do not require additional hash
        // to calculate the parent
        let proof_nodes_indices = utils::collections::difference(&sibling_indices, &layer_nodes);

        proof_indices.push(proof_nodes_indices);
        // Passing parent nodes indices to the next iteration cycle
        layer_nodes = parent_indices(&layer_nodes);
    }

    proof_indices
}

pub fn div_ceil(x: usize, y: usize) -> usize {
    x / y + if x % y != 0 { 1 } else { 0 }
}