Skip to main content

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