// SPDX-License-Identifier: GPL-2.0+

/*
 * Crossported from the same named file of btrfs-progs.
 *
 * Minor modification to include headers.
 */
#include <linux/kernel.h>
#include <linux/rbtree.h>
#include <linux/errno.h>
#include <linux/bug.h>
#include <stdlib.h>
#include "extent-cache.h"
#include "common/rbtree-utils.h"

struct cache_extent_search_range {
	u64 objectid;
	u64 start;
	u64 size;
};

static int cache_tree_comp_range(struct rb_node *node, void *data)
{
	struct cache_extent *entry;
	struct cache_extent_search_range *range;

	range = (struct cache_extent_search_range *)data;
	entry = rb_entry(node, struct cache_extent, rb_node);

	if (entry->start + entry->size <= range->start)
		return 1;
	else if (range->start + range->size <= entry->start)
		return -1;
	else
		return 0;
}

static int cache_tree_comp_nodes(struct rb_node *node1, struct rb_node *node2)
{
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	entry = rb_entry(node2, struct cache_extent, rb_node);
	range.start = entry->start;
	range.size = entry->size;

	return cache_tree_comp_range(node1, (void *)&range);
}

static int cache_tree_comp_range2(struct rb_node *node, void *data)
{
	struct cache_extent *entry;
	struct cache_extent_search_range *range;

	range = (struct cache_extent_search_range *)data;
	entry = rb_entry(node, struct cache_extent, rb_node);

	if (entry->objectid < range->objectid)
		return 1;
	else if (entry->objectid > range->objectid)
		return -1;
	else if (entry->start + entry->size <= range->start)
		return 1;
	else if (range->start + range->size <= entry->start)
		return -1;
	else
		return 0;
}

static int cache_tree_comp_nodes2(struct rb_node *node1, struct rb_node *node2)
{
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	entry = rb_entry(node2, struct cache_extent, rb_node);
	range.objectid = entry->objectid;
	range.start = entry->start;
	range.size = entry->size;

	return cache_tree_comp_range2(node1, (void *)&range);
}

void cache_tree_init(struct cache_tree *tree)
{
	tree->root = RB_ROOT;
}

static struct cache_extent *alloc_cache_extent(u64 start, u64 size)
{
	struct cache_extent *pe = malloc(sizeof(*pe));

	if (!pe)
		return pe;

	pe->objectid = 0;
	pe->start = start;
	pe->size = size;
	return pe;
}

int add_cache_extent(struct cache_tree *tree, u64 start, u64 size)
{
	struct cache_extent *pe = alloc_cache_extent(start, size);
	int ret;

	if (!pe)
		return -ENOMEM;

	ret = insert_cache_extent(tree, pe);
	if (ret)
		free(pe);

	return ret;
}

int insert_cache_extent(struct cache_tree *tree, struct cache_extent *pe)
{
	return rb_insert(&tree->root, &pe->rb_node, cache_tree_comp_nodes);
}

int insert_cache_extent2(struct cache_tree *tree, struct cache_extent *pe)
{
	return rb_insert(&tree->root, &pe->rb_node, cache_tree_comp_nodes2);
}

struct cache_extent *lookup_cache_extent(struct cache_tree *tree,
					 u64 start, u64 size)
{
	struct rb_node *node;
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	range.start = start;
	range.size = size;
	node = rb_search(&tree->root, &range, cache_tree_comp_range, NULL);
	if (!node)
		return NULL;

	entry = rb_entry(node, struct cache_extent, rb_node);
	return entry;
}

struct cache_extent *lookup_cache_extent2(struct cache_tree *tree,
					 u64 objectid, u64 start, u64 size)
{
	struct rb_node *node;
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	range.objectid = objectid;
	range.start = start;
	range.size = size;
	node = rb_search(&tree->root, &range, cache_tree_comp_range2, NULL);
	if (!node)
		return NULL;

	entry = rb_entry(node, struct cache_extent, rb_node);
	return entry;
}

struct cache_extent *search_cache_extent(struct cache_tree *tree, u64 start)
{
	struct rb_node *next;
	struct rb_node *node;
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	range.start = start;
	range.size = 1;
	node = rb_search(&tree->root, &range, cache_tree_comp_range, &next);
	if (!node)
		node = next;
	if (!node)
		return NULL;

	entry = rb_entry(node, struct cache_extent, rb_node);
	return entry;
}

struct cache_extent *search_cache_extent2(struct cache_tree *tree,
					 u64 objectid, u64 start)
{
	struct rb_node *next;
	struct rb_node *node;
	struct cache_extent *entry;
	struct cache_extent_search_range range;

	range.objectid = objectid;
	range.start = start;
	range.size = 1;
	node = rb_search(&tree->root, &range, cache_tree_comp_range2, &next);
	if (!node)
		node = next;
	if (!node)
		return NULL;

	entry = rb_entry(node, struct cache_extent, rb_node);
	return entry;
}

struct cache_extent *first_cache_extent(struct cache_tree *tree)
{
	struct rb_node *node = rb_first(&tree->root);

	if (!node)
		return NULL;
	return rb_entry(node, struct cache_extent, rb_node);
}

struct cache_extent *last_cache_extent(struct cache_tree *tree)
{
	struct rb_node *node = rb_last(&tree->root);

	if (!node)
		return NULL;
	return rb_entry(node, struct cache_extent, rb_node);
}

struct cache_extent *prev_cache_extent(struct cache_extent *pe)
{
	struct rb_node *node = rb_prev(&pe->rb_node);

	if (!node)
		return NULL;
	return rb_entry(node, struct cache_extent, rb_node);
}

struct cache_extent *next_cache_extent(struct cache_extent *pe)
{
	struct rb_node *node = rb_next(&pe->rb_node);

	if (!node)
		return NULL;
	return rb_entry(node, struct cache_extent, rb_node);
}

void remove_cache_extent(struct cache_tree *tree, struct cache_extent *pe)
{
	rb_erase(&pe->rb_node, &tree->root);
}

void cache_tree_free_extents(struct cache_tree *tree,
			     free_cache_extent free_func)
{
	struct cache_extent *ce;

	while ((ce = first_cache_extent(tree))) {
		remove_cache_extent(tree, ce);
		free_func(ce);
	}
}

static void free_extent_cache(struct cache_extent *pe)
{
	free(pe);
}

void free_extent_cache_tree(struct cache_tree *tree)
{
	cache_tree_free_extents(tree, free_extent_cache);
}

int add_merge_cache_extent(struct cache_tree *tree, u64 start, u64 size)
{
	struct cache_extent *cache;
	struct cache_extent *next = NULL;
	struct cache_extent *prev = NULL;
	int next_merged = 0;
	int prev_merged = 0;
	int ret = 0;

	if (cache_tree_empty(tree))
		goto insert;

	cache = search_cache_extent(tree, start);
	if (!cache) {
		/*
		 * Either the tree is completely empty, or the no range after
		 * start.
		 * Either way, the last cache_extent should be prev.
		 */
		prev = last_cache_extent(tree);
	} else if (start <= cache->start) {
		next = cache;
		prev = prev_cache_extent(cache);
	} else {
		prev = cache;
		next = next_cache_extent(cache);
	}

	/*
	 * Ensure the range to be inserted won't cover with existings
	 * Or we will need extra loop to do merge
	 */
	BUG_ON(next && start + size > next->start);
	BUG_ON(prev && prev->start + prev->size > start);

	if (next && start + size == next->start) {
		next_merged = 1;
		next->size = next->start + next->size - start;
		next->start = start;
	}
	if (prev && prev->start + prev->size == start) {
		prev_merged = 1;
		if (next_merged) {
			next->size = next->start + next->size - prev->start;
			next->start = prev->start;
			remove_cache_extent(tree, prev);
			free(prev);
		} else {
			prev->size = start + size - prev->start;
		}
	}
insert:
	if (!prev_merged && !next_merged)
		ret = add_cache_extent(tree, start, size);
	return ret;
}