/*
File:   cwasm_arena.c
Author: Taylor Robbins
Date:   08\28\2025
Description:
	** Defines a basic memory Arena structure and API
	** Also defines a few global memory arenas: MainWasmMemory and 2x ScratchArenas
*/

// +--------------------------------------------------------------+
// |                         Arena Types                          |
// +--------------------------------------------------------------+
typedef enum ArenaFlag ArenaFlag;
enum ArenaFlag
{
	ArenaFlag_None             = 0x00000000,
	ArenaFlag_Debug            = 0x00000001,
	ArenaFlag_IsMainWasmMemory = 0x00000002,
	ArenaFlag_DefaultAligned4  = 0x00000004,
	ArenaFlag_DefaultAligned8  = 0x00000008,
	ArenaFlag_DefaultAligned16 = 0x00000010,
	ArenaFlag_IsScratch        = 0x00000020,
	ArenaFlag_All              = 0x0000003F,
	ArenaFlag_Count            = 6,
};

typedef struct Arena Arena;
struct Arena
{
	u32 flags; //ArenaFlags
	void* base;
	u32 size;
	u32 used;
};

typedef struct ArenaMark ArenaMark;
struct ArenaMark
{
	Arena* arena;
	u32 mark;
};

// +--------------------------------------------------------------+
// |                        Global Arenas                         |
// +--------------------------------------------------------------+
static Arena MainWasmMemory = ZEROED;
static Arena ScratchArenas[2] = ZEROED;

// +--------------------------------------------------------------+
// |                          Arena API                           |
// +--------------------------------------------------------------+
__attribute__((always_inline)) Arena NewArena(u32 flags, u32 size, void* basePntr)
{
	Arena result = ZEROED;
	result.flags = flags;
	result.size = size;
	result.base = basePntr;
	return result;
}

static inline ArenaMark GetArenaMark(Arena* arena)
{
	ArenaMark result = ZEROED;
	result.arena = arena;
	result.mark = arena->used;
	return result;
}
static inline void ResetToMark(Arena* arena, u32 mark)
{
	arena->used = mark;
}
static inline void ResetToArenaMark(ArenaMark arenaMark) { ResetToMark(arenaMark.arena, arenaMark.mark); }

void* AllocMemAligned(Arena* arena, u32 numBytes, u32 alignment)
{
	Assert(arena != nullptr && arena->base != nullptr);
	u32 realAlignment = alignment;
	if      (realAlignment == UINT32_MAX) { realAlignment = 0; }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned16) != 0) { realAlignment = 16; }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned8)  != 0) { realAlignment = 8;  }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned4)  != 0) { realAlignment = 4;  }
	u32 alignOffset = AlignOffset((u8*)arena->base + arena->used, realAlignment);
	u32 newSize = arena->used + alignOffset + numBytes;
	if (newSize > arena->size)
	{
		if ((arena->flags & ArenaFlag_IsMainWasmMemory) != 0)
		{
			void* newMemory = grow_mem(newSize - arena->size);
			Assert(newMemory == (u8*)arena->base + arena->size);
			arena->size = newSize;
		}
		else { return nullptr; }
	}
	
	void* result = (u8*)arena->base + arena->used;
	arena->used += alignOffset + numBytes;
	return result;
}
static inline void* AllocMemUnaligned(Arena* arena, u32 numBytes) { return AllocMemAligned(arena, numBytes, UINT32_MAX); }
static inline void* AllocMem(Arena* arena, u32 numBytes) { return AllocMemAligned(arena, numBytes, 0); }

static inline Arena* AllocArena(Arena* sourceArena, u32 newArenaFlags, u32 newArenaSize, u32 newArenaAlignment)
{
	//TODO: Do we need to make sure the Arena structure is aligned to it's own struct alignment?
	#if LANGUAGE_IS_C
	u32 arenaStructAlignment = (u32)_Alignof(Arena);
	#else
	u32 arenaStructAlignment = (u32)std::alignment_of<Arena>();
	#endif
	u32 offsetAfterStruct = AlignOffset(sizeof(Arena), newArenaAlignment);
	Arena* result = (Arena*)AllocMemAligned(sourceArena, sizeof(Arena) + offsetAfterStruct + newArenaSize, (arenaStructAlignment >= newArenaAlignment) ? arenaStructAlignment : newArenaAlignment);
	if (result == nullptr) { return nullptr; }
	*result = NewArena(newArenaFlags, newArenaSize, ((u8*)result) + sizeof(Arena) + offsetAfterStruct);
	return result;
}

#if LANGUAGE_IS_C
#define AllocStruct(type, arenaPntr) (type*)AllocMemAligned((arenaPntr), sizeof(type), (u32)_Alignof(type))
#define AllocArray(type, arenaPntr, count) (type*)AllocMemAligned((arenaPntr), sizeof(type) * (count), (u32)_Alignof(type))
#else
#define AllocStruct(type, arenaPntr) (type*)AllocMemAligned((arenaPntr), sizeof(type), (u32)std::alignment_of<type>())
#define AllocArray(type, arenaPntr, count) (type*)AllocMemAligned((arenaPntr), sizeof(type) * (count), (u32)std::alignment_of<type>())
#endif

void* ReallocMemAligned(Arena* arena, void* oldPntr, u32 oldSize, u32 newSize, u32 newAlignment)
{
	if (oldSize == newSize) { return oldPntr; }
	if (oldPntr == nullptr) { return AllocMemAligned(arena, newSize, newAlignment); }
	Assert(arena != nullptr && arena->base != nullptr);
	
	u32 realAlignment = newAlignment;
	if      (realAlignment == UINT32_MAX) { realAlignment = 0; }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned16) != 0) { realAlignment = 16; }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned8)  != 0) { realAlignment = 8;  }
	else if (realAlignment == 0 && (arena->flags & ArenaFlag_DefaultAligned4)  != 0) { realAlignment = 4;  }
	
	if ((u8*)oldPntr + oldSize == (u8*)arena->base + arena->used && IsAlignedTo(oldPntr, realAlignment))
	{
		if (newSize > oldSize)
		{
			void* newSpace = AllocMemUnaligned(arena, newSize - oldSize);
			if (newSpace == nullptr) { return nullptr; }
			Assert(newSpace == (u8*)oldPntr + oldSize);
		}
		else { arena->used -= (oldSize - newSize); }
		return oldPntr;
	}
	else
	{
		void* newPntr = AllocMemAligned(arena, newSize, newAlignment);
		if (oldSize > 0 && newPntr != nullptr) { memcpy(newPntr, oldPntr, oldSize); }
		return newPntr;
	}
}
static inline void* ReallocMemUnaligned(Arena* arena, void* oldPntr, u32 oldSize, u32 newSize) { return ReallocMemAligned(arena, oldPntr, oldSize, newSize, UINT32_MAX); }
static inline void* ReallocMem(Arena* arena, void* oldPntr, u32 oldSize, u32 newSize) { return ReallocMemAligned(arena, oldPntr, oldSize, newSize, 0); }

void InitGlobalArenas(u32 scratchArenasSize)
{
	u32 mainInitialSize = 1024;
	void* mainPntr = grow_mem(mainInitialSize);
	MainWasmMemory = NewArena(ArenaFlag_IsMainWasmMemory|ArenaFlag_DefaultAligned16, mainInitialSize, mainPntr);
	
	for (u32 aIndex = 0; aIndex < ArrayCount(ScratchArenas); aIndex++)
	{
		void* scratchMemory = AllocMem(&MainWasmMemory, scratchArenasSize);
		ScratchArenas[aIndex] = NewArena(ArenaFlag_IsScratch, scratchArenasSize, scratchMemory);
	}
}

// +==============================+
// |         Scratch API          |
// +==============================+
ArenaMark GetScratch1(Arena* conflictArena)
{
	return GetArenaMark((conflictArena == &ScratchArenas[0]) ? &ScratchArenas[1] : &ScratchArenas[0]);
}
static inline ArenaMark GetScratch() { return GetScratch1(nullptr); }

#define ScratchBegin1(scratchName, conflictArenaPntr) ArenaMark scratchName##Mark = GetScratch1(conflictArenaPntr); Arena* scratchName = scratchName##Mark.arena
#define ScratchBegin(scratchName) ScratchBegin1(scratchName, nullptr)
#define ScratchEnd(scratchName) ResetToArenaMark(scratchName##Mark)