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;
}

struct Arena
{
	SLList!(ArenaPool) pools;
	u64                def_size;
}

struct TempArena
{
	Arena*            arena;
	Node!(ArenaPool)* start_node;
	u64               start_pos;
}

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

T[]
MAllocArray(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
MFreeArray(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[]
AllocCopy(T)(T[] target)
{
	T[] arr = Alloc!(T)(target.length);
	arr[] = target[];
	return arr;
}

T[]
AllocArray(T)(u64 count)
{
	return Alloc!(T)(count);
}

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

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.pools.first;
	for(;;)
	{
		if(n.next == null)
		{
			t.start_node = n;
			t.start_pos  = n.value.pos;
			break;
		}

		n = n.next;
	}

	return t;
}

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

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

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

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

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

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

	Node!(ArenaPool)* node = cast(Node!(ArenaPool)*)mem;

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

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

	SLLPushFront(&arena.pools, 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[]
AllocArray(T)(Arena* arena, u64 count)
{
	return Alloc!(T)(arena, count);
}

T[]
AllocCopy(T)(Arena* arena, T[] target)
{
	T[] arr = Alloc!(T)(arena, target.length);
	arr[] = target[];
	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;
	Node!(ArenaPool)* node = arena.pools.first;
	while (true)
	{
		if(node == null)
		{
			if(size > arena.def_size)
			{
				size += arena.def_size;
			}

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

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

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

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

	return ptr;
};

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

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

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

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

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 AllocArray!(T)(&g_scratch.arena, count);
}

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

unittest
{
	{
		u64[5] arr = [1, 2, 3, 4, 5];

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

		assert(arr == copy);
	}

	{
		Arena a = CreateArena(64);

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