From 9738c2a45c03034f937a0b2ce7a690f64e7662eb Mon Sep 17 00:00:00 2001
From: Patrick Walton <pcwalton@mimiga.net>
Date: Tue, 16 Apr 2013 16:11:16 -0700
Subject: [PATCH] test: Rewrite nbody and spectralnorm shootout benchmarks

---
 src/libcore/core.rc                     |   1 +
 src/libcore/iter.rs                     |   4 +
 src/libcore/prelude.rs                  |   2 +-
 src/libcore/vec.rs                      |  41 ++-
 src/test/bench/shootout-nbody.rs        | 384 +++++++++---------------
 src/test/bench/shootout-spectralnorm.rs | 104 +++----
 6 files changed, 218 insertions(+), 318 deletions(-)

diff --git a/src/libcore/core.rc b/src/libcore/core.rc
index e7a5cfbaf4b..4080c70c8ae 100644
--- a/src/libcore/core.rc
+++ b/src/libcore/core.rc
@@ -98,6 +98,7 @@ pub use vec::{ImmutableEqVector, ImmutableCopyableVector};
 pub use vec::{OwnedVector, OwnedCopyableVector};
 pub use iter::{BaseIter, ExtendedIter, EqIter, CopyableIter};
 pub use iter::{CopyableOrderedIter, CopyableNonstrictIter, Times};
+pub use iter::{ExtendedMutableIter};
 
 pub use num::{Num, NumCast};
 pub use ptr::Ptr;
diff --git a/src/libcore/iter.rs b/src/libcore/iter.rs
index a220cd520c3..3dcca0e06c2 100644
--- a/src/libcore/iter.rs
+++ b/src/libcore/iter.rs
@@ -45,6 +45,10 @@ pub trait ExtendedIter<A> {
     fn flat_map_to_vec<B,IB: BaseIter<B>>(&self, op: &fn(&A) -> IB) -> ~[B];
 }
 
+pub trait ExtendedMutableIter<A> {
+    fn eachi_mut(&mut self, blk: &fn(uint, &mut A) -> bool);
+}
+
 pub trait EqIter<A:Eq> {
     fn contains(&self, x: &A) -> bool;
     fn count(&self, x: &A) -> uint;
diff --git a/src/libcore/prelude.rs b/src/libcore/prelude.rs
index e148493ca45..95481d032a4 100644
--- a/src/libcore/prelude.rs
+++ b/src/libcore/prelude.rs
@@ -33,7 +33,7 @@ pub use container::{Container, Mutable, Map, Set};
 pub use hash::Hash;
 pub use iter::{BaseIter, ReverseIter, MutableIter, ExtendedIter, EqIter};
 pub use iter::{CopyableIter, CopyableOrderedIter, CopyableNonstrictIter};
-pub use iter::Times;
+pub use iter::{Times, ExtendedMutableIter};
 pub use num::{Num, NumCast};
 pub use path::GenericPath;
 pub use path::Path;
diff --git a/src/libcore/vec.rs b/src/libcore/vec.rs
index f6492ede9f9..3e0ba38a6e8 100644
--- a/src/libcore/vec.rs
+++ b/src/libcore/vec.rs
@@ -1388,13 +1388,19 @@ pub fn each<'r,T>(v: &'r [T], f: &fn(&'r T) -> bool) {
 /// to mutate the contents as you iterate.
 #[inline(always)]
 pub fn each_mut<'r,T>(v: &'r mut [T], f: &fn(elem: &'r mut T) -> bool) {
-    let mut i = 0;
-    let n = v.len();
-    while i < n {
-        if !f(&mut v[i]) {
-            return;
+    do vec::as_mut_buf(v) |p, n| {
+        let mut n = n;
+        let mut p = p;
+        while n > 0 {
+            unsafe {
+                let q: &'r mut T = cast::transmute_mut_region(&mut *p);
+                if !f(q) {
+                    break;
+                }
+                p = p.offset(1);
+            }
+            n -= 1;
         }
-        i += 1;
     }
 }
 
@@ -1426,6 +1432,22 @@ pub fn eachi<'r,T>(v: &'r [T], f: &fn(uint, v: &'r T) -> bool) {
     }
 }
 
+/**
+ * Iterates over a mutable vector's elements and indices
+ *
+ * Return true to continue, false to break.
+ */
+#[inline(always)]
+pub fn eachi_mut<'r,T>(v: &'r mut [T], f: &fn(uint, v: &'r mut T) -> bool) {
+    let mut i = 0;
+    for each_mut(v) |p| {
+        if !f(i, p) {
+            return;
+        }
+        i += 1;
+    }
+}
+
 /**
  * Iterates over a vector's elements in reverse
  *
@@ -2654,6 +2676,13 @@ impl<'self,A> iter::ExtendedIter<A> for &'self [A] {
     }
 }
 
+impl<'self,A> iter::ExtendedMutableIter<A> for &'self mut [A] {
+    #[inline(always)]
+    pub fn eachi_mut(&mut self, blk: &fn(uint, v: &mut A) -> bool) {
+        eachi_mut(*self, blk)
+    }
+}
+
 // FIXME(#4148): This should be redundant
 impl<A> iter::ExtendedIter<A> for ~[A] {
     pub fn eachi(&self, blk: &fn(uint, v: &A) -> bool) {
diff --git a/src/test/bench/shootout-nbody.rs b/src/test/bench/shootout-nbody.rs
index 97907025bd1..e633f307bc2 100644
--- a/src/test/bench/shootout-nbody.rs
+++ b/src/test/bench/shootout-nbody.rs
@@ -1,254 +1,150 @@
-// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
-// file at the top-level directory of this distribution and at
-// http://rust-lang.org/COPYRIGHT.
-//
-// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
-// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
-// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
-// option. This file may not be copied, modified, or distributed
-// except according to those terms.
+use core::from_str::FromStr;
+use core::uint::range;
+use core::unstable::intrinsics::sqrtf64;
 
-// based on:
-// http://shootout.alioth.debian.org/u32/benchmark.php?test=nbody&lang=java
+static PI: f64 = 3.141592653589793;
+static SOLAR_MASS: f64 = 4.0 * PI * PI;
+static YEAR: f64 = 365.24;
+static N_BODIES: uint = 5;
 
-extern mod std;
+static BODIES: [Planet, ..N_BODIES] = [
+    // Sun
+    Planet {
+        x: [ 0.0, 0.0, 0.0 ],
+        v: [ 0.0, 0.0, 0.0 ],
+        mass: SOLAR_MASS,
+    },
+    // Jupiter
+    Planet {
+        x: [
+            4.84143144246472090e+00,
+            -1.16032004402742839e+00,
+            -1.03622044471123109e-01,
+        ],
+        v: [
+            1.66007664274403694e-03 * YEAR,
+            7.69901118419740425e-03 * YEAR,
+            -6.90460016972063023e-05 * YEAR,
+        ],
+        mass: 9.54791938424326609e-04 * SOLAR_MASS,
+    },
+    // Saturn
+    Planet {
+        x: [
+            8.34336671824457987e+00,
+            4.12479856412430479e+00,
+            -4.03523417114321381e-01,
+        ],
+        v: [
+            -2.76742510726862411e-03 * YEAR,
+            4.99852801234917238e-03 * YEAR,
+            2.30417297573763929e-05 * YEAR,
+        ],
+        mass: 2.85885980666130812e-04 * SOLAR_MASS,
+    },
+    // Uranus
+    Planet {
+        x: [
+            1.28943695621391310e+01,
+            -1.51111514016986312e+01,
+            -2.23307578892655734e-01,
+        ],
+        v: [
+            2.96460137564761618e-03 * YEAR,
+            2.37847173959480950e-03 * YEAR,
+            -2.96589568540237556e-05 * YEAR,
+        ],
+        mass: 4.36624404335156298e-05 * SOLAR_MASS,
+    },
+    // Neptune
+    Planet {
+        x: [
+            1.53796971148509165e+01,
+            -2.59193146099879641e+01,
+            1.79258772950371181e-01,
+        ],
+        v: [
+            2.68067772490389322e-03 * YEAR,
+            1.62824170038242295e-03 * YEAR,
+            -9.51592254519715870e-05 * YEAR,
+        ],
+        mass: 5.15138902046611451e-05 * SOLAR_MASS,
+    },
+];
 
-use core::os;
+struct Planet {
+    x: [f64, ..3],
+    v: [f64, ..3],
+    mass: f64,
+}
 
-// Using sqrt from the standard library is way slower than using libc
-// directly even though std just calls libc, I guess it must be
-// because the the indirection through another dynamic linker
-// stub. Kind of shocking. Might be able to make it faster still with
-// an llvm intrinsic.
-mod libc {
-    #[nolink]
-    pub extern {
-        pub fn sqrt(n: float) -> float;
+fn advance(bodies: &mut [Planet, ..N_BODIES], dt: f64, steps: i32) {
+    let mut d = [ 0.0, ..3 ];
+    for (steps as uint).times {
+        for range(0, N_BODIES) |i| {
+            for range(i + 1, N_BODIES) |j| {
+                d[0] = bodies[i].x[0] - bodies[j].x[0];
+                d[1] = bodies[i].x[1] - bodies[j].x[1];
+                d[2] = bodies[i].x[2] - bodies[j].x[2];
+
+                let d2 = d[0]*d[0] + d[1]*d[1] + d[2]*d[2];
+                let mag = dt / (d2 * sqrtf64(d2));
+
+                let a_mass = bodies[i].mass, b_mass = bodies[j].mass;
+                bodies[i].v[0] -= d[0] * b_mass * mag;
+                bodies[i].v[1] -= d[1] * b_mass * mag;
+                bodies[i].v[2] -= d[2] * b_mass * mag;
+
+                bodies[j].v[0] += d[0] * a_mass * mag;
+                bodies[j].v[1] += d[1] * a_mass * mag;
+                bodies[j].v[2] += d[2] * a_mass * mag;
+            }
+        }
+
+        for vec::each_mut(*bodies) |a| {
+            a.x[0] += dt * a.v[0];
+            a.x[1] += dt * a.v[1];
+            a.x[2] += dt * a.v[2];
+        }
+    }
+}
+
+fn energy(bodies: &[Planet, ..N_BODIES]) -> f64 {
+    let mut e = 0.0;
+    let mut d = [ 0.0, ..3 ];
+    for range(0, N_BODIES) |i| {
+        for range(0, 3) |k| {
+            e += bodies[i].mass * bodies[i].v[k] * bodies[i].v[k] / 2.0;
+        }
+
+        for range(i + 1, N_BODIES) |j| {
+            for range(0, 3) |k| {
+                d[k] = bodies[i].x[k] - bodies[j].x[k];
+            }
+            let dist = sqrtf64(d[0]*d[0] + d[1]*d[1] + d[2]*d[2]);
+            e -= bodies[i].mass * bodies[j].mass / dist;
+        }
+    }
+    e
+}
+
+fn offset_momentum(bodies: &mut [Planet, ..N_BODIES]) {
+    for range(0, N_BODIES) |i| {
+        for range(0, 3) |k| {
+            bodies[0].v[k] -= bodies[i].v[k] * bodies[i].mass / SOLAR_MASS;
+        }
     }
 }
 
 fn main() {
-    let args = os::args();
-    let args = if os::getenv(~"RUST_BENCH").is_some() {
-        ~[~"", ~"4000000"]
-    } else if args.len() <= 1u {
-        ~[~"", ~"100000"]
-    } else {
-        args
-    };
-    let n = int::from_str(args[1]).get();
-    let mut bodies: ~[Body::Props] = NBodySystem::make();
-    io::println(fmt!("%f", NBodySystem::energy(bodies)));
-    let mut i = 0;
-    while i < n {
-        NBodySystem::advance(bodies, 0.01);
-        i += 1;
-    }
-    io::println(fmt!("%f", NBodySystem::energy(bodies)));
+    let n: i32 = FromStr::from_str(os::args()[1]).get();
+    let mut bodies = BODIES;
+
+    offset_momentum(&mut bodies);
+    println(fmt!("%.9f", energy(&bodies) as float));
+
+    advance(&mut bodies, 0.01, n);
+
+    println(fmt!("%.9f", energy(&bodies) as float));
 }
 
-pub mod NBodySystem {
-    use Body;
-
-    pub fn make() -> ~[Body::Props] {
-        let mut bodies: ~[Body::Props] =
-            ~[Body::sun(),
-              Body::jupiter(),
-              Body::saturn(),
-              Body::uranus(),
-              Body::neptune()];
-
-        let mut px = 0.0;
-        let mut py = 0.0;
-        let mut pz = 0.0;
-
-        let mut i = 0;
-        while i < 5 {
-            px += bodies[i].vx * bodies[i].mass;
-            py += bodies[i].vy * bodies[i].mass;
-            pz += bodies[i].vz * bodies[i].mass;
-
-            i += 1;
-        }
-
-        // side-effecting
-        Body::offset_momentum(&mut bodies[0], px, py, pz);
-
-        return bodies;
-    }
-
-    pub fn advance(bodies: &mut [Body::Props], dt: float) {
-        let mut i = 0;
-        while i < 5 {
-            let mut j = i + 1;
-            while j < 5 {
-                advance_one(&mut bodies[i],
-                            &mut bodies[j], dt);
-                j += 1;
-            }
-
-            i += 1;
-        }
-
-        i = 0;
-        while i < 5 {
-            move_(&mut bodies[i], dt);
-            i += 1;
-        }
-    }
-
-    pub fn advance_one(bi: &mut Body::Props,
-                       bj: &mut Body::Props,
-                       dt: float) {
-        unsafe {
-            let dx = bi.x - bj.x;
-            let dy = bi.y - bj.y;
-            let dz = bi.z - bj.z;
-
-            let dSquared = dx * dx + dy * dy + dz * dz;
-
-            let distance = ::libc::sqrt(dSquared);
-            let mag = dt / (dSquared * distance);
-
-            bi.vx -= dx * bj.mass * mag;
-            bi.vy -= dy * bj.mass * mag;
-            bi.vz -= dz * bj.mass * mag;
-
-            bj.vx += dx * bi.mass * mag;
-            bj.vy += dy * bi.mass * mag;
-            bj.vz += dz * bi.mass * mag;
-        }
-    }
-
-    pub fn move_(b: &mut Body::Props, dt: float) {
-        b.x += dt * b.vx;
-        b.y += dt * b.vy;
-        b.z += dt * b.vz;
-    }
-
-    pub fn energy(bodies: &[Body::Props]) -> float {
-        unsafe {
-            let mut dx;
-            let mut dy;
-            let mut dz;
-            let mut distance;
-            let mut e = 0.0;
-
-            let mut i = 0;
-            while i < 5 {
-                e +=
-                    0.5 * bodies[i].mass *
-                    (bodies[i].vx * bodies[i].vx
-                     + bodies[i].vy * bodies[i].vy
-                     + bodies[i].vz * bodies[i].vz);
-
-                let mut j = i + 1;
-                while j < 5 {
-                    dx = bodies[i].x - bodies[j].x;
-                    dy = bodies[i].y - bodies[j].y;
-                    dz = bodies[i].z - bodies[j].z;
-
-                    distance = ::libc::sqrt(dx * dx
-                                            + dy * dy
-                                            + dz * dz);
-                    e -= bodies[i].mass
-                        * bodies[j].mass / distance;
-
-                    j += 1;
-                }
-
-                i += 1;
-            }
-            return e;
-        }
-    }
-}
-
-pub mod Body {
-    use Body;
-
-    pub static PI: float = 3.141592653589793;
-    pub static SOLAR_MASS: float = 39.478417604357432;
-    // was 4 * PI * PI originally
-    pub static DAYS_PER_YEAR: float = 365.24;
-
-    pub struct Props {
-        x: float,
-        y: float,
-        z: float,
-        vx: float,
-        vy: float,
-        vz: float,
-        mass: float
-    }
-
-    pub fn jupiter() -> Body::Props {
-        return Props {
-            x: 4.84143144246472090e+00,
-            y: -1.16032004402742839e+00,
-            z: -1.03622044471123109e-01,
-            vx: 1.66007664274403694e-03 * DAYS_PER_YEAR,
-            vy: 7.69901118419740425e-03 * DAYS_PER_YEAR,
-            vz: -6.90460016972063023e-05 * DAYS_PER_YEAR,
-            mass: 9.54791938424326609e-04 * SOLAR_MASS
-        };
-    }
-
-    pub fn saturn() -> Body::Props {
-        return Props {
-            x: 8.34336671824457987e+00,
-            y: 4.12479856412430479e+00,
-            z: -4.03523417114321381e-01,
-            vx: -2.76742510726862411e-03 * DAYS_PER_YEAR,
-            vy: 4.99852801234917238e-03 * DAYS_PER_YEAR,
-            vz: 2.30417297573763929e-05 * DAYS_PER_YEAR,
-            mass: 2.85885980666130812e-04 * SOLAR_MASS
-        };
-    }
-
-    pub fn uranus() -> Body::Props {
-        return Props {
-            x: 1.28943695621391310e+01,
-            y: -1.51111514016986312e+01,
-            z: -2.23307578892655734e-01,
-            vx: 2.96460137564761618e-03 * DAYS_PER_YEAR,
-            vy: 2.37847173959480950e-03 * DAYS_PER_YEAR,
-            vz: -2.96589568540237556e-05 * DAYS_PER_YEAR,
-            mass: 4.36624404335156298e-05 * SOLAR_MASS
-        };
-    }
-
-    pub fn neptune() -> Body::Props {
-        return Props {
-            x: 1.53796971148509165e+01,
-            y: -2.59193146099879641e+01,
-            z: 1.79258772950371181e-01,
-            vx: 2.68067772490389322e-03 * DAYS_PER_YEAR,
-            vy: 1.62824170038242295e-03 * DAYS_PER_YEAR,
-            vz: -9.51592254519715870e-05 * DAYS_PER_YEAR,
-            mass: 5.15138902046611451e-05 * SOLAR_MASS
-        };
-    }
-
-    pub fn sun() -> Body::Props {
-        return Props {
-            x: 0.0,
-            y: 0.0,
-            z: 0.0,
-            vx: 0.0,
-            vy: 0.0,
-            vz: 0.0,
-            mass: SOLAR_MASS
-        };
-    }
-
-    pub fn offset_momentum(props: &mut Body::Props,
-                           px: float,
-                           py: float,
-                           pz: float) {
-        props.vx = -px / SOLAR_MASS;
-        props.vy = -py / SOLAR_MASS;
-        props.vz = -pz / SOLAR_MASS;
-    }
-
-}
diff --git a/src/test/bench/shootout-spectralnorm.rs b/src/test/bench/shootout-spectralnorm.rs
index 6e39b755b22..00e255d890b 100644
--- a/src/test/bench/shootout-spectralnorm.rs
+++ b/src/test/bench/shootout-spectralnorm.rs
@@ -1,84 +1,54 @@
-// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
-// file at the top-level directory of this distribution and at
-// http://rust-lang.org/COPYRIGHT.
-//
-// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
-// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
-// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
-// option. This file may not be copied, modified, or distributed
-// except according to those terms.
+use core::from_str::FromStr;
+use core::iter::ExtendedMutableIter;
+use core::unstable::intrinsics::sqrtf64;
 
-// Based on spectalnorm.gcc by Sebastien Loisel
-
-extern mod std;
-
-fn eval_A(i: uint, j: uint) -> float {
-    1.0/(((i+j)*(i+j+1u)/2u+i+1u) as float)
+#[inline]
+fn A(i: i32, j: i32) -> i32 {
+    (i+j) * (i+j+1) / 2 + i + 1
 }
 
-fn eval_A_times_u(u: &const [float], Au: &mut [float]) {
-    let N = vec::len(u);
-    let mut i = 0u;
-    while i < N {
-        Au[i] = 0.0;
-        let mut j = 0u;
-        while j < N {
-            Au[i] += eval_A(i, j) * u[j];
-            j += 1u;
+fn dot(v: &[f64], u: &[f64]) -> f64 {
+    let mut sum = 0.0;
+    for v.eachi |i, &v_i| {
+        sum += v_i * u[i];
+    }
+    sum
+}
+
+fn mult_Av(v: &mut [f64], out: &mut [f64]) {
+    for vec::eachi_mut(out) |i, out_i| {
+        let mut sum = 0.0;
+        for vec::eachi_mut(v) |j, &v_j| {
+            sum += v_j / (A(i as i32, j as i32) as f64);
         }
-        i += 1u;
+        *out_i = sum;
     }
 }
 
-fn eval_At_times_u(u: &const [float], Au: &mut [float]) {
-    let N = vec::len(u);
-    let mut i = 0u;
-    while i < N {
-        Au[i] = 0.0;
-        let mut j = 0u;
-        while j < N {
-            Au[i] += eval_A(j, i) * u[j];
-            j += 1u;
+fn mult_Atv(v: &mut [f64], out: &mut [f64]) {
+    for vec::eachi_mut(out) |i, out_i| {
+        let mut sum = 0.0;
+        for vec::eachi_mut(v) |j, &v_j| {
+            sum += v_j / (A(j as i32, i as i32) as f64);
         }
-        i += 1u;
+        *out_i = sum;
     }
 }
 
-fn eval_AtA_times_u(u: &const [float], AtAu: &mut [float]) {
-    let mut v = vec::from_elem(vec::len(u), 0.0);
-    eval_A_times_u(u, v);
-    eval_At_times_u(v, AtAu);
+fn mult_AtAv(v: &mut [f64], out: &mut [f64], tmp: &mut [f64]) {
+    mult_Av(v, tmp);
+    mult_Atv(tmp, out);
 }
 
+#[fixed_stack_segment]
 fn main() {
-    let args = os::args();
-    let args = if os::getenv(~"RUST_BENCH").is_some() {
-        ~[~"", ~"2000"]
-    } else if args.len() <= 1u {
-        ~[~"", ~"1000"]
-    } else {
-        args
-    };
-
-    let N = uint::from_str(args[1]).get();
-
-    let mut u = vec::from_elem(N, 1.0);
-    let mut v = vec::from_elem(N, 0.0);
-    let mut i = 0u;
-    while i < 10u {
-        eval_AtA_times_u(u, v);
-        eval_AtA_times_u(v, u);
-        i += 1u;
+    let n: uint = FromStr::from_str(os::args()[1]).get();
+    let mut u = vec::from_elem(n, 1f64), v = u.clone(), tmp = u.clone();
+    for 8.times {
+        mult_AtAv(u, v, tmp);
+        mult_AtAv(v, u, tmp);
     }
 
-    let mut vBv = 0.0;
-    let mut vv = 0.0;
-    let mut i = 0u;
-    while i < N {
-        vBv += u[i] * v[i];
-        vv += v[i] * v[i];
-        i += 1u;
-    }
-
-    io::println(fmt!("%0.9f\n", float::sqrt(vBv / vv)));
+    println(fmt!("%.9f", sqrtf64(dot(u,v) / dot(v,v)) as float));
 }
+