feat: represent forces as objects

4 files changed, 63 insertions, 8 deletions

diff --git a/src/force/drag.rs b/src/force/drag.rs
new file mode 100644
index 0000000..ef859a0
--- /dev/null
+++ b/src/force/drag.rs
@@ -0,0 +1,15 @@
+use crate::particle_system::Scalar;
+
+use super::Force;
+
+pub struct Drag {
+    pub coefficient: Scalar,
+}
+
+impl Force for Drag {
+    fn apply(&self, particles: &mut Vec<crate::particle_system::Particle>) {
+        for particle in particles {
+            particle.apply_force(-self.coefficient * particle.velocity);
+        }
+    }
+}
diff --git a/src/force/gravity.rs b/src/force/gravity.rs
new file mode 100644
index 0000000..b023e63
--- /dev/null
+++ b/src/force/gravity.rs
@@ -0,0 +1,15 @@
+use crate::particle_system::Vector;
+
+use super::Force;
+
+pub struct Gravity {
+    pub vector: Vector,
+}
+
+impl Force for Gravity {
+    fn apply(&self, particles: &mut Vec<crate::particle_system::Particle>) {
+        for particle in particles {
+            particle.apply_force(self.vector * particle.mass);
+        }
+    }
+}
diff --git a/src/force/mod.rs b/src/force/mod.rs
new file mode 100644
index 0000000..ce10f9f
--- /dev/null
+++ b/src/force/mod.rs
@@ -0,0 +1,24 @@
+use crate::particle_system::{Particle, ParticleSystem};
+
+pub mod gravity;
+pub mod drag;
+pub mod spring;
+
+pub trait Force {
+    fn apply(&self, particles: &mut Vec<Particle>);
+}
+
+impl ParticleSystem {
+    fn reset_forces(&mut self) {
+        for particle in &mut self.particles {
+            particle.reset_force();
+        }
+    }
+    pub fn apply_forces(&mut self) {
+        self.reset_forces();
+
+        for force in &self.forces {
+            force.apply(&mut self.particles)
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 718abf8..58e649f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,6 @@
 use std::path::PathBuf;
 
+use force::{drag::Drag, gravity::Gravity};
 use particle_system::{Particle, ParticleSystem, Point, Vector};
 use ppm::PPM;
 use solver::Solver;
@@ -8,6 +9,7 @@ mod constraint;
 mod particle_system;
 mod ppm;
 mod solver;
+mod force;
 
 fn main() {
     let ppm = PPM {
@@ -25,6 +27,12 @@ fn main() {
             Particle::new(Point::new(5.0, 20.0), 50.0),
         ],
         constraints: vec![],
+        forces: vec![
+            Box::new(Gravity {
+                vector: Vector::y() * -9.8,
+            }),
+            Box::new(Drag { coefficient: 0.0 }),
+        ],
         t: 0.0,
     };
 
@@ -34,15 +42,8 @@ fn main() {
     system.add_beam_constraint([1, 3]);
     system.add_beam_constraint([2, 3]);
 
-    let gravity = -9.8 * Vector::y();
-
     for i in 0..150_00 {
-        for particle in &mut system.particles {
-            particle.reset_force();
-            // Apply custom forces
-            particle.apply_force(gravity * particle.mass);
-        }
-        system.particles[0].apply_force(gravity); // Extra force on particle 0
+        system.apply_forces();
 
         if i % 10 == 0 {
             println!("Iteration #{i}");