use alloc::vec::Vec;
use crate::{
mast::{
BasicBlockNodeBuilder, DecoratorId, DynNodeBuilder, ExternalNodeBuilder, MastForest,
MastForestContributor, MastForestError, MastNode, MastNodeBuilder, MastNodeId,
node::MastNodeExt,
},
utils::LookupByIdx,
};
pub fn build_node_with_remapped_ids<NMap, DMap>(
node_id: MastNodeId,
node: MastNode,
source_forest: &MastForest,
node_remapping: &NMap,
decorator_remapping: &DMap,
) -> Result<MastNodeBuilder, MastForestError>
where
NMap: LookupByIdx<MastNodeId, MastNodeId>,
DMap: LookupByIdx<DecoratorId, DecoratorId>,
{
let map_decorator_id = |decorator_id: DecoratorId| {
decorator_remapping
.get(decorator_id)
.copied()
.ok_or(MastForestError::DecoratorIdOverflow(decorator_id, 0))
};
let map_decorators = |decorators: &[DecoratorId]| -> Result<Vec<_>, MastForestError> {
decorators.iter().copied().map(map_decorator_id).collect()
};
let before_enter_decorators = map_decorators(source_forest.before_enter_decorators(node_id))?;
let after_exit_decorators = map_decorators(source_forest.after_exit_decorators(node_id))?;
let builder = match node {
MastNode::Join(join_node) => {
let builder = join_node
.to_builder(source_forest)
.remap_children(node_remapping)
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::Join(builder)
},
MastNode::Split(split_node) => {
let builder = split_node
.to_builder(source_forest)
.remap_children(node_remapping)
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::Split(builder)
},
MastNode::Loop(loop_node) => {
let builder = loop_node
.to_builder(source_forest)
.remap_children(node_remapping)
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::Loop(builder)
},
MastNode::Call(call_node) => {
let builder = call_node
.to_builder(source_forest)
.remap_children(node_remapping)
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::Call(builder)
},
MastNode::Block(basic_block_node) => {
let builder = BasicBlockNodeBuilder::from_op_batches(
basic_block_node.op_batches().to_vec(),
basic_block_node
.indexed_decorator_iter(source_forest)
.map(|(idx, decorator_id)| {
let mapped_decorator = map_decorator_id(decorator_id)?;
Ok((idx, mapped_decorator))
})
.collect::<Result<Vec<_>, _>>()?,
basic_block_node.digest(),
)
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::BasicBlock(builder)
},
MastNode::Dyn(dyn_node) => {
let builder = if dyn_node.is_dyncall() {
DynNodeBuilder::new_dyncall()
} else {
DynNodeBuilder::new_dyn()
}
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::Dyn(builder)
},
MastNode::External(external_node) => {
let builder = ExternalNodeBuilder::new(external_node.digest())
.with_before_enter(before_enter_decorators)
.with_after_exit(after_exit_decorators);
MastNodeBuilder::External(builder)
},
};
Ok(builder)
}