// SPDX-License-Identifier: MPL-2.0
#include "memory/pez/radix.h"
#include "memory/slabs/slab_radix_node.h"
SH_STATUS sh_radix_node_read_value(sh_radix_NODE *node,sh_uint8 index,sh_page_VIRTUAL_ADDRESS *value) {
  if (node==SH_NULLPTR || index>15 || value==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_page_VIRTUAL_ADDRESS val=node->ptr[index];
  if (val==0) return SH_STATUS_NOT_FOUND;
  *value=val;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_radix_node_set_value(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_radix_NODE *node,sh_uint8 index,sh_page_VIRTUAL_ADDRESS value) {
  if (node==SH_NULLPTR || alloc==SH_NULLPTR|| index>15) return SH_STATUS_INVALID_PARAMETER;
  sh_uint16 *bitmap_ptr=sh_slab_radix_node_get_node_bitmap(alloc,node);
  if (bitmap_ptr==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  if (value!=0) {
    *bitmap_ptr|=(1<<index);
  } else {
    *bitmap_ptr&= (sh_uint16)(~(1<<index));
  }
  node->ptr[index]=value;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_radix_tree_init(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_radix_TREE *tree, sh_uint8 depth) {
    if (tree==SH_NULLPTR || alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
    if (depth>16) return SH_STATUS_INVALID_PARAMETER;
    tree->depth=depth;
    sh_radix_NODE *root_node;
    SH_STATUS status=sh_slab_radix_node_alloc(alloc,ptp,&root_node);
    if (sh_status_error(status)) return status;
    for (sh_uint8 i=0;i<16;i++) root_node->ptr[i]=0;
    tree->root_node=root_node;
    return SH_STATUS_SUCCESS;
}
SH_STATUS sh_radix_tree_get_value(sh_radix_TREE *tree,sh_uint64 key,sh_page_VIRTUAL_ADDRESS *value) {
  if (tree==SH_NULLPTR || tree->root_node==SH_NULLPTR || value==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_radix_NODE *current_node=tree->root_node;
  sh_page_VIRTUAL_ADDRESS current_val=0;
  SH_STATUS status;
  for (sh_uint8 level=0;level<tree->depth;level++) {
    sh_uint8 shift=(sh_uint8)((tree->depth-1-level)*4);
    sh_uint8 index=(sh_uint8)((key>>shift)&0x0F);
    status=sh_radix_node_read_value(current_node,index,&current_val);
    if (sh_status_error(status)) {
      return status;
    }
    if (level<(tree->depth-1)) {
      current_node=(sh_radix_NODE *)current_val;
    }
  }
  *value=current_val;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_radix_tree_insert_value(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_radix_TREE *tree,sh_uint64 key,sh_page_VIRTUAL_ADDRESS value) {
  if (tree==SH_NULLPTR || tree->root_node==SH_NULLPTR || alloc==SH_NULLPTR || ptp==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_radix_NODE *current_node=tree->root_node;
  SH_STATUS status;
  for (sh_uint8 level=0;level<tree->depth;level++) {
    sh_uint8 shift=(sh_uint8)((tree->depth-1-level)*4);
    sh_uint8 index=(sh_uint8)((key>>shift)&0x0F);
    if (level==(tree->depth-1)) {
      return sh_radix_node_set_value(alloc,current_node,index,value);
    }
    sh_page_VIRTUAL_ADDRESS next_node_va;
    status=sh_radix_node_read_value(current_node,index,&next_node_va);
    if (status==SH_STATUS_NOT_FOUND) {
      sh_radix_NODE *new_node;
      status=sh_slab_radix_node_alloc(alloc,ptp,&new_node);
      if (sh_status_error(status)) return status;
      sh_uint8 i=0;
      while(i<16) {
        new_node->ptr[i]=0;
        i++;
      }
      status=sh_radix_node_set_value(alloc,current_node,index,(sh_page_VIRTUAL_ADDRESS)new_node);
      if (sh_status_error(status)) return status;
      current_node=new_node;
    } else if (status==SH_STATUS_SUCCESS) {
      current_node=(sh_radix_NODE *)next_node_va;
    } else {
      return status;
    }
  }
  return SH_STATUS_SUCCESS;
}
static SH_STATUS sh_radix_tree_delete_recursive(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_radix_NODE *current_node,sh_uint64 key,sh_uint8 depth,sh_uint8 current_level,sh_bool *is_empty_out) {
  sh_uint8 shift=(sh_uint8)((depth-1-current_level)*4);
  sh_uint8 index=(sh_uint8)((key>>shift)&0x0F);
  SH_STATUS status;
  if (current_level==(depth-1)) {
    status=sh_radix_node_set_value(alloc,current_node,index,0);
    if (sh_status_error(status)) return status;
  } else {
    sh_page_VIRTUAL_ADDRESS next_node_va=current_node->ptr[index];
    if (next_node_va==0) return SH_STATUS_NOT_FOUND;
    sh_radix_NODE *next_node=(sh_radix_NODE *)next_node_va;
    sh_bool child_is_empty=SH_FALSE;
    status=sh_radix_tree_delete_recursive(alloc,next_node,key,depth,current_level+1,&child_is_empty);
    if (sh_status_error(status)) return status;
    if (child_is_empty) {
      status=sh_slab_radix_node_dealloc(alloc,next_node);
      if (sh_status_error(status)) return status;
      status=sh_radix_node_set_value(alloc,current_node,index,0);
      if (sh_status_error(status)) return status;
    }
  }
  sh_bool empty=SH_TRUE;
  for (sh_uint8 i=0;i<16;i++) {
    if (current_node->ptr[i]!=0) {
      empty=SH_FALSE;
      break;
    }
  }
  if (is_empty_out) *is_empty_out=empty;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_radix_tree_delete_value(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_radix_TREE *tree,sh_uint64 key) {
  if (tree==SH_NULLPTR || tree->root_node==SH_NULLPTR || alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_bool root_is_empty=SH_FALSE;
  return sh_radix_tree_delete_recursive(alloc,tree->root_node,key,tree->depth,0,&root_is_empty);
}
SH_STATUS sh_radix_tree_search_smallest_min_bound(struct sh_slab_radix_node_SLAB_ALLOCATOR *alloc,sh_radix_TREE *tree,sh_uint64 lower_bound_key,sh_page_VIRTUAL_ADDRESS *value){
  if (tree==SH_NULLPTR || tree->root_node==SH_NULLPTR || alloc==SH_NULLPTR || value==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_uint8 depth=tree->depth;
  sh_radix_NODE *node_stack[16];
  sh_uint8 idx_stack[16];
  sh_uint8 level=0;
  node_stack[0]=tree->root_node;
  while (level<depth) {
    sh_uint8 shift=(sh_uint8)((depth-1-level)*4);
    sh_uint8 index=(sh_uint8)((lower_bound_key>>shift)&0x0F);
    idx_stack[level]=index;
    if (level==depth-1) {
      sh_radix_NODE *node=node_stack[level];
      sh_uint16 *bmp=sh_slab_radix_node_get_node_bitmap(alloc,node);
      sh_uint16 mask=(sh_uint16)(0xFFFF<<index) & *bmp;
      if (mask!=0) {
        sh_uint8 best=(sh_uint8)__builtin_ctz(mask);
        *value=node->ptr[best];
        return SH_STATUS_SUCCESS;
      }
      goto backtrack;
    }
    sh_radix_NODE *node=node_stack[level];
    sh_page_VIRTUAL_ADDRESS child_va=node->ptr[index];
    if (child_va!=0) {
      node_stack[level+1]=(sh_radix_NODE *)child_va;
      level++;
      continue;
    }
    goto try_next_sibling;
  try_next_sibling:;
    sh_uint16 *bmp=sh_slab_radix_node_get_node_bitmap(alloc,node);
    sh_uint8 next_index=index+1;
    if (next_index>15) goto backtrack;
    sh_uint16 mask=(sh_uint16)(0xFFFF<<next_index) & *bmp;
    if (mask==0) goto backtrack;
    sh_uint8 best=(sh_uint8)__builtin_ctz(mask);
    node_stack[level+1]=(sh_radix_NODE *)node->ptr[best];
    idx_stack[level]=best;
    level++;
    while (level<depth) {
      sh_radix_NODE *cur=node_stack[level];
      sh_uint16 *b=sh_slab_radix_node_get_node_bitmap(alloc,cur);
      if (*b==0) goto backtrack;
      sh_uint8 min_idx=(sh_uint8)__builtin_ctz(*b);
      idx_stack[level]=min_idx;
      if (level==depth-1) {
        *value=cur->ptr[min_idx];
        return SH_STATUS_SUCCESS;
      }
      node_stack[level+1]=(sh_radix_NODE *)cur->ptr[min_idx];
      level++;
    }
  backtrack:
    if (level==0) return SH_STATUS_NOT_FOUND;
    level--;
    node=node_stack[level];
    index=idx_stack[level];
    goto try_next_sibling;
  }
  return SH_STATUS_NOT_FOUND;
}