onnx_runtime_loader/
proto.rs

1//! ONNX protobuf decoding (§19.1).
2//!
3//! The `onnx` submodule contains the `prost`-generated types compiled from the
4//! vendored `proto/onnx.proto3` (see `build.rs`). [`decode_model`] parses a
5//! serialized `ModelProto` from bytes.
6
7use prost::Message;
8
9use crate::LoaderError;
10
11/// The `prost`-generated ONNX protobuf types (package `onnx`).
12#[allow(clippy::all, missing_docs, non_snake_case)]
13pub mod onnx {
14    include!(concat!(env!("OUT_DIR"), "/onnx.rs"));
15}
16
17pub use onnx::ModelProto;
18
19/// Encoded `FileDescriptorSet` for the exact vendored ONNX schema used to
20/// generate [`onnx`]. Textual codecs use this descriptor so every present and
21/// future field in the bound proto is handled from one source of truth.
22pub const FILE_DESCRIPTOR_SET: &[u8] =
23    include_bytes!(concat!(env!("OUT_DIR"), "/onnx_descriptor.bin"));
24
25/// Decode a [`ModelProto`] from serialized protobuf bytes.
26pub fn decode_model(bytes: &[u8]) -> Result<ModelProto, LoaderError> {
27    ModelProto::decode(bytes).map_err(|e| LoaderError::ProtobufParse(e.to_string()))
28}
29
30/// Convert an ONNX protobuf **TextFormat** (`.textproto`) document into the
31/// binary protobuf wire encoding of a `ModelProto`.
32///
33/// Parsing is descriptor-driven from the exact vendored ONNX schema (via
34/// [`FILE_DESCRIPTOR_SET`]), so every bound message, field, oneof, and enum is
35/// covered — the same source of truth used to generate the prost types. The
36/// returned bytes are byte-for-byte loadable by [`decode_model`] and the
37/// weight/graph build path, so a textproto fixture flows through the identical
38/// binary-decode pipeline as a real `.onnx` model.
39///
40/// Because this yields a self-contained binary buffer with no model-directory
41/// context, textproto fixtures must inline all weights (no external
42/// `.onnx.data`).
43pub fn textproto_to_binary(text: &str) -> Result<Vec<u8>, LoaderError> {
44    use prost_reflect::{DescriptorPool, DynamicMessage};
45    use std::sync::OnceLock;
46
47    static POOL: OnceLock<DescriptorPool> = OnceLock::new();
48    let pool = POOL.get_or_init(|| {
49        DescriptorPool::decode(FILE_DESCRIPTOR_SET)
50            .expect("the generated ONNX descriptor set must be valid")
51    });
52    let descriptor = pool
53        .get_message_by_name("onnx.ModelProto")
54        .expect("the ONNX descriptor set must define onnx.ModelProto");
55    let dynamic = DynamicMessage::parse_text_format(descriptor, text)
56        .map_err(|e| LoaderError::TextProtoParse(e.to_string()))?;
57    Ok(dynamic.encode_to_vec())
58}
onnx_runtime_loader/proto.rs

onnx_runtime_loader/
proto.rs
