Upgrade to use snmalloc for rust compiler

.gitmodules

@@ -43,3 +43,7 @@
 	path = library/backtrace
 	url = https://github.com/rust-lang/backtrace-rs.git
 	shallow = true
+[submodule "library/snmalloc-edp/snmalloc"]
+	path = library/snmalloc-edp/snmalloc
+	url = https://github.com/microsoft/snmalloc.git
+        shallow = true

Cargo.lock

@@ -658,6 +658,15 @@ dependencies = [
  "rustc-semver",
 ]
 
+[[package]]
+name = "cmake"
+version = "0.1.50"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "collect-license-metadata"
 version = "0.1.0"
@@ -1178,6 +1187,12 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "elf"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4445909572dbd556c457c849c4ca58623d84b27c8fff1e74b0b4227d8b90d17b"
+
 [[package]]
 name = "elsa"
 version = "1.7.1"
@@ -5141,6 +5156,17 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 
+[[package]]
+name = "snmalloc-edp"
+version = "0.1.0"
+dependencies = [
+ "cc",
+ "cmake",
+ "compiler_builtins",
+ "elf",
+ "rustc-std-workspace-core",
+]
+
 [[package]]
 name = "socket2"
 version = "0.5.6"
@@ -5247,6 +5273,7 @@ dependencies = [
  "rand",
  "rand_xorshift",
  "rustc-demangle",
+ "snmalloc-edp",
  "std_detect",
  "unwind",
  "wasi",

library/snmalloc-edp/CMakeLists.txt

@@ -0,0 +1,10 @@
+cmake_minimum_required(VERSION 3.14)
+set(CMAKE_TRY_COMPILE_TARGET_TYPE STATIC_LIBRARY)
+project(snmalloc-edp CXX)
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED True)
+set(SNMALLOC_HEADER_ONLY_LIBRARY ON)
+add_subdirectory(snmalloc EXCLUDE_FROM_ALL)
+add_library(snmalloc-edp src/rust-sgx-snmalloc-shim.cpp)
+target_link_libraries(snmalloc-edp PRIVATE snmalloc_lib)
+target_compile_options(snmalloc-edp PRIVATE -nostdlib -ffreestanding -fno-exceptions -mrdrnd -fPIC)

library/snmalloc-edp/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "snmalloc-edp"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+build = "build.rs"
+
+[build-dependencies]
+cc = "1.0.86"
+cmake = "0.1.50"
+elf = "0.7"
+
+[dependencies]
+core = { version = "1.0.0", optional = true, package = "rustc-std-workspace-core" }
+compiler_builtins = { version = "0.1.0", optional = true }
+
+[features]
+docs = []
+rustc-dep-of-std = ["core", "compiler_builtins/rustc-dep-of-std"]

library/snmalloc-edp/build.rs

@@ -0,0 +1,71 @@
+use elf::ElfStream;
+use elf::endian::LittleEndian;
+use std::fs::{DirEntry, File};
+use std::path::{Path, PathBuf};
+
+fn files_in_dir(p: &Path) -> impl Iterator<Item = DirEntry> {
+    p.read_dir().unwrap().map(|e| e.unwrap()).filter(|e| e.file_type().unwrap().is_file())
+}
+
+fn main() {
+    let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").unwrap());
+
+    // # Use CMake to build the shim
+    let mut dst = cmake::build(".");
+    dst.push("build");
+    println!("cargo:rustc-link-search=native={}", dst.display());
+
+    // ideally, the cmake crate would have a way to output this
+    println!("cargo:rerun-if-changed=CMakeLists.txt");
+    println!("cargo:rerun-if-changed=src/rust-sgx-snmalloc-shim.cpp");
+
+    // # Extract the static library archive into a temporary directory
+    let mut objs = out_dir.clone();
+    objs.push("objs");
+    std::fs::create_dir_all(&objs).unwrap();
+    // clear existing files in the temp dir
+    for file in files_in_dir(&objs) {
+        std::fs::remove_file(file.path()).unwrap();
+    }
+
+    dst.push("libsnmalloc-edp.a");
+
+    let mut ar = cc::Build::new().get_archiver();
+    ar.args(&["x", "--output"]);
+    ar.arg(&objs);
+    ar.arg(dst);
+    assert!(ar.status().unwrap().success());
+
+    // # Read the symbols from the shim ELF object
+    assert_eq!(files_in_dir(&objs).count(), 1);
+    let f = files_in_dir(&objs).next().unwrap();
+    let mut elf = elf::ElfStream::<elf::endian::LittleEndian, _>::open_stream(File::open(f.path()).unwrap()).unwrap();
+    let (symtab, strtab) = elf.symbol_table().unwrap().unwrap();
+    let mut sn_alloc_size = None;
+    let mut sn_alloc_align = None;
+    for sym in symtab {
+        match strtab.get(sym.st_name as _).unwrap() {
+            "sn_alloc_size" => assert!(sn_alloc_size.replace(sym).is_none()),
+            "sn_alloc_align" => assert!(sn_alloc_align.replace(sym).is_none()),
+            _ => {}
+        }
+    }
+    let sn_alloc_size = sn_alloc_size.expect("sn_alloc_size");
+    let sn_alloc_align = sn_alloc_align.expect("sn_alloc_align");
+
+    let sn_alloc_size = get_u64_at_symbol(sn_alloc_size, &mut elf);
+    let sn_alloc_align = get_u64_at_symbol(sn_alloc_align, &mut elf);
+
+    // # Write the type
+    let contents = format!("#[repr(align({}), C)] pub struct Alloc {{ _0: [u8; {}] }}", sn_alloc_align, sn_alloc_size);
+    let mut alloc_type_rs = out_dir.clone();
+    alloc_type_rs.push("alloc-type.rs");
+    std::fs::write(alloc_type_rs, contents).unwrap();
+}
+
+fn get_u64_at_symbol(sym: elf::symbol::Symbol, elf: &mut ElfStream<LittleEndian, File>) -> u64 {
+    assert_eq!(sym.st_size, 8);
+    let (data, _) = elf.section_data(&elf.section_headers()[sym.st_shndx as usize].clone()).unwrap();
+    let data: &[u8; 8] = data.split_at(8).0.try_into().unwrap();
+    u64::from_le_bytes(*data)
+}

library/snmalloc-edp/src/lib.rs

@@ -0,0 +1,15 @@
+#![no_std]
+
+include!(concat!(env!("OUT_DIR"), "/alloc-type.rs"));
+
+#[link(name = "snmalloc-edp", kind = "static")]
+extern {
+    pub fn sn_global_init(heap_base: *const u8, heap_size: usize);
+    pub fn sn_thread_init(allocator: *mut Alloc);
+    pub fn sn_thread_cleanup(allocator: *mut Alloc);
+
+    pub fn sn_rust_alloc(alignment: usize, size: usize) -> *mut u8;
+    pub fn sn_rust_alloc_zeroed(alignment: usize, size: usize) -> *mut u8;
+    pub fn sn_rust_dealloc(ptr: *mut u8, alignment: usize, size: usize);
+    pub fn sn_rust_realloc(ptr: *mut u8, alignment: usize, old_size: usize, new_size: usize) -> *mut u8;
+}

library/snmalloc-edp/src/rust-sgx-snmalloc-shim.cpp

@@ -0,0 +1,181 @@
+// Copyright (c) Microsoft Corporation.
+// Copyright (c) Open Enclave SDK contributors.
+// Copyright (c) 2020 SchrodingerZhu
+// Copyright (c) Fortanix, Inc.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE
+
+#include <immintrin.h>
+#include <string.h>
+
+/***************************************************/
+/*** Imported symbols needed by snmalloc SGX PAL ***/
+/***************************************************/
+
+// from entry.S
+extern "C" size_t get_tcs_addr();
+
+// from Rust std
+extern "C" void __rust_print_err(const char* m, size_t s);
+extern "C" [[noreturn]] void __rust_abort();
+
+/*******************************************************/
+/*** Standard C functions needed by snmalloc SGX PAL ***/
+/*******************************************************/
+
+// definition needs to match GNU header
+extern "C" [[noreturn]] void abort() __THROW {
+    __rust_abort();
+}
+
+// definition needs to match GNU header and will not return an actual errno
+extern "C" inline int * __attribute_const__ __errno_location (void) __THROW {
+    static int errno;
+    return &errno;
+}
+
+/***********************************/
+/*** snmalloc SGX PAL definition ***/
+/***********************************/
+
+#define SNMALLOC_PROVIDE_OWN_CONFIG
+#define SNMALLOC_SGX
+#define SNMALLOC_USE_SMALL_CHUNKS
+#define SNMALLOC_MEMORY_PROVIDER PALEdpSgx
+
+#include "../snmalloc/src/snmalloc/pal/pal_noalloc.h"
+
+namespace snmalloc {
+void register_clean_up() {
+    // Unused on SGX
+    abort();
+}
+
+class EdpErrorHandler {
+  public:
+    static void print_stack_trace() {}
+
+    [[noreturn]] static void error(const char* const str) {
+        __rust_print_err(str, strlen(str));
+        abort();
+    }
+    static constexpr size_t address_bits = Aal::address_bits;
+    static constexpr size_t page_size = Aal::smallest_page_size;
+};
+
+using EdpBasePAL = PALNoAlloc<EdpErrorHandler>;
+
+class PALEdpSgx : public EdpBasePAL {
+   public:
+    const static size_t RAND_NUM_GEN_MAX_RETRIES = 64;
+    using ThreadIdentity = size_t;
+    static constexpr uint64_t pal_features = EdpBasePAL::pal_features | Entropy;
+
+    template <bool page_aligned = false>
+    static void zero(void *p, size_t size) noexcept {
+        memset(p, 0, size);
+    }
+
+    static inline uint64_t get_entropy64() {
+        for (size_t retry_count = 0; retry_count < RAND_NUM_GEN_MAX_RETRIES; retry_count++) {
+            long long unsigned int result;
+            if (_rdrand64_step(&result) == 1) {
+                return result;
+            }
+        }
+        EdpErrorHandler::error("no entropy available");
+    }
+
+    static inline ThreadIdentity get_tid() noexcept {
+        return (size_t)get_tcs_addr();
+    }
+};
+} // namespace snmalloc
+
+/**************************************/
+/*** Instantiation of the allocator ***/
+/**************************************/
+
+#include "../snmalloc/src/snmalloc/backend/fixedglobalconfig.h"
+#include "../snmalloc/src/snmalloc/snmalloc_core.h"
+
+using namespace snmalloc;
+
+using Globals = FixedRangeConfig<PALEdpSgx>;
+using Alloc = LocalAllocator<Globals>;
+
+/// Do global initialization for snmalloc. Should be called exactly once prior
+/// to any other snmalloc function calls.
+extern "C" void sn_global_init(void* heap_base, size_t heap_size) {
+    Globals::init(nullptr, heap_base, heap_size);
+}
+
+/// Construct a thread-local allocator object in place
+extern "C" void sn_thread_init(Alloc* allocator) {
+    new(allocator) Alloc();
+    allocator->init();
+}
+
+/// Destruct a thread-local allocator object in place
+extern "C" void sn_thread_cleanup(Alloc* allocator) {
+    allocator->teardown();
+    allocator->~Alloc();
+}
+
+extern "C" size_t sn_alloc_size = sizeof(Alloc);
+extern "C" size_t sn_alloc_align = alignof(Alloc);
+
+/// Return a pointer to a thread-local allocator object of size
+/// `sn_alloc_size` and alignment `sn_alloc_align`.
+extern "C" Alloc* __rust_get_thread_allocator();
+
+/******************************************************/
+/*** Rust-compatible shims for the global allocator ***/
+/******************************************************/
+
+extern "C" void *sn_rust_alloc(size_t alignment, size_t size) {
+    return __rust_get_thread_allocator()->alloc(aligned_size(alignment, size));
+}
+
+extern "C" void *sn_rust_alloc_zeroed(size_t alignment, size_t size) {
+    return __rust_get_thread_allocator()->alloc<YesZero>(
+        aligned_size(alignment, size));
+}
+
+extern "C" void sn_rust_dealloc(void *ptr, size_t alignment, size_t size) {
+    __rust_get_thread_allocator()->dealloc(ptr, aligned_size(alignment, size));
+}
+
+extern "C" void *sn_rust_realloc(void *ptr, size_t alignment, size_t old_size,
+                                 size_t new_size) {
+    size_t aligned_old_size = aligned_size(alignment, old_size),
+           aligned_new_size = aligned_size(alignment, new_size);
+    if (size_to_sizeclass_full(aligned_old_size).raw() ==
+        size_to_sizeclass_full(aligned_new_size).raw())
+        return ptr;
+    Alloc* allocator = __rust_get_thread_allocator();
+    void *p = allocator->alloc(aligned_new_size);
+    if (p) {
+        std::memcpy(p, ptr, old_size < new_size ? old_size : new_size);
+        allocator->dealloc(ptr, aligned_old_size);
+    }
+    return p;
+}