power_play/src/base/base_arena.h

////////////////////////////////////////////////////////////
//~ Arena types

#define ArenaHeaderSize CachelineSize
#define ArenaBlockSize  16384

Struct(Arena)
{
    u64 pos;
    u64 committed;
    u64 reserved;
#if RtcIsEnabled
    b32 readonly;
#endif
};

Struct(TempArena)
{
    Arena *arena;
    u64 start_pos;

#if RtcIsEnabled
    u64 scratch_id;
#endif
};

////////////////////////////////////////////////////////////
//~ State types

#define ScratchArenasPerCtx 2

Struct(FiberArenaCtx)
{
    Arena *perm_arena;
    Arena *scratch_arenas[ScratchArenasPerCtx];
};

Struct(SharedArenaCtx)
{
    FiberArenaCtx arena_contexts[MaxFibers];
} extern shared_arena_ctx;

////////////////////////////////////////////////////////////
//~ Arena push/pop

#define PushStruct(a, type) ((type *)PushBytes((a), sizeof(type), alignof(type)))
#define PushStructNoZero(a, type) ((type *)PushBytesNoZero((a), sizeof(type), alignof(type)))

#define PushStructs(a, type, n) ((type *)PushBytes((a), (sizeof(type) * (n)), alignof(type)))
#define PushStructsNoZero(a, type, n) ((type *)PushBytesNoZero((a), (sizeof(type) * (n)), alignof(type)))

#define PopStruct(a, type, dst) PopBytes((a), sizeof(type), dst)
#define PopStructs(a, type, n, dst) PopBytes((a), sizeof(type) * (n), dst)

#define PopStructNoCopy(a, type) PopBytesNoCopy((a), sizeof(type))
#define PopStructsNoCopy(a, type, n) PopBytesNoCopy((a), sizeof(type) * (n))

/* TOOD: Replace ArenaBase with 'ArenaFirst(type)' for dynamic-array-style uses */
#define ArenaBase(arena) ((u8 *)(arena) + ArenaHeaderSize)
#define ArenaCount(arena, type) ((arena)->pos / sizeof(type))

/* Returns a pointer to where the next push would be (at alignment of type).
 * Equivalent to PushStruct but without actually allocating anything or modifying the arena. */
#define PushDry(a, type) (type *)(_PushDry((a), alignof(type)))

void *PushBytesNoZero(Arena *arena, u64 size, u64 align);

Inline void *PushBytes(Arena *arena, u64 size, u64 align)
{
    void *p = PushBytesNoZero(arena, size, align);
    ZeroBytes(p, size);
    return p;
}

Inline void PopTo(Arena *arena, u64 pos)
{
    Assert(arena->pos >= pos);
    Assert(!arena->readonly);

    AsanPoison(ArenaBase(arena) + pos, arena->pos - pos);
    arena->pos = pos;
}

Inline void PopBytesNoCopy(Arena *arena, u64 size)
{
    Assert(arena->pos >= size);
    Assert(!arena->readonly);
    u64 new_pos = arena->pos - size;
    AsanPoison(ArenaBase(arena) + new_pos, arena->pos - new_pos);
    arena->pos = new_pos;
}

Inline void PopBytes(Arena *arena, u64 size, void *copy_dst)
{
    Assert(arena->pos >= size);
    Assert(!arena->readonly);

    u64 new_pos = arena->pos - size;
    void *src = (void *)(ArenaBase(arena) + new_pos);
    CopyBytes(copy_dst, src, size);

    AsanPoison(ArenaBase(arena) + new_pos, arena->pos - new_pos);
    arena->pos = new_pos;
}

Inline void *_PushDry(Arena *arena, u64 align)
{
    u64 aligned_start_pos = (arena->pos + (align - 1));
    aligned_start_pos -= aligned_start_pos % align;
    void *ptr = ArenaBase(arena) + aligned_start_pos;
    return ptr;
}

////////////////////////////////////////////////////////////
//~ Arena management

Arena *AcquireArena(u64 reserve);
void ReleaseArena(Arena *arena);

void CopyArena(Arena *dst, Arena *src);
void ShrinkArena(Arena *arena);
void SetArenaReadonly(Arena *arena);
void SetArenaReadWrite(Arena *arena);

Inline void *PushAlign(Arena *arena, u64 align)
{
    Assert(!arena->readonly);
    if (align > 0)
    {
        u64 aligned_start_pos = (arena->pos + (align - 1));
        aligned_start_pos -= aligned_start_pos % align;
        u64 align_bytes = aligned_start_pos - (u64)arena->pos;
        if (align_bytes > 0)
        {
            return (void *)PushStructsNoZero(arena, u8, align_bytes);
        }
        else
        {
            return (void *)(ArenaBase(arena) + arena->pos);
        }
    }
    else
    {
        /* 0 alignment */
        Assert(0);
        return (void *)(ArenaBase(arena) + arena->pos);
    }
}

Inline void *ResetArena(Arena *arena)
{
    PopTo(arena, 0);
    return (void *)ArenaBase(arena);
}

////////////////////////////////////////////////////////////
//~ Temp arena operations

Inline TempArena BeginTempArena(Arena *arena)
{
    TempArena t = ZI;
    t.arena = arena;
    t.start_pos = arena->pos;
    return t;
}

Inline void EndTempArena(TempArena temp)
{
    PopTo(temp.arena, temp.start_pos);
}

////////////////////////////////////////////////////////////
//~ Fiber arena ctx operations

Inline FiberArenaCtx *FiberArenaCtxFromId(i16 fiber_id)
{
    SharedArenaCtx *g = &shared_arena_ctx;
    FiberArenaCtx *ctx = &g->arena_contexts[fiber_id];
    if (!ctx->perm_arena)
    {
        __profn("Initialize fiber arena ctx");
        ctx->perm_arena = AcquireArena(Gibi(64));
        for (i32 i = 0; i < (i32)countof(ctx->scratch_arenas); ++i)
        {
            ctx->scratch_arenas[i] = AcquireArena(Gibi(64));
        }
    }
    return ctx;
}

#define PermArena() (FiberArenaCtxFromId(FiberId())->perm_arena)

////////////////////////////////////////////////////////////
//~ Scratch helpers

/* Any parameterized arenas in the caller's scope should be passed into this
 * function as a potential "conflict". This is to prevent friction in case the
 * passed arena is itself a scratch arena from another scope (since
 * parameterized arenas are often used to allocate persistent results for the
 * caller).
 *
 * Use `BeginScratchNoConflict` instead if there is no arena in the current
 * scope that could potentially be a scratch arena from another scope. */
Inline TempArena BeginScratch(Arena *potential_conflict)
{
    /* This function is currently hard-coded to search through 2 scratch arenas */
    StaticAssert(ScratchArenasPerCtx == 2);

    /* Use `BeginScratchNoConflict` if no conflicts are present */
    Assert(potential_conflict != 0);

    FiberArenaCtx *ctx = FiberArenaCtxFromId(FiberId());
    Arena *scratch_arena = ctx->scratch_arenas[0];
    if (scratch_arena == potential_conflict)
    {
        scratch_arena = ctx->scratch_arenas[1];
    }
    TempArena temp = BeginTempArena(scratch_arena);
    return temp;
}

/* This macro declares an unused "arena" variable that will error if an existing "arena"
 * variable is present (due to shadowing). This is for catching obvious cases of
 * `BeginScratchNoConflict` getting called when an `arena` variable already
 * exists in the caller's scope (`BeginScratch(arena)` should be called
 * instead). */
#define BeginScratchNoConflict() \
    BeginScratchNoConflict_(); \
    do { \
        u8 arena = 0; \
        LAX arena; \
    } while (0)

Inline TempArena BeginScratchNoConflict_(void)
{
    FiberArenaCtx *ctx = FiberArenaCtxFromId(FiberId());
    Arena *scratch_arena = ctx->scratch_arenas[0];
    TempArena temp = BeginTempArena(scratch_arena);
    return temp;
}

Inline void EndScratch(TempArena scratch_temp)
{
    EndTempArena(scratch_temp);
}