// SPDX-License-Identifier: GPL-2.0
#include "block-range.h"
#include "annotate.h"
#include <assert.h>
#include <stdlib.h>

struct {
	struct rb_root root;
	u64 blocks;
} block_ranges;

static void block_range__debug(void)
{
#ifndef NDEBUG
	struct rb_node *rb;
	u64 old = 0; /* NULL isn't executable */

	for (rb = rb_first(&block_ranges.root); rb; rb = rb_next(rb)) {
		struct block_range *entry = rb_entry(rb, struct block_range, node);

		assert(old < entry->start);
		assert(entry->start <= entry->end); /* single instruction block; jump to a jump */

		old = entry->end;
	}
#endif
}

struct block_range *block_range__find(u64 addr)
{
	struct rb_node **p = &block_ranges.root.rb_node;
	struct rb_node *parent = NULL;
	struct block_range *entry;

	while (*p != NULL) {
		parent = *p;
		entry = rb_entry(parent, struct block_range, node);

		if (addr < entry->start)
			p = &parent->rb_left;
		else if (addr > entry->end)
			p = &parent->rb_right;
		else
			return entry;
	}

	return NULL;
}

static inline void rb_link_left_of_node(struct rb_node *left, struct rb_node *node)
{
	struct rb_node **p = &node->rb_left;
	while (*p) {
		node = *p;
		p = &node->rb_right;
	}
	rb_link_node(left, node, p);
}

static inline void rb_link_right_of_node(struct rb_node *right, struct rb_node *node)
{
	struct rb_node **p = &node->rb_right;
	while (*p) {
		node = *p;
		p = &node->rb_left;
	}
	rb_link_node(right, node, p);
}

/**
 * block_range__create
 * @start: branch target starting this basic block
 * @end:   branch ending this basic block
 *
 * Create all the required block ranges to precisely span the given range.
 */
struct block_range_iter block_range__create(u64 start, u64 end)
{
	struct rb_node **p = &block_ranges.root.rb_node;
	struct rb_node *n, *parent = NULL;
	struct block_range *next, *entry = NULL;
	struct block_range_iter iter = { NULL, NULL };

	while (*p != NULL) {
		parent = *p;
		entry = rb_entry(parent, struct block_range, node);

		if (start < entry->start)
			p = &parent->rb_left;
		else if (start > entry->end)
			p = &parent->rb_right;
		else
			break;
	}

	/*
	 * Didn't find anything.. there's a hole at @start, however @end might
	 * be inside/behind the next range.
	 */
	if (!*p) {
		if (!entry) /* tree empty */
			goto do_whole;

		/*
		 * If the last node is before, advance one to find the next.
		 */
		n = parent;
		if (entry->end < start) {
			n = rb_next(n);
			if (!n)
				goto do_whole;
		}
		next = rb_entry(n, struct block_range, node);

		if (next->start <= end) { /* add head: [start...][n->start...] */
			struct block_range *head = malloc(sizeof(struct block_range));
			if (!head)
				return iter;

			*head = (struct block_range){
				.start		= start,
				.end		= next->start - 1,
				.is_target	= 1,
				.is_branch	= 0,
			};

			rb_link_left_of_node(&head->node, &next->node);
			rb_insert_color(&head->node, &block_ranges.root);
			block_range__debug();

			iter.start = head;
			goto do_tail;
		}

do_whole:
		/*
		 * The whole [start..end] range is non-overlapping.
		 */
		entry = malloc(sizeof(struct block_range));
		if (!entry)
			return iter;

		*entry = (struct block_range){
			.start		= start,
			.end		= end,
			.is_target	= 1,
			.is_branch	= 1,
		};

		rb_link_node(&entry->node, parent, p);
		rb_insert_color(&entry->node, &block_ranges.root);
		block_range__debug();

		iter.start = entry;
		iter.end   = entry;
		goto done;
	}

	/*
	 * We found a range that overlapped with ours, split if needed.
	 */
	if (entry->start < start) { /* split: [e->start...][start...] */
		struct block_range *head = malloc(sizeof(struct block_range));
		if (!head)
			return iter;

		*head = (struct block_range){
			.start		= entry->start,
			.end		= start - 1,
			.is_target	= entry->is_target,
			.is_branch	= 0,

			.coverage	= entry->coverage,
			.entry		= entry->entry,
		};

		entry->start		= start;
		entry->is_target	= 1;
		entry->entry		= 0;

		rb_link_left_of_node(&head->node, &entry->node);
		rb_insert_color(&head->node, &block_ranges.root);
		block_range__debug();

	} else if (entry->start == start)
		entry->is_target = 1;

	iter.start = entry;

do_tail:
	/*
	 * At this point we've got: @iter.start = [@start...] but @end can still be
	 * inside or beyond it.
	 */
	entry = iter.start;
	for (;;) {
		/*
		 * If @end is inside @entry, split.
		 */
		if (end < entry->end) { /* split: [...end][...e->end] */
			struct block_range *tail = malloc(sizeof(struct block_range));
			if (!tail)
				return iter;

			*tail = (struct block_range){
				.start		= end + 1,
				.end		= entry->end,
				.is_target	= 0,
				.is_branch	= entry->is_branch,

				.coverage	= entry->coverage,
				.taken		= entry->taken,
				.pred		= entry->pred,
			};

			entry->end		= end;
			entry->is_branch	= 1;
			entry->taken		= 0;
			entry->pred		= 0;

			rb_link_right_of_node(&tail->node, &entry->node);
			rb_insert_color(&tail->node, &block_ranges.root);
			block_range__debug();

			iter.end = entry;
			goto done;
		}

		/*
		 * If @end matches @entry, done
		 */
		if (end == entry->end) {
			entry->is_branch = 1;
			iter.end = entry;
			goto done;
		}

		next = block_range__next(entry);
		if (!next)
			goto add_tail;

		/*
		 * If @end is in beyond @entry but not inside @next, add tail.
		 */
		if (end < next->start) { /* add tail: [...e->end][...end] */
			struct block_range *tail;
add_tail:
			tail = malloc(sizeof(struct block_range));
			if (!tail)
				return iter;

			*tail = (struct block_range){
				.start		= entry->end + 1,
				.end		= end,
				.is_target	= 0,
				.is_branch	= 1,
			};

			rb_link_right_of_node(&tail->node, &entry->node);
			rb_insert_color(&tail->node, &block_ranges.root);
			block_range__debug();

			iter.end = tail;
			goto done;
		}

		/*
		 * If there is a hole between @entry and @next, fill it.
		 */
		if (entry->end + 1 != next->start) {
			struct block_range *hole = malloc(sizeof(struct block_range));
			if (!hole)
				return iter;

			*hole = (struct block_range){
				.start		= entry->end + 1,
				.end		= next->start - 1,
				.is_target	= 0,
				.is_branch	= 0,
			};

			rb_link_left_of_node(&hole->node, &next->node);
			rb_insert_color(&hole->node, &block_ranges.root);
			block_range__debug();
		}

		entry = next;
	}

done:
	assert(iter.start->start == start && iter.start->is_target);
	assert(iter.end->end == end && iter.end->is_branch);

	block_ranges.blocks++;

	return iter;
}


/*
 * Compute coverage as:
 *
 *    br->coverage / br->sym->max_coverage
 *
 * This ensures each symbol has a 100% spot, to reflect that each symbol has a
 * most covered section.
 *
 * Returns [0-1] for coverage and -1 if we had no data what so ever or the
 * symbol does not exist.
 */
double block_range__coverage(struct block_range *br)
{
	struct symbol *sym;
	struct annotated_branch *branch;

	if (!br) {
		if (block_ranges.blocks)
			return 0;

		return -1;
	}

	sym = br->sym;
	if (!sym)
		return -1;

	branch = symbol__annotation(sym)->branch;
	if (!branch)
		return -1;

	return (double)br->coverage / branch->max_coverage;
}