#include "test/jemalloc_test.h"

#include <stdlib.h>

#include "jemalloc/internal/rb.h"

#define rbtn_black_height(a_type, a_field, a_rbt, r_height) do {        \
       a_type *rbp_bh_t;                                               \
       for (rbp_bh_t = (a_rbt)->rbt_root, (r_height) = 0; rbp_bh_t !=  \
           NULL; rbp_bh_t = rbtn_left_get(a_type, a_field,             \
           rbp_bh_t)) {                                                \
               if (!rbtn_red_get(a_type, a_field, rbp_bh_t)) {         \
               (r_height)++;                                           \
               }                                                       \
       }                                                               \
} while (0)

static bool summarize_always_returns_true = false;

typedef struct node_s node_t;
struct node_s {
#define NODE_MAGIC 0x9823af7e
       uint32_t magic;
       rb_node(node_t) link;
       /* Order used by nodes. */
       uint64_t key;
       /*
        * Our made-up summary property is "specialness", with summarization
        * taking the max.
        */
       uint64_t specialness;

       /*
        * Used by some of the test randomization to avoid double-removing
        * nodes.
        */
       bool mid_remove;

       /*
        * To test searching functionality, we want to temporarily weaken the
        * ordering to allow non-equal nodes that nevertheless compare equal.
        */
       bool allow_duplicates;

       /*
        * In check_consistency, it's handy to know a node's rank in the tree;
        * this tracks it (but only there; not all tests use this).
        */
       int rank;
       int filtered_rank;

       /*
        * Replicate the internal structure of the tree, to make sure the
        * implementation doesn't miss any updates.
        */
       const node_t *summary_lchild;
       const node_t *summary_rchild;
       uint64_t summary_max_specialness;
};

static int
node_cmp(const node_t *a, const node_t *b) {
       int ret;

       expect_u32_eq(a->magic, NODE_MAGIC, "Bad magic");
       expect_u32_eq(b->magic, NODE_MAGIC, "Bad magic");

       ret = (a->key > b->key) - (a->key < b->key);
       if (ret == 0 && !a->allow_duplicates) {
               /*
                * Duplicates are not allowed in the tree, so force an
                * arbitrary ordering for non-identical items with equal keys,
                * unless the user is searching and wants to allow the
                * duplicate.
                */
               ret = (((uintptr_t)a) > ((uintptr_t)b))
                   - (((uintptr_t)a) < ((uintptr_t)b));
       }
       return ret;
}

static uint64_t
node_subtree_specialness(node_t *n, const node_t *lchild,
   const node_t *rchild) {
       uint64_t subtree_specialness = n->specialness;
       if (lchild != NULL
           && lchild->summary_max_specialness > subtree_specialness) {
               subtree_specialness = lchild->summary_max_specialness;
       }
       if (rchild != NULL
           && rchild->summary_max_specialness > subtree_specialness) {
               subtree_specialness = rchild->summary_max_specialness;
       }
       return subtree_specialness;
}

static bool
node_summarize(node_t *a, const node_t *lchild, const node_t *rchild) {
       uint64_t new_summary_max_specialness = node_subtree_specialness(
           a, lchild, rchild);
       bool changed = (a->summary_lchild != lchild)
           || (a->summary_rchild != rchild)
           || (new_summary_max_specialness != a->summary_max_specialness);
       a->summary_max_specialness = new_summary_max_specialness;
       a->summary_lchild = lchild;
       a->summary_rchild = rchild;
       return changed || summarize_always_returns_true;
}

typedef rb_tree(node_t) tree_t;
rb_summarized_proto(static, tree_, tree_t, node_t);
rb_summarized_gen(static, tree_, tree_t, node_t, link, node_cmp,
   node_summarize);

static bool
specialness_filter_node(void *ctx, node_t *node) {
       uint64_t specialness = *(uint64_t *)ctx;
       return node->specialness >= specialness;
}

static bool
specialness_filter_subtree(void *ctx, node_t *node) {
       uint64_t specialness = *(uint64_t *)ctx;
       return node->summary_max_specialness >= specialness;
}

static node_t *
tree_iterate_cb(tree_t *tree, node_t *node, void *data) {
       unsigned *i = (unsigned *)data;
       node_t *search_node;

       expect_u32_eq(node->magic, NODE_MAGIC, "Bad magic");

       /* Test rb_search(). */
       search_node = tree_search(tree, node);
       expect_ptr_eq(search_node, node,
           "tree_search() returned unexpected node");

       /* Test rb_nsearch(). */
       search_node = tree_nsearch(tree, node);
       expect_ptr_eq(search_node, node,
           "tree_nsearch() returned unexpected node");

       /* Test rb_psearch(). */
       search_node = tree_psearch(tree, node);
       expect_ptr_eq(search_node, node,
           "tree_psearch() returned unexpected node");

       (*i)++;

       return NULL;
}

TEST_BEGIN(test_rb_empty) {
       tree_t tree;
       node_t key;

       tree_new(&tree);

       expect_true(tree_empty(&tree), "Tree should be empty");
       expect_ptr_null(tree_first(&tree), "Unexpected node");
       expect_ptr_null(tree_last(&tree), "Unexpected node");

       key.key = 0;
       key.magic = NODE_MAGIC;
       expect_ptr_null(tree_search(&tree, &key), "Unexpected node");

       key.key = 0;
       key.magic = NODE_MAGIC;
       expect_ptr_null(tree_nsearch(&tree, &key), "Unexpected node");

       key.key = 0;
       key.magic = NODE_MAGIC;
       expect_ptr_null(tree_psearch(&tree, &key), "Unexpected node");

       unsigned nodes = 0;
       tree_iter_filtered(&tree, NULL, &tree_iterate_cb,
           &nodes, &specialness_filter_node, &specialness_filter_subtree,
           NULL);
       expect_u_eq(0, nodes, "");

       nodes = 0;
       tree_reverse_iter_filtered(&tree, NULL, &tree_iterate_cb,
           &nodes, &specialness_filter_node, &specialness_filter_subtree,
           NULL);
       expect_u_eq(0, nodes, "");

       expect_ptr_null(tree_first_filtered(&tree, &specialness_filter_node,
           &specialness_filter_subtree, NULL), "");
       expect_ptr_null(tree_last_filtered(&tree, &specialness_filter_node,
           &specialness_filter_subtree, NULL), "");

       key.key = 0;
       key.magic = NODE_MAGIC;
       expect_ptr_null(tree_search_filtered(&tree, &key,
           &specialness_filter_node, &specialness_filter_subtree, NULL), "");
       expect_ptr_null(tree_nsearch_filtered(&tree, &key,
           &specialness_filter_node, &specialness_filter_subtree, NULL), "");
       expect_ptr_null(tree_psearch_filtered(&tree, &key,
           &specialness_filter_node, &specialness_filter_subtree, NULL), "");
}
TEST_END

static unsigned
tree_recurse(node_t *node, unsigned black_height, unsigned black_depth) {
       unsigned ret = 0;
       node_t *left_node;
       node_t *right_node;

       if (node == NULL) {
               return ret;
       }

       left_node = rbtn_left_get(node_t, link, node);
       right_node = rbtn_right_get(node_t, link, node);

       expect_ptr_eq(left_node, node->summary_lchild,
           "summary missed a tree update");
       expect_ptr_eq(right_node, node->summary_rchild,
           "summary missed a tree update");

       uint64_t expected_subtree_specialness = node_subtree_specialness(node,
           left_node, right_node);
       expect_u64_eq(expected_subtree_specialness,
           node->summary_max_specialness, "Incorrect summary");

       if (!rbtn_red_get(node_t, link, node)) {
               black_depth++;
       }

       /* Red nodes must be interleaved with black nodes. */
       if (rbtn_red_get(node_t, link, node)) {
               if (left_node != NULL) {
                       expect_false(rbtn_red_get(node_t, link, left_node),
                               "Node should be black");
               }
               if (right_node != NULL) {
                       expect_false(rbtn_red_get(node_t, link, right_node),
                           "Node should be black");
               }
       }

       /* Self. */
       expect_u32_eq(node->magic, NODE_MAGIC, "Bad magic");

       /* Left subtree. */
       if (left_node != NULL) {
               ret += tree_recurse(left_node, black_height, black_depth);
       } else {
               ret += (black_depth != black_height);
       }

       /* Right subtree. */
       if (right_node != NULL) {
               ret += tree_recurse(right_node, black_height, black_depth);
       } else {
               ret += (black_depth != black_height);
       }

       return ret;
}

static unsigned
tree_iterate(tree_t *tree) {
       unsigned i;

       i = 0;
       tree_iter(tree, NULL, tree_iterate_cb, (void *)&i);

       return i;
}

static unsigned
tree_iterate_reverse(tree_t *tree) {
       unsigned i;

       i = 0;
       tree_reverse_iter(tree, NULL, tree_iterate_cb, (void *)&i);

       return i;
}

static void
node_remove(tree_t *tree, node_t *node, unsigned nnodes) {
       node_t *search_node;
       unsigned black_height, imbalances;

       tree_remove(tree, node);

       /* Test rb_nsearch(). */
       search_node = tree_nsearch(tree, node);
       if (search_node != NULL) {
               expect_u64_ge(search_node->key, node->key,
                   "Key ordering error");
       }

       /* Test rb_psearch(). */
       search_node = tree_psearch(tree, node);
       if (search_node != NULL) {
               expect_u64_le(search_node->key, node->key,
                   "Key ordering error");
       }

       node->magic = 0;

       rbtn_black_height(node_t, link, tree, black_height);
       imbalances = tree_recurse(tree->rbt_root, black_height, 0);
       expect_u_eq(imbalances, 0, "Tree is unbalanced");
       expect_u_eq(tree_iterate(tree), nnodes-1,
           "Unexpected node iteration count");
       expect_u_eq(tree_iterate_reverse(tree), nnodes-1,
           "Unexpected node iteration count");
}

static node_t *
remove_iterate_cb(tree_t *tree, node_t *node, void *data) {
       unsigned *nnodes = (unsigned *)data;
       node_t *ret = tree_next(tree, node);

       node_remove(tree, node, *nnodes);

       return ret;
}

static node_t *
remove_reverse_iterate_cb(tree_t *tree, node_t *node, void *data) {
       unsigned *nnodes = (unsigned *)data;
       node_t *ret = tree_prev(tree, node);

       node_remove(tree, node, *nnodes);

       return ret;
}

static void
destroy_cb(node_t *node, void *data) {
       unsigned *nnodes = (unsigned *)data;

       expect_u_gt(*nnodes, 0, "Destruction removed too many nodes");
       (*nnodes)--;
}

TEST_BEGIN(test_rb_random) {
       enum {
               NNODES = 25,
               NBAGS = 500,
               SEED = 42
       };
       sfmt_t *sfmt;
       uint64_t bag[NNODES];
       tree_t tree;
       node_t nodes[NNODES];
       unsigned i, j, k, black_height, imbalances;

       sfmt = init_gen_rand(SEED);
       for (i = 0; i < NBAGS; i++) {
               switch (i) {
               case 0:
                       /* Insert in order. */
                       for (j = 0; j < NNODES; j++) {
                               bag[j] = j;
                       }
                       break;
               case 1:
                       /* Insert in reverse order. */
                       for (j = 0; j < NNODES; j++) {
                               bag[j] = NNODES - j - 1;
                       }
                       break;
               default:
                       for (j = 0; j < NNODES; j++) {
                               bag[j] = gen_rand64_range(sfmt, NNODES);
                       }
               }

               /*
                * We alternate test behavior with a period of 2 here, and a
                * period of 5 down below, so there's no cycle in which certain
                * combinations get omitted.
                */
               summarize_always_returns_true = (i % 2 == 0);

               for (j = 1; j <= NNODES; j++) {
                       /* Initialize tree and nodes. */
                       tree_new(&tree);
                       for (k = 0; k < j; k++) {
                               nodes[k].magic = NODE_MAGIC;
                               nodes[k].key = bag[k];
                               nodes[k].specialness = gen_rand64_range(sfmt,
                                   NNODES);
                               nodes[k].mid_remove = false;
                               nodes[k].allow_duplicates = false;
                               nodes[k].summary_lchild = NULL;
                               nodes[k].summary_rchild = NULL;
                               nodes[k].summary_max_specialness = 0;
                       }

                       /* Insert nodes. */
                       for (k = 0; k < j; k++) {
                               tree_insert(&tree, &nodes[k]);

                               rbtn_black_height(node_t, link, &tree,
                                   black_height);
                               imbalances = tree_recurse(tree.rbt_root,
                                   black_height, 0);
                               expect_u_eq(imbalances, 0,
                                   "Tree is unbalanced");

                               expect_u_eq(tree_iterate(&tree), k+1,
                                   "Unexpected node iteration count");
                               expect_u_eq(tree_iterate_reverse(&tree), k+1,
                                   "Unexpected node iteration count");

                               expect_false(tree_empty(&tree),
                                   "Tree should not be empty");
                               expect_ptr_not_null(tree_first(&tree),
                                   "Tree should not be empty");
                               expect_ptr_not_null(tree_last(&tree),
                                   "Tree should not be empty");

                               tree_next(&tree, &nodes[k]);
                               tree_prev(&tree, &nodes[k]);
                       }

                       /* Remove nodes. */
                       switch (i % 5) {
                       case 0:
                               for (k = 0; k < j; k++) {
                                       node_remove(&tree, &nodes[k], j - k);
                               }
                               break;
                       case 1:
                               for (k = j; k > 0; k--) {
                                       node_remove(&tree, &nodes[k-1], k);
                               }
                               break;
                       case 2: {
                               node_t *start;
                               unsigned nnodes = j;

                               start = NULL;
                               do {
                                       start = tree_iter(&tree, start,
                                           remove_iterate_cb, (void *)&nnodes);
                                       nnodes--;
                               } while (start != NULL);
                               expect_u_eq(nnodes, 0,
                                   "Removal terminated early");
                               break;
                       } case 3: {
                               node_t *start;
                               unsigned nnodes = j;

                               start = NULL;
                               do {
                                       start = tree_reverse_iter(&tree, start,
                                           remove_reverse_iterate_cb,
                                           (void *)&nnodes);
                                       nnodes--;
                               } while (start != NULL);
                               expect_u_eq(nnodes, 0,
                                   "Removal terminated early");
                               break;
                       } case 4: {
                               unsigned nnodes = j;
                               tree_destroy(&tree, destroy_cb, &nnodes);
                               expect_u_eq(nnodes, 0,
                                   "Destruction terminated early");
                               break;
                       } default:
                               not_reached();
                       }
               }
       }
       fini_gen_rand(sfmt);
}
TEST_END

static void
expect_simple_consistency(tree_t *tree, uint64_t specialness,
   bool expected_empty, node_t *expected_first, node_t *expected_last) {
       bool empty;
       node_t *first;
       node_t *last;

       empty = tree_empty_filtered(tree, &specialness_filter_node,
           &specialness_filter_subtree, &specialness);
       expect_b_eq(expected_empty, empty, "");

       first = tree_first_filtered(tree,
           &specialness_filter_node, &specialness_filter_subtree,
           (void *)&specialness);
       expect_ptr_eq(expected_first, first, "");

       last = tree_last_filtered(tree,
           &specialness_filter_node, &specialness_filter_subtree,
           (void *)&specialness);
       expect_ptr_eq(expected_last, last, "");
}

TEST_BEGIN(test_rb_filter_simple) {
       enum {FILTER_NODES = 10};
       node_t nodes[FILTER_NODES];
       for (unsigned i = 0; i < FILTER_NODES; i++) {
               nodes[i].magic = NODE_MAGIC;
               nodes[i].key = i;
               if (i == 0) {
                       nodes[i].specialness = 0;
               } else {
                       nodes[i].specialness = ffs_u(i);
               }
               nodes[i].mid_remove = false;
               nodes[i].allow_duplicates = false;
               nodes[i].summary_lchild = NULL;
               nodes[i].summary_rchild = NULL;
               nodes[i].summary_max_specialness = 0;
       }

       summarize_always_returns_true = false;

       tree_t tree;
       tree_new(&tree);

       /* Should be empty */
       expect_simple_consistency(&tree, /* specialness */ 0, /* empty */ true,
           /* first */ NULL, /* last */ NULL);

       /* Fill in just the odd nodes. */
       for (int i = 1; i < FILTER_NODES; i += 2) {
               tree_insert(&tree, &nodes[i]);
       }

       /* A search for an odd node should succeed. */
       expect_simple_consistency(&tree, /* specialness */ 0, /* empty */ false,
           /* first */ &nodes[1], /* last */ &nodes[9]);

       /* But a search for an even one should fail. */
       expect_simple_consistency(&tree, /* specialness */ 1, /* empty */ true,
           /* first */ NULL, /* last */ NULL);

       /* Now we add an even. */
       tree_insert(&tree, &nodes[4]);
       expect_simple_consistency(&tree, /* specialness */ 1, /* empty */ false,
           /* first */ &nodes[4], /* last */ &nodes[4]);

       /* A smaller even, and a larger even. */
       tree_insert(&tree, &nodes[2]);
       tree_insert(&tree, &nodes[8]);

       /*
        * A first-search (resp. last-search) for an even should switch to the
        * lower (higher) one, now that it's been added.
        */
       expect_simple_consistency(&tree, /* specialness */ 1, /* empty */ false,
           /* first */ &nodes[2], /* last */ &nodes[8]);

       /*
        * If we remove 2, a first-search we should go back to 4, while a
        * last-search should remain unchanged.
        */
       tree_remove(&tree, &nodes[2]);
       expect_simple_consistency(&tree, /* specialness */ 1, /* empty */ false,
           /* first */ &nodes[4], /* last */ &nodes[8]);

       /* Reinsert 2, then find it again. */
       tree_insert(&tree, &nodes[2]);
       expect_simple_consistency(&tree, /* specialness */ 1, /* empty */ false,
           /* first */ &nodes[2], /* last */ &nodes[8]);

       /* Searching for a multiple of 4 should not have changed. */
       expect_simple_consistency(&tree, /* specialness */ 2, /* empty */ false,
           /* first */ &nodes[4], /* last */ &nodes[8]);

       /* And a multiple of 8 */
       expect_simple_consistency(&tree, /* specialness */ 3, /* empty */ false,
           /* first */ &nodes[8], /* last */ &nodes[8]);

       /* But not a multiple of 16 */
       expect_simple_consistency(&tree, /* specialness */ 4, /* empty */ true,
           /* first */ NULL, /* last */ NULL);
}
TEST_END

typedef struct iter_ctx_s iter_ctx_t;
struct iter_ctx_s {
       int ncalls;
       node_t *last_node;

       int ncalls_max;
       bool forward;
};

static node_t *
tree_iterate_filtered_cb(tree_t *tree, node_t *node, void *arg) {
       iter_ctx_t *ctx = (iter_ctx_t *)arg;
       ctx->ncalls++;
       expect_u64_ge(node->specialness, 1,
           "Should only invoke cb on nodes that pass the filter");
       if (ctx->last_node != NULL) {
               if (ctx->forward) {
                       expect_d_lt(node_cmp(ctx->last_node, node), 0,
                           "Incorrect iteration order");
               } else {
                       expect_d_gt(node_cmp(ctx->last_node, node), 0,
                           "Incorrect iteration order");
               }
       }
       ctx->last_node = node;
       if (ctx->ncalls == ctx->ncalls_max) {
               return node;
       }
       return NULL;
}

static int
qsort_node_cmp(const void *ap, const void *bp) {
       node_t *a = *(node_t **)ap;
       node_t *b = *(node_t **)bp;
       return node_cmp(a, b);
}

#define UPDATE_TEST_MAX 100
static void
check_consistency(tree_t *tree, node_t nodes[UPDATE_TEST_MAX], int nnodes) {
       uint64_t specialness = 1;

       bool empty;
       bool real_empty = true;
       node_t *first;
       node_t *real_first = NULL;
       node_t *last;
       node_t *real_last = NULL;
       for (int i = 0; i < nnodes; i++) {
               if (nodes[i].specialness >= specialness) {
                       real_empty = false;
                       if (real_first == NULL
                           || node_cmp(&nodes[i], real_first) < 0) {
                               real_first = &nodes[i];
                       }
                       if (real_last == NULL
                           || node_cmp(&nodes[i], real_last) > 0) {
                               real_last = &nodes[i];
                       }
               }
       }

       empty = tree_empty_filtered(tree, &specialness_filter_node,
           &specialness_filter_subtree, &specialness);
       expect_b_eq(real_empty, empty, "");

       first = tree_first_filtered(tree, &specialness_filter_node,
           &specialness_filter_subtree, &specialness);
       expect_ptr_eq(real_first, first, "");

       last = tree_last_filtered(tree, &specialness_filter_node,
           &specialness_filter_subtree, &specialness);
       expect_ptr_eq(real_last, last, "");

       for (int i = 0; i < nnodes; i++) {
               node_t *next_filtered;
               node_t *real_next_filtered = NULL;
               node_t *prev_filtered;
               node_t *real_prev_filtered = NULL;
               for (int j = 0; j < nnodes; j++) {
                       if (nodes[j].specialness < specialness) {
                               continue;
                       }
                       if (node_cmp(&nodes[j], &nodes[i]) < 0
                           && (real_prev_filtered == NULL
                           || node_cmp(&nodes[j], real_prev_filtered) > 0)) {
                               real_prev_filtered = &nodes[j];
                       }
                       if (node_cmp(&nodes[j], &nodes[i]) > 0
                           && (real_next_filtered == NULL
                           || node_cmp(&nodes[j], real_next_filtered) < 0)) {
                               real_next_filtered = &nodes[j];
                       }
               }
               next_filtered = tree_next_filtered(tree, &nodes[i],
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_next_filtered, next_filtered, "");

               prev_filtered = tree_prev_filtered(tree, &nodes[i],
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_prev_filtered, prev_filtered, "");

               node_t *search_filtered;
               node_t *real_search_filtered;
               node_t *nsearch_filtered;
               node_t *real_nsearch_filtered;
               node_t *psearch_filtered;
               node_t *real_psearch_filtered;

               /*
                * search, nsearch, psearch from a node before nodes[i] in the
                * ordering.
                */
               node_t before;
               before.magic = NODE_MAGIC;
               before.key = nodes[i].key - 1;
               before.allow_duplicates = false;
               real_search_filtered = NULL;
               search_filtered = tree_search_filtered(tree, &before,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_search_filtered, search_filtered, "");

               real_nsearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_next_filtered);
               nsearch_filtered = tree_nsearch_filtered(tree, &before,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_nsearch_filtered, nsearch_filtered, "");

               real_psearch_filtered = real_prev_filtered;
               psearch_filtered = tree_psearch_filtered(tree, &before,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_psearch_filtered, psearch_filtered, "");

               /* search, nsearch, psearch from nodes[i] */
               real_search_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : NULL);
               search_filtered = tree_search_filtered(tree, &nodes[i],
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_search_filtered, search_filtered, "");

               real_nsearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_next_filtered);
               nsearch_filtered = tree_nsearch_filtered(tree, &nodes[i],
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_nsearch_filtered, nsearch_filtered, "");

               real_psearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_prev_filtered);
               psearch_filtered = tree_psearch_filtered(tree, &nodes[i],
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_psearch_filtered, psearch_filtered, "");

               /*
                * search, nsearch, psearch from a node equivalent to but
                * distinct from nodes[i].
                */
               node_t equiv;
               equiv.magic = NODE_MAGIC;
               equiv.key = nodes[i].key;
               equiv.allow_duplicates = true;
               real_search_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : NULL);
               search_filtered = tree_search_filtered(tree, &equiv,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_search_filtered, search_filtered, "");

               real_nsearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_next_filtered);
               nsearch_filtered = tree_nsearch_filtered(tree, &equiv,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_nsearch_filtered, nsearch_filtered, "");

               real_psearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_prev_filtered);
               psearch_filtered = tree_psearch_filtered(tree, &equiv,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_psearch_filtered, psearch_filtered, "");

               /*
                * search, nsearch, psearch from a node after nodes[i] in the
                * ordering.
                */
               node_t after;
               after.magic = NODE_MAGIC;
               after.key = nodes[i].key + 1;
               after.allow_duplicates = false;
               real_search_filtered = NULL;
               search_filtered = tree_search_filtered(tree, &after,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_search_filtered, search_filtered, "");

               real_nsearch_filtered = real_next_filtered;
               nsearch_filtered = tree_nsearch_filtered(tree, &after,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_nsearch_filtered, nsearch_filtered, "");

               real_psearch_filtered = (nodes[i].specialness >= specialness ?
                   &nodes[i] : real_prev_filtered);
               psearch_filtered = tree_psearch_filtered(tree, &after,
                   &specialness_filter_node, &specialness_filter_subtree,
                   &specialness);
               expect_ptr_eq(real_psearch_filtered, psearch_filtered, "");
       }

       /* Filtered iteration test setup. */
       int nspecial = 0;
       node_t *sorted_nodes[UPDATE_TEST_MAX];
       node_t *sorted_filtered_nodes[UPDATE_TEST_MAX];
       for (int i = 0; i < nnodes; i++) {
               sorted_nodes[i] = &nodes[i];
       }
       qsort(sorted_nodes, nnodes, sizeof(node_t *), &qsort_node_cmp);
       for (int i = 0; i < nnodes; i++) {
               sorted_nodes[i]->rank = i;
               sorted_nodes[i]->filtered_rank = nspecial;
               if (sorted_nodes[i]->specialness >= 1) {
                       sorted_filtered_nodes[nspecial] = sorted_nodes[i];
                       nspecial++;
               }
       }

       node_t *iter_result;

       iter_ctx_t ctx;
       ctx.ncalls = 0;
       ctx.last_node = NULL;
       ctx.ncalls_max = INT_MAX;
       ctx.forward = true;

       /* Filtered forward iteration from the beginning. */
       iter_result = tree_iter_filtered(tree, NULL, &tree_iterate_filtered_cb,
           &ctx, &specialness_filter_node, &specialness_filter_subtree,
           &specialness);
       expect_ptr_null(iter_result, "");
       expect_d_eq(nspecial, ctx.ncalls, "");
       /* Filtered forward iteration from a starting point. */
       for (int i = 0; i < nnodes; i++) {
               ctx.ncalls = 0;
               ctx.last_node = NULL;
               iter_result = tree_iter_filtered(tree, &nodes[i],
                   &tree_iterate_filtered_cb, &ctx, &specialness_filter_node,
                   &specialness_filter_subtree, &specialness);
               expect_ptr_null(iter_result, "");
               expect_d_eq(nspecial - nodes[i].filtered_rank, ctx.ncalls, "");
       }
       /* Filtered forward iteration from the beginning, with stopping */
       for (int i = 0; i < nspecial; i++) {
               ctx.ncalls = 0;
               ctx.last_node = NULL;
               ctx.ncalls_max = i + 1;
               iter_result = tree_iter_filtered(tree, NULL,
                   &tree_iterate_filtered_cb, &ctx, &specialness_filter_node,
                   &specialness_filter_subtree, &specialness);
               expect_ptr_eq(sorted_filtered_nodes[i], iter_result, "");
               expect_d_eq(ctx.ncalls, i + 1, "");
       }
       /* Filtered forward iteration from a starting point, with stopping. */
       for (int i = 0; i < nnodes; i++) {
               for (int j = 0; j < nspecial - nodes[i].filtered_rank; j++) {
                       ctx.ncalls = 0;
                       ctx.last_node = NULL;
                       ctx.ncalls_max = j + 1;
                       iter_result = tree_iter_filtered(tree, &nodes[i],
                           &tree_iterate_filtered_cb, &ctx,
                           &specialness_filter_node,
                           &specialness_filter_subtree, &specialness);
                       expect_d_eq(j + 1, ctx.ncalls, "");
                       expect_ptr_eq(sorted_filtered_nodes[
                           nodes[i].filtered_rank + j], iter_result, "");
               }
       }

       /* Backwards iteration. */
       ctx.ncalls = 0;
       ctx.last_node = NULL;
       ctx.ncalls_max = INT_MAX;
       ctx.forward = false;

       /* Filtered backward iteration from the end. */
       iter_result = tree_reverse_iter_filtered(tree, NULL,
           &tree_iterate_filtered_cb, &ctx, &specialness_filter_node,
           &specialness_filter_subtree, &specialness);
       expect_ptr_null(iter_result, "");
       expect_d_eq(nspecial, ctx.ncalls, "");
       /* Filtered backward iteration from a starting point. */
       for (int i = 0; i < nnodes; i++) {
               ctx.ncalls = 0;
               ctx.last_node = NULL;
               iter_result = tree_reverse_iter_filtered(tree, &nodes[i],
                   &tree_iterate_filtered_cb, &ctx, &specialness_filter_node,
                   &specialness_filter_subtree, &specialness);
               expect_ptr_null(iter_result, "");
               int surplus_rank = (nodes[i].specialness >= 1 ? 1 : 0);
               expect_d_eq(nodes[i].filtered_rank + surplus_rank, ctx.ncalls,
                   "");
       }
       /* Filtered backward iteration from the end, with stopping */
       for (int i = 0; i < nspecial; i++) {
               ctx.ncalls = 0;
               ctx.last_node = NULL;
               ctx.ncalls_max = i + 1;
               iter_result = tree_reverse_iter_filtered(tree, NULL,
                   &tree_iterate_filtered_cb, &ctx, &specialness_filter_node,
                   &specialness_filter_subtree, &specialness);
               expect_ptr_eq(sorted_filtered_nodes[nspecial - i - 1],
                   iter_result, "");
               expect_d_eq(ctx.ncalls, i + 1, "");
       }
       /* Filtered backward iteration from a starting point, with stopping. */
       for (int i = 0; i < nnodes; i++) {
               int surplus_rank = (nodes[i].specialness >= 1 ? 1 : 0);
               for (int j = 0; j < nodes[i].filtered_rank + surplus_rank;
                   j++) {
                       ctx.ncalls = 0;
                       ctx.last_node = NULL;
                       ctx.ncalls_max = j + 1;
                       iter_result = tree_reverse_iter_filtered(tree,
                           &nodes[i], &tree_iterate_filtered_cb, &ctx,
                           &specialness_filter_node,
                           &specialness_filter_subtree, &specialness);
                       expect_d_eq(j + 1, ctx.ncalls, "");
                       expect_ptr_eq(sorted_filtered_nodes[
                           nodes[i].filtered_rank - j - 1 + surplus_rank],
                           iter_result, "");
               }
       }
}

static void
do_update_search_test(int nnodes, int ntrees, int nremovals,
   int nupdates) {
       node_t nodes[UPDATE_TEST_MAX];
       assert(nnodes <= UPDATE_TEST_MAX);

       sfmt_t *sfmt = init_gen_rand(12345);
       for (int i = 0; i < ntrees; i++) {
               tree_t tree;
               tree_new(&tree);
               for (int j = 0; j < nnodes; j++) {
                       nodes[j].magic = NODE_MAGIC;
                       /*
                        * In consistency checking, we increment or decrement a
                        * key and assume that the result is not a key in the
                        * tree.  This isn't a *real* concern with 64-bit keys
                        * and a good PRNG, but why not be correct anyways?
                        */
                       nodes[j].key = 2 * gen_rand64(sfmt);
                       nodes[j].specialness = 0;
                       nodes[j].mid_remove = false;
                       nodes[j].allow_duplicates = false;
                       nodes[j].summary_lchild = NULL;
                       nodes[j].summary_rchild = NULL;
                       nodes[j].summary_max_specialness = 0;
                       tree_insert(&tree, &nodes[j]);
               }
               for (int j = 0; j < nremovals; j++) {
                       int victim = (int)gen_rand64_range(sfmt, nnodes);
                       if (!nodes[victim].mid_remove) {
                               tree_remove(&tree, &nodes[victim]);
                               nodes[victim].mid_remove = true;
                       }
               }
               for (int j = 0; j < nnodes; j++) {
                       if (nodes[j].mid_remove) {
                               nodes[j].mid_remove = false;
                               nodes[j].key = 2 * gen_rand64(sfmt);
                               tree_insert(&tree, &nodes[j]);
                       }
               }
               for (int j = 0; j < nupdates; j++) {
                       uint32_t ind = gen_rand32_range(sfmt, nnodes);
                       nodes[ind].specialness = 1 - nodes[ind].specialness;
                       tree_update_summaries(&tree, &nodes[ind]);
                       check_consistency(&tree, nodes, nnodes);
               }
       }
}

TEST_BEGIN(test_rb_update_search) {
       summarize_always_returns_true = false;
       do_update_search_test(2, 100, 3, 50);
       do_update_search_test(5, 100, 3, 50);
       do_update_search_test(12, 100, 5, 1000);
       do_update_search_test(100, 1, 50, 500);
}
TEST_END

typedef rb_tree(node_t) unsummarized_tree_t;
rb_gen(static UNUSED, unsummarized_tree_, unsummarized_tree_t, node_t, link,
   node_cmp);

static node_t *
unsummarized_tree_iterate_cb(unsummarized_tree_t *tree, node_t *node,
   void *data) {
       unsigned *i = (unsigned *)data;
       (*i)++;
       return NULL;
}
/*
* The unsummarized and summarized funtionality is implemented via the same
* functions; we don't really need to do much more than test that we can exclude
* the filtered functionality without anything breaking.
*/
TEST_BEGIN(test_rb_unsummarized) {
       unsummarized_tree_t tree;
       unsummarized_tree_new(&tree);
       unsigned nnodes = 0;
       unsummarized_tree_iter(&tree, NULL, &unsummarized_tree_iterate_cb,
           &nnodes);
       expect_u_eq(0, nnodes, "");
}
TEST_END

int
main(void) {
       return test_no_reentrancy(
           test_rb_empty,
           test_rb_random,
           test_rb_filter_simple,
           test_rb_update_search,
           test_rb_unsummarized);
}

You are viewing proxied material from ftp.icm.edu.pl. The copyright of proxied material belongs to its original authors. Any comments or complaints in relation to proxied material should be directed to the original authors of the content concerned. Please see the disclaimer for more details.