use crate::math::Fix128;
use crate::solver::RigidBody;
#[cfg(not(feature = "std"))]
use alloc::vec;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SleepState {
Awake,
Sleeping,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SleepData {
pub state: SleepState,
pub idle_frames: u32,
}
impl SleepData {
#[must_use]
pub const fn new() -> Self {
Self {
state: SleepState::Awake,
idle_frames: 0,
}
}
#[inline]
#[must_use]
pub fn is_sleeping(&self) -> bool {
self.state == SleepState::Sleeping
}
#[inline]
pub fn wake(&mut self) {
self.state = SleepState::Awake;
self.idle_frames = 0;
}
}
impl Default for SleepData {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SleepConfig {
pub linear_threshold: Fix128,
pub angular_threshold: Fix128,
pub frames_to_sleep: u32,
}
impl Default for SleepConfig {
fn default() -> Self {
Self {
linear_threshold: Fix128::from_ratio(1, 100), angular_threshold: Fix128::from_ratio(1, 100), frames_to_sleep: 60, }
}
}
#[derive(Clone, Debug)]
pub struct Island {
pub bodies: Vec<usize>,
pub all_sleeping: bool,
}
pub struct IslandManager {
parent: Vec<usize>,
rank: Vec<u32>,
pub sleep_data: Vec<SleepData>,
pub config: SleepConfig,
}
impl IslandManager {
#[must_use]
pub fn new(num_bodies: usize, config: SleepConfig) -> Self {
let mut parent = Vec::with_capacity(num_bodies);
let rank = vec![0u32; num_bodies];
let sleep_data = vec![SleepData::new(); num_bodies];
for i in 0..num_bodies {
parent.push(i);
}
Self {
parent,
rank,
sleep_data,
config,
}
}
pub fn resize(&mut self, num_bodies: usize) {
while self.parent.len() < num_bodies {
let idx = self.parent.len();
self.parent.push(idx);
self.rank.push(0);
self.sleep_data.push(SleepData::new());
}
}
pub fn reset_unions(&mut self) {
for i in 0..self.parent.len() {
self.parent[i] = i;
self.rank[i] = 0;
}
}
pub fn find(&mut self, mut x: usize) -> usize {
while self.parent[x] != x {
self.parent[x] = self.parent[self.parent[x]]; x = self.parent[x];
}
x
}
pub fn union(&mut self, a: usize, b: usize) {
let ra = self.find(a);
let rb = self.find(b);
if ra == rb {
return;
}
match self.rank[ra].cmp(&self.rank[rb]) {
core::cmp::Ordering::Less => self.parent[ra] = rb,
core::cmp::Ordering::Greater => self.parent[rb] = ra,
core::cmp::Ordering::Equal => {
self.parent[rb] = ra;
self.rank[ra] += 1;
}
}
}
pub fn build_islands(&mut self) -> Vec<Island> {
let n = self.parent.len();
let mut island_map: Vec<Option<usize>> = vec![None; n];
let mut islands: Vec<Island> = Vec::new();
for i in 0..n {
let root = self.find(i);
let island_idx = island_map[root].unwrap_or_else(|| {
let idx = islands.len();
island_map[root] = Some(idx);
islands.push(Island {
bodies: Vec::new(),
all_sleeping: true,
});
idx
});
islands[island_idx].bodies.push(i);
if !self.sleep_data[i].is_sleeping() {
islands[island_idx].all_sleeping = false;
}
}
islands
}
pub fn update_sleep(&mut self, bodies: &[RigidBody]) {
for (i, body) in bodies.iter().enumerate() {
if i >= self.sleep_data.len() {
break;
}
self.update_sleep_body(i, body);
}
}
pub(crate) fn update_sleep_body(&mut self, i: usize, body: &RigidBody) {
if body.is_static() {
self.sleep_data[i].state = SleepState::Sleeping;
return;
}
let linear_speed = body.velocity.length();
let angular_speed = body.angular_velocity.length();
let is_idle = linear_speed < self.config.linear_threshold
&& angular_speed < self.config.angular_threshold;
if is_idle {
self.mark_idle(i);
} else {
self.sleep_data[i].state = SleepState::Awake;
self.sleep_data[i].idle_frames = 0;
}
}
pub(crate) fn update_sleep_idle(&mut self, i: usize, is_static: bool) {
if is_static {
self.sleep_data[i].state = SleepState::Sleeping;
} else {
self.mark_idle(i);
}
}
#[inline]
fn mark_idle(&mut self, i: usize) {
self.sleep_data[i].idle_frames += 1;
if self.sleep_data[i].idle_frames >= self.config.frames_to_sleep {
self.sleep_data[i].state = SleepState::Sleeping;
}
}
pub fn wake_island(&mut self, body_index: usize) {
let root = self.find(body_index);
for i in 0..self.parent.len() {
if self.find(i) == root {
self.sleep_data[i].wake();
}
}
}
#[inline]
pub fn wake_body(&mut self, body_index: usize) {
if body_index < self.sleep_data.len() {
self.sleep_data[body_index].wake();
}
}
#[inline]
pub fn is_sleeping(&self, body_index: usize) -> bool {
self.sleep_data
.get(body_index)
.is_some_and(SleepData::is_sleeping)
}
#[must_use]
pub fn sleeping_count(&self) -> usize {
self.sleep_data.iter().filter(|d| d.is_sleeping()).count()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::math::Vec3Fix;
#[test]
fn test_union_find() {
let config = SleepConfig::default();
let mut mgr = IslandManager::new(5, config);
mgr.union(0, 1);
mgr.union(2, 3);
assert_eq!(mgr.find(0), mgr.find(1));
assert_ne!(mgr.find(0), mgr.find(2));
assert_eq!(mgr.find(2), mgr.find(3));
}
#[test]
fn test_build_islands() {
let config = SleepConfig::default();
let mut mgr = IslandManager::new(5, config);
mgr.union(0, 1);
mgr.union(2, 3);
let islands = mgr.build_islands();
assert_eq!(islands.len(), 3, "Should have 3 islands");
let total_bodies: usize = islands.iter().map(|i| i.bodies.len()).sum();
assert_eq!(total_bodies, 5);
}
#[test]
fn test_sleep_after_threshold() {
let config = SleepConfig {
frames_to_sleep: 3,
..Default::default()
};
let mut mgr = IslandManager::new(1, config);
let still_body = RigidBody::new_static(Vec3Fix::ZERO);
let bodies = [still_body];
for _ in 0..3 {
mgr.update_sleep(&bodies);
}
assert!(mgr.is_sleeping(0));
}
#[test]
fn test_wake_on_motion() {
let config = SleepConfig {
linear_threshold: Fix128::from_ratio(1, 100),
frames_to_sleep: 2,
..Default::default()
};
let mut mgr = IslandManager::new(1, config);
let still_body = RigidBody::new(Vec3Fix::ZERO, Fix128::ONE);
mgr.update_sleep(&[still_body]);
mgr.update_sleep(&[still_body]);
assert!(mgr.is_sleeping(0));
let mut moving_body = RigidBody::new(Vec3Fix::ZERO, Fix128::ONE);
moving_body.velocity = Vec3Fix::from_int(10, 0, 0);
mgr.update_sleep(&[moving_body]);
assert!(!mgr.is_sleeping(0));
}
#[test]
fn test_wake_island() {
let config = SleepConfig::default();
let mut mgr = IslandManager::new(3, config);
for d in &mut mgr.sleep_data {
d.state = SleepState::Sleeping;
}
mgr.union(0, 1);
mgr.wake_island(0);
assert!(!mgr.is_sleeping(0));
assert!(!mgr.is_sleeping(1));
assert!(mgr.is_sleeping(2)); }
#[test]
fn test_resize() {
let config = SleepConfig::default();
let mut mgr = IslandManager::new(2, config);
mgr.resize(5);
assert_eq!(mgr.parent.len(), 5);
assert_eq!(mgr.sleep_data.len(), 5);
}
#[test]
fn test_sleeping_count() {
let config = SleepConfig::default();
let mut mgr = IslandManager::new(5, config);
mgr.sleep_data[0].state = SleepState::Sleeping;
mgr.sleep_data[2].state = SleepState::Sleeping;
mgr.sleep_data[4].state = SleepState::Sleeping;
assert_eq!(mgr.sleeping_count(), 3);
}
}