/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Pool allocator
 * Author:      Chloe M.
 */

#ifndef _EX_POOL_H_
#define _EX_POOL_H_ 1

#include <hal/page.h>
#include <stapi/status.h>
#include <ke/spinlock.h>
#include <stdef.h>

/* Minimum log2 allocation size */
#define POOL_MIN_LOG2 3

/* Compute the granularity of a specific level */
#define LEVEL_GRAN(LEVEL) \
    (1 << (POOL_MIN_LOG2 + (LEVEL)))

/* Maximum levels per pool */
#define LEVEL_COUNT 8

/*
 * When allocating pages, we sacrifice a little bit of its
 * first bytes to store some metadata. This is the true size
 * of unreserved / usable data.
 */
#define POOL_REAL_PAGESZ \
    (PAGESIZE - sizeof(MEMORY_PAGE))

/* Used to prevent corruption */
#define POOL_MAGIC 0xCA7F00D

/*
 * Valid pool types
 *
 * @POOL_NON_PAGED: Not pagable to disk
 */
typedef enum {
    POOL_NON_PAGED
} POOL_TYPE;

/*
 * Represents the header that goes before each returned
 * and allocated block.
 *
 * @Magic:      Magic numbers
 * @MemoryTag:  Tag assigned at allocation
 * @Length:     Length of memory allocation
 */
typedef struct PACKED {
    ULONG Magic;
    ULONG MemoryTag;
    USIZE Length;
} MEMORY_HEADER;

/*
 * Page data structure used internally to link together page
 * lists
 *
 * @Bitmap: Bitmap to data [0 : free, 1 : allocated]
 * @Next:   Next page in list
 */
typedef struct _MEMORY_PAGE {
    UQUAD Bitmap;
    struct _MEMORY_PAGE *Next;
} MEMORY_PAGE;

/*
 * A memory block holds a list of memory chunks of a specific
 * granularity.
 *
 * @FirstPage:  First page in list
 * @LastPage:   Last page in list
 * @PageCount:  Number of pages total
 */
typedef struct {
    MEMORY_PAGE *FirstPage;
    MEMORY_PAGE *LastPage;
    USIZE PageCount;
} MEMORY_BLOCK;

/*
 * A memory pool holds blocks of varying granularities
 *
 * @BlockLevels: Block levels [granularity : GRAN(level)]
 * @Id:          Pool ID
 * @Lock:        Lock protecting this pool
 */
typedef struct {
    MEMORY_BLOCK BlockLevels[LEVEL_COUNT];
    USIZE Id;
    SPINLOCK Lock;
} MEMORY_POOL;

/*
 * Initialize a memory pool
 *
 * @Pool: Memory pool to initialize
 */
ST_STATUS ExPoolInit(MEMORY_POOL *Pool);

/*
 * Allocate a tagged pool
 *
 * @Type:       Pool type
 * @ByteCount:  Number of bytes to allocate
 * @Tag:        Tag to assign to pool
 *
 * Returns the memory base on success, otherwise
 * NULL on failure.
 */
VOID *ExAllocatePoolWithTag(
    POOL_TYPE Type, USIZE ByteCount,
    ULONG Tag
);

#endif  /* !_EX_POOL_H_ */
