/* Copyright (C) 2021,2022 fef <owo@fef.moe>.  All rights reserved. */

#pragma once

#include <arch/atom.h>
#include <arch/page.h>

#include <gay/cdefs.h>
#include <gay/clist.h>
#include <gay/config.h>
#include <gay/systm.h>
#include <gay/types.h>

/*
 * I'm trying really hard to keep the size of struct vm_page a power of two
 * on LP64 systems, because that way we can quickly get to the page frame number
 * by shifting the byte offset of the vm_page_t in vm_page_array to the right
 * rather than doing a costly divide instruction (or store the page frame number
 * within the structure itself, which takes up precious space).
 *
 * There is insane pressure on the size of this structure, because a typical
 * system will have millions of instances of it.  Every additional byte makes
 * a significant difference in memory management overhead.
 */

union vm_page_attr {
	int _val;
	struct {
		unsigned order:8;	/**< @brief Index into `mm_zones[zone].pools` */
		bool lock:1;		/**< @brief Page is locked */
		bool rsvd:1;		/**< @brief Page is reserved and must never be touched */
		bool pcpu:1;		/**< @brief Page is in a per-cpu cache */
		bool slab:1;		/**< @brief Page is used by the slab allocator */
		unsigned zone:2;	/**< @brief Index into `mm_zones` */
	};
};
#define _PGA_ORDER_SHIFT	0
#define _PGA_ORDER_MASK		(0xf << _PGA_ORDER_SHIFT)
#define _PGA_LOCK_SHIFT		8
#define _PGA_LOCK_MASK		(1 << _PGA_LOCK_SHIFT)
#define _PGA_RSVD_SHIFT		9
#define _PGA_RSVD_MASK		(1 << _PGA_RSVD_SHIFT)
#define _PGA_PCPU_SHIFT		10
#define _PGA_PCPU_MASK		(1 << _PGA_PCPU_SHIFT)
#define _PGA_SLAB_SHIFT		11
#define _PGA_SLAB_MASK		(1 << _PGA_SLAB_SHIFT)
#define _PGA_ZONE_SHIFT		12
#define _PGA_ZONE_MASK		(3 << _PGA_ZONE_SHIFT)

typedef union vm_page_attr vm_page_attr_t;

/**
 * @brief Stores information about a single page in physical memory.
 * There is exactly one of these for every physical page, no matter what that
 * page is used for or whether it is usable at all.
 */
struct vm_page {
	/** @brief Reference count (0 = unused, < 0 = locked) */
	atom_t count;
	/** @brief Page attributes, use the macros below to access this */
	atom_t attr;
	/** @brief If the page is free, this is its freelist. */
	struct clist link;
	/**
	 * @brief Optional extra data pointer, reserved for private use.
	 * The current owner of the page may use this to track the underlying
	 * object in memory (or pretty much anything else), for example the
	 * `struct slab` if this page is currently used by the slab allocator.
	 */
	void *extra;
};

typedef struct vm_page *vm_page_t;

/** @brief Array of every single page in physical memory, indexed by page frame number. */
extern struct vm_page *const vm_page_array;

#if CFG_DEBUG_PGADDRS
extern vm_page_t _vm_page_array_end;
#define PGADDR_ASSERT(x) KASSERT(x)
#else
#define PGADDR_ASSERT(x) ({})
#endif

static inline u8 pga_order(vm_page_t page)
{
	union vm_page_attr attr = { ._val = atom_read(&page->attr) };
	return attr.order;
}

static inline bool pga_rsvd(vm_page_t page)
{
	union vm_page_attr attr = { ._val = atom_read(&page->attr) };
	return attr.rsvd;
}

static inline bool pga_pcpu(vm_page_t page)
{
	union vm_page_attr attr = { ._val = atom_read(&page->attr) };
	return attr.pcpu;
}

static inline bool pga_slab(vm_page_t page)
{
	union vm_page_attr attr = { ._val = atom_read(&page->attr) };
	return attr.slab;
}

static inline enum mm_zone_type pga_zone(vm_page_t page)
{
	union vm_page_attr attr = { ._val = atom_read(&page->attr) };
	return attr.zone;
}

static inline u8 pga_set_order(vm_page_t page, u8 order)
{
	spin_loop {
		union vm_page_attr old = { ._val = atom_read(&page->attr) };
		union vm_page_attr new = old;
		new.order = order;
		if (atom_cmp_xchg(&page->attr, old._val, new._val) == old._val)
			return old.order;
	}
}

static inline bool pga_set_pcpu(vm_page_t page, bool pcpu)
{
	if (pcpu)
		return atom_set_bit(&page->attr, _PGA_PCPU_SHIFT);
	else
		return atom_clr_bit(&page->attr, _PGA_PCPU_SHIFT);
}

static inline bool pga_set_slab(vm_page_t page, bool slab)
{
	if (slab)
		return atom_set_bit(&page->attr, _PGA_SLAB_SHIFT);
	else
		return atom_clr_bit(&page->attr, _PGA_SLAB_SHIFT);
}

static inline enum mm_zone_type pga_set_zone(vm_page_t page, enum mm_zone_type zone)
{
	spin_loop {
		union vm_page_attr old = { ._val = atom_read(&page->attr) };
		union vm_page_attr new = old;
		new.zone = zone;
		if (atom_cmp_xchg(&page->attr, old._val, new._val) == old._val)
			return old.zone;
	}
}

static __always_inline bool page_get(vm_page_t page)
{
	return atom_inc(&page->count);
}

static __always_inline bool page_put(vm_page_t page)
{
	return atom_dec(&page->count);
}

/* XXX we should probably use a wait queue for these rather than a spinlock like thing */

static inline void page_lock(vm_page_t page)
{
	spin_loop {
		if (atom_set_bit(&page->attr, _PGA_LOCK_SHIFT))
			break;
	}
}

static __always_inline void page_unlock(vm_page_t page)
{
	atom_clr_bit(&page->attr, _PGA_LOCK_SHIFT);
}

static __always_inline bool page_trylock(vm_page_t page)
{
	return atom_set_bit(&page->attr, _PGA_LOCK_SHIFT);
}

static inline void __page_set_flag(vm_page_t page, unsigned flag)
{
	atom_or(&page->attr, (int)flag);
}

static inline void __page_clr_flag(vm_page_t page, unsigned mask)
{
	atom_and(&page->attr, (int)~mask);
}

static __always_inline void page_attr_load(vm_page_attr_t *attr, vm_page_t page)
{
	attr->_val = atom_read(&page->attr);
}

static __always_inline void page_attr_copy(vm_page_attr_t *dest, const vm_page_attr_t *src)
{
	dest->_val = src->_val;
}

static __always_inline bool page_attr_cmp_xchg(vm_page_t page, const vm_page_attr_t *cmp,
					       const vm_page_attr_t *val)
{
	return atom_cmp_xchg(&page->attr, cmp->_val, val->_val);
}

/** @brief Get the page frame number of a page. */
__pure2
static inline u_long pg2pfn(vm_page_t page)
{
	PGADDR_ASSERT(page < _vm_page_array_end);
	return page - vm_page_array;
}

/**
 * @brief Get the page that a virtual address points to.
 * The address must point to the DMAP region (i.e. an address that is returned
 * by either `get_pages()` and friends, or `kmalloc()` and friends).
 */
__pure2
static inline vm_page_t vaddr2pg(void *vaddr)
{
	PGADDR_ASSERT(vaddr >= DMAP_START && vaddr < (void *)_vm_page_array_end);
	uintptr_t offset = (uintptr_t)vaddr - DMAP_OFFSET;
	return &vm_page_array[offset >> PAGE_SHIFT];
}

/**
 * @brief Get the page frame number for a virtual address.
 * The address must point to the DMAP region (i.e. an address that is returned
 * by either `get_pages()` and friends, or `kmalloc()` and friends).
 */
__pure2
static inline u_long vaddr2pfn(void *vaddr)
{
	u_long pfn = ((uintptr_t)vaddr - DMAP_OFFSET) >> PAGE_SHIFT;
	PGADDR_ASSERT(vaddr >= DMAP_START && &vm_page_array[pfn] < _vm_page_array_end);
	return pfn;
}

/** @brief Get the page frame number for a physical address. */
__pure2
static inline u_long paddr2pfn(vm_paddr_t paddr)
{
	PGADDR_ASSERT(&vm_page_array[paddr >> PAGE_SHIFT] < _vm_page_array_end);
	return paddr >> PAGE_SHIFT;
}

/** @brief Get the page that a physical address belongs to. */
__pure2
static inline vm_page_t paddr2pg(vm_paddr_t paddr)
{
	vm_page_t page = vm_page_array + (paddr >> PAGE_SHIFT);
	PGADDR_ASSERT(page < _vm_page_array_end);
	return page;
}

/**
 * @brief Translate a page frame number to its corresponding virtual address
 * in the DMAP region.
 */
__pure2
static inline void *pfn2vaddr(u_long pfn)
{
	PGADDR_ASSERT(&vm_page_array[pfn] < _vm_page_array_end);
	return DMAP_START + (pfn << PAGE_SHIFT);
}

/*
 * We have to be careful in this macro, because only the first page in the
 * order group has the correct order set.  So we can only read it once at
 * the beginning of the loop, since the page pointer is being updated.
 */

/**
 * @brief Iterate over every page in its order group.
 *
 * @param page The first `vm_page_t` in the group.
 */
#define vm_page_foreach_in_order(page)		\
	for (int __i = 1 << pga_order(page);	\
	     __i >= 0;				\
	     __i = ({ ++(page); --__i; }))