miden_core/mast/
node_builder_utils.rs

1use alloc::vec::Vec;
2
3use crate::{
4    mast::{
5        BasicBlockNodeBuilder, DecoratorId, DynNodeBuilder, ExternalNodeBuilder, MastForest,
6        MastForestContributor, MastForestError, MastNode, MastNodeBuilder, MastNodeId,
7        node::MastNodeExt,
8    },
9    utils::LookupByIdx,
10};
11
12/// Builds a node builder with remapped children and decorators using the provided mappings.
13///
14/// This is a generic helper used by both `MastForestMerger` and `MastForestBuilder` to avoid
15/// code duplication when copying nodes between forests.
16pub fn build_node_with_remapped_ids<NMap, DMap>(
17    node_id: MastNodeId,
18    node: MastNode,
19    source_forest: &MastForest,
20    node_remapping: &NMap,
21    decorator_remapping: &DMap,
22) -> Result<MastNodeBuilder, MastForestError>
23where
24    NMap: LookupByIdx<MastNodeId, MastNodeId>,
25    DMap: LookupByIdx<DecoratorId, DecoratorId>,
26{
27    let map_decorator_id = |decorator_id: DecoratorId| {
28        decorator_remapping
29            .get(decorator_id)
30            .copied()
31            .ok_or(MastForestError::DecoratorIdOverflow(decorator_id, 0))
32    };
33
34    let map_decorators = |decorators: &[DecoratorId]| -> Result<Vec<_>, MastForestError> {
35        decorators.iter().copied().map(map_decorator_id).collect()
36    };
37
38    // Get decorators from source forest and remap them
39    let before_enter_decorators = map_decorators(source_forest.before_enter_decorators(node_id))?;
40    let after_exit_decorators = map_decorators(source_forest.after_exit_decorators(node_id))?;
41
42    // Build node-specific builder with remapped children and decorators
43    let builder = match node {
44        MastNode::Join(join_node) => {
45            let builder = join_node
46                .to_builder(source_forest)
47                .remap_children(node_remapping)
48                .with_before_enter(before_enter_decorators)
49                .with_after_exit(after_exit_decorators);
50            MastNodeBuilder::Join(builder)
51        },
52        MastNode::Split(split_node) => {
53            let builder = split_node
54                .to_builder(source_forest)
55                .remap_children(node_remapping)
56                .with_before_enter(before_enter_decorators)
57                .with_after_exit(after_exit_decorators);
58            MastNodeBuilder::Split(builder)
59        },
60        MastNode::Loop(loop_node) => {
61            let builder = loop_node
62                .to_builder(source_forest)
63                .remap_children(node_remapping)
64                .with_before_enter(before_enter_decorators)
65                .with_after_exit(after_exit_decorators);
66            MastNodeBuilder::Loop(builder)
67        },
68        MastNode::Call(call_node) => {
69            let builder = call_node
70                .to_builder(source_forest)
71                .remap_children(node_remapping)
72                .with_before_enter(before_enter_decorators)
73                .with_after_exit(after_exit_decorators);
74            MastNodeBuilder::Call(builder)
75        },
76        MastNode::Block(basic_block_node) => {
77            // For BasicBlockNode, we need to remap op-indexed decorators as well
78            let builder = BasicBlockNodeBuilder::new(
79                basic_block_node.operations().copied().collect(),
80                basic_block_node
81                    .indexed_decorator_iter(source_forest)
82                    .map(|(idx, decorator_id)| {
83                        let mapped_decorator = map_decorator_id(decorator_id)?;
84                        Ok((idx, mapped_decorator))
85                    })
86                    .collect::<Result<Vec<_>, _>>()?,
87            )
88            .with_before_enter(before_enter_decorators)
89            .with_after_exit(after_exit_decorators);
90            MastNodeBuilder::BasicBlock(builder)
91        },
92        MastNode::Dyn(dyn_node) => {
93            let builder = if dyn_node.is_dyncall() {
94                DynNodeBuilder::new_dyncall()
95            } else {
96                DynNodeBuilder::new_dyn()
97            }
98            .with_before_enter(before_enter_decorators)
99            .with_after_exit(after_exit_decorators);
100            MastNodeBuilder::Dyn(builder)
101        },
102        MastNode::External(external_node) => {
103            let builder = ExternalNodeBuilder::new(external_node.digest())
104                .with_before_enter(before_enter_decorators)
105                .with_after_exit(after_exit_decorators);
106            MastNodeBuilder::External(builder)
107        },
108    };
109
110    Ok(builder)
111}
miden_core/mast/node_builder_utils.rs

miden_core/mast/
node_builder_utils.rs
