@[has_globals]
module builtin

#insert "@VEXEROOT/vlib/builtin/prealloc_atomics.h"

fn C.v_prealloc_atomic_add_i32(ptr &int, delta int) int
fn C.v_prealloc_atomic_load_i32(ptr &int) int
fn C.v_prealloc_atomic_store_i32(ptr &int, val int) int
fn C.v_prealloc_atomic_cas_i32(ptr &int, expected int, desired int) int

// With -prealloc, V calls libc's malloc to get chunks, each at least 16MB
// in size, as needed. Once a chunk is available, all malloc() calls within
// V code, that can fit inside the chunk, will use it instead, each bumping a
// pointer, till the chunk is filled. Once a chunk is filled, a new chunk will
// be allocated by calling libc's malloc, and the process continues.
// Each new chunk has a pointer to the old one. The base arena is thread-local;
// scoped arenas can be freed earlier with `prealloc_scope_end` or transferred
// and freed later with `prealloc_scope_free_after`.
// The goal of all this is to amortize the cost of calling libc's malloc,
// trading higher memory usage for a compiler (or any single threaded batch
// mode program), for a ~8-10% speed increase.

// size of the process/thread preallocated chunk
const prealloc_block_size = 16 * 1024 * 1024

// size of the first chunk for a scoped prealloc arena. Request-scoped arenas
// should not force a 16MB libc allocation for every request.
const prealloc_scope_block_size = 256 * 1024

// `malloc` has to return memory suitably aligned for any V value. Keep the
// default at the common max alignment used by libc malloc on current targets.
const prealloc_default_align = sizeof(voidptr) * 2

__global g_memory_block &VMemoryBlock
@[heap]
struct VMemoryBlock {
mut:
    current        &u8             = 0 // 8
    stop           &u8             = 0 // 8
    start          &u8             = 0 // 8
    previous       &VMemoryBlock   = 0 // 8
    next           &VMemoryBlock   = 0 // 8
    scope          &VPreallocScope = 0
    min_block_size isize
    is_scope       bool
    id             int // 4
    mallocs        int // 4
}

@[heap]
struct VPreallocScope {
mut:
    previous       &VMemoryBlock = 0
    first          &VMemoryBlock = 0
    refs           int
    free_requested int
    abandoned      int
    finalized      int
}

fn vmemory_abort_on_nil(p voidptr, bytes isize) {
    if unsafe { p == 0 } {
        C.fprintf(C.stderr, c'could not allocate %td bytes\n', bytes)
        exit(1)
    }
}

fn vmemory_effective_align(align isize) isize {
    default_align := isize(prealloc_default_align)
    if align > default_align {
        return align
    }
    return default_align
}

@[unsafe]
fn vmemory_align_up(ptr &u8, align isize) &u8 {
    if align <= 1 {
        return ptr
    }
    addr := u64(ptr)
    alignment := u64(align)
    offset := addr % alignment
    if offset == 0 {
        return ptr
    }
    return unsafe { &u8(i64(addr + alignment - offset)) }
}

fn vmemory_block_used(mb &VMemoryBlock) i64 {
    return unsafe { i64(mb.current) - i64(mb.start) }
}

fn vmemory_block_size(mb &VMemoryBlock) i64 {
    return unsafe { i64(mb.stop) - i64(mb.start) }
}

@[unsafe]
fn prealloc_trace_scope(action &char, scope &VPreallocScope) {
    $if trace_prealloc ? {
        if scope == unsafe { nil } {
            C.fprintf(C.stderr, c'[trace_prealloc] scope %s scope=%p\n', action, scope)
            return
        }
        unsafe {
            mut blocks := 0
            mut used := i64(0)
            mut size := i64(0)
            mut mallocs := 0
            mut mb := scope.first
            for mb != 0 {
                blocks++
                used += vmemory_block_used(mb)
                size += vmemory_block_size(mb)
                mallocs += mb.mallocs
                mb = mb.next
            }
            C.fprintf(C.stderr,
                c'[trace_prealloc] scope %s scope=%p previous=%p first=%p blocks=%d used=%lld size=%lld mallocs=%d\n',
                action, scope, scope.previous, scope.first, blocks, used, size, mallocs)
        }
    }
}

@[unsafe]
fn vmemory_block_new(prev &VMemoryBlock, at_least isize, align isize) &VMemoryBlock {
    return unsafe { vmemory_block_new_sized(prev, at_least, align, isize(prealloc_block_size)) }
}

@[unsafe]
fn vmemory_block_new_sized(prev &VMemoryBlock, at_least isize, align isize, min_block_size isize) &VMemoryBlock {
    vmem_block_size := sizeof(VMemoryBlock)
    mut v := unsafe { &VMemoryBlock(C.calloc(1, vmem_block_size)) }
    vmemory_abort_on_nil(v, vmem_block_size)
    if unsafe { prev != 0 } {
        v.id = prev.id + 1
    }

    v.previous = prev
    if unsafe { prev != 0 } {
        prev.next = v
        v.is_scope = prev.is_scope
    }
    effective_min_block_size := if min_block_size > 0 {
        min_block_size
    } else {
        isize(prealloc_block_size)
    }
    v.min_block_size = effective_min_block_size
    base_block_size := if at_least < effective_min_block_size {
        effective_min_block_size
    } else {
        at_least
    }
    block_size := if align > 0 {
        if base_block_size % align == 0 {
            base_block_size
        } else {
            base_block_size + align - (base_block_size % align)
        }
    } else {
        base_block_size
    }
    $if prealloc_trace_malloc ? {
        C.fprintf(C.stderr,
            c'vmemory_block_new id: %d, block_size: %lld, at_least: %lld, align: %lld\n', v.id,
            block_size, at_least, align)
    }

    fixed_align := if align <= 1 { 1 } else { align }
    $if windows {
        v.start = unsafe { C._aligned_malloc(block_size, fixed_align) }
    } $else {
        if fixed_align == 1 {
            v.start = unsafe { C.malloc(block_size) }
        } else {
            v.start = unsafe { C.aligned_alloc(fixed_align, block_size) }
        }
    }
    vmemory_abort_on_nil(v.start, block_size)
    $if prealloc_memset ? {
        unsafe { C.memset(v.start, int($d('prealloc_memset_value', 0)), block_size) }
    }
    v.stop = unsafe { &u8(i64(v.start) + block_size) }
    v.current = v.start
    $if trace_prealloc ? {
        if v.is_scope {
            C.fprintf(C.stderr,
                c'[trace_prealloc] block alloc block=%p previous=%p id=%d size=%lld at_least=%lld align=%lld start=%p stop=%p\n',
                v, prev, v.id, block_size, at_least, align, v.start, v.stop)
        }
    }
    return v
}

@[unsafe]
fn vmemory_block_malloc(n isize, align isize) &u8 {
    unsafe {
        // Lazy per-thread initialization: when g_memory_block is thread-local,
        // new threads start with a null pointer and need their own arena.
        if g_memory_block == nil {
            g_memory_block = vmemory_block_new(nil, isize(prealloc_block_size), 0)
        }
    }
    $if prealloc_trace_malloc ? {
        C.fprintf(C.stderr, c'vmemory_block_malloc g_memory_block.id: %d, n: %lld align: %d\n',
            g_memory_block.id, n, align)
    }
    unsafe {
        fixed_align := vmemory_effective_align(align)
        mut current := vmemory_align_up(g_memory_block.current, fixed_align)
        remaining := i64(g_memory_block.stop) - i64(current)
        if _unlikely_(remaining < n) {
            was_scope := g_memory_block.is_scope
            scope := g_memory_block.scope
            min_block_size := if g_memory_block.min_block_size > 0 {
                g_memory_block.min_block_size
            } else {
                isize(prealloc_block_size)
            }
            g_memory_block = vmemory_block_new_sized(g_memory_block, n, fixed_align, min_block_size)
            g_memory_block.is_scope = was_scope
            g_memory_block.scope = scope
            current = vmemory_align_up(g_memory_block.current, fixed_align)
        }
        res := &u8(current)
        g_memory_block.current = current
        g_memory_block.current += n
        $if prealloc_stats ? {
            g_memory_block.mallocs++
        } $else {
            $if trace_prealloc ? {
                g_memory_block.mallocs++
            }
        }
        $if trace_prealloc ? {
            if g_memory_block.is_scope {
                used := vmemory_block_used(g_memory_block)
                size := vmemory_block_size(g_memory_block)
                C.fprintf(C.stderr,
                    c'[trace_prealloc] alloc block=%p ptr=%p size=%lld align=%lld used=%lld/%lld mallocs=%d\n',
                    g_memory_block, res, n, fixed_align, used, size, g_memory_block.mallocs)
            }
        }
        return res
    }
}

@[unsafe]
fn vmemory_block_free(mb &VMemoryBlock) {
    $if trace_prealloc ? {
        if mb.is_scope {
            C.fprintf(C.stderr,
                c'[trace_prealloc] block free block=%p id=%d start=%p used=%lld size=%lld mallocs=%d\n',
                mb, mb.id, mb.start, vmemory_block_used(mb), vmemory_block_size(mb), mb.mallocs)
        }
    }
    $if windows {
        // Warning! On windows, we always use _aligned_free to free memory.
        C._aligned_free(mb.start)
    } $else {
        C.free(mb.start)
    }
    C.free(mb)
}

@[unsafe]
fn vmemory_block_free_after(marker &VMemoryBlock) {
    if marker == unsafe { nil } {
        return
    }
    unsafe {
        mut mb := marker.next
        marker.next = nil
        vmemory_block_free_chain(mb)
    }
}

@[unsafe]
fn vmemory_block_free_chain(first &VMemoryBlock) {
    unsafe {
        mut mb := first
        for mb != 0 {
            next := mb.next
            vmemory_block_free(mb)
            mb = next
        }
    }
}

/////////////////////////////////////////////////

@[unsafe]
fn prealloc_vinit() {
    $if prealloc_trace_vinit ? {
        C.fprintf(C.stderr, c'prealloc_vinit started\n')
    }
    unsafe {
        g_memory_block = vmemory_block_new(nil, isize(prealloc_block_size), 0)
        at_exit(prealloc_vcleanup) or {}
    }
}

@[unsafe]
fn prealloc_vcleanup() {
    $if prealloc_trace_vcleanup ? {
        C.fprintf(C.stderr, c'prealloc_vcleanup started\n')
    }
    $if prealloc_stats ? {
        // Note: we do 2 loops here, because string interpolation
        // in the first loop may still use g_memory_block
        // The second loop however should *not* allocate at all.
        mut nr_mallocs := i64(0)
        mut total_used := i64(0)
        mut mb := g_memory_block
        for unsafe { mb != 0 } {
            nr_mallocs += mb.mallocs
            used := i64(mb.current) - i64(mb.start)
            total_used += used
            remaining := i64(mb.stop) - i64(mb.current)
            size := i64(mb.stop) - i64(mb.start)
            C.fprintf(C.stderr,
                c'> freeing mb: %16p, mb.id: %3d | size: %10lld | rem: %10lld | start: %16p | current: %16p | used: %10lld bytes | mallocs: %6d\n',
                mb, mb.id, size, remaining, mb.start, mb.current, used, mb.mallocs)
            mb = mb.previous
        }
        C.fprintf(C.stderr, c'> nr_mallocs: %lld, total_used: %lld bytes\n', nr_mallocs, total_used)
    }
    $if prealloc_dump ? {
        C.fprintf(C.stderr, c'prealloc_vcleanup dumping memory contents ...\n')
        mut start := g_memory_block
        unsafe {
            for start.previous != 0 {
                start = start.previous
            }
            C.fprintf(C.stderr, c'prealloc_vcleanup      start: %p\n', start)
            C.fprintf(C.stderr, c'prealloc_vcleanup   start.id: %d\n', start.id)
            C.fprintf(C.stderr, c'prealloc_vcleanup start.next: %p\n', start.next)

            mut total_used := u64(0)
            path := $d('memdumpfile', 'memdump.bin')
            C.fprintf(C.stderr, c'prealloc_vcleanup dumping process memory to path: %s\n', path.str)
            stream := C.fopen(path.str, c'wb')
            mut mb := start
            for {
                used := u64(mb.current) - u64(mb.start)
                total_used += used
                C.fprintf(C.stderr,
                    c'prealloc_vcleanup dumping mb: %p, mb.id: %d, used: %10lld bytes\n', mb,
                    mb.id, used)

                mut ptr := mb.start
                mut remaining_bytes := isize(used)
                mut x := isize(0)
                for remaining_bytes > 0 {
                    x = isize(C.fwrite(ptr, 1, remaining_bytes, stream))
                    ptr += x
                    remaining_bytes -= x
                }

                if mb.next == 0 {
                    break
                }
                mb = mb.next
            }
            C.fclose(stream)
            C.fprintf(C.stderr, c'prealloc_vcleanup total dump size in bytes: %lld\n', total_used)
        }
    }
    unsafe {
        for g_memory_block != 0 {
            $if windows {
                // Warning! On windows, we always use _aligned_free to free memory.
                C._aligned_free(g_memory_block.start)
            } $else {
                C.free(g_memory_block.start)
            }
            tmp := g_memory_block
            g_memory_block = g_memory_block.previous
            // free the link node
            C.free(tmp)
        }
    }
}

// prealloc_scope_begin starts a nested arena on the current thread. All V
// allocations after this call use the nested arena until `prealloc_scope_end`.
// The returned scope can be passed across threads and later freed with
// `prealloc_scope_free_after`, which is useful when a response buffer outlives
// the request handler thread.
@[unsafe]
pub fn prealloc_scope_begin() voidptr {
    unsafe {
        scope := &VPreallocScope(C.calloc(1, sizeof(VPreallocScope)))
        vmemory_abort_on_nil(scope, sizeof(VPreallocScope))
        scope.previous = g_memory_block
        scope.first = vmemory_block_new_sized(scope.previous, isize(prealloc_scope_block_size), 0,
            isize(prealloc_scope_block_size))
        scope.first.is_scope = true
        scope.first.scope = scope
        g_memory_block = scope.first
        prealloc_trace_scope(c'begin', scope)
        return scope
    }
}

@[unsafe]
pub fn prealloc_scope_checkpoint(label &char) {
    $if trace_prealloc ? {
        unsafe {
            if g_memory_block == 0 || !g_memory_block.is_scope {
                return
            }
            mut blocks := 0
            mut used := i64(0)
            mut size := i64(0)
            mut mallocs := 0
            mut first := g_memory_block
            for first.previous != 0 && first.previous.is_scope {
                first = first.previous
            }
            mut mb := first
            for mb != 0 {
                blocks++
                used += vmemory_block_used(mb)
                size += vmemory_block_size(mb)
                mallocs += mb.mallocs
                mb = mb.next
            }
            C.fprintf(C.stderr,
                c'[trace_prealloc] checkpoint label=%s first=%p current=%p blocks=%d used=%lld size=%lld mallocs=%d\n',
                label, first, g_memory_block, blocks, used, size, mallocs)
        }
    }
}

@[unsafe]
fn prealloc_scope_free_blocks(scope &VPreallocScope) {
    if scope == unsafe { nil } {
        return
    }
    unsafe {
        if scope.previous != 0 {
            scope.previous.next = nil
        }
        vmemory_block_free_chain(scope.first)
    }
}

@[unsafe]
fn prealloc_scope_request_free(scope &VPreallocScope, abandoned bool) {
    if scope == unsafe { nil } {
        return
    }
    unsafe {
        if abandoned {
            C.v_prealloc_atomic_store_i32(&scope.abandoned, 1)
        }
        C.v_prealloc_atomic_store_i32(&scope.free_requested, 1)
        prealloc_scope_finish_if_ready(scope)
    }
}

@[unsafe]
fn prealloc_scope_finish_if_ready(scope &VPreallocScope) {
    if scope == unsafe { nil } {
        return
    }
    unsafe {
        if C.v_prealloc_atomic_load_i32(&scope.free_requested) == 0 {
            return
        }
        if C.v_prealloc_atomic_load_i32(&scope.refs) != 0 {
            return
        }
        if C.v_prealloc_atomic_cas_i32(&scope.finalized, 0, 1) == 0 {
            return
        }
        if C.v_prealloc_atomic_load_i32(&scope.abandoned) == 0 {
            prealloc_scope_free_blocks(scope)
        }
        C.free(scope)
    }
}

@[unsafe]
fn prealloc_scope_detach_current(scope &VPreallocScope) {
    if scope == unsafe { nil } {
        return
    }
    unsafe {
        previous := scope.previous
        if previous != 0 {
            previous.next = nil
        }
        if g_memory_block != 0 && g_memory_block.is_scope && g_memory_block.scope == scope {
            g_memory_block = previous
        }
        scope.previous = nil
    }
}

// prealloc_scope_retain_current keeps the current scoped arena alive after the
// owner calls `prealloc_scope_end`/`prealloc_scope_free_after`. It is used by
// generated `spawn` wrappers so detached threads can safely receive arguments
// allocated in a request arena.
@[unsafe]
pub fn prealloc_scope_retain_current() voidptr {
    $if prealloc {
        unsafe {
            if g_memory_block == 0 || !g_memory_block.is_scope || g_memory_block.scope == 0 {
                return nil
            }
            scope := g_memory_block.scope
            C.v_prealloc_atomic_add_i32(&scope.refs, 1)
            $if trace_prealloc ? {
                prealloc_trace_scope(c'retain', scope)
            }
            return scope
        }
    } $else {
        return unsafe { nil }
    }
}

@[unsafe]
pub fn prealloc_scope_release(scope_ptr voidptr) {
    $if prealloc {
        if scope_ptr == unsafe { nil } {
            return
        }
        unsafe {
            scope := &VPreallocScope(scope_ptr)
            C.v_prealloc_atomic_add_i32(&scope.refs, -1)
            $if trace_prealloc ? {
                prealloc_trace_scope(c'release', scope)
            }
            prealloc_scope_finish_if_ready(scope)
        }
    }
}

// prealloc_scope_end frees a nested arena and restores the current thread arena
// to the state before `prealloc_scope_begin`.
@[unsafe]
pub fn prealloc_scope_end(scope_ptr voidptr) {
    if scope_ptr == unsafe { nil } {
        return
    }
    unsafe {
        scope := &VPreallocScope(scope_ptr)
        prealloc_trace_scope(c'end', scope)
        prealloc_scope_detach_current(scope)
        prealloc_scope_request_free(scope, false)
    }
}

// prealloc_scope_leave restores the current thread arena without freeing the
// scoped blocks. Call this before another thread takes ownership of the scope.
@[unsafe]
pub fn prealloc_scope_leave(scope_ptr voidptr) {
    if scope_ptr == unsafe { nil } {
        return
    }
    unsafe {
        scope := &VPreallocScope(scope_ptr)
        prealloc_trace_scope(c'leave', scope)
        prealloc_scope_detach_current(scope)
    }
}

// prealloc_scope_abandon restores the current thread arena and intentionally
// leaks the scoped blocks. It is only for APIs that transfer request state to
// user code without providing a close hook yet.
@[unsafe]
pub fn prealloc_scope_abandon(scope_ptr voidptr) {
    if scope_ptr == unsafe { nil } {
        return
    }
    unsafe {
        scope := &VPreallocScope(scope_ptr)
        prealloc_trace_scope(c'abandon', scope)
        prealloc_scope_leave(scope_ptr)
        prealloc_scope_request_free(scope, true)
    }
}

// prealloc_scope_free_after frees a nested arena from a marker without touching
// the caller's thread-local arena pointer. Use this when another thread finishes
// sending data that was allocated in the request thread.
@[unsafe]
pub fn prealloc_scope_free_after(scope_ptr voidptr) {
    if scope_ptr == unsafe { nil } {
        return
    }
    unsafe {
        scope := &VPreallocScope(scope_ptr)
        prealloc_trace_scope(c'free-after', scope)
        prealloc_scope_request_free(scope, false)
    }
}

@[unsafe]
fn prealloc_malloc(n isize) &u8 {
    return unsafe { vmemory_block_malloc(n, 0) }
}

@[unsafe]
fn prealloc_realloc(old_data &u8, old_size isize, new_size isize) &u8 {
    new_ptr := unsafe { vmemory_block_malloc(new_size, 0) }
    min_size := if old_size < new_size { old_size } else { new_size }
    unsafe { C.memcpy(new_ptr, old_data, min_size) }
    return new_ptr
}

@[unsafe]
fn prealloc_calloc(n isize) &u8 {
    new_ptr := unsafe { vmemory_block_malloc(n, 0) }
    unsafe { C.memset(new_ptr, 0, n) }
    return new_ptr
}

@[unsafe]
fn prealloc_malloc_align(n isize, align isize) &u8 {
    return unsafe { vmemory_block_malloc(n, align) }
}