avian3d 0.6.1

use crate::{
    dynamics::{
        joints::EntityConstraint,
        solver::{
            SolverDiagnostics,
            constraint_graph::{
                COLOR_OVERFLOW_INDEX, ConstraintGraph, ContactManifoldHandle, GraphColor,
            },
            contact::ContactConstraint,
            schedule::SubstepSolverSystems,
            softness_parameters::{SoftnessCoefficients, SoftnessParameters},
            solver_body::{SolverBody, SolverBodyInertia},
        },
    },
    prelude::*,
};
use bevy::{
    ecs::{query::QueryData, system::lifetimeless::Read},
    prelude::*,
};
use core::cmp::Ordering;

/// Manages and solves contacts, joints, and other constraints.
///
/// Note that the [`ContactConstraints`] are currently generated by tbe [`NarrowPhasePlugin`].
///
/// # Implementation
///
/// Avian uses an impulse-based solver with substepping and [soft constraints](super::softness_parameters).
/// Warm starting is used to improve convergence, along with a relaxation pass to reduce overshooting.
///
/// [Speculative collision](dynamics::ccd#speculative-collision) is used by default to prevent tunneling.
/// Optional [sweep-based Continuous Collision Detection (CCD)](dynamics::ccd#swept-ccd) is handled by the [`CcdPlugin`].
///
/// [Joints](dynamics::joints) and user constraints are currently solved using [Extended Position-Based Dynamics (XPBD)](super::xpbd)
/// if the `xpbd_joints` feature is enabled. In the future, they may transition to an impulse-based approach as well.
///
/// ## Solver Bodies
///
/// The solver maintains a [`SolverBody`] for each awake dynamic and kinematic body.
/// It stores the body data needed by the solver in a more optimized format
/// with better memory locality.
///
/// Only awake dynamic bodies and kinematic bodies have an associated solver body.
/// Static bodies and sleeping dynamic bodies do not move and are not included in the solver.
///
/// The [`SolverBodyPlugin`] is added for managing solver bodies and synchronizing them with rigid body data.
///
/// # Steps
///
/// Below are the main steps of the `SolverPlugin`.
///
/// 1. [Prepare solver bodies](SolverSystems::PrepareSolverBodies)
/// 2. [Prepare joint constraints](SolverSystems::PrepareJoints)
/// 3. [Prepare contact constraints](SolverSystems::PrepareContactConstraints)
/// 4. Substepping loop (runs the [`SubstepSchedule`] [`SubstepCount`] times)
///     1. [Integrate velocities](crate::dynamics::integrator::IntegrationSystems::Velocity)
///     2. [Warm start](SubstepSolverSystems::WarmStart)
///     3. [Solve constraints with bias](SubstepSolverSystems::SolveConstraints)
///     4. [Integrate positions](crate::dynamics::integrator::IntegrationSystems::Position)
///     5. [Solve constraints without bias to relax velocities](SubstepSolverSystems::Relax)
/// 5. [Apply restitution](SolverSystems::Restitution)
/// 6. [Write back solver body data to rigid bodies](SolverSystems::Finalize)
/// 7. [Store contact impulses for next frame's warm starting](SolverSystems::StoreContactImpulses)
///
/// If the `xpbd_joints` feature is enabled, the [`XpbdSolverPlugin`] can also be added to solve joints
/// using Extended Position-Based Dynamics (XPBD).
pub struct SolverPlugin {
    length_unit: Scalar,
}

impl Default for SolverPlugin {
    fn default() -> Self {
        Self::new_with_length_unit(1.0)
    }
}

impl SolverPlugin {
    /// Creates a [`SolverPlugin`] with the given approximate dimensions of most objects.
    ///
    /// The length unit will be used for initializing the [`PhysicsLengthUnit`]
    /// resource unless it already exists.
    pub fn new_with_length_unit(unit: Scalar) -> Self {
        Self { length_unit: unit }
    }
}

impl Plugin for SolverPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<SolverConfig>()
            .init_resource::<ContactSoftnessCoefficients>()
            .init_resource::<ContactConstraints>()
            .init_resource::<ConstraintGraph>();

        if app
            .world()
            .get_resource::<PhysicsLengthUnit>()
            .is_none_or(|unit| unit.0 == 1.0)
        {
            app.insert_resource(PhysicsLengthUnit(self.length_unit));
        }

        // Get the `PhysicsSchedule`, and panic if it doesn't exist.
        let physics = app
            .get_schedule_mut(PhysicsSchedule)
            .expect("add PhysicsSchedule first");

        physics.add_systems(update_contact_softness.before(PhysicsStepSystems::NarrowPhase));

        // Prepare contact constraints before the substepping loop.
        physics.add_systems(
            prepare_contact_constraints.in_set(SolverSystems::PrepareContactConstraints),
        );

        // Apply restitution.
        physics.add_systems(solve_restitution.in_set(SolverSystems::Restitution));

        // Store the current contact impulses for the next frame's warm starting.
        physics.add_systems(store_contact_impulses.in_set(SolverSystems::StoreContactImpulses));

        // Get the `SubstepSchedule`, and panic if it doesn't exist.
        let substeps = app
            .get_schedule_mut(SubstepSchedule)
            .expect("add SubstepSchedule first");

        // Warm start the impulses.
        // This applies the impulses stored from the previous substep,
        // which improves convergence.
        substeps.add_systems(warm_start.in_set(SubstepSolverSystems::WarmStart));

        // Solve velocities using a position bias.
        substeps.add_systems(solve_contacts::<true>.in_set(SubstepSolverSystems::SolveConstraints));

        // Relax biased velocities and impulses.
        // This reduces overshooting caused by warm starting.
        substeps.add_systems(solve_contacts::<false>.in_set(SubstepSolverSystems::Relax));

        // Perform constraint damping.
        substeps.add_systems(
            (
                joint_damping::<FixedJoint>,
                joint_damping::<RevoluteJoint>,
                #[cfg(feature = "3d")]
                joint_damping::<SphericalJoint>,
                joint_damping::<PrismaticJoint>,
                joint_damping::<DistanceJoint>,
            )
                .chain()
                .in_set(SubstepSolverSystems::Damping),
        );
    }

    fn finish(&self, app: &mut App) {
        // Register timer and counter diagnostics for the solver.
        app.register_physics_diagnostics::<SolverDiagnostics>();
    }
}

// TODO: Where should this type be and which plugin should initialize it?
/// A units-per-meter scaling factor that adjusts the engine's internal properties
/// to the scale of the world.
///
/// For example, a 2D game might use pixels as units and have an average object size
/// of around 100 pixels. By setting the length unit to `100.0`, the physics engine
/// will interpret 100 pixels as 1 meter for internal thresholds, improving stability.
///
/// Note that this is *not* used to scale forces or any other user-facing inputs or outputs.
/// Instead, the value is only used to scale some internal length-based tolerances, such as
/// [`SleepingThreshold::linear`] and [`NarrowPhaseConfig::default_speculative_margin`],
/// as well as the scale used for [debug rendering](PhysicsDebugPlugin).
///
/// Choosing the appropriate length unit can help improve stability and robustness.
///
/// Default: `1.0`
///
/// # Example
///
/// The [`PhysicsLengthUnit`] can be inserted as a resource like normal,
/// but it can also be specified through the [`PhysicsPlugins`] plugin group.
///
/// ```no_run
/// # #[cfg(feature = "2d")]
/// use avian2d::prelude::*;
/// use bevy::prelude::*;
///
/// # #[cfg(feature = "2d")]
/// fn main() {
///     App::new()
///         .add_plugins((
///             DefaultPlugins,
///             // A 2D game with 100 pixels per meter
///             PhysicsPlugins::default().with_length_unit(100.0),
///         ))
///         .run();
/// }
/// # #[cfg(not(feature = "2d"))]
/// # fn main() {} // Doc test needs main
/// ```
#[derive(Resource, Clone, Debug, Deref, DerefMut, PartialEq, Reflect)]
#[reflect(Resource)]
pub struct PhysicsLengthUnit(pub Scalar);

impl Default for PhysicsLengthUnit {
    fn default() -> Self {
        Self(1.0)
    }
}

/// Configuration parameters for the constraint solver that handles
/// things like contacts and joints.
///
/// These are tuned to give good results for most applications, but can
/// be configured if more control over the simulation behavior is needed.
#[derive(Resource, Clone, Debug, PartialEq, Reflect)]
#[reflect(Resource)]
pub struct SolverConfig {
    /// The damping ratio used for contact stabilization.
    ///
    /// Lower values make contacts more compliant or "springy",
    /// allowing more visible penetration before overlap has been
    /// resolved and the contact has been stabilized.
    ///
    /// Consider using a higher damping ratio if contacts seem too soft.
    /// Note that making the value too large can cause instability.
    ///
    /// Default: `10.0`.
    pub contact_damping_ratio: Scalar,

    /// Scales the frequency used for contacts. A higher frequency
    /// makes contact responses faster and reduces visible springiness,
    /// but can hurt stability.
    ///
    /// The solver computes the frequency using the time step and substep count,
    /// and limits the maximum frequency to be at most half of the time step due to
    /// [Nyquist's theorem](https://en.wikipedia.org/wiki/Nyquist%E2%80%93Shannon_sampling_theorem).
    /// This factor scales the resulting frequency, which can lead to unstable behavior
    /// if the factor is too large.
    ///
    /// Default: `1.5`
    pub contact_frequency_factor: Scalar,

    /// The maximum speed at which overlapping bodies are pushed apart by the solver.
    ///
    /// With a small value, overlap is resolved gently and gradually, while large values
    /// can result in more snappy behavior.
    ///
    /// This is implicitly scaled by the [`PhysicsLengthUnit`].
    ///
    /// Default: `4.0`
    pub max_overlap_solve_speed: Scalar,

    /// The coefficient in the `[0, 1]` range applied to
    /// [warm start](SubstepSolverSystems::WarmStart) impulses.
    ///
    /// Warm starting uses the impulses from the previous frame as the initial
    /// solution for the current frame. This helps the solver reach the desired
    /// state much faster, meaning that *convergence* is improved.
    ///
    /// The coefficient should typically be set to `1.0`.
    ///
    /// Default: `1.0`
    pub warm_start_coefficient: Scalar,

    /// The minimum speed along the contact normal in units per second
    /// for [restitution](Restitution) to be applied.
    ///
    /// An appropriate threshold should typically be small enough that objects
    /// keep bouncing until the bounces are effectively unnoticeable,
    /// but large enough that restitution is not applied unnecessarily,
    /// improving performance and stability.
    ///
    /// This is implicitly scaled by the [`PhysicsLengthUnit`].
    ///
    /// Default: `1.0`
    pub restitution_threshold: Scalar,

    /// The number of iterations used for applying [restitution](Restitution).
    ///
    /// A higher number of iterations can result in more accurate bounces,
    /// but it only makes a difference when there are more than one contact point.
    ///
    /// For example, with just one iteration, a cube falling flat on the ground
    /// might bounce and rotate to one side, because the impulses are applied
    /// to the corners sequentially, and some of the impulses are likely to be larger
    /// than the others. With multiple iterations, the impulses are applied more evenly.
    ///
    /// Default: `1`
    pub restitution_iterations: usize,
}

impl Default for SolverConfig {
    fn default() -> Self {
        Self {
            contact_damping_ratio: 10.0,
            contact_frequency_factor: 1.5,
            max_overlap_solve_speed: 4.0,
            warm_start_coefficient: 1.0,
            restitution_threshold: 1.0,
            restitution_iterations: 1,
        }
    }
}

/// The [`SoftnessCoefficients`] used for contacts.
///
/// **Note**: This resource is updated automatically and not intended to be modified manually.
/// Use the [`SolverConfig`] resource instead for tuning contact behavior.
#[derive(Resource, Clone, Copy, PartialEq, Reflect)]
#[reflect(Resource)]
pub struct ContactSoftnessCoefficients {
    /// The [`SoftnessCoefficients`] used for contacts against dynamic bodies.
    pub dynamic: SoftnessCoefficients,
    /// The [`SoftnessCoefficients`] used for contacts against static or kinematic bodies.
    pub non_dynamic: SoftnessCoefficients,
}

impl Default for ContactSoftnessCoefficients {
    fn default() -> Self {
        Self {
            dynamic: SoftnessParameters::new(10.0, 30.0).compute_coefficients(1.0 / 60.0),
            non_dynamic: SoftnessParameters::new(10.0, 60.0).compute_coefficients(1.0 / 60.0),
        }
    }
}

fn update_contact_softness(
    mut coefficients: ResMut<ContactSoftnessCoefficients>,
    solver_config: Res<SolverConfig>,
    physics_time: Res<Time<Physics>>,
    substep_time: Res<Time<Substeps>>,
) {
    if solver_config.is_changed() || physics_time.is_changed() || substep_time.is_changed() {
        let dt = physics_time.delta_secs_f64() as Scalar;
        let h = substep_time.delta_secs_f64() as Scalar;

        // The contact frequency should at most be half of the time step due to Nyquist's theorem.
        // https://en.wikipedia.org/wiki/Nyquist%E2%80%93Shannon_sampling_theorem
        let max_hz = 1.0 / (dt * 2.0);
        let hz = solver_config.contact_frequency_factor * max_hz.min(0.25 / h);

        coefficients.dynamic = SoftnessParameters::new(solver_config.contact_damping_ratio, hz)
            .compute_coefficients(h);

        // TODO: Perhaps the non-dynamic softness should be configurable separately.
        // Make contacts against static and kinematic bodies stiffer to avoid clipping through the environment.
        coefficients.non_dynamic =
            SoftnessParameters::new(solver_config.contact_damping_ratio, 2.0 * hz)
                .compute_coefficients(h);
    }
}

/// A resource that stores the contact constraints.
#[derive(Resource, Default, Deref, DerefMut)]
pub struct ContactConstraints(pub Vec<ContactConstraint>);

#[derive(QueryData)]
pub(super) struct BodyQuery {
    pub(super) rb: Read<RigidBody>,
    pub(super) linear_velocity: Read<LinearVelocity>,
    pub(super) inertia: Option<Read<SolverBodyInertia>>,
}

fn prepare_contact_constraints(
    contact_graph: Res<ContactGraph>,
    mut constraint_graph: ResMut<ConstraintGraph>,
    mut diagnostics: ResMut<SolverDiagnostics>,
    bodies: Query<BodyQuery, RigidBodyActiveFilter>,
    narrow_phase_config: Res<NarrowPhaseConfig>,
    contact_softness: Res<ContactSoftnessCoefficients>,
) {
    let start = crate::utils::Instant::now();

    for color in constraint_graph.colors.iter_mut() {
        // TODO: Instead of clearing the vector, we could resize it, and just overwrite the old values in the loop below.
        //       Then the inner loop could be parallelized too.
        // TODO: Box2D uses an arena allocator for constraints. Might be worth looking into?
        color.contact_constraints.clear();
    }

    // Get colors with at least one active manifold handle.
    let mut active_colors = constraint_graph
        .colors
        .iter_mut()
        .filter(|color| !color.manifold_handles.is_empty())
        .collect::<Vec<&mut GraphColor>>();

    // Generate contact constraints for each contact pair, parallelizing over graph colors.
    crate::utils::par_for_each(&mut active_colors, 2, |_i, color| {
        for handle in color.manifold_handles.iter() {
            // Get the contact pair and its manifold.
            let contact_pair = contact_graph
                .get_by_id(handle.contact_id)
                .unwrap_or_else(|| {
                    panic!("Contact pair not found in graph: {:?}", handle.contact_id)
                })
                .1;
            let manifold_index = handle.manifold_index;
            let manifold = &contact_pair.manifolds[manifold_index];

            if !contact_pair.generates_constraints() {
                continue;
            }

            let (Some(body1_entity), Some(body2_entity)) = (contact_pair.body1, contact_pair.body2)
            else {
                continue;
            };

            // Get the two colliding bodies.
            let Ok(body1) = bodies.get(body1_entity) else {
                continue;
            };
            let Ok(body2) = bodies.get(body2_entity) else {
                continue;
            };

            // TODO: To skip this, we probably shouldn't have manifold handles between non-dynamic bodies
            //       in the constraint graph. Or alternatively, just don't generate contacts at all for them.
            if !body1.rb.is_dynamic() && !body2.rb.is_dynamic() {
                // If both bodies are static or kinematic, skip the contact.
                continue;
            }

            let constraint = ContactConstraint::generate(
                body1_entity,
                body2_entity,
                body1,
                body2,
                contact_pair.contact_id,
                manifold,
                manifold_index,
                narrow_phase_config.match_contacts,
                &contact_softness,
            );

            if !constraint.points.is_empty() {
                color.contact_constraints.push(constraint);
            }
        }
    });

    diagnostics.prepare_constraints += start.elapsed();
    diagnostics.contact_constraint_count = constraint_graph
        .colors
        .iter()
        .map(|color| color.contact_constraints.len())
        .sum::<usize>() as u32;
}

/// Warm starts the solver by applying the impulses from the previous frame or substep.
///
/// See [`SubstepSolverSystems::WarmStart`] for more information.
fn warm_start(
    bodies: Query<(&mut SolverBody, &SolverBodyInertia)>,
    mut constraint_graph: ResMut<ConstraintGraph>,
    solver_config: Res<SolverConfig>,
    mut diagnostics: ResMut<SolverDiagnostics>,
) {
    let start = crate::utils::Instant::now();

    // Warm start overflow constraints serially. They have lower priority, so they are solved first.
    for constraint in constraint_graph.colors[COLOR_OVERFLOW_INDEX]
        .contact_constraints
        .iter_mut()
    {
        warm_start_internal(&bodies, constraint, solver_config.warm_start_coefficient);
    }

    // Warm start constraints in each color in parallel.
    for color in constraint_graph
        .colors
        .iter_mut()
        .take(COLOR_OVERFLOW_INDEX)
        .filter(|color| !color.contact_constraints.is_empty())
    {
        crate::utils::par_for_each(&mut color.contact_constraints, 64, |_i, constraint| {
            warm_start_internal(&bodies, constraint, solver_config.warm_start_coefficient);
        });
    }

    diagnostics.warm_start += start.elapsed();
}

fn warm_start_internal(
    bodies: &Query<(&mut SolverBody, &SolverBodyInertia)>,
    constraint: &mut ContactConstraint,
    warm_start_coefficient: Scalar,
) {
    debug_assert!(!constraint.points.is_empty());

    let mut dummy_body1 = SolverBody::DUMMY;
    let mut dummy_body2 = SolverBody::DUMMY;

    let (mut body1, mut inertia1) = (&mut dummy_body1, &SolverBodyInertia::DUMMY);
    let (mut body2, mut inertia2) = (&mut dummy_body2, &SolverBodyInertia::DUMMY);

    // Get the solver bodies for the two colliding entities.
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body1) } {
        body1 = body.into_inner();
        inertia1 = inertia;
    }
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body2) } {
        body2 = body.into_inner();
        inertia2 = inertia;
    }

    // If a body has a higher dominance, it is treated as a static or kinematic body.
    match constraint.relative_dominance.cmp(&0) {
        Ordering::Greater => inertia1 = &SolverBodyInertia::DUMMY,
        Ordering::Less => inertia2 = &SolverBodyInertia::DUMMY,
        _ => {}
    }

    constraint.warm_start(body1, body2, inertia1, inertia2, warm_start_coefficient);
}

/// Solves contacts by iterating through the given contact constraints
/// and applying impulses to colliding rigid bodies.
///
/// This solve is done `iterations` times. With a substepped solver,
/// `iterations` should typically be `1`, as substeps will handle the iteration.
///
/// If `use_bias` is `true`, the impulses will be boosted to account for overlap.
/// The solver should often be run twice per frame or substep: first with the bias,
/// and then without it to *relax* the velocities and reduce overshooting caused by
/// [warm starting](SubstepSolverSystems::WarmStart).
///
/// See [`SubstepSolverSystems::SolveConstraints`] and [`SubstepSolverSystems::Relax`] for more information.
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
fn solve_contacts<const USE_BIAS: bool>(
    bodies: Query<(&mut SolverBody, &SolverBodyInertia)>,
    mut constraint_graph: ResMut<ConstraintGraph>,
    solver_config: Res<SolverConfig>,
    length_unit: Res<PhysicsLengthUnit>,
    time: Res<Time>,
    mut diagnostics: ResMut<SolverDiagnostics>,
) {
    let start = crate::utils::Instant::now();

    let delta_secs = time.delta_seconds_adjusted();
    let max_overlap_solve_speed = solver_config.max_overlap_solve_speed * length_unit.0;

    // Solve overflow constraints serially. They have lower priority, so they are solved first.
    for constraint in constraint_graph.colors[COLOR_OVERFLOW_INDEX]
        .contact_constraints
        .iter_mut()
    {
        solve_contacts_internal::<USE_BIAS>(
            &bodies,
            constraint,
            max_overlap_solve_speed,
            delta_secs,
        );
    }

    // Solve contact constraints in each color in parallel.
    for color in constraint_graph
        .colors
        .iter_mut()
        .take(COLOR_OVERFLOW_INDEX)
        .filter(|color| !color.contact_constraints.is_empty())
    {
        crate::utils::par_for_each(&mut color.contact_constraints, 64, |_i, constraint| {
            solve_contacts_internal::<USE_BIAS>(
                &bodies,
                constraint,
                max_overlap_solve_speed,
                delta_secs,
            );
        });
    }

    if USE_BIAS {
        diagnostics.solve_constraints += start.elapsed();
    } else {
        diagnostics.relax_velocities += start.elapsed();
    }
}

fn solve_contacts_internal<const USE_BIAS: bool>(
    bodies: &Query<(&mut SolverBody, &SolverBodyInertia)>,
    constraint: &mut ContactConstraint,
    max_overlap_solve_speed: Scalar,
    delta_secs: Scalar,
) {
    let mut dummy_body1 = SolverBody::DUMMY;
    let mut dummy_body2 = SolverBody::DUMMY;

    let (mut body1, mut inertia1) = (&mut dummy_body1, &SolverBodyInertia::DUMMY);
    let (mut body2, mut inertia2) = (&mut dummy_body2, &SolverBodyInertia::DUMMY);

    // Get the solver bodies for the two colliding entities.
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body1) } {
        body1 = body.into_inner();
        inertia1 = inertia;
    }
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body2) } {
        body2 = body.into_inner();
        inertia2 = inertia;
    }

    // If a body has a higher dominance, it is treated as a static or kinematic body.
    match constraint.relative_dominance.cmp(&0) {
        Ordering::Greater => inertia1 = &SolverBodyInertia::DUMMY,
        Ordering::Less => inertia2 = &SolverBodyInertia::DUMMY,
        _ => {}
    }

    constraint.solve(
        body1,
        body2,
        inertia1,
        inertia2,
        delta_secs,
        USE_BIAS,
        max_overlap_solve_speed,
    );
}

/// Iterates through contact constraints and applies impulses to account for [`Restitution`].
///
/// Note that restitution with TGS Soft and speculative contacts may not be perfectly accurate.
/// This is a tradeoff, but cheap CCD is often more important than perfect restitution.
///
/// The number of iterations can be increased with [`SolverConfig::restitution_iterations`]
/// to apply restitution for multiple contact points more evenly.
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
fn solve_restitution(
    bodies: Query<(&mut SolverBody, &SolverBodyInertia)>,
    mut constraint_graph: ResMut<ConstraintGraph>,
    solver_config: Res<SolverConfig>,
    length_unit: Res<PhysicsLengthUnit>,
    mut diagnostics: ResMut<SolverDiagnostics>,
) {
    let start = crate::utils::Instant::now();

    // The restitution threshold determining the speed required for restitution to be applied.
    let threshold = solver_config.restitution_threshold * length_unit.0;

    // Solve restitution for overflow constraints serially. They have lower priority, so they are solved first.
    for constraint in constraint_graph.colors[COLOR_OVERFLOW_INDEX]
        .contact_constraints
        .iter_mut()
    {
        solve_restitution_internal(
            &bodies,
            constraint,
            threshold,
            solver_config.restitution_iterations,
        );
    }

    // Solve restitution for contact constraints in each color in parallel.
    for color in constraint_graph
        .colors
        .iter_mut()
        .take(COLOR_OVERFLOW_INDEX)
        .filter(|color| !color.contact_constraints.is_empty())
    {
        crate::utils::par_for_each(&mut color.contact_constraints, 64, |_i, constraint| {
            solve_restitution_internal(
                &bodies,
                constraint,
                threshold,
                solver_config.restitution_iterations,
            );
        });
    }

    diagnostics.apply_restitution += start.elapsed();
}

fn solve_restitution_internal(
    bodies: &Query<(&mut SolverBody, &SolverBodyInertia)>,
    constraint: &mut ContactConstraint,
    threshold: Scalar,
    iterations: usize,
) {
    let restitution = constraint.restitution;

    if restitution == 0.0 {
        return;
    }

    let mut dummy_body1 = SolverBody::DUMMY;
    let mut dummy_body2 = SolverBody::DUMMY;

    let (mut body1, mut inertia1) = (&mut dummy_body1, &SolverBodyInertia::DUMMY);
    let (mut body2, mut inertia2) = (&mut dummy_body2, &SolverBodyInertia::DUMMY);

    // Get the solver bodies for the two colliding entities.
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body1) } {
        body1 = body.into_inner();
        inertia1 = inertia;
    }
    if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(constraint.body2) } {
        body2 = body.into_inner();
        inertia2 = inertia;
    }

    // If a body has a higher dominance, it is treated as a static or kinematic body.
    match constraint.relative_dominance.cmp(&0) {
        Ordering::Greater => inertia1 = &SolverBodyInertia::DUMMY,
        Ordering::Less => inertia2 = &SolverBodyInertia::DUMMY,
        _ => {}
    }

    // Performing multiple iterations can result in more accurate restitution,
    // but only if there are more than one contact point.
    let point_count = constraint.points.len();
    let iterations = if point_count > 1 { iterations } else { 1 };

    for _ in 0..iterations {
        constraint.apply_restitution(body1, body2, inertia1, inertia2, threshold);
    }
}

/// Copies contact impulses from [`ContactConstraints`] to the contacts in the [`ContactGraph`].
/// They will be used for [warm starting](SubstepSolverSystems::WarmStart).
fn store_contact_impulses(
    mut contact_graph: ResMut<ContactGraph>,
    mut constraint_graph: ResMut<ConstraintGraph>,
    mut diagnostics: ResMut<SolverDiagnostics>,
) {
    let start = crate::utils::Instant::now();

    for color in constraint_graph.colors.iter_mut() {
        for constraint in &mut color.contact_constraints {
            let Some(manifold) = contact_graph.get_manifold_mut(ContactManifoldHandle {
                contact_id: constraint.contact_id,
                manifold_index: constraint.manifold_index,
            }) else {
                unreachable!(
                    "Contact manifold {:?} for contact ID {:?} not found in contact graph.",
                    constraint.contact_id, constraint.manifold_index
                );
            };

            for (contact, constraint_point) in
                manifold.points.iter_mut().zip(constraint.points.iter())
            {
                contact.warm_start_normal_impulse = constraint_point.normal_part.impulse;
                contact.warm_start_tangent_impulse = constraint_point
                    .tangent_part
                    .as_ref()
                    .map_or(default(), |part| part.impulse);
                contact.normal_impulse = constraint_point.normal_part.total_impulse;
            }
        }
    }

    diagnostics.store_impulses += start.elapsed();
}

/// Applies velocity corrections caused by joint damping.
#[allow(clippy::type_complexity)]
pub fn joint_damping<T: Component + EntityConstraint<2>>(
    bodies: Query<(&mut SolverBody, &SolverBodyInertia)>,
    joints: Query<(&T, &JointDamping)>,
    time: Res<Time>,
) {
    let delta_secs = time.delta_seconds_adjusted();

    let mut dummy_body1 = SolverBody::DUMMY;
    let mut dummy_body2 = SolverBody::DUMMY;

    for (joint, damping) in &joints {
        let entities = joint.entities();

        let (mut body1, mut inertia1) = (&mut dummy_body1, &SolverBodyInertia::DUMMY);
        let (mut body2, mut inertia2) = (&mut dummy_body2, &SolverBodyInertia::DUMMY);

        // Get the solver bodies for the two jointed entities.
        if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(entities[0]) } {
            body1 = body.into_inner();
            inertia1 = inertia;
        }
        if let Ok((body, inertia)) = unsafe { bodies.get_unchecked(entities[1]) } {
            body2 = body.into_inner();
            inertia2 = inertia;
        }

        let delta_omega = (body2.angular_velocity - body1.angular_velocity)
            * (damping.angular * delta_secs).min(1.0);

        if !body1.flags.is_kinematic() {
            body1.angular_velocity += delta_omega;
        }
        if !body2.flags.is_kinematic() {
            body2.angular_velocity -= delta_omega;
        }

        let delta_v = (body2.linear_velocity - body1.linear_velocity)
            * (damping.linear * delta_secs).min(1.0);

        let w1 = inertia1.effective_inv_mass();
        let w2 = inertia2.effective_inv_mass();

        let p = delta_v * (w1 + w2).recip_or_zero();

        body1.linear_velocity += p * w1;
        body2.linear_velocity -= p * w2;
    }
}