//! A library to help grow the stack when it runs out of space.
//!
//! This is an implementation of manually instrumented segmented stacks where points in a program's
//! control flow are annotated with "maybe grow the stack here". Each point of annotation indicates
//! how far away from the end of the stack it's allowed to be, plus the amount of stack to allocate
//! if it does reach the end.
//!
//! Once a program has reached the end of its stack, a temporary stack on the heap is allocated and
//! is switched to for the duration of a closure.
//!
//! For a set of lower-level primitives, consider the `psm` crate.
//!
//! # Examples
//!
//! ```
//! // Grow the stack if we are within the "red zone" of 32K, and if we allocate
//! // a new stack allocate 1MB of stack space.
//! //
//! // If we're already in bounds, just run the provided closure on current stack.
//! stacker::maybe_grow(32 * 1024, 1024 * 1024, || {
//!     // guaranteed to have at least 32K of stack
//! });
//! ```

#![allow(improper_ctypes)]

#[macro_use]
extern crate cfg_if;
extern crate libc;
#[cfg(windows)]
extern crate winapi;
#[macro_use]
extern crate psm;

use std::cell::Cell;

/// Grows the call stack if necessary.
///
/// This function is intended to be called at manually instrumented points in a program where
/// recursion is known to happen quite a bit. This function will check to see if we're within
/// `red_zone` bytes of the end of the stack, and if so it will allocate a new stack of at least
/// `stack_size` bytes.
///
/// The closure `f` is guaranteed to run on a stack with at least `red_zone` bytes, and it will be
/// run on the current stack if there's space available.
#[inline(always)]
pub fn maybe_grow<R, F: FnOnce() -> R>(red_zone: usize, stack_size: usize, callback: F) -> R {
    // if we can't guess the remaining stack (unsupported on some platforms) we immediately grow
    // the stack and then cache the new stack size (which we do know now because we allocated it.
    let enough_space = match remaining_stack() {
        Some(remaining) => remaining >= red_zone,
        None => false,
    };
    if enough_space {
        callback()
    } else {
        grow(stack_size, callback)
    }
}

/// Always creates a new stack for the passed closure to run on.
/// The closure will still be on the same thread as the caller of `grow`.
/// This will allocate a new stack with at least `stack_size` bytes.
pub fn grow<R, F: FnOnce() -> R>(stack_size: usize, callback: F) -> R {
    // To avoid monomorphizing `_grow()` and everything it calls,
    // we convert the generic callback to a dynamic one.
    let mut opt_callback = Some(callback);
    let mut ret = None;
    let ret_ref = &mut ret;

    // This wrapper around `callback` achieves two things:
    // * It converts the `impl FnOnce` to a `dyn FnMut`.
    //   `dyn` because we want it to not be generic, and
    //   `FnMut` because we can't pass a `dyn FnOnce` around without boxing it.
    // * It eliminates the generic return value, by writing it to the stack of this function.
    //   Otherwise the closure would have to return an unsized value, which isn't possible.
    let dyn_callback: &mut dyn FnMut() = &mut || {
        let taken_callback = opt_callback.take().unwrap();
        *ret_ref = Some(taken_callback());
    };

    _grow(stack_size, dyn_callback);
    ret.unwrap()
}

/// Queries the amount of remaining stack as interpreted by this library.
///
/// This function will return the amount of stack space left which will be used
/// to determine whether a stack switch should be made or not.
pub fn remaining_stack() -> Option<usize> {
    let current_ptr = current_stack_ptr();
    get_stack_limit().map(|limit| current_ptr - limit)
}

psm_stack_information! (
    yes {
        fn current_stack_ptr() -> usize {
            psm::stack_pointer() as usize
        }
    }
    no {
        #[inline(always)]
        fn current_stack_ptr() -> usize {
            unsafe {
                let mut x = std::mem::MaybeUninit::<u8>::uninit();
                // Unlikely to be ever exercised. As a fallback we execute a volatile read to a
                // local (to hopefully defeat the optimisations that would make this local a static
                // global) and take its address. This way we get a very approximate address of the
                // current frame.
                x.as_mut_ptr().write_volatile(42);
                x.as_ptr() as usize
            }
        }
    }
);

thread_local! {
    static STACK_LIMIT: Cell<Option<usize>> = Cell::new(unsafe {
        guess_os_stack_limit()
    })
}

#[inline(always)]
fn get_stack_limit() -> Option<usize> {
    STACK_LIMIT.with(|s| s.get())
}

#[inline(always)]
#[allow(unused)]
fn set_stack_limit(l: Option<usize>) {
    STACK_LIMIT.with(|s| s.set(l))
}

psm_stack_manipulation! {
    yes {
        struct StackRestoreGuard {
            new_stack: *mut std::ffi::c_void,
            stack_bytes: usize,
            old_stack_limit: Option<usize>,
        }

        impl StackRestoreGuard {
            #[cfg(target_arch = "wasm32")]
            unsafe fn new(stack_bytes: usize, _page_size: usize) -> StackRestoreGuard {
                let layout = std::alloc::Layout::from_size_align(stack_bytes, 16).unwrap();
                let ptr = std::alloc::alloc(layout);
                assert!(!ptr.is_null(), "unable to allocate stack");
                StackRestoreGuard {
                    new_stack: ptr as *mut _,
                    stack_bytes,
                    old_stack_limit: get_stack_limit(),
                }
            }

            #[cfg(not(target_arch = "wasm32"))]
            unsafe fn new(stack_bytes: usize, page_size: usize) -> StackRestoreGuard {
                let new_stack = libc::mmap(
                    std::ptr::null_mut(),
                    stack_bytes,
                    libc::PROT_NONE,
                    libc::MAP_PRIVATE |
                    libc::MAP_ANON,
                    -1, // Some implementations assert fd = -1 if MAP_ANON is specified
                    0
                );
                if new_stack == libc::MAP_FAILED {
                    let error = std::io::Error::last_os_error();
                    panic!("allocating stack failed with: {}", error)
                }
                let guard = StackRestoreGuard {
                    new_stack,
                    stack_bytes,
                    old_stack_limit: get_stack_limit(),
                };
                let above_guard_page = new_stack.add(page_size);
                #[cfg(not(target_os = "openbsd"))]
                let result = libc::mprotect(
                    above_guard_page,
                    stack_bytes - page_size,
                    libc::PROT_READ | libc::PROT_WRITE
                );
                #[cfg(target_os = "openbsd")]
                let result = if libc::mmap(
                        above_guard_page,
                        stack_bytes - page_size,
                        libc::PROT_READ | libc::PROT_WRITE,
                        libc::MAP_FIXED | libc::MAP_PRIVATE | libc::MAP_ANON | libc::MAP_STACK,
                        -1,
                        0) == above_guard_page {
                    0
                } else {
                    -1
                };
                if result == -1 {
                    let error = std::io::Error::last_os_error();
                    drop(guard);
                    panic!("setting stack permissions failed with: {}", error)
                }
                guard
            }
        }

        impl Drop for StackRestoreGuard {
            fn drop(&mut self) {
                #[cfg(target_arch = "wasm32")]
                unsafe {
                    std::alloc::dealloc(
                        self.new_stack as *mut u8,
                        std::alloc::Layout::from_size_align_unchecked(self.stack_bytes, 16),
                    );
                }
                #[cfg(not(target_arch = "wasm32"))]
                unsafe {
                    // FIXME: check the error code and decide what to do with it.
                    // Perhaps a debug_assertion?
                    libc::munmap(self.new_stack, self.stack_bytes);
                }
                set_stack_limit(self.old_stack_limit);
            }
        }

        fn _grow(stack_size: usize, callback: &mut dyn FnMut()) {
            // Calculate a number of pages we want to allocate for the new stack.
            // For maximum portability we want to produce a stack that is aligned to a page and has
            // a size that’s a multiple of page size. Furthermore we want to allocate two extras pages
            // for the stack guard. To achieve that we do our calculations in number of pages and
            // convert to bytes last.
            let page_size = page_size();
            let requested_pages = stack_size
                .checked_add(page_size - 1)
                .expect("unreasonably large stack requested") / page_size;
            let stack_pages = std::cmp::max(1, requested_pages) + 2;
            let stack_bytes = stack_pages.checked_mul(page_size)
                .expect("unreasonably large stack requesteed");

            // Next, there are a couple of approaches to how we allocate the new stack. We take the
            // most obvious path and use `mmap`. We also `mprotect` a guard page into our
            // allocation.
            //
            // We use a guard pattern to ensure we deallocate the allocated stack when we leave
            // this function and also try to uphold various safety invariants required by `psm`
            // (such as not unwinding from the callback we pass to it).
            //
            // Other than that this code has no meaningful gotchas.
            unsafe {
                let guard = StackRestoreGuard::new(stack_bytes, page_size);
                let above_guard_page = guard.new_stack.add(page_size);
                set_stack_limit(Some(above_guard_page as usize));
                let panic = psm::on_stack(above_guard_page as *mut _, stack_size, move || {
                    std::panic::catch_unwind(std::panic::AssertUnwindSafe(callback)).err()
                });
                drop(guard);
                if let Some(p) = panic {
                    std::panic::resume_unwind(p);
                }
            }
        }

        fn page_size() -> usize {
            // FIXME: consider caching the page size.
            #[cfg(not(target_arch = "wasm32"))]
            unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as usize }
            #[cfg(target_arch = "wasm32")]
            { 65536 }
        }
    }

    no {
        #[cfg(not(windows))]
        fn _grow(stack_size: usize, callback: &mut dyn FnMut()) {
            drop(stack_size);
            callback();
        }
    }
}

cfg_if! {
    if #[cfg(windows)] {
        use std::ptr;
        use std::io;

        use winapi::shared::basetsd::*;
        use winapi::shared::minwindef::{LPVOID, BOOL};
        use winapi::shared::ntdef::*;
        use winapi::um::fibersapi::*;
        use winapi::um::memoryapi::*;
        use winapi::um::processthreadsapi::*;
        use winapi::um::winbase::*;

        // Make sure the libstacker.a (implemented in C) is linked.
        // See https://github.com/rust-lang/rust/issues/65610
        #[link(name="stacker")]
        extern {
            fn __stacker_get_current_fiber() -> PVOID;
        }

        struct FiberInfo<F> {
            callback: std::mem::MaybeUninit<F>,
            panic: Option<Box<dyn std::any::Any + Send + 'static>>,
            parent_fiber: LPVOID,
        }

        unsafe extern "system" fn fiber_proc<F: FnOnce()>(data: LPVOID) {
            // This function is the entry point to our inner fiber, and as argument we get an
            // instance of `FiberInfo`. We will set-up the "runtime" for the callback and execute
            // it.
            let data = &mut *(data as *mut FiberInfo<F>);
            let old_stack_limit = get_stack_limit();
            set_stack_limit(guess_os_stack_limit());
            let callback = data.callback.as_ptr();
            data.panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(callback.read())).err();

            // Restore to the previous Fiber
            set_stack_limit(old_stack_limit);
            SwitchToFiber(data.parent_fiber);
            return;
        }

        fn _grow(stack_size: usize, callback: &mut dyn FnMut()) {
            // Fibers (or stackful coroutines) is the only official way to create new stacks on the
            // same thread on Windows. So in order to extend the stack we create fiber and switch
            // to it so we can use it's stack. After running `callback` within our fiber, we switch
            // back to the current stack and destroy the fiber and its associated stack.
            unsafe {
                let was_fiber = IsThreadAFiber() == TRUE as BOOL;
                let mut data = FiberInfo {
                    callback: std::mem::MaybeUninit::new(callback),
                    panic: None,
                    parent_fiber: {
                        if was_fiber {
                            // Get a handle to the current fiber. We need to use a C implementation
                            // for this as GetCurrentFiber is an header only function.
                            __stacker_get_current_fiber()
                        } else {
                            // Convert the current thread to a fiber, so we are able to switch back
                            // to the current stack. Threads coverted to fibers still act like
                            // regular threads, but they have associated fiber data. We later
                            // convert it back to a regular thread and free the fiber data.
                            ConvertThreadToFiber(ptr::null_mut())
                        }
                    },
                };

                if data.parent_fiber.is_null() {
                    panic!("unable to convert thread to fiber: {}", io::Error::last_os_error());
                }

                let fiber = CreateFiber(
                    stack_size as SIZE_T,
                    Some(fiber_proc::<&mut dyn FnMut()>),
                    &mut data as *mut FiberInfo<&mut dyn FnMut()> as *mut _,
                );
                if fiber.is_null() {
                    panic!("unable to allocate fiber: {}", io::Error::last_os_error());
                }

                // Switch to the fiber we created. This changes stacks and starts executing
                // fiber_proc on it. fiber_proc will run `callback` and then switch back to run the
                // next statement.
                SwitchToFiber(fiber);
                DeleteFiber(fiber);

                // Clean-up.
                if !was_fiber {
                    if ConvertFiberToThread() == 0 {
                        // FIXME: Perhaps should not panic here?
                        panic!("unable to convert back to thread: {}", io::Error::last_os_error());
                    }
                }
                if let Some(p) = data.panic {
                    std::panic::resume_unwind(p);
                }
            }
        }

        #[inline(always)]
        fn get_thread_stack_guarantee() -> usize {
            let min_guarantee = if cfg!(target_pointer_width = "32") {
                0x1000
            } else {
                0x2000
            };
            let mut stack_guarantee = 0;
            unsafe {
                // Read the current thread stack guarantee
                // This is the stack reserved for stack overflow
                // exception handling.
                // This doesn't return the true value so we need
                // some further logic to calculate the real stack
                // guarantee. This logic is what is used on x86-32 and
                // x86-64 Windows 10. Other versions and platforms may differ
                SetThreadStackGuarantee(&mut stack_guarantee)
            };
            std::cmp::max(stack_guarantee, min_guarantee) as usize + 0x1000
        }

        #[inline(always)]
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            // Query the allocation which contains our stack pointer in order
            // to discover the size of the stack
            //
            // FIXME: we could read stack base from the TIB, specifically the 3rd element of it.
            type QueryT = winapi::um::winnt::MEMORY_BASIC_INFORMATION;
            let mut mi = std::mem::MaybeUninit::<QueryT>::uninit();
            VirtualQuery(
                psm::stack_pointer() as *const _,
                mi.as_mut_ptr(),
                std::mem::size_of::<QueryT>() as SIZE_T,
            );
            Some(mi.assume_init().AllocationBase as usize + get_thread_stack_guarantee() + 0x1000)
        }
    } else if #[cfg(any(target_os = "linux", target_os="solaris", target_os = "netbsd"))] {
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            let mut attr = std::mem::MaybeUninit::<libc::pthread_attr_t>::uninit();
            assert_eq!(libc::pthread_attr_init(attr.as_mut_ptr()), 0);
            assert_eq!(libc::pthread_getattr_np(libc::pthread_self(),
                                                attr.as_mut_ptr()), 0);
            let mut stackaddr = std::ptr::null_mut();
            let mut stacksize = 0;
            assert_eq!(libc::pthread_attr_getstack(
                attr.as_ptr(), &mut stackaddr, &mut stacksize
            ), 0);
            assert_eq!(libc::pthread_attr_destroy(attr.as_mut_ptr()), 0);
            Some(stackaddr as usize)
        }
    } else if #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))] {
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            let mut attr = std::mem::MaybeUninit::<libc::pthread_attr_t>::uninit();
            assert_eq!(libc::pthread_attr_init(attr.as_mut_ptr()), 0);
            assert_eq!(libc::pthread_attr_get_np(libc::pthread_self(), attr.as_mut_ptr()), 0);
            let mut stackaddr = std::ptr::null_mut();
            let mut stacksize = 0;
            assert_eq!(libc::pthread_attr_getstack(
                attr.as_ptr(), &mut stackaddr, &mut stacksize
            ), 0);
            assert_eq!(libc::pthread_attr_destroy(attr.as_mut_ptr()), 0);
            Some(stackaddr as usize)
        }
    } else if #[cfg(target_os = "openbsd")] {
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            let mut stackinfo = std::mem::MaybeUninit::<libc::stack_t>::uninit();
            assert_eq!(libc::pthread_stackseg_np(libc::pthread_self(), stackinfo.as_mut_ptr()), 0);
            Some(stackinfo.assume_init().ss_sp as usize - stackinfo.assume_init().ss_size)
        }
    } else if #[cfg(target_os = "macos")] {
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            Some(libc::pthread_get_stackaddr_np(libc::pthread_self()) as usize -
                libc::pthread_get_stacksize_np(libc::pthread_self()) as usize)
        }
    } else {
        // fallback for other platforms is to always increase the stack if we're on
        // the root stack. After we increased the stack once, we know the new stack
        // size and don't need this pessimization anymore
        #[inline(always)]
        unsafe fn guess_os_stack_limit() -> Option<usize> {
            None
        }
    }
}