quadtree, pancam

src/quadtree.rs

@@ -2,40 +2,159 @@ use bevy::prelude::*;
 
 pub trait Body {
 	fn mass(&self) -> f32;
-	fn center_of_mass(&self) -> (f32, Vec2);
+	fn pos(&self) -> Vec2;
+	fn add_mass(&mut self, mass: f32);
 }
 
-pub enum Node<L: Body> {
-	Node([Node; 4]),
-	Leaf(Vec<L>),
-}
-
-impl<L: Body> Body for Node<L> {
-	fn mass(&self) -> f32 {
-		match self {
-			Node::Node([n1, n2, n3, n4]) => n0.mass() + n1.mass() + n2.mass() + n3.mass(),
-			Node::Leaf(v) => v.iter().map(Body::mass).sum(),
-		}
-	}
-	fn center_of_mass(&self) {
-		match self {
-			Node::Node([n1, n2, n3, n4]) => n0.mass() + n1.mass() + n2.mass() + n3.mass(),
-			Node::Leaf(v) => {
-				let mut mass = 0.0;
-				let mut center_of_mass = Vec2::zero();
-				for(n in v) {
-					let (n_mass, n_center) = n.center_of_mass();
-					mass += n_mass;
-					center_of_mass += mass * n_center;
-				}
-				(mass, center_of_mass / mass)
-			},
-		}
-	}
+pub enum Node<L> {
+	Branch {
+		nodes: Box<[Node<L>; 4]>,
+		center: Vec2,
+		mass: f32,
+		center_of_mass: Vec2,
+		width: f32,
+	},
+	Leaf {
+		body: Option<L>,
+		pos: (Vec2, Vec2),
+	},
 }
 
 impl<L: Body> Node<L> {
-	fn add_body(&mut self, body: L) {
-		
+	pub fn new(pos: (Vec2, Vec2)) -> Self {
+		Node::Leaf { body: None, pos }
+		// let center = (pos.1 - pos.0) / 2.0;
+		// Node::Branch {
+		// 	nodes: [
+		// 		Box::new(Node::Leaf {
+		// 			body: None,
+		// 			pos: (pos.0, center),
+		// 		}),
+		// 		Box::new(Node::Leaf {
+		// 			body: None,
+		// 			pos: (Vec2::new(center.x, pos.0.y), Vec2::new(pos.1.x, center.y)),
+		// 		}),
+		// 		Box::new(Node::Leaf {
+		// 			body: None,
+		// 			pos: (Vec2::new(pos.0.x, center.y), Vec2::new(center.x, pos.1.y)),
+		// 		}),
+		// 		Box::new(Node::Leaf {
+		// 			body: None,
+		// 			pos: (center, pos.1),
+		// 		}),
+		// 	],
+		// 	center,
+		// 	mass: 0.0,
+		// 	center_of_mass: center,
+		// 	width: pos.1.x - pos.0.x,
+		// }
+	}
+	pub fn add_body(&mut self, new_body: L) {
+		match self {
+			Node::Branch {
+				nodes,
+				center,
+				mass,
+				center_of_mass,
+				..
+			} => {
+				let new_body_pos = new_body.pos();
+				let new_body_mass = new_body.mass();
+				*center_of_mass = (*center_of_mass * *mass + new_body_mass * new_body_pos)
+					/ (*mass + new_body_mass);
+				*mass += new_body_mass;
+				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)
+			}
+			Node::Leaf { body, pos } => {
+				if let Some(mut body) = body.take() {
+					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;
+					}
+					let center = (pos.0 + pos.1) / 2.0;
+					*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,
+					};
+					self.add_body(body);
+					self.add_body(new_body)
+				} else {
+					*body = Some(new_body);
+				}
+			}
+		}
+	}
+
+	pub fn apply(&self, on: Vec2, theta: f32) -> Vec2 {
+		match self {
+			Node::Branch {
+				nodes,
+				mass,
+				center_of_mass,
+				width,
+				..
+			} => {
+				if on == *center_of_mass {
+					return Vec2::ZERO;
+				}
+				let dist = on.distance(*center_of_mass);
+				if width / dist < theta {
+					*mass * (*center_of_mass - on) / (dist * dist * dist)
+				} else {
+					nodes[0].apply(on, theta)
+						+ nodes[1].apply(on, theta)
+						+ nodes[2].apply(on, theta)
+						+ nodes[3].apply(on, theta)
+				}
+			}
+			Node::Leaf { body, .. } => {
+				if let Some(body) = body {
+					if on == body.pos() {
+						return Vec2::ZERO;
+					}
+					let dist = on.distance(body.pos());
+					body.mass() * (body.pos() - on) / (dist * dist * dist)
+				} else {
+					Vec2::ZERO
+				}
+			}
+		}
 	}
 }