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}