Skip to main content

uor_addr/onnx/
value.rs

1//! ONNX typed input (IR ≤ v13) (ADR-023 amended by ADR-060).
2//!
3//! Protobuf v3 admits many byte-representations of the same logical
4//! message; this realization defines a canonical form — a **flat
5//! skeleton** — that collapses that freedom. Two ONNX models that decode
6//! to the same logical content (regardless of protobuf field order, node
7//! ordering among valid topological orderings, or whether tensor data is
8//! stored in `raw_data` or the typed-data fields) canonicalize to
9//! byte-identical skeletons and therefore to the same κ-label.
10//!
11//! ```text
12//! LE_i64(ir_version)
13//! ── opset imports, sorted by (domain, version) ──
14//!   for op: sha256(domain) || LE_i64(version)
15//! ── graph (recursive) ──
16//!   sha256(graph_name)
17//!   nodes in Kahn-topological order
18//!   (lex (name, op_type, domain, outputs..., canonical_node_digest) tie-break):
19//!     sha256(name) || sha256(op_type) || sha256(domain) || sha256(overload)
20//!       || LE_u32(n_in)  || (sha256(input_name)  × n_in)
21//!       || LE_u32(n_out) || (sha256(output_name) × n_out)
22//!       || attributes, sorted by name (GRAPH/GRAPHS recurse inline)
23//!   initializers (#5), sorted by name, each a canonical TensorProto record
24//!   graph input (#11) / output (#12) / value_info (#13), sorted by name
25//! ── model metadata ──
26//!   sha256(producer_name) || sha256(producer_version) || sha256(domain)
27//!     || LE_i64(model_version) || metadata_props sorted by key
28//! ```
29//!
30//! Under ADR-060 the **full skeleton** flows through the pipeline as a
31//! [`TermValue::Borrowed`] carrier and ψ₉ folds it through the σ-axis —
32//! there is no two-level commitment, no carrier ceiling, and no node /
33//! attribute / initializer / IO count cap. Variable-length leaves (tensor
34//! data bytes, strings, opaque sub-message payloads) are still replaced by
35//! their `sha256(...)` digest so the skeleton stays bounded by structure
36//! size, not data size, while still binding every weight byte into the
37//! κ-label.
38//!
39//! [`OnnxValue`] (the owned parsed value, `alloc`-gated) holds the
40//! skeleton; [`OnnxCarrier`] is the borrowed model-input handle the
41//! pipeline binds.
42
43use prism::operation::TermValue;
44use prism::pipeline::{
45    ConstrainedTypeShape, ConstraintRef, IntoBindingValue, PartitionProductFields,
46};
47
48// ─── OnnxCarrier — the borrowed model-input handle (no_alloc) ───────────
49
50/// Borrowed canonical-skeleton input handle (ADR-060 borrowed carrier). A
51/// thin, `Copy` borrow of the skeleton bytes produced by [`canonicalize`];
52/// `as_binding_value` returns the `Borrowed` carrier zero-copy.
53#[derive(Clone, Copy, Debug)]
54pub struct OnnxCarrier<'a>(&'a [u8]);
55
56impl<'a> OnnxCarrier<'a> {
57    /// Wrap a canonical-skeleton byte slice as a model input handle.
58    #[must_use]
59    pub fn new(skeleton: &'a [u8]) -> Self {
60        Self(skeleton)
61    }
62
63    /// Borrow the canonical-skeleton bytes.
64    #[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// ═════════════════════════════════════════════════════════════════════
91// alloc-gated parser + owned value
92// ═════════════════════════════════════════════════════════════════════
93
94#[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    // ─── ShapeViolation IRIs ─────────────────────────────────────────────
112
113    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    /// Recursion ceiling for the opaque-message field-order canonicalizer.
157    const CANON_PROTO_DEPTH_MAX: usize = 32;
158
159    /// Field-order-canonical digest of an opaque protobuf message — folds
160    /// its fields in ascending field-number order (stable within a number,
161    /// so repeated-field order is preserved), recursing into
162    /// length-delimited fields. This applies canonicalization rule 1
163    /// (field-number ordering) to sub-messages the realization otherwise
164    /// treats opaquely (`TypeProto`, `SparseTensorProto`), so two
165    /// serializations of the same logical value canonicalize identically.
166    /// Returns a 32-byte **leaf digest** (an opaque sub-message is a leaf
167    /// — its digest is appended inline, never expanded into the skeleton).
168    ///
169    /// A length-delimited field that is genuinely a string / bytes leaf
170    /// (e.g. `dim_param`) generally fails to re-parse as a well-formed
171    /// message; that case falls back to a digest of the raw payload. The
172    /// transform is deterministic either way: identical bytes always take
173    /// the same path.
174    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        // Stable sort by field number (preserves repeated-field order).
218        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    /// Fold `bytes` into the running hasher behind a mutable reference (the
243    /// `Hasher::fold_bytes` consume-and-return API is awkward inside
244    /// `FnMut` closures; this wraps the take-replace dance once).
245    #[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    // ─── Protobuf field accessors over a message body ──────────────────
252
253    /// First occurrence of `field_no` in `body`, or `None`.
254    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    /// Invoke `f` for every occurrence of `field_no` (the repeated-field
279    /// iterator). Stops and propagates the first error `f` returns.
280    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    /// A `(offset, len)` span into a parent buffer.
304    #[derive(Clone, Copy)]
305    struct Span {
306        off: usize,
307        len: usize,
308    }
309
310    /// Collect every occurrence of `field_no` (length-delimited) in `body`
311    /// as a `(offset, len)` span.
312    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    /// A parsed, canonicalized ONNX `ModelProto`. The stored bytes are the
329    /// flat canonical skeleton (see [module docs](super)). **`alloc`-gated**
330    /// — the pipeline binds the borrowed [`OnnxCarrier`](super::OnnxCarrier).
331    #[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        /// Borrow the canonical-skeleton bytes.
346        #[must_use]
347        pub fn canonical_bytes(&self) -> &[u8] {
348            &self.bytes
349        }
350
351        /// Parse an ONNX `ModelProto` wire buffer into a canonicalized
352        /// skeleton.
353        ///
354        /// # Errors
355        ///
356        /// A [`ShapeViolation`] whose `constraint_iri` names the violated
357        /// invariant (protobuf decode failure, unsupported IR version,
358        /// opset below the minimum, missing graph, a subgraph cycle, an
359        /// over-deep subgraph nesting, or an unknown tensor data type).
360        pub fn parse(raw: &[u8]) -> Result<Self, ShapeViolation> {
361            let mut out: Vec<u8> = Vec::new();
362
363            // ── ir_version (ModelProto #1) ──
364            // Accept any known IR revision (1..=ONNX_IR_VERSION_MAX); the
365            // canonical skeleton is IR-version-agnostic and binds the
366            // `ir_version` value, so distinct revisions canonicalize
367            // distinctly. Reject absent / 0 / a future unknown revision.
368            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            // ── opset imports (ModelProto #8, repeated OperatorSetIdProto) ──
375            emit_opsets(&mut out, raw)?;
376
377            // ── graph (ModelProto #7) ──
378            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            // ── model metadata ──
385            emit_model_meta(&mut out, raw)?;
386
387            Ok(Self { bytes: out })
388        }
389    }
390
391    /// Emit opset imports sorted by `(domain, version)`. Enforces at least
392    /// one default-domain (`""`) import at or above
393    /// [`ONNX_OPSET_VERSION_MIN`].
394    fn emit_opsets(out: &mut Vec<u8>, model: &[u8]) -> Result<(), ShapeViolation> {
395        let entries = collect_spans(model, 8)?;
396
397        // Default-domain minimum-version check.
398        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    /// Emit producer / domain / model_version + sorted `metadata_props`.
437    fn emit_model_meta(out: &mut Vec<u8>, model: &[u8]) -> Result<(), ShapeViolation> {
438        out.extend_from_slice(&sha256(first_bytes(model, 2)?)); // producer_name
439        out.extend_from_slice(&sha256(first_bytes(model, 3)?)); // producer_version
440        out.extend_from_slice(&sha256(first_bytes(model, 4)?)); // domain
441        out.extend_from_slice(&(first_varint(model, 5)?.unwrap_or(0) as i64).to_le_bytes()); // model_version
442        emit_string_string(out, model, 14) // metadata_props
443    }
444
445    /// Emit a repeated `StringStringEntryProto` map (`field_no`), sorted by
446    /// key, inline: for each entry `sha256(key) || sha256(value)`.
447    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    /// Emit a `GraphProto` body inline, recursing into subgraphs bounded by
471    /// [`ONNX_SUBGRAPH_DEPTH_MAX`].
472    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)?)); // graph name
482
483        // ── Nodes in Kahn-topological order (lex tie-break) ──
484        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            // Find the lex-min ready (all producers emitted), unemitted node.
492            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)?; // no ready node ⇒ cycle
512            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        // ── Initializers (#5), sorted by name, with tensor-data digests ──
518        emit_tensor_section(out, graph, 5)?;
519
520        // ── Graph IO: inputs (#11), outputs (#12), value_info (#13) ──
521        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    /// A node is ready when every input name that is *produced by another
529    /// node in this graph* has had its producer emitted.
530    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    /// Lexicographic order on
564    /// `(name, op_type, domain, outputs..., canonical_node_digest)`.
565    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    /// Emit a `NodeProto` inline: identity fields, positional inputs /
631    /// outputs, then the name-sorted attributes (which recurse into
632    /// subgraphs inline).
633    fn emit_node(out: &mut Vec<u8>, node: &[u8], depth: usize) -> Result<(), ShapeViolation> {
634        out.extend_from_slice(&sha256(first_bytes(node, 3)?)); // name
635        out.extend_from_slice(&sha256(first_bytes(node, 4)?)); // op_type
636        out.extend_from_slice(&sha256(first_bytes(node, 7)?)); // domain
637        out.extend_from_slice(&sha256(first_bytes(node, 8)?)); // overload (IR v10+)
638
639        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    /// Emit a node's `attribute` field (#5), sorted by name, inline.
661    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)?)); // name
675            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    /// Emit an attribute's value inline, dispatched on its `AttributeType`.
683    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                // FLOAT (#2, fixed32)
692                if let Some(FieldValue::Fixed32(bits)) = first_field(a, 2)? {
693                    out.extend_from_slice(&bits.to_le_bytes());
694                }
695            }
696            2 => {
697                // INT (#3, varint)
698                out.extend_from_slice(&(first_varint(a, 3)?.unwrap_or(0) as i64).to_le_bytes());
699            }
700            3 => {
701                // STRING (#4, bytes)
702                out.extend_from_slice(&sha256(first_bytes(a, 4)?));
703            }
704            4 => {
705                // TENSOR (#5)
706                emit_tensor(out, first_bytes(a, 5)?)?;
707            }
708            5 => {
709                // GRAPH (#6) — recurse inline
710                emit_canonical_graph(out, first_bytes(a, 6)?, depth + 1)?;
711            }
712            6 => {
713                // FLOATS (#7, packed fixed32)
714                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                // INTS (#8, packed varint)
725                emit_packed_varints(out, a, 8)?;
726            }
727            8 => {
728                // STRINGS (#9, repeated bytes)
729                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                // TENSORS (#10)
738                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                // GRAPHS (#11) — recurse inline
745                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)?), // SPARSE_TENSOR
751            12 => {
752                // SPARSE_TENSORS (#23)
753                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)?), // TYPE_PROTO
761            14 => {
762                // TYPE_PROTOS (#15)
763                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    /// Emit a name-sorted section of repeated `TensorProto` (initializers).
798    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    /// Emit a canonical `TensorProto` record inline: `sha256(name) ||
821    /// LE_i32(dtype) || LE_u32(rank) || (LE_i64 dim …) || tensor_data_digest`,
822    /// where `tensor_data_digest` is a 32-byte leaf digest streaming
823    /// `raw_data` if present, else the typed-data field re-encoded to the
824    /// canonical little-endian `raw_data` layout (so the two storage forms
825    /// canonicalize identically).
826    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)?)); // name
831        out.extend_from_slice(&dtype_id.to_le_bytes());
832
833        // dims (#1, repeated int64; packed or unpacked).
834        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        // data digest (a leaf — appended inline as 32 bytes).
839        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    /// Stream the tensor's data through SHA-256 in canonical `raw_data`
864    /// layout, returning the 32-byte leaf digest.
865    fn tensor_data_digest(t: &[u8], dtype: OnnxDataType) -> Result<[u8; 32], ShapeViolation> {
866        // External data (`data_location` #14 == EXTERNAL = 1): the core
867        // cannot open the referenced sibling file, so the κ-label binds the
868        // external *reference* (`external_data` #13 — location / offset /
869        // length / checksum, sorted by key) rather than the dereferenced
870        // bytes. A domain tag keeps external digests disjoint from inline
871        // ones. Hosts requiring inline≡external equivalence dereference
872        // before calling.
873        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            // metadata-style sorted digest of external_data (#13).
877            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        // raw_data (#9) takes precedence and is already canonical.
883        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            // float_data (#4) / double_data (#10): packed fixed-width — the
891            // packed payload IS the canonical raw layout.
892            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            // int64_data (#7): re-encode each varint to 8-byte LE.
896            OnnxDataType::Int64 => fold_typed_varints(t, 7, 8, &mut h)?,
897            // uint64_data (#11): UINT64 → 8-byte LE; UINT32 → 4-byte LE.
898            OnnxDataType::Uint64 => fold_typed_varints(t, 11, 8, &mut h)?,
899            OnnxDataType::Uint32 => fold_typed_varints(t, 11, 4, &mut h)?,
900            // int32_data (#5) carries INT32/INT16/INT8/UINT16/UINT8/BOOL and
901            // the bit-packed small floats — re-encode to the dtype's width.
902            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            // string_data (#6): fold each element's digest.
918            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    /// Fold the (already-canonical) packed payload of a fixed-width
931    /// repeated field directly.
932    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    /// Re-encode each varint of a packed/unpacked repeated field to `width`
949    /// little-endian bytes (the canonical `raw_data` element layout).
950    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    /// Emit a name-sorted section of repeated `ValueInfoProto` (graph
974    /// input / output / value_info). Binds the name plus a field-order-
975    /// canonical leaf digest of the `TypeProto`.
976    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)?)); // name
994            out.extend_from_slice(&canonical_proto_digest(first_bytes(body, 2)?, 0)?);
995            // type (TypeProto)
996        }
997        Ok(())
998    }
999
1000    /// Canonical skeleton as an owned `Vec<u8>`.
1001    ///
1002    /// # Errors
1003    ///
1004    /// Surfaces the [`ShapeViolation`] [`OnnxValue::parse`] would raise.
1005    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        // ── Minimal protobuf encoders for building test `ModelProto`s ──
1014
1015        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        /// Smallest valid ONNX `ModelProto`: ir_version=13, one
1045        /// default-domain opset import (version 1), and a non-empty graph.
1046        fn minimal_onnx() -> Vec<u8> {
1047            // OperatorSetIdProto { domain = "", version = 1 }
1048            let mut opset = Vec::new();
1049            field_bytes(&mut opset, 1, b""); // domain
1050            field_varint(&mut opset, 2, 1); // version
1051
1052            // GraphProto { name = "g" }
1053            let mut graph = Vec::new();
1054            field_bytes(&mut graph, 2, b"g");
1055
1056            // ModelProto
1057            let mut model = Vec::new();
1058            field_varint(&mut model, 1, ONNX_IR_VERSION_MAX as u64); // ir_version
1059            field_bytes(&mut model, 7, &graph); // graph
1060            field_bytes(&mut model, 8, &opset); // opset_import
1061            model
1062        }
1063
1064        #[test]
1065        fn parses_minimal_model() {
1066            let canon = canonicalize(&minimal_onnx()).expect("valid");
1067            // ir_version(8) + opset(domain digest 32 + version 8)
1068            //   + graph: name(32) + node_count(4) + init_count(4)
1069            //     + 3× IO counts(4 each)
1070            //   + meta: producer(32) + producer_ver(32) + domain(32)
1071            //     + model_ver(8) + metadata_props count(4)
1072            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            // IR 7 is in range (1..=13) → accepted (would reach MISSING_GRAPH);
1081            // 14 is a future/unknown revision → rejected at the IR gate.
1082            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}