// Copyright 2017 the authors. See the 'Copyright and license' section of the
// README.md file at the top-level directory of this repository.
//
// Licensed under the Apache License, Version 2.0 (the LICENSE file). This file
// may not be copied, modified, or distributed except according to those terms.

//! Bindings for the C `malloc` API to Rust allocators.
//!
//! This crate provides a mechanism to construct a C allocator - an implementation of `malloc`,
//! `free`, and related functions - that is backed by a Rust allocator (an implementation of the
//! `Alloc` trait).
//!
//! In order to create bindings, two things must be provided: an implementation of the `Alloc`
//! trait, and an implementation of the `LayoutFinder` trait (defined in this crate). Since the C
//! API does not provide size or alignment on `free`, but the Rust `Alloc` API requires both size
//! and alignment on `dealloc`, a mapping must be maintained between allocated objects and those
//! objects' size and alignment. The `LayoutFinder` provides this functionality.

#![no_std]
#![feature(allocator_api)]
#![feature(alloc)]
#![feature(core_intrinsics)]
#![feature(const_fn)]

extern crate alloc;
extern crate libc;
extern crate errno;
extern crate sysconf;
// lazy_static's macros are only used in the macros we define, so if no macros are called (which is
// the case when compiling this crate on its own), then lazy_static's macros (and thus the
// #[macro_use] attribute) will appear unused.  Due to an issue with clippy
// (https://rust-lang-nursery.github.io/rust-clippy/master/index.html#useless_attribute), this
// allow(unused_imports) directive will be seen as useless, so we suppress the useless_attribute
// warning as well.
#[cfg_attr(feature = "cargo-clippy", allow(useless_attribute))]
#[allow(unused_imports)]
#[macro_use]
extern crate lazy_static;
use alloc::allocator::{Alloc, AllocErr, Layout};

use libc::{c_void, size_t};
use core::{mem, ptr};
use core::cmp::max;

const WORD_SIZE: usize = mem::size_of::<*mut c_void>();

/// A mechanism for mapping allocated objects to their `Layout`s.
///
/// A `LayoutFinder` is an object that can store and look up the `Layout` associated with an
/// allocated object. In the functions generated by this crate, newly-allocated objects will be
/// inserted into a global `LayoutFinder` object, and this `LayoutFinder` will be used to look up
/// the `Layout`s associated with objects passed to `free` and other functions.
pub trait LayoutFinder {
    /// Get the `Layout` associated with an allocated object.
    ///
    /// `get_layout` is passed a pointer to an allocated object, and it returns a `Layout`
    /// describing that object. `ptr` is guaranteed to be an object previously allocated using one
    /// of the various C allocation functions.
    fn get_layout(&self, ptr: *mut u8) -> Layout;

    /// Insert a new object to `Layout` mapping.
    ///
    /// `insert_layout` is passed a pointer to a newly-allocated object and a `Layout` describing
    /// that object, and it stores this mapping. `insert_layout` is called immediately after
    /// allocation in all of the C allocation functions.
    ///
    /// The default implementation of `insert_layout` is a no-op, as some allocators may already
    /// keep track of the information necessary to implement `get_layout` internally.
    fn insert_layout(&self, _ptr: *mut u8, _layout: Layout) {}

    /// Delete an existing object to `Layout` mapping.
    ///
    /// `delete_layout` is passed a pointer to an object whose mapping has previously been
    /// inserted, and it deletes this mapping. `delete_layout` is called immediately after
    /// deallocation in all of the C deallocation functions.
    ///
    /// The default implementation of `delete_layout` is a no-op, as some allocators may already
    /// keep track of the information necessary to implement `get_layout` internally.
    fn delete_layout(&self, _ptr: *mut u8) {}
}

/// A wrapper for a Rust allocator providing C bindings.
///
/// `Malloc` wraps existing `Alloc` and `LayoutFinder` instances and provides methods for each of
/// the various C allocation functions. Most users should simply call the `define_malloc` or
/// `define_malloc_lazy_static` macros, which take care of constructing a `Malloc` instance and
/// defining the various `extern "C"` functions of the C allocation API. Users who wish to expose
/// only a subset of this API will need to instantiate a `Malloc` and define the `extern "C"`
/// functions manually.
pub struct Malloc<A, L: LayoutFinder>
    where for<'a> &'a A: Alloc
{
    alloc: A,
    layout_finder: L,
}

impl<A, L: LayoutFinder> Malloc<A, L>
    where for<'a> &'a A: Alloc
{
    /// Construct a new `Malloc`.
    ///
    /// `new` constructs a new `Malloc` using the provided allocator and `LayoutFinder`. Since C
    /// allocation functions can be called from many threads simultaneously, the allocator must be
    /// thread-safe. Thus, `A` (the type of the `alloc` parameter) isn't required to implement
    /// `Alloc`. Instead, `&A` must implement `Alloc` so that `Alloc`'s methods can be called
    /// concurrently.
    pub const fn new(alloc: A, layout_finder: L) -> Malloc<A, L> {
        Malloc {
            alloc,
            layout_finder,
        }
    }

    /// The C `malloc` function.
    pub unsafe fn malloc(&self, size: size_t) -> *mut c_void {
        if size == 0 {
            return ptr::null_mut();
        }

        // According to the posix_memalign manpage, "The glibc malloc(3) always returns 8-byte
        // aligned memory addresses..." Thus, we round up the size of allocations to 8 bytes in
        // order guarantee that 8 is a valid alignment (since Layout requires that the size is a
        // multiple of the alignment).
        let size = max(size, 8);

        let layout = Layout::from_size_align(size as usize, 8).unwrap();
        match (&self.alloc).alloc(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The C `free` function.
    pub unsafe fn free(&self, ptr: *mut c_void) {
        if ptr.is_null() {
            return;
        }

        let layout = self.layout_finder.get_layout(ptr as *mut u8);
        self.layout_finder.delete_layout(ptr as *mut u8);
        (&self.alloc).dealloc(ptr as *mut u8, layout);
    }

    /// The obsolete C `cfree` function.
    pub unsafe fn cfree(&self, ptr: *mut c_void) {
        // See https://linux.die.net/man/3/cfree
        self.free(ptr)
    }

    /// The C `calloc` function.
    pub unsafe fn calloc(&self, nmemb: size_t, size: size_t) -> *mut c_void {
        if nmemb == 0 || size == 0 {
            return ptr::null_mut();
        }

        // According to the posix_memalign manpage, "The glibc malloc(3) always returns 8-byte
        // aligned memory addresses..." Thus, we round up the size of allocations to 8 bytes in
        // order guarantee that 8 is a valid alignment (since Layout requires that the size is a
        // multiple of the alignment).
        let size = max(size, 8);

        let layout = Layout::from_size_align(nmemb * size as usize, 8).unwrap();
        match (&self.alloc).alloc_zeroed(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The obsolete C `valloc` function.
    pub unsafe fn valloc(&self, size: size_t) -> *mut c_void {
        if size == 0 {
            return ptr::null_mut();
        }

        let layout = Layout::from_size_align(size as usize, sysconf::page::pagesize()).unwrap();
        match (&self.alloc).alloc_zeroed(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The obsolete C `pvalloc` function (only implemented on Linux).
    #[cfg(target_os = "linux")]
    pub unsafe fn pvalloc(&self, size: size_t) -> *mut c_void {
        // See http://man7.org/linux/man-pages/man3/posix_memalign.3.html

        if size == 0 {
            return ptr::null_mut();
        }

        // TODO: round size up to the next multiple of the page size.

        let layout = Layout::from_size_align(size as usize, sysconf::page::pagesize()).unwrap();
        match (&self.alloc).alloc_zeroed(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The C `realloc` function.
    pub unsafe fn realloc(&self, ptr: *mut c_void, size: size_t) -> *mut c_void {
        // See http://man7.org/linux/man-pages/man3/malloc.3.html,
        // http://www.manpagez.com/man/3/malloc/osx-10.6.php

        if ptr.is_null() {
            return self.malloc(size);
        }

        if size == 0 {
            // According to the Linux manpage: "if size is equal to zero, and ptr is not NULL, then
            // the call is equivalent to free(ptr)." However, according to Darwin: "If size is zero
            // and ptr is not NULL, a new, minimum sized object is allocated and the original
            // object is freed." Since it is valid for malloc(0) to simply return NULL, we opt to
            // implement the Linux behavior in both cases. The only way for this to cause problems
            // is for Darwin programs to rely on the fact that the returned pointer represents the
            // "minimum sized object" instead of only assuming that, since the size passed was 0,
            // the object has 0 size. Since "minimum sized object" does not seem to be a
            // well-defined term, reliance on such behavior is erroneous.

            // TODO: What should we return?
            self.free(ptr);
            return ptr::null_mut();
        }

        // TODO: Round size up to 8 and use 8-byte alignment like in malloc/calloc?

        let layout = self.layout_finder.get_layout(ptr as *mut u8);
        // TODO: What's the right choice of alignment here?
        let new_layout = Layout::from_size_align(size as usize, 1).unwrap();
        match (&self.alloc).realloc(ptr as *mut u8, layout, new_layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.delete_layout(ptr);
                self.layout_finder.insert_layout(ptr, new_layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The C `reallocf` function (only implemented on Mac).
    #[cfg(target_os = "macos")]
    pub unsafe fn reallocf(&self, ptr: *mut c_void, size: size_t) -> *mut c_void {
        // See http://www.manpagez.com/man/3/malloc/osx-10.6.php

        if ptr.is_null() {
            return self.malloc(size);
        }

        if size == 0 {
            // According to the malloc manpage: "If size is zero and ptr is not NULL, a new,
            // minimum sized object is allocated and the original object is freed." See the
            // equivalent comment in realloc for why we do this.

            // TODO: What should we return?
            self.free(ptr);
            return ptr::null_mut();
        }

        // TODO: Round size up to 8 and use 8-byte alignment like in malloc/calloc?

        let layout = self.layout_finder.get_layout(ptr as *mut u8);
        // TODO: What's the right choice of alignment here?
        let new_layout = Layout::from_size_align(size as usize, 1).unwrap();
        match (&self.alloc).realloc(ptr as *mut u8, layout, new_layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.delete_layout(ptr);
                self.layout_finder.insert_layout(ptr, new_layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => {
                self.free(ptr);
                ptr::null_mut()
            }
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The C `reallocarray` function (only implemented on Linux).
    #[cfg(target_os = "linux")]
    pub unsafe fn reallocarray(&self,
                               ptr: *mut c_void,
                               nmemb: size_t,
                               size: size_t)
                               -> *mut c_void {
        // See http://man7.org/linux/man-pages/man3/malloc.3.html

        // According to the malloc manpage, "unlike that realloc() call, reallocarray() fails
        // safely in the case where the multiplication would overflow. If such an overflow occurs,
        // reallocarray() returns NULL, sets errno to ENOMEM, and leaves the original block of
        // memory unchanged."
        match nmemb.checked_mul(size) {
            Some(product) => self.realloc(ptr, product),
            None => {
                errno::set_errno(errno::Errno(libc::ENOMEM));
                ptr::null_mut()
            }
        }
    }


    /// The C `posix_memalign` function.
    pub unsafe fn posix_memalign(&self,
                                 memptr: *mut *mut c_void,
                                 alignment: size_t,
                                 size: size_t)
                                 -> i32 {
        // See http://man7.org/linux/man-pages/man3/posix_memalign.3.html

        // The manpage also specifies that the alignment must be a multiple of the word size, but
        // all powers of two greater than or equal to the word size are multiples of the word size,
        // so we omit that check.
        if alignment <= WORD_SIZE || !alignment.is_power_of_two() {
            return libc::EINVAL;
        }

        if size == 0 {
            *memptr = ptr::null_mut();
            return 0;
        }

        // TODO: posix_memalign does not require that size is a multiple of alignment. Thus, we
        // need to manually round up since valid Layouts must have that property. This is safe
        // because this API never takes the memory region size on deallocation, so it's fine that
        // the caller might think they have a smaller memory region than they actually do.

        let layout = Layout::from_size_align(size as usize, alignment).unwrap();
        match (&self.alloc).alloc(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                *memptr = ptr as *mut c_void;
                0
            }
            Err(AllocErr::Exhausted { .. }) => libc::ENOMEM,
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The obsolete C `memalign` function (only implemented on Linux).
    #[cfg(target_os = "linux")]
    pub unsafe fn memalign(&self, alignment: size_t, size: size_t) -> *mut c_void {
        // See http://man7.org/linux/man-pages/man3/posix_memalign.3.html

        if !alignment.is_power_of_two() {
            return ptr::null_mut();
        }

        if size == 0 {
            return ptr::null_mut();
        }

        // TODO: memalign does not require that size is a multiple of alignment. Thus, we need to
        // manually round up since valid Layouts must have that property. This is safe because this
        // API never takes the memory region size on deallocation, so it's fine that the caller
        // might think they have a smaller memory region than they actually do.

        let layout = Layout::from_size_align(size as usize, alignment).unwrap();
        match (&self.alloc).alloc(layout.clone()) {
            Ok(ptr) => {
                self.layout_finder.insert_layout(ptr, layout);
                ptr as *mut c_void
            }
            Err(AllocErr::Exhausted { .. }) => ptr::null_mut(),
            Err(AllocErr::Unsupported { .. }) => core::intrinsics::abort(),
        }
    }

    /// The C `aligned_alloc` function (only implemented on Linux).
    #[cfg(target_os = "linux")]
    pub unsafe fn aligned_alloc(&self, alignment: size_t, size: size_t) -> *mut c_void {
        // See http://man7.org/linux/man-pages/man3/posix_memalign.3.html

        // From the aligned_alloc manpage: "The function aligned_alloc() is the same as memalign(),
        // except for the added restriction that size should be a multiple of alignment."
        if size % alignment != 0 {
            return ptr::null_mut();
        }
        self.memalign(alignment, size)
    }
}

/// Define `extern "C"` functions for the C allocation API.
///
/// `define_malloc` is a convenience macro that constructs a global instance of `Malloc` and
/// defines each of the functions of the C allocation API by calling methods on that instance. One
/// function is defined for each of the methods on `Malloc`. Users who only want to define a subset
/// of the C allocation API should instead define these functions manually.
///
/// `define_malloc` takes an allocator type, an expression to construct a new instance of that
/// type, a `LayoutFinder` type, and an expression to construct a new instance of that type. Both
/// expressions must be constant expressions, as they will be used in the initialization of a
/// static variable.
#[macro_export]
macro_rules! define_malloc {
    ($alloc_ty:ty, $alloc_new:expr, $layout_finder_ty:ty, $layout_finder_new:expr) => (
        static HEAP: $crate::Malloc<$alloc_ty, $layout_finder_ty> = $crate::Malloc::new($alloc_new, $layout_finder_new);

        #[no_mangle]
        pub extern "C" fn malloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.malloc(size) }
        }

        #[no_mangle]
        pub extern "C" fn free(ptr: *mut c_void) {
            unsafe { HEAP.free(ptr) }
        }

        #[no_mangle]
        pub extern "C" fn cfree(ptr: *mut c_void) {
            unsafe { HEAP.cfree(ptr) }
        }

        #[no_mangle]
        pub extern "C" fn calloc(nmemb: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.calloc(nmemb, size) }
        }

        #[no_mangle]
        pub extern "C" fn valloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.valloc(size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn pvalloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.pvalloc(size) }
        }

        #[no_mangle]
        pub extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void {
            unsafe { HEAP.realloc(ptr, size) }
        }

        #[cfg(target_os = "macos")]
        #[no_mangle]
        pub extern "C" fn reallocf(ptr: *mut c_void, size: size_t) -> *mut c_void {
            unsafe { HEAP.reallocf(ptr, size) }
        }

        #[cfg(target_os = "linux")]
        pub extern "C" fn reallocarray(ptr: *mut c_void, nmemb: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.reallocarray(ptr, nmemb, size) }
        }

        #[no_mangle]
        pub extern "C" fn posix_memalign(memptr: *mut *mut c_void, alignment: size_t, size: size_t) -> i32 {
            unsafe { HEAP.posix_memalign(memptr, alignment, size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn memalign(alignment: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.memalign(alignment, size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn aligned_alloc(alignment: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.aligned_alloc(alignment, size) }
        }
    )
}

// This line re-exports the macros from lazy_static so that they'll be available to the code
// calling define_malloc_lazy_static. This allows define_malloc_lazy_static to be used without the
// caller needing to know about lazy_static and import its macros themselves.
//
// Credit to https://users.rust-lang.org/t/how-to-use-macro-inside-another-macro/12061/2
pub use lazy_static::*;

/// Define `extern "C"` functions for the C allocation API with non-constant initializers.
///
/// `define_malloc_lazy_static` is like `define_malloc`, except there is no requirement that the
/// initialization expressions must be constant. Instead, `lazy_static` is used to construct the
/// global `Malloc` instance.
#[macro_export]
macro_rules! define_malloc_lazy_static {
    ($alloc_ty:ty, $alloc_new:expr, $layout_finder_ty:ty, $layout_finder_new:expr) => (
        lazy_static!{
            static ref HEAP: $crate::Malloc<$alloc_ty, $layout_finder_ty> = $crate::Malloc::new($alloc_new, $layout_finder_new);
        }

        #[no_mangle]
        pub extern "C" fn malloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.malloc(size) }
        }

        #[no_mangle]
        pub extern "C" fn free(ptr: *mut c_void) {
            unsafe { HEAP.free(ptr) }
        }

        #[no_mangle]
        pub extern "C" fn cfree(ptr: *mut c_void) {
            unsafe { HEAP.cfree(ptr) }
        }

        #[no_mangle]
        pub extern "C" fn calloc(nmemb: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.calloc(nmemb, size) }
        }

        #[no_mangle]
        pub extern "C" fn valloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.valloc(size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn pvalloc(size: size_t) -> *mut c_void {
            unsafe { HEAP.pvalloc(size) }
        }

        #[no_mangle]
        pub extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void {
            unsafe { HEAP.realloc(ptr, size) }
        }

        #[cfg(target_os = "macos")]
        #[no_mangle]
        pub extern "C" fn reallocf(ptr: *mut c_void, size: size_t) -> *mut c_void {
            unsafe { HEAP.reallocf(ptr, size) }
        }

        #[cfg(target_os = "linux")]
        pub extern "C" fn reallocarray(ptr: *mut c_void, nmemb: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.reallocarray(ptr, nmemb, size) }
        }

        #[no_mangle]
        pub extern "C" fn posix_memalign(memptr: *mut *mut c_void, alignment: size_t, size: size_t) -> i32 {
            unsafe { HEAP.posix_memalign(memptr, alignment, size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn memalign(alignment: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.memalign(alignment, size) }
        }

        #[cfg(target_os = "linux")]
        #[no_mangle]
        pub extern "C" fn aligned_alloc(alignment: size_t, size: size_t) -> *mut c_void {
            unsafe { HEAP.aligned_alloc(alignment, size) }
        }
    )
}