leo_passes/common/tree_node/
mod.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17use indexmap::IndexSet;
18
19use leo_span::Symbol;
20use std::{fmt::Debug, hash::Hash};
21
22///  A binary search tree to store all paths through nested conditional blocks.
23pub type ConditionalTreeNode = TreeNode<Symbol>;
24
25/// A node in a graph.
26pub trait Node: Copy + 'static + Eq + PartialEq + Debug + Hash {}
27
28impl Node for Symbol {}
29
30/// A node in a tree.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct TreeNode<N: Node> {
33    /// The current depth.
34    pub depth: usize,
35    /// The current node.
36    pub elements: IndexSet<N>, // TODO: Can optimize with bitmap if performance is bad.
37    /// A counter.
38    pub counter: usize,
39}
40
41impl<N: Node> TreeNode<N> {
42    /// Initializes a new `TreeNode` from a vector of starting elements.
43    pub fn new(elements: IndexSet<N>) -> Self {
44        Self { depth: 0, elements, counter: 0 }
45    }
46
47    /// Adds a child to the current node.
48    pub fn create_child(&mut self) -> TreeNode<N> {
49        Self { depth: self.depth + 1, elements: self.elements.clone(), counter: self.counter }
50    }
51
52    /// Removes an element from the current node.
53    /// If the element does not exist, increment an internal counter which later used to generate an error that the user attempted to await a future twice.
54    /// Returns `true` if the element was removed but not the first one in the node.
55    pub fn remove_element(&mut self, element: &N) -> bool {
56        // Check if the element is the first one in the node.
57        let is_not_first = match self.elements.first() {
58            Some(first) => first != element,
59            None => false,
60        };
61        // Remove the element from the node.
62        if !self.elements.shift_remove(element) {
63            self.counter += 1;
64            false
65        } else {
66            is_not_first
67        }
68    }
69}