miden_core/mast/node/
split_node.rs

1use alloc::{boxed::Box, vec::Vec};
2use core::fmt;
3
4use super::{
5    MastForestContributor, MastNodeContext, MastNodeExt, fingerprint_with_child_fingerprints,
6};
7use crate::{
8    Felt, Word,
9    chiplets::hasher,
10    mast::{MastForest, MastForestError, MastNodeId},
11    operations::opcodes,
12    prettier::PrettyPrint,
13    utils::LookupByIdx,
14};
15
16// SPLIT NODE
17// ================================================================================================
18
19/// A Split node defines conditional execution. When the VM encounters a Split node it executes
20/// either the `on_true` child or `on_false` child.
21///
22/// Which child is executed is determined based on the top of the stack. If the value is `1`, then
23/// the `on_true` child is executed. If the value is `0`, then the `on_false` child is executed. If
24/// the value is neither `0` nor `1`, the execution fails.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct SplitNode {
27    branches: [MastNodeId; 2],
28    digest: Word,
29}
30
31/// Constants
32impl SplitNode {
33    /// The domain of the split node (used for control block hashing).
34    pub const DOMAIN: Felt = Felt::new_unchecked(opcodes::SPLIT as u64);
35}
36
37/// Public accessors
38impl SplitNode {
39    /// Returns the ID of the node which is to be executed if the top of the stack is `1`.
40    pub fn on_true(&self) -> MastNodeId {
41        self.branches[0]
42    }
43
44    /// Returns the ID of the node which is to be executed if the top of the stack is `0`.
45    pub fn on_false(&self) -> MastNodeId {
46        self.branches[1]
47    }
48}
49
50// PRETTY PRINTING
51// ================================================================================================
52
53impl SplitNode {
54    pub(super) fn to_display<'a>(&'a self, mast_forest: &'a MastForest) -> impl fmt::Display + 'a {
55        SplitNodePrettyPrint { split_node: self, mast_forest }
56    }
57
58    pub(super) fn to_pretty_print<'a>(
59        &'a self,
60        mast_forest: &'a MastForest,
61    ) -> impl PrettyPrint + 'a {
62        SplitNodePrettyPrint { split_node: self, mast_forest }
63    }
64}
65
66struct SplitNodePrettyPrint<'a> {
67    split_node: &'a SplitNode,
68    mast_forest: &'a MastForest,
69}
70
71impl PrettyPrint for SplitNodePrettyPrint<'_> {
72    #[rustfmt::skip]
73    fn render(&self) -> crate::prettier::Document {
74        use crate::prettier::*;
75
76        let true_branch = self.mast_forest[self.split_node.on_true()].to_pretty_print(self.mast_forest);
77        let false_branch = self.mast_forest[self.split_node.on_false()].to_pretty_print(self.mast_forest);
78
79        let mut doc = Document::Empty;
80        doc += indent(4, const_text("if.true") + nl() + true_branch.render()) + nl();
81        doc += indent(4, const_text("else") + nl() + false_branch.render());
82        doc += nl() + const_text("end");
83        doc
84    }
85}
86
87impl fmt::Display for SplitNodePrettyPrint<'_> {
88    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89        use crate::prettier::PrettyPrint;
90        self.pretty_print(f)
91    }
92}
93
94// MAST NODE TRAIT IMPLEMENTATION
95// ================================================================================================
96
97impl MastNodeExt for SplitNode {
98    /// Returns a commitment to this Split node.
99    ///
100    /// The commitment is computed as a hash of the `on_true` and `on_false` child nodes in the
101    /// domain defined by [Self::DOMAIN] - i..e,:
102    /// ```
103    /// # use miden_core::mast::SplitNode;
104    /// # use miden_crypto::{Word, hash::poseidon2::Poseidon2 as Hasher};
105    /// # let on_true_digest = Word::default();
106    /// # let on_false_digest = Word::default();
107    /// Hasher::merge_in_domain(&[on_true_digest, on_false_digest], SplitNode::DOMAIN);
108    /// ```
109    fn digest(&self) -> Word {
110        self.digest
111    }
112
113    fn to_display<'a>(&'a self, mast_forest: &'a MastForest) -> Box<dyn fmt::Display + 'a> {
114        Box::new(SplitNode::to_display(self, mast_forest))
115    }
116
117    fn to_pretty_print<'a>(&'a self, mast_forest: &'a MastForest) -> Box<dyn PrettyPrint + 'a> {
118        Box::new(SplitNode::to_pretty_print(self, mast_forest))
119    }
120
121    fn has_children(&self) -> bool {
122        true
123    }
124
125    fn append_children_to(&self, target: &mut Vec<MastNodeId>) {
126        target.push(self.on_true());
127        target.push(self.on_false());
128    }
129
130    fn for_each_child<F>(&self, mut f: F)
131    where
132        F: FnMut(MastNodeId),
133    {
134        f(self.on_true());
135        f(self.on_false());
136    }
137
138    fn domain(&self) -> Felt {
139        Self::DOMAIN
140    }
141
142    type Builder = SplitNodeBuilder;
143
144    fn to_builder(self, _forest: &MastForest) -> Self::Builder {
145        SplitNodeBuilder::new(self.branches).with_digest(self.digest)
146    }
147}
148
149// ------------------------------------------------------------------------------------------------
150/// Builder for creating [`SplitNode`] instances.
151#[derive(Debug)]
152pub struct SplitNodeBuilder {
153    branches: [MastNodeId; 2],
154    digest: Option<Word>,
155}
156
157impl SplitNodeBuilder {
158    /// Creates a new builder for a SplitNode with the specified branches.
159    pub fn new(branches: [MastNodeId; 2]) -> Self {
160        Self { branches, digest: None }
161    }
162
163    /// Builds the SplitNode.
164    pub fn build(self, context: &impl MastNodeContext) -> Result<SplitNode, MastForestError> {
165        let true_branch = context.get_node_by_id(self.branches[0]).ok_or_else(|| {
166            MastForestError::NodeIdOverflow(self.branches[0], context.node_count())
167        })?;
168        let false_branch = context.get_node_by_id(self.branches[1]).ok_or_else(|| {
169            MastForestError::NodeIdOverflow(self.branches[1], context.node_count())
170        })?;
171
172        // Use the forced digest if provided, otherwise compute the digest
173        let digest = if let Some(forced_digest) = self.digest {
174            forced_digest
175        } else {
176            let true_branch_hash = true_branch.digest();
177            let false_branch_hash = false_branch.digest();
178
179            hasher::merge_in_domain(&[true_branch_hash, false_branch_hash], SplitNode::DOMAIN)
180        };
181
182        Ok(SplitNode { branches: self.branches, digest })
183    }
184
185    pub(in crate::mast) fn build_linked(self) -> Result<SplitNode, MastForestError> {
186        Ok(SplitNode {
187            branches: self.branches,
188            digest: self.digest.ok_or(MastForestError::DigestRequiredForDeserialization)?,
189        })
190    }
191}
192
193#[cfg(any(test, feature = "arbitrary"))]
194impl SplitNodeBuilder {
195    /// Adds this builder to a mutable forest for test and arbitrary data construction.
196    pub fn add_to_forest(self, forest: &mut MastForest) -> Result<MastNodeId, MastForestError> {
197        let node = self.build(forest)?;
198        forest.nodes.push(node.into()).map_err(|_| MastForestError::TooManyNodes)
199    }
200}
201
202impl MastForestContributor for SplitNodeBuilder {
203    fn fingerprint_for_node(
204        &self,
205        context: &impl MastNodeContext,
206        hash_by_node_id: &impl LookupByIdx<MastNodeId, Word>,
207    ) -> Result<Word, MastForestError> {
208        let node_digest = if let Some(forced_digest) = self.digest {
209            forced_digest
210        } else {
211            let if_branch_hash = context
212                .get_node_by_id(self.branches[0])
213                .ok_or_else(|| {
214                    MastForestError::NodeIdOverflow(self.branches[0], context.node_count())
215                })?
216                .digest();
217            let else_branch_hash = context
218                .get_node_by_id(self.branches[1])
219                .ok_or_else(|| {
220                    MastForestError::NodeIdOverflow(self.branches[1], context.node_count())
221                })?
222                .digest();
223
224            hasher::merge_in_domain(&[if_branch_hash, else_branch_hash], SplitNode::DOMAIN)
225        };
226
227        fingerprint_with_child_fingerprints(node_digest, &self.branches, context, hash_by_node_id)
228    }
229
230    fn remap_children(self, remapping: &impl LookupByIdx<MastNodeId, MastNodeId>) -> Self {
231        SplitNodeBuilder {
232            branches: [
233                *remapping.get(self.branches[0]).unwrap_or(&self.branches[0]),
234                *remapping.get(self.branches[1]).unwrap_or(&self.branches[1]),
235            ],
236            digest: self.digest,
237        }
238    }
239
240    fn with_digest(mut self, digest: Word) -> Self {
241        self.digest = Some(digest);
242        self
243    }
244}
245
246#[cfg(any(test, feature = "arbitrary"))]
247impl proptest::prelude::Arbitrary for SplitNodeBuilder {
248    type Parameters = ();
249    type Strategy = proptest::strategy::BoxedStrategy<Self>;
250
251    fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy {
252        use proptest::prelude::*;
253
254        any::<[MastNodeId; 2]>().prop_map(Self::new).boxed()
255    }
256}
miden_core/mast/node/split_node.rs

miden_core/mast/node/
split_node.rs
