Skip to main content

miden_core/mast/
dense_builder.rs

1use alloc::vec::Vec;
2
3use crate::{
4    advice::AdviceMap,
5    mast::{
6        MastForest, MastForestError, MastForestParts, MastNode, MastNodeBuilder, MastNodeContext,
7        MastNodeId,
8    },
9    utils::{DenseIdMap, Idx, IndexVec},
10};
11
12/// Construction surface for dense MAST forests.
13///
14/// The builder may append nodes while a forest is under construction. The value returned by
15/// [`Self::finish`] is a finalized [`MastForest`] in final dense order: external nodes sorted by
16/// digest, then basic blocks in construction order, then internal nodes with children before
17/// parents and construction order as the tie-breaker.
18#[derive(Debug, Default)]
19pub struct DenseMastForestBuilder {
20    nodes: IndexVec<MastNodeId, MastNode>,
21    roots: Vec<MastNodeId>,
22    advice_map: AdviceMap,
23}
24
25impl DenseMastForestBuilder {
26    /// Returns an empty dense MAST forest builder.
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    /// Adds a node builder to this forest and returns its builder-local node ID.
32    pub fn push_node(
33        &mut self,
34        builder: impl Into<MastNodeBuilder>,
35    ) -> Result<MastNodeId, MastForestError> {
36        let node = builder.into().build(self)?;
37        self.nodes.push(node).map_err(|_| MastForestError::TooManyNodes)
38    }
39
40    /// Returns a node by its builder-local node ID.
41    pub fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
42        self.nodes.get(node_id)
43    }
44
45    /// Marks a builder-local node ID as a procedure root.
46    ///
47    /// # Panics
48    ///
49    /// Panics if `root` does not identify a node already added to this builder.
50    pub fn mark_root(&mut self, root: MastNodeId) {
51        assert!(root.to_usize() < self.nodes.len());
52
53        if !self.roots.contains(&root) {
54            self.roots.push(root);
55        }
56    }
57
58    pub(crate) fn merge_advice_map(
59        &mut self,
60        advice_map: &AdviceMap,
61    ) -> Result<(), MastForestError> {
62        self.advice_map
63            .merge(advice_map)
64            .map_err(|((key, _prev), _new)| MastForestError::AdviceMapKeyCollisionOnMerge(key))
65    }
66
67    /// Finalizes this builder into a dense [`MastForest`].
68    pub fn finish(self) -> Result<MastForest, MastForestError> {
69        self.finish_with_id_map().map(|(forest, _remapping)| forest)
70    }
71
72    /// Finalizes this builder and returns the builder-local to final node ID map.
73    pub fn finish_with_id_map(
74        self,
75    ) -> Result<(MastForest, DenseIdMap<MastNodeId, MastNodeId>), MastForestError> {
76        MastForest::from_parts_with_id_map(MastForestParts {
77            nodes: self.nodes,
78            roots: self.roots,
79            advice_map: self.advice_map,
80        })
81    }
82}
83
84impl MastNodeContext for DenseMastForestBuilder {
85    fn node_count(&self) -> usize {
86        self.nodes.len()
87    }
88
89    fn get_node_by_id(&self, node_id: MastNodeId) -> Option<&MastNode> {
90        self.nodes.get(node_id)
91    }
92}