/*
 * Monolithium Unit Tests
 * avltree.c
 *
 * Copyright (C) 2018 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <check.h>
#include <stdlib.h>
#include <sdk/avltree.h>

#define SIZE 10000

typedef int key_t;
typedef int value_t;

typedef struct
{
    avl_node_t node;
    key_t key;
    value_t value;
} map_entry_t;

avl_tree_t tree;

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

static int height(avl_node_t *root)
{
    if (!root) return 0;
    int left = height(root->left);
    int right = height(root->right);
    return 1 + (left > right ? left : right);
}

static const void *minimum(avl_node_t *root)
{
    const void *result = avl_get_keyptr(&tree, root);

    if (root->left)
    {
        const void *left = minimum(root->left);
        if (compare(left, result) < 0) result = left;
    }

    if (root->right)
    {
        const void *right = minimum(root->right);
        if (compare(right, result) < 0) result = right;
    }

    return result;
}

static const void *maximum(avl_node_t *root)
{
    const void *result = avl_get_keyptr(&tree, root);

    if (root->left)
    {
        const void *left = maximum(root->left);
        if (compare(left, result) > 0) result = left;
    }

    if (root->right)
    {
        const void *right = maximum(root->right);
        if (compare(right, result) > 0) result = right;
    }

    return result;
}

static void test_integrity(avl_node_t *root, avl_node_t *parent)
{
    if (!root) return;
    ck_assert(root->parent == parent);
    test_integrity(root->left, root);
    test_integrity(root->right, root);
}

static void test_order(avl_node_t *root)
{
    if (!root) return;
    const void *key = avl_get_keyptr(&tree, root);
    test_order(root->left);
    if (root->left) ck_assert(compare(maximum(root->left), key) < 0);
    if (root->right) ck_assert(compare(minimum(root->right), key) > 0);
    test_order(root->right);
}

static void test_balance(avl_node_t *root)
{
    if (!root) return;
    test_balance(root->left);

    int balance = height(root->right) - height(root->left);
    ck_assert(root->balance == balance);
    ck_assert(balance >= -1 && balance <= 1);

    test_balance(root->right);
}

static map_entry_t *map_lookup(key_t key)
{
    avl_node_t *node = avl_tree_lookup(&tree, &key);
    return node ? CONTAINER_OF(node, map_entry_t, node) : NULL;
}

static void map_insert(key_t key, value_t value)
{
    map_entry_t *entry = malloc(sizeof(map_entry_t));
    entry->key = key;
    entry->value = value;
    avl_tree_insert(&tree, &entry->node);
}

static void setup(void)
{
    AVL_TREE_INIT(&tree, map_entry_t, node, key, compare);
}

static void teardown(void)
{
}

START_TEST(test_avl_tree_insert_ascending)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, i * 2);

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL && entry->value == i * 2);
    }
}
END_TEST

START_TEST(test_avl_tree_insert_descending)
{
    int i;
    for (i = SIZE - 1; i >= 0; i--) map_insert(i, i * 2);

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL && entry->value == i * 2);
    }
}
END_TEST

START_TEST(test_avl_tree_insert_all_equal)
{
    int i;
    for (i = SIZE - 1; i >= 0; i--) map_insert(i, 42);

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL && entry->value == 42);
    }
}
END_TEST

START_TEST(test_avl_tree_insert_random_order)
{
    int i;
    int list[SIZE];
    for (i = 0; i < SIZE; i++) list[i] = i;

    for (i = 0; i < SIZE; i++)
    {
        int j = i + (rand() % (SIZE - i));
        int t = list[i];
        list[i] = list[j];
        list[j] = t;
    }

    for (i = 0; i < SIZE; i++) map_insert(list[i], list[i] * 2);

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL);
        ck_assert(entry->value == i * 2);
    }
}
END_TEST

START_TEST(test_avl_tree_remove_half_ascending)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, i * 2);

    for (i = 0; i < SIZE / 2; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL);
        avl_tree_remove(&tree, &entry->node);
        entry = map_lookup(i);
        ck_assert(entry == NULL);
    }

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);
}
END_TEST

START_TEST(test_avl_tree_remove_random_half)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, i * 2);

    for (i = 0; i < SIZE / 2; i++)
    {
        key_t key;
        map_entry_t *entry = NULL;

        while (!entry)
        {
            key = rand() % SIZE;
            entry = map_lookup(key);
        }

        avl_tree_remove(&tree, &entry->node);
        entry = map_lookup(key);
        ck_assert(entry == NULL);
    }

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);
}
END_TEST

START_TEST(test_avl_tree_remove_half_with_duplicates)
{
    int i;
    int counts[10] = {0};
    for (i = 0; i < SIZE; i++)
    {
        map_insert(i % 10, (i % 10) * 2);
        counts[i % 10]++;
    }

    for (i = 0; i < SIZE / 2; i++)
    {
        map_entry_t *entry = CONTAINER_OF(tree.root, map_entry_t, node);
        key_t key = entry->key;
        ck_assert(key >= 0 && key < 10);

        avl_tree_remove(&tree, &entry->node);
        counts[key]--;
        ck_assert(counts[key] >= 0);

        if (counts[key] == 0)
        {
            entry = map_lookup(key);
            ck_assert(entry == NULL);
        }
    }

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);
}
END_TEST

START_TEST(test_avl_tree_misc_traversal_forward)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, i * 2);

    avl_node_t *node = tree.root;
    while (node->left) node = node->left;

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = CONTAINER_OF(node, map_entry_t, node);
        ck_assert(entry->key == i && entry->value == i * 2);
        node = avl_get_next_node(node);
    }
}
END_TEST

START_TEST(test_avl_tree_misc_traversal_backward)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, i * 2);

    avl_node_t *node = tree.root;
    while (node->right) node = node->right;

    for (i = SIZE - 1; i >= 0; i--)
    {
        map_entry_t *entry = CONTAINER_OF(node, map_entry_t, node);
        ck_assert(entry->key == i && entry->value == i * 2);
        node = avl_get_previous_node(node);
    }
}
END_TEST

START_TEST(test_avl_tree_misc_key_changing)
{
    int i;
    for (i = 0; i < SIZE; i++) map_insert(i, (SIZE - i - 1) * 2);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL);
        key_t new_key = SIZE - i - 1;
        avl_tree_change_key(&tree, &entry->node, &new_key);
    }

    test_integrity(tree.root, NULL);
    test_order(tree.root);
    test_balance(tree.root);

    for (i = 0; i < SIZE; i++)
    {
        map_entry_t *entry = map_lookup(i);
        ck_assert(entry != NULL);
        ck_assert(entry->value == i * 2);
    }
}
END_TEST

TCase *tcase_avl_tree_insertion_only_create(void)
{
    TCase *tc = tcase_create("Insertion Only");
    tcase_add_checked_fixture(tc, setup, teardown);
    tcase_add_test(tc, test_avl_tree_insert_ascending);
    tcase_add_test(tc, test_avl_tree_insert_descending);
    tcase_add_test(tc, test_avl_tree_insert_all_equal);
    tcase_add_test(tc, test_avl_tree_insert_random_order);
    return tc;
}

TCase *tcase_avl_tree_removal_create(void)
{
    TCase *tc = tcase_create("Removal");
    tcase_add_checked_fixture(tc, setup, teardown);
    tcase_add_test(tc, test_avl_tree_remove_half_ascending);
    tcase_add_test(tc, test_avl_tree_remove_random_half);
    tcase_add_test(tc, test_avl_tree_remove_half_with_duplicates);
    return tc;
}

TCase *tcase_avl_tree_misc_create(void)
{
    TCase *tc = tcase_create("Miscellaneous");
    tcase_add_checked_fixture(tc, setup, teardown);
    tcase_add_test(tc, test_avl_tree_misc_traversal_forward);
    tcase_add_test(tc, test_avl_tree_misc_traversal_backward);
    tcase_add_test(tc, test_avl_tree_misc_key_changing);
    return tc;
}

Suite *suite_avl_tree_create(void)
{
    Suite *suite = suite_create("AVL Tree");
    suite_add_tcase(suite, tcase_avl_tree_insertion_only_create());
    suite_add_tcase(suite, tcase_avl_tree_removal_create());
    suite_add_tcase(suite, tcase_avl_tree_misc_create());
    return suite;
}

