#add_context arena      := Arena.{ proc = PagingArenaProc };
#add_context temp_arena := Arena.{ proc = TempArenaProc   };

/// Arena is an interface for an arena-based memory allocation.
Arena :: struct {
   proc: ArenaProc;
   data: *void;
}

ArenaProc :: #type (
   event:       ArenaEvent,
   arena_data: *void,

   new_size: int, old_size: int, align: int,
   new_ptr: *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void;

ArenaEvent :: enum {
   Setup;
   Teardown;

   Acquire;
   Reset;

   Save;
   Restore;
}

/// ArenaPush pushes a value of type T to the current context.arena.
ArenaPush :: ($T: Type, loc := #caller_location) -> *T {
   return ArenaPushSize(size_of(T), loc = loc).(*T);
}

/// ArenaPushSize pushes size bytes to the current context.arena.
ArenaPushSize :: (size: int, align := DefaultAlign, loc := #caller_location) -> *void {
   if size < 0
    { Panic("jc: size cannot be negative", loc = loc); }
   TrySetupArena(*context.arena);
   return context.arena.proc(.Acquire, context.arena.data, size, 0, align, null, null, loc);
}

ArenaSave :: (loc := #caller_location) -> ArenaSaveResult {
   TrySetupArena(*context.arena, loc);
   return context.arena.proc(.Save, context.arena.data, 0, 0, 0, null, null, loc).(ArenaSaveResult);
}

ArenaRestore :: (savepoint: ArenaSaveResult, loc := #caller_location) {
   context.arena.proc(.Restore, context.arena.data, 0, 0, 0, null, savepoint.(*void), loc);
}

ArenaAutoRestore :: () #expand {
   wm := ArenaSave();
   `defer ArenaRestore(wm);
}

/// ArenaSaveResult is an arena-specific value used for saving/restoring.
/// see: ArenaSave
ArenaSaveResult :: #type,distinct *void;


/// ArenaReset resets arena's state, allowing its memory to be reused.
ArenaReset :: (loc := #caller_location) {
   TrySetupArena(*context.arena);
   context.arena.proc(.Reset, context.arena.data, 0, 0, 0, null, null, loc);
}

/// ArenaRelease frees the current context.arena's memory.
ArenaRelease :: (loc := #caller_location) {
   TrySetupArena(*context.arena, loc);
   context.arena.proc(.Teardown, context.arena.data, 0, 0, 0, null, null, loc);
}

/// ArenaToAllocator wraps the given arena, allowing it to be used with Jai's builtin Allocator system.
ArenaToAllocator :: (arena: *Arena) -> Allocator {
   return .{ proc = JaiAllocatorProc, data = arena };
}

/// PanicArena is an Arena that panics when used.
PanicArena :: () -> Arena {
   return .{ proc = PanicArenaProc };
}

BumpData :: struct {
   memory: []u8;
   offset: int;
}

BumpArena :: (bump: *BumpData) -> Arena {
   return .{ proc = BumpArenaProc, data = bump };
}

PagingArena :: () -> Arena {
   return .{ proc = PagingArenaProc };
}


#scope_module

TrySetupArena :: (arena: *Arena, loc := #caller_location) #expand {
   if arena.data != null
    { return; }

   data := arena.proc(.Setup, null, 0, 0, 0, null, null, loc);
   if data != null
    { arena.data = data; }
} @jc.nodocs

PanicArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   if event == {
      case .Acquire;
         if new_size > 0
          { Panic("jc: PanicArena, acquiring memory is not allowed", loc = caller); }

      case .Save;
         Panic("jc: PanicArena, saving is not allowed", loc = caller);
      case .Restore;
         Panic("jc: PanicArena, restoring is not allowed", loc = caller);
   }

   return null;
} @jc.nodocs

BumpArenaProc :: (
   event:       ArenaEvent,
   arena_data: *void,

   new_size: int, old_size: int, align: int,
   new_ptr: *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void #no_abc {
   bump := arena_data.(*BumpData);
   if event == {
      case .Setup;
         if bump.memory.data == null
          { Panic("jc: BumpArena, no backing memory", loc = caller); }
         return arena_data;

      case .Teardown;
         // Do nothing, we don't own this memory

      case .Acquire;
         end := bump.memory.data + bump.offset;
         if new_size == 0
          { return end; }

         ptr  := MemAlignForward(end, align);
         size := new_size + (ptr - end);
         if bump.offset + size > bump.memory.count
          { Panic("jc: BumpArena, out of memory", loc = caller); }

         bump.offset += size;
         return ptr;
      case .Reset;
         bump.offset = 0;

      case .Save;
         return bump.offset.(*void);
      case .Restore;
         wm := old_ptr.(int);
         if wm < 0 || wm >= bump.memory.count
          { Panic("jc: BumpArena, restored invalid savepoint", loc = caller); }
         bump.offset = wm;
   }

   return null;
} @jc.nodocs

PagingArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   MaxPages :: 64;

   PageArenaData :: struct {
      pages:  [MaxPages]Page;
      count:  int;
      in_use: int;

      Page :: struct {
         base:   *u8;
         offset: int;
      }
   }

   // @note(judah): page 0 has a special offset because it stores the header
   Page0Offset :: #run AlignForward(size_of(PageArenaData), align_of(PageArenaData));

   NewPage :: inline (previous_base: *void) -> PageArenaData.Page {
      p: PageArenaData.Page = ---;
      p.base   = MemAcquirePage(.CommitImmediately);
      p.offset = 0;
      return p;
   }

   data := arena_data.(*PageArenaData);
   if event == {
      case .Setup;
         page := NewPage(null);
         data = page.base.(*PageArenaData);

         page.offset = Page0Offset;
         data.pages[0] = page;

         data.in_use = 1;
         data.count  = 1;

         return data;

      case .Teardown;
         // @note(judah): release in reverse oreder so we don't release the memory
         // containing the header until we don't need it anymore.
         i := data.count - 1;
         while i >= 0 {
            page := data.pages[i];
            if page.base != null
             { MemReleasePage(page.base, PageSize); }

            i -= 1;
         }

      case .Acquire;
         Assert(data.count != 0, "jc: PageArena, no pages exist");

         last := *data.pages[data.in_use - 1];
         end  := last.base + last.offset;
         if new_size == 0
          { return end; }

         ptr  := MemAlignForward(end, align);
         size := new_size + (ptr - end);

         // create or reuse page
         if last.offset + size > PageSize {
            if data.in_use >= MaxPages
             { Panic("jc: PageArena, max page count"); }

            page: *PageArenaData.Page;
            if data.in_use < data.count {
               page = *data.pages[data.in_use];
            }
            else {
               p := NewPage(data.pages[data.count - 1].base);
               data.pages[data.count] = p;
               page = *data.pages[data.count];

               data.count  += 1;
            }

            data.in_use += 1;

            ptr  = page.base;
            size = new_size;
            last = page;
         }

         last.offset += size;
         return ptr;

      case .Reset;
         for 0..data.count - 1 {
            page := *data.pages[it];

            if it == 0 {
               page.offset = Page0Offset;
            }
            else {
               page.offset = 0;
            }
         }

         data.in_use = 1;

      case .Save;
         return data.in_use.(*void);

      case .Restore;
         in_use_at_save := old_ptr.(int);

         // @note(judah): We need to restore from the savepoint to count
         // because we could've created more pages inbetween save/restore
         for in_use_at_save..data.count - 1 {
            p := *data.pages[it];
            if it == 0 {
               p.offset = Page0Offset;
            }
            else {
               p.offset = 0;
            }
         }

         data.in_use = in_use_at_save;
   }

   return null;
} @jc.nodocs

// @note(judah): TempArenaProc just wraps BumpArenaProc, but provides its own backing memory.
// This is needed so we can use TempArenaProc without forcing the user to set everything up.
TempArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   // @temp(judah)
   libc :: #library,system "libc";
   malloc :: (size: int) -> *void #c_call #foreign libc;
   free   :: (ptr: *void) #c_call #foreign libc;

   if event == {
      // @note(judah): allows the temp arena to initialize itself without relying on another arena to provide its memory.
      case .Setup;
         mcount := AlignForward(size_of(BumpData) + 32768).(int);
         memory := malloc(mcount).(*u8);
         MemZero(memory, mcount);

         bump := memory.(*BumpData);
         bump.memory.data  = memory + size_of(BumpData);
         bump.memory.count = mcount - size_of(BumpData);

         return bump;

      case .Teardown;
         free(arena_data);
   }

   return BumpArenaProc(event, arena_data, new_size, old_size, align, new_ptr, old_ptr, caller);
} @jc.nodocs


#scope_file

JaiAllocatorProc :: (event: Allocator_Mode, requested_size: s64, old_size: s64, old_memory: *void, allocator_data: *void) -> *void {
   CallerFromStackTrace :: () -> Source_Code_Location #expand {
      node := context.stack_trace;
      if node.next != null {
         node = node.next;
      }
      return node.info.location;
   }

   arena := allocator_data.(*Arena);
   if event == {
      case .STARTUP;
         TrySetupArena(arena);
         return null;

      case .ALLOCATE;
         loc := CallerFromStackTrace();
         return ArenaPushSize(requested_size, loc = loc,, arena = arena);

      case .RESIZE;
         loc := CallerFromStackTrace();
         new := ArenaPushSize(requested_size, loc = loc,, arena = arena);
         MemCopy(new, old_memory, old_size);
         return new;

      case .FREE;
         // Arenas don't free single pointers so give this a value that's obviously incorrect
         MemOverwrite(old_memory, old_size, 0xAA);
   }

   return null;
} @jc.nodocs