// SPDX-License-Identifier: MPL-2.0
#include "memory/slabs/slab_radix_node.h"
#include "std/mem.h"
#include "kernel/log.h"
#include "memory/pez/pez.h"
static sh_uint8 sig[8]=SH_SLAB_RADIX_NODE_SLAB_SIG;
SH_STATUS sh_slab_radix_node_alloc_init(struct sh_slab_radix_node_SLAB_ALLOCATOR* slab_alloc,sh_pba_PAGE_BLOCK_ALLOCATOR *pba) {
  if (slab_alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  slab_alloc->first_slab=SH_NULLPTR;
  slab_alloc->slab_count=0;
  slab_alloc->partial_head=SH_NULLPTR;
  slab_alloc->pba=pba;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_add_slab(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_slab_radix_node_SLAB** out_slab) {
  if (alloc==SH_NULLPTR || out_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_page_VIRTUAL_ADDRESS slab_va;
  SH_STATUS status=sh_pba_alloc(alloc->pba,ptp,&slab_va);
  if (sh_status_error(status)) return status;
  if (alloc->slab_count==0) {
    sh_slab_radix_node_SLAB *new_slab=(sh_slab_radix_node_SLAB*)slab_va;
    for (sh_iter64 i=0;i<sizeof(sig);i++) new_slab->sig[i]=sig[i];
    new_slab->next_slab=SH_NULLPTR;
    new_slab->prev_slab=SH_NULLPTR;
    new_slab->next_partial=SH_NULLPTR;
    new_slab->prev_partial=SH_NULLPTR;
    new_slab->slab_index=alloc->slab_count;
    new_slab->used_count=0;
    for (sh_iter64 i=0;i<16;i++) new_slab->free_bitmap[i]=0;
    new_slab->free_bitmap[0]|=0x3FFFFULL;
    sh_mem_set_8(new_slab->padding,0x00,sizeof(new_slab->padding));
    sh_mem_set_8((sh_uint8*)new_slab->node_bitmap,0x00,sizeof(new_slab->node_bitmap));
    sh_mem_set_8((sh_uint8*)new_slab->nodes,0x00,sizeof(new_slab->nodes));
    alloc->first_slab=new_slab;
    alloc->partial_head=new_slab;
    alloc->slab_count++;
    *out_slab=new_slab;
  } else {
    sh_slab_radix_node_SLAB *new_slab=(sh_slab_radix_node_SLAB*)slab_va;
    for (sh_iter64 i=0;i<sizeof(sig);i++) new_slab->sig[i]=sig[i];
    new_slab->next_slab=alloc->first_slab;
    new_slab->prev_slab=SH_NULLPTR;
    alloc->first_slab->prev_slab=new_slab;
    alloc->first_slab=new_slab;
    new_slab->next_partial=alloc->partial_head;
    new_slab->prev_partial=SH_NULLPTR;
    alloc->partial_head->prev_partial=new_slab;
    alloc->partial_head=new_slab;
    new_slab->slab_index=alloc->slab_count;
    new_slab->used_count=0;
    for (sh_iter64 i=0;i<16;i++) new_slab->free_bitmap[i]=0;
    new_slab->free_bitmap[0]|=0x3FFFFULL;
    sh_mem_set_8(new_slab->padding,0x00,sizeof(new_slab->padding));
    sh_mem_set_8((sh_uint8*)new_slab->node_bitmap,0x00,sizeof(new_slab->node_bitmap));
    sh_mem_set_8((sh_uint8*)new_slab->nodes,0x00,sizeof(new_slab->nodes));
    alloc->slab_count++;
    *out_slab=new_slab;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_get_partial_slab(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL* ptp,sh_slab_radix_node_SLAB** found_slab) {
  if (alloc==SH_NULLPTR ||ptp==SH_NULLPTR || found_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (alloc->partial_head==SH_NULLPTR) {
    *found_slab=SH_NULLPTR;
    return SH_STATUS_NOT_FOUND;
  }
  sh_slab_radix_node_SLAB *slab=alloc->partial_head;
  if (slab->used_count>=SH_SLAB_RADIX_NODE_ACTUAL_OBJECTS_PER_SLAB) {
    if (slab->next_partial!=SH_NULLPTR && slab->next_partial->used_count<SH_SLAB_RADIX_NODE_ACTUAL_OBJECTS_PER_SLAB) {
      alloc->partial_head=slab->next_partial;
      alloc->partial_head->prev_partial=SH_NULLPTR;
      slab->next_partial=SH_NULLPTR;
      slab->prev_partial=SH_NULLPTR;
      *found_slab=alloc->partial_head;
      return SH_STATUS_SUCCESS;
    } else {
      sh_slab_radix_node_SLAB *new_slab;
      SH_STATUS status=sh_slab_radix_node_add_slab(alloc,ptp,&new_slab);
      if (sh_status_error(status)) return status;
      *found_slab=new_slab;
      return SH_STATUS_NEW_SLAB_ADDED;
    }
  }
  *found_slab=slab;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_scan_slabs(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc) {
  if (alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_radix_node_SLAB* current=alloc->first_slab;
  sh_slab_radix_node_SLAB* last_partial=SH_NULLPTR;
  alloc->partial_head=SH_NULLPTR;
  sh_log_debug("Slabs rescan triggered.",SH_LOG_SOURCE_SLAB);
  while (current!=SH_NULLPTR) {
    if (*(sh_uint64*)current->sig!=SH_SLAB_RADIX_NODE_MAGIC) {
      return SH_STATUS_FOUND_CORRUPTED_SLAB;
    }
    current->used_count=0;
    for (sh_iter64 w=0;w<16;w++) current->free_bitmap[w]=0;
    current->free_bitmap[0]|=0x3FFFFULL;
    for (sh_uint16 obj_idx=18;obj_idx<SH_SLAB_RADIX_NODE_OBJECTS_PER_SLAB;obj_idx++) {
      sh_uint8* node_ptr=(sh_uint8*)&current->nodes[obj_idx-18];
      sh_bool is_free=SH_TRUE;
      sh_uint64* check_ptr=(sh_uint64*)node_ptr;
      for (sh_uint8 b=0;b<(SH_SLAB_RADIX_NODE_OBJECT_SIZE_BYTES/8);b++) {
        if (check_ptr[b]!=0) {
          is_free=SH_FALSE;
          break;
        }
      }
      sh_uint16 word=obj_idx>>6;
      sh_uint16 bit=obj_idx&63;
      if (!is_free) {
        current->free_bitmap[word]|=(1ULL<<bit);
        current->used_count++;
      }
    }
    current->next_partial=SH_NULLPTR;
    current->prev_partial=SH_NULLPTR;
    if (current->used_count<SH_SLAB_RADIX_NODE_ACTUAL_OBJECTS_PER_SLAB) {
      if (alloc->partial_head==SH_NULLPTR) {
        alloc->partial_head=current;
        last_partial=current;
      } else {
        last_partial->next_partial=current;
        current->prev_partial=last_partial;
        last_partial=current;
      }
    }
    current=current->next_slab;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_find_free_object(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_radix_NODE** out,sh_slab_radix_node_NODE_INDEX_IN_SLAB* index_in_slab) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out==SH_NULLPTR || index_in_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_radix_node_SLAB *partial_slab;
  sh_slab_radix_node_SLAB *new_slab;
  SH_STATUS status=sh_slab_radix_node_get_partial_slab(alloc,ptp,&partial_slab);
  if (status==SH_STATUS_RESCAN_NEEDED) {
    status=sh_slab_radix_node_scan_slabs(alloc);
    if (sh_status_error(status)) return status;
    status=sh_slab_radix_node_get_partial_slab(alloc,ptp,&partial_slab);
  }
  if (status==SH_STATUS_NOT_FOUND) {
    status=sh_slab_radix_node_add_slab(alloc,ptp,&new_slab);
    if (sh_status_error(status)) return status;
    status=sh_slab_radix_node_get_partial_slab(alloc,ptp,&partial_slab);
    if (sh_status_error(status)) return status;
  } else if (sh_status_error(status) && status!=SH_STATUS_NEW_SLAB_ADDED) {
    return status;
  }
  sh_uint16 found_idx=0;
  sh_bool bit_found=SH_FALSE;
  for (sh_uint8 w=0;w<16;w++) {
    if (partial_slab->free_bitmap[w]!=0xFFFFFFFFFFFFFFFFULL) {
      for (sh_uint8 b=0;b<64;b++) {
        if (!(partial_slab->free_bitmap[w] & (1ULL<<b))) {
          found_idx=(w<<6)|b;
          if (found_idx<18) continue;
          bit_found=SH_TRUE;
          break;
        }
      }
    }
    if (bit_found) break;
  }
  if (!bit_found) return SH_STATUS_NOT_FOUND;
  *out=&partial_slab->nodes[found_idx-18];
  *index_in_slab=found_idx;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_alloc(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_radix_NODE** out_obj) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out_obj==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_radix_NODE *object_ptr;
  sh_slab_radix_node_NODE_INDEX_IN_SLAB index_in_slab;
  SH_STATUS status=sh_slab_radix_node_find_free_object(alloc,ptp,&object_ptr,&index_in_slab);
  if (sh_status_error(status)) return status;
  if (object_ptr==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  for (sh_iter64 i=0;i<16;i++) object_ptr->ptr[i]=0x0;
  sh_slab_radix_node_SLAB* slab=(sh_slab_radix_node_SLAB*)((sh_uint64)object_ptr & ~0x1FFFFULL);
  if (*(sh_uint64*)slab->sig!=SH_SLAB_RADIX_NODE_MAGIC) {
    return SH_STATUS_FOUND_CORRUPTED_SLAB;
  }
  sh_uint16 word_idx=index_in_slab>>6;
  sh_uint16 bit_idx=index_in_slab&63;
  slab->free_bitmap[word_idx]|=(1ULL<<bit_idx);
  slab->node_bitmap[index_in_slab]=0;
  slab->used_count++;
  *out_obj=object_ptr;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_radix_node_dealloc(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_radix_NODE *object_ptr) {
  if (alloc==SH_NULLPTR || object_ptr==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if ((sh_uint64)object_ptr%SH_SLAB_RADIX_NODE_OBJECT_SIZE_BYTES!=0) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_radix_node_SLAB* slab=(sh_slab_radix_node_SLAB*)((sh_uint64)object_ptr & ~0x1FFFFULL);
  if (*(sh_uint64*)slab->sig!=SH_SLAB_RADIX_NODE_MAGIC) {
    return SH_STATUS_FOUND_CORRUPTED_SLAB;
  }
  sh_uint16 object_idx=(sh_uint16)(((sh_uint64)object_ptr & 0x1FFFFULL)/SH_SLAB_RADIX_NODE_OBJECT_SIZE_BYTES);
  if (object_idx<18) return SH_STATUS_INVALID_PARAMETER;
  sh_uint16 word=object_idx>>6;
  sh_uint16 bit=object_idx&63;
  if (!(slab->free_bitmap[word] & (1ULL<<bit))) return SH_STATUS_SLAB_SLOT_ALREADY_FREE;
  slab->free_bitmap[word]&= ~(1ULL<<bit);
  if (slab->used_count==SH_SLAB_RADIX_NODE_ACTUAL_OBJECTS_PER_SLAB) {
    slab->prev_partial=SH_NULLPTR;
    slab->next_partial=alloc->partial_head;
    if (alloc->partial_head!=SH_NULLPTR) {
      alloc->partial_head->prev_partial=slab;
    }
    alloc->partial_head=slab;
  }
  slab->used_count--;
  for (sh_iter64 i=0;i<16;i++) object_ptr->ptr[i]=0x0;
  slab->node_bitmap[object_idx]=0;
  return SH_STATUS_SUCCESS;
}
sh_uint16 *sh_slab_radix_node_get_node_bitmap(struct sh_slab_radix_node_SLAB_ALLOCATOR* alloc,sh_radix_NODE *object_ptr) {
  if (alloc==SH_NULLPTR || object_ptr==SH_NULLPTR) return SH_NULLPTR;
  if ((sh_uint64)object_ptr%SH_SLAB_RADIX_NODE_OBJECT_SIZE_BYTES!=0) return SH_NULLPTR;
  sh_slab_radix_node_SLAB* slab=(sh_slab_radix_node_SLAB*)((sh_uint64)object_ptr & ~0x1FFFFULL);
  if (*(sh_uint64*)slab->sig!=SH_SLAB_RADIX_NODE_MAGIC) {
    return SH_NULLPTR;
  }
  sh_uint16 object_idx=(sh_uint16)(((sh_uint64)object_ptr & 0x1FFFFULL)/SH_SLAB_RADIX_NODE_OBJECT_SIZE_BYTES);
  if (object_idx<18) return SH_NULLPTR;
  return &slab->node_bitmap[object_idx];
}
