path: root/src/dynamics/solver/interaction_groups.rs

use crate::dynamics::{BodyPair, JointGraphEdge, JointIndex, RigidBodySet};
use crate::geometry::{ContactManifold, ContactManifoldIndex};
#[cfg(feature = "simd-is-enabled")]
use {
    crate::math::{SIMD_LAST_INDEX, SIMD_WIDTH},
    vec_map::VecMap,
};

pub(crate) trait PairInteraction {
    fn body_pair(&self) -> BodyPair;
}

impl<'a> PairInteraction for &'a mut ContactManifold {
    fn body_pair(&self) -> BodyPair {
        self.data.body_pair
    }
}

impl<'a> PairInteraction for JointGraphEdge {
    fn body_pair(&self) -> BodyPair {
        BodyPair::new(self.weight.body1, self.weight.body2)
    }
}

#[cfg(feature = "parallel")]
pub(crate) struct ParallelInteractionGroups {
    bodies_color: Vec<u128>,         // Workspace.
    interaction_indices: Vec<usize>, // Workspace.
    interaction_colors: Vec<usize>,  // Workspace.
    sorted_interactions: Vec<usize>,
    groups: Vec<usize>,
}

#[cfg(feature = "parallel")]
impl ParallelInteractionGroups {
    pub fn new() -> Self {
        Self {
            bodies_color: Vec::new(),
            interaction_indices: Vec::new(),
            interaction_colors: Vec::new(),
            sorted_interactions: Vec::new(),
            groups: Vec::new(),
        }
    }

    pub fn group(&self, i: usize) -> &[usize] {
        let range = self.groups[i]..self.groups[i + 1];
        &self.sorted_interactions[range]
    }
    pub fn num_groups(&self) -> usize {
        self.groups.len() - 1
    }

    pub fn group_interactions<Interaction: PairInteraction>(
        &mut self,
        island_id: usize,
        bodies: &RigidBodySet,
        interactions: &[Interaction],
        interaction_indices: &[usize],
    ) {
        let num_island_bodies = bodies.active_island(island_id).len();
        self.bodies_color.clear();
        self.interaction_indices.clear();
        self.groups.clear();
        self.sorted_interactions.clear();
        self.interaction_colors.clear();

        let mut color_len = [0; 128];
        self.bodies_color.resize(num_island_bodies, 0u128);
        self.interaction_indices
            .extend_from_slice(interaction_indices);
        self.interaction_colors.resize(interaction_indices.len(), 0);
        let bcolors = &mut self.bodies_color;

        for (interaction_id, color) in self
            .interaction_indices
            .iter()
            .zip(self.interaction_colors.iter_mut())
        {
            let body_pair = interactions[*interaction_id].body_pair();
            let rb1 = &bodies[body_pair.body1];
            let rb2 = &bodies[body_pair.body2];

            match (rb1.is_static(), rb2.is_static()) {
                (false, false) => {
                    let color_mask =
                        bcolors[rb1.active_set_offset] | bcolors[rb2.active_set_offset];
                    *color = (!color_mask).trailing_zeros() as usize;
                    color_len[*color] += 1;
                    bcolors[rb1.active_set_offset] |= 1 << *color;
                    bcolors[rb2.active_set_offset] |= 1 << *color;
                }
                (true, false) => {
                    let color_mask = bcolors[rb2.active_set_offset];
                    *color = (!color_mask).trailing_zeros() as usize;
                    color_len[*color] += 1;
                    bcolors[rb2.active_set_offset] |= 1 << *color;
                }
                (false, true) => {
                    let color_mask = bcolors[rb1.active_set_offset];
                    *color = (!color_mask).trailing_zeros() as usize;
                    color_len[*color] += 1;
                    bcolors[rb1.active_set_offset] |= 1 << *color;
                }
                (true, true) => unreachable!(),
            }
        }

        let mut sort_offsets = [0; 128];
        let mut last_offset = 0;

        for i in 0..128 {
            if color_len[i] == 0 {
                break;
            }

            self.groups.push(last_offset);
            sort_offsets[i] = last_offset;
            last_offset += color_len[i];
        }

        self.sorted_interactions
            .resize(interaction_indices.len(), 0);

        for (interaction_id, color) in interaction_indices
            .iter()
            .zip(self.interaction_colors.iter())
        {
            self.sorted_interactions[sort_offsets[*color]] = *interaction_id;
            sort_offsets[*color] += 1;
        }

        self.groups.push(self.sorted_interactions.len());
    }
}

pub(crate) struct InteractionGroups {
    #[cfg(feature = "simd-is-enabled")]
    buckets: VecMap<([usize; SIMD_WIDTH], usize)>,
    #[cfg(feature = "simd-is-enabled")]
    body_masks: Vec<u128>,
    #[cfg(feature = "simd-is-enabled")]
    pub grouped_interactions: Vec<usize>,
    pub nongrouped_interactions: Vec<usize>,
}

impl InteractionGroups {
    pub fn new() -> Self {
        Self {
            #[cfg(feature = "simd-is-enabled")]
            buckets: VecMap::new(),
            #[cfg(feature = "simd-is-enabled")]
            body_masks: Vec::new(),
            #[cfg(feature = "simd-is-enabled")]
            grouped_interactions: Vec::new(),
            nongrouped_interactions: Vec::new(),
        }
    }

    // FIXME: there is a lot of duplicated code with group_manifolds here.
    // But we don't refactor just now because we may end up with distinct
    // grouping strategies in the future.
    #[cfg(not(feature = "simd-is-enabled"))]
    pub fn group_joints(
        &mut self,
        _island_id: usize,
        _bodies: &RigidBodySet,
        _interactions: &[JointGraphEdge],
        interaction_indices: &[JointIndex],
    ) {
        self.nongrouped_interactions
            .extend_from_slice(interaction_indices);
    }

    #[cfg(feature = "simd-is-enabled")]
    pub fn group_joints(
        &mut self,
        island_id: usize,
        bodies: &RigidBodySet,
        interactions: &[JointGraphEdge],
        interaction_indices: &[JointIndex],
    ) {
        // NOTE: in 3D we have up to 10 different joint types.
        // In 2D we only have 5 joint types.
        #[cfg(feature = "dim3")]
        const NUM_JOINT_TYPES: usize = 10;
        #[cfg(feature = "dim2")]
        const NUM_JOINT_TYPES: usize = 5;

        // The j-th bit of joint_type_conflicts[i] indicates that the
        // j-th bucket contains a joint with a type different than `i`.
        let mut joint_type_conflicts = [0u128; NUM_JOINT_TYPES];

        // Note: each bit of a body mask indicates what bucket already contains
        // a constraints involving this body.
        // FIXME: currently, this is a bit overconservative because when a bucket
        // is full, we don't clear the corresponding body mask bit. This may result
        // in less grouped constraints.
        self.body_masks
            .resize(bodies.active_island(island_id).len(), 0u128);

        // NOTE: each bit of the occupied mask indicates what bucket already
        // contains at least one constraint.
        let mut occupied_mask = 0u128;

        for interaction_i in interaction_indices {
            let interaction = &interactions[*interaction_i].weight;
            let body1 = &bodies[interaction.body1];
            let body2 = &bodies[interaction.body2];
            let is_static1 = !body1.