//! Whether an ONNX model keeps tensor data outside its own bytes (SPEC §3.1,
//! *Decoding against a verified set*).
//!
//! A pinned decode builds its session from the bytes the pin verified. An ONNX
//! tensor can instead name another file for its data (`data_location = EXTERNAL`,
//! with the file in `external_data`'s `location` entry), and a runtime given only
//! the model's bytes resolves that file itself, outside the verified set. So before
//! a session is built from verified bytes, the model is read for such a tensor and
//! refused when it has one.
//!
//! This reads the protobuf wire format of `onnx.proto` directly and visits every
//! place a tensor can sit: the graph's initializers and sparse initializers, the
//! tensors and sparse tensors of every node attribute, the subgraphs of attributes
//! (control flow), the graphs of the training information, and the nodes of the
//! model's local functions. A field read with an unexpected wire type is skipped,
//! as a protobuf parser treats it as an unknown field. Bytes that are not a
//! well-formed protobuf message are an error, so a model this reader cannot walk is
//! never taken as free of external data.

/// Deepest nesting of messages read before the model is refused as malformed.
const MAX_DEPTH: usize = 64;

/// The `location` of the first tensor in `model`, an ONNX `ModelProto`, whose data
/// is stored outside the model (`""` when it names none); `None` when every tensor
/// keeps its data inside. `Err` when `model` is not a well-formed protobuf message.
pub(crate) fn external_tensor(model: &[u8]) -> Result<Option<String>, String> {
    walk(model, Msg::Model, 0)
}

#[derive(Clone, Copy)]
enum Msg {
    Model,
    TrainingInfo,
    Function,
    Graph,
    Node,
    Attribute,
    SparseTensor,
    Tensor,
}

impl Msg {
    /// The message type of field `field` of `self`, when it can hold a tensor.
    fn child(self, field: u64) -> Option<Msg> {
        use Msg::*;
        Some(match (self, field) {
            // ModelProto: graph = 7, training_info = 20, functions = 25.
            (Model, 7) => Graph,
            (Model, 20) => TrainingInfo,
            (Model, 25) => Function,
            // TrainingInfoProto: initialization = 1, algorithm = 2.
            (TrainingInfo, 1 | 2) => Graph,
            // FunctionProto: node = 7.
            (Function, 7) => Node,
            // GraphProto: node = 1, initializer = 5, sparse_initializer = 15.
            (Graph, 1) => Node,
            (Graph, 5) => Tensor,
            (Graph, 15) => SparseTensor,
            // NodeProto: attribute = 5.
            (Node, 5) => Attribute,
            // AttributeProto: t = 5, g = 6, tensors = 10, graphs = 11,
            // sparse_tensor = 22, sparse_tensors = 23.
            (Attribute, 5 | 10) => Tensor,
            (Attribute, 6 | 11) => Graph,
            (Attribute, 22 | 23) => SparseTensor,
            // SparseTensorProto: values = 1, indices = 2.
            (SparseTensor, 1 | 2) => Tensor,
            _ => return None,
        })
    }
}

fn varint(b: &[u8], at: &mut usize) -> Result<u64, String> {
    let mut v = 0u64;
    for shift in (0..64).step_by(7) {
        let byte = *b.get(*at).ok_or("truncated varint")?;
        *at += 1;
        v |= u64::from(byte & 0x7f) << shift;
        if byte & 0x80 == 0 {
            return Ok(v);
        }
    }
    Err("varint longer than 10 bytes".into())
}

/// The bytes of a length-delimited field at `at`.
fn delimited<'a>(b: &'a [u8], at: &mut usize) -> Result<&'a [u8], String> {
    let len = usize::try_from(varint(b, at)?).map_err(|_| "length out of range")?;
    let end = at.checked_add(len).filter(|&e| e <= b.len()).ok_or("truncated field")?;
    let s = &b[*at..end];
    *at = end;
    Ok(s)
}

fn walk(b: &[u8], msg: Msg, depth: usize) -> Result<Option<String>, String> {
    if depth > MAX_DEPTH {
        return Err(format!("messages nested deeper than {MAX_DEPTH}"));
    }
    let mut at = 0;
    // For a TensorProto: whether it is external, and its `location`.
    let mut external = false;
    let mut location: Option<String> = None;
    while at < b.len() {
        let key = varint(b, &mut at)?;
        let (field, wire) = (key >> 3, key & 7);
        if field == 0 {
            return Err("field number 0".into());
        }
        match wire {
            0 => {
                let v = varint(b, &mut at)?;
                // TensorProto.data_location = 14; EXTERNAL = 1.
                if matches!(msg, Msg::Tensor) && field == 14 && v == 1 {
                    external = true;
                }
            }
            1 => at = at.checked_add(8).filter(|&e| e <= b.len()).ok_or("truncated fixed64")?,
            5 => at = at.checked_add(4).filter(|&e| e <= b.len()).ok_or("truncated fixed32")?,
            2 => {
                let body = delimited(b, &mut at)?;
                if matches!(msg, Msg::Tensor) && field == 13 {
                    // TensorProto.external_data = 13: StringStringEntryProto
                    // { key = 1, value = 2 }. Any entry marks the tensor external.
                    external = true;
                    if let Some(l) = location_entry(body)? {
                        location.get_or_insert(l);
                    }
                } else if let Some(child) = msg.child(field) {
                    if let Some(found) = walk(body, child, depth + 1)? {
                        return Ok(Some(found));
                    }
                }
            }
            _ => return Err(format!("wire type {wire} (groups are not used by onnx.proto)")),
        }
    }
    Ok(external.then(|| location.unwrap_or_default()))
}

/// The value of a `StringStringEntryProto` whose key is `location`.
fn location_entry(b: &[u8]) -> Result<Option<String>, String> {
    let mut at = 0;
    let (mut key, mut value) = (None, None);
    while at < b.len() {
        let k = varint(b, &mut at)?;
        match k & 7 {
            0 => {
                varint(b, &mut at)?;
            }
            1 => at = at.checked_add(8).filter(|&e| e <= b.len()).ok_or("truncated fixed64")?,
            5 => at = at.checked_add(4).filter(|&e| e <= b.len()).ok_or("truncated fixed32")?,
            2 => {
                let s = String::from_utf8_lossy(delimited(b, &mut at)?).into_owned();
                match k >> 3 {
                    1 => key = Some(s),
                    2 => value = Some(s),
                    _ => {}
                }
            }
            w => return Err(format!("wire type {w} in an external_data entry")),
        }
    }
    Ok((key.as_deref() == Some("location")).then(|| value.unwrap_or_default()))
}

#[cfg(test)]
pub(crate) mod tests {
    use super::external_tensor;

    fn varint(mut v: u64, out: &mut Vec<u8>) {
        loop {
            let b = (v & 0x7f) as u8;
            v >>= 7;
            if v == 0 {
                out.push(b);
                return;
            }
            out.push(b | 0x80);
        }
    }

    /// A length-delimited field.
    pub(crate) fn field(n: u64, body: &[u8]) -> Vec<u8> {
        let mut out = Vec::new();
        varint(n << 3 | 2, &mut out);
        varint(body.len() as u64, &mut out);
        out.extend_from_slice(body);
        out
    }

    /// A varint field.
    pub(crate) fn int(n: u64, v: u64) -> Vec<u8> {
        let mut out = Vec::new();
        varint(n << 3, &mut out);
        varint(v, &mut out);
        out
    }

    /// A float tensor of `dims` named `name`, its data inside (`raw_data`).
    pub(crate) fn inline_tensor(name: &str, dims: &[u64]) -> Vec<u8> {
        let n: u64 = dims.iter().product();
        let mut t: Vec<u8> = dims.iter().flat_map(|&d| int(1, d)).collect();
        t.extend(int(2, 1)); // FLOAT
        t.extend(field(8, name.as_bytes()));
        t.extend(field(9, &vec![0u8; 4 * n as usize]));
        t
    }

    /// The same tensor with its data in the file `location`.
    pub(crate) fn external_tensor_proto(name: &str, dims: &[u64], location: &str) -> Vec<u8> {
        let mut t: Vec<u8> = dims.iter().flat_map(|&d| int(1, d)).collect();
        t.extend(int(2, 1));
        t.extend(field(8, name.as_bytes()));
        let mut entry = field(1, b"location");
        entry.extend(field(2, location.as_bytes()));
        t.extend(field(13, &entry));
        t.extend(int(14, 1)); // EXTERNAL
        t
    }

    /// A ModelProto whose graph is one `Add` of input `x` and initializer `w`.
    pub(crate) fn model(initializer: &[u8]) -> Vec<u8> {
        let float_1 = {
            // TypeProto { tensor_type = 1 { elem_type = 1, shape = 2 { dim = 1 { dim_value = 1: 1 } } } }
            let shape = field(1, &int(1, 1));
            let mut tt = int(1, 1);
            tt.extend(field(2, &shape));
            field(1, &tt)
        };
        let value_info = |name: &str| {
            let mut v = field(1, name.as_bytes());
            v.extend(field(2, &float_1));
            v
        };
        let mut node = field(1, b"x");
        node.extend(field(1, b"w"));
        node.extend(field(2, b"y"));
        node.extend(field(4, b"Add"));
        let mut graph = field(1, &node);
        graph.extend(field(2, b"g"));
        graph.extend(field(5, initializer));
        graph.extend(field(11, &value_info("x")));
        graph.extend(field(12, &value_info("y")));
        let mut m = int(1, 8); // ir_version
        m.extend(field(8, &int(2, 17))); // opset_import { version = 17 }
        m.extend(field(7, &graph));
        m
    }

    #[test]
    fn a_tensor_kept_inside_the_model_is_not_external() {
        assert_eq!(external_tensor(&model(&inline_tensor("w", &[1]))), Ok(None));
    }

    #[test]
    fn an_external_initializer_is_found_with_its_location() {
        let m = model(&external_tensor_proto("w", &[1], "weights.bin"));
        assert_eq!(external_tensor(&m), Ok(Some("weights.bin".into())));
    }

    #[test]
    fn an_external_tensor_in_a_subgraph_is_found() {
        // A node whose attribute `then_branch` (g = 6) holds a graph with an
        // external initializer, as an `If` does.
        let mut sub = field(2, b"then");
        sub.extend(field(5, &external_tensor_proto("w", &[1], "../outside.bin")));
        let mut attr = field(1, b"then_branch");
        attr.extend(field(6, &sub));
        let mut node = field(4, b"If");
        node.extend(field(5, &attr));
        let graph = field(1, &node);
        let m = field(7, &graph);
        assert_eq!(external_tensor(&m), Ok(Some("../outside.bin".into())));
    }

    #[test]
    fn external_data_without_a_location_still_marks_the_tensor() {
        let mut t = int(2, 1);
        t.extend(int(14, 1));
        assert_eq!(external_tensor(&model(&t)), Ok(Some(String::new())));
    }

    #[test]
    fn bytes_that_are_not_a_protobuf_message_are_an_error() {
        let m = model(&inline_tensor("w", &[1]));
        assert!(external_tensor(&m[..m.len() - 3]).is_err());
        assert!(external_tensor(&[0x0b]).is_err()); // field 1, wire type 3 (group)
        // Nesting past the bound: attributes' subgraphs inside each other.
        let mut g = Vec::new();
        for _ in 0..40 {
            let attr = field(6, &g);
            let node = field(5, &attr);
            g = field(1, &node);
        }
        assert!(external_tensor(&field(7, &g)).is_err());
    }
}
