unity_formats/
typetree.rs

1//! Type trees: a class's serialized layout, and decoding an object's bytes by it.
2//!
3//! The rules: primitives by their type name; `string` is a 32-bit length, the UTF-8 bytes, then 4-byte alignment;
4//! `TypelessData` a 32-bit length and the bytes; `pair` its two children; a node whose first child is `Array` is a
5//! vector (a 32-bit count, then the elements, each by the array's second child); anything else is a struct of its
6//! children in order. After any node whose meta flag has `ALIGN`, the position is aligned to 4 bytes (from the
7//! object's start); an `Array` child's flag aligns its vector.
8
9use std::sync::Arc;
10
11use crate::reader::Reader;
12use crate::value::{Array, Value};
13use crate::{Error, Result, invalid};
14
15/// The meta flag that aligns the position to 4 bytes after the node.
16pub const ALIGN: u32 = 0x4000;
17
18/// One node: its type and field name, its byte size (-1: variable), version, meta flags and children (indices into
19/// the tree's node list).
20#[derive(Clone, Debug, PartialEq)]
21pub struct Node {
22    /// The node's type name (`int`, `string`, `vector`, a class name, ...).
23    pub ty: Arc<str>,
24    /// The field name.
25    pub name: Arc<str>,
26    /// The size in bytes the tree states for a value of this node, or -1 where it varies.
27    pub byte_size: i32,
28    /// The node's version, as stored.
29    pub version: i32,
30    /// The meta flags: [`ALIGN`] and others the decoder does not use.
31    pub meta_flag: u32,
32    /// Indices into [`TypeTree::nodes`], in field order; every child's index is greater than its parent's.
33    pub children: Vec<usize>,
34}
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37enum Prim {
38    Bool,
39    I8,
40    U8,
41    I16,
42    U16,
43    I32,
44    U32,
45    I64,
46    U64,
47    F32,
48    F64,
49}
50
51impl Prim {
52    fn of(ty: &str) -> Option<Prim> {
53        Some(match ty {
54            "bool" => Prim::Bool,
55            "SInt8" => Prim::I8,
56            "UInt8" | "char" => Prim::U8,
57            "short" | "SInt16" => Prim::I16,
58            "unsigned short" | "UInt16" => Prim::U16,
59            "int" | "SInt32" => Prim::I32,
60            "unsigned int" | "UInt32" | "Type*" => Prim::U32,
61            "long long" | "SInt64" => Prim::I64,
62            "unsigned long long" | "UInt64" | "FileSize" => Prim::U64,
63            "float" => Prim::F32,
64            "double" => Prim::F64,
65            _ => return None,
66        })
67    }
68
69    fn size(self) -> usize {
70        match self {
71            Prim::Bool | Prim::I8 | Prim::U8 => 1,
72            Prim::I16 | Prim::U16 => 2,
73            Prim::I32 | Prim::U32 | Prim::F32 => 4,
74            Prim::I64 | Prim::U64 | Prim::F64 => 8,
75        }
76    }
77}
78
79#[derive(Clone, Debug, PartialEq)]
80enum Kind {
81    Prim(Prim),
82    String,
83    Typeless,
84    Pair,
85    /// The element node; whether the `Array` child aligns.
86    Vector {
87        element: usize,
88        aligned: bool,
89    },
90    /// Managed references (`[SerializeReference]`), which need the file's own reference type trees.
91    Referenced,
92    Struct,
93}
94
95/// A class's type tree; node 0 is the root.
96#[derive(Clone, Debug, PartialEq)]
97pub struct TypeTree {
98    /// The nodes, the root first.
99    pub nodes: Vec<Node>,
100    kinds: Vec<Kind>,
101    /// Per node, the fewest bytes a value of it takes (alignment left out): what bounds a vector's count by the bytes
102    /// left.
103    min_sizes: Vec<usize>,
104}
105
106/// The most elements a vector of a type that takes no bytes (a struct without fields) may state: its elements cost
107/// nothing to the bytes, so the bytes cannot bound its count.
108const MAX_EMPTY_ELEMENTS: usize = 1 << 16;
109
110impl TypeTree {
111    /// Build a tree from its nodes (the root first). Refused: no nodes, a child index that is not after its parent's
112    /// or not in the list, a `pair` without two children, or a vector whose `Array` child does not have two children.
113    pub fn new(nodes: Vec<Node>) -> Result<Self> {
114        if nodes.is_empty() {
115            return Err(invalid("empty type tree"));
116        }
117        let mut kinds = Vec::with_capacity(nodes.len());
118        for (i, n) in nodes.iter().enumerate() {
119            if n.children.iter().any(|&c| c <= i || c >= nodes.len()) {
120                return Err(invalid(format!("type tree node {i} has a bad child index")));
121            }
122            let first = n.children.first().and_then(|&c| nodes.get(c));
123            let kind = if let Some(p) = Prim::of(&n.ty) {
124                Kind::Prim(p)
125            } else if &*n.ty == "string" {
126                Kind::String
127            } else if &*n.ty == "TypelessData" {
128                Kind::Typeless
129            } else if &*n.ty == "pair" {
130                if n.children.len() != 2 {
131                    return Err(invalid(format!(
132                        "pair node {i} has {} children",
133                        n.children.len()
134                    )));
135                }
136                Kind::Pair
137            } else if &*n.ty == "ReferencedObject" {
138                Kind::Referenced
139            } else if let Some(a) = first.filter(|f| &*f.ty == "Array") {
140                let &[_, element] = a.children.as_slice() else {
141                    return Err(invalid(format!(
142                        "Array node {} has {} children",
143                        n.children.first().copied().unwrap_or_default(),
144                        a.children.len()
145                    )));
146                };
147                Kind::Vector {
148                    element,
149                    aligned: a.meta_flag & ALIGN != 0,
150                }
151            } else {
152                Kind::Struct
153            };
154            kinds.push(kind);
155        }
156        // children come after their parent, so the sizes are summed from the last node up
157        let mut min_sizes = vec![0usize; nodes.len()];
158        for (i, (node, kind)) in nodes.iter().zip(&kinds).enumerate().rev() {
159            let size = match kind {
160                Kind::Prim(p) => p.size(),
161                Kind::String | Kind::Typeless | Kind::Vector { .. } => 4,
162                Kind::Pair | Kind::Struct => node
163                    .children
164                    .iter()
165                    .map(|&c| min_sizes.get(c).copied().unwrap_or(0))
166                    .fold(0, usize::saturating_add),
167                Kind::Referenced => 0,
168            };
169            if let Some(slot) = min_sizes.get_mut(i) {
170                *slot = size;
171            }
172        }
173        Ok(TypeTree {
174            nodes,
175            kinds,
176            min_sizes,
177        })
178    }
179
180    /// The root node; `None` only when [`TypeTree::nodes`] was emptied after [`TypeTree::new`] refused an empty tree.
181    pub fn root(&self) -> Option<&Node> {
182        self.nodes.first()
183    }
184
185    /// Node `i` and its kind: an error when the tree no longer holds the node (`nodes` is public, `new` checked it).
186    fn node(&self, i: usize) -> Result<(&Node, &Kind)> {
187        self.nodes
188            .get(i)
189            .zip(self.kinds.get(i))
190            .ok_or_else(|| invalid(format!("type tree node {i} of {}", self.nodes.len())))
191    }
192
193    /// Decode `data` (an object's bytes) from its start; returns the value and the bytes consumed.
194    pub fn decode(&self, data: &[u8], big_endian: bool) -> Result<(Value, usize)> {
195        let mut r = Reader::new(data, big_endian);
196        let v = self.read(0, &mut r, 0)?;
197        Ok((v, r.pos))
198    }
199
200    /// The root struct's field `name` alone: the fields before it are walked (and so must decode), the ones after it
201    /// are not read, so a name stored first costs a few bytes whatever the object holds. `None` when the root has no
202    /// such field; the value is what [`TypeTree::decode`] gives it. A root that is not a struct has no fields.
203    pub fn decode_field(&self, data: &[u8], big_endian: bool, name: &str) -> Result<Option<Value>> {
204        let (root, kind) = self.node(0)?;
205        if !matches!(kind, Kind::Struct) {
206            return Ok(None);
207        }
208        let mut r = Reader::new(data, big_endian);
209        for &c in &root.children {
210            let v = self.read(c, &mut r, 1)?;
211            if &*self.node(c)?.0.name == name {
212                return Ok(Some(v));
213            }
214        }
215        Ok(None)
216    }
217
218    fn read(&self, i: usize, r: &mut Reader, depth: u32) -> Result<Value> {
219        if depth > 128 {
220            return Err(invalid("type tree recursion deeper than 128"));
221        }
222        let (n, kind) = self.node(i)?;
223        let v = match kind {
224            Kind::Prim(p) => prim(*p, r)?,
225            Kind::String => {
226                let len = length(r, 1)?;
227                let b = r.take(len)?;
228                let s = std::str::from_utf8(b)
229                    .map_err(|_| invalid(format!("{}: string is not UTF-8", n.name)))?
230                    .to_owned();
231                r.align(4)?;
232                Value::String(s)
233            }
234            Kind::Typeless => {
235                let len = length(r, 1)?;
236                Value::Array(Array::U8(r.take(len)?.to_vec()))
237            }
238            Kind::Pair => {
239                let &[first, second] = n.children.as_slice() else {
240                    return Err(invalid(format!(
241                        "pair node {i} has {} children",
242                        n.children.len()
243                    )));
244                };
245                let a = self.read(first, r, depth.saturating_add(1))?;
246                let b = self.read(second, r, depth.saturating_add(1))?;
247                Value::Pair(Box::new((a, b)))
248            }
249            Kind::Vector { element, aligned } => {
250                let v = self.vector(*element, r, depth)?;
251                if *aligned {
252                    r.align(4)?;
253                }
254                v
255            }
256            Kind::Referenced => {
257                return Err(Error::Unsupported(format!(
258                    "{}: managed reference (needs the file's reference type trees)",
259                    n.name
260                )));
261            }
262            Kind::Struct => {
263                let mut fields = Vec::with_capacity(n.children.len());
264                for &c in &n.children {
265                    fields.push((
266                        self.node(c)?.0.name.clone(),
267                        self.read(c, r, depth.saturating_add(1))?,
268                    ));
269                }
270                Value::Struct(fields)
271            }
272        };
273        if n.meta_flag & ALIGN != 0 {
274            r.align(4)?;
275        }
276        Ok(v)
277    }
278
279    fn vector(&self, e: usize, r: &mut Reader, depth: u32) -> Result<Value> {
280        let (en, kind) = self.node(e)?;
281        // a primitive element without its own alignment: the elements packed
282        if let Kind::Prim(p) = *kind
283            && en.meta_flag & ALIGN == 0
284        {
285            let n = length(r, p.size())?;
286            return Ok(Value::Array(prim_array(p, n, r)?));
287        }
288        // the shortest an element can be, to bound the count by the bytes left
289        let min = self.min_sizes.get(e).copied().unwrap_or(0);
290        let n = length(r, min)?;
291        if min == 0 && n > MAX_EMPTY_ELEMENTS {
292            return Err(invalid(format!(
293                "a vector of {n} elements that take no bytes (at most {MAX_EMPTY_ELEMENTS})"
294            )));
295        }
296        let mut out = Vec::with_capacity(n.min(1 << 16));
297        for _ in 0..n {
298            out.push(self.read(e, r, depth.saturating_add(1))?);
299        }
300        Ok(Value::Array(Array::Values(out)))
301    }
302}
303
304/// A 32-bit length or count, checked against the bytes left at `unit` bytes per element.
305fn length(r: &mut Reader, unit: usize) -> Result<usize> {
306    let at = r.pos;
307    let n = r.i32()?;
308    let n = usize::try_from(n).map_err(|_| invalid(format!("negative length {n} at offset {at}")))?;
309    if n.saturating_mul(unit) > r.remaining() {
310        return Err(invalid(format!(
311            "length {n} at offset {at} runs past the object ({} bytes left)",
312            r.remaining()
313        )));
314    }
315    Ok(n)
316}
317
318fn prim(p: Prim, r: &mut Reader) -> Result<Value> {
319    Ok(match p {
320        Prim::Bool => Value::Bool(r.bool()?),
321        Prim::I8 => Value::Int(i64::from(r.i8()?)),
322        Prim::U8 => Value::UInt(u64::from(r.u8()?)),
323        Prim::I16 => Value::Int(i64::from(r.i16()?)),
324        Prim::U16 => Value::UInt(u64::from(r.u16()?)),
325        Prim::I32 => Value::Int(i64::from(r.i32()?)),
326        Prim::U32 => Value::UInt(u64::from(r.u32()?)),
327        Prim::I64 => Value::Int(r.i64()?),
328        Prim::U64 => Value::UInt(r.u64()?),
329        Prim::F32 => Value::Float(r.f32()?),
330        Prim::F64 => Value::Double(r.f64()?),
331    })
332}
333
334fn prim_array(p: Prim, n: usize, r: &mut Reader) -> Result<Array> {
335    macro_rules! many {
336        ($f:ident, $v:ident) => {
337            Array::$v((0..n).map(|_| r.$f()).collect::<Result<Vec<_>>>()?)
338        };
339    }
340    Ok(match p {
341        Prim::U8 => Array::U8(r.take(n)?.to_vec()),
342        Prim::Bool => many!(bool, Bool),
343        Prim::I8 => many!(i8, I8),
344        Prim::I16 => many!(i16, I16),
345        Prim::U16 => many!(u16, U16),
346        Prim::I32 => many!(i32, I32),
347        Prim::U32 => many!(u32, U32),
348        Prim::I64 => many!(i64, I64),
349        Prim::U64 => many!(u64, U64),
350        Prim::F32 => many!(f32, F32),
351        Prim::F64 => many!(f64, F64),
352    })
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    fn node(ty: &str, name: &str, meta: u32, children: Vec<usize>) -> Node {
360        Node {
361            ty: ty.into(),
362            name: name.into(),
363            byte_size: -1,
364            version: 1,
365            meta_flag: meta,
366            children,
367        }
368    }
369
370    /// A vector's count is bounded by the bytes left at the fewest bytes an element takes: a struct of fields takes the
371    /// sum of theirs, and elements that take none (a struct without fields) are capped by a constant.
372    #[test]
373    fn a_vectors_count_is_bounded_by_what_its_elements_take() {
374        let vector_of = |element: Vec<Node>| {
375            let mut nodes = vec![
376                node("Thing", "Base", 0, vec![1]),
377                node("vector", "m_V", 0, vec![2]),
378                node("Array", "Array", 0, vec![3, 4]),
379                node("int", "size", 0, vec![]),
380            ];
381            let first = nodes.len();
382            for (k, mut n) in element.into_iter().enumerate() {
383                n.children = n.children.iter().map(|c| c + first).collect();
384                if k == 0 {
385                    n.name = "data".into();
386                }
387                nodes.push(n);
388            }
389            TypeTree::new(nodes).unwrap()
390        };
391        let count = |n: i32| n.to_le_bytes();
392        // elements of a struct of two ints take 8 bytes: a count the bytes cannot hold is refused before anything is
393        // made of it
394        let pairs = vector_of(vec![
395            node("Two", "data", 0, vec![1, 2]),
396            node("int", "a", 0, vec![]),
397            node("int", "b", 0, vec![]),
398        ]);
399        let e = pairs
400            .decode(&count(1_000_000_000), false)
401            .unwrap_err()
402            .to_string();
403        assert!(e.contains("runs past the object"), "{e}");
404        let mut two = count(2).to_vec();
405        two.extend([0u8; 16]);
406        assert!(pairs.decode(&two, false).is_ok());
407        // elements of no bytes: some are fine, billions are not
408        let empty = vector_of(vec![node("Empty", "data", 0, vec![])]);
409        assert!(empty.decode(&count(3), false).is_ok());
410        let e = empty.decode(&count(i32::MAX), false).unwrap_err().to_string();
411        assert!(e.contains("take no bytes"), "{e}");
412    }
413
414    /// A struct of a bool, an aligned string, a vector of floats and a pair, decoded and aligned as the rules say.
415    #[test]
416    fn decodes_by_the_rules() {
417        let t = TypeTree::new(vec![
418            node("Thing", "Base", 0, vec![1, 2, 6, 10]),
419            node("bool", "m_On", ALIGN, vec![]),
420            node("string", "m_Name", 0, vec![3]),
421            node("Array", "Array", ALIGN, vec![4, 5]),
422            node("int", "size", 0, vec![]),
423            node("char", "data", 0, vec![]),
424            node("vector", "m_V", 0, vec![7]),
425            node("Array", "Array", 0, vec![8, 9]),
426            node("int", "size", 0, vec![]),
427            node("float", "data", 0, vec![]),
428            node("pair", "m_P", 0, vec![11, 12]),
429            node("int", "first", 0, vec![]),
430            node("UInt8", "second", ALIGN, vec![]),
431        ])
432        .unwrap();
433        let mut d = vec![1, 0, 0, 0];
434        d.extend_from_slice(&3i32.to_le_bytes());
435        d.extend_from_slice(b"abc\0");
436        d.extend_from_slice(&2i32.to_le_bytes());
437        d.extend_from_slice(&1.5f32.to_le_bytes());
438        d.extend_from_slice(&(-2.0f32).to_le_bytes());
439        d.extend_from_slice(&7i32.to_le_bytes());
440        d.extend_from_slice(&[9, 0, 0, 0]);
441        let (v, used) = t.decode(&d, false).unwrap();
442        assert_eq!(used, d.len());
443        assert_eq!(v.get("m_On"), Some(&Value::Bool(true)));
444        assert_eq!(v.get("m_Name").and_then(Value::as_str), Some("abc"));
445        assert_eq!(v.get("m_V"), Some(&Value::Array(Array::F32(vec![1.5, -2.0]))));
446        assert_eq!(
447            v.get("m_P"),
448            Some(&Value::Pair(Box::new((Value::Int(7), Value::UInt(9)))))
449        );
450        // one field alone: what decode gives it, the fields after it unread (a broken tail does not matter), none for a
451        // name the root lacks
452        for name in ["m_On", "m_Name", "m_V", "m_P"] {
453            assert_eq!(
454                t.decode_field(&d, false, name).unwrap().as_ref(),
455                v.get(name),
456                "{name}"
457            );
458        }
459        assert_eq!(t.decode_field(&d, false, "m_Missing").unwrap(), None);
460        assert_eq!(
461            t.decode_field(&d[..12], false, "m_Name")
462                .unwrap()
463                .as_ref()
464                .and_then(Value::as_str),
465            Some("abc"),
466            "the name decodes from the bytes up to it"
467        );
468        assert!(t.decode_field(&d[..10], false, "m_Name").is_err());
469        // a count past the data and a bad bool are errors
470        let mut short = d.clone();
471        short[12..16].copy_from_slice(&1000i32.to_le_bytes());
472        assert!(t.decode(&short, false).is_err());
473        let mut b = d.clone();
474        b[0] = 2;
475        assert!(t.decode(&b, false).is_err());
476    }
477}
unity_formats/typetree.rs

unity_formats/
typetree.rs
