Fix issue causing missing schema entries

src/source2/client/mod.rs

@@ -1,3 +1,3 @@
-pub use input::*;
+pub use input::KeyButton;
 
 pub mod input;

src/source2/schema_system/mod.rs

@@ -82,7 +82,7 @@ pub struct SchemaAtomicTTF {
 #[derive(Pod)]
 #[repr(C)]
 pub struct SchemaBaseClassInfoData {
-    pub offset: u32,                          // 0x0000
+    pub offset: i32,                          // 0x0000
     pad_0004: [u8; 0x4],                      // 0x0004
     pub prev: Pointer64<SchemaClassInfoData>, // 0x0008
 }
@@ -91,9 +91,9 @@ pub struct SchemaBaseClassInfoData {
 #[repr(C)]
 pub struct SchemaClassFieldData {
     pub name: Pointer64<ReprCString>,                 // 0x0000
-    pub type_: Pointer64<SchemaType>,                 // 0x0008
-    pub offset: u32,                                  // 0x0010
-    pub num_metadata: u32,                            // 0x0014
+    pub schema_type: Pointer64<SchemaType>,           // 0x0008
+    pub offset: i32,                                  // 0x0010
+    pub num_metadata: i32,                            // 0x0014
     pub metadata: Pointer64<SchemaMetadataEntryData>, // 0x0018
 }
 
@@ -101,41 +101,43 @@ pub struct SchemaClassFieldData {
 #[derive(Pod)]
 #[repr(C)]
 pub struct SchemaClassInfoData {
-    pub base: Pointer64<SchemaClassInfoData>,                // 0x0000
-    pub name: Pointer64<ReprCString>,                        // 0x0008
-    pub module_name: Pointer64<ReprCString>,                 // 0x0010
-    pub size: u32,                                           // 0x0018
-    pub num_fields: u16,                                     // 0x001C
-    pub num_static_fields: u16,                              // 0x001E
-    pub num_static_metadata: u16,                            // 0x0020
-    pub alignment: u8,                                       // 0x0022
-    pub has_base_class: u8,                                  // 0x0023
-    pub total_class_size: u16,                               // 0x0024
-    pub derived_class_size: u16,                             // 0x0026
-    pub fields: Pointer64<SchemaClassFieldData>,             // 0x0028
-    pub static_fields: Pointer64<SchemaStaticFieldData>,     // 0x0030
-    pub base_classes: Pointer64<SchemaBaseClassInfoData>,    // 0x0038
-    pad_0040: [u8; 0x8],                                     // 0x0040
-    pub static_metadata: Pointer64<SchemaMetadataEntryData>, // 0x0048
-    pub type_scope: Pointer64<SchemaSystemTypeScope>,        // 0x0050
-    pub type_: Pointer64<SchemaType>,                        // 0x0058
-    pad_0060: [u8; 0x10],                                    // 0x0060
+    pub base: Pointer64<SchemaClassInfoData>,                  // 0x0000
+    pub name: Pointer64<ReprCString>,                          // 0x0008
+    pub module_name: Pointer64<ReprCString>,                   // 0x0010
+    pub size: i32,                                             // 0x0018
+    pub num_fields: i16,                                       // 0x001C
+    pub num_static_fields: i16,                                // 0x001E
+    pub num_static_metadata: i16,                              // 0x0020
+    pub alignment: u8,                                         // 0x0022
+    pub has_base_class: u8,                                    // 0x0023
+    pub total_class_size: i16,                                 // 0x0024
+    pub derived_class_size: i16,                               // 0x0026
+    pub fields: Pointer64<[SchemaClassFieldData]>,             // 0x0028
+    pub static_fields: Pointer64<[SchemaStaticFieldData]>,     // 0x0030
+    pub base_classes: Pointer64<SchemaBaseClassInfoData>,      // 0x0038
+    pad_0040: [u8; 0x8],                                       // 0x0040
+    pub static_metadata: Pointer64<[SchemaMetadataEntryData]>, // 0x0048
+    pub type_scope: Pointer64<SchemaSystemTypeScope>,          // 0x0050
+    pub schema_type: Pointer64<SchemaType>,                    // 0x0058
+    pad_0060: [u8; 0x10],                                      // 0x0060
 }
 
+#[rustfmt::skip]
 #[derive(Pod)]
 #[repr(C)]
 pub struct SchemaEnumInfoData {
-    pub base: Pointer64<SchemaEnumInfoData>,
-    pub name: Pointer64<ReprCString>,
-    pub module_name: Pointer64<ReprCString>,
-    pub alignment: u8,
-    pad_0019: [u8; 0x3],
-    pub size: u16,
-    pub num_static_metadata: u16,
-    pub enum_info: Pointer64<SchemaEnumeratorInfoData>,
-    pub static_metadata: Pointer64<SchemaMetadataEntryData>,
-    pub type_scope: Pointer64<SchemaSystemTypeScope>,
-    pad_0038: [u8; 0x10],
+    pub base: Pointer64<SchemaEnumInfoData>,                 // 0x0000
+    pub name: Pointer64<ReprCString>,                        // 0x0008
+    pub module_name: Pointer64<ReprCString>,                 // 0x0010
+    pub size: u8,                                            // 0x0018
+    pub alignment: u8,                                       // 0x0019
+    pad_001a: u16,                                           // 0x001A
+    pub num_enumerators: i16,                                // 0x001C
+    pub num_static_metadata: i16,                            // 0x001E
+    pub enumerators: Pointer64<[SchemaEnumeratorInfoData]>,  // 0x0020
+    pub static_metadata: Pointer64<SchemaMetadataEntryData>, // 0x0028
+    pub type_scope: Pointer64<SchemaSystemTypeScope>,        // 0x0030
+    pad_0038: [u8; 0x10],                                    // 0x0038
 }
 
 #[repr(C)]

src/source2/tier1/mod.rs

@@ -1,8 +1,8 @@
-pub use interface::*;
-pub use utl_memory::*;
-pub use utl_memory_pool::*;
-pub use utl_ts_hash::*;
-pub use utl_vector::*;
+pub use interface::InterfaceReg;
+pub use utl_memory::UtlMemory;
+pub use utl_memory_pool::UtlMemoryPoolBase;
+pub use utl_ts_hash::UtlTsHash;
+pub use utl_vector::UtlVector;
 
 pub mod interface;
 pub mod utl_memory;

src/source2/tier1/utl_memory_pool.rs

@@ -1,30 +1,49 @@
+use memflow::prelude::v1::*;
+
+#[repr(u32)]
+pub enum MemoryPoolGrowType {
+    None = 0,
+    Fast,
+    Slow,
+    RbTree,
+}
+
+#[derive(Pod)]
+#[repr(C)]
+pub struct Blob {
+    pub next: Pointer64<Blob>, // 0x0000
+    pub num_bytes: i32,        // 0x0008
+    pub data: [u8; 1],         // 0x000C
+    pad_000d: [u8; 3],         // 0x000D
+}
+
+#[derive(Pod)]
+#[repr(C)]
+pub struct FreeList {
+    pub next: Pointer64<FreeList>, // 0x0000
+}
+
 /// Represents an optimized pool memory allocator.
 #[repr(C)]
-pub struct UtlMemoryPool {
-    pub block_size: i32,       // 0x0000
-    pub blocks_per_blob: i32,  // 0x0004
-    pub grow_mode: i32,        // 0x0008
-    pub blocks_alloc: i32,     // 0x000C
-    pub block_alloc_size: i32, // 0x0010
-    pub peak_alloc: i32,       // 0x0014
+pub struct UtlMemoryPoolBase {
+    pub block_size: i32,                                // 0x0000
+    pub blocks_per_blob: i32,                           // 0x0004
+    pub grow_mode: MemoryPoolGrowType,                  // 0x0008
+    pub blocks_alloc: i32,                              // 0x000C
+    pub peak_alloc: i32,                                // 0x0010
+    pub alignment: u16,                                 // 0x0014
+    pub num_blobs: u16,                                 // 0x0016
+    pub free_list_tail: Pointer64<Pointer64<FreeList>>, // 0x0018
+    pub free_list_head: Pointer64<FreeList>,            // 0x0020
+    pad_0028: [u8; 0x44],                               // 0x0028
+    pub blob_head: Pointer64<Blob>,                     // 0x0070
+    pad_0078: [u8; 0x8],                                // 0x0078
 }
 
-impl UtlMemoryPool {
-    /// Returns the size of a block.
+impl UtlMemoryPoolBase {
+    /// Returns the total size of the memory pool.
     #[inline]
-    pub fn block_size(&self) -> i32 {
-        self.block_size
-    }
-
-    /// Returns the number of allocated blocks per blob.
-    #[inline]
-    pub fn count(&self) -> i32 {
-        self.blocks_per_blob
-    }
-
-    /// Returns the maximum number of allocated blocks.
-    #[inline]
-    pub fn peak_count(&self) -> i32 {
-        self.peak_alloc
+    pub fn size(&self) -> i32 {
+        (self.num_blobs as i32 * self.blocks_per_blob) * self.block_size
     }
 }

src/source2/tier1/utl_ts_hash.rs

@@ -1,6 +1,6 @@
 use memflow::prelude::v1::*;
 
-use super::UtlMemoryPool;
+use super::UtlMemoryPoolBase;
 
 use crate::error::Result;
 use crate::mem::IsNull;
@@ -17,79 +17,98 @@ unsafe impl<D: 'static> Pod for HashAllocatedBlob<D> {}
 
 #[repr(C)]
 pub struct HashBucket<D, K> {
-    pad_0000: [u8; 0x18],                                  // 0x0000,
-    pub first: Pointer64<HashFixedData<D, K>>,             // 0x0018
-    pub first_uncommitted: Pointer64<HashFixedData<D, K>>, // 0x0020
+    pad_0000: [u8; 0x18],                                          // 0x0000,
+    pub first: Pointer64<HashFixedDataInternal<D, K>>,             // 0x0018
+    pub first_uncommitted: Pointer64<HashFixedDataInternal<D, K>>, // 0x0020
 }
 
 #[repr(C)]
-pub struct HashFixedData<D, K> {
-    pub ui_key: K,                            // 0x0000
-    pub next: Pointer64<HashFixedData<D, K>>, // 0x0008
-    pub data: D,                              // 0x0010
+pub struct HashFixedDataInternal<D, K> {
+    pub ui_key: K,                                    // 0x0000
+    pub next: Pointer64<HashFixedDataInternal<D, K>>, // 0x0008
+    pub data: D,                                      // 0x0010
 }
 
-unsafe impl<D: 'static, K: 'static> Pod for HashFixedData<D, K> {}
+unsafe impl<D: 'static, K: 'static> Pod for HashFixedDataInternal<D, K> {}
 
 /// Represents a thread-safe hash table.
 #[repr(C)]
 pub struct UtlTsHash<D, const C: usize = 256, K = u64> {
-    pub entry_mem: UtlMemoryPool,               // 0x0000
-    pad_0018: [u8; 0x8],                        // 0x0018
-    pub blobs: Pointer64<HashAllocatedBlob<D>>, // 0x0020
-    pad_0028: [u8; 0x58],                       // 0x0028
-    pub buckets: [HashBucket<D, K>; C],         // 0x0080
-    pad_2880: [u8; 0x10],                       // 0x2880
+    pub entry_mem: UtlMemoryPoolBase,   // 0x0000
+    pub buckets: [HashBucket<D, K>; C], // 0x0080
+    pub needs_commit: bool,             // 0x2880
+    pad_2881: [u8; 0xF],                // 0x2881
 }
 
 impl<D: Pod + IsNull, const C: usize, K: Pod> UtlTsHash<D, C, K> {
+    /// Returns the number of allocated blocks.
+    #[inline]
+    pub fn blocks_alloc(&self) -> i32 {
+        self.entry_mem.blocks_alloc
+    }
+
+    /// Returns the size of a block.
+    #[inline]
+    pub fn block_size(&self) -> i32 {
+        self.entry_mem.block_size
+    }
+
+    /// Returns the maximum number of allocated blocks.
+    #[inline]
+    pub fn peak_count(&self) -> i32 {
+        self.entry_mem.peak_alloc
+    }
+
     /// Returns all elements in the hash table.
     pub fn elements(&self, process: &mut IntoProcessInstanceArcBox<'_>) -> Result<Vec<D>> {
-        // TODO: Refactor this.
+        let blocks_alloc = self.blocks_alloc() as usize;
+        let peak_alloc = self.peak_count() as usize;
 
-        let mut elements: Vec<_> = self
-            .buckets
-            .iter()
-            .flat_map(|bucket| {
-                let mut cur_element = bucket.first;
+        let mut allocated_list = Vec::with_capacity(peak_alloc);
+        let mut unallocated_list = Vec::with_capacity(blocks_alloc);
 
-                let mut list = Vec::new();
+        for bucket in &self.buckets {
+            let mut cur_element = bucket.first_uncommitted;
 
-                while !cur_element.is_null() {
-                    if let Ok(element) = cur_element.read(process) {
-                        if !element.data.is_null() {
-                            list.push(element.data);
-                        }
+            while !cur_element.is_null() {
+                let element = cur_element.read(process)?;
 
-                        cur_element = element.next;
-                    }
+                if !element.data.is_null() {
+                    unallocated_list.push(element.data);
                 }
 
-                list
-            })
-            .collect();
-
-        if let Ok(blob) = self.blobs.read(process) {
-            let mut unallocated_data = blob.next;
-
-            if !unallocated_data.is_null() {
-                if !blob.data.is_null() {
-                    elements.push(blob.data);
+                // Check if we have too many elements.
+                if unallocated_list.len() >= blocks_alloc {
+                    break;
                 }
 
-                while !unallocated_data.is_null() {
-                    if let Ok(element) = unallocated_data.read(process) {
-                        if !element.data.is_null() {
-                            elements.push(element.data);
-                        }
-
-                        unallocated_data = element.next;
-                    }
-                }
+                cur_element = element.next;
             }
         }
 
-        Ok(elements)
+        let mut cur_blob =
+            Pointer64::<HashAllocatedBlob<D>>::from(self.entry_mem.free_list_head.address());
+
+        while !cur_blob.is_null() {
+            let blob = cur_blob.read(process)?;
+
+            if !blob.data.is_null() {
+                allocated_list.push(blob.data);
+            }
+
+            // Check if we have too many elements.
+            if allocated_list.len() >= peak_alloc {
+                break;
+            }
+
+            cur_blob = blob.next;
+        }
+
+        Ok(if unallocated_list.len() > allocated_list.len() {
+            unallocated_list
+        } else {
+            allocated_list
+        })
     }
 }