unity_typetree_gen/assembly.rs
1// TODO(ai-review): review for style and correctness
2//! Owning generator that serves MonoBehaviour type trees by
3//! `(assembly, namespace, class)`, so callers don't deal with dotnetdll or its
4//! borrow lifetimes.
5//!
6//! Assemblies are resolved lazily via a [`Loader`] passed to [`generate`] and
7//! parsed only on first use — a game's `Managed` directory has dozens of DLLs
8//! (BCL, engine modules) but only the few that actually define or are
9//! referenced by a MonoBehaviour are ever loaded.
10//!
11//! [`generate`]: AssemblyTypeTreeGenerator::generate
12use std::collections::{BTreeMap, HashMap};
13use std::path::Path;
14use std::sync::{Arc, Mutex};
15
16use dotnetdll::prelude::{ReadOptions, Resolution};
17
18use rabex::UnityVersion;
19
20use crate::TypeTreeNode;
21use crate::generator::Generator;
22
23/// Resolves an assembly's bytes by name
24///
25/// Usually given assembly name like `Assembly-CSharp.dll` it ould read the bytes
26/// from `game_Data/Managed/Assembly-CSharp.dll`.
27/// Can return `std::io::ErrorKind::NotFound` to signal absence.
28pub type Loader<'a> = dyn Fn(&str) -> Result<Vec<u8>, std::io::Error> + 'a;
29
30/// A parsed-assembly cache that generates MonoBehaviour type trees for a fixed
31/// Unity version.
32///
33/// Assembly bytes and parsed resolutions are leaked to obtain `'static`
34/// lifetimes; this is meant for a generator built once per game environment.
35///
36/// # Example
37///
38/// ```no_run
39/// # use std::path::Path;
40/// # use unity_typetree_gen::AssemblyTypeTreeGenerator;
41/// let version = "6000.0.0f1".parse().unwrap();
42/// let generator = AssemblyTypeTreeGenerator::new(version);
43///
44/// let managed_dir = Path::new("/path/to/Game_Data/Managed");
45/// let tree = generator
46/// .generate_from_dir(managed_dir, "Assembly-CSharp.dll", "MyGame", "PlayerController")?
47/// .expect("assembly or type not found");
48///
49/// println!("{} fields", tree.children.len());
50/// # Ok::<(), std::io::Error>(())
51/// ```
52pub struct AssemblyTypeTreeGenerator {
53 /// Assembly bytes, loaded lazily and cached (leaked) on first use.
54 bytes: Mutex<HashMap<String, &'static [u8]>>,
55 resolutions: Mutex<HashMap<String, &'static Resolution<'static>>>,
56 /// Per-assembly locks for single-flight parsing: concurrent callers for the same
57 /// assembly wait for the first to parse instead of all parsing it in parallel.
58 locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
59 unity_version: UnityVersion,
60}
61
62impl AssemblyTypeTreeGenerator {
63 pub fn new(unity_version: UnityVersion) -> Self {
64 AssemblyTypeTreeGenerator {
65 bytes: Mutex::new(HashMap::new()),
66 resolutions: Mutex::new(HashMap::new()),
67 locks: Mutex::new(HashMap::new()),
68 unity_version,
69 }
70 }
71
72 /// Bytes for `assembly_name`: cached first, otherwise via `loader` (whose
73 /// result is leaked and cached so it runs at most once per assembly).
74 /// `Ok(None)` if the loader reports the assembly is absent (`NotFound`).
75 fn assembly_bytes(
76 &self,
77 assembly_name: &str,
78 loader: &Loader,
79 ) -> Result<Option<&'static [u8]>, std::io::Error> {
80 let mut bytes = self.bytes.lock().unwrap();
81 if let Some(b) = bytes.get(assembly_name) {
82 return Ok(Some(b));
83 }
84 let loaded = match loader(assembly_name) {
85 Ok(loaded) => loaded,
86 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
87 Err(e) => return Err(e),
88 };
89 let leaked: &'static [u8] = Vec::leak(loaded);
90 bytes.insert(assembly_name.to_owned(), leaked);
91 Ok(Some(leaked))
92 }
93
94 /// Parsed resolution for `assembly_name`, parsing (and caching) it on first
95 /// access. `Ok(None)` if the assembly is absent or fails to parse.
96 pub(crate) fn resolution(
97 &self,
98 assembly_name: &str,
99 loader: &Loader,
100 ) -> Result<Option<&'static Resolution<'static>>, std::io::Error> {
101 if let Some(resolution) = self.resolutions.lock().unwrap().get(assembly_name) {
102 return Ok(Some(resolution));
103 }
104 // Single-flight: only one thread parses a given assembly. Concurrent callers for the
105 // same assembly block on this per-assembly lock, then hit the cache on the re-check
106 // below — avoiding a parse stampede when a parallel scan first touches an assembly.
107 let assembly_lock = {
108 let mut locks = self.locks.lock().unwrap();
109 Arc::clone(locks.entry(assembly_name.to_owned()).or_default())
110 };
111 let _flight = assembly_lock.lock().unwrap();
112 if let Some(resolution) = self.resolutions.lock().unwrap().get(assembly_name) {
113 return Ok(Some(resolution));
114 }
115 let Some(bytes) = self.assembly_bytes(assembly_name, loader)? else {
116 return Ok(None);
117 };
118 // Type trees only need type/field metadata, never method IL — skipping
119 // method bodies avoids the expensive part of parsing large assemblies.
120 #[allow(clippy::needless_update)] // improves patchability
121 let options = ReadOptions {
122 skip_method_bodies: true,
123 lazy_method_bodies: false,
124 lazy_method_signatures: true,
125 ..Default::default()
126 };
127 let Ok(resolution) = Resolution::parse(bytes, options) else {
128 return Ok(None);
129 };
130 let leaked: &'static Resolution<'static> = Box::leak(Box::new(resolution));
131 self.resolutions
132 .lock()
133 .unwrap()
134 .insert(assembly_name.to_owned(), leaked);
135 Ok(Some(leaked))
136 }
137
138 /// Generate the type tree for `namespace.type_name` defined in
139 /// `assembly_name`. The assembly, as well as possible references assemblies, are resolved on demand through `loader`.
140 ///
141 /// `Ok(None)` if the type can't be resolved (assembly or type absent);
142 /// `Err` if the loader itself fails (other than `NotFound`).
143 pub fn generate(
144 &self,
145 loader: &Loader,
146 assembly_name: &str,
147 namespace: &str,
148 type_name: &str,
149 ) -> Result<Option<TypeTreeNode>, std::io::Error> {
150 let Some(primary) = self.resolution(assembly_name, loader)? else {
151 return Ok(None);
152 };
153 let children = Generator::new(self, &self.unity_version, loader)
154 .read(primary, namespace, type_name)?;
155 Ok(children.map(|children| crate::assemble(children, type_name)))
156 }
157
158 /// Convenience: generate using a loader that reads `<managed_dir>/<assembly_name>`.
159 pub fn generate_from_dir(
160 &self,
161 managed_dir: &Path,
162 assembly_name: &str,
163 namespace: &str,
164 type_name: &str,
165 ) -> Result<Option<TypeTreeNode>, std::io::Error> {
166 self.generate(
167 &|name| std::fs::read(managed_dir.join(name)),
168 assembly_name,
169 namespace,
170 type_name,
171 )
172 }
173
174 /// Pre-load an assembly by name so it is available to [`monobehaviour_definitions`](Self::monobehaviour_definitions).
175 /// Returns `true` if the assembly was found and loaded (or was already cached),
176 /// `false` if the loader reports it is absent.
177 pub fn load_assembly(&self, loader: &Loader, name: &str) -> Result<bool, std::io::Error> {
178 self.resolution(name, loader).map(|r| r.is_some())
179 }
180
181 /// Returns a map from assembly name to the list of full type names for every
182 /// type in the currently-loaded assemblies that derives (directly or
183 /// transitively) from `UnityEngine.MonoBehaviour`.
184 pub fn monobehaviour_definitions(
185 &self,
186 loader: &Loader,
187 ) -> Result<BTreeMap<String, Vec<String>>, std::io::Error> {
188 let g = crate::generator::Generator::new(self, &self.unity_version, loader);
189 let resolutions: Vec<(String, &'static Resolution<'static>)> = self
190 .resolutions
191 .lock()
192 .unwrap()
193 .iter()
194 .map(|(k, v)| (k.clone(), *v))
195 .collect();
196 let mut defs: BTreeMap<String, Vec<String>> = BTreeMap::new();
197 for (asm_name, res) in &resolutions {
198 for td in &res.type_definitions {
199 // Mono.Cecil's Types (used by the reference C# impl) only yields
200 // top-level types; skip nested types to match that behaviour.
201 if td.encloser.is_some() {
202 continue;
203 }
204 if g.derives_from_monobehaviour(res, td)? {
205 defs.entry(asm_name.clone())
206 .or_default()
207 .push(td.type_name());
208 }
209 }
210 }
211 Ok(defs)
212 }
213}