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

Arena :: struct {
   proc: ArenaProc;
   data: *void;
}

ArenaProc :: #type (
   mode:       ArenaMode,
   arena_data: *void,

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

   caller: Source_Code_Location
) -> *void;

ArenaMode :: 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 an 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 an arena, allowing it to be used with Jai's builtin Allocator system.
ArenaToAllocator :: (arena: *Arena) -> Allocator {
   return .{ proc = JaiAllocatorProc, data = arena };
}


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

PanicArenaProc :: (
   mode: ArenaMode,
   arena_data: *void,

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

   caller: Source_Code_Location
) -> *void {
   if mode == {
      case .Acquire;
         if new_size > 0
          { Panic("jc: cannot acquire memory using the PanicArena", loc = caller); }

      case .Save;
         Panic("jc: cannot save using the PanicArean", loc = caller);
      case .Restore;
         Panic("jc: cannot restore using the PanicArean", loc = caller);
   }

   return null;
} @jc.nodocs

BumpArenaProc :: (
   mode:       ArenaMode,
   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 mode == {
      case .Setup;
         if bump.memory.data == null
          { Panic("jc: BumpArena has no 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 :: (
   mode: ArenaMode,
   arena_data: *void,

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

   caller: Source_Code_Location
) -> *void {
   Panic("not implemented for PagingArena", loc = caller);
   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 :: (
   mode: ArenaMode,
   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 mode == {
      // @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(mode, arena_data, new_size, old_size, align, new_ptr, old_ptr, caller);
} @jc.nodocs


#scope_file

JaiAllocatorProc :: (mode: 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 mode == {
      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