Weak AVL Tree Implementation 5

/* -*- Mode: C; indent-tabs-mode: nil; c-basic-offset: 2; tab-width: 2 -*- */
/*
  c-tree-map.c
  This file is part of Clair.

  Copyright (C) 2024 Hodong Kim <hodong@nimfsoft.art>

  Permission to use, copy, modify, and/or distribute this software for any
  purpose with or without fee is hereby granted.

  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
*/
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include <string.h>

void* c_malloc (size_t size)
{
  if (!size)
    return nullptr;

  void* mem = malloc (size);

  if (mem)
    return mem;

  perror (__PRETTY_FUNCTION__);
  abort ();
}

#define VOIDP_TO_INT(p)  ((int)   (intptr_t) (p))
#define INT_TO_VOIDP(i)  ((void*) (intptr_t) (i))
#define C_MAX(a, b)      (((a) > (b)) ? (a) : (b))
#define _RANK(node)      ((node) ? (node)->rank : -1)

typedef void (*CFreeFunc)            (void* data);
typedef int  (*CCompareFunc)         (const void* a, const void* b);
typedef void (*CTreeMapForEachFunc)  (void* k, void* v, void* user_data);

enum _CTreeMapOrder {
  C_TREE_MAP_PRE_ORDER,
  C_TREE_MAP_IN_ORDER,
  C_TREE_MAP_POST_ORDER
};
typedef enum _CTreeMapOrder  CTreeMapOrder;

typedef struct _CTreeNode  CTreeNode;
struct _CTreeNode {
  struct _CTreeNode* left;
  struct _CTreeNode* right;
  void* key;
  void* value;
  int rank;
};

typedef struct _CTreeMap  CTreeMap;
struct _CTreeMap {
  CTreeNode* root;
  CCompareFunc key_compare_func;
  CFreeFunc key_free_func;
  CFreeFunc value_free_func;
  int size;
};

static void _update_rank (CTreeNode* h)
{
  h->rank = C_MAX (_RANK (h->left), _RANK (h->right)) + 1;
}

CTreeNode* _rotate_left (CTreeNode* h)
{
  assert (h->left != nullptr);

  CTreeNode* x = h->left;
  h->left = x->right;
  x->right = h;

  _update_rank (h);
  _update_rank (x);

  return x;
}

CTreeNode* _rotate_right (CTreeNode* h)
{
  assert (h->right != nullptr);

  CTreeNode* x = h->right;
  h->right = x->left;
  x->left = h;

  _update_rank (h);
  _update_rank (x);

  return x;
}

bool c_tree_map_lookup (CTreeMap* map, const void* key, void** value)
{
  CTreeNode* node = map->root;

  while (node)
  {
    int retval = map->key_compare_func (key, node->key);
    if (retval < 0)
    {
      node = node->left;
    }
    else if (retval > 0)
    {
      node = node->right;
    }
    else
    {
      if (value)
        *value = node->value;

      return true;
    }
  }

  return false;
}

static void _foreach (CTreeNode* node,
                      CTreeMapOrder order,
                      CTreeMapForEachFunc func,
                      void* user_data)
{
  if (!node)
    return;

  if (order == C_TREE_MAP_PRE_ORDER)
    func (node->key, node->value, user_data);

  _foreach (node->left, order, func, user_data);

  if (order == C_TREE_MAP_IN_ORDER)
    func (node->key, node->value, user_data);

  _foreach (node->right, order, func, user_data);

  if (order == C_TREE_MAP_POST_ORDER)
    func (node->key, node->value, user_data);
}

void c_tree_map_foreach (CTreeMap* map,
                         CTreeMapOrder order,
                         CTreeMapForEachFunc func,
                         void* user_data)
{
  _foreach (map->root, order, func, user_data);
}

static CTreeNode* _balance (CTreeNode* h)
{
  if ((_RANK(h->left) - _RANK(h->right) > 1))
  {
    if (_RANK(h->left->right) > _RANK(h->left->left))
      h->left = _rotate_right (h->left);

    h = _rotate_left (h);
  }
  else
  {
    if ((_RANK(h->right) - _RANK(h->left) > 1))
    {
      if (_RANK(h->right->left) > _RANK(h->right->right))
        h->right = _rotate_left (h->right);

      h = _rotate_right (h);
    }
  }

  _update_rank (h);
  return h;
}

static CTreeNode* _node_new (void* key, void* value)
{
  CTreeNode* node = c_malloc (sizeof (CTreeNode));
  node->left  = nullptr;
  node->right = nullptr;
  node->key   = key;
  node->value = value;
  node->rank  = 0;
  return node;
}

static void _free_node (CTreeMap* map, CTreeNode* node)
{
  if (!node)
    return;

  _free_node (map, node->left);
  _free_node (map, node->right);

  if (map->key_free_func)
    map->key_free_func (node->key);

  if (map->value_free_func)
    map->value_free_func (node->value);

  free (node);
}

static CTreeNode* _insert (CTreeMap* map, CTreeNode* h, void* key, void* value)
{
  if (h == nullptr)
  {
    map->size++;
    return _node_new (key, value);
  }

  int retval = map->key_compare_func (key, h->key);
  if (retval < 0)
  {
    h->left = _insert (map, h->left, key, value);
  }
  else if (retval > 0)
  {
    h->right = _insert (map, h->right, key, value);
  }
  else // replace
  {
    if (map->key_free_func)
      map->key_free_func (h->key);

    if (map->value_free_func)
      map->value_free_func (h->value);

    h->key   = key;
    h->value = value;
    return h;
  }

  return _balance (h);
}

CTreeMap* c_tree_map_new (CCompareFunc key_compare_func,
                          CFreeFunc key_free_func,
                          CFreeFunc value_free_func)
{
  CTreeMap* map         = c_malloc (sizeof (CTreeMap));
  map->root             = nullptr;
  map->key_compare_func = key_compare_func;
  map->key_free_func    = key_free_func;
  map->value_free_func  = value_free_func;
  map->size             = 0;
  return map;
}

void c_tree_map_insert (CTreeMap* map, void* key, void* value)
{
  map->root = _insert (map, map->root, key, value);
}

int c_tree_map_size (CTreeMap* map)
{
  return map->size;
}

CTreeNode* _real_remove (CTreeMap* map, CTreeNode* h)
{
  if (h->left && h->right)
  {
    if (h->left->rank > h->right->rank)
    {
      h = _rotate_left (h);
      h->right = _real_remove (map, h->right);
    }
    else
    {
      h = _rotate_right (h);
      h->left = _real_remove (map, h->left);
    }
  }
  else
  {
    CTreeNode* _next = nullptr;

    if (h->left)
      _next = h->left;
    else
      _next = h->right;

    if (map->key_free_func)
      map->key_free_func (h->key);

    if (map->value_free_func)
      map->value_free_func (h->value);

    free (h);
    return _next;
  }

  return _balance (h);
}

CTreeNode* _remove (CTreeMap* map, CTreeNode* h, void* key)
{
  if (!h)
    return nullptr;

  int retval = map->key_compare_func (key, h->key);
  if (retval < 0)
  {
    h->left = _remove (map, h->left, key);
  }
  else if (retval > 0)
  {
    h->right = _remove (map, h->right, key);
  }
  else
  {
    map->size--;
    return _real_remove (map, h);
  }

  return _balance (h);
}

void c_tree_map_remove (CTreeMap* map, void* key)
{
  map->root = _remove (map, map->root, key);
}

void c_tree_map_free (CTreeMap* map)
{
  _free_node (map, map->root);
  free (map);
}

static void verify (CTreeNode* node, int* count)
{
  if (!node)
    return;

  verify (node->left, count);

  if (node->left)
    printf ("%d(%d)\t", VOIDP_TO_INT (node->left->key), node->left->rank);

  else
    printf ("\t");

  printf ("%d(%d):%s  ", VOIDP_TO_INT (node->key), node->rank,
                         (char*) node->value);
  if (node->right)
    printf ("\t%d(%d)", VOIDP_TO_INT (node->right->key), node->right->rank);
  else
    printf ("\t");

  puts ("");

  // All leaf nodes have rank 0.
  if (!node->left && !node->right)
    assert (node->rank == 0);

  if (node->left && node->right && // An internal node
      !node->left->left && !node->left->right &&
      !node->right->left && !node->right->right) //  with two leaf nodes
    assert (node->rank == 1);  // has rank 1.

  // The rank difference from the parent node is 1 or 2.
  int diff;
  diff = abs (_RANK(node) - _RANK(node->left));
  assert (diff == 1 || diff == 2);
  diff = abs (_RANK(node) - _RANK(node->right));
  assert (diff == 1 || diff == 2);

  (*count)++;

  verify (node->right, count);
}

static void tree_verify (CTreeMap* tree, int size)
{
  int count = 0;
  verify (tree->root, &count);
  assert (tree->size == size);
  assert (tree->size == count);
}

static int compare (const void* a, const void* b)
{
  return a - b;
}

static void print_kv (void* key, void* value, void* user_data)
{
  CTreeMap* map = user_data;
  printf ("%d %s\n", VOIDP_TO_INT (key), (char*) value);
  c_tree_map_remove (map, key);
}

int main ()
{
  int numbers[] = {
    29, 13, 23, 59, 25, 20,  5,  9, 33,  2, 31, 18, 11, 37,  7, 98,
    31,  2, 13, 25,  9, 18, 37, 29, 23, 98, 20, 59,  5,  7, 11, 33
  };

  char* strings[] = {
    "s29", "s13", "s23", "s59", "s25", "s20", "s5",  "s9",
    "s33", "s2",  "s31", "s18", "s11", "s37", "s7",  "s98",
    "s31", "s2",  "s13", "s25", "s9",  "s18", "s37", "s29",
    "s23", "s98", "s20", "s59", "s5",  "s7",  "s11", "s33"
  };

  CTreeMap* map = c_tree_map_new (compare, nullptr, nullptr);
  int numbers_len = sizeof (numbers) / sizeof (numbers[0]);
  int size = 0;

  for (int i = 0; i < numbers_len; i++)
  {
    void* key = INT_TO_VOIDP (numbers[i]);
    c_tree_map_insert (map, key, strings[i]);
    char* s;
    assert (c_tree_map_lookup (map, key, (void**) &s) == true);
    assert (c_tree_map_lookup (map, key + 1000, (void**) &s) == false);
    assert (strcmp (s, strings[i]) == 0);
  }

  size = numbers_len / 2;

  void* v = (void*) 111;
  c_tree_map_insert (map, INT_TO_VOIDP (1000), INT_TO_VOIDP (0));
  size++;
  assert (c_tree_map_lookup (map, INT_TO_VOIDP (1000), (void**) &v) == true);
  assert (v == 0);

  tree_verify (map, size);
  c_tree_map_remove (map, INT_TO_VOIDP (0));
  c_tree_map_remove (map, INT_TO_VOIDP (1000));
  size--;

  tree_verify (map, size);

  c_tree_map_foreach (map, C_TREE_MAP_POST_ORDER, print_kv, map);

  c_tree_map_free (map);

  return 0;
}