tracking works

Author: kermanx

Cargo.lock

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["crates/*", "packages/napi", "tasks/benchmark"]
+members = ["crates/*", "packages/napi", "tasks/benchmark", "oxc_allocator"]
 default-members = ["crates/*"]
 
 [workspace.package]
@@ -11,11 +11,12 @@ homepage = "https://github.com/kermanx/jsshaker"
 keywords = ["JavaScript", "TypeScript", "oxc"]
 license = "MIT"
 repository = "https://github.com/kermanx/jsshaker"
-rust-version = "1.89.0"                                              # Should sync with oxc
+rust-version = "1.89.0"                                            # Should sync with oxc
 description = "Code size optimizer for JavaScript"
 
 [workspace.dependencies]
 jsshaker = { path = "crates/jsshaker" }
+oxc_allocator = { path = "oxc_allocator" }
 
 bitflags = "2.6.0"
 clap = "4.5.20"
@@ -25,7 +26,6 @@ flamescope = "0.1.3"
 flate2 = "1.1.0"
 line-index = "0.1.2"
 oxc = "0.107.0"
-oxc_allocator = { "version" = "0.107.0", features = ["track_allocations"] }
 oxc_ast_visit = "0.107.0"
 oxc_ecmascript = "0.107.0"
 oxc_index = "4.1.0"

Cargo.toml

@@ -31,6 +31,7 @@ flame = { workspace = true, optional = true }
 flamescope = { workspace = true, optional = true }
 flate2 = { workspace = true }
 oxc = { workspace = true, features = ["codegen", "semantic", "minifier"] }
+oxc_allocator = { workspace = true, features = ["track_allocations"] }
 oxc_ast_visit = { workspace = true }
 oxc_ecmascript = { workspace = true }
 oxc_index = { workspace = true }

crates/jsshaker/Cargo.toml

@@ -70,6 +70,9 @@ pub fn tree_shake<F: Vfs + 'static>(options: JsShakerOptions<F>, entry: String)
       conditional_data,
       ..
     } = unsafe { &mut *(&mut analyzer as *mut _) };
+
+    allocator.print_allocation_stats_types();
+
     let mangler = Rc::new(RefCell::new(mangler));
     let mut codegen_return = FxHashMap::default();
     for module_info in mem::take(&mut modules.modules) {

crates/jsshaker/src/lib.rs

@@ -99,7 +99,7 @@ impl Alloc for Bump {
         // `track_allocations` feature, so should never be compiled in production code.
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
         unsafe {
-            crate::tracking::get_stats_ref(self).record_allocation();
+            crate::tracking::get_stats_ref(self).record_allocation::<Layout>(0);
         }
 
         self.alloc_layout(layout)

oxc_allocator/src/alloc.rs

@@ -301,7 +301,7 @@ impl Allocator {
         const { assert!(!std::mem::needs_drop::<T>(), "Cannot allocate Drop type in arena") };
 
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.stats.record_allocation();
+        self.stats.record_allocation::<T>(std::mem::size_of::<T>());
 
         self.bump.alloc(val)
     }
@@ -325,7 +325,7 @@ impl Allocator {
     #[inline(always)]
     pub fn alloc_str<'alloc>(&'alloc self, src: &str) -> &'alloc str {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.stats.record_allocation();
+        self.stats.record_allocation::<&str>(std::mem::size_of_val(src));
 
         self.bump.alloc_str(src)
     }
@@ -348,7 +348,7 @@ impl Allocator {
     #[inline(always)]
     pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> &mut [T] {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.stats.record_allocation();
+        self.stats.record_allocation::<&[T]>(std::mem::size_of_val(src));
 
         self.bump.alloc_slice_copy(src)
     }
@@ -363,7 +363,7 @@ impl Allocator {
     /// Panics if reserving space matching `layout` fails.
     pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.stats.record_allocation();
+        self.stats.record_allocation::<Layout>(0);
 
         self.bump.alloc_layout(layout)
     }
@@ -416,7 +416,7 @@ impl Allocator {
         );
 
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.stats.record_allocation();
+        self.stats.record_allocation::<[&str; N]>(total_len);
 
         // Create actual `&str` in a separate function, to ensure that `alloc_concat_strs_array`
         // is inlined, so that compiler has knowledge to remove the overflow checks above.

oxc_allocator/src/allocator.rs

@@ -29,14 +29,22 @@ pub struct AllocationStats {
     num_alloc: Cell<usize>,
     /// Number of reallocations
     num_realloc: Cell<usize>,
+
+    types: std::cell::RefCell<std::collections::HashMap<&'static str, (usize, usize)>>
 }
 
 impl AllocationStats {
     /// Record that an allocation was made.
-    pub(crate) fn record_allocation(&self) {
+    pub(crate) fn record_allocation<T>(&self, size: usize) {
         // Counter maxes out at `usize::MAX`, but if there's that many allocations,
         // the exact number is not important
         self.num_alloc.set(self.num_alloc.get().saturating_add(1));
+
+        let type_name = std::any::type_name::<T>();
+        let mut map = self.types.borrow_mut();
+        let entry = map.entry(type_name).or_insert((0, 0));
+        entry.0 += 1;    // 数量 +1
+        entry.1 += size; // 总字节数 + size
     }
 
     /// Record that a reallocation was made.
@@ -61,6 +69,15 @@ impl Allocator {
         let num_realloc = self.stats.num_realloc.get();
         (num_alloc, num_realloc)
     }
+
+    pub fn print_allocation_stats_types(&self) {
+      let map = self.stats.types.borrow();
+      let mut entries: Vec<_> = map.iter().collect();
+      entries.sort_by(|a, b| b.1.1.cmp(&a.1.1));
+      for (name, (count, size)) in entries {
+        println!("{size}\t{count}\t{name}");
+      }
+    }
 }
 
 /// Get reference to [`AllocationStats`] for a [`Bump`].