module dlib.alloc;

import dlib.aliases;
import dlib.math;
import dlib.platform;
import dlib.util;

import std.stdio;
import core.stdc.string : memset;
import core.memory;

static Scratch g_scratch;
static u64     g_scratch_index;

const DEFAULT_ALIGNMENT = (void *).sizeof * 2;

struct Scratch
{
	Arena arena;
	bool  init;
}

struct ArenaPool
{
	u8*        mem;
	u64        pos;
	u64        length;
	ArenaPool* next;
}

struct Arena
{
	ArenaPool* first, last;
	u64        def_size;
}

struct TempArena
{
	Arena*     arena;
	u64        start_pos;
	ArenaPool* start_pool;
}

extern(C)
{
	void gc_setProxy(void* p);
	void gc_clrProxy();
}

T*
MAlloc(T)()
{
	void* mem = MemAlloc(T.sizeof);
	return cast(T*)mem;
}

T[]
MAlloc(T)(u64 count)
{
	void* mem = MemAlloc(T.sizeof * count);
	return (cast(T*)mem)[0 .. count];
}

void
MFree(T)(T* ptr)
{
	MemFree(ptr, T.sizeof);
}

void
MFree(T)(T[] slice)
{
	MemFree(slice.ptr, cast(u64)slice.length * T.sizeof);
}

T*
Alloc(T)()
{
	void* mem = pureMalloc(T.sizeof);
	memset(mem, 0, T.sizeof);
	return (cast(T*)mem);
}

T[]
Alloc(T)(u64 count)
{
	void* mem = pureMalloc(T.sizeof * count);
	memset(mem, 0, T.sizeof * count);
	return (cast(T*)mem)[0 .. count];
}

T[]
Alloc(T)(T[] target)
{
	T[] arr = Alloc!(T)(target.length);
	arr[] = target[];
	return arr;
}

string
Alloc(string target)
{
	u8[] str = Alloc!(u8)(target.length);
	str[]    = cast(u8[])target[];
	return ConvToStr(str);
}

T[]
Alloc(T)(T[] target, u64 start, u64 len)
{
	T[] arr = Alloc!(T)(len);
	arr[0 .. $] = target[start .. start+len];
	return arr;
}

string
Alloc(string target, u64 start, u64 len)
{
	u8[] str    = Alloc!(u8)(len);
	str[0 .. $] = cast(u8[])target[start .. start+len];
	return ConvToStr(str);
}

T[]
Alloc(T)(u64 count, T set)
{
	T[] arr = Alloc!(T)(count);
	arr[] = set;
	return arr;
}

T[]
Realloc(T)(T[] arr, u64 count)
{
	void* mem = pureRealloc(arr.ptr, T.sizeof * count);
	return (cast(T*)mem)[0 .. count];
}

void
Free(T)(T[] arr)
{
	pureFree(arr.ptr);
}

void 
Free(T)(T* ptr)
{
	pureFree(ptr);
}

Arena 
CreateArena(u64 size)
{
	Arena arena = { 
		def_size: size,
	};

	AddArenaPool(&arena, size);

	return arena;
}

TempArena
BeginTempArena(Arena* arena)
{
	TempArena t = {
		arena: arena,
	};

	auto n = arena.first;
	for(;;)
	{
		if(n.next == null)
		{
			t.start_pool = n;
			t.start_pos  = n.pos;
			break;
		}

		n = n.next;
	}

	return t;
}

void
End(TempArena* t)
{
	bool resetting = false;
	for(auto n = t.arena.first; n != null; n = n.next)
	{
		if(t.start_pool == n)
		{
			n.pos = t.start_pos;
			resetting = true;
		}
		else if(resetting)
		{
			n.pos = 0;
		}
	}
}

T[]
Alloc(T)(TempArena* t, T[] target)
{
	return Alloc!(T)(t.arena, target);
}

T[]
Alloc(T)(TempArena* t, u64 count, T set)
{
	return Alloc!(T)(t.arena, count, set);
}

T[]
Alloc(T)(TempArena* t, u64 count)
{
	return Alloc!(T)(t.arena, count);
}

T[]
Alloc(T)(TempArena* t, T[] target)
{
	return Alloc!(T)(t.arena, target);
}

T[]
Alloc(T)(TempArena* t, T[] target, u64 start, u64 len)
{
	return Alloc!(T)(t.arena, target, start, len);
}

T*
Alloc(T)(TempArena* t)
{
	return Alloc!(T)(t.arena);
}

void
AddArenaPool(Arena* arena, u64 size)
{
	u8* mem = Alloc!(u8)(size + ArenaPool.sizeof).ptr;

	ArenaPool* node = cast(ArenaPool*)mem;

	node.mem = (cast(u8*)mem) + ArenaPool.sizeof;
	node.pos = 0;
	node.length = size;

	assert(node.mem != null, "Unable to allocate memory for arena");

	SLLPushFront(arena, node, null);
}

T[]
Alloc(T)(Arena* arena, u64 count)
{
	void* mem = AllocAlign(arena, T.sizeof * count, DEFAULT_ALIGNMENT);
	memset(mem, 0, T.sizeof * count);
	return (cast(T*)mem)[0 .. count];
}

T[]
Alloc(T)(Arena* arena, T[] target)
{
	T[] arr = Alloc!(T)(arena, target.length);
	arr[]   = target[];
	return arr;
}

string
Alloc(Arena* arena, string target)
{
	u8[] str = Alloc!(u8)(arena, target.length);
	str[]    = cast(u8[])target[];
	return ConvToStr(str);
}

T[]
Alloc(T)(Arena* arena, T[] target, u64 start, u64 len)
{
	T[] arr     = Alloc!(T)(arena, len);
	arr[0 .. $] = target[start .. start+len];
	return arr;
}

string
Alloc(Arena* arena, string target, u64 start, u64 len)
{
	u8[] str    = Alloc!(u8)(arena, len);
	str[0 .. $] = cast(u8[])target[start .. start+len];
	return ConvToStr(str);
}

T[]
Alloc(T)(Arena* arena, u64 count, T set)
{
	T[] arr = Alloc!(T)(arena, count);
	arr[]   = set;
	return arr;
}

T* 
Alloc(T)(Arena* arena)
{
	void* mem = AllocAlign(arena, T.sizeof, DEFAULT_ALIGNMENT);
	memset(mem, 0, T.sizeof);
	return cast(T*)mem;
};

void* 
AllocAlign(Arena* arena, u64 size, u64 alignment)
{
	void* ptr = null;

	uintptr mem_pos, current, offset;
	ArenaPool* node = arena.first;
	while (true)
	{
		if(node == null)
		{
			if(size > arena.def_size)
			{
				size += arena.def_size;
			}

			AddArenaPool(arena, Max(size, arena.def_size));
			node = arena.first;
		}

		mem_pos = cast(uintptr)node.mem;
		current = mem_pos + node.pos;
		offset = AlignPow2(current, alignment) - mem_pos;

		if(offset+size <= node.length)
		{
			break;
		}

		node = node.next;
	}
	
	ptr = &node.mem[offset];
	node.pos = offset+size;

	return ptr;
};

void
Reset(Arena* arena)
{
	ArenaPool* node = arena.first;
	while (node != null)
	{
		node.pos = 0;
		node = node.next;
	}
}

void
Free(Arena* arena)
{
	ArenaPool* node = arena.first;
	ArenaPool* next;
	while (node != null)
	{
		next = node.next;
		Free(node);
		node = next;
	}
}

void
ResetScratch(u64 size)
{
	if(!g_scratch.init)
	{
		g_scratch.arena = CreateArena(size);
		g_scratch.init = true;
	}

	Reset(&g_scratch.arena);
}

T*
ScratchAlloc(T)()
{
	return Alloc!(T)(&g_scratch.arena);
}

T[]
ScratchAlloc(T)(u64 count)
{
	return Alloc!(T)(&g_scratch.arena, count);
}

T[]
ScratchAlloc(T)(T[] target)
{
	T[] arr = ScratchAlloc!(T)(target.length);
	arr[] = target[];
	return arr;
}

string
ScratchAlloc(string target)
{
	u8[] str = ScratchAlloc!(u8)(target.length);
	str[]    = cast(u8[])target[];
	return ConvToStr(str);
}

T[]
ScratchAlloc(T)(T[] target, u64 start, u64 len)
{
	T[] arr     = ScratchAlloc!(T)(len);
	arr[0 .. $] = target[start .. start+len];
	return arr;
}

string
ScratchAlloc(string target, u64 start, u64 len)
{
	u8[] str    = ScratchAlloc!(u8)(len);
	str[0 .. $] = cast(u8[])target[start .. start+len];
	return ConvToStr(str);
}

T[]
ScratchAlloc(T)(u64 count, T set)
{
	T[] arr = ScratchAlloc!(T)(count);
	arr[] = set;
	return arr;
}

version(DLIB_TEST) unittest
{
	{
		u64[5] arr = [1, 2, 3, 4, 5];

		u64[] copy = Alloc!(u64)(arr);

		assert(arr == copy);
	}

	{
		Arena a = CreateArena(64);

		u64[] arr = Alloc!(u64)(&a, 128);
	}
}