For challenge 7 (multi threaded)

my-step-tsp-4m.rs

#[macro_use]
extern crate log;
extern crate crossbeam;
extern crate csv;
extern crate env_logger;
extern crate num_cpus;
extern crate ordered_float;
extern crate rustc_serialize;

use std::collections::HashSet;
use ordered_float::OrderedFloat;

fn main() {
    env_logger::init().unwrap();
    let args = std::env::args().collect::<Vec<String>>();
    assert!(args.len() > 1);
    let mut reader = csv::Reader::from_file(&args[1]).unwrap();
    let cities = reader.decode().map(|r| r.unwrap()).collect::<Vec<City>>();
    let n = cities.len();
    let tsp = TSP::new(cities);
    let answer = tsp.solve();
    debug_assert_eq!(answer.iter().cloned().collect::<HashSet<usize>>().len(), n);
    info!("Distance: {}", tsp.route_dist(&answer));
}

#[derive(RustcDecodable, Clone)]
struct City {
    x: f64,
    y: f64,
}

type Route = Vec<usize>;

#[derive(Clone)]
struct Square {
    city_indices: Vec<usize>,
    cities: Vec<City>,
    route: Route,
}

impl Square {

    fn new() -> Square {
        Square {
            city_indices: Vec::with_capacity(2048),
            cities: Vec::with_capacity(2048),
            route: Vec::new(),
        }
    }
}

impl Square {

    fn dist(&self, i: usize, j: usize) -> f64 {
        (self.cities[i].x - self.cities[j].x).hypot(self.cities[i].y - self.cities[j].y)
    }

    fn solve(&self) -> Route {
        let n = self.city_indices.len();
        let mut unvisited = (1..n).collect::<HashSet<usize>>();
        let mut current = 0;
        let mut route = vec![current];
        while let Some(&i) = unvisited.iter().min_by_key(|&&i| OrderedFloat(self.dist(current, i))) {
            route.push(i);
            unvisited.remove(&i);
            current = i;
        }
        route.iter().map(|&i| self.city_indices[i]).collect::<Route>()
    }
}

struct TSP {
    n: usize,
    cities: Vec<City>,
}

impl TSP {

    fn new(cities: Vec<City>) -> TSP {
        TSP {
            n: cities.len(),
            cities: cities,
        }
    }

    fn dist(&self, i: usize, j: usize) -> f64 {
        (self.cities[i].x - self.cities[j].x).hypot(self.cities[i].y - self.cities[j].y)
    }

    fn route_dist(&self, route: &Route) -> f64 {
        route.iter().fold(0.0, |acc, &i| acc + self.dist(route[i], route[(i + 1) % self.n]))
    }

    fn solve(&self) -> Route {
        // TODO(hayato): Avoid magic numbers
        let dx = 1600;
        let dy = 900;
        let square_size = 25;
        assert_eq!(dx % square_size, 0);
        assert_eq!(dy % square_size, 0);
        let sx = dx / square_size;
        let sy = dy / square_size;
        let squares_num = sx * sy;
        let mut squares = vec![Square::new(); squares_num];
        for (i, city) in self.cities.iter().enumerate() {
            let x = (city.x as usize) / square_size;
            let y = (city.y as usize) / square_size;
            let square_index = y * sx + x;
            squares[square_index].city_indices.push(i);
            squares[square_index].cities.push(city.clone());
        }

        let threads_num = num_cpus::get();
        debug!("threads_num: {}", threads_num);
        crossbeam::scope(|scope| {
            for chunk in squares.chunks_mut((squares_num + threads_num - 1)/ threads_num) {
                scope.spawn(move || {
                    debug!("thread: squares: {}", chunk.len());
                    for square in chunk {
                        square.route = square.solve();
                    };
                });
            }
        });

        let mut route = Vec::with_capacity(self.n);
        assert_eq!(sx % 2, 0);
        assert_eq!(sy % 2, 0);
        let hx = sx / 2;
        for y in 0..sy {
            if y % 2 == 0 {
                for x in (0..hx).rev() {
                    route.append(&mut squares[y * sx + x].route);
                }
            } else {
                for x in 0..hx {
                    route.append(&mut squares[y * sx + x].route);
                }
            }
        }
        for y in (0..sy).rev() {
            if y % 2 == 0 {
                for x in (hx..sx).rev() {
                    route.append(&mut squares[y * sx + x].route);
                }
            } else {
                for x in hx..sx {
                    route.append(&mut squares[y * sx + x].route);
                }
            }
        }
        route
    }
}