Skip to main content

miden_core/mast/node/
loop_node.rs

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