// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/hal/memory/tlsf.h"

#include "iree/base/internal/math.h"

//===----------------------------------------------------------------------===//
// Internal helpers
//===----------------------------------------------------------------------===//

// Returns a pointer to the block node at |index| via stride arithmetic.
static inline iree_hal_memory_tlsf_block_t* iree_hal_memory_tlsf_block_at(
    const iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t index) {
  return (iree_hal_memory_tlsf_block_t*)(tlsf->block_storage +
                                         (iree_host_size_t)index *
                                             tlsf->block_stride);
}

// Returns a pointer to the inline frontier of a block node.
static inline iree_async_frontier_t* iree_hal_memory_tlsf_block_frontier(
    const iree_hal_memory_tlsf_t* tlsf, iree_hal_memory_tlsf_block_t* block) {
  return (iree_async_frontier_t*)((uint8_t*)block + tlsf->frontier_offset);
}

// Grows the block node pool using 2x doubling. New nodes are linked into the
// unused_node_head free list. Uses iree_allocator_grow_array for
// overflow-checked capacity doubling and realloc.
static iree_status_t iree_hal_memory_tlsf_grow_pool(
    iree_hal_memory_tlsf_t* tlsf) {
  if (tlsf->block_capacity >= UINT32_MAX) {
    return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                            "block pool at maximum capacity (%" PRIu32 ")",
                            (uint32_t)UINT32_MAX);
  }
  iree_host_size_t old_capacity = tlsf->block_capacity;
  iree_host_size_t new_capacity = old_capacity;
  IREE_RETURN_IF_ERROR(iree_allocator_grow_array(
      tlsf->host_allocator, /*minimum_capacity=*/old_capacity + 1,
      tlsf->block_stride, &new_capacity, (void**)&tlsf->block_storage));
  // Cap at UINT32_MAX since block indices are uint32_t.
  if (new_capacity > UINT32_MAX) new_capacity = UINT32_MAX;
  // Initialize new nodes and link them into the unused list.
  for (iree_host_size_t i = old_capacity; i < new_capacity; ++i) {
    iree_hal_memory_tlsf_block_t* block =
        iree_hal_memory_tlsf_block_at(tlsf, (uint32_t)i);
    memset(block, 0, tlsf->block_stride);
    block->next_free = tlsf->unused_node_head;
    tlsf->unused_node_head = (uint32_t)i;
  }
  tlsf->block_capacity = new_capacity;
  return iree_ok_status();
}

// Allocates a block node from the pool. Grows the pool if no unused nodes
// are available.
static iree_status_t iree_hal_memory_tlsf_alloc_node(
    iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t* out_index) {
  if (tlsf->unused_node_head == IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    IREE_RETURN_IF_ERROR(iree_hal_memory_tlsf_grow_pool(tlsf));
  }
  iree_hal_memory_tlsf_block_index_t index = tlsf->unused_node_head;
  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, index);
  tlsf->unused_node_head = block->next_free;
  memset(block, 0, tlsf->block_stride);
  *out_index = index;
  return iree_ok_status();
}

// Returns a block node to the unused pool.
static void iree_hal_memory_tlsf_free_node(
    iree_hal_memory_tlsf_t* tlsf, iree_hal_memory_tlsf_block_index_t index) {
  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, index);
  memset(block, 0, tlsf->block_stride);
  block->next_free = tlsf->unused_node_head;
  tlsf->unused_node_head = index;
}

// Maps a block length to (FL, SL) indices for insertion into the free list
// matrix. The block goes into the bin whose range contains |length|.
static inline void iree_hal_memory_tlsf_mapping_insert(
    iree_device_size_t length, uint8_t* out_fl, uint8_t* out_sl) {
  // FL = position of the most significant set bit.
  int fl = 63 - iree_math_count_leading_zeros_u64(length);
  // SL = next SL_LOG2 bits below the MSB.
  int sl_shift = (fl > IREE_HAL_MEMORY_TLSF_SL_LOG2)
                     ? (fl - IREE_HAL_MEMORY_TLSF_SL_LOG2)
                     : 0;
  *out_fl = (uint8_t)fl;
  *out_sl =
      (uint8_t)((length >> sl_shift) & (IREE_HAL_MEMORY_TLSF_SL_COUNT - 1));
}

// Maps a requested allocation length to (FL, SL) indices for searching the
// free list matrix. We want the smallest bin that could contain a block >=
// |length|, so we round up within the current FL level's SL range.
static inline void iree_hal_memory_tlsf_mapping_search(
    iree_device_size_t length, uint8_t* out_fl, uint8_t* out_sl) {
  // For sizes that span SL bins within an FL level, round up to ensure we
  // find a block that is at least |length|. We add (1 << sl_shift) - 1
  // before extracting the SL index, which effectively rounds up.
  int fl = 63 - iree_math_count_leading_zeros_u64(length);
  if (fl > IREE_HAL_MEMORY_TLSF_SL_LOG2) {
    int sl_shift = fl - IREE_HAL_MEMORY_TLSF_SL_LOG2;
    iree_device_size_t rounded =
        length + (((iree_device_size_t)1 << sl_shift) - 1);
    // If rounding caused overflow into the next FL level, use that instead.
    int new_fl = 63 - iree_math_count_leading_zeros_u64(rounded);
    if (new_fl != fl) {
      *out_fl = (uint8_t)new_fl;
      *out_sl = 0;
      return;
    }
    *out_fl = (uint8_t)fl;
    *out_sl =
        (uint8_t)((rounded >> sl_shift) & (IREE_HAL_MEMORY_TLSF_SL_COUNT - 1));
  } else {
    *out_fl = (uint8_t)fl;
    *out_sl = (uint8_t)(length & (IREE_HAL_MEMORY_TLSF_SL_COUNT - 1));
  }
}

// Inserts a free block into the appropriate (FL, SL) free list and updates
// the bitmaps.
static void iree_hal_memory_tlsf_insert_free_block(
    iree_hal_memory_tlsf_t* tlsf, iree_hal_memory_tlsf_block_index_t index) {
  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, index);
  IREE_ASSERT(block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE);

  uint8_t fl = 0, sl = 0;
  iree_hal_memory_tlsf_mapping_insert(block->length, &fl, &sl);

  // Insert at the head of the free list for this (FL, SL) bin.
  iree_hal_memory_tlsf_block_index_t old_head = tlsf->free_lists[fl][sl];
  block->prev_free = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  block->next_free = old_head;
  if (old_head != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    iree_hal_memory_tlsf_block_at(tlsf, old_head)->prev_free = index;
  }
  tlsf->free_lists[fl][sl] = index;

  // Set bitmap bits.
  tlsf->fl_bitmap |= (1ull << fl);
  tlsf->sl_bitmaps[fl] |= (1u << sl);

  tlsf->free_block_count++;
}

// Removes a free block from its (FL, SL) free list and updates bitmaps if
// the list becomes empty.
static void iree_hal_memory_tlsf_remove_free_block(
    iree_hal_memory_tlsf_t* tlsf, iree_hal_memory_tlsf_block_index_t index) {
  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, index);
  IREE_ASSERT(block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE);

  uint8_t fl = 0, sl = 0;
  iree_hal_memory_tlsf_mapping_insert(block->length, &fl, &sl);

  // Unlink from the doubly-linked free list.
  if (block->prev_free != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    iree_hal_memory_tlsf_block_at(tlsf, block->prev_free)->next_free =
        block->next_free;
  } else {
    // Was the head of the list.
    tlsf->free_lists[fl][sl] = block->next_free;
  }
  if (block->next_free != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    iree_hal_memory_tlsf_block_at(tlsf, block->next_free)->prev_free =
        block->prev_free;
  }

  block->prev_free = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  block->next_free = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;

  // Clear bitmap bits if the list is now empty.
  if (tlsf->free_lists[fl][sl] == IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    tlsf->sl_bitmaps[fl] &= ~(1u << sl);
    if (tlsf->sl_bitmaps[fl] == 0) {
      tlsf->fl_bitmap &= ~(1ull << fl);
    }
  }

  tlsf->free_block_count--;
}

// Searches for a free block of at least |length| bytes using bitmap scans.
// Returns IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE if no suitable block exists.
static iree_hal_memory_tlsf_block_index_t
iree_hal_memory_tlsf_find_suitable_block(iree_hal_memory_tlsf_t* tlsf,
                                         iree_device_size_t length) {
  uint8_t fl = 0, sl = 0;
  iree_hal_memory_tlsf_mapping_search(length, &fl, &sl);

  // Try to find a block in the current FL level at or above the target SL.
  uint32_t sl_map = tlsf->sl_bitmaps[fl] & (~0u << sl);
  if (sl_map != 0) {
    int found_sl = iree_math_count_trailing_zeros_u32(sl_map);
    return tlsf->free_lists[fl][found_sl];
  }

  // No block in the current FL level; search higher FL levels.
  // The mask ~0ull << (fl + 1) is UB when fl==63 (shift by 64). Rewrite using
  // iree_shr: iree_shr(~0ull, 63 - fl) produces a mask with the bottom fl+1
  // bits set, and its complement gives us only the FL levels above fl.
  uint64_t fl_map = tlsf->fl_bitmap & ~iree_shr(~0ull, 63 - fl);
  if (fl_map == 0) {
    return IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  }
  int found_fl = iree_math_count_trailing_zeros_u64(fl_map);
  // Take the first SL bin in the found FL level.
  IREE_ASSERT(tlsf->sl_bitmaps[found_fl] != 0);
  int found_sl = iree_math_count_trailing_zeros_u32(tlsf->sl_bitmaps[found_fl]);
  return tlsf->free_lists[found_fl][found_sl];
}

// Merges |source_frontier| into |target_block|'s inline frontier. If the
// merge overflows capacity, marks the target block as tainted and increments
// the diagnostic counter.
static void iree_hal_memory_tlsf_merge_frontiers(
    iree_hal_memory_tlsf_t* tlsf, iree_hal_memory_tlsf_block_t* target_block,
    const iree_async_frontier_t* source_frontier) {
  if (!source_frontier || source_frontier->entry_count == 0) return;

  iree_async_frontier_t* target_frontier =
      iree_hal_memory_tlsf_block_frontier(tlsf, target_block);

  // If target is already tainted, nothing to do; the frontier is meaningless.
  if (target_block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED) return;

  if (!iree_async_frontier_merge(target_frontier, tlsf->frontier_capacity,
                                 source_frontier)) {
    // Overflow: mark tainted and zero the frontier.
    target_block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED;
    target_frontier->entry_count = 0;
    tlsf->tainted_coalesce_count++;
  }
}

//===----------------------------------------------------------------------===//
// API
//===----------------------------------------------------------------------===//

iree_status_t iree_hal_memory_tlsf_initialize(
    iree_hal_memory_tlsf_options_t options, iree_allocator_t host_allocator,
    iree_hal_memory_tlsf_t* out_tlsf) {
  IREE_ASSERT_ARGUMENT(out_tlsf);
  memset(out_tlsf, 0, sizeof(*out_tlsf));

  // Validate options.
  if (options.range_length == 0) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "range_length must be > 0");
  }
  if (options.alignment == 0) {
    options.alignment = IREE_HAL_MEMORY_TLSF_MIN_ALIGNMENT;
  }
  if (options.alignment < IREE_HAL_MEMORY_TLSF_MIN_ALIGNMENT) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "alignment (%" PRIdsz ") must be >= %" PRIdsz, options.alignment,
        (iree_device_size_t)IREE_HAL_MEMORY_TLSF_MIN_ALIGNMENT);
  }
  if (!iree_device_size_is_power_of_two(options.alignment)) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "alignment (%" PRIdsz ") must be a power of two",
                            options.alignment);
  }
  if (options.range_length < options.alignment) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "range_length (%" PRIdsz
                            ") must be >= alignment (%" PRIdsz ")",
                            options.range_length, options.alignment);
  }
  // Round range_length down to a multiple of alignment.
  options.range_length = options.range_length & ~(options.alignment - 1);
  if (options.frontier_capacity == 0) {
    options.frontier_capacity = IREE_HAL_MEMORY_TLSF_DEFAULT_FRONTIER_CAPACITY;
  }

  // Compute block node layout using overflow-checked struct math.
  // Each block node is: [fixed fields] [padding] [frontier header] [entries]
  // The frontier must be aligned for its entry type (8-byte aligned).
  iree_host_size_t frontier_offset = 0;
  iree_host_size_t block_stride = 0;
  IREE_RETURN_IF_ERROR(IREE_STRUCT_LAYOUT(
      sizeof(iree_hal_memory_tlsf_block_t), &block_stride,
      IREE_STRUCT_FIELD_ALIGNED(1, iree_async_frontier_t,
                                iree_alignof(iree_async_frontier_entry_t),
                                &frontier_offset),
      IREE_STRUCT_FIELD(options.frontier_capacity, iree_async_frontier_entry_t,
                        NULL)));

  // Determine initial block pool capacity.
  iree_host_size_t initial_capacity = options.initial_block_capacity;
  if (initial_capacity == 0) {
    initial_capacity = IREE_HAL_MEMORY_TLSF_DEFAULT_INITIAL_BLOCK_CAPACITY;
  }

  // Allocate block storage (overflow-checked array allocation).
  uint8_t* block_storage = NULL;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc_array(
      host_allocator, initial_capacity, block_stride, (void**)&block_storage));

  // Initialize the TLSF instance.
  out_tlsf->range_length = options.range_length;
  out_tlsf->alignment = options.alignment;
  out_tlsf->fl_min =
      (uint8_t)iree_math_count_trailing_zeros_u64(options.alignment);
  out_tlsf->frontier_capacity = options.frontier_capacity;
  out_tlsf->fl_bitmap = 0;
  memset(out_tlsf->sl_bitmaps, 0, sizeof(out_tlsf->sl_bitmaps));
  for (int fl = 0; fl < IREE_HAL_MEMORY_TLSF_FL_COUNT; ++fl) {
    for (int sl = 0; sl < IREE_HAL_MEMORY_TLSF_SL_COUNT; ++sl) {
      out_tlsf->free_lists[fl][sl] = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
    }
  }
  out_tlsf->block_storage = block_storage;
  out_tlsf->block_stride = block_stride;
  out_tlsf->frontier_offset = frontier_offset;
  out_tlsf->block_count = 0;
  out_tlsf->block_capacity = initial_capacity;
  out_tlsf->unused_node_head = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  out_tlsf->bytes_allocated = 0;
  out_tlsf->bytes_free = options.range_length;
  out_tlsf->allocation_count = 0;
  out_tlsf->free_block_count = 0;
  out_tlsf->tainted_coalesce_count = 0;
  out_tlsf->host_allocator = host_allocator;

  // Link all pre-allocated nodes into the unused list (except node 0 which
  // we will use for the initial free block).
  for (iree_host_size_t i = 1; i < initial_capacity; ++i) {
    iree_hal_memory_tlsf_block_t* block =
        iree_hal_memory_tlsf_block_at(out_tlsf, (uint32_t)i);
    block->next_free = out_tlsf->unused_node_head;
    out_tlsf->unused_node_head = (uint32_t)i;
  }
  out_tlsf->block_count = initial_capacity;

  // Create the initial free block spanning the entire range.
  iree_hal_memory_tlsf_block_t* initial_block =
      iree_hal_memory_tlsf_block_at(out_tlsf, 0);
  initial_block->offset = 0;
  initial_block->length = options.range_length;
  initial_block->prev_physical = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  initial_block->next_physical = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  initial_block->prev_free = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  initial_block->next_free = IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  initial_block->flags = IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE |
                         IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST;
  // Frontier starts empty.
  iree_async_frontier_t* initial_frontier =
      iree_hal_memory_tlsf_block_frontier(out_tlsf, initial_block);
  iree_async_frontier_initialize(initial_frontier, 0);

  // Insert the initial free block into the appropriate FL/SL bin.
  iree_hal_memory_tlsf_insert_free_block(out_tlsf, 0);

  return iree_ok_status();
}

void iree_hal_memory_tlsf_deinitialize(iree_hal_memory_tlsf_t* tlsf) {
  IREE_ASSERT_ARGUMENT(tlsf);
  if (tlsf->allocation_count > 0) {
    IREE_ASSERT(false, "TLSF deinitialize with %" PRIu32 " leaked allocations",
                tlsf->allocation_count);
  }
  iree_allocator_free(tlsf->host_allocator, tlsf->block_storage);
  memset(tlsf, 0, sizeof(*tlsf));
}

iree_status_t iree_hal_memory_tlsf_try_allocate(
    iree_hal_memory_tlsf_t* tlsf, iree_device_size_t length,
    iree_hal_memory_tlsf_allocation_t* out_allocation,
    iree_hal_memory_tlsf_allocate_result_t* out_result) {
  IREE_ASSERT_ARGUMENT(tlsf);
  IREE_ASSERT_ARGUMENT(out_allocation);
  IREE_ASSERT_ARGUMENT(out_result);
  memset(out_allocation, 0, sizeof(*out_allocation));
  *out_result = IREE_HAL_MEMORY_TLSF_ALLOCATE_EXHAUSTED;

  if (length == 0) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "allocation length must be > 0");
  }

  // Guard against overflow: if length is so large that rounding up to
  // alignment would wrap around, reject immediately. This prevents
  // near-SIZE_MAX requests from silently succeeding as tiny allocations.
  if (length > IREE_DEVICE_SIZE_MAX - (tlsf->alignment - 1)) {
    return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                            "allocation length %" PRIdsz
                            " overflows when aligned to %" PRIdsz,
                            length, tlsf->alignment);
  }

  // Round up to alignment (and ensure at least minimum block size).
  iree_device_size_t aligned_length =
      iree_device_align(length, tlsf->alignment);
  if (aligned_length < tlsf->alignment) {
    aligned_length = tlsf->alignment;
  }

  // Find a suitable free block.
  iree_hal_memory_tlsf_block_index_t block_index =
      iree_hal_memory_tlsf_find_suitable_block(tlsf, aligned_length);
  if (block_index == IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    return iree_ok_status();
  }

  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, block_index);

  iree_device_size_t remainder = block->length - aligned_length;
  iree_hal_memory_tlsf_block_index_t remainder_index =
      IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE;
  if (remainder >= tlsf->alignment) {
    // Allocate the remainder node before unlinking the selected free block so
    // metadata growth failure leaves allocator state unchanged.
    IREE_RETURN_IF_ERROR(
        iree_hal_memory_tlsf_alloc_node(tlsf, &remainder_index));
  }

  // Remove the block from the free list after all fallible split preparation.
  iree_hal_memory_tlsf_remove_free_block(tlsf, block_index);

  // Re-fetch block: alloc_node may have grown the block pool via realloc,
  // invalidating all prior block pointers into block_storage.
  block = iree_hal_memory_tlsf_block_at(tlsf, block_index);

  if (remainder_index != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    iree_hal_memory_tlsf_block_t* remainder_block =
        iree_hal_memory_tlsf_block_at(tlsf, remainder_index);
    remainder_block->offset = block->offset + aligned_length;
    remainder_block->length = remainder;
    remainder_block->prev_physical = block_index;
    remainder_block->next_physical = block->next_physical;
    remainder_block->flags = IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE;

    // Transfer LAST flag if the original block was last.
    if (block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST) {
      remainder_block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST;
      block->flags &= ~IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST;
    }

    // Update the old right neighbor's prev_physical to point to the remainder.
    if (remainder_block->next_physical !=
        IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
      iree_hal_memory_tlsf_block_at(tlsf, remainder_block->next_physical)
          ->prev_physical = remainder_index;
    }

    // Link the remainder into the physical list.
    block->next_physical = remainder_index;
    block->length = aligned_length;

    // The remainder block gets an empty frontier (it was just split from a
    // block being allocated; it has no independent usage history).
    iree_async_frontier_t* remainder_frontier =
        iree_hal_memory_tlsf_block_frontier(tlsf, remainder_block);
    iree_async_frontier_initialize(remainder_frontier, 0);

    // Insert the remainder into the free list.
    iree_hal_memory_tlsf_insert_free_block(tlsf, remainder_index);

    // Update free bytes (the remainder stays free).
    tlsf->bytes_free -= aligned_length;
  } else {
    // Cannot split; give the whole block (may be slightly over-sized).
    tlsf->bytes_free -= block->length;
  }

  // Mark the block as allocated: clear the FREE flag before populating the
  // result so the caller sees only the allocated-state flags (LAST, TAINTED).
  block->flags &= ~IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE;

  iree_async_frontier_t* frontier =
      iree_hal_memory_tlsf_block_frontier(tlsf, block);
  out_allocation->offset = block->offset;
  out_allocation->length = block->length;
  out_allocation->block_index = block_index;
  out_allocation->death_frontier =
      (frontier->entry_count > 0) ? frontier : NULL;
  out_allocation->block_flags = block->flags;

  tlsf->bytes_allocated += block->length;
  tlsf->allocation_count++;

  *out_result = IREE_HAL_MEMORY_TLSF_ALLOCATE_OK;
  return iree_ok_status();
}

iree_status_t iree_hal_memory_tlsf_allocate(
    iree_hal_memory_tlsf_t* tlsf, iree_device_size_t length,
    iree_hal_memory_tlsf_allocation_t* out_allocation) {
  iree_hal_memory_tlsf_allocate_result_t result =
      IREE_HAL_MEMORY_TLSF_ALLOCATE_EXHAUSTED;
  IREE_RETURN_IF_ERROR(
      iree_hal_memory_tlsf_try_allocate(tlsf, length, out_allocation, &result));
  if (result == IREE_HAL_MEMORY_TLSF_ALLOCATE_EXHAUSTED) {
    return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                            "no free block of %" PRIdsz
                            " bytes; largest free block is %" PRIdsz " bytes",
                            length,
                            iree_hal_memory_tlsf_largest_free_block(tlsf));
  }
  return iree_ok_status();
}

static void iree_hal_memory_tlsf_free_impl(
    iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t block_index,
    const iree_async_frontier_t* death_frontier, bool preserve_block_metadata) {
  IREE_ASSERT_ARGUMENT(tlsf);
  IREE_ASSERT(block_index != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE);

  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, block_index);
  IREE_ASSERT(!(block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE));

  // Update stats for the deallocation.
  tlsf->bytes_allocated -= block->length;
  tlsf->bytes_free += block->length;
  tlsf->allocation_count--;

  // Mark as free before coalescing. Normal free() replaces the block's
  // dependency metadata with |death_frontier|. restore() preserves the current
  // frontier and taint state so speculative candidate rejection does not
  // launder a stale/tainted block into a fresh one.
  block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE;
  if (!preserve_block_metadata) {
    block->flags &= ~IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED;

    // Copy the death frontier into the block's inline storage.
    iree_async_frontier_t* block_frontier =
        iree_hal_memory_tlsf_block_frontier(tlsf, block);
    if (death_frontier && death_frontier->entry_count > 0) {
      // Copy entries (clamped to capacity, though the caller should have sized
      // the frontier to fit). If the death frontier itself exceeds capacity,
      // we taint immediately.
      if (death_frontier->entry_count <= tlsf->frontier_capacity) {
        if (block_frontier != death_frontier) {
          memcpy(block_frontier, death_frontier,
                 sizeof(iree_async_frontier_t) +
                     (iree_host_size_t)death_frontier->entry_count *
                         sizeof(iree_async_frontier_entry_t));
        }
      } else {
        // Death frontier too large for inline storage; taint.
        iree_async_frontier_initialize(block_frontier, 0);
        block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED;
        tlsf->tainted_coalesce_count++;
      }
    } else {
      iree_async_frontier_initialize(block_frontier, 0);
    }
  }

  // Coalesce with the right neighbor first (so that the left coalesce can
  // absorb the combined block).
  if (!(block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST)) {
    iree_hal_memory_tlsf_block_index_t right_index = block->next_physical;
    if (right_index != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
      iree_hal_memory_tlsf_block_t* right_block =
          iree_hal_memory_tlsf_block_at(tlsf, right_index);
      if (right_block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE) {
        // Remove the right block from its free list.
        iree_hal_memory_tlsf_remove_free_block(tlsf, right_index);

        // Merge frontiers: right into this block.
        iree_hal_memory_tlsf_merge_frontiers(
            tlsf, block,
            iree_hal_memory_tlsf_block_frontier(tlsf, right_block));

        // Absorb geometry: extend this block to cover the right block.
        block->length += right_block->length;

        // Transfer LAST flag.
        if (right_block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST) {
          block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST;
        }

        // Transfer taint flag. A tainted right neighbor means the merged
        // block's frontier is untrustworthy; the merge_frontiers call above
        // will have been a no-op (tainted source has entry_count==0), but the
        // taint itself must propagate. Zero the frontier to prevent stale data
        // from being misinterpreted as valid entries.
        if (right_block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED) {
          block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED;
          iree_hal_memory_tlsf_block_frontier(tlsf, block)->entry_count = 0;
        }

        // Update physical links.
        block->next_physical = right_block->next_physical;
        if (right_block->next_physical !=
            IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
          iree_hal_memory_tlsf_block_at(tlsf, right_block->next_physical)
              ->prev_physical = block_index;
        }

        // Release the right block node.
        iree_hal_memory_tlsf_free_node(tlsf, right_index);
      }
    }
  }

  // Coalesce with the left neighbor.
  if (block->prev_physical != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
    iree_hal_memory_tlsf_block_index_t left_index = block->prev_physical;
    iree_hal_memory_tlsf_block_t* left_block =
        iree_hal_memory_tlsf_block_at(tlsf, left_index);
    if (left_block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_FREE) {
      // Remove the left block from its free list.
      iree_hal_memory_tlsf_remove_free_block(tlsf, left_index);

      // Merge frontiers: this block into left.
      iree_hal_memory_tlsf_merge_frontiers(
          tlsf, left_block, iree_hal_memory_tlsf_block_frontier(tlsf, block));

      // Absorb geometry: extend left to cover this block.
      left_block->length += block->length;

      // Transfer LAST flag.
      if (block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST) {
        left_block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_LAST;
      }

      // Transfer taint flag. As with right coalesce, a tainted source means
      // the merged block's frontier cannot be trusted; zero it.
      if (block->flags & IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED) {
        left_block->flags |= IREE_HAL_MEMORY_TLSF_BLOCK_FLAG_TAINTED;
        iree_hal_memory_tlsf_block_frontier(tlsf, left_block)->entry_count = 0;
      }

      // Update physical links.
      left_block->next_physical = block->next_physical;
      if (block->next_physical != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE) {
        iree_hal_memory_tlsf_block_at(tlsf, block->next_physical)
            ->prev_physical = left_index;
      }

      // Release this block node and switch to inserting the left block.
      iree_hal_memory_tlsf_free_node(tlsf, block_index);
      block_index = left_index;
      block = left_block;
    }
  }

  // Insert the (possibly merged) block into the appropriate free list.
  iree_hal_memory_tlsf_insert_free_block(tlsf, block_index);
}

void iree_hal_memory_tlsf_free(iree_hal_memory_tlsf_t* tlsf,
                               iree_hal_memory_tlsf_block_index_t block_index,
                               const iree_async_frontier_t* death_frontier) {
  iree_hal_memory_tlsf_free_impl(tlsf, block_index, death_frontier,
                                 /*preserve_block_metadata=*/false);
}

void iree_hal_memory_tlsf_restore(
    iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t block_index) {
  iree_hal_memory_tlsf_free_impl(tlsf, block_index, /*death_frontier=*/NULL,
                                 /*preserve_block_metadata=*/true);
}

void iree_hal_memory_tlsf_query_stats(const iree_hal_memory_tlsf_t* tlsf,
                                      iree_hal_memory_tlsf_stats_t* out_stats) {
  IREE_ASSERT_ARGUMENT(tlsf);
  IREE_ASSERT_ARGUMENT(out_stats);
  out_stats->bytes_allocated = tlsf->bytes_allocated;
  out_stats->bytes_free = tlsf->bytes_free;
  out_stats->allocation_count = tlsf->allocation_count;
  out_stats->free_block_count = tlsf->free_block_count;
  out_stats->tainted_coalesce_count = tlsf->tainted_coalesce_count;
}

iree_device_size_t iree_hal_memory_tlsf_largest_free_block(
    const iree_hal_memory_tlsf_t* tlsf) {
  IREE_ASSERT_ARGUMENT(tlsf);
  if (tlsf->fl_bitmap == 0) return 0;

  // Find the highest populated FL level.
  int fl = 63 - iree_math_count_leading_zeros_u64(tlsf->fl_bitmap);
  IREE_ASSERT(tlsf->sl_bitmaps[fl] != 0);
  // Find the highest populated SL within that FL level.
  int sl = 31 - iree_math_count_leading_zeros_u32(tlsf->sl_bitmaps[fl]);
  // The head block of that list is the largest (or close to it).
  iree_hal_memory_tlsf_block_index_t head = tlsf->free_lists[fl][sl];
  IREE_ASSERT(head != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE);
  return iree_hal_memory_tlsf_block_at(tlsf, head)->length;
}

const iree_async_frontier_t* iree_hal_memory_tlsf_block_death_frontier(
    const iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t block_index) {
  IREE_ASSERT_ARGUMENT(tlsf);
  IREE_ASSERT(block_index != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE);
  iree_hal_memory_tlsf_block_t* block =
      iree_hal_memory_tlsf_block_at(tlsf, block_index);
  iree_async_frontier_t* frontier =
      iree_hal_memory_tlsf_block_frontier(tlsf, block);
  return (frontier->entry_count > 0) ? frontier : NULL;
}

iree_hal_memory_tlsf_block_flags_t iree_hal_memory_tlsf_block_flags(
    const iree_hal_memory_tlsf_t* tlsf,
    iree_hal_memory_tlsf_block_index_t block_index) {
  IREE_ASSERT_ARGUMENT(tlsf);
  IREE_ASSERT(block_index != IREE_HAL_MEMORY_TLSF_BLOCK_INDEX_NONE);
  return iree_hal_memory_tlsf_block_at(tlsf, block_index)->flags;
}
