miden_core/mast/node/
join_node.rs1use 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#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct JoinNode {
23 children: [MastNodeId; 2],
24 digest: Word,
25}
26
27impl JoinNode {
29 pub const DOMAIN: Felt = Felt::new_unchecked(opcodes::JOIN as u64);
31}
32
33impl JoinNode {
35 pub fn first(&self) -> MastNodeId {
37 self.children[0]
38 }
39
40 pub fn second(&self) -> MastNodeId {
43 self.children[1]
44 }
45}
46
47impl 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
96impl MastNodeExt for JoinNode {
100 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#[derive(Debug)]
154pub struct JoinNodeBuilder {
155 children: [MastNodeId; 2],
156 digest: Option<Word>,
157}
158
159impl JoinNodeBuilder {
160 pub fn new(children: [MastNodeId; 2]) -> Self {
162 Self { children, digest: None }
163 }
164
165 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 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 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}