jsb-gravity/src/quadtree.rs

use bevy::prelude::*;

/// Un trait est la définition abstraite d'un comportement.
/// Tout ce qui implémente ce trait possède les méthodes mass, pos, add_mass.
pub trait Body {
	/// Quelle est la masse en kg du corps ?
	fn mass(&self) -> f32;
	/// Quelle est la position (x, y) du corps ?
	fn pos(&self) -> Vec2;
	/// Ajouter de la masse dans le corps.
	fn add_mass(&mut self, mass: f32);
}

/// Nœud de l'arbre, étant soit une feuille soit une branche.
/// Sur chaque branche pousse des branches et des feuilles.
/// Rien ne pousse sur une feuille.
/// Un nœud représente une portion rectangulaire de l'univers, contenant tous les objets qui sont dans ce rectangle.
pub enum Node<L> {
	/// Branche
	Branch {
		/// 4 nœuds contenus par cette branche, séparant son espace en 4 cadrants de même taille
		nodes: Box<[Node<L>; 4]>,
		/// Position du centre du rectangle
		center: Vec2,
		/// Masse cumulée de tous les objets dans le nœud
		mass: f32,
		/// Centre de masse d'ensemble des objets contenus
		center_of_mass: Vec2,
		/// Largeur de notre rectangle
		width: f32,
	},
	/// Feuille
	Leaf {
		/// Contient soit un corps Some(body) soit rien None
		body: Option<L>,
		/// Position de la feuille (pas celle du corps, qui peut être à des endroits différents dans la feuille)
		pos: (Vec2, Vec2),
	},
}

/// Ajoutons des fonctions aux nœuds
impl<L: Body> Node<L> {
	/// Création d'un nouveau nœud vide
	pub fn new(pos: (Vec2, Vec2)) -> Self {
		Node::Leaf { body: None, pos }
	}
	/// Ajouter un corps dans le nœud
	pub fn add_body(&mut self, new_body: L) {
		match self {
			// Si le nœud est une branche...
			Node::Branch {
				nodes,
				center,
				mass,
				center_of_mass,
				..
			} => {
				let new_body_pos = new_body.pos();
				let new_body_mass = new_body.mass();
				// Calculer le centre de masse des corps dans la branche plus le nouveau corps
				*center_of_mass = (*center_of_mass * *mass + new_body_mass * new_body_pos)
					/ (*mass + new_body_mass);
				*mass += new_body_mass;
				// Trouver le bon cadrant où ajouter le corps
				nodes[if new_body_pos.x < center.x {
					if new_body_pos.y < center.y {
						0
					} else {
						2
					}
				} else {
					if new_body_pos.y < center.y {
						1
					} else {
						3
					}
				}]
				.add_body(new_body)
			}
			// Si le nœud est une feuille...
			Node::Leaf { body, pos } => {
				// Si la feuille contient un corps...
				if let Some(mut body) = body.take() {
					// Si les deux corps sont très proches, on évite de créer plein de branches ce qui ferait tout planter.
					// Dans ce cas on fusionne les deux
					if body.pos().distance_squared(new_body.pos()) < 1.0 {
						body.add_mass(new_body.mass());
						*self = Node::Leaf {
							body: Some(body),
							pos: *pos,
						};
						return;
					}
					// Cette feuille va devenir une branche.
					// On calcule donc son centre.
					let center = (pos.0 + pos.1) / 2.0;
					// Et on la remplace par une branche, pour l'instant avec ses 4 cadrants vides.
					*self = Node::Branch {
						nodes: Box::new([
							Node::Leaf {
								body: None,
								pos: (pos.0, center),
							},
							Node::Leaf {
								body: None,
								pos: (Vec2::new(center.x, pos.0.y), Vec2::new(pos.1.x, center.y)),
							},
							Node::Leaf {
								body: None,
								pos: (Vec2::new(pos.0.x, center.y), Vec2::new(center.x, pos.1.y)),
							},
							Node::Leaf {
								body: None,
								pos: (center, pos.1),
							},
						]),
						center,
						mass: 0.0,
						center_of_mass: center,
						width: pos.1.x - pos.0.x,
					};
					
					// On ajoute les deux corps dans la branche.
					self.add_body(body);
					self.add_body(new_body)
				} else {
					// Ici est le cas où la feuille était vide, alors on met juste le nouveau corps dedans.
					*body = Some(new_body);
				}
			}
		}
	}

	/// Calculer la force de gravité s'appliquant sur le point `on`
	/// `theta` est un nombre entre 0.0 et 1.0, qui détermine si on veut plus de précision ou plus de rapidité.
	pub fn apply(&self, on: Vec2, theta: f32) -> Vec2 {
		match self {
			// Si le nœud est une branche...
			Node::Branch {
				nodes,
				mass,
				center_of_mass,
				width,
				..
			} => {
				if on == *center_of_mass {
					// Dans ce cas la distance est nulle, donc on évite de diviser par zéro.
					return Vec2::ZERO;
				}
				let dist = on.distance(*center_of_mass);
				if width / dist < theta {
					// On est dans le cas où on est assez loin pour pouvoir faire une approximation.
					// On fait comme si le nœud était un gros corps, la fusion de tous les corps qu'il contient.
					*mass * (*center_of_mass - on) / (dist * dist * dist)
				} else {
					// On est dans le cas où on est trop près pour faire l'approximation.
					// On applique alors récursivement pour chaque sous-nœud.
					nodes[0].apply(on, theta)
						+ nodes[1].apply(on, theta)
						+ nodes[2].apply(on, theta)
						+ nodes[3].apply(on, theta)
				}
			}
			// Si le nœud est une feuille...
			Node::Leaf { body, .. } => {
				if let Some(body) = body {
					// La feuille contient un corps.
					if on == body.pos() {
						return Vec2::ZERO;
					}
					let dist = on.distance(body.pos());
					body.mass() * (body.pos() - on) / (dist * dist * dist)
				} else {
					// La feuille est vide.
					Vec2::ZERO
				}
			}
		}
	}
}