package mem

// The Rollback Stack Allocator was designed for the test runner to be fast,
// able to grow, and respect the Tracking Allocator's requirement for
// individual frees. It is not overly concerned with fragmentation, however.
//
// It has support for expansion when configured with a block allocator and
// limited support for out-of-order frees.
//
// Allocation has constant-time best and usual case performance.
// At worst, it is linear according to the number of memory blocks.
//
// Allocation follows a first-fit strategy when there are multiple memory
// blocks.
//
// Freeing has constant-time best and usual case performance.
// At worst, it is linear according to the number of memory blocks and number
// of freed items preceding the last item in a block.
//
// Resizing has constant-time performance, if it's the last item in a block, or
// the new size is smaller. Naturally, this becomes linear-time if there are
// multiple blocks to search for the pointer's owning block. Otherwise, the
// allocator defaults to a combined alloc & free operation internally.
//
// Out-of-order freeing is accomplished by collapsing a run of freed items
// from the last allocation backwards.
//
// Each allocation has an overhead of 8 bytes and any extra bytes to satisfy
// the requested alignment.

import "base:runtime"

ROLLBACK_STACK_DEFAULT_BLOCK_SIZE :: 4 * Megabyte

// This limitation is due to the size of `prev_ptr`, but it is only for the
// head block; any allocation in excess of the allocator's `block_size` is
// valid, so long as the block allocator can handle it.
//
// This is because allocations over the block size are not split up if the item
// within is freed; they are immediately returned to the block allocator.
ROLLBACK_STACK_MAX_HEAD_BLOCK_SIZE :: 2 * Gigabyte


Rollback_Stack_Header :: bit_field u64 {
	prev_offset:  uintptr | 32,
	is_free:         bool |  1,
	prev_ptr:     uintptr | 31,
}

Rollback_Stack_Block :: struct {
	next_block: ^Rollback_Stack_Block,
	last_alloc: rawptr,
	offset: uintptr,
	buffer: []byte,
}

Rollback_Stack :: struct {
	head: ^Rollback_Stack_Block,
	block_size: int,
	block_allocator: Allocator,
}


@(private="file", require_results)
rb_ptr_in_bounds :: proc(block: ^Rollback_Stack_Block, ptr: rawptr) -> bool {
	start := raw_data(block.buffer)
	end   := start[block.offset:]
	return start < ptr && ptr <= end
}

@(private="file", require_results)
rb_find_ptr :: proc(stack: ^Rollback_Stack, ptr: rawptr) -> (
	parent: ^Rollback_Stack_Block,
	block:  ^Rollback_Stack_Block,
	header: ^Rollback_Stack_Header,
	err: Allocator_Error,
) {
	for block = stack.head; block != nil; block = block.next_block {
		if rb_ptr_in_bounds(block, ptr) {
			header = cast(^Rollback_Stack_Header)(cast(uintptr)ptr - size_of(Rollback_Stack_Header))
			return
		}
		parent = block
	}
	return nil, nil, nil, .Invalid_Pointer
}

@(private="file", require_results)
rb_find_last_alloc :: proc(stack: ^Rollback_Stack, ptr: rawptr) -> (
	block: ^Rollback_Stack_Block,
	header: ^Rollback_Stack_Header,
	ok: bool,
) {
	for block = stack.head; block != nil; block = block.next_block {
		if block.last_alloc == ptr {
			header = cast(^Rollback_Stack_Header)(cast(uintptr)ptr - size_of(Rollback_Stack_Header))
			return block, header, true
		}
	}
	return nil, nil, false
}

@(private="file")
rb_rollback_block :: proc(block: ^Rollback_Stack_Block, header: ^Rollback_Stack_Header) {
	header := header
	for block.offset > 0 && header.is_free {
		block.offset = header.prev_offset
		block.last_alloc = raw_data(block.buffer)[header.prev_ptr:]
		header = cast(^Rollback_Stack_Header)(raw_data(block.buffer)[header.prev_ptr - size_of(Rollback_Stack_Header):])
	}
}

@(private="file", require_results)
rb_free :: proc(stack: ^Rollback_Stack, ptr: rawptr) -> Allocator_Error {
	parent, block, header := rb_find_ptr(stack, ptr) or_return
	if header.is_free {
		return .Invalid_Pointer
	}
	header.is_free = true
	if block.last_alloc == ptr {
		block.offset = header.prev_offset
		rb_rollback_block(block, header)
	}
	if parent != nil && block.offset == 0 {
		parent.next_block = block.next_block
		runtime.mem_free_with_size(block, size_of(Rollback_Stack_Block) + len(block.buffer), stack.block_allocator)
	}
	return nil
}

@(private="file")
rb_free_all :: proc(stack: ^Rollback_Stack) {
	for block := stack.head.next_block; block != nil; /**/ {
		next_block := block.next_block
		runtime.mem_free_with_size(block, size_of(Rollback_Stack_Block) + len(block.buffer), stack.block_allocator)
		block = next_block
	}

	stack.head.next_block = nil
	stack.head.last_alloc = nil
	stack.head.offset = 0
}

@(private="file", require_results)
rb_resize :: proc(stack: ^Rollback_Stack, ptr: rawptr, old_size, size, alignment: int) -> (result: []byte, err: Allocator_Error) {
	if ptr != nil {
		if block, _, ok := rb_find_last_alloc(stack, ptr); ok {
			// `block.offset` should never underflow because it is contingent
			// on `old_size` in the first place, assuming sane arguments.
			assert(block.offset >= cast(uintptr)old_size, "Rollback Stack Allocator received invalid `old_size`.")

			if block.offset + cast(uintptr)size - cast(uintptr)old_size < cast(uintptr)len(block.buffer) {
				// Prevent singleton allocations from fragmenting by forbidding
				// them to shrink, removing the possibility of overflow bugs.
				if len(block.buffer) <= stack.block_size {
					block.offset += cast(uintptr)size - cast(uintptr)old_size
				}
				#no_bounds_check return (cast([^]byte)ptr)[:size], nil
			}
		}
	}

	result = rb_alloc(stack, size, alignment) or_return
	runtime.mem_copy_non_overlapping(raw_data(result), ptr, old_size)
	err = rb_free(stack, ptr)

	return
}

@(private="file", require_results)
rb_alloc :: proc(stack: ^Rollback_Stack, size, alignment: int) -> (result: []byte, err: Allocator_Error) {
	parent: ^Rollback_Stack_Block
	for block := stack.head; /**/; block = block.next_block {
		when !ODIN_DISABLE_ASSERT {
			allocated_new_block: bool
		}

		if block == nil {
			if stack.block_allocator.procedure == nil {
				return nil, .Out_Of_Memory
			}

			minimum_size_required := size_of(Rollback_Stack_Header) + size + alignment - 1
			new_block_size := max(minimum_size_required, stack.block_size)
			block = rb_make_block(new_block_size, stack.block_allocator) or_return
			parent.next_block = block
			when !ODIN_DISABLE_ASSERT {
				allocated_new_block = true
			}
		}

		start := raw_data(block.buffer)[block.offset:]
		padding := cast(uintptr)calc_padding_with_header(cast(uintptr)start, cast(uintptr)alignment, size_of(Rollback_Stack_Header))

		if block.offset + padding + cast(uintptr)size > cast(uintptr)len(block.buffer) {
			when !ODIN_DISABLE_ASSERT {
				if allocated_new_block {
					panic("Rollback Stack Allocator allocated a new block but did not use it.")
				}
			}
			parent = block
			continue
		}

		header := cast(^Rollback_Stack_Header)(start[padding - size_of(Rollback_Stack_Header):])
		ptr := start[padding:]

		header^ = {
			prev_offset = block.offset,
			prev_ptr = uintptr(0) if block.last_alloc == nil else cast(uintptr)block.last_alloc - cast(uintptr)raw_data(block.buffer),
			is_free = false,
		}

		block.last_alloc = ptr
		block.offset += padding + cast(uintptr)size

		if len(block.buffer) > stack.block_size {
			// This block exceeds the allocator's standard block size and is considered a singleton.
			// Prevent any further allocations on it.
			block.offset = cast(uintptr)len(block.buffer)
		}
		
		#no_bounds_check return ptr[:size], nil
	}

	return nil, .Out_Of_Memory
}

@(private="file", require_results)
rb_make_block :: proc(size: int, allocator: Allocator) -> (block: ^Rollback_Stack_Block, err: Allocator_Error) {
	buffer := runtime.mem_alloc(size_of(Rollback_Stack_Block) + size, align_of(Rollback_Stack_Block), allocator) or_return

	block = cast(^Rollback_Stack_Block)raw_data(buffer)
	#no_bounds_check block.buffer = buffer[size_of(Rollback_Stack_Block):]
	return
}


rollback_stack_init_buffered :: proc(stack: ^Rollback_Stack, buffer: []byte, location := #caller_location) {
	MIN_SIZE :: size_of(Rollback_Stack_Block) + size_of(Rollback_Stack_Header) + size_of(rawptr)
	assert(len(buffer) >= MIN_SIZE, "User-provided buffer to Rollback Stack Allocator is too small.", location)

	block := cast(^Rollback_Stack_Block)raw_data(buffer)
	block^ = {}
	#no_bounds_check block.buffer = buffer[size_of(Rollback_Stack_Block):]

	stack^ = {}
	stack.head = block
	stack.block_size = len(block.buffer)
}

rollback_stack_init_dynamic :: proc(
	stack: ^Rollback_Stack,
	block_size : int = ROLLBACK_STACK_DEFAULT_BLOCK_SIZE,
	block_allocator := context.allocator,
	location := #caller_location,
) -> Allocator_Error {
	assert(block_size >= size_of(Rollback_Stack_Header) + size_of(rawptr), "Rollback Stack Allocator block size is too small.", location)
	when size_of(int) > 4 {
		// It's impossible to specify an argument in excess when your integer
		// size is insufficient; check only on platforms with big enough ints.
		assert(block_size <= ROLLBACK_STACK_MAX_HEAD_BLOCK_SIZE, "Rollback Stack Allocators cannot support head blocks larger than 2 gigabytes.", location)
	}

	block := rb_make_block(block_size, block_allocator) or_return

	stack^ = {}
	stack.head = block
	stack.block_size = block_size
	stack.block_allocator = block_allocator

	return nil
}

rollback_stack_init :: proc {
	rollback_stack_init_buffered,
	rollback_stack_init_dynamic,
}

rollback_stack_destroy :: proc(stack: ^Rollback_Stack) {
	if stack.block_allocator.procedure != nil {
		rb_free_all(stack)
		free(stack.head, stack.block_allocator)
	}
	stack^ = {}
}

@(require_results)
rollback_stack_allocator :: proc(stack: ^Rollback_Stack) -> Allocator {
	return Allocator {
		data = stack,
		procedure = rollback_stack_allocator_proc,
	}
}

@(require_results)
rollback_stack_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
                                      size, alignment: int,
                                      old_memory: rawptr, old_size: int, location := #caller_location,
) -> (result: []byte, err: Allocator_Error) {
	stack := cast(^Rollback_Stack)allocator_data

	switch mode {
	case .Alloc, .Alloc_Non_Zeroed:
		assert(size >= 0, "Size must be positive or zero.", location)
		assert(is_power_of_two(cast(uintptr)alignment), "Alignment must be a power of two.", location)
		result = rb_alloc(stack, size, alignment) or_return

		if mode == .Alloc {
			zero_slice(result)
		}

	case .Free:
		err = rb_free(stack, old_memory)

	case .Free_All:
		rb_free_all(stack)

	case .Resize, .Resize_Non_Zeroed:
		assert(size >= 0, "Size must be positive or zero.", location)
		assert(old_size >= 0, "Old size must be positive or zero.", location)
		assert(is_power_of_two(cast(uintptr)alignment), "Alignment must be a power of two.", location)
		result = rb_resize(stack, old_memory, old_size, size, alignment) or_return

		#no_bounds_check if mode == .Resize && size > old_size {
			zero_slice(result[old_size:])
		}

	case .Query_Features:
		set := (^Allocator_Mode_Set)(old_memory)
		if set != nil {
			set^ = {.Alloc, .Alloc_Non_Zeroed, .Free, .Free_All, .Resize, .Resize_Non_Zeroed}
		}
		return nil, nil

	case .Query_Info:
		return nil, .Mode_Not_Implemented
	}

	return
}