An idea I had for how storage could be managed in nphysics.

world_storage.rs

//! The idea of this system is to abstract the Set type, remove the need for `pop_insertion_event`
//! and `pop_removal_event`, and abstract the maintain functions of both worlds. The only issue I
//! have had so far is trying to figure out what WorldStorage should be responsible over compared to
//! GeometricalWorld.
//!
//! Goal:
//!  - Give ecs control over how storages are maintained
//!  - Make the implementation of Set consistent for bodies, colliders, and constraints
//!  - Bundle all of the storages into one convenient struct (low priority)
#![allow(missing_docs)]

use std::collections::{hash_map, HashMap};

use nalgebra::RealField;

use crate::joint::{DefaultJointConstraintHandle, DefaultJointConstraintSet, JointConstraint};
use crate::object::{
    Body, BodyHandle, BodySet, Collider, ColliderAnchor, ColliderHandle, DefaultBodyHandle,
    DefaultBodySet, DefaultColliderHandle, DefaultColliderSet,
};
use crate::volumetric::Volumetric;

/// A set of unknown elements corresponding to handles.
pub trait Set<Handle> {
    type Item: ?Sized;
    fn get(&self, handle: Handle) -> Option<&Self::Item>;
    fn get_mut(&mut self, handle: Handle) -> Option<&mut Self::Item>;
    fn contains(&self, handle: Handle) -> bool;
    fn foreach(&self, f: impl FnMut(Handle, &Self::Item));
    fn foreach_mut(&mut self, f: impl FnMut(Handle, &mut Self::Item));
}

/// A set for managing entries.
/// I originally created this trait so that I could implement the add/remove body/collider/constraint
/// methods inside the WorldStorage trait and handle everything there so it wouldn't need to be
/// implemented again every time WorldStorage is used. However, I would have needed to either add
/// more arguments to add/remove methods to pass in details from the geometric world or add a bunch
/// of one-time use functions to the WorldStorage to be used exclusivly by those specific
/// implementations of the add/remove methods.
///
/// TL;DR Not currently needed, but if I continue this idea I might add it back in later.
pub trait MutableSet<Handle>: Set<Handle> {
    type InsertType;
    /// Add a new element to this set. This function is unsafe because it has the potential to cause
    /// unknown issues when called directly. Any insertion for a Collider, Body, or Constraint must
    /// be accompanied with checks in at least one other set to avoid situations such as colliders
    /// without bodies. Use the corresponding insert method in `WorldStorage` instead.
    unsafe fn insert(&mut self, item: Self::InsertType) -> Handle;
    unsafe fn remove(&mut self, handle: Handle);
}

/// Needs to cover the abilities of a BodySet, ColliderSet and JointConstraintSet
pub trait WorldStorage<
    N: RealField,
    Handle: BodyHandle,
    CollHandle: ColliderHandle,
    ConstrHandle: Copy,
> where
    <Self::Bodies as Set<Handle>>::Item: Body<N>,
    <Self::Constraints as Set<ConstrHandle>>::Item: JointConstraint<N, Self::Bodies>,
{
    // The internal data storage types
    type Bodies: MutableSet<Handle> + BodySet<N>;
    // TODO: Remove BodySet<N> when possible
    type Colliders: MutableSet<CollHandle> + Set<CollHandle, Item = Collider<N, Handle>>;
    type Constraints: MutableSet<ConstrHandle>;

    fn bodies(&self) -> &Self::Bodies;
    fn colliders(&self) -> &Self::Colliders;
    fn constraints(&self) -> &Self::Constraints;

    fn bodies_mut(&mut self) -> &mut Self::Bodies;
    fn colliders_mut(&mut self) -> &mut Self::Colliders;
    fn constraints_mut(&mut self) -> &mut Self::Constraints;

    /// Break up the world storage into its unique components
    fn split(
        &mut self,
    ) -> (
        &mut Self::Bodies,
        &mut Self::Colliders,
        &mut Self::Constraints,
    );

    // insert entities
    fn insert_body(&mut self, body: <Self::Bodies as MutableSet<Handle>>::InsertType) -> Handle;
    fn insert_collider(
        &mut self,
        body: <Self::Colliders as MutableSet<CollHandle>>::InsertType,
    ) -> CollHandle;
    fn insert_constraint(
        &mut self,
        body: <Self::Constraints as MutableSet<ConstrHandle>>::InsertType,
    ) -> ConstrHandle;

    // remove entities
    fn remove_body(&mut self, handle: Handle);
    fn remove_collider(&mut self, handle: CollHandle);
    fn remove_constraint(&mut self, handle: ConstrHandle);

    // Due to ecs being fundamentally different in how it manages relations between entities, this method
    // will likely look very different in an ecs implementation.
    fn body_colliders(&self, body: Handle) -> Option<&[CollHandle]>;
}

/// Basically just wraps the three storage structs and provides all the basic functionality plus
/// `MechanicalWorld::maintain` and `GeometricalWorld::maintain` built in.
///
/// This could potentially be generified over the set type in the future
pub struct DefaultWorldStorage<N: RealField> {
    pub bodies: DefaultBodySet<N>,
    pub colliders: DefaultColliderSet<N, DefaultBodyHandle>,
    pub constraints: DefaultJointConstraintSet<N, DefaultBodySet<N>>,
    pub(crate) body_colliders: HashMap<DefaultBodyHandle, Vec<DefaultColliderHandle>>,
}

impl<N: RealField> Default for DefaultWorldStorage<N> {
    fn default() -> Self {
        Self {
            bodies: DefaultBodySet::new(),
            colliders: DefaultColliderSet::new(),
            constraints: DefaultJointConstraintSet::new(),
            body_colliders: HashMap::default(),
        }
    }
}

impl<N: RealField> DefaultWorldStorage<N> {
    pub fn new() -> Self {
        Self::default()
    }
}

// Note: I ended up just ignoring MutableSet for the ease of making this example
impl<N: RealField>
    WorldStorage<N, DefaultBodyHandle, DefaultColliderHandle, DefaultJointConstraintHandle>
    for DefaultWorldStorage<N>
{
    type Bodies = DefaultBodySet<N>;
    type Colliders = DefaultColliderSet<N, DefaultBodyHandle>;
    type Constraints = DefaultJointConstraintSet<N, DefaultBodySet<N>>;

    fn bodies(&self) -> &Self::Bodies {
        &self.bodies
    }

    fn colliders(&self) -> &Self::Colliders {
        &self.colliders
    }

    fn constraints(&self) -> &Self::Constraints {
        &self.constraints
    }

    fn bodies_mut(&mut self) -> &mut Self::Bodies {
        &mut self.bodies
    }

    fn colliders_mut(&mut self) -> &mut Self::Colliders {
        &mut self.colliders
    }

    fn constraints_mut(&mut self) -> &mut Self::Constraints {
        &mut self.constraints
    }

    fn split(
        &mut self,
    ) -> (
        &mut Self::Bodies,
        &mut Self::Colliders,
        &mut Self::Constraints,
    ) {
        (&mut self.bodies, &mut self.colliders, &mut self.constraints)
    }

    fn insert_body(
        &mut self,
        body: <Self::Bodies as MutableSet<DefaultBodyHandle>>::InsertType,
    ) -> DefaultBodyHandle
    {
        self.bodies.insert_boxed(body)
    }

    fn insert_collider(
        &mut self,
        mut collider: <Self::Colliders as MutableSet<DefaultColliderHandle>>::InsertType,
    ) -> DefaultColliderHandle
    {
        // TODO: Figure out how to call `register_collider` here, or move responsibilities

        match collider.anchor() {
            ColliderAnchor::OnBodyPart {
                body_part,
                position_wrt_body_part,
            } => {
                let body = self
                    .bodies
                    .get_mut(body_part.0)
                    .expect("Invalid parent body part handle.");

                // Update the parent body's inertia.
                if !collider.density().is_zero() {
                    let (com, inertia) = collider
                        .shape()
                        .transformed_mass_properties(collider.density(), position_wrt_body_part);
                    body.add_local_inertia_and_com(body_part.1, com, inertia);
                }

                // Set the position and ndofs (note this will be done in the `sync_collider` too.
                let ndofs = body.status_dependent_ndofs();
                let part = body
                    .part(body_part.1)
                    .expect("Invalid parent body part handle.");
                let pos = part.position() * position_wrt_body_part;

                collider.set_body_status_dependent_ndofs(ndofs);
                collider.set_position(pos);
            }
            _ => {}
        };

        self.colliders.insert(collider)
    }

    fn insert_constraint(
        &mut self,
        constr: <Self::Constraints as MutableSet<DefaultJointConstraintHandle>>::InsertType,
    ) -> DefaultJointConstraintHandle
    {
        let (body1, body2) = constr.anchors();

        if let Some(body1) = self.bodies.get_mut(body1.0) {
            body1.activate()
        }

        if let Some(body2) = self.bodies.get_mut(body2.0) {
            body2.activate()
        }

        self.constraints.insert_boxed(constr)
    }

    fn remove_body(&mut self, handle: DefaultBodyHandle) {
        let colliders = match self.body_colliders.get(&handle) {
            Some(x) => x.to_owned(),
            None => Vec::new(),
        };

        for collider in colliders {
            self.remove_collider(collider);
        }

        let mut constraints_to_remove = Vec::new();
        self.constraints.foreach(|h, c| {
            let (b1, b2) = c.anchors();
            if !Set::contains(&self.bodies, b1.0) || !Set::contains(&self.bodies, b2.0) {
                constraints_to_remove.push(h)
            }
        });

        for to_remove in constraints_to_remove {
            self.remove_constraint(to_remove);
        }
    }

    fn remove_collider(&mut self, handle: DefaultColliderHandle) {
        if let Some(collider) = self.colliders.get_mut(handle) {
            // removal_data will always return Some() so this is currently safe to unwrap
            let removed = collider.removal_data().unwrap();

            // Activate the body the deleted collider was attached to.
            if let Some(body) = self.bodies.get_mut(removed.anchor.body()) {
                // Update the parent body's inertia.
                if !removed.density.is_zero() {
                    if let ColliderAnchor::OnBodyPart {
                        body_part,
                        position_wrt_body_part,
                    } = &removed.anchor
                    {
                        let (com, inertia) = removed
                            .shape
                            .transformed_mass_properties(removed.density, position_wrt_body_part);
                        body.add_local_inertia_and_com(body_part.1, -com, -inertia)
                    }
                }

                body.activate()
            }

            // Remove the collider from the list of colliders for this body.
            match self.body_colliders.entry(removed.anchor.body()) {
                hash_map::Entry::Occupied(mut e) => {
                    if let Some(i) = e.get().iter().position(|h| *h == handle) {
                        let _ = e.get_mut().swap_remove(i);
                    }

                    if e.get().is_empty() {
                        let _ = e.remove_entry();
                    }
                }
                hash_map::Entry::Vacant(_) => {}
            }

            // TODO: Remove proxies and handle the graph index remapping.
        }
        let _ = self.colliders.remove(handle);
    }

    fn remove_constraint(&mut self, handle: DefaultJointConstraintHandle) {
        if let Some(joint) = self.constraints.get(handle) {
            let (body1, body2) = joint.anchors();

            if let Some(body1) = self.bodies.get_mut(body1.0) {
                body1.activate()
            }

            if let Some(body2) = self.bodies.get_mut(body2.0) {
                body2.activate()
            }
        }
        let _ = self.constraints.remove(handle);
    }

    fn body_colliders(&self, body: DefaultBodyHandle) -> Option<&[DefaultColliderHandle]> {
        self.body_colliders.get(&body).map(|c| &c[..])
    }
}