// #TODO: Consider how to approach ExpandableArenas

#load "arena_platform.jai";

// Counted in the arena's own pages, so this is 64 KB on a 4 KB-page arena and 32 MB on one
// backed by 2 MB large pages. Rounded up to a power of two by arena_init, which is what lets the
// arena store it as a log2.
ARENA_DEFAULT_COMMIT_PAGE_COUNT :: 16;
ARENA_DEFAULT_ALIGNMENT :: 8; // CPU register width
PLATFORM_MEMORY_PAGE_SIZE :: 4096;
PLATFORM_MEMORY_PAGE_SIZE_LOG2 :: 12;
#run assert((1 << PLATFORM_MEMORY_PAGE_SIZE_LOG2) == PLATFORM_MEMORY_PAGE_SIZE);
PLATFORM_CACHE_LINE_SIZE :: 64; // #Planned: unused today. Reserved for cache-line padding of
                                // per-thread or hot arena fields once those exist.

Arena_Sizes :: s64.[
    64 * 1024,
     2 * 1024 * 1024,
    64 * 1024 * 1024,
     2 * 1024 * 1024 * 1024,
    64 * 1024 * 1024 * 1024,
     2 * 1024 * 1024 * 1024 * 1024,
    64 * 1024 * 1024 * 1024 * 1024,
];

Arena_Reserve :: enum u8 {
  Size_64K :: 0;
  Size_2M  :: 1;
  Size_64M :: 2;
  Size_2G  :: 3;
  Size_64G :: 4;
  Size_2T  :: 5;
  Size_64T :: 6;
}

#run assert(Arena_Sizes.count == type_info(Arena_Reserve).values.count);

Arena_Flags :: enum_flags u8 {
  None            :: 0;
  Chained         :: 0b00000001; // #Planned: chained/expandable arenas. NO_CHAIN is the only behaviour.
  Is_Bootstrapped :: 0b00000010; // Live: set by arena_set_bootstrap_flag; read by arena_delete
                                 // and free_pages_down_to.
  Large_Pages     :: 0b01000000; // Live: pass it to arena_init or bootstrap_arena. The reservation
                                 // is backed by the platform's large/huge pages and the arena's
                                 // page granularity becomes platform_large_page_size() rather
                                 // than PLATFORM_MEMORY_PAGE_SIZE. Requesting it is a hint: if the
                                 // OS refuses, the arena falls back to 4 KB pages and records the
                                 // miss in Arena_Status.Large_Pages_Failed rather than failing to initialize.
                                 //
                                 // On Windows and macOS a granted request commits and wires the
                                 // WHOLE reservation up front -- there is no reserve-now-commit-
                                 // later form -- so the reserve class is what becomes resident.
                                 // Those platforms also refuse partial decommit, which makes
                                 // free_pages_down_to a no-op for such an arena.
                                 //
                                 // Large pages are already non-pageable, so a .Large_Pages arena
                                 // that is also .Secure_Arena skips the per-commit page lock: the
                                 // guarantee holds without it.
  Secure_Arena    :: 0b10000000; // Live: pass it to arena_init or bootstrap_arena. Every page the
                                 // arena commits is locked into physical memory as it is committed,
                                 // and unlocked before it is decommitted. Locking a range once at
                                 // the top level does not survive this arena's growth or its
                                 // release cycle, which is why it belongs on the commit path.
}

Arena_Status :: enum_flags u8 {
  None               :: 0;
  Secure_Lock_Failed :: 0b00000001; // Sticky: a range that spent time unlocked cannot be made retroactively safe by a later success.
  Large_Pages_Failed :: 0b00000010; // Sticky: pages already handed out at 4 KB granularity do not become large ones later.
}

arena_status_is_set :: (arena: *Arena, flag: Arena_Status) -> bool #expand {
  return (arena.status & flag) == flag;
}

arena_is_secure :: (arena: *Arena) -> bool #expand {
  return (arena.flags & .Secure_Arena) == .Secure_Arena;
}

// Whether a secure arena has actually held its guarantee.
arena_secure_lock_intact :: (arena: *Arena) -> bool {
  if !is_valid(arena) return false;
  return arena_is_secure(arena) && !arena_status_is_set(arena, .Secure_Lock_Failed);
}

arena_is_large_pages :: (arena: *Arena) -> bool #expand {
  return (arena.flags & .Large_Pages) == .Large_Pages;
}

// Whether a large-page arena actually got large pages, as opposed to falling back to 4 KB.
arena_large_pages_active :: (arena: *Arena) -> bool {
  if !is_valid(arena) return false;
  return arena_is_large_pages(arena) && !arena_status_is_set(arena, .Large_Pages_Failed);
}

arena_page_size :: (arena: *Arena) -> s64 #expand {
  return cast(s64) 1 << arena.page_size_log2;
}

arena_alignment :: (arena: *Arena) -> s64 #expand {
  return cast(s64) 1 << arena.alignment_log2;
}

arena_initial_commit_pages :: (arena: *Arena) -> s64 #expand {
  return cast(s64) 1 << arena.initial_commit_page_count_log2;
}

Arena :: struct {
  current_point:*u8;
  memory_base: *u8;
  first_uncommitted_page: *u8;
  reserve: Arena_Reserve;
  flags: Arena_Flags;
  status: Arena_Status;
  alignment_log2: u8;
  page_size_log2: u8 = PLATFORM_MEMORY_PAGE_SIZE_LOG2; // Set by arena_platform_init. The ground truth for every page calculation.
  initial_commit_page_count_log2: u8;
  #if ARENA_DEBUG {
    high_water_mark: s64;
    label: string;
    file_path: string;
    function_name: string;
    line_number: s32;
  }
}

bootstrap_arena :: (new_reserve: Arena_Reserve, default_commit_page_count: s32 = ARENA_DEFAULT_COMMIT_PAGE_COUNT, flags: Arena_Flags = .None, label:="", loc := #caller_location)
       -> *Arena #expand {
  arena_ptr := bootstrap_arena_internal(new_reserve, default_commit_page_count, flags, label, loc);
  
  #if ARENA_DEBUG {
    if arena_ptr {
      push_allocator(context.default_allocator);
      arena_ptr.label     = copy_string(label);
      arena_ptr.file_path = copy_string(loc.fully_pathed_filename);
      this_allocation_is_not_a_leak(arena_ptr.label.data);
      this_allocation_is_not_a_leak(arena_ptr.file_path.data);
      arena_ptr.line_number = cast(s32)loc.line_number;
      arena_ptr.function_name = #procedure_name();
    }
  }
  
  return arena_ptr;
}

bootstrap_arena_internal :: (new_reserve: Arena_Reserve, default_commit_page_count: s32, flags: Arena_Flags, label: string, loc: Source_Code_Location) -> *Arena {
  new_arena: Arena;
  success, pages_committed := arena_init(*new_arena, new_reserve, default_commit_page_count, flags);
  if !success || !pages_committed {
    arena_delete(*new_arena);
    return null;
  }
  
  arena_ptr: *Arena = xx new_arena.memory_base;
  memcpy(arena_ptr, *new_arena, size_of(Arena));
  arena_ptr.current_point = arena_start(arena_ptr);
  
  arena_set_bootstrap_flag(arena_ptr);

  return arena_ptr;
}

arena_allocator_proc :: (mode: Allocator_Mode, requested_size: s64, old_size: s64, old_memory: *void, allocator_data: *void) -> *void {
  arena := cast(*Arena)allocator_data;
  
  if #complete mode == {
    case .FREE;
      return null;
    case .RESIZE;     #through;
    case .ALLOCATE;
      result := arena_alloc(arena, requested_size);
      if mode == .RESIZE {
          size_to_copy := ifx old_size < requested_size then old_size else requested_size;
          if result && size_to_copy  memcpy(result, old_memory, size_to_copy);
      }
      return result;

    case .STARTUP;      #through;
    case .SHUTDOWN;
      return null;

    case .THREAD_START; #through;
    case .THREAD_STOP;
      return null;

    case .CREATE_HEAP; #through;
    case .DESTROY_HEAP;
      return null;

    case .IS_THIS_YOURS;
      if !arena return null;
      // The contract is 'allocated by your heap and is currently live'.
      if old_memory >= arena_start(arena) && old_memory < arena.current_point {
        return cast(*void)true;
      }
      return null;

    case .CAPS;
      if old_memory { (cast(*string)old_memory).* = "modules/Basic arena_allocator_proc"; }
      return cast(*void)(Allocator_Caps.IS_THIS_YOURS | .HINT_I_AM_A_FAST_BUMP_ALLOCATOR);
  }
}

arena_init :: (using arena: *Arena, new_reserve: Arena_Reserve, commit_page_count:=ARENA_DEFAULT_COMMIT_PAGE_COUNT, new_flags: Arena_Flags = .None)
      -> (success: bool, pages_committed: s64 = 0) {
  if arena == null {
    log_error("[Basic/Arena] Attempting to initialize a null arena!");
    return false;
  }
  
  if (arena.memory_base == null) {
    // Before arena_platform_init, so .Large_Pages reaches the reserve call, and before the first
    // commit, so .Secure_Arena locks through the same path every later growth uses.
    flags |= new_flags;

    if !arena_platform_init(arena, new_reserve) {
      return false;
    }
    current_point = arena_start(arena);
    arena_set_alignment(arena, ARENA_DEFAULT_ALIGNMENT);

    commit_pages := ifx commit_page_count > 0
      then cast(s64) 1 << ceil_log2(commit_page_count)
      else cast(s64) commit_page_count;

    // Page size is not known until the reserve has happened, so the commit size is sized here.
    commit_size := commit_pages * arena_page_size(arena);
    if commit_size < 0 {
      log_warning("[Basic/Arena] Invalid commit_page_count passed to `arena_init`. Using default parameter: % pages",
        ARENA_DEFAULT_COMMIT_PAGE_COUNT);
      commit_size = ARENA_DEFAULT_COMMIT_PAGE_COUNT * arena_page_size(arena);
    }

    bytes_committed := arena_commit_first_pages(arena, commit_size);

    committed_pages := bytes_committed >> page_size_log2;
    if committed_pages <= 0  return true, 0; // The commit failed; bootstrap_arena_internal checks for this.

    initial_commit_page_count_log2 = log2_of_power_of_two(committed_pages);

    return true, committed_pages;
  }

  return true, arena_initial_commit_pages(arena); // already initialized
}

// start_offset_bytes is for when we're doing arena-backed structures such as Arena_Array and we may need more 
// space than `commit_size` because we're using some space in the first page. 
// #Planned: no caller passes it yet -- Arena_Array does not exist. It is implemented and guarded,
// so it works if you pass it, but it is untested and should be treated as such until it has a user.
arena_commit_first_pages :: (using arena: *Arena, commit_size: s64, start_offset_bytes: s64 = 0) 
    -> (bytes_committed: s64) {
  assert(start_offset_bytes >= 0); 
  if (start_offset_bytes < 0) {
    log_error("[Basic/Arena] Invalid start_offset_bytes parameter provided to arena_commit_first_pages! Given: %", start_offset_bytes);
    return 0;
  }
  
  commit_size_final := commit_size;
  
  // Clamp commit size:
  if commit_size <= 0 {
    commit_size_final = ARENA_DEFAULT_COMMIT_PAGE_COUNT * arena_page_size(arena);
    log_warning("[Basic/Arena] Requested commit size is invalid! Got % expected 0 < commit_size <= %. Using default of %!", 
        commit_size, reserve_size_bytes(arena), commit_size_final);
  }
  if commit_size > reserve_size_bytes(arena) {
    log_warning("[Basic/Arena] Requested commit size is out of expected range! Got % expected 0 < commit_size <= %. Clamping to max size!", 
        commit_size, reserve_size_bytes(arena));
    commit_size_final = reserve_size_bytes(arena);
  }
  
  final_address := Align(memory_base + commit_size_final + start_offset_bytes, arena_page_size(arena));
  
  commit_size_final = cast(s64)(final_address - memory_base);
  
  if (commit_size_final > reserve_size_bytes(arena)) {
    log_warning("[Basic/Arena] Requested commit size larger than available address space! Got % expected <= %. Clamping to max size!", 
        commit_size_final, reserve_size_bytes(arena));
    commit_size_final = reserve_size_bytes(arena);
  }
  
  final_address = memory_base + commit_size_final;
  
  if !extend_committed_pages(arena, final_address) {
    return 0;
  }
  
  return arena_committed_bytes(arena);
}

arena_reset_keeping_memory :: arena_reset;
arena_reset :: (using arena: *Arena) {
  if !is_valid(arena) return;
  current_point = arena_start(arena);
}

arena_release :: (using arena: *Arena) {
  if !is_valid(arena) return;
  current_point = arena_start(arena);
  
  free_pages_down_to(arena, arena_initial_commit_pages(arena));
}

arena_reset_overwriting_memory :: (using arena: *Arena, $wipe_function:=memory_secure_overwrite) {
  if !is_valid(arena) return;
  arena_wipe(arena_start(arena), current_point, wipe_function);
  arena_reset(arena);
}

arena_alloc :: (using arena: *Arena, byte_count: s64) -> *void {
  if !arena {
    log_error("[Basic/Arena] Passed a null pointer to arena_alloc!");
    return null;
  }
  
  if arena.memory_base == null {
    log_error("Attempting to allocate using an uninitialized Arena!\n Initializing with a 64GB capacity.");
    
    arena.reserve = .Size_64G;
    success := arena_init(arena, arena.reserve);
    
    if !success return null;
  }
  
  if byte_count < 0 {
    log_error("[Basic/Arena] Arena %: negative allocation size %.\n", arena, byte_count);
    return null;
  }
  
  result := Align(current_point, arena_alignment(arena));

  if byte_count > arena_address_limit(arena) - result {
    log_error("[Basic/Arena] Attempt to allocate, but Arena % is full! (Allocation size %; alignment %; address space size %; already allocated %.)\n",
        arena, byte_count, arena_alignment(arena), reserve_size_bytes(arena), current_point - memory_base);
    return null;
  }
  
  result_end := byte_count + result;
  if result_end > first_uncommitted_page {
    if !extend_committed_pages(arena, result_end)  return null;
  }
  
  current_point = result_end;
  
  update_high_water_mark(arena);
  
  return result;
}

arena_secure_overwrite_starting_from :: (arena: *Arena, start: *u8) {
  arena_wipe(start, arena.first_uncommitted_page);
}

Arena_Secure_Overwrite_Endpoint :: enum s32 {
  Current_Point; // erase from arena_start to the current point
  First_Uncommitted_Page;
}

arena_secure_overwrite :: (arena: *Arena, $endpoint: Arena_Secure_Overwrite_Endpoint=.Current_Point) {
  #if endpoint == {
    case .Current_Point;  
      arena_wipe(arena_start(arena), arena.current_point);
    case .First_Uncommitted_Page;
      arena_wipe(arena_start(arena), arena.first_uncommitted_page);
  }
}

// Arena @Utility
arena_set_alignment :: inline (arena: *Arena, alignment: s64) {
  assert(is_valid(arena));
  assert(alignment > 0 && (alignment & (alignment - 1)) == 0,
         "Arena alignment must be a non-zero power of two; got %.", alignment);
  if alignment <= 0 || (alignment & (alignment - 1)) != 0 {
    log_error("[Basic/Arena][arena_set_alignment] Attempting to set an invalid alignment `%`! Keeping original value `%`",
            alignment, arena_alignment(arena));
    return;
  }

  limit := reserve_size_bytes(arena);
  if alignment > limit {
    log_error("[Basic/Arena][arena_set_alignment] Attempting to set an alignment `%` wider than the % byte reservation! Clamping to the reservation",
            alignment, limit);
    arena.alignment_log2 = log2_of_power_of_two(limit);
    return;
  }

  arena.alignment_log2 = log2_of_power_of_two(alignment);
}

arena_set_bootstrap_flag :: (arena: *Arena) #expand {
  arena.flags = (arena.flags | .Is_Bootstrapped);
}

arena_is_bootstrapped :: (arena: *Arena) -> bool #expand {
  return (arena.flags & .Is_Bootstrapped == .Is_Bootstrapped);
}

update_high_water_mark :: (arena: *Arena) #expand {
  #if ARENA_DEBUG {
    current_usage := arena_usage_bytes(arena);
    if (current_usage > arena.high_water_mark) {
      arena.high_water_mark = current_usage;
    }
  }
}

// Every arena reserves `size_of(Arena)` bytes at its base, whether or not it is bootstrapped.
//
// The reason is not the cost of testing `arena_is_bootstrapped` here. This is a macro, and
// `flags` shares a cache line with `memory_base`, so the test would be a byte load, an AND and a
// CMOV, predicted perfectly because the bit never changes after init, less than `arena_alloc`
// already spends around this call. What a branch would cost instead:
//
// - `arena_start` would stop being a pure function of `memory_base`, which is written once at
//   init. Every caller that compares two of its results across an arena's life: `arena_usage_bytes`,
//   the IS_THIS_YOURS range test, the reset paths, the `Arena_State` marks
//   depends on it not moving, and `flags` is mutable and public.
//
// - The flag is not set yet at the two places that need it most. `arena_init` runs on a stack
//   local before anything knows the arena will be bootstrapped, and `bootstrap_arena_internal`
//   computes `current_point` from here one line before `arena_set_bootstrap_flag`, so a
//   branching version would park `current_point` on top of the header it means to skip.
//
// Note that `size_of(Arena)` is larger under ARENA_DEBUG, so `arena_start` does not return the
// same offset in debug and release builds. *Never persist an offset into arena memory across a
// change of build configuration*
arena_start :: (arena: *Arena) -> *u8 #expand {
  return Align(arena.memory_base + size_of(Arena), ARENA_DEFAULT_ALIGNMENT);
}

arena_address_limit :: (arena: *Arena) -> *u8 #expand {
  return (arena.memory_base + reserve_size_bytes(arena));
}

arena_usage_bytes :: (arena: *Arena) -> s64 #expand {
  return cast(s64)(arena.current_point - arena_start(arena));
}

arena_committed_bytes :: (arena: *Arena) -> s64 #expand {
  return cast(s64)(arena.first_uncommitted_page - arena.memory_base);
}

reserve_size_bytes :: (arena_reserve: Arena_Reserve) -> s64 {
  return Arena_Sizes[cast(s32)arena_reserve];
}

reserve_size_bytes :: (arena: *Arena) -> s64 #expand {
  return reserve_size_bytes(arena.reserve);
}

is_valid :: (arena: *Arena) -> bool {
  return (arena != null && arena.memory_base != null);
}

get_allocator :: (arena: *Arena) -> Allocator #expand {
  return .{arena_allocator_proc, arena};
}

next_reserve :: (size: s64, alignment: s64) -> Arena_Reserve {
  size_with_overhead := Align(size + size_of(Arena), alignment);
  
  for i: 0..enum_count(Arena_Reserve)-1 {
    if size_with_overhead <= Arena_Sizes[i] {
      return cast(Arena_Reserve)i;
    }
  }
  
  return .Size_64T;
}

// @Utility `push_allocator` style macros
push_arena :: (arena: *Arena) #expand #no_debug {
  assert(is_valid(arena));
  
  old_allocator := context.allocator;
  context.allocator = .{arena_allocator_proc, arena};
  
  `defer context.allocator = old_allocator;
}

push_alignment :: (arena: *Arena, alignment: s64) #expand #no_debug {
  assert(is_valid(arena));
  assert(alignment > 0 && (alignment & (alignment - 1)) == 0,
         "Arena alignment must be a non-zero power of two; got %.", alignment);
  old_alignment := arena_alignment(arena);
  if alignment <= 0 || (alignment & (alignment - 1)) != 0 {
    log_error("[Basic/Arena][push_alignment] Attempting to set an invalid alignment `%`! Keeping original value `%`",
            alignment, old_alignment);
    return;
  }
  arena_set_alignment(arena, alignment);

  `defer arena_set_alignment(arena, old_alignment);
}

auto_reset_overwriting_memory :: (arena: *Arena, $wipe_function:=memory_secure_overwrite) #expand #no_debug {
  __mark := get_arena_state(arena);
  `defer restore_arena_state(arena, __mark, false, wipe_function);  
}

auto_release_overwriting_memory :: (arena: *Arena, $wipe_function:=memory_secure_overwrite) #expand #no_debug {
  __mark := get_arena_state(arena);
  `defer restore_arena_state(arena, __mark, true, wipe_function);
}

auto_reset :: (arena: *Arena) #expand {
  __mark := get_arena_state(arena);
  `defer restore_arena_state(arena, __mark, false, wipe_function=null);
}

auto_release :: (arena: *Arena) #expand {
  __mark := get_arena_state(arena);
  `defer restore_arena_state(arena, __mark, true, wipe_function=null);
}

// Temporary for debugging: 
arena_details :: (arena: *Arena, loc := #caller_location) -> string {
  push_allocator(temp);

  builder: String_Builder;

  if arena == null {
      print_to_builder(*builder, "(arena_details) Passed a null pointer @ %\n", loc);
      return builder_to_string(*builder);
  }
  
  print_to_builder(*builder, "Details for Arena % with memory_base %\n", arena, arena.memory_base);
  
  if arena.memory_base == null {
      return builder_to_string(*builder);
  }
  
  reserve := arena_address_limit(arena) - arena.memory_base;
  committed_size := arena.first_uncommitted_page - arena.memory_base;
  used_size := arena.current_point - arena_start(arena);

  print_to_builder(*builder, " - committed: %, used: %, reserve: %\n", format_bytes(committed_size), format_bytes(used_size), format_bytes(reserve));
  print_to_builder(*builder, " - page size: %\n", format_bytes(arena_page_size(arena)));
  if arena_is_large_pages(arena) {
    print_to_builder(*builder, " - large pages: % (requested)\n",
      ifx arena_status_is_set(arena, .Large_Pages_Failed) then "unavailable, fell back" else "active");
  }

  return builder_to_string(*builder);
}

arena_wipe :: (start: *u8, end: *u8, $wipe_function:=memory_secure_overwrite) #expand {
  #if wipe_function {
    if end <= start  return;
    wipe_function(start, cast(s64)(end - start));
  }
}
#scope_module
Arena_State :: struct {
    starting_arena : *Arena;
    starting_point : *u8;
}

get_arena_state :: inline (arena: *Arena) -> Arena_State {
    assert(arena != null);
    
    state: Arena_State = ---;
    state.starting_arena = arena;
    state.starting_point = state.starting_arena.current_point;
    
    return state;
}

restore_arena_state :: inline (arena: *Arena, state: Arena_State, $free_pages: bool, $wipe_function:=memory_secure_overwrite) {
  assert(arena == state.starting_arena);
  assert(is_valid(state.starting_arena));
  
  #if wipe_function {
    arena_wipe(state.starting_point, state.starting_arena.current_point, wipe_function);
  }
  
  state.starting_arena.current_point = state.starting_point;
      
  #if free_pages {
    using state.starting_arena;
    
    last_page := cast(*void)Align(current_point, arena_page_size(arena));
    keep_page_count := cast(s64)(last_page - memory_base) >> page_size_log2;
    initial_pages := arena_initial_commit_pages(arena);
    if keep_page_count < initial_pages  keep_page_count = initial_pages;
    
    free_pages_down_to(arena, keep_page_count);
  }
}