<commit_before>use std::io::Write;
use std::io::Result as IoResult;
use std::net::{TcpListener, TcpStream, SocketAddr, ToSocketAddrs};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use std::ops::RangeFrom;
use std::collections::HashMap;
use std::f64::consts::PI;

use bincode::{serialize, deserialize_from, Infinite, Bounded};

use state::{WorldState, GameState, Player, Unit, UnitId};
use network::{Message, Command};

/// A `Server` instance holds global server state.
pub struct Server {
    socket_addr: SocketAddr,
    world: Arc<Mutex<WorldState>>,
    game: Arc<Mutex<GameState>>,
    /// Generator that returns sequential unit IDs
    unit_id_generator: Arc<Mutex<RangeFrom<u32>>>,
    /// Generator that returns sequential client IDs
    client_id_generator: Arc<Mutex<RangeFrom<u32>>>,

    /// Map with active unit move commands
    unit_targets: Arc<Mutex<HashMap<UnitId, [f64; 2]>>>,
}

impl Server {
    pub fn new<T: ToSocketAddrs>(addr: T,
                                world_size: (f64, f64))
                                -> IoResult<Server> {
        let addr = try!(addr.to_socket_addrs()).next().unwrap();
        let world = Arc::new(Mutex::new(WorldState::new(world_size.0, world_size.1)));
        let game = Arc::new(Mutex::new(GameState::new()));
        Ok(Server {
            socket_addr: addr,
            world: world,
            game: game,
            client_id_generator: Arc::new(Mutex::new(0..)),
            unit_id_generator: Arc::new(Mutex::new(0..)),
            unit_targets: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    pub fn serve(&self) {
        let tcp_listener = TcpListener::bind(self.socket_addr).unwrap();
        println!("Start server: {:?}", tcp_listener);

        let game_clone = self.game.clone();
        let unit_targets_clone = self.unit_targets.clone();
        thread::spawn(move || {
            update_world(game_clone, unit_targets_clone);
        });

        for stream in tcp_listener.incoming() {
            match stream {
                Ok(stream) => {
                    let world_clone = self.world.clone();
                    let game_clone = self.game.clone();
                    let client_id_generator_clone = self.client_id_generator.clone();
                    let unit_id_generator_clone = self.unit_id_generator.clone();
                    let unit_targets = self.unit_targets.clone();
                    println!("Spawning thread...");
                    thread::spawn(move || {
                        handle_client(stream, world_clone, game_clone,
                                      client_id_generator_clone, unit_id_generator_clone, unit_targets);
                    });
                }
                Err(e) => {
                    println!("{:?}", e);
                }
            }
        }
    }
}

pub type SafeWorldState = Arc<Mutex<WorldState>>;
pub type SafeUnitTargets = Arc<Mutex<HashMap<UnitId, [f64; 2]>>>;

pub fn handle_client(mut stream: TcpStream,
                     world: SafeWorldState,
                     game: Arc<Mutex<GameState>>,
                     client_id_generator: Arc<Mutex<RangeFrom<u32>>>,
                     unit_id_generator: Arc<Mutex<RangeFrom<u32>>>,
                     unit_targets: SafeUnitTargets) {

    // handle client hello
    let client_message = deserialize_from(&mut stream, Bounded(128));
    match client_message {
        Ok(message) => {
            match message {
                Message::ClientHello => {
                    // Get exclusive world access
                    let mut game_lock = game.lock().unwrap();

                    // Create new player for the newly connected client
                    let client_id = client_id_generator
                        .lock().expect("Could not lock client_id_generator mutex")
                        .next().expect("No more client IDs available!");
                    let mut player = Player::new(client_id);

                    // Create four initial units for the player
                    let coords = [
                        [50.0f64, 50.0f64], [50.0f64, 100.0f64], [100.0f64, 50.0f64], [100.0f64, 100.0f64],
                    ];
                    for coord in coords.iter() {
                        let unit_id = unit_id_generator
                            .lock().expect("Could not lock unit_id_generator mutex")
                            .next().expect("No more unit IDs available!");
                        player.units.push(Unit::new(unit_id, *coord));
                    }

                    // Add player to the world
                    let player_id = player.id;
                    game_lock.players.push(player);

                    // Send ServerHello message
                    let encoded: Vec<u8> = serialize(
                        &Message::ServerHello(player_id, world.lock().unwrap().clone()),
                        Infinite
                    ).unwrap();
                    stream.write(&encoded).unwrap();
                },
                Message::ClientReconnect(id) => {
                    // Get exclusive world access
                    let world_lock = world.lock().unwrap();
                    let game_lock = game.lock().unwrap();

                    // Find player with specified id
                    match game_lock.players.iter().find(|player| player.id == id) {
                        Some(_) => {
                            println!("Found you :)");

                            // Send ServerHello message
                            let encoded: Vec<u8> = serialize(
                                &Message::ServerHello(id, world_lock.clone()),
                                Infinite
                            ).unwrap();
                            stream.write(&encoded).unwrap();
                        },
                        None => {
                            println!("Reconnect to id {} not possible", id);

                            // Send Error message
                            let encoded: Vec<u8> = serialize(
                                &Message::Error,
                                Infinite).unwrap();
                            stream.write(&encoded).unwrap();
                            return  // Don't enter game loop
                        }
                    }
                },
                _ => {
                    println!("Did not receive ClientHello: {:?}", message);
                    let encoded: Vec<u8> = serialize(&Message::Error, Infinite).unwrap();
                    stream.write(&encoded).unwrap();
                    return  // Don't enter game loop
                }
            }
        }
        Err(e) => {
            println!("Error: {:?}", e);
            return  // Don't enter game loop
        }
    }

    let mut command_stream = stream.try_clone().unwrap();
    let world_clone = world.clone();
    let game_clone = game.clone();
    let unit_targets_clone = unit_targets.clone();
    // Command receiver loop
    thread::spawn(move || {
        loop {
            let client_message = deserialize_from(&mut command_stream, Bounded(128));
            match client_message {
                Ok(message) => {
                    match message {
                        Message::Command(command) => {
                            let world_lock = world_clone.lock().unwrap();
                            let mut game_lock = game_clone.lock().unwrap();
                            let mut unit_targets_lock = unit_targets_clone.lock().unwrap();
                            handle_command(&world_lock, &mut game_lock, &mut unit_targets_lock, &command);
                        },
                        _ => {
                            println!("Did receive unexpected message: {:?}", message);
                            let encoded: Vec<u8> = serialize(&Message::Error, Infinite).unwrap();
                            command_stream.write(&encoded).unwrap();
                            return
                        },
                    }
                },
                Err(e) => {
                    println!("Error: {:?}", e);
                    return;
                }
            };
        }
    });

    // GameState loop
    loop {
        let encoded: Vec<u8> = {
            let game_lock = game.lock().unwrap();
            serialize(&*game_lock, Infinite).unwrap()
        };
        match stream.write(&encoded) {
            Err(e) => {
                println!("Error: {:?}", e);
                return;
            }
            _ => thread::sleep(Duration::from_millis(10)),
        };
    }
}

pub fn handle_command(world: &WorldState,
                      game: &mut GameState,
                      unit_targets: &mut HashMap<UnitId, [f64; 2]>, command: &Command) {
    println!("Did receive command {:?}", command);
    match command {
        &Command::Move(id, move_target) => {
            for player in game.players.iter_mut() {
                for unit in player.units.iter_mut() {
                    if unit.id == id {
                        println!("Found it :)");
                        let mut target = [0.0; 2];
                        target[0] = if move_target[0] > world.x {
                            world.x
                        } else if move_target[0] < 0.0 {
                            0.0
                        } else {
                            move_target[0]
                        };
                        target[1] = if move_target[1] > world.y {
                            world.y
                        } else if move_target[1] < 0.0 {
                            0.0
                        } else {
                            move_target[1]
                        };
                        let dx = target[0] - unit.position[0];
                        let dy = target[1] - unit.position[1];
                        if dx.is_sign_negative() {
                            unit.angle = (dy / dx).atan() + PI;
                        } else {
                            unit.angle = (dy / dx).atan();
                        }
                        unit_targets.insert(id, target);
                    }
                }
            }
            println!("Move {} to {:?}!", id, move_target);
        }
        &Command::Shoot(id) => {
            game.shoot(id);
        }
    }
}

pub fn update_world(game: Arc<Mutex<GameState>>, unit_targets: SafeUnitTargets) {
    loop {
        {
            let mut game_lock = game.lock().unwrap();
            let unit_targets = unit_targets.lock().unwrap();
            game_lock.update_targets(&unit_targets);
            game_lock.update(5.0);
        }
        thread::sleep(Duration::from_millis(5));
    }
}
<commit_msg>server: Do not use a Mutex for the WorldState<commit_after>use std::io::Write;
use std::io::Result as IoResult;
use std::net::{TcpListener, TcpStream, SocketAddr, ToSocketAddrs};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use std::ops::RangeFrom;
use std::collections::HashMap;
use std::f64::consts::PI;

use bincode::{serialize, deserialize_from, Infinite, Bounded};

use state::{WorldState, GameState, Player, Unit, UnitId};
use network::{Message, Command};

/// A `Server` instance holds global server state.
pub struct Server {
    socket_addr: SocketAddr,
    world: Arc<WorldState>,
    game: Arc<Mutex<GameState>>,
    /// Generator that returns sequential unit IDs
    unit_id_generator: Arc<Mutex<RangeFrom<u32>>>,
    /// Generator that returns sequential client IDs
    client_id_generator: Arc<Mutex<RangeFrom<u32>>>,

    /// Map with active unit move commands
    unit_targets: Arc<Mutex<HashMap<UnitId, [f64; 2]>>>,
}

impl Server {
    pub fn new<T: ToSocketAddrs>(addr: T,
                                world_size: (f64, f64))
                                -> IoResult<Server> {
        let addr = try!(addr.to_socket_addrs()).next().unwrap();
        let world = Arc::new(WorldState::new(world_size.0, world_size.1));
        let game = Arc::new(Mutex::new(GameState::new()));
        Ok(Server {
            socket_addr: addr,
            world: world,
            game: game,
            client_id_generator: Arc::new(Mutex::new(0..)),
            unit_id_generator: Arc::new(Mutex::new(0..)),
            unit_targets: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    pub fn serve(&self) {
        let tcp_listener = TcpListener::bind(self.socket_addr).unwrap();
        println!("Start server: {:?}", tcp_listener);

        let game_clone = self.game.clone();
        let unit_targets_clone = self.unit_targets.clone();
        thread::spawn(move || {
            update_world(game_clone, unit_targets_clone);
        });

        for stream in tcp_listener.incoming() {
            match stream {
                Ok(stream) => {
                    let world_clone = self.world.clone();
                    let game_clone = self.game.clone();
                    let client_id_generator_clone = self.client_id_generator.clone();
                    let unit_id_generator_clone = self.unit_id_generator.clone();
                    let unit_targets = self.unit_targets.clone();
                    println!("Spawning thread...");
                    thread::spawn(move || {
                        handle_client(stream, world_clone, game_clone,
                                      client_id_generator_clone, unit_id_generator_clone, unit_targets);
                    });
                }
                Err(e) => {
                    println!("{:?}", e);
                }
            }
        }
    }
}

pub type SafeUnitTargets = Arc<Mutex<HashMap<UnitId, [f64; 2]>>>;

pub fn handle_client(mut stream: TcpStream,
                     world: Arc<WorldState>,
                     game: Arc<Mutex<GameState>>,
                     client_id_generator: Arc<Mutex<RangeFrom<u32>>>,
                     unit_id_generator: Arc<Mutex<RangeFrom<u32>>>,
                     unit_targets: SafeUnitTargets) {

    // handle client hello
    let client_message = deserialize_from(&mut stream, Bounded(128));
    match client_message {
        Ok(message) => {
            match message {
                Message::ClientHello => {
                    // Get exclusive world access
                    let mut game_lock = game.lock().unwrap();

                    // Create new player for the newly connected client
                    let client_id = client_id_generator
                        .lock().expect("Could not lock client_id_generator mutex")
                        .next().expect("No more client IDs available!");
                    let mut player = Player::new(client_id);

                    // Create four initial units for the player
                    let coords = [
                        [50.0f64, 50.0f64], [50.0f64, 100.0f64], [100.0f64, 50.0f64], [100.0f64, 100.0f64],
                    ];
                    for coord in coords.iter() {
                        let unit_id = unit_id_generator
                            .lock().expect("Could not lock unit_id_generator mutex")
                            .next().expect("No more unit IDs available!");
                        player.units.push(Unit::new(unit_id, *coord));
                    }

                    // Add player to the world
                    let player_id = player.id;
                    game_lock.players.push(player);

                    // Send ServerHello message
                    let encoded: Vec<u8> = serialize(
                        &Message::ServerHello(player_id, (*world).clone()),
                        Infinite
                    ).unwrap();
                    stream.write(&encoded).unwrap();
                },
                Message::ClientReconnect(id) => {
                    // Get exclusive world access
                    let game_lock = game.lock().unwrap();

                    // Find player with specified id
                    match game_lock.players.iter().find(|player| player.id == id) {
                        Some(_) => {
                            println!("Found you :)");

                            // Send ServerHello message
                            let encoded: Vec<u8> = serialize(
                                &Message::ServerHello(id, (*world).clone()),
                                Infinite
                            ).unwrap();
                            stream.write(&encoded).unwrap();
                        },
                        None => {
                            println!("Reconnect to id {} not possible", id);

                            // Send Error message
                            let encoded: Vec<u8> = serialize(
                                &Message::Error,
                                Infinite).unwrap();
                            stream.write(&encoded).unwrap();
                            return  // Don't enter game loop
                        }
                    }
                },
                _ => {
                    println!("Did not receive ClientHello: {:?}", message);
                    let encoded: Vec<u8> = serialize(&Message::Error, Infinite).unwrap();
                    stream.write(&encoded).unwrap();
                    return  // Don't enter game loop
                }
            }
        }
        Err(e) => {
            println!("Error: {:?}", e);
            return  // Don't enter game loop
        }
    }

    let mut command_stream = stream.try_clone().unwrap();
    let world_clone = world.clone();
    let game_clone = game.clone();
    let unit_targets_clone = unit_targets.clone();
    // Command receiver loop
    thread::spawn(move || {
        loop {
            let client_message = deserialize_from(&mut command_stream, Bounded(128));
            match client_message {
                Ok(message) => {
                    match message {
                        Message::Command(command) => {
                            let mut game_lock = game_clone.lock().unwrap();
                            let mut unit_targets_lock = unit_targets_clone.lock().unwrap();
                            handle_command(&world_clone.clone(), &mut game_lock, &mut unit_targets_lock, &command);
                        },
                        _ => {
                            println!("Did receive unexpected message: {:?}", message);
                            let encoded: Vec<u8> = serialize(&Message::Error, Infinite).unwrap();
                            command_stream.write(&encoded).unwrap();
                            return
                        },
                    }
                },
                Err(e) => {
                    println!("Error: {:?}", e);
                    return;
                }
            };
        }
    });

    // GameState loop
    loop {
        let encoded: Vec<u8> = {
            let game_lock = game.lock().unwrap();
            serialize(&*game_lock, Infinite).unwrap()
        };
        match stream.write(&encoded) {
            Err(e) => {
                println!("Error: {:?}", e);
                return;
            }
            _ => thread::sleep(Duration::from_millis(10)),
        };
    }
}

pub fn handle_command(world: &WorldState,
                      game: &mut GameState,
                      unit_targets: &mut HashMap<UnitId, [f64; 2]>, command: &Command) {
    println!("Did receive command {:?}", command);
    match command {
        &Command::Move(id, move_target) => {
            for player in game.players.iter_mut() {
                for unit in player.units.iter_mut() {
                    if unit.id == id {
                        println!("Found it :)");
                        let mut target = [0.0; 2];
                        target[0] = if move_target[0] > world.x {
                            world.x
                        } else if move_target[0] < 0.0 {
                            0.0
                        } else {
                            move_target[0]
                        };
                        target[1] = if move_target[1] > world.y {
                            world.y
                        } else if move_target[1] < 0.0 {
                            0.0
                        } else {
                            move_target[1]
                        };
                        let dx = target[0] - unit.position[0];
                        let dy = target[1] - unit.position[1];
                        if dx.is_sign_negative() {
                            unit.angle = (dy / dx).atan() + PI;
                        } else {
                            unit.angle = (dy / dx).atan();
                        }
                        unit_targets.insert(id, target);
                    }
                }
            }
            println!("Move {} to {:?}!", id, move_target);
        }
        &Command::Shoot(id) => {
            game.shoot(id);
        }
    }
}

pub fn update_world(game: Arc<Mutex<GameState>>, unit_targets: SafeUnitTargets) {
    loop {
        {
            let mut game_lock = game.lock().unwrap();
            let unit_targets = unit_targets.lock().unwrap();
            game_lock.update_targets(&unit_targets);
            game_lock.update(5.0);
        }
        thread::sleep(Duration::from_millis(5));
    }
}
<|endoftext|>


<commit_before>use latch::{LatchProbe, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use job::StackJob;
use registry::{self, WorkerThread};
use std::any::Any;
use unwind;

#[cfg(test)]
mod test;

/// The `join` function takes two closures and *potentially* runs them
/// in parallel. It returns a pair of the results from those closures.
///
/// Conceptually, calling `join()` is similar to spawning two threads,
/// one executing each of the two closures. However, the
/// implementation is quite different and incurs very low
/// overhead. The underlying technique is called "work stealing": the
/// Rayon runtime uses a fixed pool of worker threads and attempts to
/// only execute code in parallel when there are idle CPUs to handle
/// it.
///
/// When `join` is called from outside the thread pool, the calling
/// thread will block while the closures execute in the pool.  When
/// `join` is called within the pool, the calling thread still actively
/// participates in the thread pool.  It will try to execute its own
/// closures first, but if one is stolen then it will look for other
/// work while waiting for that result to come back.
///
/// ### Warning about blocking I/O
///
/// The assumption is that the closures given to `join()` are
/// CPU-bound tasks that to do not perform I/O or other blocking
/// operations. If you do perform I/O, and that I/O should block
/// (e.g., waiting for a network request), the overall performance may
/// be poor.  Moreover, if you cause one closure to be blocked waiting
/// on another (for example, using a channel), that could lead to a
/// deadlock.
///
/// ### Panics
///
/// No matter what happens, both closures will always be executed.  If
/// a single closure panics, whether it be the first or second
/// closure, that panic will be propagated and hence `join()` will
/// panic with the same panic value. If both closures panic, `join()`
/// will panic with the panic value from the first closure.
pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
    where A: FnOnce() -> RA + Send,
          B: FnOnce() -> RB + Send,
          RA: Send,
          RB: Send
{
    registry::in_worker(|worker_thread| unsafe {
        log!(Join { worker: worker_thread.index() });

        // Create virtual wrapper for task b; this all has to be
        // done here so that the stack frame can keep it all live
        // long enough.
        let job_b = StackJob::new(oper_b, SpinLatch::new());
        let job_b_ref = job_b.as_job_ref();
        worker_thread.push(job_b_ref);

        // Execute task a; hopefully b gets stolen in the meantime.
        let result_a = match unwind::halt_unwinding(oper_a) {
            Ok(v) => v,
            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err),
        };

        // Now that task A has finished, try to pop job B from the
        // local stack.  It may already have been popped by job A; it
        // may also have been stolen. There may also be some tasks
        // pushed on top of it in the stack, and we will have to pop
        // those off to get to it.
        while !job_b.latch.probe() {
            if let Some(job) = worker_thread.pop() {
                if job == job_b_ref {
                    // Found it! Let's run it.
                    //
                    // Note that this could panic, but it's ok if we unwind here.
                    log!(PoppedRhs { worker: worker_thread.index() });
                    let result_b = job_b.run_inline();
                    return (result_a, result_b);
                } else {
                    log!(PoppedJob { worker: worker_thread.index() });
                    worker_thread.execute(job);
                }
            } else {
                // Local deque is empty. Time to steal from other
                // threads.
                log!(LostJob { worker: worker_thread.index() });
                worker_thread.wait_until(&job_b.latch);
                debug_assert!(job_b.latch.probe());
                break;
            }
        }

        return (result_a, job_b.into_result());
    })
}

/// If job A panics, we still cannot return until we are sure that job
/// B is complete. This is because it may contain references into the
/// enclosing stack frame(s).
#[cold] // cold path
unsafe fn join_recover_from_panic(worker_thread: &WorkerThread,
                                  job_b_latch: &SpinLatch,
                                  err: Box<Any + Send>)
                                  -> !
{
    worker_thread.wait_until(job_b_latch);
    unwind::resume_unwinding(err)
}
<commit_msg>be specific about which closure executes locally<commit_after>use latch::{LatchProbe, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use job::StackJob;
use registry::{self, WorkerThread};
use std::any::Any;
use unwind;

#[cfg(test)]
mod test;

/// The `join` function takes two closures and *potentially* runs them
/// in parallel. It returns a pair of the results from those closures.
///
/// Conceptually, calling `join()` is similar to spawning two threads,
/// one executing each of the two closures. However, the
/// implementation is quite different and incurs very low
/// overhead. The underlying technique is called "work stealing": the
/// Rayon runtime uses a fixed pool of worker threads and attempts to
/// only execute code in parallel when there are idle CPUs to handle
/// it.
///
/// When `join` is called from outside the thread pool, the calling
/// thread will block while the closures execute in the pool.  When
/// `join` is called within the pool, the calling thread still actively
/// participates in the thread pool. It will begin by executing closure
/// A (on the current thread). While it is doing that, it will advertise
/// closure B as being available for other threads to execute. Once closure A
/// has completed, the current thread will try to execute closure B;
/// if however closure B has been stolen, then it will look for other work
/// while waiting for the thief to fully execute closure B. (This is the
/// typical work-stealing strategy).
///
/// ### Warning about blocking I/O
///
/// The assumption is that the closures given to `join()` are
/// CPU-bound tasks that to do not perform I/O or other blocking
/// operations. If you do perform I/O, and that I/O should block
/// (e.g., waiting for a network request), the overall performance may
/// be poor.  Moreover, if you cause one closure to be blocked waiting
/// on another (for example, using a channel), that could lead to a
/// deadlock.
///
/// ### Panics
///
/// No matter what happens, both closures will always be executed.  If
/// a single closure panics, whether it be the first or second
/// closure, that panic will be propagated and hence `join()` will
/// panic with the same panic value. If both closures panic, `join()`
/// will panic with the panic value from the first closure.
pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
    where A: FnOnce() -> RA + Send,
          B: FnOnce() -> RB + Send,
          RA: Send,
          RB: Send
{
    registry::in_worker(|worker_thread| unsafe {
        log!(Join { worker: worker_thread.index() });

        // Create virtual wrapper for task b; this all has to be
        // done here so that the stack frame can keep it all live
        // long enough.
        let job_b = StackJob::new(oper_b, SpinLatch::new());
        let job_b_ref = job_b.as_job_ref();
        worker_thread.push(job_b_ref);

        // Execute task a; hopefully b gets stolen in the meantime.
        let result_a = match unwind::halt_unwinding(oper_a) {
            Ok(v) => v,
            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err),
        };

        // Now that task A has finished, try to pop job B from the
        // local stack.  It may already have been popped by job A; it
        // may also have been stolen. There may also be some tasks
        // pushed on top of it in the stack, and we will have to pop
        // those off to get to it.
        while !job_b.latch.probe() {
            if let Some(job) = worker_thread.pop() {
                if job == job_b_ref {
                    // Found it! Let's run it.
                    //
                    // Note that this could panic, but it's ok if we unwind here.
                    log!(PoppedRhs { worker: worker_thread.index() });
                    let result_b = job_b.run_inline();
                    return (result_a, result_b);
                } else {
                    log!(PoppedJob { worker: worker_thread.index() });
                    worker_thread.execute(job);
                }
            } else {
                // Local deque is empty. Time to steal from other
                // threads.
                log!(LostJob { worker: worker_thread.index() });
                worker_thread.wait_until(&job_b.latch);
                debug_assert!(job_b.latch.probe());
                break;
            }
        }

        return (result_a, job_b.into_result());
    })
}

/// If job A panics, we still cannot return until we are sure that job
/// B is complete. This is because it may contain references into the
/// enclosing stack frame(s).
#[cold] // cold path
unsafe fn join_recover_from_panic(worker_thread: &WorkerThread,
                                  job_b_latch: &SpinLatch,
                                  err: Box<Any + Send>)
                                  -> !
{
    worker_thread.wait_until(job_b_latch);
    unwind::resume_unwinding(err)
}
<|endoftext|>


<commit_before>extern crate image;

use vector::*;

pub type Color = Vec3;

pub enum Object {
    Sphere { center: Point, radius: f64 }
}

pub struct Material {
    pub ambient: Color,
    pub diffuse: Color,
    pub specular: Color,
    pub specular_exponent: f64
}

pub type MaterialObject = (Object, Material);

pub struct Light {
    pub position: Point,
    pub color: Color,
    pub intensity: f64
}

pub struct Camera {
    pub position: Point,
    pub direction: Vec3,
    pub up: Vec3
}

pub struct Scene {
    pub camera: Camera,
    pub lights: Vec<Light>,
    pub objects: Vec<MaterialObject>,
}

impl Color {
    pub fn to_rgb(&self) -> image::Rgb<u8> {
        // TODO: implement tonemapping (by parameterizing color on a trait?)
        image::Rgb { data: [(self.x.min(1.0) * 255.0) as u8,
                            (self.y.min(1.0) * 255.0) as u8,
                            (self.z.min(1.0) * 255.0) as u8] }
    }
}

impl Scene {
    pub fn intersects(&self, ray: &Ray) -> Option<(&MaterialObject, f64)> {
        self.objects.iter().fold(None, |result, materialobject| {
            let (ref object, _) = *materialobject;
            match intersects(&ray, object) {
                Some(t) => {
                    match result {
                        Some((_, t_prev)) => {
                            if t < t_prev {
                                Some((&materialobject, t))
                            }
                            else {
                                result
                            }
                        }
                        None => {
                            Some((&materialobject, t))
                        }
                    }
                }
                None => {
                    result
                }
            }
        })
    }
}

const T_THRESHOLD: f64 = 0.00001;

pub fn intersects(ray: &Ray, object: &Object) -> Option<f64> {
    let Ray { origin: e, direction: d } = *ray;
    match *object {
        Object::Sphere { center: c, radius: r } => {
            let discrim = d.dot(e - c).powi(2) -
                (d.dot(d)) * ((e - c).dot(e-c) - r.powi(2));
            if discrim > T_THRESHOLD {
                let sqrtd = discrim.sqrt();
                let t1 = ((-d).dot(e - c) - sqrtd) / (d.dot(d));
                let t2 = ((-d).dot(e - c) + sqrtd) / (d.dot(d));
                match (t1 >= T_THRESHOLD, t2 >= T_THRESHOLD, t1 >= t2) {
                    (true, false, true) |
                    (true, true, false) |
                    (true, false, false) => {
                        Some(t1)
                    }
                    (false, true, true) |
                    (false, true, false) |
                    (true, true, true) => {
                        Some(t2)
                    }
                    _ => {
                        None
                    }
                }
            }
            else {
                None
            }
        }
    }
}
<commit_msg>Lower collision time threshhold<commit_after>extern crate image;

use vector::*;

pub type Color = Vec3;

pub enum Object {
    Sphere { center: Point, radius: f64 }
}

pub struct Material {
    pub ambient: Color,
    pub diffuse: Color,
    pub specular: Color,
    pub specular_exponent: f64
}

pub type MaterialObject = (Object, Material);

pub struct Light {
    pub position: Point,
    pub color: Color,
    pub intensity: f64
}

pub struct Camera {
    pub position: Point,
    pub direction: Vec3,
    pub up: Vec3
}

pub struct Scene {
    pub camera: Camera,
    pub lights: Vec<Light>,
    pub objects: Vec<MaterialObject>,
}

impl Color {
    pub fn to_rgb(&self) -> image::Rgb<u8> {
        // TODO: implement tonemapping (by parameterizing color on a trait?)
        image::Rgb { data: [(self.x.min(1.0) * 255.0) as u8,
                            (self.y.min(1.0) * 255.0) as u8,
                            (self.z.min(1.0) * 255.0) as u8] }
    }
}

impl Scene {
    pub fn intersects(&self, ray: &Ray) -> Option<(&MaterialObject, f64)> {
        self.objects.iter().fold(None, |result, materialobject| {
            let (ref object, _) = *materialobject;
            match intersects(&ray, object) {
                Some(t) => {
                    match result {
                        Some((_, t_prev)) => {
                            if t < t_prev {
                                Some((&materialobject, t))
                            }
                            else {
                                result
                            }
                        }
                        None => {
                            Some((&materialobject, t))
                        }
                    }
                }
                None => {
                    result
                }
            }
        })
    }
}

const T_THRESHOLD: f64 = 0.00000001;

pub fn intersects(ray: &Ray, object: &Object) -> Option<f64> {
    let Ray { origin: e, direction: d } = *ray;
    match *object {
        Object::Sphere { center: c, radius: r } => {
            let discrim = d.dot(e - c).powi(2) -
                (d.dot(d)) * ((e - c).dot(e-c) - r.powi(2));
            if discrim > T_THRESHOLD {
                let sqrtd = discrim.sqrt();
                let t1 = ((-d).dot(e - c) - sqrtd) / (d.dot(d));
                let t2 = ((-d).dot(e - c) + sqrtd) / (d.dot(d));
                match (t1 >= T_THRESHOLD, t2 >= T_THRESHOLD, t1 >= t2) {
                    (true, false, true) |
                    (true, true, false) |
                    (true, false, false) => {
                        Some(t1)
                    }
                    (false, true, true) |
                    (false, true, false) |
                    (true, true, true) => {
                        Some(t2)
                    }
                    _ => {
                        None
                    }
                }
            }
            else {
                None
            }
        }
    }
}
<|endoftext|>


<commit_before>mod data;

pub use self::data::AirportType;

use self::data::DataType;
use crate::path::FilePath;
use anyhow::{anyhow, Result};
use chrono::{Duration, NaiveDate, Utc};
use serde_derive::Serialize;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter};
use std::path::{Path, PathBuf};

const BASE_URL: &str = "https://ourairports.com/data";

fn validated_dir() -> Result<PathBuf> {
    FilePath::LocalData.validated_subdir("airport_data/")
}

pub fn ensure_updated() -> Result<()> {
    let base_dir = validated_dir()?;
    let mut last_update = LastUpdate::load(&base_dir);

    if !last_update.needs_update() {
        return Ok(());
    }

    download_latest(&base_dir)?;
    last_update.set_to_today()?;

    Ok(())
}

fn download_latest<P>(dir: P) -> Result<()>
where
    P: AsRef<Path>,
{
    println!("updating airport data");

    for dtype in &DataType::ALL {
        download_file(dtype.filename(), &dir)?;
    }

    println!("finished updating airport data");
    Ok(())
}

fn download_file<S, P>(name: S, dir: P) -> Result<()>
where
    S: AsRef<str>,
    P: AsRef<Path>,
{
    let name = name.as_ref();
    let dir = dir.as_ref();

    println!(".. downloading {}", name);

    let path = dir.join(name);
    let backup = Backup::create(&path)?;

    let resp = ureq::get(&format!("{}/{}", BASE_URL, name))
        .timeout_connect(15_000)
        .timeout_read(15_000)
        .call();

    if let Some(err) = resp.synthetic_error() {
        backup.restore()?;
        return Err(anyhow!("{}", err));
    }

    let file = match File::create(path) {
        Ok(file) => file,
        Err(err) => {
            backup.restore()?;
            return Err(err.into());
        }
    };

    let mut reader = resp.into_reader();
    let mut writer = BufWriter::new(file);

    match io::copy(&mut reader, &mut writer) {
        Ok(_) => Ok(()),
        Err(err) => {
            backup.restore()?;
            Err(err.into())
        }
    }
}

struct LastUpdate(PathBuf);

impl LastUpdate {
    const DATE_FORMAT: &'static str = "%Y-%m-%d";

    fn load<P>(base_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        let mut path = base_dir.into();
        path.push(".last_updated");

        Self(path)
    }

    fn needs_update(&self) -> bool {
        if !self.path().exists() {
            return true;
        }

        let updated_at = fs::read_to_string(self.path())
            .ok()
            .and_then(|updated_at| NaiveDate::parse_from_str(&updated_at, Self::DATE_FORMAT).ok());

        let updated_at = match updated_at {
            Some(updated_at) => updated_at,
            None => return true,
        };

        let today = Utc::now().naive_utc().date();

        if today - updated_at >= Duration::weeks(2) {
            return true;
        }

        false
    }

    fn set_to_today(&mut self) -> Result<()> {
        let today = Utc::now().naive_utc().date();
        fs::write(self.path(), today.format(Self::DATE_FORMAT).to_string())?;
        Ok(())
    }

    #[inline(always)]
    fn path(&self) -> &Path {
        &self.0
    }
}

enum Backup {
    Exists(PathBuf),
    DoesntExist,
}

impl Backup {
    fn create<P>(file: P) -> Result<Self>
    where
        P: AsRef<Path>,
    {
        let file = file.as_ref();

        if !file.exists() {
            return Ok(Self::DoesntExist);
        }

        let extension = match file.extension() {
            Some(ext) => format!("{}.bak", ext.to_string_lossy()).into(),
            None => Cow::Borrowed(".bak"),
        };

        let mut backup_path = PathBuf::from(file);
        backup_path.set_extension(extension.as_ref());

        fs::rename(file, &backup_path)?;

        Ok(Self::Exists(backup_path))
    }

    fn restore(self) -> Result<()> {
        match self {
            Self::Exists(path) => {
                if !path.exists() {
                    return Ok(());
                }

                let mut original_path = path.clone();
                original_path.set_extension("");

                if original_path.exists() {
                    fs::remove_file(&original_path)?;
                }

                fs::rename(path, original_path)?;
                Ok(())
            }
            Self::DoesntExist => Ok(()),
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Airport {
    pub icao: String,
    #[serde(rename = "type")]
    pub class: AirportType,
    pub position: Position,
    pub runways: Vec<Runway>,
    pub frequencies: HashMap<FrequencyType, String>,
    #[serde(rename = "countryName")]
    pub country_name: String,
}

impl Airport {
    pub fn load_all() -> Result<Vec<Self>> {
        let dir = validated_dir()?;

        let airports = data::Airport::from_dir(&dir)?;
        let mut runways = data::Runway::from_dir(&dir)?;
        let mut frequencies = data::Frequency::from_dir(&dir)?;
        let countries = data::Country::from_dir(&dir)?;

        let mut results = Vec::with_capacity(airports.len());

        for airport in airports {
            let runways = match runways.remove(&airport.id) {
                Some(runways) => runways,
                None => continue,
            };

            let frequencies = match frequencies.remove(&airport.id) {
                Some(frequencies) => frequencies,
                None => HashMap::new(),
            };

            let country = match countries.iter().find(|c| c.code == airport.country_code) {
                Some(country) => country,
                None => continue,
            };

            let result = Self {
                icao: airport.icao,
                class: airport.class,
                position: Position::new(airport.lat_deg, airport.lon_deg),
                runways: runways.into_iter().filter_map(Runway::from_data).collect(),
                frequencies: frequencies
                    .into_iter()
                    .map(|(freq_type, freq)| (freq_type.into(), freq.mhz))
                    .collect(),
                country_name: country.name.clone(),
            };

            results.push(result);
        }

        results.shrink_to_fit();
        Ok(results)
    }
}

#[derive(Debug, Serialize)]
pub struct Position {
    #[serde(rename = "latitudeDeg")]
    latitude_deg: f32,
    #[serde(rename = "longitudeDeg")]
    longitude_deg: f32,
}

impl Position {
    #[inline(always)]
    pub fn new(latitude_deg: f32, longitude_deg: f32) -> Self {
        Self {
            latitude_deg,
            longitude_deg,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Runway {
    #[serde(rename = "lengthFT")]
    pub length_ft: Option<i32>,
    #[serde(rename = "widthFT")]
    pub width_ft: Option<i32>,
    #[serde(rename = "northMarker")]
    pub north_marker: Option<RunwayMarker>,
    #[serde(rename = "southMarker")]
    pub south_marker: Option<RunwayMarker>,
}

impl Runway {
    fn from_data(source: data::Runway) -> Option<Self> {
        let north_marker = match (source.le_ident, source.le_lat_deg, source.le_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        let south_marker = match (source.he_ident, source.he_lat_deg, source.he_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        Some(Self {
            length_ft: source.length_ft,
            width_ft: source.width_ft,
            north_marker,
            south_marker,
        })
    }
}

#[derive(Debug, Serialize)]
pub struct RunwayMarker {
    pub name: String,
    pub position: Position,
}

impl RunwayMarker {
    #[inline(always)]
    pub fn new(name: String, lat_deg: f32, lon_deg: f32) -> Self {
        Self {
            name,
            position: Position::new(lat_deg, lon_deg),
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum FrequencyType {
    Atis,
    Arrival,
    Departure,
    #[serde(rename = "arrivalDeparture")]
    ArrivalDeparture,
    Ground,
    Tower,
    Unicom,
}

// TODO: return None when the source type is the Other variant
impl From<data::FrequencyType> for FrequencyType {
    fn from(source: data::FrequencyType) -> Self {
        use data::FrequencyType as DFT;

        match source {
            DFT::Atis => Self::Atis,
            DFT::Arrival => Self::Arrival,
            DFT::Departure => Self::Departure,
            DFT::ArrivalDeparture => Self::ArrivalDeparture,
            DFT::Ground => Self::Ground,
            DFT::Tower => Self::Tower,
            DFT::Unicom | DFT::Other => Self::Unicom,
        }
    }
}
<commit_msg>backend: Rename north_marker & south_marker Runway fields to he_marker and le_marker, respectively.<commit_after>mod data;

pub use self::data::AirportType;

use self::data::DataType;
use crate::path::FilePath;
use anyhow::{anyhow, Result};
use chrono::{Duration, NaiveDate, Utc};
use serde_derive::Serialize;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter};
use std::path::{Path, PathBuf};

const BASE_URL: &str = "https://ourairports.com/data";

fn validated_dir() -> Result<PathBuf> {
    FilePath::LocalData.validated_subdir("airport_data/")
}

pub fn ensure_updated() -> Result<()> {
    let base_dir = validated_dir()?;
    let mut last_update = LastUpdate::load(&base_dir);

    if !last_update.needs_update() {
        return Ok(());
    }

    download_latest(&base_dir)?;
    last_update.set_to_today()?;

    Ok(())
}

fn download_latest<P>(dir: P) -> Result<()>
where
    P: AsRef<Path>,
{
    println!("updating airport data");

    for dtype in &DataType::ALL {
        download_file(dtype.filename(), &dir)?;
    }

    println!("finished updating airport data");
    Ok(())
}

fn download_file<S, P>(name: S, dir: P) -> Result<()>
where
    S: AsRef<str>,
    P: AsRef<Path>,
{
    let name = name.as_ref();
    let dir = dir.as_ref();

    println!(".. downloading {}", name);

    let path = dir.join(name);
    let backup = Backup::create(&path)?;

    let resp = ureq::get(&format!("{}/{}", BASE_URL, name))
        .timeout_connect(15_000)
        .timeout_read(15_000)
        .call();

    if let Some(err) = resp.synthetic_error() {
        backup.restore()?;
        return Err(anyhow!("{}", err));
    }

    let file = match File::create(path) {
        Ok(file) => file,
        Err(err) => {
            backup.restore()?;
            return Err(err.into());
        }
    };

    let mut reader = resp.into_reader();
    let mut writer = BufWriter::new(file);

    match io::copy(&mut reader, &mut writer) {
        Ok(_) => Ok(()),
        Err(err) => {
            backup.restore()?;
            Err(err.into())
        }
    }
}

struct LastUpdate(PathBuf);

impl LastUpdate {
    const DATE_FORMAT: &'static str = "%Y-%m-%d";

    fn load<P>(base_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        let mut path = base_dir.into();
        path.push(".last_updated");

        Self(path)
    }

    fn needs_update(&self) -> bool {
        if !self.path().exists() {
            return true;
        }

        let updated_at = fs::read_to_string(self.path())
            .ok()
            .and_then(|updated_at| NaiveDate::parse_from_str(&updated_at, Self::DATE_FORMAT).ok());

        let updated_at = match updated_at {
            Some(updated_at) => updated_at,
            None => return true,
        };

        let today = Utc::now().naive_utc().date();

        if today - updated_at >= Duration::weeks(2) {
            return true;
        }

        false
    }

    fn set_to_today(&mut self) -> Result<()> {
        let today = Utc::now().naive_utc().date();
        fs::write(self.path(), today.format(Self::DATE_FORMAT).to_string())?;
        Ok(())
    }

    #[inline(always)]
    fn path(&self) -> &Path {
        &self.0
    }
}

enum Backup {
    Exists(PathBuf),
    DoesntExist,
}

impl Backup {
    fn create<P>(file: P) -> Result<Self>
    where
        P: AsRef<Path>,
    {
        let file = file.as_ref();

        if !file.exists() {
            return Ok(Self::DoesntExist);
        }

        let extension = match file.extension() {
            Some(ext) => format!("{}.bak", ext.to_string_lossy()).into(),
            None => Cow::Borrowed(".bak"),
        };

        let mut backup_path = PathBuf::from(file);
        backup_path.set_extension(extension.as_ref());

        fs::rename(file, &backup_path)?;

        Ok(Self::Exists(backup_path))
    }

    fn restore(self) -> Result<()> {
        match self {
            Self::Exists(path) => {
                if !path.exists() {
                    return Ok(());
                }

                let mut original_path = path.clone();
                original_path.set_extension("");

                if original_path.exists() {
                    fs::remove_file(&original_path)?;
                }

                fs::rename(path, original_path)?;
                Ok(())
            }
            Self::DoesntExist => Ok(()),
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Airport {
    pub icao: String,
    #[serde(rename = "type")]
    pub class: AirportType,
    pub position: Position,
    pub runways: Vec<Runway>,
    pub frequencies: HashMap<FrequencyType, String>,
    #[serde(rename = "countryName")]
    pub country_name: String,
}

impl Airport {
    pub fn load_all() -> Result<Vec<Self>> {
        let dir = validated_dir()?;

        let airports = data::Airport::from_dir(&dir)?;
        let mut runways = data::Runway::from_dir(&dir)?;
        let mut frequencies = data::Frequency::from_dir(&dir)?;
        let countries = data::Country::from_dir(&dir)?;

        let mut results = Vec::with_capacity(airports.len());

        for airport in airports {
            let runways = match runways.remove(&airport.id) {
                Some(runways) => runways,
                None => continue,
            };

            let frequencies = match frequencies.remove(&airport.id) {
                Some(frequencies) => frequencies,
                None => HashMap::new(),
            };

            let country = match countries.iter().find(|c| c.code == airport.country_code) {
                Some(country) => country,
                None => continue,
            };

            let result = Self {
                icao: airport.icao,
                class: airport.class,
                position: Position::new(airport.lat_deg, airport.lon_deg),
                runways: runways.into_iter().filter_map(Runway::from_data).collect(),
                frequencies: frequencies
                    .into_iter()
                    .map(|(freq_type, freq)| (freq_type.into(), freq.mhz))
                    .collect(),
                country_name: country.name.clone(),
            };

            results.push(result);
        }

        results.shrink_to_fit();
        Ok(results)
    }
}

#[derive(Debug, Serialize)]
pub struct Position {
    #[serde(rename = "latitudeDeg")]
    latitude_deg: f32,
    #[serde(rename = "longitudeDeg")]
    longitude_deg: f32,
}

impl Position {
    #[inline(always)]
    pub fn new(latitude_deg: f32, longitude_deg: f32) -> Self {
        Self {
            latitude_deg,
            longitude_deg,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Runway {
    #[serde(rename = "lengthFT")]
    pub length_ft: Option<i32>,
    #[serde(rename = "widthFT")]
    pub width_ft: Option<i32>,
    #[serde(rename = "heMarker")]
    pub he_marker: Option<RunwayMarker>,
    #[serde(rename = "leMarker")]
    pub le_marker: Option<RunwayMarker>,
}

impl Runway {
    fn from_data(source: data::Runway) -> Option<Self> {
        let he_marker = match (source.le_ident, source.le_lat_deg, source.le_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        let le_marker = match (source.he_ident, source.he_lat_deg, source.he_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        Some(Self {
            length_ft: source.length_ft,
            width_ft: source.width_ft,
            he_marker,
            le_marker,
        })
    }
}

#[derive(Debug, Serialize)]
pub struct RunwayMarker {
    pub name: String,
    pub position: Position,
}

impl RunwayMarker {
    #[inline(always)]
    pub fn new(name: String, lat_deg: f32, lon_deg: f32) -> Self {
        Self {
            name,
            position: Position::new(lat_deg, lon_deg),
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum FrequencyType {
    Atis,
    Arrival,
    Departure,
    #[serde(rename = "arrivalDeparture")]
    ArrivalDeparture,
    Ground,
    Tower,
    Unicom,
}

// TODO: return None when the source type is the Other variant
impl From<data::FrequencyType> for FrequencyType {
    fn from(source: data::FrequencyType) -> Self {
        use data::FrequencyType as DFT;

        match source {
            DFT::Atis => Self::Atis,
            DFT::Arrival => Self::Arrival,
            DFT::Departure => Self::Departure,
            DFT::ArrivalDeparture => Self::ArrivalDeparture,
            DFT::Ground => Self::Ground,
            DFT::Tower => Self::Tower,
            DFT::Unicom | DFT::Other => Self::Unicom,
        }
    }
}
<|endoftext|>


<commit_before>//! Rust entry point
//!
//! Boot is split into two parts
//!
//! + Early boot where we walk hardware data structures and setup initial
//!   machine state before
//! + Late boot where we are operating in the final kernel window and can
//!   setup the initial user environment
//!
//! Early boot is extremely architecture specific, late boot is mostly
//! generic.

use plat::*;
use vspace::*;
use panic::*;
use ::config::BootConfig;
use ::core::fmt::Write;
use ::core::marker::PhantomData;
use ::core::default::Default;
use super::halt::halt;
use super::vspace::*;
extern crate multiboot;

extern {
    /// Represent the start of the kernel image region. The type is not
    /// correct as there is no way to properly represent the type we want
    /// in Rust
    static kernel_image_start: u8;
    static kernel_image_end: u8;
}

/// Package of state that is passed to the early boot function
struct EarlyBootState<'h, 'l> {
    /// Referenece to the high window. Objects created from here can
    /// persist forever, and as such references to them can be returned
    /// in `PostEarlyBootState`
    high_window : &'h BootHighWindow<'h>,
    /// Reference to the low window. References to objects from here
    /// cannot be returned in `PostEarlyBootState`
    low_window : &'l BootLowWindow<'l>,
    /// The value of EAX passed from the assembly entry. This is checked
    /// to ensure we were multiboot loaded
    mbi_magic: usize,
    /// Raw physical pointer that should point to the multiboot structure
    mbi: *const usize,
}

/// Package of state that is returned as a result of early boot
struct PostEarlyBootState<'a> {
    /// An initialized platform interface
    plat: PlatInterfaceType,
    /// Currently the lifetime 'a is unused, so have some `PhantomData` to get
    /// around that
    phantom: PhantomData<&'a usize>,
}

/// Debug function to print out the contents of the multiboot information
fn display_multiboot<'a>(plat: &mut PlatInterfaceType, mbi: &'a multiboot::Multiboot<'a>) {
    write!(plat, "Multiboot information:\n").unwrap();
    if let Some(low) = mbi.lower_memory_bound() {
        write!(plat,"\t{}kb of low memory\n", low).unwrap();
    }
    if let Some(high) = mbi.upper_memory_bound() {
        write!(plat,"\t{}mb of high memory\n", high / 1024).unwrap();
    }
    if let Some(boot) = mbi.boot_device() {
        write!(plat,"\tBoot device {:?}\n", boot).unwrap();
    }
    if let Some(line) = mbi.command_line() {
        write!(plat,"\tCommand line \"{}\"\n", line).unwrap();
    }
    if let Some(modules) = mbi.modules() {
        write!(plat,"Multiboot modules:\n").unwrap();
        for m in modules {
            write!(plat,"\t{:?}\n", m).unwrap();
        }
    }
    if let Some(memory) = mbi.memory_regions() {
        write!(plat,"Memory regions:\n").unwrap();
        for m in memory {
            write!(plat,"\t{:?}\n", m).unwrap();
        }
    }
}

/// We create a wrapper struct because I don't know how else
/// to get the lifetime of the return value of callback function
/// to line up with the multiboot lifetime
struct MbiWrapper<'a> {
    mbi: Option<multiboot::Multiboot<'a>>,
    callback: &'a Fn(u64, usize) -> Option<&'a [u8]>,
}

fn make_bootconfig<'a>(mbi: &MbiWrapper<'a>) -> BootConfig<'a> {
    let mut cmdline = "";
    if let &Some(ref mbi) = &mbi.mbi {
        if let Some(cmd) = mbi.command_line() {
            cmdline = cmd;
        }
    }
    return BootConfig::new(cmdline);
}

/// Convert the kernel image start and end variables from the linker script
/// into useful values
fn get_kernel_image_region() -> (usize, usize) {
    (&kernel_image_start as *const u8 as usize, &kernel_image_end as *const u8 as usize)
}

/// Try and boot the system, potentially returning an error.
/// Since if we fail there is no way to display or do anything
/// with the error we do do not bother to have an error type
/// This is specifically the 'early' boot as it happens before
/// we switch to the final kernel address space
fn try_early_boot_system<'h, 'l>(init: EarlyBootState<'h, 'l>) -> Result<PostEarlyBootState<'h>, Option<PlatInterfaceType>> {
    /* check that we are multi-booted */
    if init.mbi_magic as u64 != multiboot::SIGNATURE_RAX {
        return Err(None);
    }
    /* construct a reference to the mbi to get the command line */
    let mut mbi = MbiWrapper{mbi: None, callback: &|p, s| unsafe {
        Some(init.low_window.make_slice(p as usize, s))
        }};
    unsafe {
        /* Thanks to the stupidy in this function interface of requiring
         * an unsafe function we cannot pass a lambda expression. Since
         * a lambda cannot be made unsafe and we have no other way of
         * passing our state we just transmute the lambda. Unfortunately
         * this removes all type checking on our function :( */
        mbi.mbi = match multiboot::Multiboot::new(init.mbi as multiboot::PAddr, mbi.callback) {
            Some(mbi) => Some(mbi),
            None => {
                    return Err(None);
                }
            };
    }
    let bootconfig = make_bootconfig(&mbi);
    /* Initial the serial output of our platform first so that
     * we can get debugging output. */
    let mut plat = unsafe {get_platform(&bootconfig)};
    plat.init_serial();
    write!(plat, "R4: In early setup\n").unwrap();
    let (ki_start, ki_end) = get_kernel_image_region();
    write!(plat, "Kernel imagine region {:x} {:x}\n", ki_start, ki_end).unwrap();
    if !init.high_window.range_valid(ki_start, ki_end - ki_start) {
        write!(plat, "Kernel image outside boot window!").unwrap();
        return Err(Some(plat));
    }
    /* Initialize the panic function so we can see anything
     * really bad that happens */
    panic_set_plat(&mut plat);
    /* Now we can continue with the rest of init */
    if let &Some(ref mbi) = &mbi.mbi {
        display_multiboot(&mut plat, &mbi);
    }
    Ok(PostEarlyBootState{ plat: plat, phantom: PhantomData })
}

/// Rust entry point for the kernel. This expects two parameters, one the
/// boot info magic, and the other a raw pointer to the boot info structure.
/// Additionally it expects that both the boot kernel windows are configured
/// correctly in the active address space root
#[no_mangle]
pub extern fn boot_system(magic: usize, mbi: *const usize) -> ! {
    /* This *will* be our final kernel window, but it is not our window
     * yet. We create it here so that our temporary high kernel window,
     * which is a subset of the final kernel window, can be constructed
     * from it as having the same lifetime. This allows any references
     * created from the boot high kernel window being able to live on
     * into the final kernel window */
    let final_window = unsafe{KernelWindow::default()};
    /* this variable will hold our system state as returned by early boot */
    let mut boot;
    {
        /* Construct our system state for boot */
        let boot_high_window;
        unsafe {
            boot_high_window = final_window.subwindow(BootHighWindow::default());
        }
        /* whilst doing early boot we can also access things in low memory */
        let boot_low_window = unsafe {BootLowWindow::default()};
        let boot_state = EarlyBootState {
            high_window: &boot_high_window,
            low_window: &boot_low_window,
            mbi_magic: magic,
            mbi: mbi,
        };
        boot = match try_early_boot_system(boot_state) {
            Err(Some(mut plat)) => {
                    write!(plat, "Failed early init. hlt'ing\n").unwrap();
                    halt();
                },
            Err(None) =>
                halt(),
            Ok(b) => b,
        };
        /* The 'plat' definition got moved into boot. Reset the panic location */
        panic_set_plat(&mut boot.plat);
    }
    unimplemented!()
}
<commit_msg>Don't return a platform under error during early boot. It adds nothing and just provides complication<commit_after>//! Rust entry point
//!
//! Boot is split into two parts
//!
//! + Early boot where we walk hardware data structures and setup initial
//!   machine state before
//! + Late boot where we are operating in the final kernel window and can
//!   setup the initial user environment
//!
//! Early boot is extremely architecture specific, late boot is mostly
//! generic.

use plat::*;
use vspace::*;
use panic::*;
use ::config::BootConfig;
use ::core::fmt::Write;
use ::core::marker::PhantomData;
use ::core::default::Default;
use super::halt::halt;
use super::vspace::*;
extern crate multiboot;

extern {
    /// Represent the start of the kernel image region. The type is not
    /// correct as there is no way to properly represent the type we want
    /// in Rust
    static kernel_image_start: u8;
    static kernel_image_end: u8;
}

/// Package of state that is passed to the early boot function
struct EarlyBootState<'h, 'l> {
    /// Referenece to the high window. Objects created from here can
    /// persist forever, and as such references to them can be returned
    /// in `PostEarlyBootState`
    high_window : &'h BootHighWindow<'h>,
    /// Reference to the low window. References to objects from here
    /// cannot be returned in `PostEarlyBootState`
    low_window : &'l BootLowWindow<'l>,
    /// The value of EAX passed from the assembly entry. This is checked
    /// to ensure we were multiboot loaded
    mbi_magic: usize,
    /// Raw physical pointer that should point to the multiboot structure
    mbi: *const usize,
}

/// Package of state that is returned as a result of early boot
struct PostEarlyBootState<'a> {
    /// An initialized platform interface
    plat: PlatInterfaceType,
    /// Currently the lifetime 'a is unused, so have some `PhantomData` to get
    /// around that
    phantom: PhantomData<&'a usize>,
}

/// Debug function to print out the contents of the multiboot information
fn display_multiboot<'a>(plat: &mut PlatInterfaceType, mbi: &'a multiboot::Multiboot<'a>) {
    write!(plat, "Multiboot information:\n").unwrap();
    if let Some(low) = mbi.lower_memory_bound() {
        write!(plat,"\t{}kb of low memory\n", low).unwrap();
    }
    if let Some(high) = mbi.upper_memory_bound() {
        write!(plat,"\t{}mb of high memory\n", high / 1024).unwrap();
    }
    if let Some(boot) = mbi.boot_device() {
        write!(plat,"\tBoot device {:?}\n", boot).unwrap();
    }
    if let Some(line) = mbi.command_line() {
        write!(plat,"\tCommand line \"{}\"\n", line).unwrap();
    }
    if let Some(modules) = mbi.modules() {
        write!(plat,"Multiboot modules:\n").unwrap();
        for m in modules {
            write!(plat,"\t{:?}\n", m).unwrap();
        }
    }
    if let Some(memory) = mbi.memory_regions() {
        write!(plat,"Memory regions:\n").unwrap();
        for m in memory {
            write!(plat,"\t{:?}\n", m).unwrap();
        }
    }
}

/// We create a wrapper struct because I don't know how else
/// to get the lifetime of the return value of callback function
/// to line up with the multiboot lifetime
struct MbiWrapper<'a> {
    mbi: Option<multiboot::Multiboot<'a>>,
    callback: &'a Fn(u64, usize) -> Option<&'a [u8]>,
}

fn make_bootconfig<'a>(mbi: &MbiWrapper<'a>) -> BootConfig<'a> {
    let mut cmdline = "";
    if let &Some(ref mbi) = &mbi.mbi {
        if let Some(cmd) = mbi.command_line() {
            cmdline = cmd;
        }
    }
    return BootConfig::new(cmdline);
}

/// Convert the kernel image start and end variables from the linker script
/// into useful values
fn get_kernel_image_region() -> (usize, usize) {
    (&kernel_image_start as *const u8 as usize, &kernel_image_end as *const u8 as usize)
}

/// Try and boot the system, potentially returning an error.
/// Since if we fail there is no way to display or do anything
/// with the error we do do not bother to have an error type
/// This is specifically the 'early' boot as it happens before
/// we switch to the final kernel address space
fn try_early_boot_system<'h, 'l>(init: EarlyBootState<'h, 'l>) -> Result<PostEarlyBootState<'h>, ()> {
    /* check that we are multi-booted */
    if init.mbi_magic as u64 != multiboot::SIGNATURE_RAX {
        return Err(());
    }
    /* construct a reference to the mbi to get the command line */
    let mut mbi = MbiWrapper{mbi: None, callback: &|p, s| unsafe {
        Some(init.low_window.make_slice(p as usize, s))
        }};
    unsafe {
        /* Thanks to the stupidy in this function interface of requiring
         * an unsafe function we cannot pass a lambda expression. Since
         * a lambda cannot be made unsafe and we have no other way of
         * passing our state we just transmute the lambda. Unfortunately
         * this removes all type checking on our function :( */
        mbi.mbi = match multiboot::Multiboot::new(init.mbi as multiboot::PAddr, mbi.callback) {
            Some(mbi) => Some(mbi),
            None => {
                    return Err(());
                }
            };
    }
    let bootconfig = make_bootconfig(&mbi);
    /* Initial the serial output of our platform first so that
     * we can get debugging output. */
    let mut plat = unsafe {get_platform(&bootconfig)};
    plat.init_serial();
    write!(plat, "R4: In early setup\n").unwrap();
    let (ki_start, ki_end) = get_kernel_image_region();
    write!(plat, "Kernel imagine region {:x} {:x}\n", ki_start, ki_end).unwrap();
    if !init.high_window.range_valid(ki_start, ki_end - ki_start) {
        write!(plat, "Kernel image outside boot window!").unwrap();
        return Err(());
    }
    /* Initialize the panic function so we can see anything
     * really bad that happens */
    panic_set_plat(&mut plat);
    /* Now we can continue with the rest of init */
    if let &Some(ref mbi) = &mbi.mbi {
        display_multiboot(&mut plat, &mbi);
    }
    Ok(PostEarlyBootState{ plat: plat, phantom: PhantomData })
}

/// Rust entry point for the kernel. This expects two parameters, one the
/// boot info magic, and the other a raw pointer to the boot info structure.
/// Additionally it expects that both the boot kernel windows are configured
/// correctly in the active address space root
#[no_mangle]
pub extern fn boot_system(magic: usize, mbi: *const usize) -> ! {
    /* This *will* be our final kernel window, but it is not our window
     * yet. We create it here so that our temporary high kernel window,
     * which is a subset of the final kernel window, can be constructed
     * from it as having the same lifetime. This allows any references
     * created from the boot high kernel window being able to live on
     * into the final kernel window */
    let final_window = unsafe{KernelWindow::default()};
    /* this variable will hold our system state as returned by early boot */
    let mut boot;
    {
        /* Construct our system state for boot */
        let boot_high_window;
        unsafe {
            boot_high_window = final_window.subwindow(BootHighWindow::default());
        }
        /* whilst doing early boot we can also access things in low memory */
        let boot_low_window = unsafe {BootLowWindow::default()};
        let boot_state = EarlyBootState {
            high_window: &boot_high_window,
            low_window: &boot_low_window,
            mbi_magic: magic,
            mbi: mbi,
        };
        boot = match try_early_boot_system(boot_state) {
            Err(_) => halt(),
            Ok(b) => b,
        };
        /* The 'plat' definition got moved into boot. Reset the panic location */
        panic_set_plat(&mut boot.plat);
    }
    unimplemented!()
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate byteorder;
extern crate clap;
extern crate pbr;
extern crate point_viewer;
extern crate scoped_pool;
extern crate prost;

use pbr::ProgressBar;
use point_viewer::{InternalIterator, Point};
use point_viewer::errors::*;
use point_viewer::math::{Cube, Cuboid, CuboidLike};
use point_viewer::octree;
use point_viewer::ply::PlyIterator;
use point_viewer::proto;
use point_viewer::pts::PtsIterator;
use prost::Message;
use scoped_pool::{Pool, Scope};
use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{Stdout, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc;

const UPDATE_COUNT: i64 = 100000;
const MAX_POINTS_PER_NODE: i64 = 100000;

struct SplittedNode {
    node: octree::Node,
    num_points: i64,
}

fn split<P>(
    output_directory: &Path,
    resolution: f64,
    node: &octree::Node,
    stream: P,
) -> Vec<SplittedNode>
    where P: InternalIterator
{
    let mut children: Vec<Option<octree::NodeWriter>> =
        vec![None, None, None, None, None, None, None, None];
    match stream.size_hint() {
        Some(size) => {
            println!(
                "Splitting {} which has {} points ({:.2}x MAX_POINTS_PER_NODE).",
                node.id,
                size,
                size as f64 / MAX_POINTS_PER_NODE as f64
            )
        }
        None => {
            println!(
                "Splitting {} which has an unknown number of points.",
                node.id
            )
        }
    };

    stream.for_each(
        |p| {
            let child_index = node.get_child_id_containing_point(&p.position);
            let array_index = child_index.as_u8() as usize;
            if children[array_index].is_none() {
                children[array_index] = Some(
                    octree::NodeWriter::new(
                        output_directory,
                        &node.get_child(child_index),
                        resolution,
                    )
                );
            }
            children[array_index].as_mut().unwrap().write(&p);
        }
    );

    // Remove the node file on disk by reopening the node and immediately dropping it again without
    // writing a point. This only saves some disk space during processing - all nodes will be
    // rewritten by subsampling the children in the second step anyways. We also ignore file
    // removing error. For example, we never write out the root, so it cannot be removed.
    octree::NodeWriter::new(output_directory, &node, resolution);

    let mut rv = Vec::new();
    for (child_index, c) in children.into_iter().enumerate() {
        if c.is_none() {
            continue;
        }
        let c = c.unwrap();

        rv.push(
            SplittedNode {
                node: node.get_child(octree::ChildIndex::from_u8(child_index as u8)),
                num_points: c.num_written(),
            }
        );
    }
    rv
}

fn split_node<'a, 'b: 'a, P>(
    scope: &Scope<'a>,
    output_directory: &'b Path,
    resolution: f64,
    splitted_node: SplittedNode,
    stream: P,
    leaf_nodes_sender: mpsc::Sender<octree::Node>,
) where P: InternalIterator
{
    let children = split(output_directory, resolution, &splitted_node.node, stream);
    let (leaf_nodes, split_nodes): (Vec<_>, Vec<_>) =
        children
            .into_iter()
            .partition(|n| n.num_points < MAX_POINTS_PER_NODE);

    for child in split_nodes {
        let leaf_nodes_sender_clone = leaf_nodes_sender.clone();
        scope.recurse(
            move |scope| {
                let stream = octree::NodeIterator::from_disk(output_directory, &child.node.id)
                    .unwrap();
                split_node(
                    scope,
                    output_directory,
                    resolution,
                    child,
                    stream,
                    leaf_nodes_sender_clone,
                );
            }
        );
    }

    for splitted_node in leaf_nodes {
        leaf_nodes_sender.send(splitted_node.node).unwrap();
    }
}

fn subsample_children_into(
    output_directory: &Path,
    node: &octree::Node,
    resolution: f64,
) -> Result<()> {
    let mut parent_writer = octree::NodeWriter::new(output_directory, &node, resolution);
    println!("Creating {} from subsampling children.", &node.id);
    for i in 0..8 {
        let child = node.get_child(octree::ChildIndex::from_u8(i));
        let node_iterator = match octree::NodeIterator::from_disk(output_directory, &child.id) {
            Ok(node_iterator) => node_iterator,
            Err(Error(ErrorKind::NodeNotFound, _)) => continue,
            Err(err) => return Err(err),
        };

        // We read all points into memory, because the new node writer will rewrite this child's
        // file(s).
        let mut points = Vec::with_capacity(node_iterator.size_hint().unwrap());
        node_iterator.for_each(|p| points.push((*p).clone()));

        let mut child_writer = octree::NodeWriter::new(output_directory, &child, resolution);
        for (idx, p) in points.into_iter().enumerate() {
            if idx % 8 == 0 {
                parent_writer.write(&p);
            } else {
                child_writer.write(&p);
            }
        }
    }
    Ok(())
}

#[derive(Debug)]
enum InputFile {
    Ply(PathBuf),
    Pts(PathBuf),
}

enum InputFileIterator {
    Ply(PlyIterator),
    Pts(PtsIterator),
}

impl InternalIterator for InputFileIterator {
    fn size_hint(&self) -> Option<usize> {
        match *self {
            InputFileIterator::Ply(ref p) => p.size_hint(),
            InputFileIterator::Pts(ref p) => p.size_hint(),
        }
    }

    fn for_each<F: FnMut(&Point)>(self, f: F) {
        match self {
            InputFileIterator::Ply(p) => p.for_each(f),
            InputFileIterator::Pts(p) => p.for_each(f),
        }
    }
}

fn make_stream(input: &InputFile) -> (InputFileIterator, Option<pbr::ProgressBar<Stdout>>) {
    let stream = match *input {
        InputFile::Ply(ref filename) => InputFileIterator::Ply(PlyIterator::new(filename).unwrap()),
        InputFile::Pts(ref filename) => InputFileIterator::Pts(PtsIterator::new(filename)),
    };

    let progress_bar = match stream.size_hint() {
        Some(size) => Some(ProgressBar::new(size as u64)),
        None => None,
    };
    (stream, progress_bar)
}

/// Returns the bounding_cube and the number of the points in 'input'.
fn find_bounding_cube(input: &InputFile) -> (Cube, i64) {
    let mut num_points = 0i64;
    let mut bounding_cube = Cuboid::new();
    let (stream, mut progress_bar) = make_stream(input);
    progress_bar
        .as_mut()
        .map(|pb| pb.message("Determining bounding box: "));

    stream.for_each(
        |p: &Point| {
            bounding_cube.update(&p.position);
            num_points += 1;
            if num_points % UPDATE_COUNT == 0 {
                progress_bar.as_mut().map(|pb| pb.add(UPDATE_COUNT as u64));
            }
        }
    );
    progress_bar.map(|mut f| f.finish());
    (bounding_cube.to_cube(), num_points)
}

fn main() {
    let matches = clap::App::new("build_octree")
        .args(
            &[
                clap::Arg::with_name("output_directory")
                    .help("Output directory to write the octree into.")
                    .long("output_directory")
                    .required(true)
                    .takes_value(true),
                clap::Arg::with_name("resolution")
                    .help(
                        "Minimal precision that this point cloud should have. This decides \
                           on the number of bits used to encode each node."
                    )
                    .long("resolution")
                    .default_value("0.001"),
                clap::Arg::with_name("input")
                    .help("PLY/PTS file to parse for the points.")
                    .index(1)
                    .required(true),
            ]
        )
        .get_matches();

    let output_directory = &PathBuf::from(matches.value_of("output_directory").unwrap());
    let resolution = matches
        .value_of("resolution")
        .unwrap()
        .parse::<f64>()
        .expect("resolution could not be parsed as float.");

    let input = {
        let filename = PathBuf::from(matches.value_of("input").unwrap());
        match filename.extension().and_then(|s| s.to_str()) {
            Some("ply") => InputFile::Ply(filename.clone()),
            Some("pts") => InputFile::Pts(filename.clone()),
            other => panic!("Unknown input file format: {:?}", other),
        }
    };

    let (bounding_cube, num_points) = find_bounding_cube(&input);

    // Ignore errors, maybe directory is already there.
    let _ = fs::create_dir(output_directory);

    let meta = proto::Meta {
        bounding_cube: Some(
            proto::BoundingCube {
                min: Some(
                    proto::Vector3f {
                        x: Some(bounding_cube.min().x),
                        y: Some(bounding_cube.min().y),
                        z: Some(bounding_cube.min().z),
                    }
                ),
                edge_length: Some(bounding_cube.edge_length()),
            }
        ),
        resolution: Some(resolution),
        version: Some(octree::CURRENT_VERSION),
    };
    let mut encoded = Vec::new();
    meta.encode(&mut encoded).unwrap();
    File::create(&output_directory.join("meta.pb"))
        .unwrap()
        .write_all(&encoded)
        .unwrap();

    println!("Creating octree structure.");
    let pool = Pool::new(10);

    let (leaf_nodes_sender, leaf_nodes_receiver) = mpsc::channel();
    pool.scoped(
        move |scope| {
            let (root_stream, _) = make_stream(&input);
            let root = SplittedNode {
                node: octree::Node::root_with_bounding_cube(bounding_cube),
                num_points: num_points,
            };
            split_node(
                scope,
                output_directory,
                resolution,
                root,
                root_stream,
                leaf_nodes_sender.clone(),
            );
        }
    );

    let mut deepest_level = 0usize;
    let mut nodes_to_subsample = Vec::<octree::Node>::new();
    for leaf_node in leaf_nodes_receiver.into_iter() {
        deepest_level = std::cmp::max(deepest_level, leaf_node.level());
        nodes_to_subsample.push(leaf_node);
    }

    // We start on the deepest level and work our way up the tree.
    for current_level in (0..deepest_level + 1).rev() {
        // All nodes on the same level can be subsampled in parallel.
        let res = nodes_to_subsample
            .into_iter()
            .partition(|n| n.level() == current_level);
        nodes_to_subsample = res.1;

        let mut parent_ids = HashSet::new();
        let mut subsample_nodes = Vec::new();
        for node in res.0 {
            let maybe_parent = node.parent();
            if maybe_parent.is_none() {
                // Only the root has no parents.
                assert_eq!(node.level(), 0);
                continue;
            }
            let parent = maybe_parent.unwrap();
            if parent_ids.contains(&parent.id) {
                continue;
            }
            parent_ids.insert(parent.id);
            subsample_nodes.push(parent);
        }

        pool.scoped(
            |scope| for node in &subsample_nodes {
                scope.execute(
                    move || {
                        subsample_children_into(output_directory, node, resolution).unwrap();
                    }
                );
            }
        );

        // The nodes that were just now created through sub-sampling will be required to create
        // their parents.
        nodes_to_subsample.extend(subsample_nodes.into_iter());
    }
}
<commit_msg>Do not split a node that is smaller than the desired resolution. (#42)<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate byteorder;
extern crate clap;
extern crate pbr;
extern crate point_viewer;
extern crate scoped_pool;
extern crate prost;

use pbr::ProgressBar;
use point_viewer::{InternalIterator, Point};
use point_viewer::errors::*;
use point_viewer::math::{Cube, Cuboid, CuboidLike};
use point_viewer::octree;
use point_viewer::ply::PlyIterator;
use point_viewer::proto;
use point_viewer::pts::PtsIterator;
use prost::Message;
use scoped_pool::{Pool, Scope};
use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{Stdout, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc;

const UPDATE_COUNT: i64 = 100000;
const MAX_POINTS_PER_NODE: i64 = 100000;

struct SplittedNode {
    node: octree::Node,
    num_points: i64,
}

fn split<P>(
    output_directory: &Path,
    resolution: f64,
    node: &octree::Node,
    stream: P,
) -> Vec<SplittedNode>
    where P: InternalIterator
{
    let mut children: Vec<Option<octree::NodeWriter>> =
        vec![None, None, None, None, None, None, None, None];
    match stream.size_hint() {
        Some(size) => {
            println!(
                "Splitting {} which has {} points ({:.2}x MAX_POINTS_PER_NODE).",
                node.id,
                size,
                size as f64 / MAX_POINTS_PER_NODE as f64
            )
        }
        None => {
            println!(
                "Splitting {} which has an unknown number of points.",
                node.id
            )
        }
    };

    stream.for_each(
        |p| {
            let child_index = node.get_child_id_containing_point(&p.position);
            let array_index = child_index.as_u8() as usize;
            if children[array_index].is_none() {
                children[array_index] = Some(
                    octree::NodeWriter::new(
                        output_directory,
                        &node.get_child(child_index),
                        resolution,
                    )
                );
            }
            children[array_index].as_mut().unwrap().write(&p);
        }
    );

    // Remove the node file on disk by reopening the node and immediately dropping it again without
    // writing a point. This only saves some disk space during processing - all nodes will be
    // rewritten by subsampling the children in the second step anyways. We also ignore file
    // removing error. For example, we never write out the root, so it cannot be removed.
    octree::NodeWriter::new(output_directory, &node, resolution);

    let mut rv = Vec::new();
    for (child_index, c) in children.into_iter().enumerate() {
        if c.is_none() {
            continue;
        }
        let c = c.unwrap();

        rv.push(
            SplittedNode {
                node: node.get_child(octree::ChildIndex::from_u8(child_index as u8)),
                num_points: c.num_written(),
            }
        );
    }
    rv
}

fn should_split_node(node: &SplittedNode, resolution: f64) -> bool {
    if node.num_points <= MAX_POINTS_PER_NODE {
        return false;
    }
    if node.node.bounding_cube.edge_length() as f64 <= resolution {
        // TODO(hrapp): If the data has billion of points in this small spot, performance will
        // greatly suffer if we display it. Drop points?
        println!(
            "Node {} which has {} points ({:.2}x MAX_POINTS_PER_NODE) \
            is too small to be split, keeping all points.",
            node.node.id,
            node.num_points,
            node.num_points as f64 / MAX_POINTS_PER_NODE as f64
        );
        return false;
    }
    true
}

fn split_node<'a, 'b: 'a, P>(
    scope: &Scope<'a>,
    output_directory: &'b Path,
    resolution: f64,
    splitted_node: SplittedNode,
    stream: P,
    leaf_nodes_sender: mpsc::Sender<octree::Node>,
) where P: InternalIterator
{
    let children = split(output_directory, resolution, &splitted_node.node, stream);
    let (leaf_nodes, split_nodes): (Vec<_>, Vec<_>) =
        children
            .into_iter()
            .partition(|n| !should_split_node(n, resolution));

    for child in split_nodes {
        let leaf_nodes_sender_clone = leaf_nodes_sender.clone();
        scope.recurse(
            move |scope| {
                let stream = octree::NodeIterator::from_disk(output_directory, &child.node.id)
                    .unwrap();
                split_node(
                    scope,
                    output_directory,
                    resolution,
                    child,
                    stream,
                    leaf_nodes_sender_clone,
                );
            }
        );
    }

    for splitted_node in leaf_nodes {
        leaf_nodes_sender.send(splitted_node.node).unwrap();
    }
}

fn subsample_children_into(
    output_directory: &Path,
    node: &octree::Node,
    resolution: f64,
) -> Result<()> {
    let mut parent_writer = octree::NodeWriter::new(output_directory, &node, resolution);
    println!("Creating {} from subsampling children.", &node.id);
    for i in 0..8 {
        let child = node.get_child(octree::ChildIndex::from_u8(i));
        let node_iterator = match octree::NodeIterator::from_disk(output_directory, &child.id) {
            Ok(node_iterator) => node_iterator,
            Err(Error(ErrorKind::NodeNotFound, _)) => continue,
            Err(err) => return Err(err),
        };

        // We read all points into memory, because the new node writer will rewrite this child's
        // file(s).
        let mut points = Vec::with_capacity(node_iterator.size_hint().unwrap());
        node_iterator.for_each(|p| points.push((*p).clone()));

        let mut child_writer = octree::NodeWriter::new(output_directory, &child, resolution);
        for (idx, p) in points.into_iter().enumerate() {
            if idx % 8 == 0 {
                parent_writer.write(&p);
            } else {
                child_writer.write(&p);
            }
        }
    }
    Ok(())
}

#[derive(Debug)]
enum InputFile {
    Ply(PathBuf),
    Pts(PathBuf),
}

enum InputFileIterator {
    Ply(PlyIterator),
    Pts(PtsIterator),
}

impl InternalIterator for InputFileIterator {
    fn size_hint(&self) -> Option<usize> {
        match *self {
            InputFileIterator::Ply(ref p) => p.size_hint(),
            InputFileIterator::Pts(ref p) => p.size_hint(),
        }
    }

    fn for_each<F: FnMut(&Point)>(self, f: F) {
        match self {
            InputFileIterator::Ply(p) => p.for_each(f),
            InputFileIterator::Pts(p) => p.for_each(f),
        }
    }
}

fn make_stream(input: &InputFile) -> (InputFileIterator, Option<pbr::ProgressBar<Stdout>>) {
    let stream = match *input {
        InputFile::Ply(ref filename) => InputFileIterator::Ply(PlyIterator::new(filename).unwrap()),
        InputFile::Pts(ref filename) => InputFileIterator::Pts(PtsIterator::new(filename)),
    };

    let progress_bar = match stream.size_hint() {
        Some(size) => Some(ProgressBar::new(size as u64)),
        None => None,
    };
    (stream, progress_bar)
}

/// Returns the bounding_cube and the number of the points in 'input'.
fn find_bounding_cube(input: &InputFile) -> (Cube, i64) {
    let mut num_points = 0i64;
    let mut bounding_cube = Cuboid::new();
    let (stream, mut progress_bar) = make_stream(input);
    progress_bar
        .as_mut()
        .map(|pb| pb.message("Determining bounding box: "));

    stream.for_each(
        |p: &Point| {
            bounding_cube.update(&p.position);
            num_points += 1;
            if num_points % UPDATE_COUNT == 0 {
                progress_bar.as_mut().map(|pb| pb.add(UPDATE_COUNT as u64));
            }
        }
    );
    progress_bar.map(|mut f| f.finish());
    (bounding_cube.to_cube(), num_points)
}

fn main() {
    let matches = clap::App::new("build_octree")
        .args(
            &[
                clap::Arg::with_name("output_directory")
                    .help("Output directory to write the octree into.")
                    .long("output_directory")
                    .required(true)
                    .takes_value(true),
                clap::Arg::with_name("resolution")
                    .help(
                        "Minimal precision that this point cloud should have. This decides \
                           on the number of bits used to encode each node."
                    )
                    .long("resolution")
                    .default_value("0.001"),
                clap::Arg::with_name("input")
                    .help("PLY/PTS file to parse for the points.")
                    .index(1)
                    .required(true),
            ]
        )
        .get_matches();

    let output_directory = &PathBuf::from(matches.value_of("output_directory").unwrap());
    let resolution = matches
        .value_of("resolution")
        .unwrap()
        .parse::<f64>()
        .expect("resolution could not be parsed as float.");

    let input = {
        let filename = PathBuf::from(matches.value_of("input").unwrap());
        match filename.extension().and_then(|s| s.to_str()) {
            Some("ply") => InputFile::Ply(filename.clone()),
            Some("pts") => InputFile::Pts(filename.clone()),
            other => panic!("Unknown input file format: {:?}", other),
        }
    };

    let (bounding_cube, num_points) = find_bounding_cube(&input);

    // Ignore errors, maybe directory is already there.
    let _ = fs::create_dir(output_directory);

    let meta = proto::Meta {
        bounding_cube: Some(
            proto::BoundingCube {
                min: Some(
                    proto::Vector3f {
                        x: Some(bounding_cube.min().x),
                        y: Some(bounding_cube.min().y),
                        z: Some(bounding_cube.min().z),
                    }
                ),
                edge_length: Some(bounding_cube.edge_length()),
            }
        ),
        resolution: Some(resolution),
        version: Some(octree::CURRENT_VERSION),
    };
    let mut encoded = Vec::new();
    meta.encode(&mut encoded).unwrap();
    File::create(&output_directory.join("meta.pb"))
        .unwrap()
        .write_all(&encoded)
        .unwrap();

    println!("Creating octree structure.");
    let pool = Pool::new(10);

    let (leaf_nodes_sender, leaf_nodes_receiver) = mpsc::channel();
    pool.scoped(
        move |scope| {
            let (root_stream, _) = make_stream(&input);
            let root = SplittedNode {
                node: octree::Node::root_with_bounding_cube(bounding_cube),
                num_points: num_points,
            };
            split_node(
                scope,
                output_directory,
                resolution,
                root,
                root_stream,
                leaf_nodes_sender.clone(),
            );
        }
    );

    let mut deepest_level = 0usize;
    let mut nodes_to_subsample = Vec::<octree::Node>::new();
    for leaf_node in leaf_nodes_receiver.into_iter() {
        deepest_level = std::cmp::max(deepest_level, leaf_node.level());
        nodes_to_subsample.push(leaf_node);
    }

    // We start on the deepest level and work our way up the tree.
    for current_level in (0..deepest_level + 1).rev() {
        // All nodes on the same level can be subsampled in parallel.
        let res = nodes_to_subsample
            .into_iter()
            .partition(|n| n.level() == current_level);
        nodes_to_subsample = res.1;

        let mut parent_ids = HashSet::new();
        let mut subsample_nodes = Vec::new();
        for node in res.0 {
            let maybe_parent = node.parent();
            if maybe_parent.is_none() {
                // Only the root has no parents.
                assert_eq!(node.level(), 0);
                continue;
            }
            let parent = maybe_parent.unwrap();
            if parent_ids.contains(&parent.id) {
                continue;
            }
            parent_ids.insert(parent.id);
            subsample_nodes.push(parent);
        }

        pool.scoped(
            |scope| for node in &subsample_nodes {
                scope.execute(
                    move || {
                        subsample_children_into(output_directory, node, resolution).unwrap();
                    }
                );
            }
        );

        // The nodes that were just now created through sub-sampling will be required to create
        // their parents.
        nodes_to_subsample.extend(subsample_nodes.into_iter());
    }
}
<|endoftext|>