/*
 * 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 <ptdef.h>

/* Maximum levels per region */
#define POOL_LEVEL_MAX 8

/* Minimum log2 granularity */
#define POOL_MIN_LOG2  3

/* Compute the granularity at a specific level */
#define POOL_LEVEL_GRAN(X) \
    (1 << ((X) + POOL_MIN_LOG2))

/*
 * Valid types that can be assigned to an allocated region
 *
 * @POOL_NON_PAGED:  This type is not paged to disk
 */
typedef enum {
    POOL_NON_PAGED
} POOL_TYPE;

/*
 * A pool page represents a single allocatable
 * unit within a pool segment.
 *
 * @Bitmap: Bitmap within pool page [0=free, 1=allocated]
 * @Next:   Next page within list
 */
typedef struct _POOL_PAGE {
    UQUAD Bitmap;
    struct _POOL_PAGE *Next;
} POOL_PAGE;

/*
 * A pool segment represents a list of pages at a specific
 * granularity.
 *
 * @FirstPage:  Last page in segment
 * @PageCount: Number of pages in this segment
 */
typedef struct {
    POOL_PAGE *FirstPage;
    USIZE PageCount;
} POOL_SEGMENT;

/*
 * A pool region represents a whole area in which several
 * pool levels can reside.
 *
 * @Level: Segment levels
 */
typedef struct {
    POOL_SEGMENT Level[POOL_LEVEL_MAX];
} POOL_REGION;

/*
 * A pool header is prepended to the start of allocated
 * memory.
 *
 * @Gran:       Granularity of page bitmap
 * @BitBase:    Allocation bitbase
 * @SizeBytes:  Allocation size in bytes
 * @Tag:        Tag assigned to allocation
 * @Page:       Page we reside in
 */
typedef struct {
    USHORT Gran;
    UCHAR BitBase : 6;
    USIZE SizeBytes;
    ULONG Tag;
    POOL_PAGE *Page;
} POOL_HEADER;

/*
 * Initialize a pool region
 *
 * @Region: Pool region to initialize
 */
VOID ExPoolRegionInit(POOL_REGION *Region);

/*
 * Free a pool with an associated tag
 *
 * @Ptr:  Base of memory to free
 * @Tag:  Tag associated with allocated memory
 */
VOID ExFreePoolWithTag(VOID *Ptr, ULONG Tag);

/*
 * Allocate a memory pool of a specific type and assign a tag
 * to it for verifying frees among other debugging needs.
 *
 * @Type:       Pool type to assign
 * @ByteCount:  Number of bytes to allocate
 * @Tag:        Tag to assign to allocated memory
 *
 * Returns NULL upon failure.
 */
VOID *ExAllocatePoolWithTag(
    POOL_TYPE Type, USIZE ByteCount,
    ULONG Tag
);

#endif  /* !_EX_POOL_H_ */
