/* SPDX-License-Identifier: GPL-2.0 */
#ifndef _LINUX_MIN_HEAP_H
#define _LINUX_MIN_HEAP_H

#include <linux/bug.h>
#include <linux/string.h>
#include <linux/types.h>




#define MIN_HEAP_PREALLOCATED(_type, _name, _nr)	\
struct _name {	\
	size_t nr;	\
	size_t size;	\
	_type *data;	\
	_type preallocated[_nr];	\
}

#define DEFINE_MIN_HEAP(_type, _name) MIN_HEAP_PREALLOCATED(_type, _name, 0)

typedef DEFINE_MIN_HEAP(char, min_heap_char) min_heap_char;

#define __minheap_cast(_heap)		(typeof((_heap)->data[0]) *)
#define __minheap_obj_size(_heap)	sizeof((_heap)->data[0])


struct min_heap_callbacks {
	bool (*less)(const void *lhs, const void *rhs, void *args);
	void (*swp)(void *lhs, void *rhs, void *args);
};


__attribute_const__ __always_inline
static bool is_aligned(const void *base, size_t size, unsigned char align)
{
	unsigned char lsbits = (unsigned char)size;

	(void)base;
#ifndef CONFIG_HAVE_EFFICIENT_UNALIGNED_ACCESS
	lsbits |= (unsigned char)(uintptr_t)base;
#endif
	return (lsbits & (align - 1)) == 0;
}


static __always_inline
void swap_words_32(void *a, void *b, size_t n)
{
	do {
		u32 t = *(u32 *)(a + (n -= 4));
		*(u32 *)(a + n) = *(u32 *)(b + n);
		*(u32 *)(b + n) = t;
	} while (n);
}


static __always_inline
void swap_words_64(void *a, void *b, size_t n)
{
	do {
#ifdef CONFIG_64BIT
		u64 t = *(u64 *)(a + (n -= 8));
		*(u64 *)(a + n) = *(u64 *)(b + n);
		*(u64 *)(b + n) = t;
#else
		
		u32 t = *(u32 *)(a + (n -= 4));
		*(u32 *)(a + n) = *(u32 *)(b + n);
		*(u32 *)(b + n) = t;

		t = *(u32 *)(a + (n -= 4));
		*(u32 *)(a + n) = *(u32 *)(b + n);
		*(u32 *)(b + n) = t;
#endif
	} while (n);
}


static __always_inline
void swap_bytes(void *a, void *b, size_t n)
{
	do {
		char t = ((char *)a)[--n];
		((char *)a)[n] = ((char *)b)[n];
		((char *)b)[n] = t;
	} while (n);
}


#define SWAP_WORDS_64 ((void (*)(void *, void *, void *))0)
#define SWAP_WORDS_32 ((void (*)(void *, void *, void *))1)
#define SWAP_BYTES    ((void (*)(void *, void *, void *))2)


static __always_inline
void *select_swap_func(const void *base, size_t size)
{
	if (is_aligned(base, size, 8))
		return SWAP_WORDS_64;
	else if (is_aligned(base, size, 4))
		return SWAP_WORDS_32;
	else
		return SWAP_BYTES;
}

static __always_inline
void do_swap(void *a, void *b, size_t size, void (*swap_func)(void *lhs, void *rhs, void *args),
	     void *priv)
{
	if (swap_func == SWAP_WORDS_64)
		swap_words_64(a, b, size);
	else if (swap_func == SWAP_WORDS_32)
		swap_words_32(a, b, size);
	else if (swap_func == SWAP_BYTES)
		swap_bytes(a, b, size);
	else
		swap_func(a, b, priv);
}


__attribute_const__ __always_inline
static size_t parent(size_t i, unsigned int lsbit, size_t size)
{
	i -= size;
	i -= size & -(i & lsbit);
	return i / 2;
}


static __always_inline
void __min_heap_init_inline(min_heap_char *heap, void *data, size_t size)
{
	heap->nr = 0;
	heap->size = size;
	if (data)
		heap->data = data;
	else
		heap->data = heap->preallocated;
}

#define min_heap_init_inline(_heap, _data, _size)	\
	__min_heap_init_inline(container_of(&(_heap)->nr, min_heap_char, nr), _data, _size)


static __always_inline
void *__min_heap_peek_inline(struct min_heap_char *heap)
{
	return heap->nr ? heap->data : NULL;
}

#define min_heap_peek_inline(_heap)	\
	(__minheap_cast(_heap)	\
	 __min_heap_peek_inline(container_of(&(_heap)->nr, min_heap_char, nr)))


static __always_inline
bool __min_heap_full_inline(min_heap_char *heap)
{
	return heap->nr == heap->size;
}

#define min_heap_full_inline(_heap)	\
	__min_heap_full_inline(container_of(&(_heap)->nr, min_heap_char, nr))


static __always_inline
void __min_heap_sift_down_inline(min_heap_char *heap, size_t pos, size_t elem_size,
				 const struct min_heap_callbacks *func, void *args)
{
	const unsigned long lsbit = elem_size & -elem_size;
	void *data = heap->data;
	void (*swp)(void *lhs, void *rhs, void *args) = func->swp;
	
	size_t a = pos * elem_size;
	size_t b, c, d;
	size_t n = heap->nr * elem_size;

	if (!swp)
		swp = select_swap_func(data, elem_size);

	
	for (b = a; c = 2 * b + elem_size, (d = c + elem_size) < n;)
		b = func->less(data + c, data + d, args) ? c : d;

	
	if (d == n)
		b = c;

	
	while (b != a && func->less(data + a, data + b, args))
		b = parent(b, lsbit, elem_size);

	
	c = b;
	while (b != a) {
		b = parent(b, lsbit, elem_size);
		do_swap(data + b, data + c, elem_size, swp, args);
	}
}

#define min_heap_sift_down_inline(_heap, _pos, _func, _args)	\
	__min_heap_sift_down_inline(container_of(&(_heap)->nr, min_heap_char, nr), _pos,	\
				    __minheap_obj_size(_heap), _func, _args)


static __always_inline
void __min_heap_sift_up_inline(min_heap_char *heap, size_t elem_size, size_t idx,
			       const struct min_heap_callbacks *func, void *args)
{
	const unsigned long lsbit = elem_size & -elem_size;
	void *data = heap->data;
	void (*swp)(void *lhs, void *rhs, void *args) = func->swp;
	
	size_t a = idx * elem_size, b;

	if (!swp)
		swp = select_swap_func(data, elem_size);

	while (a) {
		b = parent(a, lsbit, elem_size);
		if (func->less(data + b, data + a, args))
			break;
		do_swap(data + a, data + b, elem_size, swp, args);
		a = b;
	}
}

#define min_heap_sift_up_inline(_heap, _idx, _func, _args)	\
	__min_heap_sift_up_inline(container_of(&(_heap)->nr, min_heap_char, nr),	\
				  __minheap_obj_size(_heap), _idx, _func, _args)


static __always_inline
void __min_heapify_all_inline(min_heap_char *heap, size_t elem_size,
			      const struct min_heap_callbacks *func, void *args)
{
	ssize_t i;

	for (i = heap->nr / 2 - 1; i >= 0; i--)
		__min_heap_sift_down_inline(heap, i, elem_size, func, args);
}

#define min_heapify_all_inline(_heap, _func, _args)	\
	__min_heapify_all_inline(container_of(&(_heap)->nr, min_heap_char, nr),	\
				 __minheap_obj_size(_heap), _func, _args)


static __always_inline
bool __min_heap_pop_inline(min_heap_char *heap, size_t elem_size,
			   const struct min_heap_callbacks *func, void *args)
{
	void *data = heap->data;

	if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap"))
		return false;

	
	heap->nr--;
	memcpy(data, data + (heap->nr * elem_size), elem_size);
	__min_heap_sift_down_inline(heap, 0, elem_size, func, args);

	return true;
}

#define min_heap_pop_inline(_heap, _func, _args)	\
	__min_heap_pop_inline(container_of(&(_heap)->nr, min_heap_char, nr),	\
			      __minheap_obj_size(_heap), _func, _args)


static __always_inline
void __min_heap_pop_push_inline(min_heap_char *heap, const void *element, size_t elem_size,
				const struct min_heap_callbacks *func, void *args)
{
	memcpy(heap->data, element, elem_size);
	__min_heap_sift_down_inline(heap, 0, elem_size, func, args);
}

#define min_heap_pop_push_inline(_heap, _element, _func, _args)	\
	__min_heap_pop_push_inline(container_of(&(_heap)->nr, min_heap_char, nr), _element,	\
				   __minheap_obj_size(_heap), _func, _args)


static __always_inline
bool __min_heap_push_inline(min_heap_char *heap, const void *element, size_t elem_size,
			    const struct min_heap_callbacks *func, void *args)
{
	void *data = heap->data;
	size_t pos;

	if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap"))
		return false;

	
	pos = heap->nr;
	memcpy(data + (pos * elem_size), element, elem_size);
	heap->nr++;

	
	__min_heap_sift_up_inline(heap, elem_size, pos, func, args);

	return true;
}

#define min_heap_push_inline(_heap, _element, _func, _args)	\
	__min_heap_push_inline(container_of(&(_heap)->nr, min_heap_char, nr), _element,	\
					    __minheap_obj_size(_heap), _func, _args)


static __always_inline
bool __min_heap_del_inline(min_heap_char *heap, size_t elem_size, size_t idx,
			   const struct min_heap_callbacks *func, void *args)
{
	void *data = heap->data;
	void (*swp)(void *lhs, void *rhs, void *args) = func->swp;

	if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap"))
		return false;

	if (!swp)
		swp = select_swap_func(data, elem_size);

	
	heap->nr--;
	if (idx == heap->nr)
		return true;
	do_swap(data + (idx * elem_size), data + (heap->nr * elem_size), elem_size, swp, args);
	__min_heap_sift_up_inline(heap, elem_size, idx, func, args);
	__min_heap_sift_down_inline(heap, idx, elem_size, func, args);

	return true;
}

#define min_heap_del_inline(_heap, _idx, _func, _args)	\
	__min_heap_del_inline(container_of(&(_heap)->nr, min_heap_char, nr),	\
			      __minheap_obj_size(_heap), _idx, _func, _args)

void __min_heap_init(min_heap_char *heap, void *data, size_t size);
void *__min_heap_peek(struct min_heap_char *heap);
bool __min_heap_full(min_heap_char *heap);
void __min_heap_sift_down(min_heap_char *heap, size_t pos, size_t elem_size,
			  const struct min_heap_callbacks *func, void *args);
void __min_heap_sift_up(min_heap_char *heap, size_t elem_size, size_t idx,
			const struct min_heap_callbacks *func, void *args);
void __min_heapify_all(min_heap_char *heap, size_t elem_size,
		       const struct min_heap_callbacks *func, void *args);
bool __min_heap_pop(min_heap_char *heap, size_t elem_size,
		    const struct min_heap_callbacks *func, void *args);
void __min_heap_pop_push(min_heap_char *heap, const void *element, size_t elem_size,
			 const struct min_heap_callbacks *func, void *args);
bool __min_heap_push(min_heap_char *heap, const void *element, size_t elem_size,
		     const struct min_heap_callbacks *func, void *args);
bool __min_heap_del(min_heap_char *heap, size_t elem_size, size_t idx,
		    const struct min_heap_callbacks *func, void *args);

#define min_heap_init(_heap, _data, _size)	\
	__min_heap_init(container_of(&(_heap)->nr, min_heap_char, nr), _data, _size)
#define min_heap_peek(_heap)	\
	(__minheap_cast(_heap) __min_heap_peek(container_of(&(_heap)->nr, min_heap_char, nr)))
#define min_heap_full(_heap)	\
	__min_heap_full(container_of(&(_heap)->nr, min_heap_char, nr))
#define min_heap_sift_down(_heap, _pos, _func, _args)	\
	__min_heap_sift_down(container_of(&(_heap)->nr, min_heap_char, nr), _pos,	\
			     __minheap_obj_size(_heap), _func, _args)
#define min_heap_sift_up(_heap, _idx, _func, _args)	\
	__min_heap_sift_up(container_of(&(_heap)->nr, min_heap_char, nr),	\
			   __minheap_obj_size(_heap), _idx, _func, _args)
#define min_heapify_all(_heap, _func, _args)	\
	__min_heapify_all(container_of(&(_heap)->nr, min_heap_char, nr),	\
			  __minheap_obj_size(_heap), _func, _args)
#define min_heap_pop(_heap, _func, _args)	\
	__min_heap_pop(container_of(&(_heap)->nr, min_heap_char, nr),	\
		       __minheap_obj_size(_heap), _func, _args)
#define min_heap_pop_push(_heap, _element, _func, _args)	\
	__min_heap_pop_push(container_of(&(_heap)->nr, min_heap_char, nr), _element,	\
			    __minheap_obj_size(_heap), _func, _args)
#define min_heap_push(_heap, _element, _func, _args)	\
	__min_heap_push(container_of(&(_heap)->nr, min_heap_char, nr), _element,	\
			__minheap_obj_size(_heap), _func, _args)
#define min_heap_del(_heap, _idx, _func, _args)	\
	__min_heap_del(container_of(&(_heap)->nr, min_heap_char, nr),	\
		       __minheap_obj_size(_heap), _idx, _func, _args)

#endif 
