Rust generic value memory store implementation

Question

I am working on an API for an object store, and I am struggling to figure out a way to implement it in Rust. It is a bit frustrating, since I have a good idea of how this could be done in C++, so maybe the design is fundamentally not suitable for Rust, but I hope people here can provide some useful insights.

I would like to be able to do the following:

// This is clone & copy
struct Id<Identifiable> {
  // What I think would be needed to get this entire system 
  // to work, but it is not fixed like this so alternate 
  // ideas are welcome.
  raw_id: usize,
  _marker: PhantomData<Identifiable>,
}

struct Store {
  // ... don't know how to implement this
}

trait Object {
  // ... some behavior
}

struct ObjectA {}
impl Object for ObjectA {}

struct ObjectB {}
impl Object for ObjectB {}

impl Store {
  pub fn create_object_a(&mut self) -> Id<ObjectA> { todo!() }
  pub fn create_object_b(&mut self) -> Id<ObjectB> { todo!() }

  pub fn get<'a, Obj: Object>(&self, id: Id<Obj>) -> &'a Obj { todo!() }
}

This would be used as follows:

let store = Store::new();
let handle_a = store.create_object_a();
let handle_b = store.create_object_b();

let obj_a = store.get(handle_a); // obj_a is of type &ObjectA
let obj_b = store.get(handle_b); // obj_b is of type &ObjectB

Since the concrete types that can be put into the store are known statically (they can only be in the store if they are constructed through a create_something() method), I feel like there should be enough information in the type system to be able to do this. One thing I desperately want to avoid is something like Vec<Box<dyn Any>>, because that introduces extra indirection.

I realize that this is likely not possible in safe Rust, although I feel like it should be possible using unsafe Rust. My thought is that it is somewhat similar to how the Bevy ECS implementation stores components (components of the same type are stored contiguously in memory, a property I would like to see here as well), although I am struggling to understand exactly how that works.

Hopefully someone has ideas on how to implement this, or has an alternate design which would suite Rust better. Thank you!

Answer 1

You can create a trait generic over Obj that specifies how to retrieve arenas (Vecs) of Obj from the store, and implement it as applied to ObjectA and ObjectB. Then Store::get() uses this implementation to retrieve the values.

// Your store type:
struct Store {
    a_store: Vec<ObjectA>,
    b_store: Vec<ObjectB>,
}

// Trait family that specifies how to obtain slice of T's from Store
trait StoreGet<T> {
    fn raw_storage(&self) -> &[T];
}

// Implementation of the trait applied to ObjectA and ObjectB on Store:
impl StoreGet<ObjectA> for Store {
    fn raw_storage(&self) -> &[ObjectA] {
        &self.a_store
    }
}
impl StoreGet<ObjectB> for Store {
    fn raw_storage(&self) -> &[ObjectB] {
        &self.b_store
    }
}

With that in place, Store would be implemented as follows:

impl Store {
    pub fn create_object_a(&mut self) -> Id<ObjectA> {
        let raw_id = self.a_store.len();
        self.a_store.push(ObjectA {});
        Id {
            raw_id,
            _marker: PhantomData,
        }
    }
    pub fn create_object_b(&mut self) -> Id<ObjectB> {
        // ... essentially the same as create_object_a ...
    }

    pub fn get<Obj>(&self, id: Id<Obj>) -> &Obj
    where
        Obj: Object,
        Self: StoreGet<Obj>,
    {
        let slice = self.raw_storage();
        &slice[id.raw_id]
    }
}

Playground

I recommend looking into an arena crate rather than implementing your own arena (though it's not a tragedy if you do implement it yourself - it's just generally a good idea to reuse).

Answer 2

The key question is, is your store intended to store objects of a fixed number of types known at compile time, or is the same store object supposed to store objects of any type?

Fixed number of types

This design is possible statically, and user4815162342's answer shows a good way of doing it.

Storing objects of any type, dynamically

This design is fundamentally dynamic and not typesafe. When accessing an element of the store, you are accessing memory that could be of any type; the only way we know that it's of a particular type is the mapping between objects T and identifiers Id<T>, which is not a mapping known to the compiler. You're essentially implementing a dynamic type system since you're maintaining the types of objects as part of the store.

Dynamic doesn't mean impossible, though; as a rule of thumb, what's possible in C++ is always possible in Rust -- typically by turning off the compiler by using unsafe code or code that can panic at runtime.

Illustrative implementation with Box<dyn Any>

As you point out, perhaps we don't want to use Box<dyn Any> because it's a small runtime overhead, but it's easiest to see how to accomplish this design first with Box<dyn Any> to abstract a raw pointer-to-anything and using .downcast_ref() to cast these pointers to concrete types. Here is how that works. I've replaced your Object trait with Default as that was the only required functionality in the below code. Running the code shows that _obj_a and _obj_b have the correct type and that the code does not panic (showing that our dynamic type management is working).

use std::any::Any;
use std::collections::HashMap;
use std::marker::PhantomData;

pub struct Id<T> {
  raw_id: usize,
  _marker: PhantomData<T>,
}

// To implement Store we use Box<dyn Any>. If it is truly crucial
// to avoid the unwrap call at runtime, a pointer could be used instead
// and downcasted dynamically, but this is fundamentally not typesafe.
pub struct Store {
    next_id: usize,
    stash: HashMap<usize, Box<dyn Any>>
}

#[derive(Default)]
pub struct ObjectA {}

#[derive(Default)]
pub struct ObjectB {}

impl Store {
    pub fn new() -> Self {
        Self { next_id: 0, stash: HashMap::new() }
    }
    fn new_id(&mut self) -> usize {
        let id = self.next_id;
        self.next_id += 1;
        id
    }
    pub fn create<Obj: Default + 'static>(&mut self) -> Id<Obj> {
        let id = self.new_id();
        self.stash.insert(id, Box::new(Obj::default()));
        Id {raw_id: id, _marker: PhantomData }
    }
    pub fn get<'a, Obj: 'static>(&'a self, id: Id<Obj>) -> &'a Obj {
        // Note that because we manage Id<Obj>s, the following unwraps
        // should never panic in reality unless the programmer
        // does something bad like modify an Id<Obj>.
        self.stash.get(&id.raw_id).unwrap().downcast_ref().unwrap()
    }
}

fn main() {
    // Example usage

    let mut store = Store::new();
    let handle_a = store.create::<ObjectA>();
    let handle_b = store.create::<ObjectB>();

    let _obj_a = store.get(handle_a); // obj_a is of type &ObjectA
    let _obj_b = store.get(handle_b); // obj_b is of type &ObjectB
}

Avoiding the runtime check

In Rust nightly, we have downcast_ref_unchecked() on Any which does exactly what we want: it casts the value to a concrete type at runtime without performing the check. We simply replace .downcast_ref() with the unchecked version in the above code -- make sure to do some performance checks if you are using this in a real project to make sure the unsafety is actually worth it and helping performance (which I somewhat doubt, since unchecked operations don't necessarily actually improve performance of Rust code -- see e.g. this paper).

If you don't want to use Nightly, what you need is an (unsafe) void pointer. One way to get void pointers is c_void from the C FFI. Another way is to just cast *const T to *const () (the () is just a placeholder type) and then uncast later using std::mem::transmute.