<commit_before>use futures::sync::oneshot;
use futures::{future, Future};
use std::borrow::Cow;
use std::mem;
use std::sync::mpsc::{RecvError, TryRecvError};
use std::thread;
use std;

use core::config::{Bitrate, PlayerConfig};
use core::session::Session;
use core::util::{self, SpotifyId, Subfile};

use audio_backend::Sink;
use audio::{AudioFile, AudioDecrypt};
use audio::{VorbisDecoder, VorbisPacket};
use metadata::{FileFormat, Track, Metadata};
use mixer::AudioFilter;

#[derive(Clone)]
pub struct Player {
    commands: std::sync::mpsc::Sender<PlayerCommand>,
}

struct PlayerInternal {
    session: Session,
    config: PlayerConfig,
    commands: std::sync::mpsc::Receiver<PlayerCommand>,

    state: PlayerState,
    sink: Box<Sink>,
    audio_filter: Option<Box<AudioFilter + Send>>,
}

enum PlayerCommand {
    Load(SpotifyId, bool, u32, oneshot::Sender<()>),
    Play,
    Pause,
    Stop,
    Seek(u32),
}

impl Player {
    pub fn new<F>(config: PlayerConfig, session: Session,
                  audio_filter: Option<Box<AudioFilter + Send>>,
                  sink_builder: F) -> Player
        where F: FnOnce() -> Box<Sink> + Send + 'static
    {
        let (cmd_tx, cmd_rx) = std::sync::mpsc::channel();

        thread::spawn(move || {
            debug!("new Player[{}]", session.session_id());

            let internal = PlayerInternal {
                session: session,
                config: config,
                commands: cmd_rx,

                state: PlayerState::Stopped,
                sink: sink_builder(),
                audio_filter: audio_filter,
            };

            internal.run();
        });

        Player {
            commands: cmd_tx,
        }
    }

    fn command(&self, cmd: PlayerCommand) {
        self.commands.send(cmd).unwrap();
    }

    pub fn load(&self, track: SpotifyId, start_playing: bool, position_ms: u32)
        -> oneshot::Receiver<()>
    {
        let (tx, rx) = oneshot::channel();
        self.command(PlayerCommand::Load(track, start_playing, position_ms, tx));

        rx
    }

    pub fn play(&self) {
        self.command(PlayerCommand::Play)
    }

    pub fn pause(&self) {
        self.command(PlayerCommand::Pause)
    }

    pub fn stop(&self) {
        self.command(PlayerCommand::Stop)
    }

    pub fn seek(&self, position_ms: u32) {
        self.command(PlayerCommand::Seek(position_ms));
    }
}

type Decoder = VorbisDecoder<Subfile<AudioDecrypt<AudioFile>>>;
enum PlayerState {
    Stopped,
    Paused {
        decoder: Decoder,
        end_of_track: oneshot::Sender<()>,
    },
    Playing {
        decoder: Decoder,
        end_of_track: oneshot::Sender<()>,
    },

    Invalid,
}

impl PlayerState {
    fn is_playing(&self) -> bool {
        use self::PlayerState::*;
        match *self {
            Stopped | Paused { .. } => false,
            Playing { .. } => true,
            Invalid => panic!("invalid state"),
        }
    }

    fn decoder(&mut self) -> Option<&mut Decoder> {
        use self::PlayerState::*;
        match *self {
            Stopped => None,
            Paused { ref mut decoder, .. } |
            Playing { ref mut decoder, .. } => Some(decoder),
            Invalid => panic!("invalid state"),
        }
    }

    fn signal_end_of_track(self) {
        use self::PlayerState::*;
        match self {
            Paused { end_of_track, .. } |
            Playing { end_of_track, .. } => {
                end_of_track.complete(())
            }

            Stopped => warn!("signal_end_of_track from stopped state"),
            Invalid => panic!("invalid state"),
        }
    }

    fn paused_to_playing(&mut self) {
        use self::PlayerState::*;
        match ::std::mem::replace(self, Invalid) {
            Paused { decoder, end_of_track } => {
                *self = Playing {
                    decoder: decoder,
                    end_of_track: end_of_track,
                };
            }
            _ => panic!("invalid state"),
        }
    }

    fn playing_to_paused(&mut self) {
        use self::PlayerState::*;
        match ::std::mem::replace(self, Invalid) {
            Playing { decoder, end_of_track } => {
                *self = Paused {
                    decoder: decoder,
                    end_of_track: end_of_track,
                };
            }
            _ => panic!("invalid state"),
        }
    }
}

impl PlayerInternal {
    fn run(mut self) {
        loop {
            let cmd = if self.state.is_playing() {
                match self.commands.try_recv() {
                    Ok(cmd) => Some(cmd),
                    Err(TryRecvError::Empty) => None,
                    Err(TryRecvError::Disconnected) => return,
                }
            } else {
                match self.commands.recv() {
                    Ok(cmd) => Some(cmd),
                    Err(RecvError) => return,
                }
            };

            if let Some(cmd) = cmd {
                self.handle_command(cmd);
            }

            let packet = if let PlayerState::Playing { ref mut decoder, .. } = self.state {
                Some(decoder.next_packet().expect("Vorbis error"))
            } else { None };

            if let Some(packet) = packet {
                self.handle_packet(packet);
            }
        }
    }

    fn handle_packet(&mut self, packet: Option<VorbisPacket>) {
        match packet {
            Some(mut packet) => {
                if let Some(ref editor) = self.audio_filter {
                    editor.modify_stream(&mut packet.data_mut())
                };

                self.sink.write(&packet.data()).unwrap();
            }

            None => {
                self.sink.stop().unwrap();
                self.run_onstop();

                let old_state = mem::replace(&mut self.state, PlayerState::Stopped);
                old_state.signal_end_of_track();
            }
        }
    }

    fn handle_command(&mut self, cmd: PlayerCommand) {
        debug!("command={:?}", cmd);
        match cmd {
            PlayerCommand::Load(track_id, play, position, end_of_track) => {
                if self.state.is_playing() {
                    self.sink.stop().unwrap();
                }

                match self.load_track(track_id, position as i64) {
                    Some(decoder) => {
                        if play {
                            if !self.state.is_playing() {
                                self.run_onstart();
                            }
                            self.sink.start().unwrap();

                            self.state = PlayerState::Playing {
                                decoder: decoder,
                                end_of_track: end_of_track,
                            };
                        } else {
                            if self.state.is_playing() {
                                self.run_onstop();
                            }

                            self.state = PlayerState::Paused {
                                decoder: decoder,
                                end_of_track: end_of_track,
                            };
                        }
                    }

                    None => {
                        end_of_track.complete(());
                        if self.state.is_playing() {
                            self.run_onstop();
                        }
                    }
                }
            }

            PlayerCommand::Seek(position) => {
                if let Some(decoder) = self.state.decoder() {
                    match decoder.seek(position as i64) {
                        Ok(_) => (),
                        Err(err) => error!("Vorbis error: {:?}", err),
                    }
                } else {
                    warn!("Player::seek called from invalid state");
                }
            }

            PlayerCommand::Play => {
                if let PlayerState::Paused { .. } = self.state {
                    self.state.paused_to_playing();

                    self.run_onstart();
                    self.sink.start().unwrap();
                } else {
                    warn!("Player::play called from invalid state");
                }
            }

            PlayerCommand::Pause => {
                if let PlayerState::Playing { .. } = self.state {
                    self.state.playing_to_paused();

                    self.sink.stop().unwrap();
                    self.run_onstop();
                } else {
                    warn!("Player::pause called from invalid state");
                }
            }

            PlayerCommand::Stop => {
                match self.state {
                    PlayerState::Playing { .. } => {
                        self.sink.stop().unwrap();
                        self.run_onstop();
                        self.state = PlayerState::Stopped;
                    }
                    PlayerState::Paused { .. } => {
                        self.state = PlayerState::Stopped;
                    },
                    PlayerState::Stopped => {
                        warn!("Player::stop called from invalid state");
                    }
                    PlayerState::Invalid => panic!("invalid state"),
                }
            }
        }
    }

    fn run_onstart(&self) {
        if let Some(ref program) = self.config.onstart {
            util::run_program(program)
        }
    }

    fn run_onstop(&self) {
        if let Some(ref program) = self.config.onstop {
            util::run_program(program)
        }
    }

    fn find_available_alternative<'a>(&self, track: &'a Track) -> Option<Cow<'a, Track>> {
        if track.available {
            Some(Cow::Borrowed(track))
        } else {
            let alternatives = track.alternatives
                .iter()
                .map(|alt_id| {
                    Track::get(&self.session, *alt_id)
                });
            let alternatives = future::join_all(alternatives).wait().unwrap();

            alternatives.into_iter().find(|alt| alt.available).map(Cow::Owned)
        }
    }

    fn load_track(&self, track_id: SpotifyId, position: i64) -> Option<Decoder> {
        let track = Track::get(&self.session, track_id).wait().unwrap();

        info!("Loading track \"{}\"", track.name);

        let track = match self.find_available_alternative(&track) {
            Some(track) => track,
            None => {
                warn!("Track \"{}\" is not available", track.name);
                return None;
            }
        };

        let format = match self.config.bitrate {
            Bitrate::Bitrate96 => FileFormat::OGG_VORBIS_96,
            Bitrate::Bitrate160 => FileFormat::OGG_VORBIS_160,
            Bitrate::Bitrate320 => FileFormat::OGG_VORBIS_320,
        };

        let file_id = match track.files.get(&format) {
            Some(&file_id) => file_id,
            None => {
                warn!("Track \"{}\" is not available in format {:?}", track.name, format);
                return None;
            }
        };

        let key = self.session.audio_key().request(track.id, file_id).wait().unwrap();

        let encrypted_file = AudioFile::open(&self.session, file_id).wait().unwrap();
        let audio_file = Subfile::new(AudioDecrypt::new(key, encrypted_file), 0xa7);

        let mut decoder = VorbisDecoder::new(audio_file).unwrap();

        match decoder.seek(position) {
            Ok(_) => (),
            Err(err) => error!("Vorbis error: {:?}", err),
        }

        info!("Track \"{}\" loaded", track.name);

        Some(decoder)
    }
}

impl Drop for PlayerInternal {
    fn drop(&mut self) {
        debug!("drop Player[{}]", self.session.session_id());
    }
}

impl ::std::fmt::Debug for PlayerCommand {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match *self {
            PlayerCommand::Load(track, play, position, _) => {
                f.debug_tuple("Load")
                 .field(&track)
                 .field(&play)
                 .field(&position)
                 .finish()
            }
            PlayerCommand::Play => {
                f.debug_tuple("Play").finish()
            }
            PlayerCommand::Pause => {
                f.debug_tuple("Pause").finish()
            }
            PlayerCommand::Stop => {
                f.debug_tuple("Stop").finish()
            }
            PlayerCommand::Seek(position) => {
                f.debug_tuple("Seek")
                 .field(&position)
                 .finish()
            }
        }
    }
}
<commit_msg>Join the player thread when the player is dropped<commit_after>use futures::sync::oneshot;
use futures::{future, Future};
use std::borrow::Cow;
use std::mem;
use std::sync::mpsc::{RecvError, TryRecvError};
use std::thread;
use std;

use core::config::{Bitrate, PlayerConfig};
use core::session::Session;
use core::util::{self, SpotifyId, Subfile};

use audio_backend::Sink;
use audio::{AudioFile, AudioDecrypt};
use audio::{VorbisDecoder, VorbisPacket};
use metadata::{FileFormat, Track, Metadata};
use mixer::AudioFilter;

pub struct Player {
    commands: Option<std::sync::mpsc::Sender<PlayerCommand>>,
    thread_handle: Option<thread::JoinHandle<()>>,
}

struct PlayerInternal {
    session: Session,
    config: PlayerConfig,
    commands: std::sync::mpsc::Receiver<PlayerCommand>,

    state: PlayerState,
    sink: Box<Sink>,
    audio_filter: Option<Box<AudioFilter + Send>>,
}

enum PlayerCommand {
    Load(SpotifyId, bool, u32, oneshot::Sender<()>),
    Play,
    Pause,
    Stop,
    Seek(u32),
}

impl Player {
    pub fn new<F>(config: PlayerConfig, session: Session,
                  audio_filter: Option<Box<AudioFilter + Send>>,
                  sink_builder: F) -> Player
        where F: FnOnce() -> Box<Sink> + Send + 'static
    {
        let (cmd_tx, cmd_rx) = std::sync::mpsc::channel();

        let handle = thread::spawn(move || {
            debug!("new Player[{}]", session.session_id());

            let internal = PlayerInternal {
                session: session,
                config: config,
                commands: cmd_rx,

                state: PlayerState::Stopped,
                sink: sink_builder(),
                audio_filter: audio_filter,
            };

            internal.run();
        });

        Player {
            commands: Some(cmd_tx),
            thread_handle: Some(handle),
        }
    }

    fn command(&self, cmd: PlayerCommand) {
        self.commands.as_ref().unwrap().send(cmd).unwrap();
    }

    pub fn load(&self, track: SpotifyId, start_playing: bool, position_ms: u32)
        -> oneshot::Receiver<()>
    {
        let (tx, rx) = oneshot::channel();
        self.command(PlayerCommand::Load(track, start_playing, position_ms, tx));

        rx
    }

    pub fn play(&self) {
        self.command(PlayerCommand::Play)
    }

    pub fn pause(&self) {
        self.command(PlayerCommand::Pause)
    }

    pub fn stop(&self) {
        self.command(PlayerCommand::Stop)
    }

    pub fn seek(&self, position_ms: u32) {
        self.command(PlayerCommand::Seek(position_ms));
    }
}

impl Drop for Player {
    fn drop(&mut self) {
        debug!("Shutting down player thread ...");
        self.commands = None;
        if let Some(handle) = self.thread_handle.take() {
            match handle.join() {
                Ok(_) => (),
                Err(_) => error!("Player thread panicked!")
            }
        }
    }
}

type Decoder = VorbisDecoder<Subfile<AudioDecrypt<AudioFile>>>;
enum PlayerState {
    Stopped,
    Paused {
        decoder: Decoder,
        end_of_track: oneshot::Sender<()>,
    },
    Playing {
        decoder: Decoder,
        end_of_track: oneshot::Sender<()>,
    },

    Invalid,
}

impl PlayerState {
    fn is_playing(&self) -> bool {
        use self::PlayerState::*;
        match *self {
            Stopped | Paused { .. } => false,
            Playing { .. } => true,
            Invalid => panic!("invalid state"),
        }
    }

    fn decoder(&mut self) -> Option<&mut Decoder> {
        use self::PlayerState::*;
        match *self {
            Stopped => None,
            Paused { ref mut decoder, .. } |
            Playing { ref mut decoder, .. } => Some(decoder),
            Invalid => panic!("invalid state"),
        }
    }

    fn signal_end_of_track(self) {
        use self::PlayerState::*;
        match self {
            Paused { end_of_track, .. } |
            Playing { end_of_track, .. } => {
                end_of_track.complete(())
            }

            Stopped => warn!("signal_end_of_track from stopped state"),
            Invalid => panic!("invalid state"),
        }
    }

    fn paused_to_playing(&mut self) {
        use self::PlayerState::*;
        match ::std::mem::replace(self, Invalid) {
            Paused { decoder, end_of_track } => {
                *self = Playing {
                    decoder: decoder,
                    end_of_track: end_of_track,
                };
            }
            _ => panic!("invalid state"),
        }
    }

    fn playing_to_paused(&mut self) {
        use self::PlayerState::*;
        match ::std::mem::replace(self, Invalid) {
            Playing { decoder, end_of_track } => {
                *self = Paused {
                    decoder: decoder,
                    end_of_track: end_of_track,
                };
            }
            _ => panic!("invalid state"),
        }
    }
}

impl PlayerInternal {
    fn run(mut self) {
        loop {
            let cmd = if self.state.is_playing() {
                match self.commands.try_recv() {
                    Ok(cmd) => Some(cmd),
                    Err(TryRecvError::Empty) => None,
                    Err(TryRecvError::Disconnected) => return,
                }
            } else {
                match self.commands.recv() {
                    Ok(cmd) => Some(cmd),
                    Err(RecvError) => return,
                }
            };

            if let Some(cmd) = cmd {
                self.handle_command(cmd);
            }

            let packet = if let PlayerState::Playing { ref mut decoder, .. } = self.state {
                Some(decoder.next_packet().expect("Vorbis error"))
            } else { None };

            if let Some(packet) = packet {
                self.handle_packet(packet);
            }
        }
    }

    fn handle_packet(&mut self, packet: Option<VorbisPacket>) {
        match packet {
            Some(mut packet) => {
                if let Some(ref editor) = self.audio_filter {
                    editor.modify_stream(&mut packet.data_mut())
                };

                self.sink.write(&packet.data()).unwrap();
            }

            None => {
                self.sink.stop().unwrap();
                self.run_onstop();

                let old_state = mem::replace(&mut self.state, PlayerState::Stopped);
                old_state.signal_end_of_track();
            }
        }
    }

    fn handle_command(&mut self, cmd: PlayerCommand) {
        debug!("command={:?}", cmd);
        match cmd {
            PlayerCommand::Load(track_id, play, position, end_of_track) => {
                if self.state.is_playing() {
                    self.sink.stop().unwrap();
                }

                match self.load_track(track_id, position as i64) {
                    Some(decoder) => {
                        if play {
                            if !self.state.is_playing() {
                                self.run_onstart();
                            }
                            self.sink.start().unwrap();

                            self.state = PlayerState::Playing {
                                decoder: decoder,
                                end_of_track: end_of_track,
                            };
                        } else {
                            if self.state.is_playing() {
                                self.run_onstop();
                            }

                            self.state = PlayerState::Paused {
                                decoder: decoder,
                                end_of_track: end_of_track,
                            };
                        }
                    }

                    None => {
                        end_of_track.complete(());
                        if self.state.is_playing() {
                            self.run_onstop();
                        }
                    }
                }
            }

            PlayerCommand::Seek(position) => {
                if let Some(decoder) = self.state.decoder() {
                    match decoder.seek(position as i64) {
                        Ok(_) => (),
                        Err(err) => error!("Vorbis error: {:?}", err),
                    }
                } else {
                    warn!("Player::seek called from invalid state");
                }
            }

            PlayerCommand::Play => {
                if let PlayerState::Paused { .. } = self.state {
                    self.state.paused_to_playing();

                    self.run_onstart();
                    self.sink.start().unwrap();
                } else {
                    warn!("Player::play called from invalid state");
                }
            }

            PlayerCommand::Pause => {
                if let PlayerState::Playing { .. } = self.state {
                    self.state.playing_to_paused();

                    self.sink.stop().unwrap();
                    self.run_onstop();
                } else {
                    warn!("Player::pause called from invalid state");
                }
            }

            PlayerCommand::Stop => {
                match self.state {
                    PlayerState::Playing { .. } => {
                        self.sink.stop().unwrap();
                        self.run_onstop();
                        self.state = PlayerState::Stopped;
                    }
                    PlayerState::Paused { .. } => {
                        self.state = PlayerState::Stopped;
                    },
                    PlayerState::Stopped => {
                        warn!("Player::stop called from invalid state");
                    }
                    PlayerState::Invalid => panic!("invalid state"),
                }
            }
        }
    }

    fn run_onstart(&self) {
        if let Some(ref program) = self.config.onstart {
            util::run_program(program)
        }
    }

    fn run_onstop(&self) {
        if let Some(ref program) = self.config.onstop {
            util::run_program(program)
        }
    }

    fn find_available_alternative<'a>(&self, track: &'a Track) -> Option<Cow<'a, Track>> {
        if track.available {
            Some(Cow::Borrowed(track))
        } else {
            let alternatives = track.alternatives
                .iter()
                .map(|alt_id| {
                    Track::get(&self.session, *alt_id)
                });
            let alternatives = future::join_all(alternatives).wait().unwrap();

            alternatives.into_iter().find(|alt| alt.available).map(Cow::Owned)
        }
    }

    fn load_track(&self, track_id: SpotifyId, position: i64) -> Option<Decoder> {
        let track = Track::get(&self.session, track_id).wait().unwrap();

        info!("Loading track \"{}\"", track.name);

        let track = match self.find_available_alternative(&track) {
            Some(track) => track,
            None => {
                warn!("Track \"{}\" is not available", track.name);
                return None;
            }
        };

        let format = match self.config.bitrate {
            Bitrate::Bitrate96 => FileFormat::OGG_VORBIS_96,
            Bitrate::Bitrate160 => FileFormat::OGG_VORBIS_160,
            Bitrate::Bitrate320 => FileFormat::OGG_VORBIS_320,
        };

        let file_id = match track.files.get(&format) {
            Some(&file_id) => file_id,
            None => {
                warn!("Track \"{}\" is not available in format {:?}", track.name, format);
                return None;
            }
        };

        let key = self.session.audio_key().request(track.id, file_id).wait().unwrap();

        let encrypted_file = AudioFile::open(&self.session, file_id).wait().unwrap();
        let audio_file = Subfile::new(AudioDecrypt::new(key, encrypted_file), 0xa7);

        let mut decoder = VorbisDecoder::new(audio_file).unwrap();

        match decoder.seek(position) {
            Ok(_) => (),
            Err(err) => error!("Vorbis error: {:?}", err),
        }

        info!("Track \"{}\" loaded", track.name);

        Some(decoder)
    }
}

impl Drop for PlayerInternal {
    fn drop(&mut self) {
        debug!("drop Player[{}]", self.session.session_id());
    }
}

impl ::std::fmt::Debug for PlayerCommand {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match *self {
            PlayerCommand::Load(track, play, position, _) => {
                f.debug_tuple("Load")
                 .field(&track)
                 .field(&play)
                 .field(&position)
                 .finish()
            }
            PlayerCommand::Play => {
                f.debug_tuple("Play").finish()
            }
            PlayerCommand::Pause => {
                f.debug_tuple("Pause").finish()
            }
            PlayerCommand::Stop => {
                f.debug_tuple("Stop").finish()
            }
            PlayerCommand::Seek(position) => {
                f.debug_tuple("Seek")
                 .field(&position)
                 .finish()
            }
        }
    }
}
<|endoftext|>


<commit_before>use nom::{alphanumeric, digit, multispace};
use nom::{IResult, Err, ErrorKind, Needed};
use std::collections::{HashSet, VecDeque};
use std::str;

use column::Column;
use common::{binary_comparison_operator, column_identifier, Operator};

#[derive(Clone, Debug, Hash, PartialEq)]
pub enum ConditionBase {
    Field(Column),
    Literal(String),
    Placeholder,
}

#[derive(Clone, Debug, Hash, PartialEq)]
pub struct ConditionTree {
    pub operator: Operator,
    pub left: Box<ConditionExpression>,
    pub right: Box<ConditionExpression>,
}

impl<'a> ConditionTree {
    pub fn contained_columns(&'a self) -> HashSet<&'a Column> {
        let mut s = HashSet::new();
        let mut q = VecDeque::<&'a ConditionTree>::new();
        q.push_back(self);
        while let Some(ref ct) = q.pop_front() {
            match *ct.left.as_ref() {
                ConditionExpression::Base(ConditionBase::Field(ref c)) => {
                    s.insert(c);
                }
                ConditionExpression::LogicalOp(ref ct) |
                ConditionExpression::ComparisonOp(ref ct) => q.push_back(ct),
                _ => (),
            }
            match *ct.right.as_ref() {
                ConditionExpression::Base(ConditionBase::Field(ref c)) => {
                    s.insert(c);
                }
                ConditionExpression::LogicalOp(ref ct) |
                ConditionExpression::ComparisonOp(ref ct) => q.push_back(ct),
                _ => (),
            }
        }
        s
    }
}

#[derive(Clone, Debug, Hash, PartialEq)]
pub enum ConditionExpression {
    ComparisonOp(ConditionTree),
    LogicalOp(ConditionTree),
    NegationOp(Box<ConditionExpression>),
    Base(ConditionBase),
}

/// Parse a conditional expression into a condition tree structure
named!(pub condition_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               left: and_expr ~
               caseless_tag!("or") ~
               multispace ~
               right: condition_expr,
               || {
                   ConditionExpression::LogicalOp(
                       ConditionTree {
                           operator: Operator::Or,
                           left: Box::new(left),
                           right: Box::new(right),
                       }
                   )
               }
           )
       |   and_expr)
);

named!(pub and_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               left: parenthetical_expr ~
               caseless_tag!("and") ~
               multispace ~
               right: and_expr,
               || {
                   ConditionExpression::LogicalOp(
                       ConditionTree {
                           operator: Operator::And,
                           left: Box::new(left),
                           right: Box::new(right),
                       }
                   )
               }
           )
       |   parenthetical_expr)
);

named!(pub parenthetical_expr<&[u8], ConditionExpression>,
       alt_complete!(
           delimited!(tag!("("), condition_expr, chain!(tag!(")") ~ multispace?, ||{}))
       |   not_expr)
);

named!(pub not_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               caseless_tag!("not") ~
               multispace ~
               right: parenthetical_expr,
               || {
                   ConditionExpression::NegationOp(Box::new(right))
               }
           )
       |   boolean_primary)
);

named!(boolean_primary<&[u8], ConditionExpression>,
    chain!(
        left: predicate ~
        multispace? ~
        op: binary_comparison_operator ~
        multispace? ~
        right: predicate,
        || {
            ConditionExpression::ComparisonOp(
                ConditionTree {
                    operator: op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            )

        }
    )
);

named!(predicate<&[u8], ConditionExpression>,
    alt_complete!(
            chain!(
                delimited!(opt!(multispace), tag!("?"), opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Placeholder
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace), digit, opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Literal(String::from(str::from_utf8(field).unwrap()))
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace),
                                  alt_complete!(
                                        delimited!(tag!("\""), opt!(alphanumeric), tag!("\""))
                                      | delimited!(tag!("'"), opt!(alphanumeric), tag!("'"))
                                  ),
                                  opt!(multispace)),
                || {
                    let field = field.unwrap_or("".as_bytes());
                    ConditionExpression::Base(
                        ConditionBase::Literal(String::from(str::from_utf8(field).unwrap()))
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace), column_identifier, opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Field(field)
                    )
                }
            )
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use column::Column;
    use common::Operator;

    fn flat_condition_tree(op: Operator,
                           l: ConditionBase,
                           r: ConditionBase)
                           -> ConditionExpression {
        ConditionExpression::ComparisonOp(ConditionTree {
            operator: op,
            left: Box::new(ConditionExpression::Base(l)),
            right: Box::new(ConditionExpression::Base(r)),
        })
    }

    #[test]
    fn ct_contained_columns() {
        use std::collections::HashSet;

        let cond = "a.foo = ? and b.bar = 42";

        let res = condition_expr(cond.as_bytes());
        let c1 = Column::from("a.foo");
        let c2 = Column::from("b.bar");
        let mut expected_cols = HashSet::new();
        expected_cols.insert(&c1);
        expected_cols.insert(&c2);
        match res.unwrap().1 {
            ConditionExpression::LogicalOp(ct) => {
                assert_eq!(ct.contained_columns(), expected_cols);
            }
            _ => panic!(),
        }
    }

    #[test]
    fn equality_placeholder() {
        let cond = "foo = ?";

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Placeholder));
    }

    #[test]
    fn equality_literals() {
        let cond1 = "foo = 42";
        let cond2 = "foo = \"hello\"";

        let res1 = condition_expr(cond1.as_bytes());
        assert_eq!(res1.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("42"))));

        let res2 = condition_expr(cond2.as_bytes());
        assert_eq!(res2.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("hello"))));
    }

    #[test]
    fn inequality_literals() {
        let cond1 = "foo >= 42";
        let cond2 = "foo <= 5";

        let res1 = condition_expr(cond1.as_bytes());
        assert_eq!(res1.unwrap().1,
                   flat_condition_tree(Operator::GreaterOrEqual,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("42"))));

        let res2 = condition_expr(cond2.as_bytes());
        assert_eq!(res2.unwrap().1,
                   flat_condition_tree(Operator::LessOrEqual,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("5"))));
    }

    #[test]
    fn empty_string_literal() {
        let cond = "foo = ''";

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from(""))));
    }

    #[test]
    fn parenthesis() {
        let cond = "(foo = ? or bar = 12) and foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let a = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foo".into()))),
            right: Box::new(Base(Placeholder)),
        });

        let b = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        });

        let left = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(a),
            right: Box::new(b),
        });

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::And,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

    #[test]
    fn order_of_operations() {
        let cond = "foo = ? and bar = 12 or foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let a = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foo".into()))),
            right: Box::new(Base(Placeholder)),
        });

        let b = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        });

        let left = LogicalOp(ConditionTree {
            operator: Operator::And,
            left: Box::new(a),
            right: Box::new(b),
        });

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

    #[test]
    fn negation() {
        let cond = "not bar = 12 or foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let left = NegationOp(Box::new(ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        })));

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

}
<commit_msg>String literals can contain any character<commit_after>use nom::{digit, multispace};
use nom::{IResult, Err, ErrorKind, Needed};
use std::collections::{HashSet, VecDeque};
use std::str;

use column::Column;
use common::{binary_comparison_operator, column_identifier, Operator};

#[derive(Clone, Debug, Hash, PartialEq)]
pub enum ConditionBase {
    Field(Column),
    Literal(String),
    Placeholder,
}

#[derive(Clone, Debug, Hash, PartialEq)]
pub struct ConditionTree {
    pub operator: Operator,
    pub left: Box<ConditionExpression>,
    pub right: Box<ConditionExpression>,
}

impl<'a> ConditionTree {
    pub fn contained_columns(&'a self) -> HashSet<&'a Column> {
        let mut s = HashSet::new();
        let mut q = VecDeque::<&'a ConditionTree>::new();
        q.push_back(self);
        while let Some(ref ct) = q.pop_front() {
            match *ct.left.as_ref() {
                ConditionExpression::Base(ConditionBase::Field(ref c)) => {
                    s.insert(c);
                }
                ConditionExpression::LogicalOp(ref ct) |
                ConditionExpression::ComparisonOp(ref ct) => q.push_back(ct),
                _ => (),
            }
            match *ct.right.as_ref() {
                ConditionExpression::Base(ConditionBase::Field(ref c)) => {
                    s.insert(c);
                }
                ConditionExpression::LogicalOp(ref ct) |
                ConditionExpression::ComparisonOp(ref ct) => q.push_back(ct),
                _ => (),
            }
        }
        s
    }
}

#[derive(Clone, Debug, Hash, PartialEq)]
pub enum ConditionExpression {
    ComparisonOp(ConditionTree),
    LogicalOp(ConditionTree),
    NegationOp(Box<ConditionExpression>),
    Base(ConditionBase),
}

/// Parse a conditional expression into a condition tree structure
named!(pub condition_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               left: and_expr ~
               caseless_tag!("or") ~
               multispace ~
               right: condition_expr,
               || {
                   ConditionExpression::LogicalOp(
                       ConditionTree {
                           operator: Operator::Or,
                           left: Box::new(left),
                           right: Box::new(right),
                       }
                   )
               }
           )
       |   and_expr)
);

named!(pub and_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               left: parenthetical_expr ~
               caseless_tag!("and") ~
               multispace ~
               right: and_expr,
               || {
                   ConditionExpression::LogicalOp(
                       ConditionTree {
                           operator: Operator::And,
                           left: Box::new(left),
                           right: Box::new(right),
                       }
                   )
               }
           )
       |   parenthetical_expr)
);

named!(pub parenthetical_expr<&[u8], ConditionExpression>,
       alt_complete!(
           delimited!(tag!("("), condition_expr, chain!(tag!(")") ~ multispace?, ||{}))
       |   not_expr)
);

named!(pub not_expr<&[u8], ConditionExpression>,
       alt_complete!(
           chain!(
               caseless_tag!("not") ~
               multispace ~
               right: parenthetical_expr,
               || {
                   ConditionExpression::NegationOp(Box::new(right))
               }
           )
       |   boolean_primary)
);

named!(boolean_primary<&[u8], ConditionExpression>,
    chain!(
        left: predicate ~
        multispace? ~
        op: binary_comparison_operator ~
        multispace? ~
        right: predicate,
        || {
            ConditionExpression::ComparisonOp(
                ConditionTree {
                    operator: op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            )

        }
    )
);

named!(predicate<&[u8], ConditionExpression>,
    alt_complete!(
            chain!(
                delimited!(opt!(multispace), tag!("?"), opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Placeholder
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace), digit, opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Literal(String::from(str::from_utf8(field).unwrap()))
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace),
                                  alt_complete!(
                                        delimited!(tag!("\""), opt!(take_until!("\"")), tag!("\""))
                                      | delimited!(tag!("'"), opt!(take_until!("'")), tag!("'"))
                                  ),
                                  opt!(multispace)),
                || {
                    let field = field.unwrap_or("".as_bytes());
                    ConditionExpression::Base(
                        ConditionBase::Literal(String::from(str::from_utf8(field).unwrap()))
                    )
                }
            )
        |   chain!(
                field: delimited!(opt!(multispace), column_identifier, opt!(multispace)),
                || {
                    ConditionExpression::Base(
                        ConditionBase::Field(field)
                    )
                }
            )
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use column::Column;
    use common::Operator;

    fn flat_condition_tree(op: Operator,
                           l: ConditionBase,
                           r: ConditionBase)
                           -> ConditionExpression {
        ConditionExpression::ComparisonOp(ConditionTree {
            operator: op,
            left: Box::new(ConditionExpression::Base(l)),
            right: Box::new(ConditionExpression::Base(r)),
        })
    }

    #[test]
    fn ct_contained_columns() {
        use std::collections::HashSet;

        let cond = "a.foo = ? and b.bar = 42";

        let res = condition_expr(cond.as_bytes());
        let c1 = Column::from("a.foo");
        let c2 = Column::from("b.bar");
        let mut expected_cols = HashSet::new();
        expected_cols.insert(&c1);
        expected_cols.insert(&c2);
        match res.unwrap().1 {
            ConditionExpression::LogicalOp(ct) => {
                assert_eq!(ct.contained_columns(), expected_cols);
            }
            _ => panic!(),
        }
    }

    #[test]
    fn equality_placeholder() {
        let cond = "foo = ?";

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Placeholder));
    }

    #[test]
    fn equality_literals() {
        let cond1 = "foo = 42";
        let cond2 = "foo = \"hello\"";

        let res1 = condition_expr(cond1.as_bytes());
        assert_eq!(res1.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("42"))));

        let res2 = condition_expr(cond2.as_bytes());
        assert_eq!(res2.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("hello"))));
    }

    #[test]
    fn inequality_literals() {
        let cond1 = "foo >= 42";
        let cond2 = "foo <= 5";

        let res1 = condition_expr(cond1.as_bytes());
        assert_eq!(res1.unwrap().1,
                   flat_condition_tree(Operator::GreaterOrEqual,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("42"))));

        let res2 = condition_expr(cond2.as_bytes());
        assert_eq!(res2.unwrap().1,
                   flat_condition_tree(Operator::LessOrEqual,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from("5"))));
    }

    #[test]
    fn empty_string_literal() {
        let cond = "foo = ''";

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1,
                   flat_condition_tree(Operator::Equal,
                                       ConditionBase::Field(Column::from("foo")),
                                       ConditionBase::Literal(String::from(""))));
    }

    #[test]
    fn parenthesis() {
        let cond = "(foo = ? or bar = 12) and foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let a = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foo".into()))),
            right: Box::new(Base(Placeholder)),
        });

        let b = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        });

        let left = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(a),
            right: Box::new(b),
        });

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::And,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

    #[test]
    fn order_of_operations() {
        let cond = "foo = ? and bar = 12 or foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let a = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foo".into()))),
            right: Box::new(Base(Placeholder)),
        });

        let b = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        });

        let left = LogicalOp(ConditionTree {
            operator: Operator::And,
            left: Box::new(a),
            right: Box::new(b),
        });

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

    #[test]
    fn negation() {
        let cond = "not bar = 12 or foobar = 'a'";

        use ConditionExpression::*;
        use ConditionBase::*;

        let left = NegationOp(Box::new(ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("bar".into()))),
            right: Box::new(Base(Literal("12".into()))),
        })));

        let right = ComparisonOp(ConditionTree {
            operator: Operator::Equal,
            left: Box::new(Base(Field("foobar".into()))),
            right: Box::new(Base(Literal("a".into()))),
        });

        let complete = LogicalOp(ConditionTree {
            operator: Operator::Or,
            left: Box::new(left),
            right: Box::new(right),
        });

        let res = condition_expr(cond.as_bytes());
        assert_eq!(res.unwrap().1, complete);
    }

}
<|endoftext|>


<commit_before>//! Virtual registers.
//!
//! A virtual register is a set of related SSA values whose live ranges don't interfere. If all the
//! values in a virtual register are assigned to the same location, fewer copies will result in the
//! output.
//!
//! A virtual register is typically built by merging together SSA values that are "phi-related" -
//! that is, one value is passed as an EBB argument to a branch and the other is the EBB parameter
//! value itself.
//!
//! If any values in a virtual register are spilled, they will use the same stack slot. This avoids
//! memory-to-memory copies when a spilled value is passed as an EBB argument.

use entity::{EntityList, ListPool};
use entity::{PrimaryMap, EntityMap, Keys};
use ir::Value;
use packed_option::PackedOption;
use ref_slice::ref_slice;

/// A virtual register reference.
#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct VirtReg(u32);
entity_impl!(VirtReg, "vreg");

type ValueList = EntityList<Value>;

/// Collection of virtual registers.
///
/// Each virtual register is a list of values. Also maintain a map from values to their unique
/// virtual register, if any.
pub struct VirtRegs {
    /// Memory pool for the value lists.
    pool: ListPool<Value>,

    /// The primary table of virtual registers.
    vregs: PrimaryMap<VirtReg, ValueList>,

    /// Each value belongs to at most one virtual register.
    value_vregs: EntityMap<Value, PackedOption<VirtReg>>,
}

#[allow(dead_code)]
impl VirtRegs {
    /// Create a new virtual register collection.
    pub fn new() -> Self {
        Self {
            pool: ListPool::new(),
            vregs: PrimaryMap::new(),
            value_vregs: EntityMap::new(),
        }
    }

    /// Clear all virtual registers.
    pub fn clear(&mut self) {
        self.vregs.clear();
        self.value_vregs.clear();
        self.pool.clear();
    }

    /// Get the virtual register containing `value`, if any.
    pub fn get(&self, value: Value) -> Option<VirtReg> {
        self.value_vregs[value].into()
    }

    /// Get the list of values in `vreg`.
    /// The values are topologically ordered according dominance of their definition points.
    pub fn values(&self, vreg: VirtReg) -> &[Value] {
        self.vregs[vreg].as_slice(&self.pool)
    }

    /// Get an iterator over all virtual registers.
    pub fn all_virtregs(&self) -> Keys<VirtReg> {
        self.vregs.keys()
    }

    /// Get the congruence class of `value`.
    ///
    /// If `value` belongs to a virtual register, the congruence class is the values of the virtual
    /// register. Otherwise it is just the value itself.
    pub fn congruence_class<'a, 'b>(&'a self, value: &'b Value) -> &'b [Value]
    where
        'a: 'b,
    {
        self.get(*value).map(|vr| self.values(vr)).unwrap_or(
            ref_slice(value),
        )
    }

    /// Check if `a` and `b` belong to the same congruence class.
    pub fn same_class(&self, a: Value, b: Value) -> bool {
        match (self.get(a), self.get(b)) {
            (Some(va), Some(vb)) => va == vb,
            _ => a == b,
        }
    }

    /// Unify `values` into a single virtual register.
    ///
    /// The values in the slice can be singletons or they can belong to a virtual register already.
    /// If a value belongs to a virtual register, all of the values in that register must be
    /// present.
    ///
    /// The values are assumed to already be in topological order.
    pub fn unify(&mut self, values: &[Value]) -> VirtReg {
        // Start by clearing all virtual registers involved.
        // Pick a virtual register to reuse (the smallest number) or allocate a new one.
        let mut singletons = 0;
        let mut cleared = 0;
        let vreg = values
            .iter()
            .filter_map(|&v| {
                let vr = self.get(v);
                match vr {
                    None => singletons += 1,
                    Some(vr) => {
                        if !self.vregs[vr].is_empty() {
                            cleared += self.vregs[vr].len(&self.pool);
                            self.vregs[vr].clear(&mut self.pool);
                        }
                    }
                }
                vr
            })
            .min()
            .unwrap_or_else(|| self.vregs.push(Default::default()));

        assert_eq!(
            values.len(),
            singletons + cleared,
            "Can't unify partial virtual registers"
        );

        self.vregs[vreg].extend(values.iter().cloned(), &mut self.pool);
        for &v in values {
            self.value_vregs[v] = vreg.into();
        }

        vreg
    }
}
<commit_msg>Add support for textbook union-find to VirtRegs.<commit_after>//! Virtual registers.
//!
//! A virtual register is a set of related SSA values whose live ranges don't interfere. If all the
//! values in a virtual register are assigned to the same location, fewer copies will result in the
//! output.
//!
//! A virtual register is typically built by merging together SSA values that are "phi-related" -
//! that is, one value is passed as an EBB argument to a branch and the other is the EBB parameter
//! value itself.
//!
//! If any values in a virtual register are spilled, they will use the same stack slot. This avoids
//! memory-to-memory copies when a spilled value is passed as an EBB argument.

use dbg::DisplayList;
use entity::{EntityList, ListPool};
use entity::{PrimaryMap, EntityMap, Keys};
use entity::EntityRef;
use ir::Value;
use packed_option::PackedOption;
use ref_slice::ref_slice;
use std::cmp::Ordering;
use std::fmt;

/// A virtual register reference.
#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct VirtReg(u32);
entity_impl!(VirtReg, "vreg");

type ValueList = EntityList<Value>;

/// Collection of virtual registers.
///
/// Each virtual register is a list of values. Also maintain a map from values to their unique
/// virtual register, if any.
pub struct VirtRegs {
    /// Memory pool for the value lists.
    pool: ListPool<Value>,

    /// The primary table of virtual registers.
    vregs: PrimaryMap<VirtReg, ValueList>,

    /// Allocated virtual register numbers that are no longer in use.
    unused_vregs: Vec<VirtReg>,

    /// Each value belongs to at most one virtual register.
    value_vregs: EntityMap<Value, PackedOption<VirtReg>>,

    /// Table used during the union-find phase while `vregs` is empty.
    union_find: EntityMap<Value, i32>,

    /// Values that have been activated in the `union_find` table, but not yet added to any virtual
    /// registers by the `finish_union_find()` function.
    pending_values: Vec<Value>,
}

impl VirtRegs {
    /// Create a new virtual register collection.
    pub fn new() -> Self {
        Self {
            pool: ListPool::new(),
            vregs: PrimaryMap::new(),
            unused_vregs: Vec::new(),
            value_vregs: EntityMap::new(),
            union_find: EntityMap::new(),
            pending_values: Vec::new(),
        }
    }

    /// Clear all virtual registers.
    pub fn clear(&mut self) {
        self.vregs.clear();
        self.unused_vregs.clear();
        self.value_vregs.clear();
        self.pool.clear();
        self.union_find.clear();
        self.pending_values.clear();
    }

    /// Get the virtual register containing `value`, if any.
    pub fn get(&self, value: Value) -> Option<VirtReg> {
        self.value_vregs[value].into()
    }

    /// Get the list of values in `vreg`.
    pub fn values(&self, vreg: VirtReg) -> &[Value] {
        self.vregs[vreg].as_slice(&self.pool)
    }

    /// Get an iterator over all virtual registers.
    pub fn all_virtregs(&self) -> Keys<VirtReg> {
        self.vregs.keys()
    }

    /// Get the congruence class of `value`.
    ///
    /// If `value` belongs to a virtual register, the congruence class is the values of the virtual
    /// register. Otherwise it is just the value itself.
    pub fn congruence_class<'a, 'b>(&'a self, value: &'b Value) -> &'b [Value]
    where
        'a: 'b,
    {
        self.get(*value).map(|vr| self.values(vr)).unwrap_or(
            ref_slice(value),
        )
    }

    /// Check if `a` and `b` belong to the same congruence class.
    pub fn same_class(&self, a: Value, b: Value) -> bool {
        match (self.get(a), self.get(b)) {
            (Some(va), Some(vb)) => va == vb,
            _ => a == b,
        }
    }

    /// Sort the values in `vreg` according to `compare`.
    ///
    /// If the ordering defined by `compare` is not total, value numbers are used as a last resort
    /// tie-breaker. This makes it possible to use an unstable sorting algorithm which can be
    /// faster because it doesn't allocate memory.
    ///
    /// Returns the slice of sorted values which `values(vreg)` will also return from now on.
    pub fn sort_values<F>(&mut self, vreg: VirtReg, mut compare: F) -> &[Value]
    where
        F: FnMut(Value, Value) -> Ordering,
    {
        let s = self.vregs[vreg].as_mut_slice(&mut self.pool);
        s.sort_unstable_by(|&a, &b| compare(a, b).then(a.cmp(&b)));
        s
    }

    /// Remove a virtual register.
    ///
    /// The values in `vreg` become singletons, and the virtual register number may be reused in
    /// the future.
    pub fn remove(&mut self, vreg: VirtReg) {
        // Start by reassigning all the values.
        for &v in self.vregs[vreg].as_slice(&self.pool) {
            let old = self.value_vregs[v].take();
            debug_assert_eq!(old, Some(vreg));
        }

        self.vregs[vreg].clear(&mut self.pool);
        self.unused_vregs.push(vreg);
    }

    /// Allocate a new empty virtual register.
    fn alloc(&mut self) -> VirtReg {
        self.unused_vregs.pop().unwrap_or_else(|| {
            self.vregs.push(Default::default())
        })
    }

    /// Unify `values` into a single virtual register.
    ///
    /// The values in the slice can be singletons or they can belong to a virtual register already.
    /// If a value belongs to a virtual register, all of the values in that register must be
    /// present.
    ///
    /// The values are assumed to already be in topological order.
    pub fn unify(&mut self, values: &[Value]) -> VirtReg {
        // Start by clearing all virtual registers involved.
        // Pick a virtual register to reuse (the smallest number) or allocate a new one.
        let mut singletons = 0;
        let mut cleared = 0;
        let vreg = values
            .iter()
            .filter_map(|&v| {
                let vr = self.get(v);
                match vr {
                    None => singletons += 1,
                    Some(vr) => {
                        if !self.vregs[vr].is_empty() {
                            cleared += self.vregs[vr].len(&self.pool);
                            self.vregs[vr].clear(&mut self.pool);
                        }
                    }
                }
                vr
            })
            .min()
            .unwrap_or_else(|| self.vregs.push(Default::default()));

        assert_eq!(
            values.len(),
            singletons + cleared,
            "Can't unify partial virtual registers"
        );

        self.vregs[vreg].extend(values.iter().cloned(), &mut self.pool);
        for &v in values {
            self.value_vregs[v] = vreg.into();
        }

        vreg
    }
}

impl fmt::Display for VirtRegs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for vreg in self.all_virtregs() {
            write!(f, "\n{} = {}", vreg, DisplayList(self.values(vreg)))?;
        }
        Ok(())
    }
}

/// Expanded version of a union-find table entry.
enum UFEntry {
    /// This value is a a set leader. The embedded number is the set's rank.
    Rank(u32),

    /// This value belongs to the same set as the linked value.
    Link(Value),
}

/// The `union_find` table contains `i32` entries that are interpreted as follows:
///
/// x = 0: The value belongs to its own singleton set.
/// x > 0: The value is the leader of a set with rank x.
/// x < 0: The value belongs to the same set as the value numbered !x.
///
/// The rank of a set is an upper bound on the number of links that must be followed from a member
/// of the set to the set leader.
///
/// A singleton set is the same as a set with rank 0. It contains only the leader value.
impl UFEntry {
    /// Decode a table entry.
    fn decode(x: i32) -> UFEntry {
        if x < 0 {
            UFEntry::Link(Value::new((!x) as usize))
        } else {
            UFEntry::Rank(x as u32)
        }
    }

    /// Encode a link entry.
    fn encode_link(v: Value) -> i32 {
        !(v.index() as i32)
    }
}

/// Union-find algorithm for building virtual registers.
///
/// Before values are added to virtual registers, it is possible to use a union-find algorithm to
/// construct virtual registers efficiently. This support implemented here is used as follows:
///
/// 1. Repeatedly call the `union(a, b)` method to request that `a` and `b` are placed in the same
///    virtual register.
/// 2. When done, call `finish_union_find()` to construct the virtual register sets based on the
///    `union()` calls.
///
/// The values that were passed to `union(a, b)` mist not belong to any existing virtual registers
/// by the time `finish_union_find()` is called.
///
/// For more information on the algorithm implemented here, see Chapter 21 "Data Structures for
/// Disjoint Sets" of Cormen, Leiserson, Rivest, Stein, "Introduction to algorithms", 3rd Ed.
///
/// The [Wikipedia entry on disjoint-set data
/// structures](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) is also good.
impl VirtRegs {
    /// Find the leader value and rank of the set containing `v`.
    /// Compress the path if needed.
    fn find(&mut self, val: Value) -> (Value, u32) {
        match UFEntry::decode(self.union_find[val]) {
            UFEntry::Rank(rank) => (val, rank),
            UFEntry::Link(parent) => {
                // TODO: This recursion would be more efficient as an iteration that pushes
                // elements onto a SmallVector.
                let found = self.find(parent);
                // Compress the path if needed.
                if found.0 != parent {
                    self.union_find[val] = UFEntry::encode_link(found.0);
                }
                found
            }
        }
    }

    /// Union the two sets containing `a` and `b`.
    ///
    /// This ensures that `a` and `b` will belong to the same virtual register after calling
    /// `finish_union_find()`.
    pub fn union(&mut self, a: Value, b: Value) {
        let (leader_a, rank_a) = self.find(a);
        let (leader_b, rank_b) = self.find(b);

        if leader_a == leader_b {
            return;
        }

        // The first time we see a value, its rank will be 0. Add it to the list of pending values.
        if rank_a == 0 {
            debug_assert_eq!(a, leader_a);
            self.pending_values.push(a);
        }
        if rank_b == 0 {
            debug_assert_eq!(b, leader_b);
            self.pending_values.push(b);
        }

        // Merge into the set with the greater rank. This preserves the invariant that the rank is
        // an upper bound on the number of links to the leader.
        match rank_a.cmp(&rank_b) {
            Ordering::Less => {
                self.union_find[leader_a] = UFEntry::encode_link(leader_b);
            }
            Ordering::Greater => {
                self.union_find[leader_b] = UFEntry::encode_link(leader_a);
            }
            Ordering::Equal => {
                // When the two sets have the same rank, we arbitrarily pick the a-set to preserve.
                // We need to increase the rank by one since the elements in the b-set are now one
                // link further away from the leader.
                self.union_find[leader_a] += 1;
                self.union_find[leader_b] = UFEntry::encode_link(leader_a);
            }
        }
    }

    /// Compute virtual registers based on previous calls to `union(a, b)`.
    ///
    /// This terminates the union-find algorithm, so the next time `union()` is called, it is for a
    /// new independent batch of values.
    ///
    /// The values in each virtual register will be ordered according to when they were first
    /// passed to `union()`, but backwards. It is expected that `sort_values()` will be used to
    /// create a more sensible value order.
    ///
    /// The new virtual registers will be appended to `new_vregs`, if present.
    pub fn finish_union_find(&mut self, mut new_vregs: Option<&mut Vec<VirtReg>>) {
        debug_assert_eq!(
            self.pending_values.iter().find(|&&v| self.get(v).is_some()),
            None,
            "Values participating in union-find must not belong to existing virtual registers"
        );

        while let Some(val) = self.pending_values.pop() {
            let (leader, _) = self.find(val);

            // Get the vreg for `leader`, or create it.
            let vreg = self.get(leader).unwrap_or_else(|| {
                // Allocate a vreg for `leader`, but leave it empty.
                let vr = self.alloc();
                if let &mut Some(ref mut vec) = &mut new_vregs {
                    vec.push(vr);
                }
                self.value_vregs[leader] = vr.into();
                vr
            });

            // Push values in `pending_values` order, including when `v == leader`.
            self.vregs[vreg].push(val, &mut self.pool);
            self.value_vregs[val] = vreg.into();

            // Clear the entry in the union-find table. The `find(val)` call may still look at this
            // entry in a future iteration, but that it ok. It will return a rank 0 leader that has
            // already been assigned to the correct virtual register.
            self.union_find[val] = 0;
        }

        // We do *not* call `union_find.clear()` table here because re-initializing the table for
        // sparse use takes time linear in the number of values in the function. Instead we reset
        // the entries that are known to be non-zero in the loop above.
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use entity::EntityRef;
    use ir::Value;

    #[test]
    fn empty_union_find() {
        let mut vregs = VirtRegs::new();
        vregs.finish_union_find(None);
        assert_eq!(vregs.all_virtregs().count(), 0);
    }

    #[test]
    fn union_self() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        vregs.union(v1, v1);
        vregs.finish_union_find(None);
        assert_eq!(vregs.get(v1), None);
        assert_eq!(vregs.all_virtregs().count(), 0);
    }

    #[test]
    fn union_pair() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        let v2 = Value::new(2);
        vregs.union(v1, v2);
        vregs.finish_union_find(None);
        assert_eq!(vregs.congruence_class(&v1), &[v2, v1]);
        assert_eq!(vregs.congruence_class(&v2), &[v2, v1]);
        assert_eq!(vregs.all_virtregs().count(), 1);
    }

    #[test]
    fn union_pair_backwards() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        let v2 = Value::new(2);
        vregs.union(v2, v1);
        vregs.finish_union_find(None);
        assert_eq!(vregs.congruence_class(&v1), &[v1, v2]);
        assert_eq!(vregs.congruence_class(&v2), &[v1, v2]);
        assert_eq!(vregs.all_virtregs().count(), 1);
    }

    #[test]
    fn union_tree() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        let v2 = Value::new(2);
        let v3 = Value::new(3);
        let v4 = Value::new(4);

        vregs.union(v2, v4);
        vregs.union(v3, v1);
        // Leaders: v2, v3
        vregs.union(v4, v1);
        vregs.finish_union_find(None);
        assert_eq!(vregs.congruence_class(&v1), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v2), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v3), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v4), &[v1, v3, v4, v2]);
        assert_eq!(vregs.all_virtregs().count(), 1);
    }

    #[test]
    fn union_two() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        let v2 = Value::new(2);
        let v3 = Value::new(3);
        let v4 = Value::new(4);

        vregs.union(v2, v4);
        vregs.union(v3, v1);
        // Leaders: v2, v3
        vregs.finish_union_find(None);
        assert_eq!(vregs.congruence_class(&v1), &[v1, v3]);
        assert_eq!(vregs.congruence_class(&v2), &[v4, v2]);
        assert_eq!(vregs.congruence_class(&v3), &[v1, v3]);
        assert_eq!(vregs.congruence_class(&v4), &[v4, v2]);
        assert_eq!(vregs.all_virtregs().count(), 2);
    }

    #[test]
    fn union_uneven() {
        let mut vregs = VirtRegs::new();
        let v1 = Value::new(1);
        let v2 = Value::new(2);
        let v3 = Value::new(3);
        let v4 = Value::new(4);

        vregs.union(v2, v4); // Rank 0-0
        vregs.union(v3, v2); // Rank 0-1
        vregs.union(v2, v1); // Rank 1-0
        vregs.finish_union_find(None);
        assert_eq!(vregs.congruence_class(&v1), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v2), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v3), &[v1, v3, v4, v2]);
        assert_eq!(vregs.congruence_class(&v4), &[v1, v3, v4, v2]);
        assert_eq!(vregs.all_virtregs().count(), 1);
    }
}
<|endoftext|>


<commit_before>//! Read a las file.
//!
//! We don't always need to get all the points into memory at once, so this interface enables
//! sequential access.

use std::fs::File;
use std::io::{BufReader, Read, Seek, SeekFrom};
use std::path::Path;

use byteorder::{LittleEndian, ReadBytesExt};

use Result;
use header::Header;
use io::read_full;
use point::{Classification, NumberOfReturns, Point, ReturnNumber, ScanDirection};
use scale::scale;
use vlr::Vlr;

/// A reader of las points.
#[derive(Debug)]
pub struct Reader<R: Read> {
    header: Header,
    vlrs: Vec<Vlr>,
    reader: R,
    nread: u32,
}

impl Reader<BufReader<File>> {
    /// Opens a new reader for the given path.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Reader<BufReader<File>>> {
        let reader = BufReader::new(try!(File::open(path)));
        Reader::new(reader)
    }
}

impl<R: Read + Seek> Reader<R> {
    /// Creates a new reader for a given `Read` object, consuming the read.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::File;
    /// use las::reader::Reader;
    /// let file = File::open("data/1.0_0.las").unwrap();
    /// let reader = Reader::new(file);
    /// ```
    pub fn new(mut reader: R) -> Result<Reader<R>> {
        let header = try!(Header::read_from(&mut reader));
        let mut vlrs = Vec::with_capacity(header.number_of_variable_length_records as usize);
        let _ = try!(reader.seek(SeekFrom::Start(header.header_size as u64)));
        for _ in 0..header.number_of_variable_length_records {
            vlrs.push(try!(Vlr::read_from(&mut reader)));
        }

        let _ = try!(reader.seek(SeekFrom::Start(header.offset_to_point_data as u64)));

        Ok(Reader {
            header: header,
            vlrs: vlrs,
            reader: reader,
            nread: 0,
        })
    }

    /// Returns the next point in this reader.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let mut reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let point = reader.next_point().unwrap();
    /// assert!(point.is_some());
    /// let point = reader.next_point().unwrap();
    /// assert!(point.is_none());
    /// ```
    pub fn next_point(&mut self) -> Result<Option<Point>> {
        if self.eof() {
            return Ok(None);
        }
        let mut point = Point::new();
        let start = try!(self.reader.seek(SeekFrom::Current(0)));
        point.x = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.x_scale_factor,
                        self.header.x_offset);
        point.y = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.y_scale_factor,
                        self.header.y_offset);
        point.z = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.z_scale_factor,
                        self.header.x_offset);
        point.intensity = try!(self.reader.read_u16::<LittleEndian>());
        let byte = try!(self.reader.read_u8());
        point.return_number = try!(ReturnNumber::from_u8(byte & 0b00000111));
        point.number_of_returns = try!(NumberOfReturns::from_u8((byte >> 3) & 0b00000111));
        point.scan_direction = ScanDirection::from((byte >> 6) & 0b00000001 == 1);
        point.edge_of_flight_line = byte >> 7 == 1;
        let byte = try!(self.reader.read_u8());
        point.classification = try!(Classification::from_u8(byte & 0b00011111));
        point.synthetic = (byte >> 5) & 0b00000001 == 1;
        point.key_point = (byte >> 6) & 0b00000001 == 1;
        point.withheld = byte >> 7 == 1;
        point.scan_angle_rank = try!(self.reader.read_i8());
        point.user_data = try!(self.reader.read_u8());
        point.point_source_id = try!(self.reader.read_u16::<LittleEndian>());
        if self.header.point_data_format.has_time() {
            point.gps_time = Some(try!(self.reader.read_f64::<LittleEndian>()));
        }
        if self.header.point_data_format.has_color() {
            point.red = Some(try!(self.reader.read_u16::<LittleEndian>()));
            point.green = Some(try!(self.reader.read_u16::<LittleEndian>()));
            point.blue = Some(try!(self.reader.read_u16::<LittleEndian>()));
        }
        let bytes_read = try!(self.reader.seek(SeekFrom::Current(0))) - start;

        if bytes_read < self.header.point_data_record_length as u64 {
            let mut buf =
                vec![0; (self.header.point_data_record_length as u64 - bytes_read) as usize];
            try!(read_full(&mut self.reader, &mut buf[..]));
            point.extra_bytes = Some(buf);
        }

        self.nread += 1;

        Ok(Some(point))
    }
}

impl<R: Read> Reader<R> {
    /// Returns this reader's las header.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let header = reader.header();
    pub fn header(&self) -> Header {
        self.header
    }

    /// Returns a reference to this reader's vlrs.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let vlrs = reader.vlrs();
    pub fn vlrs(&self) -> &Vec<Vlr> {
        &self.vlrs
    }

    /// Returns the number of points in this lasfile, as specified by the header.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// assert_eq!(1, reader.npoints());
    /// ```
    pub fn npoints(&self) -> u32 {
        self.header.number_of_point_records
    }

    /// Returns true if this reader is at the end of the file.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let mut reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// assert!(!reader.eof());
    /// let _ = reader.next_point().unwrap();
    /// assert!(reader.eof());
    /// ```
    pub fn eof(&self) -> bool {
        self.nread == self.npoints()
    }
}
<commit_msg>Implement IntoIterator for Reader<commit_after>//! Read a las file.
//!
//! We don't always need to get all the points into memory at once, so this interface enables
//! sequential access.

use std::fs::File;
use std::io::{BufReader, Read, Seek, SeekFrom};
use std::iter::IntoIterator;
use std::path::Path;

use byteorder::{LittleEndian, ReadBytesExt};

use Result;
use header::Header;
use io::read_full;
use point::{Classification, NumberOfReturns, Point, ReturnNumber, ScanDirection};
use scale::scale;
use vlr::Vlr;

/// A reader of las points.
#[derive(Debug)]
pub struct Reader<R: Read> {
    header: Header,
    vlrs: Vec<Vlr>,
    reader: R,
    nread: u32,
}

impl Reader<BufReader<File>> {
    /// Opens a new reader for the given path.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Reader<BufReader<File>>> {
        let reader = BufReader::new(try!(File::open(path)));
        Reader::new(reader)
    }
}

impl<R: Read + Seek> Reader<R> {
    /// Creates a new reader for a given `Read` object, consuming the read.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::File;
    /// use las::reader::Reader;
    /// let file = File::open("data/1.0_0.las").unwrap();
    /// let reader = Reader::new(file);
    /// ```
    pub fn new(mut reader: R) -> Result<Reader<R>> {
        let header = try!(Header::read_from(&mut reader));
        let mut vlrs = Vec::with_capacity(header.number_of_variable_length_records as usize);
        let _ = try!(reader.seek(SeekFrom::Start(header.header_size as u64)));
        for _ in 0..header.number_of_variable_length_records {
            vlrs.push(try!(Vlr::read_from(&mut reader)));
        }

        let _ = try!(reader.seek(SeekFrom::Start(header.offset_to_point_data as u64)));

        Ok(Reader {
            header: header,
            vlrs: vlrs,
            reader: reader,
            nread: 0,
        })
    }

    /// Returns the next point in this reader.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let mut reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let point = reader.next_point().unwrap();
    /// assert!(point.is_some());
    /// let point = reader.next_point().unwrap();
    /// assert!(point.is_none());
    /// ```
    pub fn next_point(&mut self) -> Result<Option<Point>> {
        if self.eof() {
            return Ok(None);
        }
        let mut point = Point::new();
        let start = try!(self.reader.seek(SeekFrom::Current(0)));
        point.x = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.x_scale_factor,
                        self.header.x_offset);
        point.y = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.y_scale_factor,
                        self.header.y_offset);
        point.z = scale(try!(self.reader.read_i32::<LittleEndian>()),
                        self.header.z_scale_factor,
                        self.header.x_offset);
        point.intensity = try!(self.reader.read_u16::<LittleEndian>());
        let byte = try!(self.reader.read_u8());
        point.return_number = try!(ReturnNumber::from_u8(byte & 0b00000111));
        point.number_of_returns = try!(NumberOfReturns::from_u8((byte >> 3) & 0b00000111));
        point.scan_direction = ScanDirection::from((byte >> 6) & 0b00000001 == 1);
        point.edge_of_flight_line = byte >> 7 == 1;
        let byte = try!(self.reader.read_u8());
        point.classification = try!(Classification::from_u8(byte & 0b00011111));
        point.synthetic = (byte >> 5) & 0b00000001 == 1;
        point.key_point = (byte >> 6) & 0b00000001 == 1;
        point.withheld = byte >> 7 == 1;
        point.scan_angle_rank = try!(self.reader.read_i8());
        point.user_data = try!(self.reader.read_u8());
        point.point_source_id = try!(self.reader.read_u16::<LittleEndian>());
        if self.header.point_data_format.has_time() {
            point.gps_time = Some(try!(self.reader.read_f64::<LittleEndian>()));
        }
        if self.header.point_data_format.has_color() {
            point.red = Some(try!(self.reader.read_u16::<LittleEndian>()));
            point.green = Some(try!(self.reader.read_u16::<LittleEndian>()));
            point.blue = Some(try!(self.reader.read_u16::<LittleEndian>()));
        }
        let bytes_read = try!(self.reader.seek(SeekFrom::Current(0))) - start;

        if bytes_read < self.header.point_data_record_length as u64 {
            let mut buf =
                vec![0; (self.header.point_data_record_length as u64 - bytes_read) as usize];
            try!(read_full(&mut self.reader, &mut buf[..]));
            point.extra_bytes = Some(buf);
        }

        self.nread += 1;

        Ok(Some(point))
    }
}

impl<R: Read> Reader<R> {
    /// Returns this reader's las header.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let header = reader.header();
    pub fn header(&self) -> Header {
        self.header
    }

    /// Returns a reference to this reader's vlrs.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// let vlrs = reader.vlrs();
    pub fn vlrs(&self) -> &Vec<Vlr> {
        &self.vlrs
    }

    /// Returns the number of points in this lasfile, as specified by the header.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// assert_eq!(1, reader.npoints());
    /// ```
    pub fn npoints(&self) -> u32 {
        self.header.number_of_point_records
    }

    /// Returns true if this reader is at the end of the file.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::reader::Reader;
    /// let mut reader = Reader::from_path("data/1.0_0.las").unwrap();
    /// assert!(!reader.eof());
    /// let _ = reader.next_point().unwrap();
    /// assert!(reader.eof());
    /// ```
    pub fn eof(&self) -> bool {
        self.nread == self.npoints()
    }
}

impl<R: Read + Seek> IntoIterator for Reader<R> {
    type Item = Point;
    type IntoIter = ReaderIterator<R>;
    fn into_iter(self) -> Self::IntoIter {
        ReaderIterator { reader: self }
    }
}

/// An iterator over the reader's points.
///
/// # Panics
///
/// This iterator will panic if something goes wrong while reading the points. If you need to catch
/// those errors, use `Reader::next_point`.
#[derive(Debug)]
pub struct ReaderIterator<R: Read + Seek> {
    reader: Reader<R>,
}

impl<R: Read + Seek> Iterator for ReaderIterator<R> {
    type Item = Point;
    fn next(&mut self) -> Option<Self::Item> {
        self.reader.next_point().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reader_as_iter() {
        let reader = Reader::from_path("data/1.0_0.las").unwrap();
        let points: Vec<_> = reader.into_iter().collect();
        assert_eq!(1, points.len());
    }
}
<|endoftext|>


<commit_before>use std::fs::{self, File};
use std::io::prelude::*;

use cargotest::support::{project, execs};
use hamcrest::{assert_that, existing_file, is_not};

#[test]
fn adding_and_removing_packages() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let toml = p.root().join("Cargo.toml");
    let lock1 = p.read_lockfile();

    // add a dep
    File::create(&toml).unwrap().write_all(br#"
        [package]
        name = "foo"
        authors = []
        version = "0.0.1"

        [dependencies.bar]
        path = "bar"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock2 = p.read_lockfile();
    assert_ne!(lock1, lock2);

    // change the dep
    File::create(&p.root().join("bar/Cargo.toml")).unwrap().write_all(br#"
        [package]
        name = "bar"
        authors = []
        version = "0.0.2"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock3 = p.read_lockfile();
    assert_ne!(lock1, lock3);
    assert_ne!(lock2, lock3);

    // remove the dep
    println!("lock4");
    File::create(&toml).unwrap().write_all(br#"
        [package]
        name = "foo"
        authors = []
        version = "0.0.1"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock4 = p.read_lockfile();
    assert_eq!(lock1, lock4);
}

#[test]
fn no_index_update() {
    use cargotest::ChannelChanger;
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"

            [dependencies]
            serde = "1.0"
        "#)
        .file("src/main.rs", "fn main() {}")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_stdout("")
                    .with_stderr_contains("    Updating registry `https://github.com/rust-lang/crates.io-index`"));

    assert_that(p.cargo("generate-lockfile").masquerade_as_nightly_cargo().arg("-Zno-index-update"),
                execs().with_status(0).with_stdout("")
                    .with_stderr_does_not_contain("    Updating registry `https://github.com/rust-lang/crates.io-index`"));
}

#[test]
fn preserve_metadata() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let metadata = r#"
[metadata]
bar = "baz"
foo = "bar"
"#;
    let lockfile = p.root().join("Cargo.lock");
    let lock = p.read_lockfile();
    let data = lock + metadata;
    File::create(&lockfile).unwrap().write_all(data.as_bytes()).unwrap();

    // Build and make sure the metadata is still there
    assert_that(p.cargo("build"),
                execs().with_status(0));
    let lock = p.read_lockfile();
    assert!(lock.contains(metadata.trim()), "{}", lock);

    // Update and make sure the metadata is still there
    assert_that(p.cargo("update"),
                execs().with_status(0));
    let lock = p.read_lockfile();
    assert!(lock.contains(metadata.trim()), "{}", lock);
}

#[test]
fn preserve_line_endings_issue_2076() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    let lockfile = p.root().join("Cargo.lock");
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    assert_that(&lockfile,
                existing_file());
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let lock0 = p.read_lockfile();

    assert!(lock0.starts_with("[[package]]\n"));

    let lock1 = lock0.replace("\n", "\r\n");
    {
        File::create(&lockfile).unwrap().write_all(lock1.as_bytes()).unwrap();
    }

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let lock2 = p.read_lockfile();

    assert!(lock2.starts_with("[[package]]\r\n"));
    assert_eq!(lock1, lock2);
}

#[test]
fn cargo_update_generate_lockfile() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .build();

    let lockfile = p.root().join("Cargo.lock");
    assert_that(&lockfile, is_not(existing_file()));
    assert_that(p.cargo("update"), execs().with_status(0).with_stdout(""));
    assert_that(&lockfile, existing_file());

    fs::remove_file(p.root().join("Cargo.lock")).unwrap();

    assert_that(&lockfile, is_not(existing_file()));
    assert_that(p.cargo("update"), execs().with_status(0).with_stdout(""));
    assert_that(&lockfile, existing_file());

}
<commit_msg>Auto merge of #5119 - alexcrichton:no-clone-index-in-tests, r=matklad<commit_after>use std::fs::{self, File};
use std::io::prelude::*;

use cargotest::support::{project, execs};
use cargotest::support::registry::Package;
use cargotest::ChannelChanger;
use hamcrest::{assert_that, existing_file, is_not};

#[test]
fn adding_and_removing_packages() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let toml = p.root().join("Cargo.toml");
    let lock1 = p.read_lockfile();

    // add a dep
    File::create(&toml).unwrap().write_all(br#"
        [package]
        name = "foo"
        authors = []
        version = "0.0.1"

        [dependencies.bar]
        path = "bar"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock2 = p.read_lockfile();
    assert_ne!(lock1, lock2);

    // change the dep
    File::create(&p.root().join("bar/Cargo.toml")).unwrap().write_all(br#"
        [package]
        name = "bar"
        authors = []
        version = "0.0.2"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock3 = p.read_lockfile();
    assert_ne!(lock1, lock3);
    assert_ne!(lock2, lock3);

    // remove the dep
    println!("lock4");
    File::create(&toml).unwrap().write_all(br#"
        [package]
        name = "foo"
        authors = []
        version = "0.0.1"
    "#).unwrap();
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    let lock4 = p.read_lockfile();
    assert_eq!(lock1, lock4);
}

#[test]
fn no_index_update() {
    Package::new("serde", "1.0.0").publish();

    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"

            [dependencies]
            serde = "1.0"
        "#)
        .file("src/main.rs", "fn main() {}")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_stderr("\
[UPDATING] registry `[..]`
"));

    assert_that(p.cargo("generate-lockfile").masquerade_as_nightly_cargo().arg("-Zno-index-update"),
                execs().with_status(0).with_stdout("").with_stderr(""));
}

#[test]
fn preserve_metadata() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let metadata = r#"
[metadata]
bar = "baz"
foo = "bar"
"#;
    let lockfile = p.root().join("Cargo.lock");
    let lock = p.read_lockfile();
    let data = lock + metadata;
    File::create(&lockfile).unwrap().write_all(data.as_bytes()).unwrap();

    // Build and make sure the metadata is still there
    assert_that(p.cargo("build"),
                execs().with_status(0));
    let lock = p.read_lockfile();
    assert!(lock.contains(metadata.trim()), "{}", lock);

    // Update and make sure the metadata is still there
    assert_that(p.cargo("update"),
                execs().with_status(0));
    let lock = p.read_lockfile();
    assert!(lock.contains(metadata.trim()), "{}", lock);
}

#[test]
fn preserve_line_endings_issue_2076() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .file("bar/Cargo.toml", r#"
            [package]
            name = "bar"
            authors = []
            version = "0.0.1"
        "#)
        .file("bar/src/lib.rs", "")
        .build();

    let lockfile = p.root().join("Cargo.lock");
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));
    assert_that(&lockfile,
                existing_file());
    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let lock0 = p.read_lockfile();

    assert!(lock0.starts_with("[[package]]\n"));

    let lock1 = lock0.replace("\n", "\r\n");
    {
        File::create(&lockfile).unwrap().write_all(lock1.as_bytes()).unwrap();
    }

    assert_that(p.cargo("generate-lockfile"),
                execs().with_status(0));

    let lock2 = p.read_lockfile();

    assert!(lock2.starts_with("[[package]]\r\n"));
    assert_eq!(lock1, lock2);
}

#[test]
fn cargo_update_generate_lockfile() {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.1"
        "#)
        .file("src/main.rs", "fn main() {}")
        .build();

    let lockfile = p.root().join("Cargo.lock");
    assert_that(&lockfile, is_not(existing_file()));
    assert_that(p.cargo("update"), execs().with_status(0).with_stdout(""));
    assert_that(&lockfile, existing_file());

    fs::remove_file(p.root().join("Cargo.lock")).unwrap();

    assert_that(&lockfile, is_not(existing_file()));
    assert_that(p.cargo("update"), execs().with_status(0).with_stdout(""));
    assert_that(&lockfile, existing_file());

}
<|endoftext|>