unreal-sdk/src/lib.rs

#![feature(
    const_trait_impl,
    const_ptr_as_ref,
    const_nonnull_new,
    let_chains,
    if_let_guard
)]

// mod core_types;
pub mod fname;
pub mod global_tables;
pub mod tarray;
// pub mod types;
pub mod v2_types;

pub mod sdk {
    pub mod output;
    use std::{
        borrow::Cow,
        collections::{hash_map::Entry, BTreeMap, HashMap, HashSet},
        fmt::Display,
        sync::Mutex,
    };

    use anyhow::Context;
    use itertools::Itertools;
    use rayon::prelude::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};

    use crate::{
        global_tables::objects::{FindClass, GOBJECTS},
        v2_types::{
            actor_static_class,
            any_type::{self, AnyField, AnyObject, AnyProperty, AnyStruct},
            traits::{
                AsUObject, StaticClass, UArrayPropertyTrait, UBytePropertyTrait,
                UEnumPropertyTrait, UEnumTrait, UFunctionTrait, UObjectNonConst,
                UObjectPropertyBaseTrait, UObjectTrait, UPropertyTrait, UStructNonConst,
                UStructPropertyTrait, UStructTrait,
            },
            EFunctionFlags, EPropertyFlags, UClass, UEnum, UFunction, UObject, UStruct,
        },
    };

    use self::output::rust::inject_coreuobject_types;

    struct ClassCache {
        classes: Mutex<HashMap<String, UClass>>,
    }

    impl ClassCache {
        #[allow(dead_code)]
        pub fn new() -> Self {
            Self {
                classes: Mutex::new(HashMap::new()),
            }
        }
    }

    impl FindClass for ClassCache {
        fn find_class<S>(&self, class_name: S) -> Option<UClass>
        where
            S: Into<String>,
        {
            let class_name = class_name.into();

            let class = self.classes.lock().unwrap().get(&class_name).cloned();

            match class {
                class @ Some(_) => class,
                None => match GOBJECTS
                    .read()
                    .unwrap()
                    .as_objects()
                    .unwrap()
                    .find_class(&class_name)
                {
                    Some(class) => {
                        self.classes
                            .lock()
                            .unwrap()
                            .insert(class_name, class.clone());
                        Some(class)
                    }
                    None => None,
                },
            }
        }
    }

    pub struct Package {
        package: UObject,
        objects: Vec<UObject>,
    }

    #[derive(Debug)]
    pub struct ProcessedPackage {
        pub package: UObject,
        pub types: HashMap<UObject, Types>,
        pub package_dependencies: HashMap<UObject, Vec<UObject>>,
    }

    impl Package {
        pub fn new(package: UObject, objects: Vec<UObject>) -> Self {
            Self { package, objects }
        }

        pub fn process(self) -> ProcessedPackage {
            let types = self
                .objects
                .par_iter()
                .filter_map(|&object| {
                    match AnyObject::from_object(object) {
                        AnyObject::Field(field) => match AnyField::from_field(field) {
                            AnyField::Enum(enm) => {
                                if let Ok(enm) = Self::process_enum(enm) {
                                    return Some((object, Types::Enum(enm)));
                                }
                            }
                            AnyField::Struct(strt) => match AnyStruct::from_struct(strt) {
                                strt @ AnyStruct::ScriptStruct(_) | strt @ AnyStruct::Class(_) => {
                                    if let Ok(class) = self.process_struct(unsafe { strt.cast() }) {
                                        return Some((object, Types::Class(class)));
                                    }
                                }
                                _ => {}
                            },
                            _ => {}
                        },
                        _ => {}
                    }

                    None
                })
                .collect::<HashMap<_, _>>();

            let mut dependencies = types
                .iter()
                .filter_map(|(_, ty)| match ty {
                    Types::Class(class) => Some(class.referenced_types()),
                    _ => None,
                })
                .map(|refs| refs.into_iter())
                .flatten()
                .fold_into_packages();

            // remove any objects in Self
            dependencies.remove(&self.package);

            ProcessedPackage {
                package: self.package,
                types,
                package_dependencies: dependencies,
            }
        }

        fn process_struct(&self, strct: UStruct) -> anyhow::Result<Class> {
            let name = strct
                .get_name()
                .context("failed to get struct or class name")?;

            let full_name = strct
                .get_full_name()
                .context("failed to get struct or class full name")?;

            let is_actor = strct
                .iter_super_structs()
                .contains(&unsafe { actor_static_class().unwrap().cast() });
            let is_class = strct.is_a(&UClass::static_class().unwrap());

            let super_struct = if let Some(spr) = *strct.super_field() {
                if spr != strct {
                    Some(spr)
                } else {
                    None
                }
            } else {
                None
            };

            let (fields, methods) = self.process_children(strct)?;

            Ok(Class {
                is_class,
                size: *strct.property_size() as u32,
                name: format!(
                    "{}{name}",
                    if is_actor {
                        "A"
                    } else if is_class {
                        "U"
                    } else {
                        "F"
                    }
                ),
                super_class: super_struct,
                fields,
                methods,
                full_name,
            })
        }

        fn find_type(prop: AnyProperty) -> anyhow::Result<Type> {
            let ty = match prop.clone() {
                numeric @ AnyProperty::Numeric(_) => {
                    match numeric.as_any_numeric_property().unwrap() {
                        any_type::AnyNumericProperty::U8(prop) => match prop.uenum() {
                            Some(enm) => Type::Enum {
                                underlying: Box::new(Type::Primitive(PrimitiveType::U8)),
                                enum_type: *enm,
                            },
                            None => Type::Primitive(PrimitiveType::U8),
                        },
                        any_type::AnyNumericProperty::U16(_) => Type::Primitive(PrimitiveType::U16),
                        any_type::AnyNumericProperty::U32(_) => Type::Primitive(PrimitiveType::U32),
                        any_type::AnyNumericProperty::U64(_) => Type::Primitive(PrimitiveType::U64),
                        any_type::AnyNumericProperty::I8(_) => Type::Primitive(PrimitiveType::I8),
                        any_type::AnyNumericProperty::I16(_) => Type::Primitive(PrimitiveType::I16),
                        any_type::AnyNumericProperty::I32(_) => Type::Primitive(PrimitiveType::I32),
                        any_type::AnyNumericProperty::I64(_) => Type::Primitive(PrimitiveType::I64),
                        any_type::AnyNumericProperty::F32(_) => Type::Primitive(PrimitiveType::F32),
                        any_type::AnyNumericProperty::F64(_) => Type::Primitive(PrimitiveType::F64),
                        any_type::AnyNumericProperty::Other(_) => {
                            return Err(anyhow::anyhow!("unhandled numeric property"));
                        }
                    }
                }
                AnyProperty::Bool(_) => Type::Primitive(PrimitiveType::Bool),
                AnyProperty::Interface(_) => {
                    return Err(anyhow::anyhow!("skipping interfaces for now"))
                }
                object @ AnyProperty::ObjectBase(_) => {
                    match object.as_any_object_base_property().unwrap() {
                        any_type::AnyObjectBaseProperty::Object(obj) => {
                            Type::Ptr(Box::new(Type::Class(unsafe {
                                // safety: any uclass is also a ustruct
                                obj.property_class()
                                    .context("object property missing properry class")?
                                    .cast()
                            })))
                        }
                        any_type::AnyObjectBaseProperty::WeakObject(obj) => Type::WeakPtr(
                            obj.property_class()
                                .context("weak ptr property missing property class.")?,
                        ),
                        any_type::AnyObjectBaseProperty::LazyObject(obj) => Type::LazyPtr(
                            obj.property_class()
                                .context("lazy ptr property missing property class.")?,
                        ),
                        any_type::AnyObjectBaseProperty::SoftObject(obj) => Type::SoftPtr(
                            obj.property_class()
                                .context("soft ptr property missing property class.")?,
                        ),
                        asset @ any_type::AnyObjectBaseProperty::AssetObject(_) => {
                            match asset.as_any_asset_object_property().unwrap() {
                                any_type::AnyAssetObjectProperty::Class(class) => Type::AssetPtr(
                                    class
                                        .property_class()
                                        .context("asset object property missing properry class")?,
                                ),
                                any_type::AnyAssetObjectProperty::Object(_) => {
                                    return Err(anyhow::anyhow!(
                                        "unhandled asset object property (NOT AN ERROR)"
                                    ));
                                }
                            }
                        }
                        any_type::AnyObjectBaseProperty::Other(_) => {
                            return Err(anyhow::anyhow!("unhandled object base property"));
                        }
                    }
                }
                AnyProperty::Array(array) => {
                    Type::Array(Box::new(Self::find_type(AnyProperty::from_prop(
                        array
                            .inner()
                            .context("array property inner type missing.")?,
                    ))?))
                }
                AnyProperty::Map(_) => unreachable!("not used in ARK"),
                AnyProperty::Str(_) => Type::String,
                AnyProperty::Text(_) => Type::Text,
                AnyProperty::Name(_) => Type::Name,
                AnyProperty::Delegate(_) => {
                    return Err(anyhow::anyhow!("skipping delegates for now"));
                }
                AnyProperty::MulticastDelegate(_) => {
                    return Err(anyhow::anyhow!("skipping multicast delegates for now"));
                }
                AnyProperty::Enum(enm) => Type::Enum {
                    underlying: Box::new(Self::find_type(AnyProperty::from_prop(
                        enm.underlying_type()
                            .context("enum property was missing underlying type")?,
                    ))?),
                    enum_type: enm.uenum().context("enum property missing enum type")?,
                },
                AnyProperty::Struct(class) => Type::Class(
                    class
                        .ustruct()
                        .context("struct property had no inner struct")?,
                ),
                AnyProperty::Other(_) => {
                    return Err(anyhow::anyhow!("unhandled property."));
                }
            };

            if *prop.array_dim() > 1 {
                Ok(Type::RawArray {
                    ty: Box::new(ty),
                    len: *prop.array_dim() as u32,
                })
            } else {
                Ok(ty)
            }
        }

        fn process_children(
            &self,
            strct: UStruct,
        ) -> anyhow::Result<(Vec<ClassField>, Vec<ClassMethod>)> {
            log::debug!("{} children:", strct.get_full_name_or_default());

            let mut fields = Vec::new();
            let mut methods = Vec::new();

            for child in strct
                .iter_children()
                .map(|field| any_type::AnyField::from_field(field))
            {
                match child {
                    any_type::AnyField::Property(prop) => {
                        match Self::find_type(any_type::AnyProperty::from_prop(prop)) {
                            Ok(ty) => {
                                log::debug!("field: {ty:?}: {prop}");
                                fields.push(
                                ClassField {
                                    offset: *prop.offset() as u32,
                                    size: *prop.element_size() as u32,
                                    name: canonicalize_name(&prop.name().get_name().context("failed to retrieve field name")?).to_string(),
                                    ty,
                                });
                            },
                            Err(e) => {
                                log::warn!("skipping field with offset {}: {e}", prop.offset());
                            },
                        }
                    },
                    strt @ any_type::AnyField::Struct(_) if let Some(any_type::AnyStruct::Function(func)) = strt.as_any_struct() => {
                        log::debug!("function: {func}");

                        if let Ok(method) = Self::process_function(func) {
                            methods.push(method);
                        }
                    },
                    _ => {
                    }
                }
            }

            Ok((fields, methods))
        }

        fn process_function(func: UFunction) -> anyhow::Result<ClassMethod> {
            let full_name = func
                .get_full_name()
                .context("could not get full function name")?;
            let rust_name =
                canonicalize_name(&func.get_name().context("could not get function name")?)
                    .to_string();

            let is_static = func.function_flags().contains(EFunctionFlags::Static);
            let is_native = func.function_flags().contains(EFunctionFlags::Native);

            let mut params = Vec::new();
            let mut return_types = Vec::new();

            for child in func
                .iter_children()
                .map(|field| any_type::AnyField::from_field(field))
            {
                match child {
                    any_type::AnyField::Property(prop) => {
                        match Self::find_type(any_type::AnyProperty::from_prop(prop)) {
                            Ok(ty) => {
                                log::debug!("field: {ty:?}: {prop}");

                                let param = MethodParameter {
                                    name: prop
                                        .get_name()
                                        .context("failed to get parameter name")?,
                                    ty,
                                };

                                let param = if prop
                                    .property_flags()
                                    .contains(EPropertyFlags::ReturnParm)
                                {
                                    return_types.push(param.clone());
                                    Some(ParameterKind::Return(param))
                                } else if prop.property_flags().contains(EPropertyFlags::OutParm) {
                                    if prop.property_flags().contains(EPropertyFlags::ConstParm) {
                                        Some(ParameterKind::Default(param))
                                    } else {
                                        Some(ParameterKind::Out(param))
                                    }
                                } else if prop.property_flags().contains(EPropertyFlags::Parm) {
                                    Some(ParameterKind::Default(param))
                                } else {
                                    None
                                };

                                params.extend(param);
                            }
                            Err(e) => {
                                log::warn!("skipping field with offset {}: {e}", prop.offset());
                            }
                        }
                    }
                    _ => {}
                }
            }

            Ok(ClassMethod {
                name: rust_name,
                parameters: params,
                return_type: return_types.into_iter().next(),
                full_name,
                is_native,
                is_static,
            })
        }

        fn process_enum(enm: UEnum) -> anyhow::Result<Enum> {
            // get all the variants
            let values = enm
                .names()
                .iter()
                .enumerate()
                // canonicalize the names into valid rust
                .map(|(value, name)| {
                    let name = name.get_name().unwrap_or("AnonymousVariant".to_string());
                    let name = canonicalize_name(&name).to_string();
                    (value, name)
                })
                // store conflicts next to each other
                .fold(
                    HashMap::<String, Vec<usize>>::new(),
                    |mut acc, (value, name)| {
                        match acc.entry(name) {
                            Entry::Occupied(mut entry) => {
                                entry.get_mut().push(value);
                            }
                            Entry::Vacant(entry) => {
                                entry.insert(vec![value]);
                            }
                        }
                        acc
                    },
                );

            // sort by value
            let mut variants = BTreeMap::new();
            values.into_iter().for_each(|(name, values)| {
                if values.len() > 0 {
                    for (i, value) in values.into_iter().enumerate() {
                        variants.insert(value, format!("{name}{i}"));
                    }
                } else {
                    variants.insert(values.into_iter().next().unwrap(), name);
                }
            });

            let name = enm
                .as_uobject()
                .get_full_name()
                .context("enum name could not be found")?;
            let name = canonicalize_name(&name).to_string();

            Ok(Enum {
                name,
                values: variants.into_iter().map(|(_, name)| name).collect(),
            })
        }

        pub fn package_object(&self) -> &UObject {
            &self.package
        }
    }

    trait FoldIntoPackages: Iterator<Item = UObject> + Sized {
        fn fold_into_packages(self) -> HashMap<UObject, Vec<UObject>> {
            self.map(|obj| (obj.package_object(), obj)).fold(
                HashMap::<UObject, Vec<UObject>>::new(),
                |mut acc, (pkg, obj)| {
                    match acc.entry(pkg) {
                        Entry::Occupied(mut entry) => {
                            entry.get_mut().push(obj);
                        }
                        Entry::Vacant(entry) => {
                            entry.insert(vec![obj]);
                        }
                    }

                    acc
                },
            )
        }
    }

    impl<T> FoldIntoPackages for T where T: Iterator<Item = UObject> {}

    pub struct Sdk {
        pub packages: HashMap<UObject, ProcessedPackage>,
    }

    impl Sdk {
        pub fn new() -> Self {
            let packages = Self::package_objects();

            let packages = packages
                .into_par_iter()
                .map(|pkg| pkg.process())
                .map(|pkg| (pkg.package, pkg))
                .collect();

            Self { packages }
        }

        pub fn patch(mut self) -> Self {
            inject_coreuobject_types(&mut self);
            self
        }

        pub fn find_type(&self, obj: UObject) -> Option<&Types> {
            self.packages
                .get(&obj.package_object())
                .and_then(|pkg| pkg.types.get(&obj))
        }

        pub fn inject_type(&mut self, class: Class) -> Option<UClass> {
            if let Some(class_obj) = GOBJECTS
                .read()
                .unwrap()
                .as_objects()
                .unwrap()
                .find_class(&class.full_name)
            {
                let package = class_obj.package_object();
                match self.packages.entry(package.clone()) {
                    Entry::Occupied(mut entry) => {
                        entry
                            .get_mut()
                            .types
                            .insert(class_obj.as_uobject(), Types::Class(class));
                    }
                    Entry::Vacant(entry) => {
                        let mut types = HashMap::new();
                        types.insert(class_obj.as_uobject(), Types::Class(class));
                        entry.insert(ProcessedPackage {
                            package,
                            types,
                            package_dependencies: HashMap::new(),
                        });
                    }
                }

                Some(class_obj)
            } else {
                None
            }
        }

        pub fn package_objects() -> Vec<Package> {
            let gobjects = GOBJECTS.read().unwrap();
            let objects = gobjects.as_objects().unwrap();

            let sorted_objects = objects
                .iter()
                // ensure item contains object
                .filter_map(|item| item.object())
                // get package object
                .fold_into_packages();

            sorted_objects
                .into_iter()
                .map(|(package, objects)| Package::new(package, objects))
                .collect::<Vec<_>>()
        }
    }

    fn keywords() -> HashSet<&'static str> {
        let mut keywords = HashSet::new();

        // rust keywords
        keywords.extend([
            "as", "break", "const", "continue", "crate", "else", "enum", "extern", "false", "fn",
            "for", "if", "impl", "in", "let", "loop", "match", "mod", "move", "mut", "pub", "ref",
            "return", "self", "Self", "static", "struct", "super", "trait", "true", "type",
            "unsafe", "use", "where", "while", "async", "await", "dyn", "abstract", "become",
            "box", "do", "final", "macro", "override", "priv", "typeof", "unsized", "virtual",
            "yield", "try",
        ]);

        // rust types
        keywords.extend([
            "bool", "f64", "f32", "str", "char", "u8", "u16", "u32", "u64", "u128", "i8", "i16",
            "i32", "i64", "i128", "usize", "isize",
        ]);

        keywords
    }

    fn token_chars() -> HashSet<char> {
        let mut chars = HashSet::new();

        chars.extend([
            ' ', '?', '+', '-', ':', '/', '^', '(', ')', '[', ']', '<', '>', '&', '.', '#', '\'',
            '"', '%',
        ]);

        chars
    }

    struct SplitResult<'a> {
        start: Option<&'a str>,
        middle: usize,
        end: Option<&'a str>,
    }

    impl<'a> SplitResult<'a> {
        pub fn is_valid(&self) -> bool {
            self.start.is_some() && self.middle == 0 && self.end.is_none()
        }

        pub fn into_valid(self) -> Cow<'a, str> {
            if self.is_valid() {
                Cow::Borrowed(self.start.unwrap())
            } else {
                let mut valid = self.start.map(|s| s.to_string()).unwrap_or_default();
                valid.extend(core::iter::repeat('_').take(self.middle));

                match self.end {
                    Some(end) => {
                        valid.push_str(&split_at_illegal_char(end).into_valid());
                    }
                    None => {}
                }

                Cow::Owned(valid)
            }
        }
    }

    fn empty_or_some(s: &str) -> Option<&str> {
        if s.is_empty() {
            None
        } else {
            Some(s)
        }
    }

    fn split_at_illegal_char(input: &str) -> SplitResult {
        let illegal_chars = token_chars();
        if let Some(pos) = input.find(|c| illegal_chars.contains(&c)) {
            let start = empty_or_some(&input[..pos]);
            // skip the illegal char
            let rest = &input[pos + 1..];

            if let Some(pos2) = rest.find(|c| !illegal_chars.contains(&c)) {
                SplitResult {
                    start,
                    middle: pos2 + 1,
                    end: empty_or_some(&rest[pos2..]),
                }
            } else {
                SplitResult {
                    start,
                    middle: 1,
                    end: empty_or_some(rest),
                }
            }
        } else {
            SplitResult {
                start: empty_or_some(input),
                middle: 0,
                end: None,
            }
        }
    }

    fn canonicalize_name(name: &str) -> Cow<str> {
        let valid = split_at_illegal_char(name).into_valid();
        if keywords().contains(valid.as_ref()) {
            Cow::Owned(format!("_{}", &valid))
        } else {
            valid
        }
    }

    #[derive(Debug, Clone, Copy)]
    pub enum PrimitiveType {
        Bool,
        U8,
        U16,
        U32,
        U64,
        I8,
        I16,
        I32,
        I64,
        F32,
        F64,
        Custom(&'static str),
    }

    #[derive(Debug, Clone)]
    pub enum Type {
        Ptr(Box<Type>),
        Ref(Box<Type>),
        WeakPtr(UClass),
        SoftPtr(UClass),
        LazyPtr(UClass),
        AssetPtr(UClass),
        Array(Box<Type>),
        Primitive(PrimitiveType),
        RawArray {
            ty: Box<Type>,
            len: u32,
        },
        Name,
        String,
        Text,
        Enum {
            underlying: Box<Type>,
            enum_type: UEnum,
        },
        Class(UStruct),
    }

    impl Type {
        pub fn referenced_type(&self) -> Option<UObject> {
            match self {
                Type::Ptr(t) | Type::Ref(t) | Type::Array(t) | Type::RawArray { ty: t, .. } => {
                    t.referenced_type()
                }
                Type::WeakPtr(o) | Type::SoftPtr(o) | Type::LazyPtr(o) | Type::AssetPtr(o) => {
                    Some(o.as_uobject())
                }
                Type::Class(o) => Some(o.as_uobject()),
                Type::Primitive(_) | Type::Name | Type::String | Type::Text => None,
                Type::Enum { enum_type, .. } => Some(enum_type.as_uobject()),
            }
        }
    }

    #[derive(Debug)]
    pub struct ClassField {
        pub offset: u32,
        pub size: u32,
        pub name: String,
        pub ty: Type,
    }

    #[derive(Debug, Clone)]
    pub struct MethodParameter {
        pub name: String,
        pub ty: Type,
    }

    #[derive(Debug)]
    pub struct ClassMethod {
        pub name: String,
        pub full_name: String,
        pub is_native: bool,
        pub is_static: bool,
        pub parameters: Vec<ParameterKind>,
        pub return_type: Option<MethodParameter>,
    }

    #[derive(Debug)]
    pub struct Class {
        pub is_class: bool,
        pub size: u32,
        pub name: String,
        pub full_name: String,
        pub super_class: Option<UStruct>,
        pub fields: Vec<ClassField>,
        pub methods: Vec<ClassMethod>,
    }

    impl Class {
        pub fn rust_name(&self) -> &String {
            &self.name
        }

        pub fn referenced_types(&self) -> Vec<UObject> {
            let mut types = Vec::new();
            self.super_class.map(|obj| types.push(obj.as_uobject()));

            for field in &self.fields {
                types.extend(field.ty.referenced_type());
            }

            for method in &self.methods {
                types.extend(
                    method
                        .return_type
                        .as_ref()
                        .and_then(|param| param.ty.referenced_type()),
                );
                types.extend(
                    method
                        .parameters
                        .iter()
                        .map(|param| &param.as_param().ty)
                        .filter_map(|ty| ty.referenced_type()),
                );
            }

            types
        }
    }

    #[derive(Debug)]
    pub enum Types {
        Class(Class),
        Enum(Enum),
    }

    #[derive(Debug)]
    pub enum ParameterKind {
        Return(MethodParameter),
        Default(MethodParameter),
        Out(MethodParameter),
    }

    impl ParameterKind {
        pub fn as_param(&self) -> &MethodParameter {
            match self {
                ParameterKind::Return(param) => param,
                ParameterKind::Default(param) => param,
                ParameterKind::Out(param) => param,
            }
        }
    }

    #[derive(Debug)]
    pub struct Enum {
        pub name: String,
        pub values: Vec<String>,
    }

    impl Display for Enum {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            writeln!(f, "enum {} {{", self.name)?;
            for (i, name) in self.values.iter().enumerate() {
                writeln!(f, "\t{} = {},", name, i)?;
            }
            writeln!(f, "}}\n")
        }
    }
}