unity_asset_binary/typetree/
types.rs

1//! TypeTree data structures
2//!
3//! This module defines the core data structures for Unity TypeTree processing.
4//! TypeTree provides dynamic type information for Unity objects.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// A node in the Unity TypeTree
10///
11/// Each node represents a field or type in the Unity object structure,
12/// forming a tree that describes the complete object layout.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct TypeTreeNode {
15    /// Type name (e.g., "int", "string", "GameObject")
16    pub type_name: String,
17    /// Field name (e.g., "m_Name", "m_IsActive")
18    pub name: String,
19    /// Size in bytes (-1 for variable size)
20    pub byte_size: i32,
21    /// Index in the type tree
22    pub index: i32,
23    /// Type flags
24    pub type_flags: i32,
25    /// Version of this type
26    pub version: i32,
27    /// Meta flags (alignment, etc.)
28    pub meta_flags: i32,
29    /// Depth level in the tree
30    pub level: i32,
31    /// Offset in type string buffer
32    pub type_str_offset: u32,
33    /// Offset in name string buffer
34    pub name_str_offset: u32,
35    /// Reference type hash
36    pub ref_type_hash: u64,
37    /// Child nodes
38    pub children: Vec<TypeTreeNode>,
39}
40
41impl TypeTreeNode {
42    /// Create a new TypeTree node
43    pub fn new() -> Self {
44        Self {
45            type_name: String::new(),
46            name: String::new(),
47            byte_size: 0,
48            index: 0,
49            type_flags: 0,
50            version: 0,
51            meta_flags: 0,
52            level: 0,
53            type_str_offset: 0,
54            name_str_offset: 0,
55            ref_type_hash: 0,
56            children: Vec::new(),
57        }
58    }
59
60    /// Create a new node with basic information
61    pub fn with_info(type_name: String, name: String, byte_size: i32) -> Self {
62        Self {
63            type_name,
64            name,
65            byte_size,
66            ..Default::default()
67        }
68    }
69
70    /// Check if this node represents an array
71    pub fn is_array(&self) -> bool {
72        self.type_name == "Array" || self.type_name.starts_with("vector")
73    }
74
75    /// Check if this node is aligned
76    pub fn is_aligned(&self) -> bool {
77        (self.meta_flags & 0x4000) != 0
78    }
79
80    /// Get the size of this type
81    pub fn size(&self) -> i32 {
82        self.byte_size
83    }
84
85    /// Check if this is a primitive type
86    pub fn is_primitive(&self) -> bool {
87        matches!(
88            self.type_name.as_str(),
89            "bool"
90                | "char"
91                | "SInt8"
92                | "UInt8"
93                | "SInt16"
94                | "UInt16"
95                | "SInt32"
96                | "UInt32"
97                | "SInt64"
98                | "UInt64"
99                | "float"
100                | "double"
101                | "int"
102                | "string"
103        )
104    }
105
106    /// Check if this is a string type
107    pub fn is_string(&self) -> bool {
108        self.type_name == "string"
109    }
110
111    /// Check if this is a numeric type
112    pub fn is_numeric(&self) -> bool {
113        matches!(
114            self.type_name.as_str(),
115            "SInt8"
116                | "UInt8"
117                | "SInt16"
118                | "UInt16"
119                | "SInt32"
120                | "UInt32"
121                | "SInt64"
122                | "UInt64"
123                | "float"
124                | "double"
125                | "int"
126        )
127    }
128
129    /// Check if this is a boolean type
130    pub fn is_boolean(&self) -> bool {
131        self.type_name == "bool"
132    }
133
134    /// Find a child node by name
135    pub fn find_child(&self, name: &str) -> Option<&TypeTreeNode> {
136        self.children.iter().find(|child| child.name == name)
137    }
138
139    /// Find a child node by name (mutable)
140    pub fn find_child_mut(&mut self, name: &str) -> Option<&mut TypeTreeNode> {
141        self.children.iter_mut().find(|child| child.name == name)
142    }
143
144    /// Get all child names
145    pub fn child_names(&self) -> Vec<&str> {
146        self.children
147            .iter()
148            .map(|child| child.name.as_str())
149            .collect()
150    }
151
152    /// Add a child node
153    pub fn add_child(&mut self, child: TypeTreeNode) {
154        self.children.push(child);
155    }
156
157    /// Remove a child node by name
158    pub fn remove_child(&mut self, name: &str) -> Option<TypeTreeNode> {
159        if let Some(pos) = self.children.iter().position(|child| child.name == name) {
160            Some(self.children.remove(pos))
161        } else {
162            None
163        }
164    }
165
166    /// Get the depth of this node in the tree
167    pub fn depth(&self) -> i32 {
168        self.level
169    }
170
171    /// Check if this node has children
172    pub fn has_children(&self) -> bool {
173        !self.children.is_empty()
174    }
175
176    /// Get the number of children
177    pub fn child_count(&self) -> usize {
178        self.children.len()
179    }
180
181    /// Validate the node structure
182    pub fn validate(&self) -> Result<(), String> {
183        if self.type_name.is_empty() {
184            return Err("Type name cannot be empty".to_string());
185        }
186
187        if self.byte_size < -1 {
188            return Err("Invalid byte size".to_string());
189        }
190
191        // Validate children
192        for (i, child) in self.children.iter().enumerate() {
193            child
194                .validate()
195                .map_err(|e| format!("Child {}: {}", i, e))?;
196        }
197
198        Ok(())
199    }
200}
201
202impl Default for TypeTreeNode {
203    fn default() -> Self {
204        Self::new()
205    }
206}
207
208/// Complete TypeTree structure
209///
210/// This structure contains the complete type information for a Unity object,
211/// including all field definitions and their relationships.
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct TypeTree {
214    /// Root nodes of the type tree
215    pub nodes: Vec<TypeTreeNode>,
216    /// String buffer for type and field names
217    pub string_buffer: Vec<u8>,
218    /// Version of the type tree format
219    pub version: u32,
220    /// Platform target
221    pub platform: u32,
222    /// Whether type tree has type dependencies
223    pub has_type_dependencies: bool,
224}
225
226impl TypeTree {
227    /// Create a new empty TypeTree
228    pub fn new() -> Self {
229        Self {
230            nodes: Vec::new(),
231            string_buffer: Vec::new(),
232            version: 0,
233            platform: 0,
234            has_type_dependencies: false,
235        }
236    }
237
238    /// Create a TypeTree with initial capacity
239    pub fn with_capacity(capacity: usize) -> Self {
240        Self {
241            nodes: Vec::with_capacity(capacity),
242            string_buffer: Vec::new(),
243            version: 0,
244            platform: 0,
245            has_type_dependencies: false,
246        }
247    }
248
249    /// Check if the TypeTree is empty
250    pub fn is_empty(&self) -> bool {
251        self.nodes.is_empty()
252    }
253
254    /// Get the number of root nodes
255    pub fn node_count(&self) -> usize {
256        self.nodes.len()
257    }
258
259    /// Add a root node
260    pub fn add_node(&mut self, node: TypeTreeNode) {
261        self.nodes.push(node);
262    }
263
264    /// Find a root node by name
265    pub fn find_node(&self, name: &str) -> Option<&TypeTreeNode> {
266        self.nodes.iter().find(|node| node.name == name)
267    }
268
269    /// Find a root node by name (mutable)
270    pub fn find_node_mut(&mut self, name: &str) -> Option<&mut TypeTreeNode> {
271        self.nodes.iter_mut().find(|node| node.name == name)
272    }
273
274    /// Get all root node names
275    pub fn node_names(&self) -> Vec<&str> {
276        self.nodes.iter().map(|node| node.name.as_str()).collect()
277    }
278
279    /// Return the minimal root-child prefix length needed to read an object name field.
280    ///
281    /// This is inspired by UnityPy's `peek_name()` strategy: we parse only the first N root fields,
282    /// stopping at `m_Name`/`name` instead of reading the full object.
283    pub fn name_peek_prefix(&self) -> Option<(usize, String)> {
284        let root = self.nodes.first()?;
285        for (i, child) in root.children.iter().enumerate() {
286            if child.name == "m_Name" || child.name == "name" {
287                return Some((i + 1, child.name.clone()));
288            }
289        }
290        None
291    }
292
293    /// Clear all nodes
294    pub fn clear(&mut self) {
295        self.nodes.clear();
296        self.string_buffer.clear();
297    }
298
299    /// Get string from buffer at offset
300    pub fn get_string(&self, offset: u32) -> Option<String> {
301        if offset as usize >= self.string_buffer.len() {
302            return None;
303        }
304
305        let start = offset as usize;
306        let end = self.string_buffer[start..]
307            .iter()
308            .position(|&b| b == 0)
309            .map(|pos| start + pos)
310            .unwrap_or(self.string_buffer.len());
311
312        String::from_utf8(self.string_buffer[start..end].to_vec()).ok()
313    }
314
315    /// Add string to buffer and return offset
316    pub fn add_string(&mut self, s: &str) -> u32 {
317        let offset = self.string_buffer.len() as u32;
318        self.string_buffer.extend_from_slice(s.as_bytes());
319        self.string_buffer.push(0); // Null terminator
320        offset
321    }
322
323    /// Validate the entire TypeTree
324    pub fn validate(&self) -> Result<(), String> {
325        if self.nodes.is_empty() {
326            return Err("TypeTree has no nodes".to_string());
327        }
328
329        for (i, node) in self.nodes.iter().enumerate() {
330            node.validate()
331                .map_err(|e| format!("Root node {}: {}", i, e))?;
332        }
333
334        Ok(())
335    }
336
337    /// Get TypeTree statistics
338    pub fn statistics(&self) -> TypeTreeStatistics {
339        let mut stats = (0usize, 0i32, 0usize, 0usize); // (total_nodes, max_depth, primitive_count, array_count)
340
341        fn count_nodes(node: &TypeTreeNode, depth: i32, stats: &mut (usize, i32, usize, usize)) {
342            stats.0 += 1; // total_nodes
343            stats.1 = stats.1.max(depth); // max_depth
344
345            if node.is_primitive() {
346                stats.2 += 1; // primitive_count
347            }
348            if node.is_array() {
349                stats.3 += 1; // array_count
350            }
351
352            for child in &node.children {
353                count_nodes(child, depth + 1, stats);
354            }
355        }
356
357        for node in &self.nodes {
358            count_nodes(node, 0, &mut stats);
359        }
360
361        TypeTreeStatistics {
362            total_nodes: stats.0,
363            root_nodes: self.nodes.len(),
364            max_depth: stats.1,
365            primitive_count: stats.2,
366            array_count: stats.3,
367            string_buffer_size: self.string_buffer.len(),
368        }
369    }
370}
371
372impl Default for TypeTree {
373    fn default() -> Self {
374        Self::new()
375    }
376}
377
378/// TypeTree statistics
379#[derive(Debug, Clone, Serialize, Deserialize)]
380pub struct TypeTreeStatistics {
381    pub total_nodes: usize,
382    pub root_nodes: usize,
383    pub max_depth: i32,
384    pub primitive_count: usize,
385    pub array_count: usize,
386    pub string_buffer_size: usize,
387}
388
389/// Type information for Unity classes
390#[derive(Debug, Clone, Serialize, Deserialize)]
391pub struct TypeInfo {
392    pub class_id: i32,
393    pub class_name: String,
394    pub type_tree: TypeTree,
395    pub script_type_index: Option<i16>,
396    pub script_id: [u8; 16],
397    pub old_type_hash: [u8; 16],
398}
399
400impl TypeInfo {
401    /// Create new type info
402    pub fn new(class_id: i32, class_name: String) -> Self {
403        Self {
404            class_id,
405            class_name,
406            type_tree: TypeTree::new(),
407            script_type_index: None,
408            script_id: [0; 16],
409            old_type_hash: [0; 16],
410        }
411    }
412
413    /// Check if this is a script type
414    pub fn is_script_type(&self) -> bool {
415        self.script_type_index.is_some()
416    }
417}
418
419/// Type registry for managing multiple types
420#[derive(Debug, Clone, Default)]
421pub struct TypeRegistry {
422    types: HashMap<i32, TypeInfo>,
423}
424
425impl TypeRegistry {
426    /// Create a new type registry
427    pub fn new() -> Self {
428        Self {
429            types: HashMap::new(),
430        }
431    }
432
433    /// Add a type to the registry
434    pub fn add_type(&mut self, type_info: TypeInfo) {
435        self.types.insert(type_info.class_id, type_info);
436    }
437
438    /// Get a type by class ID
439    pub fn get_type(&self, class_id: i32) -> Option<&TypeInfo> {
440        self.types.get(&class_id)
441    }
442
443    /// Get all registered class IDs
444    pub fn class_ids(&self) -> Vec<i32> {
445        self.types.keys().copied().collect()
446    }
447
448    /// Check if a class ID is registered
449    pub fn has_type(&self, class_id: i32) -> bool {
450        self.types.contains_key(&class_id)
451    }
452
453    /// Clear all types
454    pub fn clear(&mut self) {
455        self.types.clear();
456    }
457
458    /// Get the number of registered types
459    pub fn len(&self) -> usize {
460        self.types.len()
461    }
462
463    /// Check if the registry is empty
464    pub fn is_empty(&self) -> bool {
465        self.types.is_empty()
466    }
467}