mtb-entity-slab 0.2.4

Slab-style entity storage: stable IDs, internal mutability; not a full ECS.
Documentation
#![cfg(test)]

use std::{cell::Cell, collections::HashSet};

use crate::*;

#[derive(Debug, Clone)]
#[entity_id(InstID, policy = 512, allocator_type = InstAlloc, backend = index)]
struct Inst {
    head: Cell<EntityListNodeHead<InstID>>,
    order: Cell<usize>,
    opcode: usize,
    _operands: [u64; 8],
    _heap_data: String,
}

impl IBasicEntityListID for InstID {
    fn obj_load_head(obj: &Self::ObjectT) -> EntityListNodeHead<Self> {
        obj.head.get()
    }
    fn obj_store_head(obj: &Self::ObjectT, head: EntityListNodeHead<Self>) {
        obj.head.set(head);
    }

    fn obj_is_sentinel(obj: &Self::ObjectT) -> bool {
        obj.opcode == usize::MAX
    }
    fn new_sentinel_obj() -> Self::ObjectT {
        Inst {
            head: Cell::new(EntityListNodeHead::none()),
            order: Cell::new(0),
            opcode: usize::MAX,
            _operands: [0; 8],
            _heap_data: String::new(),
        }
    }
    fn on_push_prev(self, prev: Self, a: &InstAlloc) -> EntityListRes<Self> {
        let opcode = self.deref_alloc(a).opcode;
        let prev_op = prev.deref_alloc(a).opcode;
        println!(
            "Signal `on_push_prev` invoked on pushing inst {prev_op} before {{ opcode: {opcode}, id: {prev:?} }}"
        );
        Ok(())
    }
    fn on_push_next(self, next: Self, a: &InstAlloc) -> EntityListRes<Self> {
        let opcode = self.deref_alloc(a).opcode;
        let next_op = next.deref_alloc(a).opcode;
        println!("Signal `on_push_next` invoked on pushing inst {next_op} after {opcode}");
        Ok(())
    }
    fn on_unplug(self, a: &InstAlloc) -> EntityListRes<Self> {
        let opcode = self.deref_alloc(a).opcode;
        println!("Signal `on_unplug` invoked on inst {opcode}");
        Ok(())
    }
}

impl IEntityListNodeID for InstID {}

impl IOrderCachedListNodeID for InstID {
    const ORDER_REPR: OrderRepr = OrderRepr::Relative;

    fn obj_load_order(obj: &Self::ObjectT) -> usize {
        obj.order.get()
    }
    fn obj_store_order(obj: &Self::ObjectT, order: usize) {
        obj.order.set(order);
    }
}

impl Inst {
    fn new(opcode: usize) -> Self {
        Self {
            head: Cell::new(EntityListNodeHead::none()),
            order: Cell::new(0),
            opcode,
            _operands: [0; 8],
            _heap_data: String::new(),
        }
    }

    fn new_id(alloc: &IDBoundAlloc<InstID>, opcode: usize) -> InstID {
        let back = IndexedID::allocate_from(alloc, Inst::new(opcode));
        InstID::from_backend(back)
    }
}

fn do_test_allocation_deallocation<P: policy::IAllocPolicy>() {
    type PrimaryID<E, P> = IndexedID<E, P>;
    let mut alloc: EntityAlloc<Inst, P> = EntityAlloc::with_capacity(1024);
    let wave0_len = P::CHUNK_SIZE * 4 - 257;
    let wave1_len = P::CHUNK_SIZE * 4 + 257;

    let mut ptrs = Vec::with_capacity(wave0_len);
    for i in 0..wave0_len {
        let inst = Inst::new(i);
        let prim = PrimaryID::allocate_from(&alloc, inst);
        ptrs.push(prim);
        prim.deref_alloc_mut(&mut alloc)._operands[0] = i as u64;
    }

    let wave0_inst_index = wave0_len / 2;
    let inst = ptrs[wave0_inst_index].deref_alloc(&alloc);
    for i in 0..wave1_len {
        let inst = Inst::new(i + wave0_len);
        PrimaryID::allocate_from(&alloc, inst);
    }
    let indexed = ptrs[wave0_inst_index].to_index(&alloc).unwrap();
    assert_eq!(inst.opcode, indexed.indexed.real_index());

    let mut kept_indices = Vec::new();
    alloc.free_if(|inst, _, i| {
        assert_eq!(i.get_order(), inst.opcode);
        if inst.opcode % 3 == 0 {
            kept_indices.push(i);
            false
        } else {
            true
        }
    });

    let mut alive_indices = Vec::with_capacity(kept_indices.len());
    let mut traversed = HashSet::new();
    for (i, _, inst) in alloc.iter() {
        alive_indices.push(i);
        assert_eq!(i.get_order(), inst.opcode);
        assert!(
            traversed.insert(i),
            "duplicate instance found during iteration"
        );
    }
    assert_eq!(kept_indices, alive_indices)
}

#[test]
fn test_allocation_deallocation() {
    do_test_allocation_deallocation::<AllocPolicy128>();
    if cfg!(not(miri)) {
        // these tests cost a lot of time and memory, so we skip them under miri
        do_test_allocation_deallocation::<AllocPolicy256>();
        do_test_allocation_deallocation::<AllocPolicy512>();
        do_test_allocation_deallocation::<AllocPolicy1024>();
        do_test_allocation_deallocation::<AllocPolicy2048>();
        do_test_allocation_deallocation::<AllocPolicy4096>();
    }
}

#[test]
fn test_list_op() {
    let alloc = InstAlloc::new();
    let list = EntityList::new(&alloc);

    for i in 0..100 {
        let inst = Inst::new_id(&alloc, i);
        list.push_back_id(inst, &alloc).unwrap();
    }

    let mut nodes = HashSet::new();

    for (node_id, node) in list.iter(&alloc) {
        println!("node {node_id:?} is inst {}", node.opcode);
        assert!(nodes.insert(node_id), "list iteration has duplicate nodes");
    }
}

#[test]
fn test_uaf() {
    let mut alloc = InstAlloc::new();

    let old_id = Inst::new_id(&alloc, 42);
    old_id.0.free(&mut alloc);

    let new_id = Inst::new_id(&alloc, 43);

    let old_gen = old_id.0.get_generation();
    let new_gen = new_id.0.get_generation();
    assert_ne!(
        old_gen, new_gen,
        "UAF checking failed: generation numbers are equal"
    );
    println!("Old generation: {}, New generation: {}", old_gen, new_gen);

    let None = old_id.0.try_deref(&alloc) else {
        panic!("UAF checking failed: old_id should not be dereferenceable");
    };
    let Some(inst) = new_id.0.try_deref(&alloc) else {
        panic!("Newly allocated id should be dereferenceable");
    };
    assert_eq!(
        inst.opcode, 43,
        "Newly allocated instance has incorrect opcode"
    );
}

#[test]
fn test_ordered_list() {
    let alloc = InstAlloc::new();
    let list = OrderCachedList::new(&alloc);

    for i in 0..100 {
        let inst = Inst::new_id(&alloc, i);
        list.push_back_id(inst, &alloc).unwrap();
    }
    let mut expected_order = 1;
    for (node_id, _) in list.iter(&alloc) {
        let order = list.get_node_order(&alloc, node_id);
        assert_eq!(order, expected_order, "Node order mismatch");
        expected_order += 1;
    }
}