Draft: Add custom allocator support (mimalloc)

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["mozjs-sys", "mozjs"]
+members = ["mozjs-sys", "mozjs", "examples/embedder_allocator"]
 resolver = "2"
 
 [workspace.package]

examples/embedder_allocator/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "embedder_allocator"
+version = "0.1.0"
+repository.workspace = true
+license.workspace = true
+edition.workspace = true
+description = "Example usage of mozjs, with the embedder (this crate) providing a custom allocator for mozjs"
+publish = false
+
+[dependencies]
+mozjs = { path =  "../../mozjs", features = ["custom-alloc"] }
+mimalloc = "0.1.48"

examples/embedder_allocator/Readme.md

@@ -0,0 +1,2 @@
+# Example: Using mozjs with an embedder-provided allocator
+

examples/embedder_allocator/build.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+# This file
+
+import os
+import pathlib
+import subprocess
+import sys
+from typing import Mapping
+
+def create_env() -> Mapping[str, str]:
+    env = os.environ.copy()
+    mimalloc_include_dir = pathlib.Path(__file__).parent.joinpath('mimalloc/include')
+    assert mimalloc_include_dir.is_dir(), "Could not find mimalloc include directory"
+    env['SERVO_CUSTOM_ALLOC_INCLUDE_DIR'] = mimalloc_include_dir.as_posix()
+    return env
+
+
+def main():
+    completed_process = subprocess.run(sys.argv[1:], env=create_env())
+    sys.exit(completed_process.returncode)
+
+
+if __name__ == '__main__':
+    main()

examples/embedder_allocator/build.rs

@@ -0,0 +1,4 @@
+fn main() {
+    // todo: Should we do this here or in the mozjs-sys build-script?
+    println!("cargo:rustc-link-lib=mimalloc");
+}

examples/embedder_allocator/mimalloc/include/servo_embedder_malloc_prefix.h

@@ -0,0 +1,5 @@
+#pragma once
+
+/// Defines the prefix for all malloc functions, i.e.
+/// mi_malloc, mi_calloc, mi_realloc, mi_free etc.
+#define SERVO_EMBEDDER_MALLOC_PREFIX mi_

examples/embedder_allocator/mimalloc/include/servo_embedder_memory_wrap.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <mimalloc.h>
+#include "servo_embedder_malloc_prefix.h"
+
+#define SERVO_CONCAT(x, y) x ## y
+#define SERVO_CONCAT2(x, y) SERVO_CONCAT(x, y)
+
+#define mozmem_malloc_impl(fn) SERVO_CONCAT2(SERVO_EMBEDDER_MALLOC_PREFIX, fn)
+#define mozmem_dup_impl(fn) SERVO_CONCAT2(SERVO_EMBEDDER_MALLOC_PREFIX, fn)

examples/embedder_allocator/src/main.rs

@@ -0,0 +1,76 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use std::ptr;
+use std::sync::mpsc::channel;
+use std::thread;
+
+use mozjs::jsapi::GCContext;
+use mozjs::jsapi::JSCLASS_FOREGROUND_FINALIZE;
+use mozjs::jsapi::{JSClass, JSClassOps, JSObject, OnNewGlobalHookOption};
+use mozjs::realm::AutoRealm;
+use mozjs::rooted;
+use mozjs::rust::wrappers2::{JS_NewGlobalObject, JS_NewObject};
+use mozjs::rust::{JSEngine, RealmOptions, Runtime, SIMPLE_GLOBAL_CLASS};
+
+fn main() {
+    println!("Hello, world!");
+    let engine = JSEngine::init().expect("Could not init JSEngine");
+    println!("JSEngine initialized");
+    let mut runtime = Runtime::new(engine.handle());
+    println!("Runtime initialized");
+    let context = runtime.cx();
+    let h_option = OnNewGlobalHookOption::FireOnNewGlobalHook;
+    let c_option = RealmOptions::default();
+
+    unsafe {
+        rooted!(&in(context) let global = JS_NewGlobalObject(
+            context,
+            &SIMPLE_GLOBAL_CLASS,
+            ptr::null_mut(),
+            h_option,
+            &*c_option,
+        ));
+        let mut realm = AutoRealm::new_from_handle(context, global.handle());
+        let context = &mut realm;
+        rooted!(&in(context) let _object = JS_NewObject(context, &CLASS as *const _));
+    }
+
+    let parent = runtime.prepare_for_new_child();
+    let (sender, receiver) = channel();
+    thread::spawn(move || {
+        let runtime = unsafe { Runtime::create_with_parent(parent) };
+        assert!(Runtime::get().is_some());
+        drop(runtime);
+        let _ = sender.send(());
+    });
+    let _ = receiver.recv();
+    println!("Example ran without issues...");
+}
+
+unsafe extern "C" fn finalize(_fop: *mut GCContext, _object: *mut JSObject) {
+    assert!(Runtime::get().is_some());
+}
+
+static CLASS_OPS: JSClassOps = JSClassOps {
+    addProperty: None,
+    delProperty: None,
+    enumerate: None,
+    newEnumerate: None,
+    resolve: None,
+    mayResolve: None,
+    finalize: Some(finalize),
+    call: None,
+    construct: None,
+    trace: None,
+};
+
+static CLASS: JSClass = JSClass {
+    name: c"EventTargetPrototype".as_ptr(),
+    flags: JSCLASS_FOREGROUND_FINALIZE,
+    cOps: &CLASS_OPS as *const JSClassOps,
+    spec: ptr::null(),
+    ext: ptr::null(),
+    oOps: ptr::null(),
+};

mozjs-sys/Cargo.toml

@@ -16,6 +16,8 @@ name = "mozjs_sys"
 doctest = false
 
 [features]
+default = []
+custom-alloc = []
 debugmozjs = []
 profilemozjs = []
 jit = []

mozjs-sys/build.rs

@@ -138,6 +138,7 @@ fn main() {
 
 fn build_spidermonkey(build_dir: &Path) {
     let target = env::var("TARGET").unwrap();
+    let cargo_manifest_dir = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap());
     let make;
 
     #[cfg(windows)]
@@ -210,24 +211,38 @@ fn build_spidermonkey(build_dir: &Path) {
     }
 
     cppflags.push(get_cc_rs_env_os("CPPFLAGS").unwrap_or_default());
-    cmd.env("CPPFLAGS", cppflags);
 
     if let Some(makeflags) = env::var_os("CARGO_MAKEFLAGS") {
         cmd.env("MAKEFLAGS", makeflags);
     }
 
     let mut cxxflags = vec![];
 
+    if env::var_os("CARGO_FEATURE_CUSTOM_ALLOC").is_some() {
+        let mut flags = vec![];
+        let include_dir_str = env::var("SERVO_CUSTOM_ALLOC_INCLUDE_DIR").expect("Required variable not set with feature custom-alloc");
+        let include_dir = Path::new(&include_dir_str);
+        assert!(include_dir.is_dir(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must be set to a valid directory");
+        assert!(include_dir.join("servo_embedder_memory_wrap.h").is_file(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must contain header `servo_embedder_memory_wrap.h`");
+        flags.push(format!("-I{}", &include_dir_str.replace("\\", "/")));
+        flags.push("-DSERVO_EMBEDDER_MEMORY".to_string());
+        println!("cargo:rerun-if-changed={}", include_dir_str);
+
+        cppflags.extend(flags.iter().map(|s| OsString::from(s)));
+        cxxflags.extend(flags);
+    }
+
     if target.contains("apple") || target.contains("freebsd") || target.contains("ohos") {
         cxxflags.push(String::from("-stdlib=libc++"));
     }
 
+    cmd.env("CPPFLAGS", cppflags);
+
     let base_cxxflags = env::var("CXXFLAGS").unwrap_or_default();
     let mut cxxflags = cxxflags.join(" ");
     cxxflags.push_str(&base_cxxflags);
     cmd.env("CXXFLAGS", cxxflags);
 
-    let cargo_manifest_dir = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap());
     let result = cmd
         .args(&["-R", "-f"])
         .arg(cargo_manifest_dir.join("makefile.cargo"))
@@ -292,6 +307,16 @@ fn build(build_dir: &Path, target: BuildTarget) {
 
     build.flag(include_file_flag(build.get_compiler().is_like_msvc()));
     build.flag(&js_config_path(build_dir));
+    let cargo_manifest_dir = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap());
+
+    if env::var_os("CARGO_FEATURE_CUSTOM_ALLOC").is_some() {
+        let include_dir_str = env::var("SERVO_CUSTOM_ALLOC_INCLUDE_DIR").expect("Required variable not set with feature custom-alloc");
+        let include_dir = Path::new(&include_dir_str);
+        assert!(include_dir.is_dir(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must be set to a valid directory");
+        assert!(include_dir.join("servo_embedder_memory_wrap.h").is_file(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must contain header `servo_embedder_memory_wrap.h`");
+        build.include(include_dir);
+        build.define("SERVO_EMBEDDER_MEMORY", "");
+    }
 
     for path in target.include_paths(build_dir) {
         build.include(path);
@@ -313,6 +338,8 @@ fn build_bindings(build_dir: &Path, target: BuildTarget) {
     config &= !CodegenConfig::DESTRUCTORS;
     config &= !CodegenConfig::METHODS;
 
+    let cargo_manifest_dir = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap());
+
     let mut builder = bindgen::builder()
         .rust_target(minimum_rust_target())
         .header(target.path())
@@ -332,6 +359,21 @@ fn build_bindings(build_dir: &Path, target: BuildTarget) {
             .clang_arg(env::var("WASI_SYSROOT").unwrap().to_string());
     }
 
+    let custom_alloc_flags = if env::var_os("CARGO_FEATURE_CUSTOM_ALLOC").is_some() {
+        let mut flags = vec![];
+        let include_dir_str = env::var("SERVO_CUSTOM_ALLOC_INCLUDE_DIR").expect("Required variable not set with feature custom-alloc");
+        let include_dir = Path::new(&include_dir_str);
+        assert!(include_dir.is_dir(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must be set to a valid directory");
+        assert!(include_dir.join("servo_embedder_memory_wrap.h").is_file(), "SERVO_CUSTOM_ALLOC_INCLUDE_DIR must contain header `servo_embedder_memory_wrap.h`");
+        flags.push(format!("-I{}", &include_dir_str.replace("\\", "/")));
+        flags.push("-DSERVO_EMBEDDER_MEMORY".to_string());
+        flags
+    } else {
+      vec![]
+    };
+    builder = builder.clang_args(custom_alloc_flags);
+
+
     if target == BuildTarget::JSGlue {
         builder = builder
             .parse_callbacks(Box::new(JSGlueCargoCallbacks::default()))
@@ -432,6 +474,8 @@ fn link_static_lib_binaries(build_dir: &Path) {
         // needing to use the WASI-SDK's clang for linking, which is annoying.
         println!("cargo:rustc-link-lib=stdc++")
     }
+    // TODO: link against lib from env
+    // println!("cargo:rustc-link-lib=mimalloc");
 
     if target.contains("wasi") {
         println!("cargo:rustc-link-lib=wasi-emulated-getpid");
@@ -463,6 +507,9 @@ fn should_build_from_source() -> bool {
     } else if env::var_os("CARGO_FEATURE_INTL").is_none() {
         println!("intl feature is disabled. Building from source directly.");
         true
+    } else if env::var_os("CARGO_FEATURE_CUSTOM_ALLOC").is_some() {
+        println!("custom-alloc feature is enabled. Building from source directly.");
+        true
     } else if !env::var_os("CARGO_FEATURE_JIT").is_some() {
         println!("jit feature is NOT enabled. Building from source directly.");
         true