archimedes.rs

use avian3d::prelude::*;
use bevy::prelude::*;

pub struct LiquidsPlugin;

#[derive(Default, Component, Reflect, Debug)]
#[reflect(Default, Component, Debug)]
pub struct Liquid {
    pub kind: LiquidKind,
}

#[derive(Default, Reflect, Debug, Eq, PartialEq)]
#[reflect(Default, PartialEq)]
pub enum LiquidKind {
    #[default]
    Water,
}

impl Liquid {
    pub fn damping_factor(&self) -> f32 {
        match self.kind {
            LiquidKind::Water => 0.25,
        }
    }
}

#[derive(Component, Reflect, Debug)]
#[reflect(Component, Debug)]
#[component(storage = "SparseSet")]
pub struct Submerged {
    pub entity: Entity,
    pub estimated_percent: f32,
    pub estimated_volume: f32,
    pub delta_percent: f32,
    pub delta_volume: f32,
}

impl Plugin for LiquidsPlugin {
    fn build(&self, app: &mut App) {
        app.register_type::<LiquidKind>()
            .register_type::<Liquid>()
            .register_type::<Submerged>()
            .add_systems(
                Update,
                (update_submerged, update_estimated_percent_and_volume).chain(),
            )
            .add_systems(
                FixedUpdate,
                apply_buoyancy_and_drag.in_set(PhysicsStepSet::First),
            );
    }
}

fn update_submerged(
    mut commands: Commands,
    mut collision_started: EventReader<CollisionStarted>,
    mut collision_ended: EventReader<CollisionEnded>,
    liquid_query: Query<&Liquid>,
    rigid_body_query: Query<&RigidBody>,
    submerged_query: Query<&Submerged>,
) {
    for event in collision_ended.read() {
        if liquid_query.contains(event.0)
            && submerged_query
                .get(event.1)
                .is_ok_and(|submerged| submerged.entity == event.0)
        {
            commands.entity(event.1).remove::<Submerged>();
        } else if liquid_query.contains(event.1)
            && submerged_query
                .get(event.0)
                .is_ok_and(|submerged| submerged.entity == event.1)
        {
            commands.entity(event.0).remove::<Submerged>();
        }
    }

    for event in collision_started.read() {
        if liquid_query.contains(event.0) && rigid_body_query.get(event.1).is_ok() {
            commands.entity(event.1).try_insert(Submerged {
                entity: event.0,
                estimated_percent: 0.0,
                estimated_volume: 0.0,
                delta_percent: 0.0,
                delta_volume: 0.0,
            });
        } else if liquid_query.contains(event.1) && rigid_body_query.get(event.0).is_ok() {
            commands.entity(event.0).try_insert(Submerged {
                entity: event.1,
                estimated_percent: 0.0,
                estimated_volume: 0.0,
                delta_percent: 0.0,
                delta_volume: 0.0,
            });
        }
    }
}

fn update_estimated_percent_and_volume(
    collisions: Res<Collisions>,
    mut submerged_query: Query<(Entity, &mut Submerged, &Collider, &GlobalTransform)>,
) {
    for (entity, mut submerged, collider, global_transform) in &mut submerged_query {
        let Some(contacts) = collisions.get(entity, submerged.entity) else {
            continue;
        };

        let aabb = collider.aabb(global_transform.translation(), global_transform.rotation());
        let size = aabb.size();

        if let Some(deepest_contact) = contacts.find_deepest_contact() {
            let contact = if contacts.entity1 == entity {
                deepest_contact
            } else {
                &deepest_contact.flipped()
            };

            let estimated_percent = ((contact.penetration
                * contact.global_normal1(&Rotation::from(global_transform.rotation())))
            .abs()
                / size)
                .length()
                .min(1.0);

            // with a density of 1.0, the mass equals the volume
            let estimated_volume =
                collider.shape_scaled().mass_properties(1.0).mass() * estimated_percent;

            submerged.delta_percent = estimated_percent - submerged.estimated_percent;
            submerged.delta_volume = estimated_volume - submerged.estimated_volume;

            submerged.estimated_percent = estimated_percent;

            submerged.estimated_volume = estimated_volume;
        }
    }
}

fn apply_buoyancy_and_drag(
    mut commands: Commands,
    mut submerged_rigid_body_query: Query<(
        Entity,
        &RigidBody,
        &Submerged,
        &mut LinearVelocity,
        &mut AngularVelocity,
    )>,
    liquid_query: Query<(&Liquid, &ColliderDensity)>,
    gravity: Res<Gravity>,
    time: Res<Time>,
) {
    for (entity, rigid_body, submerged, mut linear_velocity, mut angular_velocity) in
        &mut submerged_rigid_body_query
    {
        if !rigid_body.is_dynamic() {
            continue;
        }

        let Ok((liquid, liquid_density)) = liquid_query.get(submerged.entity) else {
            continue;
        };

        if submerged.delta_percent > 0.0 {
            linear_velocity.0 *= 1.0 - submerged.delta_percent;
        }

        let damping_factor = liquid.damping_factor();

        linear_velocity.0 *= damping_factor.powf(time.delta_secs() * submerged.estimated_percent);
        angular_velocity.0 *= damping_factor.powf(time.delta_secs() * submerged.estimated_percent);

        commands.entity(entity).try_insert(
            ExternalForce::new(liquid_density.0 * submerged.estimated_volume * -gravity.0)
                .with_persistence(false),
        );
    }
}