1use prism::operation::TermValue;
44use prism::pipeline::{
45 ConstrainedTypeShape, ConstraintRef, IntoBindingValue, PartitionProductFields,
46};
47
48#[derive(Clone, Copy, Debug)]
54pub struct OnnxCarrier<'a>(&'a [u8]);
55
56impl<'a> OnnxCarrier<'a> {
57 #[must_use]
59 pub fn new(skeleton: &'a [u8]) -> Self {
60 Self(skeleton)
61 }
62
63 #[must_use]
65 pub fn canonical_bytes(&self) -> &'a [u8] {
66 self.0
67 }
68}
69
70impl ConstrainedTypeShape for OnnxCarrier<'_> {
71 const IRI: &'static str = "https://uor.foundation/addr/OnnxValue";
72 const SITE_COUNT: usize = 1;
73 const CONSTRAINTS: &'static [ConstraintRef] = &[];
74 const CYCLE_SIZE: u64 = u64::MAX;
75}
76
77impl prism::uor_foundation::pipeline::__sdk_seal::Sealed for OnnxCarrier<'_> {}
78
79impl<'a> IntoBindingValue<'a> for OnnxCarrier<'a> {
80 fn as_binding_value<const INLINE_BYTES: usize>(&self) -> TermValue<'a, INLINE_BYTES> {
81 TermValue::borrowed(self.0)
82 }
83}
84
85impl PartitionProductFields for OnnxCarrier<'_> {
86 const FIELDS: &'static [(u32, u32)] = &[];
87 const FIELD_NAMES: &'static [&'static str] = &[];
88}
89
90#[cfg(feature = "alloc")]
95pub use alloc_impl::{canonicalize, OnnxValue};
96
97#[cfg(feature = "alloc")]
98mod alloc_impl {
99 use alloc::vec::Vec;
100
101 use prism::crypto::Sha256Hasher;
102 use prism::pipeline::{ShapeViolation, ViolationKind};
103 use prism::vocabulary::Hasher;
104
105 use crate::onnx::dtype::OnnxDataType;
106 use crate::onnx::protobuf::{read_varint, FieldValue, MessageReader};
107 use crate::onnx::shapes::bounds::{
108 ONNX_IR_VERSION_MAX, ONNX_OPSET_VERSION_MIN, ONNX_SUBGRAPH_DEPTH_MAX,
109 };
110
111 macro_rules! violation {
114 ($name:ident, $constraint:literal, $kind:expr) => {
115 const $name: ShapeViolation = ShapeViolation {
116 shape_iri: "https://uor.foundation/addr/OnnxValue",
117 constraint_iri: concat!("https://uor.foundation/addr/OnnxValue/", $constraint),
118 property_iri: concat!("https://uor.foundation/addr/OnnxValue/", $constraint),
119 expected_range: "http://www.w3.org/2001/XMLSchema#nonNegativeInteger",
120 min_count: 0,
121 max_count: 1,
122 kind: $kind,
123 };
124 };
125 }
126
127 violation!(PROTOBUF_FAILURE, "validProtobuf", ViolationKind::ValueCheck);
128 violation!(
129 UNSUPPORTED_IR,
130 "supportedIrVersion",
131 ViolationKind::ValueCheck
132 );
133 violation!(OPSET_TOO_OLD, "opsetVersionMin", ViolationKind::ValueCheck);
134 violation!(MISSING_GRAPH, "graphPresent", ViolationKind::ValueCheck);
135 violation!(
136 SUBGRAPH_DEPTH,
137 "subgraphDepthBound",
138 ViolationKind::CardinalityViolation
139 );
140 violation!(GRAPH_CYCLE, "acyclicGraph", ViolationKind::ValueCheck);
141 violation!(
142 UNKNOWN_DTYPE,
143 "knownTensorDataType",
144 ViolationKind::ValueCheck
145 );
146
147 fn from_wire(_e: crate::onnx::protobuf::WireError) -> ShapeViolation {
148 PROTOBUF_FAILURE
149 }
150
151 #[inline]
152 fn sha256(bytes: &[u8]) -> [u8; 32] {
153 Sha256Hasher::initial().fold_bytes(bytes).finalize()
154 }
155
156 const CANON_PROTO_DEPTH_MAX: usize = 32;
158
159 fn canonical_proto_digest(body: &[u8], depth: usize) -> Result<[u8; 32], ShapeViolation> {
175 #[derive(Clone, Copy)]
176 struct F {
177 number: u64,
178 wt: u8,
179 off: usize,
180 len: usize,
181 val: u64,
182 }
183 let mut fs: Vec<F> = Vec::new();
184 let mut r = MessageReader::new(body);
185 while let Some(f) = r.next_field().map_err(from_wire)? {
186 fs.push(match f.value {
187 FieldValue::Varint(v) => F {
188 number: f.number,
189 wt: 0,
190 off: 0,
191 len: 0,
192 val: v,
193 },
194 FieldValue::Fixed64(v) => F {
195 number: f.number,
196 wt: 1,
197 off: 0,
198 len: 0,
199 val: v,
200 },
201 FieldValue::Fixed32(v) => F {
202 number: f.number,
203 wt: 5,
204 off: 0,
205 len: 0,
206 val: u64::from(v),
207 },
208 FieldValue::Bytes(b) => F {
209 number: f.number,
210 wt: 2,
211 off: b.as_ptr() as usize - body.as_ptr() as usize,
212 len: b.len(),
213 val: 0,
214 },
215 });
216 }
217 fs.sort_by_key(|f| f.number);
219
220 let mut h = Sha256Hasher::initial();
221 for f in fs.iter() {
222 fold(&mut h, &f.number.to_le_bytes());
223 fold(&mut h, &[f.wt]);
224 match f.wt {
225 0 | 1 => fold(&mut h, &f.val.to_le_bytes()),
226 5 => fold(&mut h, &(f.val as u32).to_le_bytes()),
227 _ => {
228 let payload = &body[f.off..f.off + f.len];
229 let sub = if depth < CANON_PROTO_DEPTH_MAX && !payload.is_empty() {
230 canonical_proto_digest(payload, depth + 1)
231 .unwrap_or_else(|_| sha256(payload))
232 } else {
233 sha256(payload)
234 };
235 fold(&mut h, &sub);
236 }
237 }
238 }
239 Ok(h.finalize())
240 }
241
242 #[inline]
246 fn fold(h: &mut Sha256Hasher, bytes: &[u8]) {
247 let cur = core::mem::replace(h, Sha256Hasher::initial());
248 *h = cur.fold_bytes(bytes);
249 }
250
251 fn first_field(body: &[u8], field_no: u64) -> Result<Option<FieldValue<'_>>, ShapeViolation> {
255 let mut r = MessageReader::new(body);
256 while let Some(f) = r.next_field().map_err(from_wire)? {
257 if f.number == field_no {
258 return Ok(Some(f.value));
259 }
260 }
261 Ok(None)
262 }
263
264 fn first_varint(body: &[u8], field_no: u64) -> Result<Option<u64>, ShapeViolation> {
265 Ok(match first_field(body, field_no)? {
266 Some(FieldValue::Varint(v)) => Some(v),
267 _ => None,
268 })
269 }
270
271 fn first_bytes(body: &[u8], field_no: u64) -> Result<&[u8], ShapeViolation> {
272 Ok(match first_field(body, field_no)? {
273 Some(FieldValue::Bytes(b)) => b,
274 _ => &[],
275 })
276 }
277
278 fn for_each_field(
281 body: &[u8],
282 field_no: u64,
283 mut f: impl FnMut(FieldValue<'_>) -> Result<(), ShapeViolation>,
284 ) -> Result<(), ShapeViolation> {
285 let mut r = MessageReader::new(body);
286 while let Some(field) = r.next_field().map_err(from_wire)? {
287 if field.number == field_no {
288 f(field.value)?;
289 }
290 }
291 Ok(())
292 }
293
294 fn count_field(body: &[u8], field_no: u64) -> Result<usize, ShapeViolation> {
295 let mut n = 0;
296 for_each_field(body, field_no, |_| {
297 n += 1;
298 Ok(())
299 })?;
300 Ok(n)
301 }
302
303 #[derive(Clone, Copy)]
305 struct Span {
306 off: usize,
307 len: usize,
308 }
309
310 fn collect_spans(body: &[u8], field_no: u64) -> Result<Vec<Span>, ShapeViolation> {
313 let mut spans: Vec<Span> = Vec::new();
314 let mut r = MessageReader::new(body);
315 while let Some(f) = r.next_field().map_err(from_wire)? {
316 if f.number == field_no {
317 if let FieldValue::Bytes(b) = f.value {
318 spans.push(Span {
319 off: b.as_ptr() as usize - body.as_ptr() as usize,
320 len: b.len(),
321 });
322 }
323 }
324 }
325 Ok(spans)
326 }
327
328 #[derive(Clone, PartialEq, Eq)]
332 pub struct OnnxValue {
333 bytes: Vec<u8>,
334 }
335
336 impl core::fmt::Debug for OnnxValue {
337 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
338 f.debug_struct("OnnxValue")
339 .field("canonical_len", &self.bytes.len())
340 .finish_non_exhaustive()
341 }
342 }
343
344 impl OnnxValue {
345 #[must_use]
347 pub fn canonical_bytes(&self) -> &[u8] {
348 &self.bytes
349 }
350
351 pub fn parse(raw: &[u8]) -> Result<Self, ShapeViolation> {
361 let mut out: Vec<u8> = Vec::new();
362
363 let ir_version = first_varint(raw, 1)?.ok_or(UNSUPPORTED_IR)? as i64;
369 if !(1..=ONNX_IR_VERSION_MAX).contains(&ir_version) {
370 return Err(UNSUPPORTED_IR);
371 }
372 out.extend_from_slice(&ir_version.to_le_bytes());
373
374 emit_opsets(&mut out, raw)?;
376
377 let graph = first_bytes(raw, 7)?;
379 if graph.is_empty() {
380 return Err(MISSING_GRAPH);
381 }
382 emit_canonical_graph(&mut out, graph, 0)?;
383
384 emit_model_meta(&mut out, raw)?;
386
387 Ok(Self { bytes: out })
388 }
389 }
390
391 fn emit_opsets(out: &mut Vec<u8>, model: &[u8]) -> Result<(), ShapeViolation> {
395 let entries = collect_spans(model, 8)?;
396
397 let mut ok_min = false;
399 for e in &entries {
400 let body = &model[e.off..e.off + e.len];
401 let domain = first_bytes(body, 1)?;
402 let version = first_varint(body, 2)?.unwrap_or(0) as i64;
403 if domain.is_empty() && version >= ONNX_OPSET_VERSION_MIN {
404 ok_min = true;
405 }
406 }
407 if !ok_min && !entries.is_empty() {
408 return Err(OPSET_TOO_OLD);
409 }
410
411 let mut order: Vec<usize> = (0..entries.len()).collect();
412 order.sort_by(|&a, &b| {
413 let ea = &model[entries[a].off..entries[a].off + entries[a].len];
414 let eb = &model[entries[b].off..entries[b].off + entries[b].len];
415 let ka = (
416 first_bytes(ea, 1).unwrap_or(&[]),
417 first_varint(ea, 2).ok().flatten().unwrap_or(0),
418 );
419 let kb = (
420 first_bytes(eb, 1).unwrap_or(&[]),
421 first_varint(eb, 2).ok().flatten().unwrap_or(0),
422 );
423 ka.cmp(&kb)
424 });
425
426 for &idx in &order {
427 let body = &model[entries[idx].off..entries[idx].off + entries[idx].len];
428 let domain = first_bytes(body, 1)?;
429 let version = first_varint(body, 2)?.unwrap_or(0) as i64;
430 out.extend_from_slice(&sha256(domain));
431 out.extend_from_slice(&version.to_le_bytes());
432 }
433 Ok(())
434 }
435
436 fn emit_model_meta(out: &mut Vec<u8>, model: &[u8]) -> Result<(), ShapeViolation> {
438 out.extend_from_slice(&sha256(first_bytes(model, 2)?)); out.extend_from_slice(&sha256(first_bytes(model, 3)?)); out.extend_from_slice(&sha256(first_bytes(model, 4)?)); out.extend_from_slice(&(first_varint(model, 5)?.unwrap_or(0) as i64).to_le_bytes()); emit_string_string(out, model, 14) }
444
445 fn emit_string_string(
448 out: &mut Vec<u8>,
449 body: &[u8],
450 field_no: u64,
451 ) -> Result<(), ShapeViolation> {
452 let entries = collect_spans(body, field_no)?;
453 let mut order: Vec<usize> = (0..entries.len()).collect();
454 order.sort_by(|&a, &b| {
455 let ka = first_bytes(&body[entries[a].off..entries[a].off + entries[a].len], 1)
456 .unwrap_or(&[]);
457 let kb = first_bytes(&body[entries[b].off..entries[b].off + entries[b].len], 1)
458 .unwrap_or(&[]);
459 ka.cmp(kb)
460 });
461 out.extend_from_slice(&(order.len() as u32).to_le_bytes());
462 for &idx in &order {
463 let e = &body[entries[idx].off..entries[idx].off + entries[idx].len];
464 out.extend_from_slice(&sha256(first_bytes(e, 1)?));
465 out.extend_from_slice(&sha256(first_bytes(e, 2)?));
466 }
467 Ok(())
468 }
469
470 fn emit_canonical_graph(
473 out: &mut Vec<u8>,
474 graph: &[u8],
475 depth: usize,
476 ) -> Result<(), ShapeViolation> {
477 if depth > ONNX_SUBGRAPH_DEPTH_MAX {
478 return Err(SUBGRAPH_DEPTH);
479 }
480
481 out.extend_from_slice(&sha256(first_bytes(graph, 2)?)); let nodes = collect_spans(graph, 1)?;
485 let node_count = nodes.len();
486 out.extend_from_slice(&(node_count as u32).to_le_bytes());
487
488 let mut emitted: Vec<bool> = alloc::vec![false; node_count];
489 let mut node_digest_cache: Vec<Option<[u8; 32]>> = alloc::vec![None; node_count];
490 for _ in 0..node_count {
491 let mut best: Option<usize> = None;
493 for (cand, node) in nodes.iter().enumerate() {
494 if emitted[cand] {
495 continue;
496 }
497 if !node_ready(graph, &nodes, &emitted, node)? {
498 continue;
499 }
500 best = Some(match best {
501 None => cand,
502 Some(b) => {
503 if node_lex_le(graph, &nodes, cand, b, &mut node_digest_cache)? {
504 cand
505 } else {
506 b
507 }
508 }
509 });
510 }
511 let pick = best.ok_or(GRAPH_CYCLE)?; let body = &graph[nodes[pick].off..nodes[pick].off + nodes[pick].len];
513 emit_node(out, body, depth)?;
514 emitted[pick] = true;
515 }
516
517 emit_tensor_section(out, graph, 5)?;
519
520 emit_value_info(out, graph, 11)?;
522 emit_value_info(out, graph, 12)?;
523 emit_value_info(out, graph, 13)?;
524
525 Ok(())
526 }
527
528 fn node_ready(
531 graph: &[u8],
532 nodes: &[Span],
533 emitted: &[bool],
534 node: &Span,
535 ) -> Result<bool, ShapeViolation> {
536 let body = &graph[node.off..node.off + node.len];
537 let mut ready = true;
538 for_each_field(body, 1, |v| {
539 if let FieldValue::Bytes(name) = v {
540 if !name.is_empty() {
541 for (k, prod) in nodes.iter().enumerate() {
542 let pbody = &graph[prod.off..prod.off + prod.len];
543 let mut produces = false;
544 for_each_field(pbody, 2, |ov| {
545 if let FieldValue::Bytes(on) = ov {
546 if on == name {
547 produces = true;
548 }
549 }
550 Ok(())
551 })?;
552 if produces && !emitted[k] {
553 ready = false;
554 }
555 }
556 }
557 }
558 Ok(())
559 })?;
560 Ok(ready)
561 }
562
563 fn node_lex_le(
566 graph: &[u8],
567 nodes: &[Span],
568 a_idx: usize,
569 b_idx: usize,
570 node_digest_cache: &mut [Option<[u8; 32]>],
571 ) -> Result<bool, ShapeViolation> {
572 let a = &nodes[a_idx];
573 let b = &nodes[b_idx];
574 let ba = &graph[a.off..a.off + a.len];
575 let bb = &graph[b.off..b.off + b.len];
576 let ka = (
577 first_bytes(ba, 3)?,
578 first_bytes(ba, 4)?,
579 first_bytes(ba, 7)?,
580 );
581 let kb = (
582 first_bytes(bb, 3)?,
583 first_bytes(bb, 4)?,
584 first_bytes(bb, 7)?,
585 );
586 if ka != kb {
587 return Ok(ka <= kb);
588 }
589
590 let oa = collect_spans(ba, 2)?;
591 let ob = collect_spans(bb, 2)?;
592 let common = core::cmp::min(oa.len(), ob.len());
593 for i in 0..common {
594 let a_out = &ba[oa[i].off..oa[i].off + oa[i].len];
595 let b_out = &bb[ob[i].off..ob[i].off + ob[i].len];
596 match a_out.cmp(b_out) {
597 core::cmp::Ordering::Less => return Ok(true),
598 core::cmp::Ordering::Greater => return Ok(false),
599 core::cmp::Ordering::Equal => {}
600 }
601 }
602 if oa.len() != ob.len() {
603 return Ok(oa.len() <= ob.len());
604 }
605
606 let da = cached_node_digest(graph, nodes, a_idx, node_digest_cache)?;
607 let db = cached_node_digest(graph, nodes, b_idx, node_digest_cache)?;
608
609 Ok(da <= db)
610 }
611
612 fn cached_node_digest(
613 graph: &[u8],
614 nodes: &[Span],
615 idx: usize,
616 node_digest_cache: &mut [Option<[u8; 32]>],
617 ) -> Result<[u8; 32], ShapeViolation> {
618 match node_digest_cache[idx] {
619 Some(d) => Ok(d),
620 None => {
621 let node = &nodes[idx];
622 let body = &graph[node.off..node.off + node.len];
623 let d = canonical_proto_digest(body, 0)?;
624 node_digest_cache[idx] = Some(d);
625 Ok(d)
626 }
627 }
628 }
629
630 fn emit_node(out: &mut Vec<u8>, node: &[u8], depth: usize) -> Result<(), ShapeViolation> {
634 out.extend_from_slice(&sha256(first_bytes(node, 3)?)); out.extend_from_slice(&sha256(first_bytes(node, 4)?)); out.extend_from_slice(&sha256(first_bytes(node, 7)?)); out.extend_from_slice(&sha256(first_bytes(node, 8)?)); let n_in = count_field(node, 1)?;
640 out.extend_from_slice(&(n_in as u32).to_le_bytes());
641 for_each_field(node, 1, |v| {
642 if let FieldValue::Bytes(name) = v {
643 out.extend_from_slice(&sha256(name));
644 }
645 Ok(())
646 })?;
647
648 let n_out = count_field(node, 2)?;
649 out.extend_from_slice(&(n_out as u32).to_le_bytes());
650 for_each_field(node, 2, |v| {
651 if let FieldValue::Bytes(name) = v {
652 out.extend_from_slice(&sha256(name));
653 }
654 Ok(())
655 })?;
656
657 emit_attributes(out, node, depth)
658 }
659
660 fn emit_attributes(out: &mut Vec<u8>, node: &[u8], depth: usize) -> Result<(), ShapeViolation> {
662 let attrs = collect_spans(node, 5)?;
663 let mut order: Vec<usize> = (0..attrs.len()).collect();
664 order.sort_by(|&a, &b| {
665 let na =
666 first_bytes(&node[attrs[a].off..attrs[a].off + attrs[a].len], 1).unwrap_or(&[]);
667 let nb =
668 first_bytes(&node[attrs[b].off..attrs[b].off + attrs[b].len], 1).unwrap_or(&[]);
669 na.cmp(nb)
670 });
671 out.extend_from_slice(&(order.len() as u32).to_le_bytes());
672 for &idx in &order {
673 let a = &node[attrs[idx].off..attrs[idx].off + attrs[idx].len];
674 out.extend_from_slice(&sha256(first_bytes(a, 1)?)); let atype = first_varint(a, 20)?.unwrap_or(0) as i32;
676 out.extend_from_slice(&atype.to_le_bytes());
677 emit_attribute_value(out, a, atype, depth)?;
678 }
679 Ok(())
680 }
681
682 fn emit_attribute_value(
684 out: &mut Vec<u8>,
685 a: &[u8],
686 atype: i32,
687 depth: usize,
688 ) -> Result<(), ShapeViolation> {
689 match atype {
690 1 => {
691 if let Some(FieldValue::Fixed32(bits)) = first_field(a, 2)? {
693 out.extend_from_slice(&bits.to_le_bytes());
694 }
695 }
696 2 => {
697 out.extend_from_slice(&(first_varint(a, 3)?.unwrap_or(0) as i64).to_le_bytes());
699 }
700 3 => {
701 out.extend_from_slice(&sha256(first_bytes(a, 4)?));
703 }
704 4 => {
705 emit_tensor(out, first_bytes(a, 5)?)?;
707 }
708 5 => {
709 emit_canonical_graph(out, first_bytes(a, 6)?, depth + 1)?;
711 }
712 6 => {
713 for_each_field(a, 7, |v| {
715 if let FieldValue::Bytes(p) = v {
716 out.extend_from_slice(&sha256(p));
717 } else if let FieldValue::Fixed32(b) = v {
718 out.extend_from_slice(&b.to_le_bytes());
719 }
720 Ok(())
721 })?;
722 }
723 7 => {
724 emit_packed_varints(out, a, 8)?;
726 }
727 8 => {
728 for_each_field(a, 9, |v| {
730 if let FieldValue::Bytes(s) = v {
731 out.extend_from_slice(&sha256(s));
732 }
733 Ok(())
734 })?;
735 }
736 9 => {
737 let spans = collect_spans(a, 10)?;
739 for s in &spans {
740 emit_tensor(out, &a[s.off..s.off + s.len])?;
741 }
742 }
743 10 => {
744 let spans = collect_spans(a, 11)?;
746 for s in &spans {
747 emit_canonical_graph(out, &a[s.off..s.off + s.len], depth + 1)?;
748 }
749 }
750 11 => out.extend_from_slice(&canonical_proto_digest(first_bytes(a, 22)?, 0)?), 12 => {
752 for_each_field(a, 23, |v| {
754 if let FieldValue::Bytes(s) = v {
755 out.extend_from_slice(&canonical_proto_digest(s, 0)?);
756 }
757 Ok(())
758 })?;
759 }
760 13 => out.extend_from_slice(&canonical_proto_digest(first_bytes(a, 14)?, 0)?), 14 => {
762 for_each_field(a, 15, |v| {
764 if let FieldValue::Bytes(s) = v {
765 out.extend_from_slice(&canonical_proto_digest(s, 0)?);
766 }
767 Ok(())
768 })?;
769 }
770 _ => {}
771 }
772 Ok(())
773 }
774
775 fn emit_packed_varints(
776 out: &mut Vec<u8>,
777 body: &[u8],
778 field_no: u64,
779 ) -> Result<(), ShapeViolation> {
780 for_each_field(body, field_no, |v| {
781 match v {
782 FieldValue::Bytes(p) => {
783 let mut pos = 0;
784 while pos < p.len() {
785 let (val, np) = read_varint(p, pos).map_err(from_wire)?;
786 out.extend_from_slice(&(val as i64).to_le_bytes());
787 pos = np;
788 }
789 }
790 FieldValue::Varint(val) => out.extend_from_slice(&(val as i64).to_le_bytes()),
791 _ => {}
792 }
793 Ok(())
794 })
795 }
796
797 fn emit_tensor_section(
799 out: &mut Vec<u8>,
800 graph: &[u8],
801 field_no: u64,
802 ) -> Result<(), ShapeViolation> {
803 let spans = collect_spans(graph, field_no)?;
804 let mut order: Vec<usize> = (0..spans.len()).collect();
805 order.sort_by(|&a, &b| {
806 let na =
807 first_bytes(&graph[spans[a].off..spans[a].off + spans[a].len], 8).unwrap_or(&[]);
808 let nb =
809 first_bytes(&graph[spans[b].off..spans[b].off + spans[b].len], 8).unwrap_or(&[]);
810 na.cmp(nb)
811 });
812 out.extend_from_slice(&(order.len() as u32).to_le_bytes());
813 for &idx in &order {
814 let body = &graph[spans[idx].off..spans[idx].off + spans[idx].len];
815 emit_tensor(out, body)?;
816 }
817 Ok(())
818 }
819
820 fn emit_tensor(out: &mut Vec<u8>, t: &[u8]) -> Result<(), ShapeViolation> {
827 let dtype_id = first_varint(t, 2)?.unwrap_or(0) as i32;
828 let dtype = OnnxDataType::from_i32(dtype_id).ok_or(UNKNOWN_DTYPE)?;
829
830 out.extend_from_slice(&sha256(first_bytes(t, 8)?)); out.extend_from_slice(&dtype_id.to_le_bytes());
832
833 let rank = count_dims(t)?;
835 out.extend_from_slice(&(rank as u32).to_le_bytes());
836 emit_packed_varints(out, t, 1)?;
837
838 out.extend_from_slice(&tensor_data_digest(t, dtype)?);
840 Ok(())
841 }
842
843 fn count_dims(t: &[u8]) -> Result<usize, ShapeViolation> {
844 let mut n = 0;
845 for_each_field(t, 1, |v| {
846 match v {
847 FieldValue::Bytes(p) => {
848 let mut pos = 0;
849 while pos < p.len() {
850 let (_, np) = read_varint(p, pos).map_err(from_wire)?;
851 n += 1;
852 pos = np;
853 }
854 }
855 FieldValue::Varint(_) => n += 1,
856 _ => {}
857 }
858 Ok(())
859 })?;
860 Ok(n)
861 }
862
863 fn tensor_data_digest(t: &[u8], dtype: OnnxDataType) -> Result<[u8; 32], ShapeViolation> {
866 if first_varint(t, 14)?.unwrap_or(0) == 1 {
874 let mut h = Sha256Hasher::initial();
875 fold(&mut h, b"onnx:external-data:v1");
876 let mut sub: Vec<u8> = Vec::new();
878 emit_string_string(&mut sub, t, 13)?;
879 fold(&mut h, &sub);
880 return Ok(h.finalize());
881 }
882 if let Some(FieldValue::Bytes(raw)) = first_field(t, 9)? {
884 if !raw.is_empty() {
885 return Ok(sha256(raw));
886 }
887 }
888 let mut h = Sha256Hasher::initial();
889 match dtype {
890 OnnxDataType::Float => fold_fixed_payload(t, 4, &mut h)?,
893 OnnxDataType::Double | OnnxDataType::Complex128 => fold_fixed_payload(t, 10, &mut h)?,
894 OnnxDataType::Complex64 => fold_fixed_payload(t, 4, &mut h)?,
895 OnnxDataType::Int64 => fold_typed_varints(t, 7, 8, &mut h)?,
897 OnnxDataType::Uint64 => fold_typed_varints(t, 11, 8, &mut h)?,
899 OnnxDataType::Uint32 => fold_typed_varints(t, 11, 4, &mut h)?,
900 OnnxDataType::Int32 => fold_typed_varints(t, 5, 4, &mut h)?,
903 OnnxDataType::Int16
904 | OnnxDataType::Uint16
905 | OnnxDataType::Float16
906 | OnnxDataType::Bfloat16 => fold_typed_varints(t, 5, 2, &mut h)?,
907 OnnxDataType::Int8
908 | OnnxDataType::Uint8
909 | OnnxDataType::Bool
910 | OnnxDataType::Float8E4M3Fn
911 | OnnxDataType::Float8E4M3Fnuz
912 | OnnxDataType::Float8E5M2
913 | OnnxDataType::Float8E5M2Fnuz
914 | OnnxDataType::Int4
915 | OnnxDataType::Uint4
916 | OnnxDataType::Float4E2M1 => fold_typed_varints(t, 5, 1, &mut h)?,
917 OnnxDataType::String => {
919 for_each_field(t, 6, |v| {
920 if let FieldValue::Bytes(s) = v {
921 fold(&mut h, &sha256(s));
922 }
923 Ok(())
924 })?;
925 }
926 }
927 Ok(h.finalize())
928 }
929
930 fn fold_fixed_payload(
933 body: &[u8],
934 field_no: u64,
935 h: &mut Sha256Hasher,
936 ) -> Result<(), ShapeViolation> {
937 for_each_field(body, field_no, |v| {
938 match v {
939 FieldValue::Bytes(p) => fold(h, p),
940 FieldValue::Fixed32(b) => fold(h, &b.to_le_bytes()),
941 FieldValue::Fixed64(b) => fold(h, &b.to_le_bytes()),
942 _ => {}
943 }
944 Ok(())
945 })
946 }
947
948 fn fold_typed_varints(
951 body: &[u8],
952 field_no: u64,
953 width: usize,
954 h: &mut Sha256Hasher,
955 ) -> Result<(), ShapeViolation> {
956 for_each_field(body, field_no, |v| {
957 match v {
958 FieldValue::Bytes(p) => {
959 let mut pos = 0;
960 while pos < p.len() {
961 let (val, np) = read_varint(p, pos).map_err(from_wire)?;
962 fold(h, &val.to_le_bytes()[..width]);
963 pos = np;
964 }
965 }
966 FieldValue::Varint(val) => fold(h, &val.to_le_bytes()[..width]),
967 _ => {}
968 }
969 Ok(())
970 })
971 }
972
973 fn emit_value_info(
977 out: &mut Vec<u8>,
978 graph: &[u8],
979 field_no: u64,
980 ) -> Result<(), ShapeViolation> {
981 let spans = collect_spans(graph, field_no)?;
982 let mut order: Vec<usize> = (0..spans.len()).collect();
983 order.sort_by(|&a, &b| {
984 let na =
985 first_bytes(&graph[spans[a].off..spans[a].off + spans[a].len], 1).unwrap_or(&[]);
986 let nb =
987 first_bytes(&graph[spans[b].off..spans[b].off + spans[b].len], 1).unwrap_or(&[]);
988 na.cmp(nb)
989 });
990 out.extend_from_slice(&(order.len() as u32).to_le_bytes());
991 for &idx in &order {
992 let body = &graph[spans[idx].off..spans[idx].off + spans[idx].len];
993 out.extend_from_slice(&sha256(first_bytes(body, 1)?)); out.extend_from_slice(&canonical_proto_digest(first_bytes(body, 2)?, 0)?);
995 }
997 Ok(())
998 }
999
1000 pub fn canonicalize(raw: &[u8]) -> Result<Vec<u8>, ShapeViolation> {
1006 Ok(OnnxValue::parse(raw)?.bytes)
1007 }
1008
1009 #[cfg(test)]
1010 mod tests {
1011 use super::*;
1012
1013 fn put_varint(out: &mut Vec<u8>, mut v: u64) {
1016 loop {
1017 let mut byte = (v & 0x7f) as u8;
1018 v >>= 7;
1019 if v != 0 {
1020 byte |= 0x80;
1021 }
1022 out.push(byte);
1023 if v == 0 {
1024 break;
1025 }
1026 }
1027 }
1028
1029 fn tag(out: &mut Vec<u8>, field_no: u64, wire: u64) {
1030 put_varint(out, (field_no << 3) | wire);
1031 }
1032
1033 fn field_varint(out: &mut Vec<u8>, field_no: u64, v: u64) {
1034 tag(out, field_no, 0);
1035 put_varint(out, v);
1036 }
1037
1038 fn field_bytes(out: &mut Vec<u8>, field_no: u64, b: &[u8]) {
1039 tag(out, field_no, 2);
1040 put_varint(out, b.len() as u64);
1041 out.extend_from_slice(b);
1042 }
1043
1044 fn minimal_onnx() -> Vec<u8> {
1047 let mut opset = Vec::new();
1049 field_bytes(&mut opset, 1, b""); field_varint(&mut opset, 2, 1); let mut graph = Vec::new();
1054 field_bytes(&mut graph, 2, b"g");
1055
1056 let mut model = Vec::new();
1058 field_varint(&mut model, 1, ONNX_IR_VERSION_MAX as u64); field_bytes(&mut model, 7, &graph); field_bytes(&mut model, 8, &opset); model
1062 }
1063
1064 #[test]
1065 fn parses_minimal_model() {
1066 let canon = canonicalize(&minimal_onnx()).expect("valid");
1067 assert_eq!(
1073 canon.len(),
1074 8 + 40 + (32 + 4 + 4 + 12) + (32 + 32 + 32 + 8 + 4)
1075 );
1076 }
1077
1078 #[test]
1079 fn rejects_out_of_range_ir() {
1080 let mut model = Vec::new();
1083 field_varint(&mut model, 1, (ONNX_IR_VERSION_MAX + 1) as u64);
1084 let err = OnnxValue::parse(&model).expect_err("unsupported ir");
1085 assert_eq!(err.constraint_iri, UNSUPPORTED_IR.constraint_iri);
1086 }
1087
1088 #[test]
1089 fn rejects_missing_graph() {
1090 let mut opset = Vec::new();
1091 field_bytes(&mut opset, 1, b"");
1092 field_varint(&mut opset, 2, 1);
1093 let mut model = Vec::new();
1094 field_varint(&mut model, 1, ONNX_IR_VERSION_MAX as u64);
1095 field_bytes(&mut model, 8, &opset);
1096 let err = OnnxValue::parse(&model).expect_err("no graph");
1097 assert_eq!(err.constraint_iri, MISSING_GRAPH.constraint_iri);
1098 }
1099
1100 #[test]
1101 fn deterministic() {
1102 let a = canonicalize(&minimal_onnx()).expect("valid");
1103 let b = canonicalize(&minimal_onnx()).expect("valid");
1104 assert_eq!(a, b);
1105 }
1106 }
1107}