package container_pool

import "base:intrinsics"
import "base:runtime"
import "base:sanitizer"
import "core:sync"

_ :: sanitizer
_ :: sync

DEFAULT_BLOCK_SIZE :: _DEFAULT_BLOCK_SIZE

Pool_Arena :: _Pool_Arena

/*
A thread-safe (between init and destroy) object pool backed by virtual growing arena returning stable pointers.
The element type requires an intrusive link node.

Example:
	Elem :: struct {
		link: ^Elem,
	}

	p: pool.Pool(Elem)
	pool.init(&p, "link")
*/
Pool :: struct($T: typeid) {
	arena:           Pool_Arena,
	num_outstanding: int,
	num_ready:       int,
	link_off:        uintptr,
	// Guards free_list.  An untagged Treiber stack is vulnerable to ABA:
	// between get's head load and its CAS, another thread can pop the head,
	// reuse it, and push it back — the CAS then succeeds and installs a
	// stale next pointer, handing an in-flight element to two owners.
	// A tagged head would keep this lock-free but requires a double-width
	// CAS, which is not portably available.
	mu:              sync.Mutex,
	free_list:       ^T,
}

@(require_results)
init :: proc(p: ^Pool($T), $link_field: string, block_size: uint = DEFAULT_BLOCK_SIZE) -> (err: runtime.Allocator_Error)
	where intrinsics.type_has_field(T, link_field),
	      intrinsics.type_field_type(T, link_field) == ^T {
	p.link_off = offset_of_by_string(T, link_field)
	return _pool_arena_init(&p.arena, block_size)
}

destroy :: proc(p: ^Pool($T)) {
	elem := sync.atomic_exchange_explicit(&p.free_list, nil, .Acquire)

	sync.atomic_store_explicit(&p.num_ready, 0, .Relaxed)

	when .Address in ODIN_SANITIZER_FLAGS {
		for ; elem != nil; elem = _get_next(p, elem) {
			_unpoison_elem(p, elem)
		}
	} else {
		_ = elem
	}

	_pool_arena_destroy(&p.arena)
	p.arena = {}
}

@(require_results)
get :: proc(p: ^Pool($T)) -> (elem: ^T, err: runtime.Allocator_Error) #optional_allocator_error {
	defer sync.atomic_add_explicit(&p.num_outstanding, 1, .Relaxed)

	sync.mutex_lock(&p.mu)
	elem = p.free_list
	if elem == nil {
		sync.mutex_unlock(&p.mu)
		// NOTE: pool arena has an internal lock.
		return new(T, _pool_arena_allocator(&p.arena))
	}
	p.free_list = _get_next(p, elem)
	sync.mutex_unlock(&p.mu)

	_set_next(p, elem, nil)
	_unpoison_elem(p, elem)
	sync.atomic_sub_explicit(&p.num_ready, 1, .Relaxed)
	return
}

put :: proc(p: ^Pool($T), elem: ^T) {
	intrinsics.mem_zero(elem, size_of(T))
	_poison_elem(p, elem)

	defer sync.atomic_sub_explicit(&p.num_outstanding, 1, .Relaxed)
	defer sync.atomic_add_explicit(&p.num_ready, 1, .Relaxed)

	sync.mutex_lock(&p.mu)
	_set_next(p, elem, p.free_list)
	p.free_list = elem
	sync.mutex_unlock(&p.mu)
}

num_outstanding :: proc(p: ^Pool($T)) -> int {
	return sync.atomic_load(&p.num_outstanding)
}

num_ready :: proc(p: ^Pool($T)) -> int {
	return sync.atomic_load(&p.num_ready)
}

cap :: proc(p: ^Pool($T)) -> int {
	return sync.atomic_load(&p.num_ready) + sync.atomic_load(&p.num_outstanding)
}

_get_next :: proc(p: ^Pool($T), elem: ^T) -> ^T {
	return (^^T)(uintptr(elem) + p.link_off)^
}

_set_next :: proc(p: ^Pool($T), elem: ^T, next: ^T) {
	(^^T)(uintptr(elem) + p.link_off)^ = next
}

_poison_elem :: proc(p: ^Pool($T), elem: ^T) {
	when .Address in ODIN_SANITIZER_FLAGS {
		if p.link_off > 0 {
			sanitizer.address_poison_rawptr(elem, int(p.link_off))
		}

		len := size_of(T) - p.link_off - size_of(rawptr)
		if len > 0 {
			ptr := rawptr(uintptr(elem) + p.link_off + size_of(rawptr))
			sanitizer.address_poison_rawptr(ptr, int(len))
		}
	}
}

_unpoison_elem :: proc(p: ^Pool($T), elem: ^T) {
	when .Address in ODIN_SANITIZER_FLAGS {
		if p.link_off > 0 {
			sanitizer.address_unpoison_rawptr(elem, int(p.link_off))
		}

		len := size_of(T) - p.link_off - size_of(rawptr)
		if len > 0 {
			ptr := rawptr(uintptr(elem) + p.link_off + size_of(rawptr))
			sanitizer.address_unpoison_rawptr(ptr, int(len))
		}
	}
}