Can rust guarantee I free an object with the right object pool?

Question

Say I have defined my own object pool struct. Internally it keeps a Vec of all the objects and some data structure that lets it know which items in the vector are currently handed out and which are free. It has an allocate method that returns an index of an unused item in the vector, and a free method to tell the pool at an index in the vector is available to be used again.

Is it possible for me to define the API of my object pool in such a way that the type system and borrow checker will guarantee that I free an object back to the correct pool? This is assuming a situation where I might have multiple instances of the pool that are all the same type. It seems to me that for the regular global allocator rust doesn't have to worry about this problem because there is only one global allocator.

usage example:

fn foo() {
    let new_obj1 = global_pool1.allocate();
    let new_obj2 = global_pool2.allocate();

    // do stuff with objects

    global_pool1.free(new_obj2); // oops!, giving back to the wrong pool
    global_pool2.free(new_obj1); // oops!, giving back to the wrong pool
}

Answer 1

You can use a Zero Sized Type (ZST, for short) to get the API you want, without the overhead of another pointer.

Here is an implementation for 2 pools, which can be extended to support any number of pools using a macro to generate the "marker" struct (P1, P2, etc). A major downside is that forgetting to free using the pool will "leak" the memory.

This Ferrous Systems blog post has a number of possible improvements that might interest you, especially if you statically allocate the pools, and they also have a number of tricks for playing with the visibility of P1 so that users wont be able to misuse the API.

use std::marker::PhantomData;
use std::{cell::RefCell, mem::size_of};

struct Index<D>(usize, PhantomData<D>);
struct Pool<D> {
    data: Vec<[u8; 4]>,
    free_list: RefCell<Vec<bool>>,
    marker: PhantomData<D>,
}

impl<D> Pool<D> {
    fn new() -> Pool<D> {
        Pool {
            data: vec![[0,0,0,0]],
            free_list: vec![true].into(),
            marker: PhantomData::default(),
        }
    }
    
    fn allocate(&self) -> Index<D> {
        self.free_list.borrow_mut()[0] = false;
        
        Index(0, self.marker)
    }
    
    fn free<'a>(&self, item: Index<D>) {
        self.free_list.borrow_mut()[item.0] = true;
    }
}

struct P1;
fn create_pool1() -> Pool<P1> {
    assert_eq!(size_of::<Index<P1>>(), size_of::<usize>());
    Pool::new()
}

struct P2;
fn create_pool2() -> Pool<P2> {
    Pool::new()
}


fn main() {
    
    let global_pool1 = create_pool1();
    let global_pool2 = create_pool2();
    
    let new_obj1 = global_pool1.allocate();
    let new_obj2 = global_pool2.allocate();

    // do stuff with objects

    global_pool1.free(new_obj1);
    global_pool2.free(new_obj2);

    global_pool1.free(new_obj2); // oops!, giving back to the wrong pool
    global_pool2.free(new_obj1); // oops!, giving back to the wrong pool
}

Trying to free using the wrong pool results in:

error[E0308]: mismatched types
  --> zpool\src\main.rs:57:23
   |
57 |     global_pool1.free(new_obj2); // oops!, giving back to the wrong pool
   |                       ^^^^^^^^ expected struct `P1`, found struct `P2`
   |
   = note: expected struct `Index<P1>`
              found struct `Index<P2>`

Link to playground

This can be improved a little so that the borrow checker will enforce that Index will not out-live the Pool, using:

fn allocate(&self) -> Index<&'_ D> {
    self.free_list.borrow_mut()[0] = false;
        
    Index(0, Default::default())
}

So you get this error if the pool is dropped while an Index is alive:

error[E0505]: cannot move out of `global_pool1` because it is borrowed
  --> zpool\src\main.rs:54:10
   |
49 |     let new_obj1 = global_pool1.allocate();
   |                    ------------ borrow of `global_pool1` occurs here
...
54 |     drop(global_pool1);
   |          ^^^^^^^^^^^^ move out of `global_pool1` occurs here
...
58 |     println!("{}", new_obj1.0);
   |                    ---------- borrow later used here

Link to playground

Also, a link to playground with Item API (returning an Item, vs only and Index)