dlib/dlib/alloc.d

module dlib.alloc;

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

import core.stdc.string;

static if(NativeTarget)
{
	public import core.memory : malloc = pureMalloc, realloc = pureRealloc, free = pureFree;
}

static Scratch g_scratch;
static usize   g_scratch_index;

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

struct Scratch
{
	Arena arena;
	bool  init;
}

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

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

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

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

@nogc:

version(WebAssembly)
{

pragma(LDC_intrinsic, "llvm.wasm.memory.grow.i32")
extern(C) usize grow(usize index, usize size) nothrow @nogc;

pragma(LDC_intrinsic, "llvm.wasm.memory.size.i32")
extern(C) usize size(usize index) nothrow @nogc;


extern(C) u32
WasmGrow(u32 size) nothrow @nogc
{
	return grow(0, size);
}

extern(C) u32
WasmGrow2(u32 index, u32 size)
{
	return grow(0, size);
}

extern(C) u32
WasmSize() nothrow @nogc
{
	return size(0);
}

}

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

T[]
MAlloc(T)(usize 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(usize)slice.length * T.sizeof);
}

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

T[]
Alloc(T)(usize count)
{
	void* mem = malloc(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);
}

string
AllocEscaped(char[] target)
{
	u8[] str                = Alloc!(u8)(target.length+1);
	str[0 .. target.length] = str[];
	str[str.length-1]       = '\0';
	return cast(string)str[0 .. str.length-1];
}

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

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

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

T[]
Realloc(T)(T[] arr, usize count)
{
	static if(NativeTarget)
	{
		void* mem = realloc(arr.ptr, T.sizeof * count);
	}
	else
	{
		void* mem = alloc(T.sizeof*count);
		memcpy(mem, arr.ptr, T.sizeof*arr.length);
		Free(arr);
	}

	return (cast(T*)mem)[0 .. count];
}

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

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

Arena 
CreateArena(usize 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, usize count, T set)
{
	return Alloc!(T)(t.arena, count, set);
}

T[]
Alloc(T)(TempArena* t, usize 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, usize start, usize len)
{
	return Alloc!(T)(t.arena, target, start, len);
}

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

void
AddArenaPool(Arena* arena, usize 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, usize count)
{
	void* mem = AllocAlign(arena, T.sizeof * count, DEFAULT_ALIGNMENT);
	memset(mem, 0, T.sizeof * count);
	return (cast(T*)mem)[0 .. cast(usize)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, usize start, usize len)
{
	T[] arr     = Alloc!(T)(arena, len);
	arr[0 .. $] = target[start .. start+len];
	return arr;
}

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

T[]
Alloc(T)(Arena* arena, usize 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, usize size, usize alignment)
{
	void* ptr = null;

	usize pool_alloc_size = size;
	if(pool_alloc_size > arena.def_size)
	{
		pool_alloc_size += arena.def_size;
	}

	usize mem_pos, current, offset;
	ArenaPool* node = arena.first;
	while(true)
	{
		if(node == null)
		{
			AddArenaPool(arena, Max(pool_alloc_size, arena.def_size));
			node = arena.first;
		}

		mem_pos = cast(usize)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(usize 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)(usize 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, usize start, usize len)
{
	T[] arr     = ScratchAlloc!(T)(len);
	arr[0 .. $] = target[start .. start+len];
	return arr;
}

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

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

version(DLIB_TEST)
{
	void WriteToArray(T)(T[] arr, T factor)
	{
		foreach(i; 0 .. arr.length)
		{
			arr[i] = i*factor;
		}
	}

	void AssertArrayValues(T)(T[] arr, T factor)
	{
		foreach(i; 0 .. arr.length)
		{
			assert(arr[i] == i*factor);
		}
	}

	void DLibTestAlloc()
	{
		{
			u64[5] arr0 = [1, 2, 3, 4, 5];

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

			assert(arr0 == copy);

			u64[] arr1 = Alloc!(u64)(64);

			WriteToArray(arr1, 5);

			assert(arr0 == copy);
			AssertArrayValues(arr1,  5);
		}

		{
			Arena a = CreateArena(64);

			u64[] arr0 = Alloc!(u64)(&a, 128);
			u64[] arr1 = Alloc!(u64)(&a, 256);

			WriteToArray(arr0, 2);
			WriteToArray(arr1, 3);
			WriteToArray(arr0, 2);

			AssertArrayValues(arr0, 2);
			AssertArrayValues(arr1, 3);

			Reset(&a);

			arr0 = Alloc!(u64)(&a, 256);
			arr1 = Alloc!(u64)(&a, 128);

			WriteToArray(arr0, 4);
			WriteToArray(arr1, 5);
			WriteToArray(arr0, 4);

			AssertArrayValues(arr0, 4);
			AssertArrayValues(arr1, 5);
		}
	}

	unittest
	{
		DLibTestAlloc();
	}
}

version(WebAssembly) nothrow @nogc:

struct MemBlock
{
	void*     addr;
	MemBlock* next;
	usize     size;
}

struct Heap
{
	MemBlock* free;
	MemBlock* used;
	MemBlock* fresh;
	usize     top;
}

extern extern(C) u8 __heap_base;
extern extern(C) u8 __data_end;

__gshared Heap* heap;
__gshared void* heap_limit;
__gshared usize heap_split_thresh;
__gshared usize heap_alignment;
__gshared usize heap_max_blocks;
__gshared usize current_pages;

enum bool MALLOC_COMPACT = true;
enum bool MALLOC_SPLIT   = true;

usize WasmPageSize(usize x) => x*1024*64;

/**
 * If compaction is enabled, inserts block
 * into free list, sorted by addr.
 * If disabled, add block has new head of
 * the free list.
 */
void 
InsertBlock(MemBlock *block) 
{
	static if(MALLOC_COMPACT)
	{
		MemBlock *ptr  = heap.free;
		MemBlock *prev = null;
		while(ptr != null) 
		{
			if(cast(usize)block.addr <= cast(usize)ptr.addr) 
			{
				break;
			}
			prev = ptr;
			ptr  = ptr.next;
		}

		if(prev != null) 
		{
			prev.next = block;
		} 
		else 
		{
			heap.free = block;
		}
		block.next = ptr;
	}
	else
	{
		block.next = heap.free;
		heap.free  = block;
	}
}

void 
ReleaseBlocks(MemBlock *scan, MemBlock *to) 
{
	MemBlock *scan_next;
	while(scan != to) 
	{
		scan_next  = scan.next;
		scan.next  = heap.fresh;
		heap.fresh = scan;
		scan.addr  = null;
		scan.size  = 0;
		scan       = scan_next;
	}
}

void 
Compact() 
{
	MemBlock *ptr = heap.free;
	MemBlock *prev;
	MemBlock *scan;
	while(ptr != null) 
	{
		prev = ptr;
		scan = ptr.next;

		while(scan != null && cast(usize)prev.addr + prev.size == cast(usize)scan.addr) 
		{
			prev = scan;
			scan = scan.next;
		}

		if(prev != ptr) 
		{
			usize new_size = cast(usize)prev.addr - cast(usize)ptr.addr + prev.size;
			ptr.size   = new_size;
			MemBlock *next = prev.next;
			// make merged blocks available
			ReleaseBlocks(ptr.next, prev.next);
			// relink
			ptr.next = next;
		}

		ptr = ptr.next;
	}
}

void 
MallocInit(const usize heap_blocks, const usize split_thresh, const usize alignment) 
{
	current_pages = WasmSize();
	void* limit   = cast(void *)WasmPageSize(current_pages);

	heap              = cast(Heap *)&__heap_base;
	heap_limit        = limit;
	heap_split_thresh = split_thresh;
	heap_alignment    = alignment;
	heap_max_blocks   = heap_blocks;

	heap.free   = null;
	heap.used   = null;
	heap.fresh  = cast(MemBlock*)(heap + 1);
	heap.top    = cast(usize)(heap.fresh + heap_blocks);

	MemBlock *block = heap.fresh;
	usize     i     = heap_max_blocks - 1;
	while(i--) 
	{
		block.next = block + 1;
		block++;
	}

	block.next = null;
}

void 
free(void *ptr) 
{
	MemBlock *block = heap.used;
	MemBlock *prev  = null;
	while(block != null) 
	{
		if(ptr == block.addr) 
		{
			if(prev) 
			{
				prev.next = block.next;
			} 
			else 
			{
				heap.used = block.next;
			}
			InsertBlock(block);

			static if(MALLOC_COMPACT)
			{
				Compact();
			}
		}
		prev  = block;
		block = block.next;
	}
}

usize
PtrInfo(void *ptr)
{
	usize result = 0;

	MemBlock *block = heap.used;
	while(block != null)
	{
		if(ptr == block.addr)
		{
			result = block.size;
			break;
		}
	}

	return result;
}

MemBlock*
AllocBlock(usize num) 
{
	MemBlock *ptr  = heap.free;
	MemBlock *prev = null;
	usize top  = heap.top;
	num         = (num + heap_alignment - 1) & -heap_alignment;
	while(ptr != null) 
	{
		const i32 is_top = (cast(usize)ptr.addr + ptr.size >= top) && (cast(usize)ptr.addr + num <= cast(usize)heap_limit);
		if(is_top || ptr.size >= num) 
		{
			if(prev != null) 
			{
				prev.next = ptr.next;
			} 
			else 
			{
				heap.free = ptr.next;
			}
			ptr.next  = heap.used;
			heap.used = ptr;

			static if(!MALLOC_SPLIT)
			{
				if(is_top) 
				{
					ptr.size = num;
					heap.top = cast(usize)ptr.addr + num;
					static if(MALLOC_COMPACT)
					{
						Compact();
					}
				}
			}
			else
			{
				if(is_top) 
				{
					ptr.size = num;
					heap.top = cast(usize)ptr.addr + num;
				} 
				else if(heap.fresh != null) 
				{
					usize excess = ptr.size - num;
					if(excess >= heap_split_thresh) 
					{
						ptr.size        = num;
						MemBlock *split = heap.fresh;
						heap.fresh      = split.next;
						split.addr      = cast(void *)(cast(usize)ptr.addr + num);
						split.size      = excess;
						InsertBlock(split);
						static if(MALLOC_COMPACT)
						{
							Compact();
						}
					}
				}
			}

			return ptr;
		}

		prev = ptr;
		ptr  = ptr.next;
	}

	// no matching free blocks
	// see if any other blocks available
	usize new_top = top + num;
	if(heap.fresh != null && new_top <= cast(usize)heap_limit) 
	{
		ptr         = heap.fresh;
		heap.fresh = ptr.next;
		ptr.addr   = cast(void *)top;
		ptr.next   = heap.used;
		ptr.size   = num;
		heap.used  = ptr;
		heap.top   = new_top;

		return ptr;
	}

	return null;
}

usize
BytesToPages(usize count)
{
	usize pages = count/(1024*64);
	return pages > 0 ? pages : 1;
}

void*
malloc(usize num)
{
	MemBlock *block = AllocBlock(num);

	if(block == null)
	{
		WasmGrow(BytesToPages(num));
		usize new_pages = WasmSize();
		usize page_diff = new_pages - current_pages;
		if(page_diff > 0)
		{
			current_pages = new_pages;
			heap_limit    = cast(void*)WasmPageSize(current_pages);

			block = AllocBlock(num);
		}
	}

	return block != null ? block.addr : null;
}

void 
MemClear(void *ptr, usize num) 
{
	usize *ptrw = cast(usize*)ptr;
	usize numw  = (num & -usize.sizeof) / usize.sizeof;
	while(numw--) 
	{
		*ptrw++ = 0;
	}
	num     &= (usize.sizeof - 1);
	u8* ptrb = cast(u8*)ptrw;
	while(num--) 
	{
		*ptrb++ = 0;
	}
}

void*
calloc(usize num, usize size) 
{
	num *= size;
	MemBlock *block = AllocBlock(num);

	{
		MemClear(block.addr, num);
		return block.addr;
	}

	return null;
}

void*
realloc(void *ptr, usize size)
{
	void  *new_ptr  = malloc(size);
	usize prev_size = PtrInfo(ptr);

	memcpy(new_ptr, ptr, prev_size);

	free(ptr);

	return new_ptr;
}

usize 
CountBlocks(MemBlock *ptr) 
{
	usize num = 0;
	while(ptr != null) 
	{
		num++;
		ptr = ptr.next;
	}
	return num;
}

usize MallocFreeCount() 
{
	return CountBlocks(heap.free);
}

usize MallocUsedCount() 
{
	return CountBlocks(heap.used);
}

usize MallocMallocFreshCount() 
{
	return CountBlocks(heap.fresh);
}

bool MallocCheck() 
{
	return heap_max_blocks == MallocFreeCount() + MallocUsedCount() + MallocMallocFreshCount();
}

unittest
{
	{
		string zeroed = AllocZeroed("Test");
		assert(zeroed.ptr[4] == '\0');
	}
}