Implement a Fixed-Size Pool Allocator in Rust

A common pattern in performance-sensitive applications, especially games and embedded systems, is the use of pool allocators. These allocators manage a pre-allocated block of memory divided into fixed-size chunks, allowing for very fast allocation and deallocation by simply managing a free list. This challenge asks you to implement a basic fixed-size pool allocator in Rust.

Problem Description

You need to implement a PoolAllocator<T> struct that can allocate and deallocate objects of a fixed size T. The allocator should manage a contiguous block of memory. When an object is allocated, a chunk of memory from this block is returned. When an object is deallocated, its memory chunk is returned to the allocator's pool.

Key Requirements:

Fixed-Size Chunks: The allocator must deal with objects of a single, fixed size determined at initialization. This size should be sufficient to hold T and potentially some metadata for managing the free list.
Initialization: The allocator must be initialized with a specified capacity, meaning it can hold a certain number of objects concurrently.
Allocation (allocate):
- Should return a mutable reference to an allocated object of type T.
- If the pool is exhausted, it should return None.
- The returned reference must be valid for the lifetime of the allocator.
Deallocation (deallocate):
- Should take a mutable reference to an object previously allocated by this allocator and return it to the pool.
- If the reference provided was not allocated by this allocator, or has already been deallocated, the behavior is undefined (though panic or a safe return is preferred for robustness).
Memory Management: The allocator should manage its own memory internally. It does not rely on the global allocator for individual object allocations.
No T initialization: The allocator does not need to call T::new() or similar. It will be responsible for providing raw memory, and the user will be responsible for constructing T in that memory (e.g., using std::mem::MaybeUninit).

Expected Behavior:

The allocator should be able to allocate up to its capacity.
After deallocating objects, those chunks should be available for subsequent allocations.
Attempting to allocate when the pool is full should result in None.

Edge Cases:

Initializing with zero capacity.
Allocating and deallocating in a non-sequential order.
Attempting to deallocate an object that has already been deallocated or was not allocated by this pool.
What happens if T is larger than the default chunk size? (This should be handled by the size_of::<T>() calculation).

Examples

Example 1: Basic Allocation and Deallocation

// Assume PoolAllocator is implemented
let mut allocator = PoolAllocator::<u32>::new(2); // Capacity of 2 u32s

// Allocate first object
let obj1_ptr = allocator.allocate(); // Returns Some(ptr_to_u32)
// Construct u32 in the allocated memory
unsafe {
    let obj1_ref = &mut *(obj1_ptr.unwrap() as *mut u32);
    *obj1_ref = 10;
    println!("Allocated obj1: {}", *obj1_ref); // Output: Allocated obj1: 10
}


// Allocate second object
let obj2_ptr = allocator.allocate(); // Returns Some(ptr_to_u32)
// Construct u32
unsafe {
    let obj2_ref = &mut *(obj2_ptr.unwrap() as *mut u32);
    *obj2_ref = 20;
    println!("Allocated obj2: {}", *obj2_ref); // Output: Allocated obj2: 20
}

// Attempt to allocate when full
let obj3_ptr = allocator.allocate(); // Returns None
assert!(obj3_ptr.is_none());

// Deallocate first object
// We need a way to get the pointer back for deallocation.
// For simplicity in this example, imagine we stored the pointer.
// In a real scenario, you'd likely return a pointer from allocate
// and accept a pointer to deallocate.
// Let's simulate this by assuming we have the original pointer.
let obj1_raw_ptr = obj1_ptr.unwrap();
allocator.deallocate(obj1_raw_ptr);
println!("Deallocated obj1.");

// Allocate again, should get the chunk back
let obj4_ptr = allocator.allocate(); // Returns Some(ptr_to_u32, which should be obj1's old spot)
unsafe {
    let obj4_ref = &mut *(obj4_ptr.unwrap() as *mut u32);
    *obj4_ref = 30;
    println!("Allocated obj4: {}", *obj4_ref); // Output: Allocated obj4: 30
}

Explanation: The allocator is created with a capacity of 2 u32s. Two u32s are successfully allocated. An attempt to allocate a third u32 fails because the pool is full. When the first u32 is deallocated, its memory chunk becomes available, and a subsequent allocation reuses that chunk.

Example 2: Handling Different Data Types

#[derive(Debug, Clone, Copy, PartialEq)]
struct MyData {
    id: u32,
    value: f32,
}

// Assume PoolAllocator is implemented
let mut allocator = PoolAllocator::<MyData>::new(3); // Capacity of 3 MyData

let data1 = MyData { id: 1, value: 1.1 };
let ptr1 = allocator.allocate().unwrap();
unsafe {
    let ref1 = &mut *(ptr1 as *mut MyData);
    *ref1 = data1;
    println!("Allocated: {:?}", *ref1); // Output: Allocated: MyData { id: 1, value: 1.1 }
}

let data2 = MyData { id: 2, value: 2.2 };
let ptr2 = allocator.allocate().unwrap();
unsafe {
    let ref2 = &mut *(ptr2 as *mut MyData);
    *ref2 = data2;
    println!("Allocated: {:?}", *ref2); // Output: Allocated: MyData { id: 2, value: 2.2 }
}

allocator.deallocate(ptr1);
allocator.deallocate(ptr2);
println!("Deallocated two items.");

// Allocate again, should reuse memory
let ptr3 = allocator.allocate().unwrap();
unsafe {
    let ref3 = &mut *(ptr3 as *mut MyData);
    *ref3 = MyData { id: 3, value: 3.3 };
    println!("Allocated: {:?}", *ref3); // Output: Allocated: MyData { id: 3, value: 3.3 }
}

Explanation: Demonstrates that the allocator works with a custom struct type. It allocates two instances, deallocates them, and then successfully allocates a third, reusing the memory.

Constraints

Capacity: The capacity of the pool allocator (number of elements it can hold) can range from 0 to 1,000,000.
Data Type T: T must be Sized. It can be any type that fits within reasonable memory limits (e.g., less than 1MB per object).
Performance: Allocation and deallocation operations should be O(1) on average.
Memory: The total memory managed by the allocator should not exceed capacity * size_of::<T>() significantly, plus a small overhead for managing the free list.
Safety: Your implementation should aim for safety where possible, but the allocate and deallocate methods will likely require unsafe blocks due to pointer manipulation. The API should clearly indicate where unsafe is necessary.

Notes

Consider how you will represent the "free list" within the allocated memory itself. A common technique is to use the memory of freed chunks to store pointers to the next free chunk.
You'll need to use std::alloc::Layout to determine the size and alignment requirements for your chunks.
std::mem::size_of::<T>() and std::mem::align_of::<T>() will be crucial.
The allocate method should return a raw pointer (*mut u8 or *mut T) which the user will then cast and use unsafe to create a mutable reference.
The deallocate method should accept a raw pointer.
For the free list, you might need to allocate slightly larger chunks than size_of::<T>() to accommodate a pointer if T is small, or use a sentinel value. A common approach is to ensure the chunk size is at least max(size_of::<T>(), size_of::<*mut FreeListNode>()).
Think about the lifetime of the allocated memory. The references returned by allocate should be valid as long as the PoolAllocator itself is alive and the memory hasn't been deallocated. This is a challenging aspect in Rust. For this exercise, returning raw pointers is a simplification. A more advanced version might involve unsafe impl<'a> Drop for a RAII wrapper.
Consider the alignment requirements of T. The allocated chunks must be properly aligned.