1use alloc::{
16 collections::{BTreeMap, BTreeSet},
17 format,
18 sync::Arc,
19 vec::Vec,
20};
21
22#[cfg(feature = "serde")]
23use serde::{Deserialize, Serialize};
24
25use super::{
26 DataChunk, DeferredError, DeferredState, Digest, Node, NodeType, PrecompileError,
27 PrecompileRegistry, TRUE_DIGEST, Tag,
28};
29use crate::{
30 Felt, ZERO,
31 serde::{
32 BudgetedReader, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
33 SliceReader,
34 },
35};
36
37pub const TRUE_INDEX: u32 = 0;
42
43#[derive(Debug, Clone, PartialEq, Eq)]
52#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
53pub enum WireEntry {
54 Data { tag: Tag, chunks: Vec<DataChunk> },
59 Join { tag: Tag, lhs: u32, rhs: u32 },
61 PairList { tag: Tag, pairs: Vec<(u32, u32)> },
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Default)]
76#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
77pub struct DeferredStateWire {
78 pub entries: Vec<WireEntry>,
79}
80
81impl DeferredStateWire {
82 pub(crate) fn from_state(state: &DeferredState) -> Result<Self, IntegrityError> {
84 let mut build = WireEncoder::default();
85 build.visit_state_digest(state, state.root())?;
86 Ok(Self { entries: build.entries })
87 }
88
89 pub(crate) fn rehydrate(
91 &self,
92 precompiles: Arc<PrecompileRegistry>,
93 max_elements: usize,
94 ) -> Result<DeferredState, IntegrityError> {
95 let (entries, root) = WireDecoder::new(self, precompiles.as_ref())?.decode()?;
96 let mut state = DeferredState::new(Arc::clone(&precompiles), max_elements)?;
97
98 for (digest, node) in entries {
102 let registered = state.register(node)?;
103 if registered != digest {
104 return Err(IntegrityError::InvalidStructure);
105 }
106 }
107
108 state.root = root;
109
110 if state.to_wire()? != *self {
114 return Err(IntegrityError::InvalidStructure);
115 }
116
117 if state.evaluate_digest(root)? != TRUE_DIGEST {
118 return Err(IntegrityError::RootNotTrue);
119 }
120
121 Ok(state)
122 }
123}
124
125#[derive(Debug, thiserror::Error)]
135pub enum IntegrityError {
136 #[error("invalid or non-canonical deferred wire/state structure")]
138 InvalidStructure,
139 #[error("deferred root failed evaluation: {0}")]
141 EvaluationFailed(#[source] PrecompileError),
142 #[error("deferred root evaluated to a non-TRUE canonical form")]
144 RootNotTrue,
145 #[error("deferred insertion requires {num_elements} elements but only {max} remain")]
147 DeferredStateTooLarge { num_elements: usize, max: usize },
148}
149
150impl From<PrecompileError> for IntegrityError {
151 fn from(err: PrecompileError) -> Self {
152 if let PrecompileError::Other(DeferredError::DeferredStateTooLarge { num_elements, max }) =
153 err.root()
154 {
155 Self::DeferredStateTooLarge { num_elements: *num_elements, max: *max }
156 } else {
157 Self::EvaluationFailed(err)
158 }
159 }
160}
161
162struct WireDecoder<'a> {
166 wire: &'a DeferredStateWire,
167 precompiles: &'a PrecompileRegistry,
168 entries: Vec<(Digest, Node)>,
169 index_to_digest: Vec<Digest>,
170 seen_digests: BTreeSet<Digest>,
171}
172
173impl<'a> WireDecoder<'a> {
174 fn new(
175 wire: &'a DeferredStateWire,
176 precompiles: &'a PrecompileRegistry,
177 ) -> Result<Self, IntegrityError> {
178 let total_nodes =
179 1usize.checked_add(wire.entries.len()).ok_or(IntegrityError::InvalidStructure)?;
180
181 let mut index_to_digest = Vec::with_capacity(total_nodes);
182 let mut seen_digests = BTreeSet::new();
183 index_to_digest.push(TRUE_DIGEST);
184 seen_digests.insert(TRUE_DIGEST);
185
186 Ok(Self {
187 wire,
188 precompiles,
189 entries: Vec::with_capacity(wire.entries.len()),
190 index_to_digest,
191 seen_digests,
192 })
193 }
194
195 fn decode(mut self) -> Result<(Vec<(Digest, Node)>, Digest), IntegrityError> {
196 for entry in &self.wire.entries {
197 let node = match entry {
198 WireEntry::Data { tag, chunks } => self.decode_data_entry(*tag, chunks)?,
199 WireEntry::Join { tag, lhs, rhs } => self.decode_join_entry(*tag, *lhs, *rhs)?,
200 WireEntry::PairList { tag, pairs } => self.decode_pair_list_entry(*tag, pairs)?,
201 };
202 self.push_entry(node)?;
203 }
204
205 let root = *self.index_to_digest.last().expect("digest table is seeded with TRUE_DIGEST");
206 Ok((self.entries, root))
207 }
208
209 fn decode_data_entry(&self, tag: Tag, chunks: &[DataChunk]) -> Result<Node, IntegrityError> {
210 let node_type = self
213 .precompiles
214 .decode_node_type(tag)
215 .map_err(|_| IntegrityError::InvalidStructure)?;
216 let NodeType::Data = node_type else {
217 return Err(IntegrityError::InvalidStructure);
218 };
219 let node = if tag == Tag::CHUNKS {
220 Node::chunks(chunks.to_vec()).map_err(|_| IntegrityError::InvalidStructure)?
221 } else {
222 Node::try_data(tag, chunks.to_vec()).map_err(|_| IntegrityError::InvalidStructure)?
223 };
224 node_type.validate_node(&node).map_err(|_| IntegrityError::InvalidStructure)?;
225 Ok(node)
226 }
227
228 fn decode_join_entry(&self, tag: Tag, lhs: u32, rhs: u32) -> Result<Node, IntegrityError> {
229 let lhs = self.resolve_index(lhs)?;
230 let rhs = self.resolve_index(rhs)?;
231 let node = if tag == Tag::AND {
232 Node::and(lhs, rhs)
233 } else {
234 Node::join(tag, lhs, rhs).map_err(|_| IntegrityError::InvalidStructure)?
235 };
236 let node_type = self
237 .precompiles
238 .decode_node_type(node.tag())
239 .map_err(|_| IntegrityError::InvalidStructure)?;
240 node_type.validate_node(&node).map_err(|_| IntegrityError::InvalidStructure)?;
241 match node_type {
242 NodeType::Join => Ok(node),
243 NodeType::True | NodeType::Data | NodeType::PairList => {
244 Err(IntegrityError::InvalidStructure)
245 },
246 }
247 }
248
249 fn decode_pair_list_entry(
250 &self,
251 tag: Tag,
252 pairs: &[(u32, u32)],
253 ) -> Result<Node, IntegrityError> {
254 let node_type = self
255 .precompiles
256 .decode_node_type(tag)
257 .map_err(|_| IntegrityError::InvalidStructure)?;
258 let NodeType::PairList = node_type else {
259 return Err(IntegrityError::InvalidStructure);
260 };
261
262 let pairs = pairs
263 .iter()
264 .map(|(lhs, rhs)| Ok((self.resolve_index(*lhs)?, self.resolve_index(*rhs)?)))
265 .collect::<Result<Vec<_>, IntegrityError>>()?;
266 let node = Node::try_pair_list(tag, pairs).map_err(|_| IntegrityError::InvalidStructure)?;
267 node_type.validate_node(&node).map_err(|_| IntegrityError::InvalidStructure)?;
268 Ok(node)
269 }
270
271 fn resolve_index(&self, idx: u32) -> Result<Digest, IntegrityError> {
272 self.index_to_digest
273 .get(idx as usize)
274 .copied()
275 .ok_or(IntegrityError::InvalidStructure)
276 }
277
278 fn push_entry(&mut self, node: Node) -> Result<(), IntegrityError> {
279 if node.is_true() {
280 return Err(IntegrityError::InvalidStructure);
281 }
282
283 let digest = node.digest();
284 if !self.seen_digests.insert(digest) {
285 return Err(IntegrityError::InvalidStructure);
286 }
287
288 let index = self.index_to_digest.len();
289 if index > u32::MAX as usize {
290 return Err(IntegrityError::InvalidStructure);
291 }
292
293 self.entries.push((digest, node));
294 self.index_to_digest.push(digest);
295 Ok(())
296 }
297}
298
299#[derive(Default)]
304struct WireEncoder {
305 seen: BTreeSet<Digest>,
306 by_digest: BTreeMap<Digest, u32>,
307 entries: Vec<WireEntry>,
308}
309
310impl WireEncoder {
311 fn visit_state_digest(
312 &mut self,
313 state: &DeferredState,
314 digest: Digest,
315 ) -> Result<(), IntegrityError> {
316 let mut pending = Vec::new();
317 pending.push(WireEncodeStep::Visit(digest));
318
319 while let Some(step) = pending.pop() {
320 match step {
321 WireEncodeStep::Visit(digest) => {
322 self.schedule_digest(state, digest, &mut pending)?
323 },
324 WireEncodeStep::Emit(digest) => {
325 let entry = self.entry_for_digest(state, digest)?;
326 self.push_entry(digest, entry)?;
327 },
328 }
329 }
330
331 Ok(())
332 }
333
334 fn schedule_digest(
335 &mut self,
336 state: &DeferredState,
337 digest: Digest,
338 pending: &mut Vec<WireEncodeStep>,
339 ) -> Result<(), IntegrityError> {
340 if digest == TRUE_DIGEST || !self.seen.insert(digest) {
341 return Ok(());
342 }
343
344 let node = self.validated_node(state, digest)?;
345 pending.push(WireEncodeStep::Emit(digest));
346
347 match self.node_type(state, node)? {
348 NodeType::Data => {},
349 NodeType::Join => {
350 let (lhs, rhs) =
351 node.payload().as_join().map_err(|_| IntegrityError::InvalidStructure)?;
352 pending.push(WireEncodeStep::Visit(rhs));
353 pending.push(WireEncodeStep::Visit(lhs));
354 },
355 NodeType::PairList => {
356 let pairs =
357 node.payload().as_pair_list().map_err(|_| IntegrityError::InvalidStructure)?;
358 for (lhs, rhs) in pairs.iter().rev() {
359 pending.push(WireEncodeStep::Visit(*rhs));
360 pending.push(WireEncodeStep::Visit(*lhs));
361 }
362 },
363 NodeType::True => return Err(IntegrityError::InvalidStructure),
364 };
365
366 Ok(())
367 }
368
369 fn entry_for_digest(
370 &self,
371 state: &DeferredState,
372 digest: Digest,
373 ) -> Result<WireEntry, IntegrityError> {
374 let node = self.validated_node(state, digest)?;
375
376 Ok(match self.node_type(state, node)? {
377 NodeType::Data => WireEntry::Data {
378 tag: node.tag(),
379 chunks: node
380 .payload()
381 .as_data()
382 .map_err(|_| IntegrityError::InvalidStructure)?
383 .to_vec(),
384 },
385 NodeType::Join => {
386 let (lhs, rhs) =
387 node.payload().as_join().map_err(|_| IntegrityError::InvalidStructure)?;
388 let lhs = self.index_for(lhs)?;
389 let rhs = self.index_for(rhs)?;
390 WireEntry::Join { tag: node.tag(), lhs, rhs }
391 },
392 NodeType::PairList => {
393 let pairs =
394 node.payload().as_pair_list().map_err(|_| IntegrityError::InvalidStructure)?;
395 let pairs = pairs
396 .iter()
397 .map(|(lhs, rhs)| Ok((self.index_for(*lhs)?, self.index_for(*rhs)?)))
398 .collect::<Result<Vec<_>, IntegrityError>>()?;
399 WireEntry::PairList { tag: node.tag(), pairs }
400 },
401 NodeType::True => return Err(IntegrityError::InvalidStructure),
402 })
403 }
404
405 fn validated_node<'a>(
406 &self,
407 state: &'a DeferredState,
408 digest: Digest,
409 ) -> Result<&'a Node, IntegrityError> {
410 let node = state.get_node(&digest).ok_or(IntegrityError::InvalidStructure)?;
411 self.node_type(state, node)?
412 .validate_node(node)
413 .map_err(|_| IntegrityError::InvalidStructure)?;
414 Ok(node)
415 }
416
417 fn node_type(&self, state: &DeferredState, node: &Node) -> Result<NodeType, IntegrityError> {
418 state
419 .registry()
420 .decode_node_type(node.tag())
421 .map_err(|_| IntegrityError::InvalidStructure)
422 }
423
424 fn index_for(&self, digest: Digest) -> Result<u32, IntegrityError> {
425 if digest == TRUE_DIGEST {
426 return Ok(TRUE_INDEX);
427 }
428 self.by_digest.get(&digest).copied().ok_or(IntegrityError::InvalidStructure)
429 }
430
431 fn push_entry(&mut self, digest: Digest, entry: WireEntry) -> Result<(), IntegrityError> {
432 let next_index =
433 self.entries.len().checked_add(1).ok_or(IntegrityError::InvalidStructure)?;
434 let next_index = u32::try_from(next_index).map_err(|_| IntegrityError::InvalidStructure)?;
435 self.entries.push(entry);
436 self.by_digest.insert(digest, next_index);
437 Ok(())
438 }
439}
440
441enum WireEncodeStep {
442 Visit(Digest),
443 Emit(Digest),
444}
445
446impl Serializable for WireEntry {
450 fn write_into<W: ByteWriter>(&self, target: &mut W) {
451 match self {
452 Self::Data { tag, chunks } => {
453 target.write_u8(0);
454 tag.write_into(target);
455 target.write_usize(chunks.len());
456 for chunk in chunks {
457 for felt in chunk {
458 felt.write_into(target);
459 }
460 }
461 },
462 Self::Join { tag, lhs, rhs } => {
463 target.write_u8(1);
464 tag.write_into(target);
465 target.write_u32(*lhs);
466 target.write_u32(*rhs);
467 },
468 Self::PairList { tag, pairs } => {
469 target.write_u8(2);
470 tag.write_into(target);
471 target.write_usize(pairs.len());
472 for (lhs, rhs) in pairs {
473 target.write_u32(*lhs);
474 target.write_u32(*rhs);
475 }
476 },
477 }
478 }
479}
480
481impl Deserializable for WireEntry {
482 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
483 let discriminant = source.read_u8()?;
484 match discriminant {
485 0 => {
486 let tag = Tag::read_from(source)?;
487 let chunk_count = source.read_usize()?;
488 let chunks = source
489 .read_many_iter::<WireDataChunk>(chunk_count)?
490 .map(|chunk| chunk.map(|chunk| chunk.0))
491 .collect::<Result<_, _>>()?;
492 Ok(Self::Data { tag, chunks })
493 },
494 1 => {
495 let tag = Tag::read_from(source)?;
496 let lhs = source.read_u32()?;
497 let rhs = source.read_u32()?;
498 Ok(Self::Join { tag, lhs, rhs })
499 },
500 2 => {
501 let tag = Tag::read_from(source)?;
502 let pair_count = source.read_usize()?;
503 let pairs = source
504 .read_many_iter::<WirePair>(pair_count)?
505 .map(|pair| pair.map(|pair| pair.0))
506 .collect::<Result<_, _>>()?;
507 Ok(Self::PairList { tag, pairs })
508 },
509 other => Err(DeserializationError::InvalidValue(format!(
510 "invalid deferred wire entry discriminant: {other}"
511 ))),
512 }
513 }
514
515 fn min_serialized_size() -> usize {
516 1
517 }
518}
519
520struct WirePair((u32, u32));
521
522impl Deserializable for WirePair {
523 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
524 Ok(Self((source.read_u32()?, source.read_u32()?)))
525 }
526
527 fn min_serialized_size() -> usize {
528 u32::min_serialized_size() * 2
529 }
530}
531
532struct WireDataChunk(DataChunk);
533
534impl Deserializable for WireDataChunk {
535 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
536 let mut chunk = [ZERO; Node::DATA_CHUNK_FELT_LEN];
537 for felt in &mut chunk {
538 *felt = Felt::read_from(source)?;
539 }
540 Ok(Self(chunk))
541 }
542
543 fn min_serialized_size() -> usize {
544 Node::DATA_CHUNK_FELT_LEN * Felt::min_serialized_size()
545 }
546}
547
548impl Serializable for DeferredStateWire {
549 fn write_into<W: ByteWriter>(&self, target: &mut W) {
550 target.write_usize(self.entries.len());
551 for entry in &self.entries {
552 entry.write_into(target);
553 }
554 }
555}
556
557impl Deserializable for DeferredStateWire {
558 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
559 let entry_count = source.read_usize()?;
560 let entries = source.read_many_iter::<WireEntry>(entry_count)?.collect::<Result<_, _>>()?;
561 Ok(Self { entries })
562 }
563
564 fn read_from_bytes(bytes: &[u8]) -> Result<Self, DeserializationError> {
565 let mut reader = BudgetedReader::new(SliceReader::new(bytes), bytes.len());
566 Self::read_from(&mut reader)
567 }
568}
569
570#[cfg(test)]
571mod tests {
572 use alloc::vec::Vec;
573
574 use super::*;
575 use crate::{
576 Felt,
577 deferred::{DeferredContext, Payload, Precompile, precompile_id},
578 serde::{ByteWriter, Serializable},
579 };
580
581 #[derive(Debug, Clone, Copy)]
582 struct PairListFixture;
583
584 impl PairListFixture {
585 const NAME: &'static str = "wire-pair-list-fixture";
586
587 fn tag() -> Tag {
588 Tag::precompile(precompile_id(Self::NAME), [ZERO; 3])
589 .expect("fixture id is precompile-owned")
590 }
591 }
592
593 impl Precompile for PairListFixture {
594 fn name(&self) -> &'static str {
595 Self::NAME
596 }
597
598 fn id(&self) -> Felt {
599 precompile_id(Self::NAME)
600 }
601
602 fn decode(&self, args: [Felt; 3]) -> Option<NodeType> {
603 (args == [ZERO; 3]).then_some(NodeType::PairList)
604 }
605
606 fn evaluate(
607 &self,
608 _args: [Felt; 3],
609 _payload: &Payload,
610 _context: &mut DeferredContext<'_>,
611 ) -> Result<Node, PrecompileError> {
612 Ok(Node::TRUE)
613 }
614 }
615
616 fn felts(seed: u64) -> [Felt; 8] {
617 core::array::from_fn(|i| Felt::new_unchecked(seed + i as u64))
618 }
619
620 fn tag(seed: u64) -> Tag {
621 Tag::from_word(felts(seed)[..4].try_into().unwrap())
622 }
623
624 fn wire(entries: Vec<WireEntry>) -> DeferredStateWire {
625 DeferredStateWire { entries }
626 }
627
628 fn assert_wire_round_trips(wire: DeferredStateWire) {
629 let decoded = DeferredStateWire::read_from_bytes(&wire.to_bytes()).unwrap();
630 assert_eq!(decoded, wire);
631 }
632
633 #[test]
634 fn wire_decoder_accepts_exact_framework_chunks_data() {
635 let registry = PrecompileRegistry::new();
636 let chunks = alloc::vec![felts(10), felts(20)];
637 let wire = wire(alloc::vec![WireEntry::Data { tag: Tag::CHUNKS, chunks: chunks.clone() }]);
638 let node = Node::chunks(chunks).unwrap();
639
640 let (entries, root) = WireDecoder::new(&wire, ®istry).unwrap().decode().unwrap();
641
642 assert_eq!(entries, alloc::vec![(node.digest(), node.clone())]);
643 assert_eq!(root, node.digest());
644 }
645
646 #[test]
647 fn rehydration_rejects_empty_data_and_pair_list_entries() {
648 let empty_data =
649 wire(alloc::vec![WireEntry::Data { tag: Tag::CHUNKS, chunks: Vec::new() }]);
650 assert!(matches!(
651 DeferredState::from_wire(Arc::new(PrecompileRegistry::new()), &empty_data, usize::MAX,),
652 Err(IntegrityError::InvalidStructure)
653 ));
654
655 let empty_pairs = wire(alloc::vec![WireEntry::PairList {
656 tag: PairListFixture::tag(),
657 pairs: Vec::new(),
658 }]);
659 assert!(matches!(
660 DeferredState::from_wire(
661 Arc::new(PrecompileRegistry::new().with_precompile(PairListFixture)),
662 &empty_pairs,
663 usize::MAX,
664 ),
665 Err(IntegrityError::InvalidStructure)
666 ));
667 }
668
669 #[test]
670 fn wire_decoder_rejects_malformed_framework_chunks_data() {
671 let registry = PrecompileRegistry::new();
672 let malformed = Tag::from_word([Tag::CHUNKS.id(), Felt::new_unchecked(1), ZERO, ZERO]);
673 let wire = wire(alloc::vec![WireEntry::Data {
674 tag: malformed,
675 chunks: alloc::vec![felts(10)],
676 }]);
677
678 assert!(matches!(
679 WireDecoder::new(&wire, ®istry).unwrap().decode(),
680 Err(IntegrityError::InvalidStructure)
681 ));
682 }
683
684 #[test]
686 fn wire_serialize_round_trip_all_entries() {
687 assert_wire_round_trips(wire(alloc::vec![
688 WireEntry::Data {
689 tag: tag(1),
690 chunks: alloc::vec![felts(10)]
691 },
692 WireEntry::Data {
693 tag: tag(2),
694 chunks: alloc::vec![felts(20), felts(30)],
695 },
696 WireEntry::Join { tag: tag(3), lhs: 1, rhs: TRUE_INDEX },
697 WireEntry::PairList {
698 tag: tag(5),
699 pairs: alloc::vec![(1, 2), (TRUE_INDEX, 3)],
700 },
701 ]));
702 assert_wire_round_trips(DeferredStateWire::default());
703 }
704
705 #[test]
706 fn wire_encoder_handles_deep_roots_iteratively() {
707 let mut state = DeferredState::default();
708 for _ in 0..4_096 {
709 state.log_statement(TRUE_DIGEST).unwrap();
710 }
711
712 let root = state.root();
713 let wire = state.to_wire().unwrap();
714
715 assert_eq!(wire.entries.len(), 4_096);
716 assert_eq!(
717 wire.entries.last(),
718 Some(&WireEntry::Join {
719 tag: Tag::AND,
720 lhs: 4_095,
721 rhs: TRUE_INDEX,
722 })
723 );
724 assert_eq!(
725 DeferredState::from_wire(Arc::new(PrecompileRegistry::new()), &wire, usize::MAX)
726 .unwrap()
727 .root(),
728 root
729 );
730 }
731
732 fn encoded_entry_count(entry_count: usize) -> Vec<u8> {
733 let mut bytes = Vec::new();
734 bytes.write_usize(entry_count);
735 bytes
736 }
737
738 #[test]
739 fn wire_rejects_over_budget_entry_count() {
740 assert!(DeferredStateWire::read_from_bytes(&encoded_entry_count(usize::MAX)).is_err());
741 }
742
743 #[test]
744 fn wire_rejects_over_budget_data_chunk_count() {
745 let mut bytes = Vec::new();
746 bytes.write_usize(1);
747 bytes.write_u8(0); tag(1).write_into(&mut bytes);
749 bytes.write_usize(usize::MAX);
750
751 assert!(DeferredStateWire::read_from_bytes(&bytes).is_err());
752 }
753
754 #[test]
755 fn wire_rejects_over_budget_pair_count() {
756 let mut bytes = Vec::new();
757 bytes.write_usize(1);
758 bytes.write_u8(2); tag(1).write_into(&mut bytes);
760 bytes.write_usize(usize::MAX);
761
762 assert!(DeferredStateWire::read_from_bytes(&bytes).is_err());
763 }
764}