<commit_before>use super::*;
use std::collections::HashMap;
use termion::event as keyevent;
use nalgebra::{distance, Vector2};
use tui::{layout::{Direction, Group, Rect, Size}, style::{Color, Style},
          widgets::{Block, Borders, Paragraph, Row, SelectableList, Table, Widget, canvas::Canvas}};

use utils::Point;
use astronomicals::system::System;
use entities::Faction;
use player::Player;

lazy_static! {
    /// Color mapping for each faction.
    static ref FACTION_COLORS: HashMap<Faction, Color> = {
        let mut m = HashMap::new();
        m.insert(Faction::Empire, Color::Red);
        m.insert(Faction::Federation, Color::Yellow);
        m.insert(Faction::Cartel, Color::Magenta);
        m.insert(Faction::Independent, Color::LightGreen);
        m
    };

    /// Styling for unselected item.
    static ref DEFAULT_STYLE: Style = Style::default();
}

/// The minimum distance within which the gui will snap to the closest system.
const MIN_SNAP_DIST: f64 = 0.9;

/// Displays the map tab.
pub struct MapTab {
    state: Arc<Game>,
    sender: Sender<Event>,
    search_mode: bool,
    search_str: String,
    selected: Option<Point>,
    route: Option<(u32, Vec<Point>)>,
    cursor: Point,
    map_scale: f64,
}

impl MapTab {
    /// Attempts to find a route to the selected system.
    fn find_route(&mut self) {
        let galaxy = &self.state.galaxy.lock().unwrap();
        let player = &mut self.state.player.lock().unwrap();
        let range = match player.ship() {
            &Some(ref ship) => ship.range(),
            &None => 0.,
        };
        let max_jumps = match player.ship() {
            &Some(ref ship) => ship.fuel(),
            &None => 0,
        };
        // Plan route if possible.
        self.route = galaxy.route(
            &player.location(),
            &self.selected.unwrap(),
            range,
            max_jumps,
        );
    }

    /// Moves the player's location to the selected system.
    fn travel_to_selected(&mut self) {
        let player = &mut self.state.player.lock().unwrap();

        // Only travel if the selected system is the same as the cursor and
        // and the final destination for the route.
        if let Some((ref jumps, ref route)) = self.route {
            if self.selected.is_some() && self.selected.unwrap() == self.cursor
                && self.selected.unwrap() == *route.last().unwrap()
            {
                player.set_route(route.clone());
                self.sender.send(Event::Travel);
            }
        }

        // Reset route.
        self.route = None;
    }

    /// Draws the event box in the given terminal and area.
    pub fn draw_search(&self, term: &mut Terminal<MouseBackend>, area: &Rect) {
        let draw_str = match self.search_mode {
            true => self.search_str.as_str(),
            false => "Press '/' to search for a system",
        };
        Paragraph::default()
            .block(Block::default().borders(Borders::ALL))
            .style(Style::default().fg(Color::Yellow))
            .text(draw_str)
            .render(term, &area);
    }

    /// Draw system ship information for the selected system, if any.
    fn draw_system_info(
        &self,
        player_loc: &Point,
        selected_system: Option<&System>,
        term: &mut Terminal<MouseBackend>,
        area: &Rect,
    ) {
        // Do not draw anything if no system is selected.
        if selected_system.is_none() {
            return;
        }
        let system = selected_system.unwrap();

        let system_data = vec![
            format!("Faction:       {}", system.faction.to_string()),
            format!("State:         {}", system.state.to_string()),
            format!("Security:      {}", system.security.to_string()),
            format!("Reputation:    {}", system.reputation.to_string()),
            format!(
                "Distance:      {:.1} ly",
                distance(player_loc, &system.location)
            ),
            format!("Star mass:     {:.1} M", system.star.mass),
            format!("Star type:     {}", system.star.startype.to_string()),
            format!("Bodies:        {}", system.satelites.len()),
        ];

        Group::default()
            .direction(Direction::Vertical)
            .sizes(&[Size::Fixed(9), Size::Min(1)])
            .render(term, area, |term, chunks| {
                SelectableList::default()
                    .items(&system_data)
                    .block(Block::default().title(&format!("{} System", system.name)))
                    .style(Style::default().fg(Color::Yellow))
                    .render(term, &chunks[0]);
                Table::new(
                    // Prepending empty character to get alignment with list above.
                    [" Planet", "Mass", "Population", "Temperature", "Type"].into_iter(),
                    system.satelites.iter().map(|ref planet| {
                        let style: &Style = &DEFAULT_STYLE;
                        Row::StyledData(
                            vec![
                                format!(" {}", planet.name.clone()),
                                format!("{:.1}", planet.mass),
                                format!("{:.1} M", planet.population),
                                format!("{:.1}", planet.surface_temperature),
                                planet.planet_type.to_string(),
                            ].into_iter(),
                            &style,
                        )
                    }),
                ).block(Block::default().title("Planets"))
                    .header_style(Style::default().fg(Color::Yellow))
                    .widths(&[15, 5, 15, 15, 10])
                    .render(term, &chunks[1]);
            });
    }

    /// Draw the galaxy map.
    fn draw_galaxy_map(
        &self,
        player: &Player,
        systems: &Vec<&System>,
        term: &mut Terminal<MouseBackend>,
        area: &Rect,
    ) {
        let player_loc = player.location();
        // Scale map to not overlap systems.
        let map_scaling = 20. * self.map_scale;
        let (max_x, max_y) = systems.iter().fold((0., 0.), |(x_max, y_max), s| {
            (
                (s.location.coords.x / map_scaling).abs().max(x_max),
                (s.location.coords.y / map_scaling).abs().max(y_max),
            )
        });
        Canvas::default()
            .block(Block::default().title("Systems").borders(Borders::ALL))
            .paint(|ctx| {
                for system in systems.iter() {
                    let color = FACTION_COLORS.get(&system.faction).unwrap().clone();
                    ctx.print(
                        system.location.coords.x,
                        system.location.coords.y,
                        ".",
                        color,
                    );
                }
                // Draw player location.
                ctx.print(player_loc.coords.x, player_loc.coords.y, "X", Color::White);

                // Draw the cursor.
                ctx.print(
                    self.cursor.coords.x,
                    self.cursor.coords.y,
                    "*",
                    Color::Yellow,
                );

                // Draw route if available.
                if let Some((_, ref route)) = self.route {
                    for system in route {
                        ctx.print(system.coords.x, system.coords.y, "X", Color::Yellow);
                    }
                    ctx.print(player_loc.coords.x, player_loc.coords.y, "S", Color::Yellow);
                    ctx.print(
                        route.last().unwrap().coords.x,
                        route.last().unwrap().coords.y,
                        "G",
                        Color::Yellow,
                    );
                }
                // Draw currently travelling route if available.
                if let Some(ref route) = player.route() {
                    for system in route {
                        ctx.print(system.coords.x, system.coords.y, "X", Color::White);
                    }
                    ctx.print(player_loc.coords.x, player_loc.coords.y, "S", Color::White);
                    ctx.print(
                        route.last().unwrap().coords.x,
                        route.last().unwrap().coords.y,
                        "G",
                        Color::White,
                    );
                }
            })
            .x_bounds([self.cursor.coords.x - max_x, self.cursor.coords.x + max_x])
            .y_bounds([self.cursor.coords.y - max_y, self.cursor.coords.y + max_y])
            .render(term, &area);
    }
}

impl Tab for MapTab {
    /// Creates a map tab.
    fn new(state: Arc<Game>, send_handle: Sender<Event>) -> Box<Self> {
        let cursor = state.player.lock().unwrap().location().clone();
        Box::new(MapTab {
            state: state,
            sender: send_handle,
            selected: Some(cursor.clone()),
            search_mode: false,
            search_str: String::new(),
            route: None,
            cursor: cursor,
            map_scale: 1.,
        })
    }

    /// Returns the title string describing the tab.
    fn title(&self) -> String {
        String::from("Galaxy Map")
    }

    /// Handles the user provided event.
    fn handle_event(&mut self, event: Event) {
        match event {
            Event::Input(input) => {
                match input {
                    keyevent::Key::Char(' ') if self.selected.is_some() => {
                        match self.route {
                            Some(_) => self.travel_to_selected(),
                            None => self.find_route(),
                        };
                    }
                    // Start search mode.
                    keyevent::Key::Char('/') => {
                        self.search_mode = true;
                        return;
                    }
                    // Quit search mode.
                    keyevent::Key::Esc => {
                        self.search_str.clear();
                        self.search_mode = false;
                        return;
                    }
                    _ => {}
                };

                if self.search_mode {
                    match input {
                        keyevent::Key::Char('\n') => {
                            let galaxy = self.state.galaxy.lock().unwrap();

                            // Set cursor to the closest matching system if
                            // possible.
                            match galaxy.search_name(&self.search_str) {
                                Some(system) => self.cursor = system.location.clone(),
                                None => {}
                            };

                            // Clear input.
                            self.search_str.clear();
                            self.search_mode = false;
                        }
                        keyevent::Key::Char(e) => {
                            self.search_str.push(e);
                            // Early exit.
                            return;
                        }
                        keyevent::Key::Backspace => {
                            self.search_str.pop();
                        }
                        _ => {}
                    };
                }

                self.map_scale *= match input {
                    // Zoom out.
                    keyevent::Key::Char('u') => 0.5,
                    // Zoom in.
                    keyevent::Key::Char('i') => 2.,
                    // No zooming.
                    _ => 1.,
                };

                // Prevent zooming too far in.
                self.map_scale = self.map_scale.min(4.);

                let translation = match input {
                    // Move up.
                    keyevent::Key::Char('k') => Vector2::new(0., 1. / self.map_scale),
                    // Move down.
                    keyevent::Key::Char('j') => Vector2::new(0., -1. / self.map_scale),
                    // Move right.
                    keyevent::Key::Char('l') => Vector2::new(1. / self.map_scale, 0.),
                    // Move left.
                    keyevent::Key::Char('h') => Vector2::new(-1. / self.map_scale, 0.),
                    _ => Vector2::new(0., 0.),
                };

                // Move out of snapping system if currently snapped.
                self.cursor += match self.selected {
                    Some(_) => Vector2::new(
                        translation.x * 1.1 * self.map_scale * MIN_SNAP_DIST,
                        translation.y * 1.1 * self.map_scale * MIN_SNAP_DIST,
                    ),
                    None => translation,
                };
                self.selected = None;

                // Check if cursor should snap to closest system.
                if let Some(neighbor) = self.state.galaxy.lock().unwrap().nearest(&self.cursor) {
                    if distance(&self.cursor, &neighbor) < MIN_SNAP_DIST {
                        self.cursor = neighbor.clone();
                        self.selected = Some(neighbor.clone());
                    }
                }
            }
            _ => {}
        };
    }

    /// Draws the tab in the given terminal and area.
    fn draw(&self, term: &mut Terminal<MouseBackend>, area: &Rect) {
        let galaxy = self.state.galaxy.lock().unwrap();
        Group::default()
            .direction(Direction::Horizontal)
            .sizes(&[Size::Fixed(75), Size::Min(1)])
            .render(term, area, |term, chunks| {
                // TODO: Draw system detailed information.
                let systems = &galaxy.systems();
                let player = &self.state.player.lock().unwrap();
                //let player_loc = &self.state.player.lock().unwrap().location().clone();
                // Draw sidebar.
                Group::default()
                    .direction(Direction::Vertical)
                    .sizes(&[Size::Min(1), Size::Fixed(3)])
                    .render(term, &chunks[0], |term, sidebar_chunk| {
                        self.draw_system_info(
                            &player.location(),
                            self.selected.map(|point| galaxy.system(&point).unwrap()),
                            term,
                            &sidebar_chunk[0],
                        );
                        self.draw_search(term, &sidebar_chunk[1]);
                    });
                self.draw_galaxy_map(player, &systems, term, &chunks[1]);
            });
    }
}
<commit_msg>Use enter key for travel instead of space (#109)<commit_after>use super::*;
use std::collections::HashMap;
use termion::event as keyevent;
use nalgebra::{distance, Vector2};
use tui::{layout::{Direction, Group, Rect, Size}, style::{Color, Style},
          widgets::{Block, Borders, Paragraph, Row, SelectableList, Table, Widget, canvas::Canvas}};

use utils::Point;
use astronomicals::system::System;
use entities::Faction;
use player::Player;

lazy_static! {
    /// Color mapping for each faction.
    static ref FACTION_COLORS: HashMap<Faction, Color> = {
        let mut m = HashMap::new();
        m.insert(Faction::Empire, Color::Red);
        m.insert(Faction::Federation, Color::Yellow);
        m.insert(Faction::Cartel, Color::Magenta);
        m.insert(Faction::Independent, Color::LightGreen);
        m
    };

    /// Styling for unselected item.
    static ref DEFAULT_STYLE: Style = Style::default();
}

/// The minimum distance within which the gui will snap to the closest system.
const MIN_SNAP_DIST: f64 = 0.9;

/// Displays the map tab.
pub struct MapTab {
    state: Arc<Game>,
    sender: Sender<Event>,
    search_mode: bool,
    search_str: String,
    selected: Option<Point>,
    route: Option<(u32, Vec<Point>)>,
    cursor: Point,
    map_scale: f64,
}

impl MapTab {
    /// Attempts to find a route to the selected system.
    fn find_route(&mut self) {
        let galaxy = &self.state.galaxy.lock().unwrap();
        let player = &mut self.state.player.lock().unwrap();
        let range = match player.ship() {
            &Some(ref ship) => ship.range(),
            &None => 0.,
        };
        let max_jumps = match player.ship() {
            &Some(ref ship) => ship.fuel(),
            &None => 0,
        };
        // Plan route if possible.
        self.route = galaxy.route(
            &player.location(),
            &self.selected.unwrap(),
            range,
            max_jumps,
        );
    }

    /// Moves the player's location to the selected system.
    fn travel_to_selected(&mut self) {
        let player = &mut self.state.player.lock().unwrap();

        // Only travel if the selected system is the same as the cursor and
        // and the final destination for the route.
        if let Some((ref jumps, ref route)) = self.route {
            if self.selected.is_some() && self.selected.unwrap() == self.cursor
                && self.selected.unwrap() == *route.last().unwrap()
            {
                player.set_route(route.clone());
                self.sender.send(Event::Travel);
            }
        }

        // Reset route.
        self.route = None;
    }

    /// Draws the event box in the given terminal and area.
    pub fn draw_search(&self, term: &mut Terminal<MouseBackend>, area: &Rect) {
        let draw_str = match self.search_mode {
            true => self.search_str.as_str(),
            false => "Press '/' to search for a system",
        };
        Paragraph::default()
            .block(Block::default().borders(Borders::ALL))
            .style(Style::default().fg(Color::Yellow))
            .text(draw_str)
            .render(term, &area);
    }

    /// Draw system ship information for the selected system, if any.
    fn draw_system_info(
        &self,
        player_loc: &Point,
        selected_system: Option<&System>,
        term: &mut Terminal<MouseBackend>,
        area: &Rect,
    ) {
        // Do not draw anything if no system is selected.
        if selected_system.is_none() {
            return;
        }
        let system = selected_system.unwrap();

        let system_data = vec![
            format!("Faction:       {}", system.faction.to_string()),
            format!("State:         {}", system.state.to_string()),
            format!("Security:      {}", system.security.to_string()),
            format!("Reputation:    {}", system.reputation.to_string()),
            format!(
                "Distance:      {:.1} ly",
                distance(player_loc, &system.location)
            ),
            format!("Star mass:     {:.1} M", system.star.mass),
            format!("Star type:     {}", system.star.startype.to_string()),
            format!("Bodies:        {}", system.satelites.len()),
        ];

        Group::default()
            .direction(Direction::Vertical)
            .sizes(&[Size::Fixed(9), Size::Min(1)])
            .render(term, area, |term, chunks| {
                SelectableList::default()
                    .items(&system_data)
                    .block(Block::default().title(&format!("{} System", system.name)))
                    .style(Style::default().fg(Color::Yellow))
                    .render(term, &chunks[0]);
                Table::new(
                    // Prepending empty character to get alignment with list above.
                    [" Planet", "Mass", "Population", "Temperature", "Type"].into_iter(),
                    system.satelites.iter().map(|ref planet| {
                        let style: &Style = &DEFAULT_STYLE;
                        Row::StyledData(
                            vec![
                                format!(" {}", planet.name.clone()),
                                format!("{:.1}", planet.mass),
                                format!("{:.1} M", planet.population),
                                format!("{:.1}", planet.surface_temperature),
                                planet.planet_type.to_string(),
                            ].into_iter(),
                            &style,
                        )
                    }),
                ).block(Block::default().title("Planets"))
                    .header_style(Style::default().fg(Color::Yellow))
                    .widths(&[15, 5, 15, 15, 10])
                    .render(term, &chunks[1]);
            });
    }

    /// Draw the galaxy map.
    fn draw_galaxy_map(
        &self,
        player: &Player,
        systems: &Vec<&System>,
        term: &mut Terminal<MouseBackend>,
        area: &Rect,
    ) {
        let player_loc = player.location();
        // Scale map to not overlap systems.
        let map_scaling = 20. * self.map_scale;
        let (max_x, max_y) = systems.iter().fold((0., 0.), |(x_max, y_max), s| {
            (
                (s.location.coords.x / map_scaling).abs().max(x_max),
                (s.location.coords.y / map_scaling).abs().max(y_max),
            )
        });
        Canvas::default()
            .block(Block::default().title("Systems").borders(Borders::ALL))
            .paint(|ctx| {
                for system in systems.iter() {
                    let color = FACTION_COLORS.get(&system.faction).unwrap().clone();
                    ctx.print(
                        system.location.coords.x,
                        system.location.coords.y,
                        ".",
                        color,
                    );
                }
                // Draw player location.
                ctx.print(player_loc.coords.x, player_loc.coords.y, "X", Color::White);

                // Draw the cursor.
                ctx.print(
                    self.cursor.coords.x,
                    self.cursor.coords.y,
                    "*",
                    Color::Yellow,
                );

                // Draw route if available.
                if let Some((_, ref route)) = self.route {
                    for system in route {
                        ctx.print(system.coords.x, system.coords.y, "X", Color::Yellow);
                    }
                    ctx.print(player_loc.coords.x, player_loc.coords.y, "S", Color::Yellow);
                    ctx.print(
                        route.last().unwrap().coords.x,
                        route.last().unwrap().coords.y,
                        "G",
                        Color::Yellow,
                    );
                }
                // Draw currently travelling route if available.
                if let Some(ref route) = player.route() {
                    for system in route {
                        ctx.print(system.coords.x, system.coords.y, "X", Color::White);
                    }
                    ctx.print(player_loc.coords.x, player_loc.coords.y, "S", Color::White);
                    ctx.print(
                        route.last().unwrap().coords.x,
                        route.last().unwrap().coords.y,
                        "G",
                        Color::White,
                    );
                }
            })
            .x_bounds([self.cursor.coords.x - max_x, self.cursor.coords.x + max_x])
            .y_bounds([self.cursor.coords.y - max_y, self.cursor.coords.y + max_y])
            .render(term, &area);
    }
}

impl Tab for MapTab {
    /// Creates a map tab.
    fn new(state: Arc<Game>, send_handle: Sender<Event>) -> Box<Self> {
        let cursor = state.player.lock().unwrap().location().clone();
        Box::new(MapTab {
            state: state,
            sender: send_handle,
            selected: Some(cursor.clone()),
            search_mode: false,
            search_str: String::new(),
            route: None,
            cursor: cursor,
            map_scale: 1.,
        })
    }

    /// Returns the title string describing the tab.
    fn title(&self) -> String {
        String::from("Galaxy Map")
    }

    /// Handles the user provided event.
    fn handle_event(&mut self, event: Event) {
        match event {
            Event::Input(input) => {
                match input {
                    keyevent::Key::Char('\n') if self.search_mode => {
                        let galaxy = self.state.galaxy.lock().unwrap();

                        // Set cursor to the closest matching system if
                        // possible.
                        match galaxy.search_name(&self.search_str) {
                            Some(system) => self.cursor = system.location.clone(),
                            None => {}
                        };

                        // Clear input.
                        self.search_str.clear();
                        self.search_mode = false;
                    }
                    keyevent::Key::Char(e) if self.search_mode => {
                        self.search_str.push(e);
                        // Early exit.
                        return;
                    }
                    keyevent::Key::Backspace if self.search_mode => {
                        self.search_str.pop();
                    }
                    keyevent::Key::Char('\n') if self.selected.is_some() => {
                        match self.route {
                            Some(_) => self.travel_to_selected(),
                            None => self.find_route(),
                        };
                    }
                    // Start search mode.
                    keyevent::Key::Char('/') => {
                        self.search_mode = true;
                        return;
                    }
                    // Quit search mode.
                    keyevent::Key::Esc => {
                        self.search_str.clear();
                        self.search_mode = false;
                        return;
                    }
                    _ => {}
                };

                self.map_scale *= match input {
                    // Zoom out.
                    keyevent::Key::Char('u') => 0.5,
                    // Zoom in.
                    keyevent::Key::Char('i') => 2.,
                    // No zooming.
                    _ => 1.,
                };

                // Prevent zooming too far in.
                self.map_scale = self.map_scale.min(4.);

                let translation = match input {
                    // Move up.
                    keyevent::Key::Char('k') => Vector2::new(0., 1. / self.map_scale),
                    // Move down.
                    keyevent::Key::Char('j') => Vector2::new(0., -1. / self.map_scale),
                    // Move right.
                    keyevent::Key::Char('l') => Vector2::new(1. / self.map_scale, 0.),
                    // Move left.
                    keyevent::Key::Char('h') => Vector2::new(-1. / self.map_scale, 0.),
                    _ => Vector2::new(0., 0.),
                };

                // Move out of snapping system if currently snapped.
                self.cursor += match self.selected {
                    Some(_) => Vector2::new(
                        translation.x * 1.1 * self.map_scale * MIN_SNAP_DIST,
                        translation.y * 1.1 * self.map_scale * MIN_SNAP_DIST,
                    ),
                    None => translation,
                };
                self.selected = None;

                // Check if cursor should snap to closest system.
                if let Some(neighbor) = self.state.galaxy.lock().unwrap().nearest(&self.cursor) {
                    if distance(&self.cursor, &neighbor) < MIN_SNAP_DIST {
                        self.cursor = neighbor.clone();
                        self.selected = Some(neighbor.clone());
                    }
                }
            }
            _ => {}
        };
    }

    /// Draws the tab in the given terminal and area.
    fn draw(&self, term: &mut Terminal<MouseBackend>, area: &Rect) {
        let galaxy = self.state.galaxy.lock().unwrap();
        Group::default()
            .direction(Direction::Horizontal)
            .sizes(&[Size::Fixed(75), Size::Min(1)])
            .render(term, area, |term, chunks| {
                // TODO: Draw system detailed information.
                let systems = &galaxy.systems();
                let player = &self.state.player.lock().unwrap();
                //let player_loc = &self.state.player.lock().unwrap().location().clone();
                // Draw sidebar.
                Group::default()
                    .direction(Direction::Vertical)
                    .sizes(&[Size::Min(1), Size::Fixed(3)])
                    .render(term, &chunks[0], |term, sidebar_chunk| {
                        self.draw_system_info(
                            &player.location(),
                            self.selected.map(|point| galaxy.system(&point).unwrap()),
                            term,
                            &sidebar_chunk[0],
                        );
                        self.draw_search(term, &sidebar_chunk[1]);
                    });
                self.draw_galaxy_map(player, &systems, term, &chunks[1]);
            });
    }
}
<|endoftext|>


<commit_before>//! `HyperLogLog` implementation.
use bytecount;
use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{BuildHasher, Hash, Hasher};

use hash_utils::MyBuildHasherDefault;

/// A simple implementation of a [HyperLogLog](https://en.wikipedia.org/wiki/HyperLogLog)
#[derive(Clone)]
pub struct HyperLogLog<B = MyBuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    registers: Vec<u8>,
    b: usize,
    buildhasher: B,
}

impl HyperLogLog {
    /// Creates a new, empty HyperLogLog.
    ///
    /// - `b` number of bits used for register selection, number of registers within the
    ///   HyperLogLog will be `2^b`. `b` must be in `[4, 16]`
    ///
    /// Panics when `b` is out of bounds.
    pub fn new(b: usize) -> Self {
        let bh = MyBuildHasherDefault::<DefaultHasher>::default();
        Self::with_hash(b, bh)
    }
}

impl<B> HyperLogLog<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Same as `new` but with a specific `BuildHasher`.
    pub fn with_hash(b: usize, buildhasher: B) -> Self {
        assert!(
            (b >= 4) & (b <= 16),
            "b ({}) must be larger or equal than 4 and smaller or equal than 16",
            b
        );

        let m = (1 as usize) << b;
        let registers = vec![0; m];
        Self {
            registers,
            b,
            buildhasher,
        }
    }

    /// Get number of bits used for register selection.
    pub fn b(&self) -> usize {
        self.b
    }

    /// Get number of registers.
    pub fn m(&self) -> usize {
        self.registers.len()
    }

    /// Get `BuildHasher`.
    pub fn buildhasher(&self) -> &B {
        &self.buildhasher
    }

    /// Adds an element to the HyperLogLog.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        let w = h >> self.b;
        let j = h - (w << self.b);
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u8);
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 / self.registers
            .iter()
            .map(|&x| 2f64.powi(-(i32::from(x))))
            .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = bytecount::count(&self.registers, 0);
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -(2f64.powi(32)) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }

    /// Merge w/ another HyperLogLog.
    ///
    /// This HyperLogLog will then have the same state as if all elements seen by `other` where
    /// directly added to `self`.
    ///
    /// Panics when `b` or `buildhasher` parameter of `self` and `other` do not match.
    pub fn merge(&mut self, other: &Self) {
        assert_eq!(
            self.b, other.b,
            "b must be equal (left={}, right={})",
            self.b, other.b
        );
        assert!(
            self.buildhasher == other.buildhasher,
            "buildhasher must be equal"
        );

        self.registers = self.registers
            .iter()
            .zip(other.registers.iter())
            .map(|x| cmp::max(x.0, x.1))
            .cloned()
            .collect();
    }

    /// Empties the HyperLogLog.
    pub fn clear(&mut self) {
        self.registers = vec![0; self.registers.len()];
    }

    /// Checks whether the HyperLogLog has never seen an element.
    pub fn is_empty(&self) -> bool {
        self.registers.iter().all(|&x| x == 0)
    }
}

impl fmt::Debug for HyperLogLog {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HyperLogLog {{ b: {} }}", self.b)
    }
}

impl<T> Extend<T> for HyperLogLog
where
    T: Hash,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(&elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::HyperLogLog;
    use hash_utils::BuildHasherSeeded;

    #[test]
    #[should_panic(expected = "b (3) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b3() {
        HyperLogLog::new(3);
    }

    #[test]
    #[should_panic(expected = "b (17) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b17() {
        HyperLogLog::new(17);
    }

    #[test]
    fn getter() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.b(), 8);
        assert_eq!(hll.m(), 1 << 8);
        hll.buildhasher();
    }

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
        assert!(!hll.is_empty());
    }

    #[test]
    fn clear() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        hll.clear();
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn clone() {
        let mut hll1 = HyperLogLog::new(12);
        for i in 0..500 {
            hll1.add(&i);
        }
        let c1a = hll1.count();

        let hll2 = hll1.clone();
        assert_eq!(hll2.count(), c1a);

        for i in 501..1000 {
            hll1.add(&i);
        }
        let c1b = hll1.count();
        assert_ne!(c1b, c1a);
        assert_eq!(hll2.count(), c1a);
    }

    #[test]
    fn merge() {
        let mut hll1 = HyperLogLog::new(12);
        let mut hll2 = HyperLogLog::new(12);
        let mut hll = HyperLogLog::new(12);
        for i in 0..500 {
            hll.add(&i);
            hll1.add(&i);
        }
        for i in 501..1000 {
            hll.add(&i);
            hll2.add(&i);
        }
        assert_ne!(hll.count(), hll1.count());
        assert_ne!(hll.count(), hll2.count());

        hll1.merge(&hll2);
        assert_eq!(hll.count(), hll1.count());
    }

    #[test]
    #[should_panic(expected = "b must be equal (left=5, right=12)")]
    fn merge_panics_p() {
        let mut hll1 = HyperLogLog::new(5);
        let hll2 = HyperLogLog::new(12);
        hll1.merge(&hll2);
    }

    #[test]
    #[should_panic(expected = "buildhasher must be equal")]
    fn merge_panics_buildhasher() {
        let mut hll1 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(0));
        let hll2 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(1));
        hll1.merge(&hll2);
    }

    #[test]
    fn debug() {
        let hll = HyperLogLog::new(12);
        assert_eq!(format!("{:?}", hll), "HyperLogLog { b: 12 }");
    }

    #[test]
    fn extend() {
        let mut hll = HyperLogLog::new(4);
        hll.extend(0..1000);
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }
}
<commit_msg>hyperloglog: improve docs<commit_after>//! HyperLogLog implementation.
use bytecount;
use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{BuildHasher, Hash, Hasher};

use hash_utils::MyBuildHasherDefault;

/// A HyperLogLog is a data structure to count unique elements on a data stream.
///
/// # Examples
/// ```
/// use pdatastructs::hyperloglog::HyperLogLog;
///
/// // set up filter
/// let address_bits = 4;  // so we store 2^4 = 16 registers in total
/// let mut hll = HyperLogLog::new(address_bits);
///
/// // add some data
/// hll.add(&"my super long string");
/// hll.add(&"another super long string");
/// hll.add(&"my super long string");  // again
///
/// // later
/// assert_eq!(hll.count(), 2);
/// ```
///
/// # Applications
/// - an approximative `COUNT(DISTINCT x)` in SQL
/// - count distinct elements in a data stream
///
/// # How It Works
/// The HyperLogLog consists of `2^b` 8bit counters. Each counter is initialized to 0.
///
/// During insertion, a hash `h(x)` is calculated. The first `b` bits of the hash function are used
/// to address a register, the other bits are used to create a number `p` which essentially counts
/// the number of leading 0-bits (or in other words: the leftmost 1-bit). The addressed register is
/// then updated to the maximum of its current value and `p`.
///
/// The calculation of the count is based on `1 / Sum_0^{2^b} (2^-register_i)` with a bunch of
/// factors a corrections applied (see paper or source code).
///
/// # Implementation
/// In contrast to the official HyperLogLog, a 64 bit instead of a 32 bit hash function is used. The
/// registers always allocate 8 bits and are not compressed.
///
/// # See Also
/// - `std::collections::HashSet`: can be used to get the exact count but requires you to store
///   each and every element
/// - `pdatastructs::bloomfilter::BloomFilter` and `pdatastructs::cuckoofilter::CuckooFilter`: when
///   you also want to check if a single element is in the observed set
///
/// # References
/// - ["HyperLogLog: the analysis of a near-optimal cardinality estimation algorithm", Philippe
///   Flajolet, Éric Fusy, Olivier Gandouet, Frédéric Meunier, 2007](http://dmod.eu/deca/ft_gateway.cfm.pdf)
/// - [Wikipedia: HyperLogLog](https://en.wikipedia.org/wiki/HyperLogLog)
#[derive(Clone)]
pub struct HyperLogLog<B = MyBuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    registers: Vec<u8>,
    b: usize,
    buildhasher: B,
}

impl HyperLogLog {
    /// Creates a new, empty HyperLogLog.
    ///
    /// - `b` number of bits used for register selection, number of registers within the
    ///   HyperLogLog will be `2^b`. `b` must be in `[4, 16]`
    ///
    /// Panics when `b` is out of bounds.
    pub fn new(b: usize) -> Self {
        let bh = MyBuildHasherDefault::<DefaultHasher>::default();
        Self::with_hash(b, bh)
    }
}

impl<B> HyperLogLog<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Same as `new` but with a specific `BuildHasher`.
    pub fn with_hash(b: usize, buildhasher: B) -> Self {
        assert!(
            (b >= 4) & (b <= 16),
            "b ({}) must be larger or equal than 4 and smaller or equal than 16",
            b
        );

        let m = (1 as usize) << b;
        let registers = vec![0; m];
        Self {
            registers,
            b,
            buildhasher,
        }
    }

    /// Get number of bits used for register selection.
    pub fn b(&self) -> usize {
        self.b
    }

    /// Get number of registers.
    pub fn m(&self) -> usize {
        self.registers.len()
    }

    /// Get `BuildHasher`.
    pub fn buildhasher(&self) -> &B {
        &self.buildhasher
    }

    /// Adds an element to the HyperLogLog.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        // split h into:
        //  - w = 64 - b upper bits
        //  - j = b lower bits
        let w = h >> self.b;
        let j = h - (w << self.b);  // no 1 as in the paper since register indices are 0-based

        // p = leftmost bit (1-based count)
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u8);
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 / self.registers
            .iter()
            .map(|&x| 2f64.powi(-(i32::from(x))))
            .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = bytecount::count(&self.registers, 0);
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -(2f64.powi(32)) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }

    /// Merge w/ another HyperLogLog.
    ///
    /// This HyperLogLog will then have the same state as if all elements seen by `other` where
    /// directly added to `self`.
    ///
    /// Panics when `b` or `buildhasher` parameter of `self` and `other` do not match.
    pub fn merge(&mut self, other: &Self) {
        assert_eq!(
            self.b, other.b,
            "b must be equal (left={}, right={})",
            self.b, other.b
        );
        assert!(
            self.buildhasher == other.buildhasher,
            "buildhasher must be equal"
        );

        self.registers = self.registers
            .iter()
            .zip(other.registers.iter())
            .map(|x| cmp::max(x.0, x.1))
            .cloned()
            .collect();
    }

    /// Empties the HyperLogLog.
    pub fn clear(&mut self) {
        self.registers = vec![0; self.registers.len()];
    }

    /// Checks whether the HyperLogLog has never seen an element.
    pub fn is_empty(&self) -> bool {
        self.registers.iter().all(|&x| x == 0)
    }
}

impl fmt::Debug for HyperLogLog {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HyperLogLog {{ b: {} }}", self.b)
    }
}

impl<T> Extend<T> for HyperLogLog
where
    T: Hash,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(&elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::HyperLogLog;
    use hash_utils::BuildHasherSeeded;

    #[test]
    #[should_panic(expected = "b (3) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b3() {
        HyperLogLog::new(3);
    }

    #[test]
    #[should_panic(expected = "b (17) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b17() {
        HyperLogLog::new(17);
    }

    #[test]
    fn getter() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.b(), 8);
        assert_eq!(hll.m(), 1 << 8);
        hll.buildhasher();
    }

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
        assert!(!hll.is_empty());
    }

    #[test]
    fn clear() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        hll.clear();
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn clone() {
        let mut hll1 = HyperLogLog::new(12);
        for i in 0..500 {
            hll1.add(&i);
        }
        let c1a = hll1.count();

        let hll2 = hll1.clone();
        assert_eq!(hll2.count(), c1a);

        for i in 501..1000 {
            hll1.add(&i);
        }
        let c1b = hll1.count();
        assert_ne!(c1b, c1a);
        assert_eq!(hll2.count(), c1a);
    }

    #[test]
    fn merge() {
        let mut hll1 = HyperLogLog::new(12);
        let mut hll2 = HyperLogLog::new(12);
        let mut hll = HyperLogLog::new(12);
        for i in 0..500 {
            hll.add(&i);
            hll1.add(&i);
        }
        for i in 501..1000 {
            hll.add(&i);
            hll2.add(&i);
        }
        assert_ne!(hll.count(), hll1.count());
        assert_ne!(hll.count(), hll2.count());

        hll1.merge(&hll2);
        assert_eq!(hll.count(), hll1.count());
    }

    #[test]
    #[should_panic(expected = "b must be equal (left=5, right=12)")]
    fn merge_panics_p() {
        let mut hll1 = HyperLogLog::new(5);
        let hll2 = HyperLogLog::new(12);
        hll1.merge(&hll2);
    }

    #[test]
    #[should_panic(expected = "buildhasher must be equal")]
    fn merge_panics_buildhasher() {
        let mut hll1 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(0));
        let hll2 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(1));
        hll1.merge(&hll2);
    }

    #[test]
    fn debug() {
        let hll = HyperLogLog::new(12);
        assert_eq!(format!("{:?}", hll), "HyperLogLog { b: 12 }");
    }

    #[test]
    fn extend() {
        let mut hll = HyperLogLog::new(4);
        hll.extend(0..1000);
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Read, Cursor};
use std::mem;
use message::Message;
use piece_field::PieceField;
use byteorder::{BigEndian, ReadBytesExt};

pub struct Reader {
    state: ReadState,
    blocks_read: usize,
    download_speed: f64,
    can_receive_bf: bool,
    received_handshake: bool,
}

enum ReadState {
    Idle,
    ReadingHandshake { data: [u8; 68], idx: u8 },
    ReadingLen { data: [u8; 17], idx: u8 },
    ReadingId { data: [u8; 17], len: u32 },
    ReadingMsg { data: [u8; 17], idx: u8, len: u32 },
    ReadingPiece { prefix: [u8; 17], data: Box<[u8; 16384]>, idx: usize },
    ReadingBitfield { data: Vec<u8>, idx: usize },
}

impl ReadState {
    fn next_state<R: Read>(self, conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        // I don't think this could feasibly stack overflow, but possibility should be considered.
        match self {
            ReadState::ReadingHandshake { mut data, mut idx } => {
                idx += conn.read(&mut data[idx as usize..])? as u8;
                if idx == data.len() as u8 - 1 {
                    if &data[1..20] != b"BitTorrent protocol" {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid protocol used in handshake"));
                    }
                    let mut reserved = [0; 8];
                    reserved.clone_from_slice(&data[20..28]);
                    let mut hash = [0; 20];
                    hash.clone_from_slice(&data[28..48]);
                    let mut pid = [0; 20];
                    pid.clone_from_slice(&data[48..68]);
                    Ok(Ok(Message::Handshake(reserved, hash, pid)))
                } else {
                    Ok(Err(ReadState::ReadingHandshake { data: data, idx: idx }))
                }
            }
            ReadState::Idle => {
                let mut data = [0; 17];
                let amnt = conn.read(&mut data[0..4])? as u8;
                if amnt != 4 {
                    Ok(Err(ReadState::ReadingLen { data: data, idx: amnt }))
                } else {
                    ReadState::process_len(data, conn)
                }
            }
            ReadState::ReadingLen { mut data, mut idx } => {
                idx += conn.read(&mut data[(idx as usize)..4])? as u8;
                if idx != 4 {
                    Ok(Err(ReadState::ReadingLen { data: data, idx: idx }))
                } else {
                    ReadState::process_len(data, conn)
                }
            }
            ReadState::ReadingId { mut data, len } => {
                let amnt = conn.read(&mut data[4..5])? as u8;
                if amnt != 1 {
                    Ok(Err(ReadState::ReadingId { data: data, len: len }))
                } else {
                    ReadState::process_id(data, len, conn)
                }
            }
            ReadState::ReadingMsg { mut data, mut idx, len } => {
                idx += conn.read(&mut data[(idx as usize)..(len + 4) as usize])? as u8;
                if idx - 4 == len as u8 {
                    match ReadState::process_message(data, len) {
                        Ok(msg) => Ok(Ok(msg)),
                        Err(e) => Err(e),
                    }
                } else {
                    Ok(Err(ReadState::ReadingMsg { data: data, idx: idx, len: len }))
                }
            }
            ReadState::ReadingPiece { mut prefix, mut data, mut idx } => {
                if idx < 13 {
                    idx += conn.read(&mut prefix[idx as usize..13])?;
                    if idx == 13 {
                        ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }.next_state(conn)
                    } else {
                        Ok(Err(ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }))
                    }
                } else {
                    idx += conn.read(&mut data[(idx - 13)..])?;
                    if idx - 13 == data.len() {
                        let idx = (&prefix[5..9]).read_u32::<BigEndian>()?;
                        let beg = (&prefix[9..13]).read_u32::<BigEndian>()?;
                        Ok(Ok(Message::Piece(idx, beg, data)))
                    } else {
                        Ok(Err(ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }))
                    }
                }
            }
            ReadState::ReadingBitfield { mut data, mut idx } => {
                idx += conn.read(&mut data[idx as usize..])?;
                let len = data.len();
                if idx == len {
                    Ok(Ok(Message::Bitfield(PieceField::from(data.into_boxed_slice(), len as u32 * 8))))
                } else {
                    Ok(Err(ReadState::ReadingBitfield { data: data, idx: idx }))
                }
            }
        }
    }

    fn process_len<R: Read>(buf: [u8; 17], conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        let len = (&buf[0..4]).read_u32::<BigEndian>()?;
        if len == 0 {
            Ok(Ok(Message::KeepAlive))
        } else {
            ReadState::ReadingId { data: buf, len: len }.next_state(conn)
        }
    }

    fn process_id<R: Read>(buf: [u8; 17], len: u32, conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        let id = buf[4];
        match id {
            0 => Ok(Ok(Message::Choke)),
            1 => Ok(Ok(Message::Unchoke)),
            2 => Ok(Ok(Message::Interested)),
            3 => Ok(Ok(Message::Uninterested)),
            5 => {
                ReadState::ReadingBitfield { data: vec![0; len as usize - 1], idx: 0 }.next_state(conn)
            },
            7 => {
                if len != 16393 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Only piece sizes of 16384 are accepted"));
                }
                ReadState::ReadingPiece { prefix: buf, data: Box::new([0u8; 16384]), idx: len as usize}.next_state(conn)
            }
            _ => {
                ReadState::ReadingMsg { data: buf, idx: 5, len: len }.next_state(conn)
            },
        }
    }

    fn process_message(buf: [u8; 17], len: u32) -> io::Result<Message> {
        match buf[4] {
            4 => {
                if len != 5 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Have message must be of len 5"));
                }
                Ok(Message::Have((&buf[5..9]).read_u32::<BigEndian>()?))
            }
            6 => {
                if len != 13 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Request message must be of len 13"));
                }
                let idx = (&buf[5..9]).read_u32::<BigEndian>()?;
                let beg = (&buf[9..13]).read_u32::<BigEndian>()?;
                let len = (&buf[13..17]).read_u32::<BigEndian>()?;
                Ok(Message::Request(idx, beg, len))
            }
            8 => {
                if len != 13 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Cancel message must be of len 13"));
                }
                let idx = (&buf[5..9]).read_u32::<BigEndian>()?;
                let beg = (&buf[9..13]).read_u32::<BigEndian>()?;
                let len = (&buf[13..17]).read_u32::<BigEndian>()?;
                Ok(Message::Cancel(idx, beg, len))
            }
            _ => {
                Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid message ID"))
            }
        }
    }
}

impl Reader {
    pub fn new() -> Reader {
        Reader {
            state: ReadState::ReadingHandshake { data: [0u8; 68] , idx: 0 },
            blocks_read: 0,
            download_speed: 0.0,
            can_receive_bf: true,
            received_handshake: false,
        }
    }

    /// Attempts to read a single message from the connection
    pub fn readable<R: Read>(&mut self, conn: &mut R) -> io::Result<Option<Message>> {
        let state = mem::replace(&mut self.state, ReadState::Idle);
        match state.next_state(conn)? {
            Ok(msg) => {
                if !msg.is_handshake() && !self.received_handshake {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Must receive handshake first!"));
                } else if msg.is_handshake() {
                    self.received_handshake = true;
                }
                if msg.is_bitfield() {
                    if !self.can_receive_bf {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "Bitfield cannot be received twice!"));
                    } else {
                        self.can_receive_bf = false;
                    }
                } else if self.received_handshake {
                    self.can_receive_bf = false;
                }
                self.state = ReadState::Idle;
                Ok(Some(msg))
            }
            Err(new_state) => {
                self.state = new_state;
                Ok(None)
            }
        }
    }
}

#[test]
fn test_read_keepalive() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 0];
    // Test one shot
    if let Message::KeepAlive = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }

    // Test split up
    if let None = r.readable(&mut &data[0..2]).unwrap() {
    } else {
        unreachable!();
    }
    if let Message::KeepAlive = r.readable(&mut &data[2..4]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_choke() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 0];
    // Test one shot
    if let Message::Choke = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }

    // Test split up
    if let None = r.readable(&mut &data[0..4]).unwrap() {
    } else {
        unreachable!();
    }
    // Simulate spurious read
    if let None = r.readable(&mut &data[4..4]).unwrap() {
    } else {
        unreachable!();
    }
    if let Message::Choke = r.readable(&mut &data[4..5]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_unchoke() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 1];
    // Test one shot
    if let Message::Unchoke = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_interested() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 2];
    // Test one shot
    if let Message::Interested = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_uninterested() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 3];
    // Test one shot
    if let Message::Uninterested = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_have() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 5, 4, 0, 0, 0, 1];
    // Test one shot
    match r.readable(&mut &data[..]).unwrap().unwrap() {
        Message::Have(piece) => {
            if piece != 1 {
                unreachable!();
            }
        }
        _ => {
            unreachable!();
        }
    }

    // Test split up
    if let None = r.readable(&mut &data[0..6]).unwrap() {
    } else {
        unreachable!();
    }
    // Simulate spurious read
    if let None = r.readable(&mut &data[6..6]).unwrap() {
    } else {
        unreachable!();
    }
    match r.readable(&mut &data[6..9]).unwrap().unwrap() {
        Message::Have(piece) => {
            if piece != 1 {
                unreachable!();
            }
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_bitfield() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = true;
    r.received_handshake = true;
    let mut data = Cursor::new(vec![0u8, 0, 0, 5, 5, 0xff, 0xff, 0xff, 0xff]);
    // Test one shot
    match r.readable(&mut data).unwrap().unwrap() {
        Message::Bitfield(pf) => {
            for i in 0..32 {
                assert!(pf.has_piece(i as u32));
            }
        }
        _ => {
            unreachable!();
        }
    }
}
<commit_msg>Complete test suite for protocol parsing<commit_after>use std::io::{self, Read, Cursor};
use std::mem;
use message::Message;
use piece_field::PieceField;
use byteorder::{BigEndian, ReadBytesExt};

pub struct Reader {
    state: ReadState,
    blocks_read: usize,
    download_speed: f64,
    can_receive_bf: bool,
    received_handshake: bool,
}

enum ReadState {
    Idle,
    ReadingHandshake { data: [u8; 68], idx: u8 },
    ReadingLen { data: [u8; 17], idx: u8 },
    ReadingId { data: [u8; 17], len: u32 },
    ReadingMsg { data: [u8; 17], idx: u8, len: u32 },
    ReadingPiece { prefix: [u8; 17], data: Box<[u8; 16384]>, idx: usize },
    ReadingBitfield { data: Vec<u8>, idx: usize },
}

impl ReadState {
    fn next_state<R: Read>(self, conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        // I don't think this could feasibly stack overflow, but possibility should be considered.
        match self {
            ReadState::ReadingHandshake { mut data, mut idx } => {
                idx += conn.read(&mut data[idx as usize..])? as u8;
                if idx == data.len() as u8 {
                    if &data[1..20] != b"BitTorrent protocol" {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid protocol used in handshake"));
                    }
                    let mut reserved = [0; 8];
                    reserved.clone_from_slice(&data[20..28]);
                    let mut hash = [0; 20];
                    hash.clone_from_slice(&data[28..48]);
                    let mut pid = [0; 20];
                    pid.clone_from_slice(&data[48..68]);
                    Ok(Ok(Message::Handshake(reserved, hash, pid)))
                } else {
                    Ok(Err(ReadState::ReadingHandshake { data: data, idx: idx }))
                }
            }
            ReadState::Idle => {
                let mut data = [0; 17];
                let amnt = conn.read(&mut data[0..4])? as u8;
                if amnt != 4 {
                    Ok(Err(ReadState::ReadingLen { data: data, idx: amnt }))
                } else {
                    ReadState::process_len(data, conn)
                }
            }
            ReadState::ReadingLen { mut data, mut idx } => {
                idx += conn.read(&mut data[(idx as usize)..4])? as u8;
                if idx != 4 {
                    Ok(Err(ReadState::ReadingLen { data: data, idx: idx }))
                } else {
                    ReadState::process_len(data, conn)
                }
            }
            ReadState::ReadingId { mut data, len } => {
                let amnt = conn.read(&mut data[4..5])? as u8;
                if amnt != 1 {
                    Ok(Err(ReadState::ReadingId { data: data, len: len }))
                } else {
                    ReadState::process_id(data, len, conn)
                }
            }
            ReadState::ReadingMsg { mut data, mut idx, len } => {
                idx += conn.read(&mut data[(idx as usize)..(len + 4) as usize])? as u8;
                if idx - 4 == len as u8 {
                    match ReadState::process_message(data, len) {
                        Ok(msg) => Ok(Ok(msg)),
                        Err(e) => Err(e),
                    }
                } else {
                    Ok(Err(ReadState::ReadingMsg { data: data, idx: idx, len: len }))
                }
            }
            ReadState::ReadingPiece { mut prefix, mut data, mut idx } => {
                if idx < 13 {
                    idx += conn.read(&mut prefix[idx as usize..13])?;
                    if idx == 13 {
                        ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }.next_state(conn)
                    } else {
                        Ok(Err(ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }))
                    }
                } else {
                    idx += conn.read(&mut data[(idx - 13)..])?;
                    if idx - 13 == data.len() {
                        let idx = (&prefix[5..9]).read_u32::<BigEndian>()?;
                        let beg = (&prefix[9..13]).read_u32::<BigEndian>()?;
                        Ok(Ok(Message::Piece(idx, beg, data)))
                    } else {
                        Ok(Err(ReadState::ReadingPiece { prefix: prefix, data: data, idx: idx }))
                    }
                }
            }
            ReadState::ReadingBitfield { mut data, mut idx } => {
                idx += conn.read(&mut data[idx as usize..])?;
                let len = data.len();
                if idx == len {
                    Ok(Ok(Message::Bitfield(PieceField::from(data.into_boxed_slice(), len as u32 * 8))))
                } else {
                    Ok(Err(ReadState::ReadingBitfield { data: data, idx: idx }))
                }
            }
        }
    }

    fn process_len<R: Read>(buf: [u8; 17], conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        let len = (&buf[0..4]).read_u32::<BigEndian>()?;
        if len == 0 {
            Ok(Ok(Message::KeepAlive))
        } else {
            ReadState::ReadingId { data: buf, len: len }.next_state(conn)
        }
    }

    fn process_id<R: Read>(buf: [u8; 17], len: u32, conn: &mut R) -> io::Result<Result<Message, ReadState>> {
        let id = buf[4];
        match id {
            0 => Ok(Ok(Message::Choke)),
            1 => Ok(Ok(Message::Unchoke)),
            2 => Ok(Ok(Message::Interested)),
            3 => Ok(Ok(Message::Uninterested)),
            5 => {
                ReadState::ReadingBitfield { data: vec![0; len as usize - 1], idx: 0 }.next_state(conn)
            },
            7 => {
                if len != 16393 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Only piece sizes of 16384 are accepted"));
                }
                ReadState::ReadingPiece { prefix: buf, data: Box::new([0u8; 16384]), idx: 5}.next_state(conn)
            }
            _ => {
                ReadState::ReadingMsg { data: buf, idx: 5, len: len }.next_state(conn)
            },
        }
    }

    fn process_message(buf: [u8; 17], len: u32) -> io::Result<Message> {
        match buf[4] {
            4 => {
                if len != 5 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Have message must be of len 5"));
                }
                Ok(Message::Have((&buf[5..9]).read_u32::<BigEndian>()?))
            }
            6 => {
                if len != 13 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Request message must be of len 13"));
                }
                let idx = (&buf[5..9]).read_u32::<BigEndian>()?;
                let beg = (&buf[9..13]).read_u32::<BigEndian>()?;
                let len = (&buf[13..17]).read_u32::<BigEndian>()?;
                Ok(Message::Request(idx, beg, len))
            }
            8 => {
                if len != 13 {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Cancel message must be of len 13"));
                }
                let idx = (&buf[5..9]).read_u32::<BigEndian>()?;
                let beg = (&buf[9..13]).read_u32::<BigEndian>()?;
                let len = (&buf[13..17]).read_u32::<BigEndian>()?;
                Ok(Message::Cancel(idx, beg, len))
            }
            _ => {
                Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid message ID"))
            }
        }
    }
}

impl Reader {
    pub fn new() -> Reader {
        Reader {
            state: ReadState::ReadingHandshake { data: [0u8; 68] , idx: 0 },
            blocks_read: 0,
            download_speed: 0.0,
            can_receive_bf: true,
            received_handshake: false,
        }
    }

    /// Attempts to read a single message from the connection
    pub fn readable<R: Read>(&mut self, conn: &mut R) -> io::Result<Option<Message>> {
        let state = mem::replace(&mut self.state, ReadState::Idle);
        match state.next_state(conn)? {
            Ok(msg) => {
                if !msg.is_handshake() && !self.received_handshake {
                    return Err(io::Error::new(io::ErrorKind::InvalidData, "Must receive handshake first!"));
                } else if msg.is_handshake() {
                    self.received_handshake = true;
                }
                if msg.is_bitfield() {
                    if !self.can_receive_bf {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "Bitfield cannot be received twice!"));
                    } else {
                        self.can_receive_bf = false;
                    }
                } else if self.received_handshake {
                    self.can_receive_bf = false;
                }
                self.state = ReadState::Idle;
                Ok(Some(msg))
            }
            Err(new_state) => {
                self.state = new_state;
                Ok(None)
            }
        }
    }
}

#[test]
fn test_read_keepalive() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 0];
    // Test one shot
    if let Message::KeepAlive = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }

    // Test split up
    if let None = r.readable(&mut &data[0..2]).unwrap() {
    } else {
        unreachable!();
    }
    if let Message::KeepAlive = r.readable(&mut &data[2..4]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_choke() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 0];
    // Test one shot
    if let Message::Choke = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }

    // Test split up
    if let None = r.readable(&mut &data[0..4]).unwrap() {
    } else {
        unreachable!();
    }
    // Simulate spurious read
    if let None = r.readable(&mut &data[4..4]).unwrap() {
    } else {
        unreachable!();
    }
    if let Message::Choke = r.readable(&mut &data[4..5]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_unchoke() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 1];
    // Test one shot
    if let Message::Unchoke = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_interested() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 2];
    // Test one shot
    if let Message::Interested = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_uninterested() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 1, 3];
    // Test one shot
    if let Message::Uninterested = r.readable(&mut &data[..]).unwrap().unwrap() {
    } else {
        unreachable!();
    }
}

#[test]
fn test_read_have() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = false;
    r.received_handshake = true;
    let data = vec![0u8, 0, 0, 5, 4, 0, 0, 0, 1];
    // Test one shot
    match r.readable(&mut &data[..]).unwrap().unwrap() {
        Message::Have(piece) => {
            if piece != 1 {
                unreachable!();
            }
        }
        _ => {
            unreachable!();
        }
    }

    // Test split up
    if let None = r.readable(&mut &data[0..6]).unwrap() {
    } else {
        unreachable!();
    }
    // Simulate spurious read
    if let None = r.readable(&mut &data[6..6]).unwrap() {
    } else {
        unreachable!();
    }
    match r.readable(&mut &data[6..9]).unwrap().unwrap() {
        Message::Have(piece) => {
            if piece != 1 {
                unreachable!();
            }
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_bitfield() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = true;
    r.received_handshake = true;
    let mut data = Cursor::new(vec![0u8, 0, 0, 5, 5, 0xff, 0xff, 0xff, 0xff]);
    // Test one shot
    match r.readable(&mut data).unwrap().unwrap() {
        Message::Bitfield(pf) => {
            for i in 0..32 {
                assert!(pf.has_piece(i as u32));
            }
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_request() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = true;
    r.received_handshake = true;
    let mut data = Cursor::new(vec![0u8, 0, 0, 13, 6, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]);
    // Test one shot
    match r.readable(&mut data).unwrap().unwrap() {
        Message::Request(i, b, l) => {
            assert_eq!(i, 1);
            assert_eq!(b, 1);
            assert_eq!(l, 1);
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_piece() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = true;
    r.received_handshake = true;
    let mut info = Cursor::new(vec![0u8, 0, 0x40, 0x09, 7, 0, 0, 0, 1, 0, 0, 0, 1]);
    let mut data = Cursor::new(vec![1u8; 16384]);
    // Test partial read
    assert!(r.readable(&mut info).unwrap().is_none());
    match r.readable(&mut data).unwrap().unwrap() {
        Message::Piece(i, b, d) => {
            assert_eq!(i, 1);
            assert_eq!(b, 1);
            for i in 0..16384 {
                assert_eq!(1, d[i]);
            }
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_cancel() {
    let mut r = Reader::new();
    r.state = ReadState::Idle;
    r.can_receive_bf = true;
    r.received_handshake = true;
    let mut data = Cursor::new(vec![0u8, 0, 0, 13, 8, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]);
    // Test one shot
    match r.readable(&mut data).unwrap().unwrap() {
        Message::Cancel(i, b, l) => {
            assert_eq!(i, 1);
            assert_eq!(b, 1);
            assert_eq!(l, 1);
        }
        _ => {
            unreachable!();
        }
    }
}

#[test]
fn test_read_handshake() {
    use ::PEER_ID;
    let mut r = Reader::new();
    let m = Message::Handshake([0; 8], [0; 20], *PEER_ID);
    let mut data = vec![0; 68];
    m.encode(&mut data[..]);
    // Test one shot
    match r.readable(&mut Cursor::new(data)).unwrap().unwrap() {
        Message::Handshake(r, hash, pid) => {
            assert_eq!(r, [0; 8]);
            assert_eq!(hash, [0; 20]);
            assert_eq!(pid, *PEER_ID);
        }
        _ => {
            unreachable!();
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(all(test, feature = "nightly_test"))]
extern crate compiletest_rs as compiletest;

#[cfg(all(test, feature = "nightly_test"))]
fn run_mode(mode: &'static str) {
    use std::path::PathBuf;

    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from(format!("tests/{}", mode));
    config.target_rustcflags = Some("-L target/debug".to_string());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(all(test, feature = "nightly_test"))]
fn compile_test() {
    run_mode("compile-fail");
}
<commit_msg>Compile-fail tests: Added -L target/debug/deps to rustcflags<commit_after>#[cfg(all(test, feature = "nightly_test"))]
extern crate compiletest_rs as compiletest;

#[cfg(all(test, feature = "nightly_test"))]
fn run_mode(mode: &'static str) {
    use std::path::PathBuf;

    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from(format!("tests/{}", mode));
    config.target_rustcflags = Some("-L target/debug -L target/debug/deps".to_string());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(all(test, feature = "nightly_test"))]
fn compile_test() {
    run_mode("compile-fail");
}
<|endoftext|>


<commit_before>use analysis;
use analysis::rust_type::rust_type;
use analysis::type_kind::TypeKind;
use env::Env;
use library;
use traits::*;

pub trait TranslateFromGlib {
    fn translate_from_glib_as_function(&self, env: &Env,
        func: &analysis::functions::Info) -> (String, String);
}

impl TranslateFromGlib for library::Parameter {
    fn translate_from_glib_as_function(&self,
        env: &Env, func: &analysis::functions::Info) -> (String, String) {
        let kind = TypeKind::of(&env.library, self.typ);
        if func.kind == library::FunctionKind::Constructor {
            let rust_type = rust_type(env, self.typ);
            if rust_type.as_str() != func.class_name.as_str() {
                let from_glib_xxx = from_glib_xxx(self.transfer);
                (
                    format!("{}::{}", rust_type.as_str(), from_glib_xxx.0),
                    format!("{}.downcast_unchecked()",from_glib_xxx.1)
                )
            } else {
                ("TODO:constructors_body ".into(), String::new())
            }
        } else {
            match kind {
                TypeKind::Converted => ("from_glib(".into(), ")".into()),
                TypeKind::Direct |
                    TypeKind::Enumeration => (String::new(), String::new()),
                //TODO: check gtk_dialog_get_content_area <type name="Box" c:type="GtkWidget*"/>
                TypeKind::Object => from_glib_xxx(self.transfer),
                _ => ("TODO:".into(), String::new()),
            }
        }
    }
}

impl TranslateFromGlib for Option<library::Parameter> {
    fn translate_from_glib_as_function(&self, env: &Env,
        func: &analysis::functions::Info) -> (String, String) {
        match self {
            &Some(ref par) => par.translate_from_glib_as_function(env, func),
            &None => (String::new(), ";".into())
        }
    }
}

fn from_glib_xxx(transfer: library::Transfer) -> (String, String) {
    use library::Transfer::*;
    match transfer {
        None => ("from_glib_none(".into(), ")".into()),
        Full => ("from_glib_full(".into(), ")".into()),
        Container => ("TODO:".into(), String::new()),
    }
}
<commit_msg>Expanded generated TODO<commit_after>use analysis;
use analysis::rust_type::rust_type;
use analysis::type_kind::TypeKind;
use env::Env;
use library;
use traits::*;

pub trait TranslateFromGlib {
    fn translate_from_glib_as_function(&self, env: &Env,
        func: &analysis::functions::Info) -> (String, String);
}

impl TranslateFromGlib for library::Parameter {
    fn translate_from_glib_as_function(&self,
        env: &Env, func: &analysis::functions::Info) -> (String, String) {
        let kind = TypeKind::of(&env.library, self.typ);
        if func.kind == library::FunctionKind::Constructor {
            let rust_type = rust_type(env, self.typ);
            if rust_type.as_str() != func.class_name.as_str() {
                let from_glib_xxx = from_glib_xxx(self.transfer);
                (
                    format!("{}::{}", rust_type.as_str(), from_glib_xxx.0),
                    format!("{}.downcast_unchecked()",from_glib_xxx.1)
                )
            } else {
                ("TODO:constructors_body ".into(), String::new())
            }
        } else {
            match kind {
                TypeKind::Converted => ("from_glib(".into(), ")".into()),
                TypeKind::Direct |
                    TypeKind::Enumeration => (String::new(), String::new()),
                //TODO: check gtk_dialog_get_content_area <type name="Box" c:type="GtkWidget*"/>
                TypeKind::Object => from_glib_xxx(self.transfer),
                _ => (format!("TODO {:?}:", kind), String::new()),
            }
        }
    }
}

impl TranslateFromGlib for Option<library::Parameter> {
    fn translate_from_glib_as_function(&self, env: &Env,
        func: &analysis::functions::Info) -> (String, String) {
        match self {
            &Some(ref par) => par.translate_from_glib_as_function(env, func),
            &None => (String::new(), ";".into())
        }
    }
}

fn from_glib_xxx(transfer: library::Transfer) -> (String, String) {
    use library::Transfer::*;
    match transfer {
        None => ("from_glib_none(".into(), ")".into()),
        Full => ("from_glib_full(".into(), ")".into()),
        Container => ("TODO:".into(), String::new()),
    }
}
<|endoftext|>