/*-
 * Copyright (c) 2015 Taylor R. Campbell
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#ifndef	RBTREE_H
#define	RBTREE_H

#include <stdint.h>

#include "wb.h"

typedef int (*rb_compare_nodes_fn)(void *, const void *, const void *);
typedef int (*rb_compare_key_fn)(void *, const void *, const void *);

struct rb_tree_ops {
	rb_compare_nodes_fn	rbto_compare_nodes;
	rb_compare_key_fn	rbto_compare_key;
	size_t			rbto_node_offset;
	void			*rbto_context;
};
typedef struct rb_tree_ops rb_tree_ops_t;

struct rb_tree {
	struct wb			*rbt_root;
	const struct rb_tree_ops	*rbt_ops;
};
typedef struct rb_tree rb_tree_t;

struct rb_node {
	struct wb	rbn_node;
};
typedef struct rb_node rb_node_t;

static inline void
rb_tree_init(struct rb_tree *rbt, const struct rb_tree_ops *ops)
{

	rbt->rbt_root = NULL;
	rbt->rbt_ops = ops;
}

static inline void
rb_tree_fini(struct rb_tree *rbt)
{

	rbt->rbt_ops = NULL;
}

static inline void *
_rb_tree_ops_wb2vp(const struct rb_tree_ops *ops, const struct wb *wb)
{

	return (void *)((uintptr_t)wb - ops->rbto_node_offset);
}

static inline struct wb *
_rb_tree_ops_vp2wb(const struct rb_tree_ops *ops, const void *vp)
{

	return (struct wb *)((uintptr_t)vp + ops->rbto_node_offset);
}

static inline void *
_rb_tree_wb2vp(const struct rb_tree *rbt, struct wb *wb)
{

	return _rb_tree_ops_wb2vp(rbt->rbt_ops, wb);
}

static inline struct wb *
_rb_tree_vp2wb(const struct rb_tree *rbt, void *vp)
{

	return _rb_tree_ops_vp2wb(rbt->rbt_ops, vp);
}

static inline void *
RB_TREE_MIN(const struct rb_tree *rbt)
{
	struct wb *wb;

	wb = wb_first(rbt->rbt_root);
	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

static inline void *
RB_TREE_MAX(const struct rb_tree *rbt)
{
	struct wb *wb;

	wb = wb_last(rbt->rbt_root);
	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

#define	RB_DIR_LEFT	WB_ITER_DESC
#define	RB_DIR_RIGHT	WB_ITER_ASC

static inline void *
rb_tree_iterate(struct rb_tree *rbt, void *vp, unsigned dir)
{
	struct wb *wb;

	if (vp == NULL)
		wb = wb_iter_start(rbt->rbt_root, 1^dir);
	else
		wb = wb_iter(_rb_tree_vp2wb(rbt, vp), dir);

	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

#define	RB_TREE_FOREACH(N, T)						\
	for ((N) = RB_TREE_MIN(T);					\
	     (N) != NULL;						\
	     (N) = rb_tree_iterate((T), (N), RB_DIR_RIGHT))

#define	RB_TREE_FOREACH_SAFE(N, T, N0)					\
	for ((N) = RB_TREE_MIN(T);					\
	     ((N) != NULL &&						\
		((N0) = rb_tree_iterate((T), (N), RB_DIR_RIGHT), 1));	\
	     (N) = (N0))

#define	RB_TREE_FOREACH_REVERSE(N, T)					\
	for ((N) = RB_TREE_MAX(T);					\
	     (N) != NULL;						\
	     (N) = rb_tree_iterate((T), (N), RB_DIR_LEFT))

#define	RB_TREE_FOREACH_REVERSE_SAFE(N, T, N0)				\
	for ((N) = RB_TREE_MAX(T);					\
	     ((N) != NULL &&						\
		((N0) = rb_tree_iterate((T), (N), RB_DIR_LEFT), 1));	\
	     (N) = (N0))

static int
_rb_tree_compare_key(void *cookie, const void *key, const struct wb *wb)
{
	const struct rb_tree_ops *const ops = cookie;

	return -(*ops->rbto_compare_key)(ops->rbto_context,
	    _rb_tree_ops_wb2vp(ops, wb), key);
}

static int
_rb_tree_compare_nodes(void *cookie, const struct wb *wa, const struct wb *wb)
{
	const struct rb_tree_ops *const ops = cookie;

	return (*ops->rbto_compare_nodes)(ops->rbto_context,
	    _rb_tree_ops_wb2vp(ops, wa), _rb_tree_ops_wb2vp(ops, wb));
}

static inline void *
rb_tree_insert_node(struct rb_tree *rbt, void *vp)
{
	struct wb *wb, *wb0;

	wb = _rb_tree_vp2wb(rbt, vp);
	wb0 = wb_insert_key_compare3(&rbt->rbt_root, wb,
	    &_rb_tree_compare_nodes, (void *)(uintptr_t)rbt->rbt_ops);

	return _rb_tree_wb2vp(rbt, wb0);
}

static inline void
rb_tree_remove_node(struct rb_tree *rbt, void *vp)
{

	wb_delete(&rbt->rbt_root, _rb_tree_vp2wb(rbt, vp));
}

static inline void *
rb_tree_find_node(const struct rb_tree *rbt, const void *key)
{
	struct wb *wb;

	wb = wb_find_key_compare3(rbt->rbt_root, key, &_rb_tree_compare_key,
	    (void *)(uintptr_t)rbt->rbt_ops);

	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

static inline void *
rb_tree_find_node_leq(const struct rb_tree *rbt, const void *key)
{
	struct wb *wb;

	wb = wb_find_key_leq_compare3(rbt->rbt_root, key,
	    &_rb_tree_compare_key, (void *)(uintptr_t)rbt->rbt_ops);

	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

static inline void *
rb_tree_find_node_geq(const struct rb_tree *rbt, const void *key)
{
	struct wb *wb;

	wb = wb_find_key_geq_compare3(rbt->rbt_root, key,
	    &_rb_tree_compare_key, (void *)(uintptr_t)rbt->rbt_ops);

	return wb == NULL ? NULL : _rb_tree_wb2vp(rbt, wb);
}

#endif	/* RBTREE_H */