sdk-builder: using tokens instead of strings to transport type idents/paths

sdk-builder/src/main.rs

@@ -135,15 +135,127 @@ fn empty_or_some(s: &str) -> Option<&str> {
     }
 }
 
-pub struct CanonicalNames {
+pub mod path_helper {
-    /// canonicalized type names for lookup when handling return types and parameters.
+    use quote::{format_ident, quote, ToTokens, TokenStreamExt};
-    types: BTreeMap<ObjectRef, String>,
+    use unreal_sdk::sdk::repr::Type;
+
+    use crate::rust::Builder;
+
+    pub struct Type2<'a> {
+        cache: &'a Builder,
+        inner: Type,
+    }
+
+    impl<'a> Type2<'a> {
+        pub fn new(cache: &'a Builder, inner: Type) -> Self {
+            Self { cache, inner }
+        }
+
+        fn child(&self, inner: Type) -> Self {
+            Self {
+                inner,
+                cache: self.cache,
+            }
+        }
+    }
+
+    impl<'a> ToTokens for Type2<'a> {
+        fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+            match &self.inner {
+                Type::Ptr(inner) | Type::Ref(inner) => {
+                    let inner = self.child(*inner.clone());
+
+                    tokens.extend(quote! {
+                        ::core::option::Option<NonNull<#inner>>
+                    });
+                }
+                Type::WeakPtr(inner) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&inner)
+                        .unwrap_or(quote!(crate::engine::UObject));
+                    tokens.extend(quote!(
+                        crate::engine::TWeakObjectPtr<#inner>
+                    ));
+                }
+                Type::SoftPtr(inner) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&inner)
+                        .unwrap_or(quote!(crate::engine::UObject));
+                    tokens.extend(quote!(
+                        crate::engine::TSoftObjectPtr<#inner>
+                    ));
+                }
+                Type::LazyPtr(inner) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&inner)
+                        .unwrap_or(quote!(crate::engine::UObject));
+
+                    tokens.extend(quote!(
+                        crate::engine::TLazyObjectPtr<#inner>
+                    ));
+                }
+                Type::AssetPtr(inner) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&inner)
+                        .unwrap_or(quote!(crate::engine::UObject));
+
+                    tokens.extend(quote!(
+                        crate::engine::TAssetPtr<#inner>,
+                    ));
+                }
+                Type::Array(inner) => {
+                    let inner = self.child(*inner.clone());
+                    tokens.extend(quote!(
+                        crate::engine::TArray<#inner>,
+                    ));
+                }
+                Type::Primitive(prim) => {
+                    tokens.append(format_ident!("{prim}"));
+                }
+                Type::RawArray { ty, len } => {
+                    let ty = self.child(*ty.clone());
+                    quote!([#ty; #len]);
+                }
+                Type::Name => tokens.extend(quote!(crate::engine::FName)),
+                Type::String => tokens.extend(quote!(crate::engine::FString)),
+                Type::Text => tokens.extend(quote!(crate::engine::FText)),
+                Type::Enum { enum_type, .. } => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&enum_type)
+                        .unwrap_or(quote!(u8));
+
+                    tokens.extend(inner);
+                }
+                Type::Class(class) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&class)
+                        .unwrap_or(quote!(crate::engine::UObject));
+
+                    tokens.extend(quote!(::core::option::Option<#inner>));
+                }
+                Type::Struct(class) => {
+                    let inner = self
+                        .cache
+                        .get_full_type_tokens(&class)
+                        .unwrap_or(quote!(()));
+
+                    tokens.extend(inner);
+                }
+            };
+        }
+    }
 }
 
 pub mod rust {
     use std::{
         borrow::Cow,
-        collections::{BTreeMap, BTreeSet, HashMap, HashSet},
+        collections::{BTreeMap, BTreeSet},
         path::Path,
     };
 
@@ -151,13 +263,12 @@ pub mod rust {
     use itertools::Itertools;
     use proc_macro2::{Ident, TokenStream};
     use quote::{format_ident, quote};
-    use rayon::prelude::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
     use unreal_sdk::sdk::repr::{
         Class, ClassField, ClassMethod, Enum, ObjectRef, PackageRef, PrimitiveType,
         ProcessedPackage, Sdk, StructKind, Type, UnrealType,
     };
 
-    use crate::split_at_illegal_char;
+    use crate::{path_helper::Type2, split_at_illegal_char};
 
     // const KEYWORDS: [&'static str; 51] = [
     //     "as", "break", "const", "continue", "crate", "else", "enum", "extern", "false", "fn",
@@ -192,6 +303,10 @@ pub mod rust {
         sdk: Sdk,
     }
 
+    fn canonicalize_ident(name: &str) -> Ident {
+        format_ident!("{}", canonicalize_name(name))
+    }
+
     fn canonicalize_name<'a>(name: &'a str) -> Cow<'a, str> {
         let valid = split_at_illegal_char(name, &CHARS).into_valid(&CHARS);
         if WORDS.contains(&valid.as_ref()) || valid.starts_with(|c: char| !c.is_alphabetic()) {
@@ -221,9 +336,17 @@ pub mod rust {
 
         /// returns the absolute path of a type with the assumption that all
         /// packages are children of the path `crate::sdk`
-        fn get_type_package_path(&self, key: &ObjectRef) -> Option<String> {
+        pub fn get_type_package_path(&self, key: &ObjectRef) -> Option<TokenStream> {
-            let pkg = &self.sdk.packages.get(&key.package)?.name;
+            let pkg = format_ident!("{}", &self.sdk.packages.get(&key.package)?.name);
-            Some(format!("crate::sdk::{pkg}"))
+            Some(quote!(crate::sdk::#pkg))
+        }
+
+        pub fn get_full_type_tokens(&self, key: &ObjectRef) -> Option<TokenStream> {
+            let pkg = self.get_type_package_path(key)?;
+            let name = format_ident!("{}", self.get_type_name(key)?);
+            Some(quote! {
+                #pkg::#name
+            })
         }
 
         /// returns the absolute path of a type with the assumption that all
@@ -234,6 +357,12 @@ pub mod rust {
                 .map(|name| format!("crate::sdk::{pkg}::{name}"))
         }
 
+        /// returns the precached, prefixed and cannonicalized (for this
+        /// language, Rust) `Ident` for this object-ref
+        fn get_type_ident(&self, key: &ObjectRef) -> Option<Ident> {
+            Some(format_ident!("{}", self.type_name_cache.get(key)?))
+        }
+
         /// returns the precached, prefixed and cannonicalized (for this
         /// language, Rust) name for this object-ref
         fn get_type_name(&self, key: &ObjectRef) -> Option<String> {
@@ -421,11 +550,11 @@ pub mod rust {
 
         fn generate_enum(&self, enum0: &Enum) -> anyhow::Result<TokenStream> {
             let name = self
-                .get_type_name(&enum0.obj_ref)
+                .get_type_ident(&enum0.obj_ref)
                 .context("enum name was not previously canonicalized and cached.")?;
 
             let variants = enum0.values.iter().map(|(&value, name)| {
-                let name = canonicalize_name(&name);
+                let name = canonicalize_ident(&name);
                 quote! {
                     #name = #value,
                 }
@@ -448,19 +577,17 @@ pub mod rust {
         fn generate_object(
             &self,
             _class: &Class,
-            name: &str,
+            name: &Ident,
         ) -> anyhow::Result<(TokenStream, TokenStream)> {
-            let ident = format_ident!("{name}");
-
             let typedef = quote! {
                 #[derive(Eq, PartialEq, Copy, Clone)]
-                pub struct #ident(pub ::core::ptr::NonNull<u8>);
+                pub struct #name(pub ::core::ptr::NonNull<u8>);
             };
 
             let static_class_impl: TokenStream = Self::generate_find_object(name);
 
             let impls = quote! {
-                impl crate::engine::AsUObject for #ident {
+                impl crate::engine::AsUObject for #name {
                     fn as_uobject(&self) -> crate::engine::UObject {
                         crate::engine::UObject(self.0)
                     }
@@ -470,7 +597,7 @@ pub mod rust {
                     }
                 }
 
-                impl crate::engine::AsPtr for #ident {
+                impl crate::engine::AsPtr for #name {
                     fn as_ptr(&self) -> *const u8 {
                         unsafe { self.0.as_ref().get() as _ }
                     }
@@ -480,7 +607,7 @@ pub mod rust {
                     }
                 }
 
-                impl crate::engine::StaticClass for #ident {
+                impl crate::engine::StaticClass for #name {
                     fn get_static_class() -> ::core::option::Option<crate::engine::UClass> {
                         let class: ::core::option::Option<crate::engine::UClass> =
                             #static_class_impl;
@@ -498,30 +625,30 @@ pub mod rust {
         fn generate_struct(
             &self,
             class: &Class,
-            name: &str,
+            name: &Ident,
             ctor: Option<TokenStream>,
         ) -> anyhow::Result<(TokenStream, TokenStream)> {
             let size = class.size;
-            let ident = format_ident!("{name}");
+
             let typedef = quote! {
-                pub struct #ident(pub ::core::cell::UnsafeCell<u8; #size>);
+                pub struct #name(pub ::core::cell::UnsafeCell<u8; #size>);
             };
 
             let impls = quote! {
-                impl Eq for #ident {}
+                impl Eq for #name {}
-                impl PartialEq for #ident {
+                impl PartialEq for #name {
                     fn eq(&self, other: &Self) -> bool {
                         unsafe {(&*self.0.get()).eq(&*other.0.get())}
                     }
                 }
 
-                impl Clone for #ident {
+                impl Clone for #name {
                     fn clone(&self) -> Self {
                         Self(::core::cell::UnsafeCell::new(unsafe {&*self.0.get()}.clone()))
                     }
                 }
 
-                impl crate::engine::AsPtr for #ident {
+                impl crate::engine::AsPtr for #name {
                     fn as_ptr(&self) -> *const u8 {
                         self.0.get().cast()
                     }
@@ -531,7 +658,7 @@ pub mod rust {
                     }
                 }
 
-                impl #ident {
+                impl #name {
                     pub fn zeroed() -> Self {
                         unsafe {
                             ::core::mem::MaybeUninit::<Self>::zeroed().assume_init()
@@ -552,7 +679,7 @@ pub mod rust {
         fn generate_struct_methods(
             &self,
             class: &Class,
-            name: &str,
+            name: &Ident,
         ) -> anyhow::Result<(Vec<TokenStream>, Vec<TokenStream>)> {
             let methods = class
                 .methods
@@ -570,18 +697,18 @@ pub mod rust {
         /// - the method wrapper.
         fn generate_method(
             &self,
-            struct_name: &str,
+            struct_name: &Ident,
             method: &ClassMethod,
         ) -> anyhow::Result<(TokenStream, TokenStream)> {
-            let method_name = canonicalize_name(&method.unique_name());
+            let method_name = canonicalize_ident(&method.unique_name());
 
             // all parameters collected as (parameter, canonicalized_name, type_ident)
             let parameters = method
                 .parameters
                 .iter()
                 .map(|parameter| {
-                    let name = canonicalize_name(&parameter.unique_name());
+                    let name = canonicalize_ident(&parameter.unique_name());
-                    let type_name = self.type_name(&parameter.ty)?;
+                    let type_name = Type2::new(self, parameter.ty.clone());
 
                     anyhow::Ok((parameter, name, type_name))
                 })
@@ -590,11 +717,7 @@ pub mod rust {
             // all parameters converted into "arg: Type" format of tokens.
             let all_params = parameters
                 .iter()
-                .map(|(param, name, ty)| {
+                .map(|(param, name, ty)| (param, quote! {#name: #ty}))
-                    let name = format_ident!("{name}");
-                    let ty = format_ident!("{ty}");
-                    (param, quote! {#name: #ty})
-                })
                 .collect::<Vec<_>>();
 
             // params that the function will accept as arguments.
@@ -611,7 +734,6 @@ pub mod rust {
             // param token streams for setting the fields of the params struct
             // with the arguments of the function.
             let init_params = parameters.iter().map(|(_, name, _)| {
-                let name = format_ident!("{name}");
                 quote! {params.#name = #name;}
             });
 
@@ -622,9 +744,6 @@ pub mod rust {
                         param.is_return_param() || (param.is_out_param() && !param.is_const_param())
                     })
                     .map(|(_, name, ty)| {
-                        let name = format_ident!("{name}");
-                        let ty = format_ident!("{ty}");
-
                         (
                             quote! {
                                 #name
@@ -680,8 +799,8 @@ pub mod rust {
         fn generate_field_accessors(
             &self,
             field: &ClassField,
-            field_name: &Cow<str>,
+            field_name: &Ident,
-            type_name: &String,
+            type_name: &Type2,
         ) -> TokenStream {
             let setter = format_ident!("set_{}", field_name);
             let getter = format_ident!("get_{}", field_name);
@@ -753,14 +872,14 @@ pub mod rust {
         fn generate_struct_ctor(
             &self,
             _class: &Class,
-            type_name: &str,
+            type_ident: &Ident,
-            fields: &Vec<(&ClassField, Cow<str>, String)>,
+            fields: &Vec<(&ClassField, Ident, Type2)>,
         ) -> TokenStream {
             let fields_defs = fields.iter().map(|(_, name, ty)| quote! {#name: #ty});
 
             let this_field_asignments = fields.iter().map(|(_, name, _ty)| {
                 let setter = format_ident!("set_{}", name);
-                let field_trait = format_ident!("{type_name}Fields");
+                let field_trait = format_ident!("{type_ident}Fields");
 
                 quote! {<Self as #field_trait>::#setter(this, #name);}
             });
@@ -785,14 +904,14 @@ pub mod rust {
         fn generate_struct_fields(
             &self,
             class: &Class,
-            name: &str,
+            name: &Ident,
         ) -> anyhow::Result<(TokenStream, Option<TokenStream>)> {
             let fields = class
                 .fields
                 .iter()
                 .map(|field| {
-                    let name = canonicalize_name(&field.unique_name());
+                    let name = canonicalize_ident(&field.unique_name());
-                    let ty = self.type_name(&field.ty)?;
+                    let ty = Type2::new(self, field.ty.clone());
 
                     anyhow::Ok((field, name, ty))
                 })
@@ -821,7 +940,8 @@ pub mod rust {
             Ok((fields_trait, ctor))
         }
 
-        fn generate_find_object(name: &str) -> TokenStream {
+        fn generate_find_object<S: ToString>(name: S) -> TokenStream {
+            let name = name.to_string();
             let not_found = format!("static object \"{name}\" not found!");
             quote! {
                 static OBJECT: ::once_cell::sync::OnceCell<::core::option::Option<UObject>> = ::once_cell::sync::OnceCell::new();
@@ -856,7 +976,7 @@ pub mod rust {
 
         fn generate_class(&self, class: &Class) -> anyhow::Result<TokenStream> {
             let name = &self
-                .get_type_name(&class.obj_ref)
+                .get_type_ident(&class.obj_ref)
                 .context("enum name was not previously canonicalized and cached.")?;
 
             let (field_trait, ctor) = self.generate_struct_fields(class, name)?;
@@ -924,7 +1044,7 @@ pub mod rust {
             pkg: &ProcessedPackage,
             feature_gate: bool,
         ) -> anyhow::Result<TokenStream> {
-            let pkg_name = canonicalize_name(&pkg.name);
+            let pkg_name = canonicalize_ident(&pkg.name);
             log::info!(
                 "generating package \"{pkg_name}\" with {} types..",
                 pkg.types.len()