unrust_codegen/
codegen.rs

1use anyhow::Context;
2use anyhow::Result;
3use genco::fmt;
4use genco::prelude::*;
5use std::ffi::OsStr;
6use std::fs::File;
7use std::path::Path;
8
9pub fn generate_csharp(path: &str, base_folder: &str) -> Result<()> {
10    clear_contents(base_folder, "cs")?;
11
12    let contents = std::fs::read_to_string(path)?;
13    let ast = syn::parse_file(&contents)?;
14
15    let custom_comps = generate_components_csharp(ast.clone(), base_folder)?;
16    let custom_states = generate_states_csharp(ast.clone(), base_folder)?;
17    generate_prefabs_csharp(ast.clone(), base_folder)?;
18
19    generate_hooks(base_folder, custom_comps, custom_states)?;
20
21    Ok(())
22}
23
24fn generate_hooks(
25    base_folder: &str,
26    custom_comps: csharp::Tokens,
27    custom_states: csharp::Tokens,
28) -> Result<()> {
29    let runtime_initialize = &csharp::import("UnityEngine", "RuntimeInitializeOnLoadMethod");
30    let runtime_initialize_load = &csharp::import("UnityEngine", "RuntimeInitializeLoadType");
31    let unrust_native = &csharp::import("unrust.runtime", "NativeWrapper");
32    let entity_manager = &csharp::import("Unity.Entities", "EntityManager");
33    let entity = &csharp::import("Unity.Entities", "Entity");
34    let create_callback = &csharp::import("unrust.runtime", "CustomCreateCallback");
35
36    let hooks: csharp::Tokens = quote! {
37        namespace unrust.userland
38        {
39            public static class UnrustHooks
40            {
41                [$runtime_initialize($runtime_initialize_load.BeforeSceneLoad)]
42                static void Initialize()
43                {
44                    $unrust_native.CustomCreates = CreateCallback;
45                }
46
47                public static unsafe ulong CreateCallback($entity_manager manager, $entity entity, $create_callback cb)
48                {
49                    var count = 0;
50                    var arr = new CustomData[CustomComponents.ComponentCount];
51
52                    $(custom_comps)
53
54                    var stateCount = 0;
55                    var stateArr = new CustomState[CustomState.CustomStateCount];
56
57                    $(custom_states)
58
59                    fixed (void* ptr = arr)
60                    {
61                        fixed (void* state = stateArr)
62                        {
63                            return cb(ptr, (nuint)count, state, (nuint)stateCount);
64                        }
65                    }
66                }
67            }
68        }
69    };
70
71    write_tokens_to_file(base_folder, "UnrustHooks.cs", hooks)
72}
73
74fn generate_prefabs_csharp(ast: syn::File, base_folder: &str) -> Result<()> {
75    let monobehaviour = &csharp::import("UnityEngine", "MonoBehaviour");
76    let baker = &csharp::import("Unity.Entities", "Baker");
77    let spawnable = &csharp::import("unrust.runtime", "UnrustSpawnable");
78
79    let enums = ast
80        .items
81        .iter()
82        .filter_map(|item| match item {
83            syn::Item::Enum(s) => Some(s),
84            _ => None,
85        })
86        .filter(|item| {
87            item.attrs
88                .iter()
89                .any(|attr| attr.path().is_ident("unity_prefab"))
90        })
91        .map(|item| {
92            (
93                item.ident.to_string(),
94                item.variants.iter().map(|v| v.ident.to_string()),
95            )
96        })
97        .enumerate()
98        .map(|(index, (enum_name, variants))| {
99            let variant_fields = variants.clone().map(|name| {
100                quote! {
101                    $['\r']public GameObject $(&name);
102                }
103            });
104
105            let count = index;
106
107            let author_fields = variants.map(|name| {
108                quote! {
109                    buffer.Add(GetEntity(authoring.$(&name), TransformUsageFlags.Dynamic));
110                }
111            });
112
113            let comp: csharp::Tokens = quote! {
114                namespace unrust.userland
115                {
116                    public class $(&enum_name)Authoring : $monobehaviour {
117                        $(for n in variant_fields => $n)
118
119                        public const int RESOURCE_ID = $count;
120
121                        class Baker : $baker<$(&enum_name)Authoring>
122                        {
123                            public override void Bake($(&enum_name)Authoring authoring)
124                            {
125                                var containerEntity = GetEntity(TransformUsageFlags.None);
126                                var buffer = AddBuffer<$spawnable>(containerEntity).Reinterpret<Entity>();
127
128                                $(for n in author_fields => $n)
129
130                                AddComponent<UnrustResourceID>(containerEntity, new UnrustResourceID { Value =  $count});
131                            }
132                        }
133                    }
134                }
135            };
136
137            (enum_name, comp)
138        });
139
140    enums.clone().try_for_each(|(name, comp)| {
141        write_tokens_to_file(base_folder, &format!("{}Authoring.cs", name), comp)
142    })?;
143
144    Ok(())
145}
146
147fn generate_states_csharp(ast: syn::File, base_folder: &str) -> Result<csharp::Tokens> {
148    let struct_layout = &csharp::import("System.Runtime.InteropServices", "StructLayout");
149    let layout_kind = &csharp::import("System.Runtime.InteropServices", "LayoutKind");
150    let monobehaviour = &csharp::import("UnityEngine", "MonoBehaviour");
151    let component_data = &csharp::import("Unity.Entities", "IComponentData");
152
153    let enums = ast
154        .items
155        .iter()
156        .filter_map(|item| match item {
157            syn::Item::Enum(s) => Some(s),
158            _ => None,
159        })
160        .filter(|item| {
161            item.attrs
162                .iter()
163                .any(|attr| attr.path().is_ident("bevy_state"))
164        })
165        .map(|item| (item.ident.to_string(), &item.variants))
166        .map(|(enum_name, enum_variants)| {
167            let enum_fields = enum_variants.iter();
168
169            let name_list = enum_fields.clone().enumerate().map(|(index, v)| {
170                quote! {
171                    $['\r']$(v.ident.to_string()) = $index,
172                }
173            });
174
175            let comp: csharp::Tokens = quote! {
176                namespace unrust.userland
177                {
178                    [$(struct_layout)($layout_kind.Sequential)]
179                    public struct $(&enum_name) : $component_data
180                    {
181                        public sbyte Value;
182                    }
183
184                    public class $(&enum_name)Authoring : $monobehaviour
185                    {
186                        public enum ENUM_$(&enum_name) : sbyte
187                        {
188                            $(for n in name_list => $n)
189                        }
190
191                        [SerializeField]
192                        public ENUM_$(&enum_name) $(&enum_name);
193
194                        class Baker : Baker<$(&enum_name)Authoring>
195                        {
196                            public override void Bake($(&enum_name)Authoring authoring)
197                            {
198                                var entity = GetEntity(TransformUsageFlags.None);
199                                AddComponent(entity, new $(&enum_name)
200                                {
201                                        Value = (sbyte)authoring.$(&enum_name)
202                                });
203                            }
204                        }
205                    }
206                }
207            };
208
209            (comp, enum_name)
210        });
211
212    enums.clone().try_for_each(|(comp, enum_name)| {
213        write_tokens_to_file(base_folder, &format!("{}Authoring.cs", enum_name), comp)
214    })?;
215
216    let enum_types = enums.clone().enumerate().map(|(index, (_, name))| {
217        quote! {
218            $(&name) = $index
219        }
220    });
221
222    let count = enums.clone().count();
223
224    let generated_comps: csharp::Tokens = quote! {
225        namespace unrust.userland
226        {
227            [$struct_layout($layout_kind.Sequential)]
228            public struct CustomState
229            {
230                public const int CustomStateCount = $count;
231                public CustomStateType ty;
232                public sbyte value;
233            }
234
235            public enum CustomStateType: byte
236            {
237                $(for n in enum_types => $n)
238            }
239        }
240    };
241
242    write_tokens_to_file(base_folder, "UnrustState.cs", generated_comps)?;
243
244    let add_enums = enums.clone().map(|(_, name)| {
245        quote! {
246            if (manager.HasComponent<$(&name)>(entity))
247            {
248                stateArr[stateCount] = new CustomState
249                {
250                    ty = CustomStateType.$(&name),
251                    value = manager.GetComponentData<$(&name)>(entity).Value,
252                };
253                stateCount++;
254            }
255
256
257        }
258    });
259
260    Ok(quote! {
261        $(for n in add_enums => $n)
262    })
263}
264
265fn write_tokens_to_file(base: &str, name: &str, tokens: csharp::Tokens) -> Result<()> {
266    let fmt = fmt::Config::from_lang::<Csharp>().with_indentation(fmt::Indentation::Space(4));
267    let config = csharp::Config::default();
268
269    let path = Path::new(base);
270
271    let file = File::create(path.join(name)).context("failed to open file")?;
272    let mut w = fmt::IoWriter::new(file);
273    tokens
274        .format_file(&mut w.as_formatter(&fmt), &config)
275        .context("could not write to file")
276}
277
278fn generate_components_csharp(ast: syn::File, base_folder: &str) -> Result<csharp::Tokens> {
279    let structs = find_structs_with_attr(ast, "unity_authoring")
280        .into_iter()
281        .map(generate_components_with_authoring);
282
283    let fmt = fmt::Config::from_lang::<Csharp>().with_indentation(fmt::Indentation::Space(4));
284    let config = csharp::Config::default();
285
286    let path = Path::new(base_folder);
287
288    structs.clone().try_for_each(|(comp, name)| {
289        write_tokens_to_file(base_folder, &format!("{}Authoring.cs", name), comp)
290    })?;
291
292    let count = structs.clone().count();
293
294    let gen_comps = structs.clone().map(|(_, name)| {
295        quote! {
296            $['\r'][FieldOffset(0)] public $(&name) $(&name);
297
298        }
299    });
300
301    let enum_types = structs.clone().enumerate().map(|(index, (_, name))| {
302        quote! {
303            $['\r']$name = $index,
304        }
305    });
306
307    let struct_layout = &csharp::import("System.Runtime.InteropServices", "StructLayout");
308    let layout_kind = &csharp::import("System.Runtime.InteropServices", "LayoutKind");
309
310    let generated_comps: csharp::Tokens = quote! {
311        namespace unrust.userland
312        {
313            [$struct_layout($layout_kind.Sequential)]
314            public struct CustomData
315            {
316                public CustomType ty;
317                public CustomComponents value;
318            }
319
320            public enum CustomType : byte
321            {
322                $(for n in enum_types  => $n)
323            }
324
325            [$struct_layout($layout_kind.Explicit)]
326            public struct CustomComponents
327            {
328                public const int ComponentCount = $count;
329                $(for n in gen_comps => $n)
330            }
331        }
332    };
333
334    let file = File::create(path.join("UnrustComponent.cs")).context("failed to open file")?;
335    let mut w = fmt::IoWriter::new(file);
336    generated_comps
337        .format_file(&mut w.as_formatter(&fmt), &config)
338        .context("could not write to file")?;
339
340    let add_comps = structs.clone().map(|(_, name)| {
341        quote! {
342
343
344            if (manager.HasComponent<$(&name)>(entity))
345            {
346                arr[count] = new CustomData
347                {
348                    ty = CustomType.$(&name),
349                    value = new CustomComponents { $(&name) = manager.GetComponentData<$(&name)>(entity) },
350                };
351                count++;
352            }
353        }
354    });
355
356    let add_comps: csharp::Tokens = quote! {
357        $(for n in add_comps  => $n)
358    };
359
360    Ok(add_comps)
361}
362
363fn generate_components_with_authoring(
364    (struct_name, fields): (String, Vec<(String, syn::Type)>),
365) -> (csharp::Tokens, String) {
366    let component_data = &csharp::import("Unity.Entities", "IComponentData");
367    let monobehaviour = &csharp::import("UnityEngine", "MonoBehaviour");
368    let struct_layout = &csharp::import("System.Runtime.InteropServices", "StructLayout");
369    let layout_kind = &csharp::import("System.Runtime.InteropServices", "LayoutKind");
370
371    let component_fields = fields.iter().map(|(field_name, field_type)| {
372        let field_type = map_rust_type(field_type);
373
374        quote! {
375            $['\r']public $field_type $field_name;
376        }
377    });
378
379    let authoring_name = format!("{struct_name}Authoring");
380
381    let authoring_fields = fields.iter().map(|(field_name, _)| {
382        quote! {
383            $['\r']$field_name = authoring.$field_name,
384        }
385    });
386
387    let tokens: csharp::Tokens = quote! {
388        namespace unrust.userland
389        {
390            [$struct_layout($layout_kind.Sequential)]
391            public struct $(&struct_name) : $component_data
392            {
393                $(for n in component_fields.clone() => $n)
394            }
395
396            public class $(&authoring_name) : $monobehaviour
397            {
398                $(for n in component_fields => $n)
399
400                class Baker : Baker<$(&authoring_name)>
401                {
402                    public override void Bake($(&authoring_name) authoring)
403                    {
404                        var entity = GetEntity(TransformUsageFlags.Dynamic);
405                        AddComponent(entity, new $(&struct_name)
406                            {
407                                $(for n in authoring_fields => $n)
408                            });
409                    }
410                }
411            }
412        }
413    };
414
415    (tokens, struct_name)
416}
417
418fn find_structs_with_attr(
419    ast: syn::File,
420    expected: &str,
421) -> Vec<(String, Vec<(String, syn::Type)>)> {
422    ast.items
423        .iter()
424        .filter_map(|item| match item {
425            syn::Item::Struct(s) => Some(s),
426            _ => None,
427        })
428        .filter(|item| item.attrs.iter().any(|attr| attr.path().is_ident(expected)))
429        .map(|item| {
430            let fields = match &item.fields {
431                syn::Fields::Named(fields) => fields
432                    .named
433                    .iter()
434                    .map(|f| {
435                        let field_name = f.ident.clone().unwrap().to_string();
436                        let field_type = f.ty.clone();
437                        (field_name, field_type)
438                    })
439                    .collect::<Vec<(String, syn::Type)>>(),
440                _ => vec![],
441            };
442
443            (item.ident.to_string(), fields)
444        })
445        .collect()
446}
447
448fn clear_contents(path: &str, extension: &str) -> Result<()> {
449    let files = std::fs::read_dir(path)?;
450    files
451        .filter_map(|p| p.ok())
452        .filter_map(|p| {
453            let path = p.path();
454            let Some(ext) = path.extension() else {
455                return None;
456            };
457
458            if ext == OsStr::new(extension) {
459                Some(p.path())
460            } else {
461                None
462            }
463        })
464        .try_for_each(std::fs::remove_file)?;
465
466    Ok(())
467}
468
469fn map_rust_type(ty: &syn::Type) -> String {
470    let syn::Type::Path(path) = ty else {
471        panic!("expected rust path type");
472    };
473
474    let path = path
475        .path
476        .get_ident()
477        .expect("expected simple path type")
478        .to_string();
479
480    match path.as_str() {
481        "f32" => "float",
482        "f64" => "double",
483        "i32" => "int",
484        "i64" => "long",
485        "u32" => "uint",
486        "u64" => "ulong",
487        _ => panic!("unsupported base type"),
488    }
489    .to_string()
490}