eh

arihant2math · arihant2math · commit 226c1d544730 · 2025-04-11T22:37:19.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -40,9 +40,6 @@ dirs = { package = "dirs-next", version = "2.0.0" }
 env_logger = { version = "0.9.0", default-features = false, features = ["atty", "termcolor"] }
 flamescope = { version = "0.1.2", optional = true }
 
-[target.'cfg(all(any(target_os = "linux", target_os = "macos", target_os = "windows"), not(any(target_env = "musl", target_env = "sgx"))))'.dependencies]
-mimalloc = "0.1"
-
 [target.'cfg(windows)'.dependencies]
 libc = { workspace = true }
 
diff --git a/src/main.rs b/src/main.rs
@@ -1,13 +1,7 @@
-#[cfg(all(
-    any(target_os = "linux", target_os = "macos", target_os = "windows"),
-    not(any(target_env = "musl", target_env = "sgx"))
-))]
-mod _alloc {
-    use mimalloc::MiMalloc;
+use rustpython_vm::RustPythonAllocator;
 
-    #[global_allocator]
-    static GLOBAL: MiMalloc = MiMalloc;
-}
+#[global_allocator]
+static ALLOCATOR: RustPythonAllocator = RustPythonAllocator::new();
 
 pub fn main() -> std::process::ExitCode {
     rustpython::run(|_vm| {})
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -108,6 +108,7 @@ widestring = { workspace = true }
 [target.'cfg(all(any(target_os = "linux", target_os = "macos", target_os = "windows"), not(any(target_env = "musl", target_env = "sgx"))))'.dependencies]
 libffi = { workspace = true, features = ["system"] }
 libloading = "0.8"
+mimalloc = "0.1"
 
 [target.'cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))'.dependencies]
 num_cpus = "1.13.1"
diff --git a/vm/src/alloc.rs b/vm/src/alloc.rs
@@ -0,0 +1,129 @@
+//! Configurable allocator for RustPython.
+//! Currently it supports `mimalloc` and `system` allocators,
+//! whereas cpython uses `pymalloc`` for most operations.
+
+#[cfg(all(
+    any(target_os = "linux", target_os = "macos", target_os = "windows"),
+    not(any(target_env = "musl", target_env = "sgx"))
+))]
+mod inner {
+    use std::alloc::{GlobalAlloc, Layout, System};
+
+    pub enum RustPythonAllocator {
+        System(System),
+        Mimalloc(mimalloc::MiMalloc),
+    }
+
+    impl RustPythonAllocator {
+        pub fn new(allocator: &str) -> Self {
+            match allocator {
+                "system" | "malloc" => RustPythonAllocator::System(System),
+                "pymalloc" | "mimalloc" | "default" => {
+                    RustPythonAllocator::Mimalloc(mimalloc::MiMalloc)
+                }
+                _ => RustPythonAllocator::System(System),
+            }
+        }
+    }
+
+    unsafe impl GlobalAlloc for RustPythonAllocator {
+        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+            unsafe {
+                match self {
+                    RustPythonAllocator::System(system) => system.alloc(layout),
+                    RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.alloc(layout),
+                }
+            }
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+            unsafe {
+                match self {
+                    RustPythonAllocator::System(system) => system.dealloc(ptr, layout),
+                    RustPythonAllocator::Mimalloc(mimalloc) => mimalloc.dealloc(ptr, layout),
+                }
+            }
+        }
+    }
+}
+
+#[cfg(not(all(
+    any(target_os = "linux", target_os = "macos", target_os = "windows"),
+    not(any(target_env = "musl", target_env = "sgx"))
+)))]
+mod inner {
+    use std::alloc::{GlobalAlloc, Layout, System};
+
+    pub enum RustPythonAllocator {
+        System(System),
+    }
+
+    impl RustPythonAllocator {
+        pub fn new(_allocator: &str) -> Self {
+            RustPythonAllocator::System(System)
+        }
+    }
+
+    impl GlobalAlloc for RustPythonAllocator {
+        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+            unsafe { self.0.alloc(layout) }
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+            unsafe { self.0.dealloc(ptr, layout) }
+        }
+    }
+}
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::UnsafeCell;
+
+pub use inner::RustPythonAllocator as InternalAllocator;
+
+pub struct RustPythonAllocator {
+    inner: UnsafeCell<InternalAllocator>,
+}
+
+unsafe impl Send for RustPythonAllocator {}
+unsafe impl Sync for RustPythonAllocator {}
+
+impl RustPythonAllocator {
+    /// Create a new allocator based on the PYTHONMALLOC environment variable
+    /// or the default allocator if not set.
+    /// If this is not intended, use [`InternalAllocator::new`] directly.
+    pub const fn new() -> Self {
+        Self {
+            inner: UnsafeCell::new(None),
+        }
+    }
+}
+
+impl RustPythonAllocator {
+    unsafe fn get_or_init(&self) -> &InternalAllocator {
+        unsafe {
+            let inner = self.inner.get();
+            if *inner.is_none() {
+                let env = std::env::var("PYTHONMALLOC").unwrap_or_default(); 
+                let allocator = InternalAllocator::new(&env);
+                *inner = allocator;
+            }
+            inner as *const InternalAllocator as _
+        }
+    }
+}
+
+unsafe impl GlobalAlloc for RustPythonAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        unsafe {
+            let inner = self.get_or_init();
+            inner.alloc(layout)
+        }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        unsafe {
+            let inner = self.get_or_init();
+            inner.dealloc(ptr, layout)
+        }
+    }
+}
diff --git a/vm/src/lib.rs b/vm/src/lib.rs
@@ -40,6 +40,7 @@ pub use rustpython_derive::*;
 #[macro_use]
 pub(crate) mod macros;
 
+mod alloc;
 mod anystr;
 pub mod buffer;
 pub mod builtins;
@@ -85,6 +86,7 @@ pub mod warn;
 #[cfg(windows)]
 pub mod windows;
 
+pub use self::alloc::RustPythonAllocator;
 pub use self::compiler::source;
 pub use self::convert::{TryFromBorrowedObject, TryFromObject};
 pub use self::object::{
