/* SPDX-License-Identifier: MIT */
/*
 * Copyright © 2021 Intel Corporation
 */

#ifndef __GPU_BUDDY_H__
#define __GPU_BUDDY_H__

#include <linux/bitops.h>
#include <linux/list.h>
#include <linux/slab.h>
#include <linux/sched.h>
#include <linux/rbtree.h>
#include <linux/rbtree_augmented.h>

/**
 * GPU_BUDDY_RANGE_ALLOCATION - Allocate within a specific address range
 *
 * When set, allocation is restricted to the range [start, end) specified
 * in gpu_buddy_alloc_blocks(). Without this flag, start/end are ignored
 * and allocation can use any free space.
 */
#define GPU_BUDDY_RANGE_ALLOCATION		BIT(0)

/**
 * GPU_BUDDY_TOPDOWN_ALLOCATION - Allocate from top of address space
 *
 * Allocate starting from high addresses and working down. Useful for
 * separating different allocation types (e.g., kernel vs userspace)
 * to reduce fragmentation.
 */
#define GPU_BUDDY_TOPDOWN_ALLOCATION		BIT(1)

/**
 * GPU_BUDDY_CONTIGUOUS_ALLOCATION - Require physically contiguous blocks
 *
 * The allocation must be satisfied with a single contiguous block.
 * If the requested size cannot be allocated contiguously, the
 * allocation fails with -ENOSPC.
 */
#define GPU_BUDDY_CONTIGUOUS_ALLOCATION		BIT(2)

/**
 * GPU_BUDDY_CLEAR_ALLOCATION - Prefer pre-cleared (zeroed) memory
 *
 * Attempt to allocate outside dirty-tracked ranges first. If insufficient
 * clear memory is available, falls back to dirty memory. Useful when the
 * caller needs zeroed memory and wants to avoid GPU clear operations.
 */
#define GPU_BUDDY_CLEAR_ALLOCATION		BIT(3)

/**
 * GPU_BUDDY_CLEARED - Mark returned blocks as cleared
 *
 * Used with gpu_buddy_free_list() to indicate that the memory being
 * freed has been cleared (zeroed). The blocks will be removed from the
 * dirty tracker for future GPU_BUDDY_CLEAR_ALLOCATION requests.
 */
#define GPU_BUDDY_CLEARED			BIT(4)

/**
 * GPU_BUDDY_TRIM_DISABLE - Disable automatic block trimming
 *
 * By default, if an allocation is smaller than the allocated block,
 * excess memory is trimmed and returned to the free pool. This flag
 * disables trimming, keeping the full power-of-two block size.
 */
#define GPU_BUDDY_TRIM_DISABLE			BIT(5)

/*
 * Clear/dirty state of a free block. Ordered so a numerically larger value
 * is "more clear" (DIRTY < MIXED < CLEAR) which lets subtree_block_state be
 * maintained as a simple max-augment over the per-order free tree.
 */
enum gpu_block_state {
	GPU_BLOCK_DIRTY = 0,
	GPU_BLOCK_MIXED = 1,
	GPU_BLOCK_CLEAR = 2,
};

/**
 * struct gpu_buddy_block - Block within a buddy allocator
 *
 * Each block in the buddy allocator is represented by this structure.
 * Blocks are organized in a binary tree where each parent block can be
 * split into two children (left and right buddies). The allocator manages
 * blocks at various orders (power-of-2 sizes) from chunk_size up to the
 * largest contiguous region.
 *
 * @private: Private data owned by the allocator user (e.g., driver-specific data)
 * @link: List node for user ownership while block is allocated
 */
struct gpu_buddy_block {
/* private: */
	/*
	 * Header bit layout:
	 * - Bits 63:12: block offset within the address space
	 * - Bits 11:10: state (ALLOCATED, FREE, or SPLIT)
	 * - Bit 9: clear bit (1 if memory is zeroed)
	 * - Bits 8:6: reserved
	 * - Bits 5:0: order (log2 of size relative to chunk_size)
	 */
#define GPU_BUDDY_HEADER_OFFSET GENMASK_ULL(63, 12)
#define GPU_BUDDY_HEADER_STATE  GENMASK_ULL(11, 10)
#define   GPU_BUDDY_ALLOCATED	   (1 << 10)
#define   GPU_BUDDY_FREE	   (2 << 10)
#define   GPU_BUDDY_SPLIT	   (3 << 10)
/*
 * GPU_BUDDY_HEADER_CLEAR has two roles:
 *  - FREE state:      set when the block's full range is cleared (dirty
 *                     tracker confirmed no overlap).
 *  - ALLOCATED state: set when the block was served from cleared memory,
 *                     informing the caller that no GPU clear pass is needed.
 */
#define GPU_BUDDY_HEADER_CLEAR  GENMASK_ULL(9, 9)
/* Free to be used, if needed in the future */
#define GPU_BUDDY_HEADER_UNUSED GENMASK_ULL(8, 6)
#define GPU_BUDDY_HEADER_ORDER  GENMASK_ULL(5, 0)
	u64 header;

	struct gpu_buddy_block *left;
	struct gpu_buddy_block *right;
	struct gpu_buddy_block *parent;
/* public: */
	void *private; /* owned by creator */

	/*
	 * While the block is allocated by the user through gpu_buddy_alloc*,
	 * the user has ownership of the link, for example to maintain within
	 * a list, if so desired. As soon as the block is freed with
	 * gpu_buddy_free* ownership is given back to the mm.
	 */
	union {
/* private: */
		struct rb_node rb;
/* public: */
		struct list_head link;
	};
/* private: */
	enum gpu_block_state subtree_block_state;
	unsigned int subtree_max_alignment;
	struct list_head tmp_link;
	bool has_clear;
};

/* Order-zero must be at least SZ_4K */
#define GPU_BUDDY_MAX_ORDER (63 - 12)

/**
 * struct gpu_dirty_extent - a contiguous dirty address range
 *
 * Tracks a single contiguous address range whose memory content is known
 * to be dirty.  Extents are non-overlapping and stored in an augmented
 * red-black tree sorted by @start.  The augmented value @subtree_max_size
 * allows O(log N) search for an extent of at least a given size.
 */
struct gpu_dirty_extent {
/* private: */
	struct rb_node	rb;
	u64		start;
	u64		end;
	u64		subtree_max_size;
};

/**
 * struct gpu_dirty_tracker - tracks dirty address intervals
 *
 * Maintains a set of non-overlapping dirty extents as an augmented
 * red-black tree.
 */
struct gpu_dirty_tracker {
/* private: */
	struct rb_root	root;
	/* Total bytes of dirty memory currently tracked. */
	u64		total_dirty;
};

/**
 * struct gpu_buddy - GPU binary buddy allocator
 *
 * The buddy allocator provides efficient power-of-two memory allocation
 * with fast allocation and free operations. It is commonly used for GPU
 * memory management where allocations can be split into power-of-two
 * block sizes.
 *
 * Locking should be handled by the user; a simple mutex around
 * gpu_buddy_alloc_blocks() and gpu_buddy_free_block()/gpu_buddy_free_list()
 * should suffice.
 *
 * @n_roots: Number of root blocks in the roots array.
 * @max_order: Maximum block order (log2 of largest block size / chunk_size).
 * @chunk_size: Minimum allocation granularity in bytes. Must be at least SZ_4K.
 * @size: Total size of the address space managed by this allocator in bytes.
 * @avail: Total free space currently available for allocation in bytes.
 * @lock_dep_map: Annotates gpu_buddy API with a driver provided lock.
 */
struct gpu_buddy {
/* private: */
	/* Tracker of dirty address ranges (decoupled from free_tree). */
	struct gpu_dirty_tracker dirty;
	/*
	 * One RB-tree per order containing all free blocks (clear and
	 * dirty alike).  The augment field subtree_block_state (a max over
	 * the subtree of each block's state) lets clear allocations
	 * find the right-most fully-clear or mixed block in O(log N).
	 * Dirty free blocks coexist here but are also indexed by the
	 * @dirty tracker for fast dirty allocation lookups.
	 */
	struct rb_root *free_tree;
	/*
	 * Array of root blocks representing the top-level blocks of the
	 * binary tree(s). Multiple roots exist when the total size is not
	 * a power of two, with each root being the largest power-of-two
	 * that fits in the remaining space.
	 */
	struct gpu_buddy_block **roots;
	/*
	 * Per-order free block scoreboard: free_scoreboard[order] holds the
	 * number of blocks of that order currently in the free state.
	 * Incremented in mark_free(), decremented wherever rbtree_remove() is
	 * called on a free block.
	 */
	u64 *free_scoreboard;
	/*
	 * Per-order used block scoreboard: used_scoreboard[order] holds the
	 * number of blocks of that order currently in the allocated state.
	 * Incremented in mark_allocated(), decremented in mark_free() (guarded
	 * by gpu_buddy_block_is_allocated()) and in __gpu_buddy_free() when an
	 * allocated block is consumed directly during buddy coalescing.
	 */
	u64 *used_scoreboard;
/* public: */
	unsigned int n_roots;
	unsigned int max_order;
	u64 chunk_size;
	u64 size;
	u64 avail;
#ifdef CONFIG_LOCKDEP
	struct lockdep_map *lock_dep_map;
#endif
};

#ifdef CONFIG_LOCKDEP
/**
 * gpu_buddy_driver_set_lock() - Set the lock protecting accesses to GPU BUDDY
 * @mm: Pointer to GPU buddy structure.
 * @lock: the lock used to protect the gpu buddy. The locking primitive
 * must contain a dep_map field.
 *
 * Call this to annotate gpu_buddy APIs which access/modify gpu_buddy manager
 */
#define gpu_buddy_driver_set_lock(mm, lock) \
	do { \
		struct gpu_buddy *__mm = (mm); \
		if (!WARN(__mm->lock_dep_map, "GPU BUDDY MM lock should be set only once.")) \
			__mm->lock_dep_map = &(lock)->dep_map; \
	} while (0)
#else
#define gpu_buddy_driver_set_lock(mm, lock) do { (void)(mm); (void)(lock); } while (0)
#endif

#ifdef CONFIG_LOCKDEP
/**
 * gpu_buddy_driver_lock_held() - Assert GPU BUDDY manager lock is held
 * @mm: Pointer to the GPU BUDDY structure.
 *
 * Ensure driver lock is held.
 */
static inline void gpu_buddy_driver_lock_held(const struct gpu_buddy *mm)
{
	if (mm->lock_dep_map)
		lockdep_assert(lock_is_held_type(mm->lock_dep_map, 0));
}
#else
static inline void gpu_buddy_driver_lock_held(const struct gpu_buddy *mm)
{
}
#endif

/**
 * gpu_buddy_clear_avail - free space that is clear (zeroed), in bytes
 * @mm: gpu buddy allocator
 *
 * A subset of @mm->avail. Derived on demand as @mm->avail minus the bytes
 * the dirty tracker records as dirty, so it is always consistent with the
 * tracker without a cached field to keep in sync. Zero for a fresh pool,
 * which is fully dirty.
 */
static inline u64 gpu_buddy_clear_avail(const struct gpu_buddy *mm)
{
	gpu_buddy_driver_lock_held(mm);
	return mm->avail - mm->dirty.total_dirty;
}

static inline u64
gpu_buddy_block_offset(const struct gpu_buddy_block *block)
{
	return block->header & GPU_BUDDY_HEADER_OFFSET;
}

static inline unsigned int
gpu_buddy_block_order(struct gpu_buddy_block *block)
{
	return block->header & GPU_BUDDY_HEADER_ORDER;
}

static inline bool
gpu_buddy_block_is_free(struct gpu_buddy_block *block)
{
	return (block->header & GPU_BUDDY_HEADER_STATE) == GPU_BUDDY_FREE;
}

static inline bool
gpu_buddy_block_is_clear(struct gpu_buddy_block *block)
{
	return block->header & GPU_BUDDY_HEADER_CLEAR;
}

static inline u64
gpu_buddy_block_size(struct gpu_buddy *mm,
		     struct gpu_buddy_block *block)
{
	return mm->chunk_size << gpu_buddy_block_order(block);
}

int gpu_buddy_init(struct gpu_buddy *mm, u64 size, u64 chunk_size);

void gpu_buddy_fini(struct gpu_buddy *mm);

int gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
			   u64 start, u64 end, u64 size,
			   u64 min_page_size,
			   struct list_head *blocks,
			   unsigned long flags);

int gpu_buddy_block_trim(struct gpu_buddy *mm,
			 u64 *start,
			 u64 new_size,
			 struct list_head *blocks);

void gpu_buddy_reset_clear(struct gpu_buddy *mm, bool is_clear);

void gpu_buddy_free_block(struct gpu_buddy *mm, struct gpu_buddy_block *block);

struct gpu_buddy_block *gpu_buddy_allocated_addr_to_block(struct gpu_buddy *mm, u64 addr);

void gpu_buddy_free_list(struct gpu_buddy *mm,
			 struct list_head *objects,
			 unsigned int flags);

void gpu_buddy_print(struct gpu_buddy *mm);
void gpu_buddy_block_print(struct gpu_buddy *mm,
			   struct gpu_buddy_block *block);
#endif