miden_core/mast/node/
loop_node.rs1use 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#[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
37impl LoopNode {
39 pub const DOMAIN: Felt = Felt::new_unchecked(opcodes::LOOP as u64);
41}
42
43impl LoopNode {
44 pub fn body(&self) -> MastNodeId {
46 self.body
47 }
48}
49
50impl 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
88impl MastNodeExt for LoopNode {
92 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#[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 (any::<MastNodeId>(), any::<[u64; 4]>())
154 .prop_map(|(body, digest_array)| {
155 let digest = Word::from(digest_array.map(Felt::new_unchecked));
157 LoopNode {
159 body,
160 digest,
161 }
162 })
163 .no_shrink() .boxed()
165 }
166
167 type Strategy = proptest::prelude::BoxedStrategy<Self>;
168}
169
170#[derive(Debug)]
173pub struct LoopNodeBuilder {
174 body: MastNodeId,
175 digest: Option<Word>,
176}
177
178impl LoopNodeBuilder {
179 pub fn new(body: MastNodeId) -> Self {
181 Self { body, digest: None }
182 }
183
184 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 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 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}