src/test/bench/shootout-nbody.rs

   1 // The Computer Language Benchmarks Game
   2 // http://benchmarksgame.alioth.debian.org/
   3 //
   4 // contributed by the Rust Project Developers
   5
   6 // Copyright (c) 2011-2014 The Rust Project Developers
   7 //
   8 // All rights reserved.
   9 //
  10 // Redistribution and use in source and binary forms, with or without
  11 // modification, are permitted provided that the following conditions
  12 // are met:
  13 //
  14 // - Redistributions of source code must retain the above copyright
  15 //   notice, this list of conditions and the following disclaimer.
  16 //
  17 // - Redistributions in binary form must reproduce the above copyright
  18 //   notice, this list of conditions and the following disclaimer in
  19 //   the documentation and/or other materials provided with the
  20 //   distribution.
  21 //
  22 // - Neither the name of "The Computer Language Benchmarks Game" nor
  23 //   the name of "The Computer Language Shootout Benchmarks" nor the
  24 //   names of its contributors may be used to endorse or promote
  25 //   products derived from this software without specific prior
  26 //   written permission.
  27 //
  28 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  29 // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  30 // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
  31 // FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
  32 // COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
  33 // INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  34 // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  35 // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  36 // HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
  37 // STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  38 // ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
  39 // OF THE POSSIBILITY OF SUCH DAMAGE.
  40
  41 use std::mem;
  42 use std::ops::{Add, Sub, Mul};
  43
  44 const PI: f64 = 3.141592653589793;
  45 const SOLAR_MASS: f64 = 4.0 * PI * PI;
  46 const YEAR: f64 = 365.24;
  47 const N_BODIES: usize = 5;
  48 const N_PAIRS: usize = N_BODIES * (N_BODIES - 1) / 2;
  49
  50 const BODIES: [Planet; N_BODIES] = [
  51     // Sun
  52     Planet {
  53         pos: Vec3(0.0, 0.0, 0.0),
  54         vel: Vec3(0.0, 0.0, 0.0),
  55         mass: SOLAR_MASS,
  56     },
  57     // Jupiter
  58     Planet {
  59         pos: Vec3(4.84143144246472090e+00,
  60                   -1.16032004402742839e+00,
  61                   -1.03622044471123109e-01),
  62         vel: Vec3(1.66007664274403694e-03 * YEAR,
  63                   7.69901118419740425e-03 * YEAR,
  64                   -6.90460016972063023e-05 * YEAR),
  65         mass: 9.54791938424326609e-04 * SOLAR_MASS,
  66     },
  67     // Saturn
  68     Planet {
  69         pos: Vec3(8.34336671824457987e+00,
  70                   4.12479856412430479e+00,
  71                   -4.03523417114321381e-01),
  72         vel: Vec3(-2.76742510726862411e-03 * YEAR,
  73                   4.99852801234917238e-03 * YEAR,
  74                   2.30417297573763929e-05 * YEAR),
  75         mass: 2.85885980666130812e-04 * SOLAR_MASS,
  76     },
  77     // Uranus
  78     Planet {
  79         pos: Vec3(1.28943695621391310e+01,
  80                   -1.51111514016986312e+01,
  81                   -2.23307578892655734e-01),
  82         vel: Vec3(2.96460137564761618e-03 * YEAR,
  83                   2.37847173959480950e-03 * YEAR,
  84                   -2.96589568540237556e-05 * YEAR),
  85         mass: 4.36624404335156298e-05 * SOLAR_MASS,
  86     },
  87     // Neptune
  88     Planet {
  89         pos: Vec3(1.53796971148509165e+01,
  90                   -2.59193146099879641e+01,
  91                   1.79258772950371181e-01),
  92         vel: Vec3(2.68067772490389322e-03 * YEAR,
  93                   1.62824170038242295e-03 * YEAR,
  94                   -9.51592254519715870e-05 * YEAR),
  95         mass: 5.15138902046611451e-05 * SOLAR_MASS,
  96     },
  97 ];
  98
  99 /// A 3d Vector type with oveloaded operators to improve readability.
 100 #[derive(Clone, Copy)]
 101 struct Vec3(pub f64, pub f64, pub f64);
 102
 103 impl Vec3 {
 104     fn zero() -> Self { Vec3(0.0, 0.0, 0.0) }
 105
 106     fn norm(&self) -> f64 { self.squared_norm().sqrt() }
 107
 108     fn squared_norm(&self) -> f64 {
 109         self.0 * self.0 + self.1 * self.1 + self.2 * self.2
 110     }
 111 }
 112
 113 impl Add for Vec3 {
 114     type Output = Self;
 115     fn add(self, rhs: Self) -> Self {
 116         Vec3(self.0 + rhs.0, self.1 + rhs.1, self.2 + rhs.2)
 117     }
 118 }
 119
 120 impl Sub for Vec3 {
 121     type Output = Self;
 122     fn sub(self, rhs: Self) -> Self {
 123         Vec3(self.0 - rhs.0, self.1 - rhs.1, self.2 - rhs.2)
 124     }
 125 }
 126
 127 impl Mul<f64> for Vec3 {
 128     type Output = Self;
 129     fn mul(self, rhs: f64) -> Self {
 130         Vec3(self.0 * rhs, self.1 * rhs, self.2 * rhs)
 131     }
 132 }
 133
 134 #[derive(Clone, Copy)]
 135 struct Planet {
 136     pos: Vec3,
 137     vel: Vec3,
 138     mass: f64,
 139 }
 140
 141 /// Computes all pairwise position differences between the planets.
 142 fn pairwise_diffs(bodies: &[Planet; N_BODIES], diff: &mut [Vec3; N_PAIRS]) {
 143     let mut bodies = bodies.iter();
 144     let mut diff = diff.iter_mut();
 145     while let Some(bi) = bodies.next() {
 146         for bj in bodies.clone() {
 147             *diff.next().unwrap() = bi.pos - bj.pos;
 148         }
 149     }
 150 }
 151
 152 /// Computes the magnitude of the force between each pair of planets.
 153 fn magnitudes(diff: &[Vec3; N_PAIRS], dt: f64, mag: &mut [f64; N_PAIRS]) {
 154     for (mag, diff) in mag.iter_mut().zip(diff.iter()) {
 155         let d2 = diff.squared_norm();
 156         *mag = dt / (d2 * d2.sqrt());
 157     }
 158 }
 159
 160 /// Updates the velocities of the planets by computing their gravitational
 161 /// accelerations and performing one step of Euler integration.
 162 fn update_velocities(bodies: &mut [Planet; N_BODIES], dt: f64,
 163                      diff: &mut [Vec3; N_PAIRS], mag: &mut [f64; N_PAIRS]) {
 164     pairwise_diffs(bodies, diff);
 165     magnitudes(&diff, dt, mag);
 166
 167     let mut bodies = &mut bodies[..];
 168     let mut mag = mag.iter();
 169     let mut diff = diff.iter();
 170     while let Some(bi) = shift_mut_ref(&mut bodies) {
 171         for bj in bodies.iter_mut() {
 172             let diff = *diff.next().unwrap();
 173             let mag = *mag.next().unwrap();
 174             bi.vel = bi.vel - diff * (bj.mass * mag);
 175             bj.vel = bj.vel + diff * (bi.mass * mag);
 176         }
 177     }
 178 }
 179
 180 /// Advances the solar system by one timestep by first updating the
 181 /// velocities and then integrating the positions using the updated velocities.
 182 ///
 183 /// Note: the `diff` & `mag` arrays are effectively scratch space. They're
 184 /// provided as arguments to avoid re-zeroing them every time `advance` is
 185 /// called.
 186 fn advance(mut bodies: &mut [Planet; N_BODIES], dt: f64,
 187            diff: &mut [Vec3; N_PAIRS], mag: &mut [f64; N_PAIRS]) {
 188     update_velocities(bodies, dt, diff, mag);
 189     for body in bodies.iter_mut() {
 190         body.pos = body.pos + body.vel * dt;
 191     }
 192 }
 193
 194 /// Computes the total energy of the solar system.
 195 fn energy(bodies: &[Planet; N_BODIES]) -> f64 {
 196     let mut e = 0.0;
 197     let mut bodies = bodies.iter();
 198     while let Some(bi) = bodies.next() {
 199         e += bi.vel.squared_norm() * bi.mass / 2.0
 200            - bi.mass * bodies.clone()
 201                              .map(|bj| bj.mass / (bi.pos - bj.pos).norm())
 202                              .fold(0.0, |a, b| a + b);
 203     }
 204     e
 205 }
 206
 207 /// Offsets the sun's velocity to make the overall momentum of the system zero.
 208 fn offset_momentum(bodies: &mut [Planet; N_BODIES]) {
 209     let p = bodies.iter().fold(Vec3::zero(), |v, b| v + b.vel * b.mass);
 210     bodies[0].vel = p * (-1.0 / bodies[0].mass);
 211 }
 212
 213 fn main() {
 214     let n = if std::env::var_os("RUST_BENCH").is_some() {
 215         5000000
 216     } else {
 217         std::env::args().nth(1)
 218             .and_then(|arg| arg.parse().ok())
 219             .unwrap_or(1000)
 220     };
 221     let mut bodies = BODIES;
 222     let mut diff = [Vec3::zero(); N_PAIRS];
 223     let mut mag = [0.0f64; N_PAIRS];
 224
 225     offset_momentum(&mut bodies);
 226     println!("{:.9}", energy(&bodies));
 227
 228     for _ in (0..n) {
 229         advance(&mut bodies, 0.01, &mut diff, &mut mag);
 230     }
 231
 232     println!("{:.9}", energy(&bodies));
 233 }
 234
 235 /// Pop a mutable reference off the head of a slice, mutating the slice to no
 236 /// longer contain the mutable reference. This is a safe operation because the
 237 /// two mutable borrows are entirely disjoint.
 238 fn shift_mut_ref<'a, T>(r: &mut &'a mut [T]) -> Option<&'a mut T> {
 239     let res = mem::replace(r, &mut []);
 240     if res.is_empty() { return None }
 241     let (a, b) = res.split_at_mut(1);
 242     *r = b;
 243     Some(&mut a[0])
 244 }