<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Chirag B Jadwani <chirag.jadwani@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

use clap::{crate_version, Arg, ArgMatches, Command};
use std::fs::File;
use std::io::{self, stdin, stdout, BufRead, BufReader, BufWriter, Read, Write};
use std::path::Path;
use std::str::FromStr;
use strum_macros::{AsRefStr, EnumString};
use uucore::display::Quotable;
use uucore::error::{FromIo, UResult, USimpleError};
use uucore::format_usage;

static ABOUT: &str = "Report or omit repeated lines.";
const USAGE: &str = "{} [OPTION]... [INPUT [OUTPUT]]...";
pub mod options {
    pub static ALL_REPEATED: &str = "all-repeated";
    pub static CHECK_CHARS: &str = "check-chars";
    pub static COUNT: &str = "count";
    pub static IGNORE_CASE: &str = "ignore-case";
    pub static REPEATED: &str = "repeated";
    pub static SKIP_FIELDS: &str = "skip-fields";
    pub static SKIP_CHARS: &str = "skip-chars";
    pub static UNIQUE: &str = "unique";
    pub static ZERO_TERMINATED: &str = "zero-terminated";
    pub static GROUP: &str = "group";
}

static ARG_FILES: &str = "files";

#[derive(PartialEq, Clone, Copy, AsRefStr, EnumString)]
#[strum(serialize_all = "snake_case")]
enum Delimiters {
    Append,
    Prepend,
    Separate,
    Both,
    None,
}

struct Uniq {
    repeats_only: bool,
    uniques_only: bool,
    all_repeated: bool,
    delimiters: Delimiters,
    show_counts: bool,
    skip_fields: Option<usize>,
    slice_start: Option<usize>,
    slice_stop: Option<usize>,
    ignore_case: bool,
    zero_terminated: bool,
}

macro_rules! write_line_terminator {
    ($writer:expr, $line_terminator:expr) => {
        $writer
            .write_all(&[$line_terminator])
            .map_err_context(|| "Could not write line terminator".to_string())
    };
}

impl Uniq {
    pub fn print_uniq<R: Read, W: Write>(
        &self,
        reader: &mut BufReader<R>,
        writer: &mut BufWriter<W>,
    ) -> UResult<()> {
        let mut first_line_printed = false;
        let mut group_count = 1;
        let line_terminator = self.get_line_terminator();
        let mut lines = reader.split(line_terminator).map(get_line_string);
        let mut line = match lines.next() {
            Some(l) => l?,
            None => return Ok(()),
        };

        // compare current `line` with consecutive lines (`next_line`) of the input
        // and if needed, print `line` based on the command line options provided
        for next_line in lines {
            let next_line = next_line?;
            if self.cmp_keys(&line, &next_line) {
                if (group_count == 1 && !self.repeats_only)
                    || (group_count > 1 && !self.uniques_only)
                {
                    self.print_line(writer, &line, group_count, first_line_printed)?;
                    first_line_printed = true;
                }
                line = next_line;
                group_count = 1;
            } else {
                if self.all_repeated {
                    self.print_line(writer, &line, group_count, first_line_printed)?;
                    first_line_printed = true;
                    line = next_line;
                }
                group_count += 1;
            }
        }
        if (group_count == 1 && !self.repeats_only) || (group_count > 1 && !self.uniques_only) {
            self.print_line(writer, &line, group_count, first_line_printed)?;
            first_line_printed = true;
        }
        if (self.delimiters == Delimiters::Append || self.delimiters == Delimiters::Both)
            && first_line_printed
        {
            write_line_terminator!(writer, line_terminator)?;
        }
        Ok(())
    }

    fn skip_fields<'a>(&self, line: &'a str) -> &'a str {
        if let Some(skip_fields) = self.skip_fields {
            let mut i = 0;
            let mut char_indices = line.char_indices();
            for _ in 0..skip_fields {
                if char_indices.find(|(_, c)| !c.is_whitespace()) == None {
                    return "";
                }
                match char_indices.find(|(_, c)| c.is_whitespace()) {
                    None => return "",

                    Some((next_field_i, _)) => i = next_field_i,
                }
            }
            &line[i..]
        } else {
            line
        }
    }

    fn get_line_terminator(&self) -> u8 {
        if self.zero_terminated {
            0
        } else {
            b'\n'
        }
    }

    fn cmp_keys(&self, first: &str, second: &str) -> bool {
        self.cmp_key(first, |first_iter| {
            self.cmp_key(second, |second_iter| first_iter.ne(second_iter))
        })
    }

    fn cmp_key<F>(&self, line: &str, mut closure: F) -> bool
    where
        F: FnMut(&mut dyn Iterator<Item = char>) -> bool,
    {
        let fields_to_check = self.skip_fields(line);
        let len = fields_to_check.len();
        let slice_start = self.slice_start.unwrap_or(0);
        let slice_stop = self.slice_stop.unwrap_or(len);
        if len > 0 {
            // fast path: avoid doing any work if there is no need to skip or map to lower-case
            if !self.ignore_case && slice_start == 0 && slice_stop == len {
                return closure(&mut fields_to_check.chars());
            }

            // fast path: avoid skipping
            if self.ignore_case && slice_start == 0 && slice_stop == len {
                return closure(&mut fields_to_check.chars().flat_map(|c| c.to_uppercase()));
            }

            // fast path: we can avoid mapping chars to upper-case, if we don't want to ignore the case
            if !self.ignore_case {
                return closure(&mut fields_to_check.chars().skip(slice_start).take(slice_stop));
            }

            closure(
                &mut fields_to_check
                    .chars()
                    .skip(slice_start)
                    .take(slice_stop)
                    .flat_map(|c| c.to_uppercase()),
            )
        } else {
            closure(&mut fields_to_check.chars())
        }
    }

    fn should_print_delimiter(&self, group_count: usize, first_line_printed: bool) -> bool {
        // if no delimiter option is selected then no other checks needed
        self.delimiters != Delimiters::None
            // print delimiter only before the first line of a group, not between lines of a group
            && group_count == 1
            // if at least one line has been output before current group then print delimiter
            && (first_line_printed
                // or if we need to prepend delimiter then print it even at the start of the output
                || self.delimiters == Delimiters::Prepend
                // the 'both' delimit mode should prepend and append delimiters
                || self.delimiters == Delimiters::Both)
    }

    fn print_line<W: Write>(
        &self,
        writer: &mut BufWriter<W>,
        line: &str,
        count: usize,
        first_line_printed: bool,
    ) -> UResult<()> {
        let line_terminator = self.get_line_terminator();

        if self.should_print_delimiter(count, first_line_printed) {
            write_line_terminator!(writer, line_terminator)?;
        }

        if self.show_counts {
            writer.write_all(format!("{:7} {}", count, line).as_bytes())
        } else {
            writer.write_all(line.as_bytes())
        }
        .map_err_context(|| "Failed to write line".to_string())?;

        write_line_terminator!(writer, line_terminator)
    }
}

fn get_line_string(io_line: io::Result<Vec<u8>>) -> UResult<String> {
    let line_bytes = io_line.map_err_context(|| "failed to split lines".to_string())?;
    String::from_utf8(line_bytes)
        .map_err(|e| USimpleError::new(1, format!("failed to convert line to utf8: {}", e)))
}

fn opt_parsed<T: FromStr>(opt_name: &str, matches: &ArgMatches) -> UResult<Option<T>> {
    Ok(match matches.value_of(opt_name) {
        Some(arg_str) => Some(arg_str.parse().map_err(|_| {
            USimpleError::new(
                1,
                format!(
                    "Invalid argument for {}: {}",
                    opt_name,
                    arg_str.maybe_quote()
                ),
            )
        })?),
        None => None,
    })
}

fn get_long_usage() -> String {
    String::from(
        "Filter adjacent matching lines from INPUT (or standard input),\n\
        writing to OUTPUT (or standard output).
        Note: 'uniq' does not detect repeated lines unless they are adjacent.\n\
        You may want to sort the input first, or use 'sort -u' without 'uniq'.\n",
    )
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let long_usage = get_long_usage();

    let matches = uu_app().after_help(&long_usage[..]).get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let (in_file_name, out_file_name) = match files.len() {
        0 => ("-".to_owned(), "-".to_owned()),
        1 => (files[0].clone(), "-".to_owned()),
        2 => (files[0].clone(), files[1].clone()),
        _ => {
            unreachable!() // Cannot happen as clap will fail earlier
        }
    };

    let uniq = Uniq {
        repeats_only: matches.is_present(options::REPEATED)
            || matches.is_present(options::ALL_REPEATED),
        uniques_only: matches.is_present(options::UNIQUE),
        all_repeated: matches.is_present(options::ALL_REPEATED)
            || matches.is_present(options::GROUP),
        delimiters: get_delimiter(&matches),
        show_counts: matches.is_present(options::COUNT),
        skip_fields: opt_parsed(options::SKIP_FIELDS, &matches)?,
        slice_start: opt_parsed(options::SKIP_CHARS, &matches)?,
        slice_stop: opt_parsed(options::CHECK_CHARS, &matches)?,
        ignore_case: matches.is_present(options::IGNORE_CASE),
        zero_terminated: matches.is_present(options::ZERO_TERMINATED),
    };
    uniq.print_uniq(
        &mut open_input_file(&in_file_name)?,
        &mut open_output_file(&out_file_name)?,
    )
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::ALL_REPEATED)
                .short('D')
                .long(options::ALL_REPEATED)
                .possible_values(&[
                    "none",
                    "prepend",
                    "separate"
                ])
                .help("print all duplicate lines. Delimiting is done with blank lines. [default: none]")
                .value_name("delimit-method")
                .min_values(0)
                .max_values(1),
        )
        .arg(
            Arg::new(options::GROUP)
                .long(options::GROUP)
                .possible_values(&[
                    "separate",
                    "prepend",
                    "append",
                    "both",
                ])
                .help("show all items, separating groups with an empty line. [default: separate]")
                .value_name("group-method")
                .min_values(0)
                .max_values(1)
                .conflicts_with_all(&[
                    options::REPEATED,
                    options::ALL_REPEATED,
                    options::UNIQUE,
                ]),
        )
        .arg(
            Arg::new(options::CHECK_CHARS)
                .short('w')
                .long(options::CHECK_CHARS)
                .help("compare no more than N characters in lines")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::COUNT)
                .short('c')
                .long(options::COUNT)
                .help("prefix lines by the number of occurrences"),
        )
        .arg(
            Arg::new(options::IGNORE_CASE)
                .short('i')
                .long(options::IGNORE_CASE)
                .help("ignore differences in case when comparing"),
        )
        .arg(
            Arg::new(options::REPEATED)
                .short('d')
                .long(options::REPEATED)
                .help("only print duplicate lines"),
        )
        .arg(
            Arg::new(options::SKIP_CHARS)
                .short('s')
                .long(options::SKIP_CHARS)
                .help("avoid comparing the first N characters")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::SKIP_FIELDS)
                .short('f')
                .long(options::SKIP_FIELDS)
                .help("avoid comparing the first N fields")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::UNIQUE)
                .short('u')
                .long(options::UNIQUE)
                .help("only print unique lines"),
        )
        .arg(
            Arg::new(options::ZERO_TERMINATED)
                .short('z')
                .long(options::ZERO_TERMINATED)
                .help("end lines with 0 byte, not newline"),
        )
        .arg(
            Arg::new(ARG_FILES)
                .multiple_occurrences(true)
                .takes_value(true)
                .max_values(2),
        )
}

fn get_delimiter(matches: &ArgMatches) -> Delimiters {
    let value = matches
        .value_of(options::ALL_REPEATED)
        .or_else(|| matches.value_of(options::GROUP));
    if let Some(delimiter_arg) = value {
        Delimiters::from_str(delimiter_arg).unwrap() // All possible values for ALL_REPEATED are Delimiters (of type `&str`)
    } else if matches.is_present(options::GROUP) {
        Delimiters::Separate
    } else {
        Delimiters::None
    }
}

fn open_input_file(in_file_name: &str) -> UResult<BufReader<Box<dyn Read + 'static>>> {
    let in_file = if in_file_name == "-" {
        Box::new(stdin()) as Box<dyn Read>
    } else {
        let path = Path::new(in_file_name);
        let in_file = File::open(&path)
            .map_err_context(|| format!("Could not open {}", in_file_name.maybe_quote()))?;
        Box::new(in_file) as Box<dyn Read>
    };
    Ok(BufReader::new(in_file))
}

fn open_output_file(out_file_name: &str) -> UResult<BufWriter<Box<dyn Write + 'static>>> {
    let out_file = if out_file_name == "-" {
        Box::new(stdout()) as Box<dyn Write>
    } else {
        let path = Path::new(out_file_name);
        let out_file = File::create(&path)
            .map_err_context(|| format!("Could not create {}", out_file_name.maybe_quote()))?;
        Box::new(out_file) as Box<dyn Write>
    };
    Ok(BufWriter::new(out_file))
}
<commit_msg>uniq: print warning when both `-D` and `-c` are passed<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Chirag B Jadwani <chirag.jadwani@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

use clap::{crate_version, Arg, ArgMatches, Command};
use std::fs::File;
use std::io::{self, stdin, stdout, BufRead, BufReader, BufWriter, Read, Write};
use std::path::Path;
use std::str::FromStr;
use strum_macros::{AsRefStr, EnumString};
use uucore::display::Quotable;
use uucore::error::{FromIo, UResult, USimpleError, UUsageError};
use uucore::format_usage;

static ABOUT: &str = "Report or omit repeated lines.";
const USAGE: &str = "{} [OPTION]... [INPUT [OUTPUT]]...";
pub mod options {
    pub static ALL_REPEATED: &str = "all-repeated";
    pub static CHECK_CHARS: &str = "check-chars";
    pub static COUNT: &str = "count";
    pub static IGNORE_CASE: &str = "ignore-case";
    pub static REPEATED: &str = "repeated";
    pub static SKIP_FIELDS: &str = "skip-fields";
    pub static SKIP_CHARS: &str = "skip-chars";
    pub static UNIQUE: &str = "unique";
    pub static ZERO_TERMINATED: &str = "zero-terminated";
    pub static GROUP: &str = "group";
}

static ARG_FILES: &str = "files";

#[derive(PartialEq, Clone, Copy, AsRefStr, EnumString)]
#[strum(serialize_all = "snake_case")]
enum Delimiters {
    Append,
    Prepend,
    Separate,
    Both,
    None,
}

struct Uniq {
    repeats_only: bool,
    uniques_only: bool,
    all_repeated: bool,
    delimiters: Delimiters,
    show_counts: bool,
    skip_fields: Option<usize>,
    slice_start: Option<usize>,
    slice_stop: Option<usize>,
    ignore_case: bool,
    zero_terminated: bool,
}

macro_rules! write_line_terminator {
    ($writer:expr, $line_terminator:expr) => {
        $writer
            .write_all(&[$line_terminator])
            .map_err_context(|| "Could not write line terminator".to_string())
    };
}

impl Uniq {
    pub fn print_uniq<R: Read, W: Write>(
        &self,
        reader: &mut BufReader<R>,
        writer: &mut BufWriter<W>,
    ) -> UResult<()> {
        let mut first_line_printed = false;
        let mut group_count = 1;
        let line_terminator = self.get_line_terminator();
        let mut lines = reader.split(line_terminator).map(get_line_string);
        let mut line = match lines.next() {
            Some(l) => l?,
            None => return Ok(()),
        };

        // compare current `line` with consecutive lines (`next_line`) of the input
        // and if needed, print `line` based on the command line options provided
        for next_line in lines {
            let next_line = next_line?;
            if self.cmp_keys(&line, &next_line) {
                if (group_count == 1 && !self.repeats_only)
                    || (group_count > 1 && !self.uniques_only)
                {
                    self.print_line(writer, &line, group_count, first_line_printed)?;
                    first_line_printed = true;
                }
                line = next_line;
                group_count = 1;
            } else {
                if self.all_repeated {
                    self.print_line(writer, &line, group_count, first_line_printed)?;
                    first_line_printed = true;
                    line = next_line;
                }
                group_count += 1;
            }
        }
        if (group_count == 1 && !self.repeats_only) || (group_count > 1 && !self.uniques_only) {
            self.print_line(writer, &line, group_count, first_line_printed)?;
            first_line_printed = true;
        }
        if (self.delimiters == Delimiters::Append || self.delimiters == Delimiters::Both)
            && first_line_printed
        {
            write_line_terminator!(writer, line_terminator)?;
        }
        Ok(())
    }

    fn skip_fields<'a>(&self, line: &'a str) -> &'a str {
        if let Some(skip_fields) = self.skip_fields {
            let mut i = 0;
            let mut char_indices = line.char_indices();
            for _ in 0..skip_fields {
                if char_indices.find(|(_, c)| !c.is_whitespace()) == None {
                    return "";
                }
                match char_indices.find(|(_, c)| c.is_whitespace()) {
                    None => return "",

                    Some((next_field_i, _)) => i = next_field_i,
                }
            }
            &line[i..]
        } else {
            line
        }
    }

    fn get_line_terminator(&self) -> u8 {
        if self.zero_terminated {
            0
        } else {
            b'\n'
        }
    }

    fn cmp_keys(&self, first: &str, second: &str) -> bool {
        self.cmp_key(first, |first_iter| {
            self.cmp_key(second, |second_iter| first_iter.ne(second_iter))
        })
    }

    fn cmp_key<F>(&self, line: &str, mut closure: F) -> bool
    where
        F: FnMut(&mut dyn Iterator<Item = char>) -> bool,
    {
        let fields_to_check = self.skip_fields(line);
        let len = fields_to_check.len();
        let slice_start = self.slice_start.unwrap_or(0);
        let slice_stop = self.slice_stop.unwrap_or(len);
        if len > 0 {
            // fast path: avoid doing any work if there is no need to skip or map to lower-case
            if !self.ignore_case && slice_start == 0 && slice_stop == len {
                return closure(&mut fields_to_check.chars());
            }

            // fast path: avoid skipping
            if self.ignore_case && slice_start == 0 && slice_stop == len {
                return closure(&mut fields_to_check.chars().flat_map(|c| c.to_uppercase()));
            }

            // fast path: we can avoid mapping chars to upper-case, if we don't want to ignore the case
            if !self.ignore_case {
                return closure(&mut fields_to_check.chars().skip(slice_start).take(slice_stop));
            }

            closure(
                &mut fields_to_check
                    .chars()
                    .skip(slice_start)
                    .take(slice_stop)
                    .flat_map(|c| c.to_uppercase()),
            )
        } else {
            closure(&mut fields_to_check.chars())
        }
    }

    fn should_print_delimiter(&self, group_count: usize, first_line_printed: bool) -> bool {
        // if no delimiter option is selected then no other checks needed
        self.delimiters != Delimiters::None
            // print delimiter only before the first line of a group, not between lines of a group
            && group_count == 1
            // if at least one line has been output before current group then print delimiter
            && (first_line_printed
                // or if we need to prepend delimiter then print it even at the start of the output
                || self.delimiters == Delimiters::Prepend
                // the 'both' delimit mode should prepend and append delimiters
                || self.delimiters == Delimiters::Both)
    }

    fn print_line<W: Write>(
        &self,
        writer: &mut BufWriter<W>,
        line: &str,
        count: usize,
        first_line_printed: bool,
    ) -> UResult<()> {
        let line_terminator = self.get_line_terminator();

        if self.should_print_delimiter(count, first_line_printed) {
            write_line_terminator!(writer, line_terminator)?;
        }

        if self.show_counts {
            writer.write_all(format!("{:7} {}", count, line).as_bytes())
        } else {
            writer.write_all(line.as_bytes())
        }
        .map_err_context(|| "Failed to write line".to_string())?;

        write_line_terminator!(writer, line_terminator)
    }
}

fn get_line_string(io_line: io::Result<Vec<u8>>) -> UResult<String> {
    let line_bytes = io_line.map_err_context(|| "failed to split lines".to_string())?;
    String::from_utf8(line_bytes)
        .map_err(|e| USimpleError::new(1, format!("failed to convert line to utf8: {}", e)))
}

fn opt_parsed<T: FromStr>(opt_name: &str, matches: &ArgMatches) -> UResult<Option<T>> {
    Ok(match matches.value_of(opt_name) {
        Some(arg_str) => Some(arg_str.parse().map_err(|_| {
            USimpleError::new(
                1,
                format!(
                    "Invalid argument for {}: {}",
                    opt_name,
                    arg_str.maybe_quote()
                ),
            )
        })?),
        None => None,
    })
}

fn get_long_usage() -> String {
    String::from(
        "Filter adjacent matching lines from INPUT (or standard input),\n\
        writing to OUTPUT (or standard output).
        Note: 'uniq' does not detect repeated lines unless they are adjacent.\n\
        You may want to sort the input first, or use 'sort -u' without 'uniq'.\n",
    )
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let long_usage = get_long_usage();

    let matches = uu_app().after_help(&long_usage[..]).get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let (in_file_name, out_file_name) = match files.len() {
        0 => ("-".to_owned(), "-".to_owned()),
        1 => (files[0].clone(), "-".to_owned()),
        2 => (files[0].clone(), files[1].clone()),
        _ => {
            unreachable!() // Cannot happen as clap will fail earlier
        }
    };

    let uniq = Uniq {
        repeats_only: matches.is_present(options::REPEATED)
            || matches.is_present(options::ALL_REPEATED),
        uniques_only: matches.is_present(options::UNIQUE),
        all_repeated: matches.is_present(options::ALL_REPEATED)
            || matches.is_present(options::GROUP),
        delimiters: get_delimiter(&matches),
        show_counts: matches.is_present(options::COUNT),
        skip_fields: opt_parsed(options::SKIP_FIELDS, &matches)?,
        slice_start: opt_parsed(options::SKIP_CHARS, &matches)?,
        slice_stop: opt_parsed(options::CHECK_CHARS, &matches)?,
        ignore_case: matches.is_present(options::IGNORE_CASE),
        zero_terminated: matches.is_present(options::ZERO_TERMINATED),
    };

    if uniq.show_counts && uniq.all_repeated {
        return Err(UUsageError::new(
            1,
            "printing all duplicated lines and repeat counts is meaningless",
        ));
    }

    uniq.print_uniq(
        &mut open_input_file(&in_file_name)?,
        &mut open_output_file(&out_file_name)?,
    )
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::ALL_REPEATED)
                .short('D')
                .long(options::ALL_REPEATED)
                .possible_values(&[
                    "none",
                    "prepend",
                    "separate"
                ])
                .help("print all duplicate lines. Delimiting is done with blank lines. [default: none]")
                .value_name("delimit-method")
                .min_values(0)
                .max_values(1),
        )
        .arg(
            Arg::new(options::GROUP)
                .long(options::GROUP)
                .possible_values(&[
                    "separate",
                    "prepend",
                    "append",
                    "both",
                ])
                .help("show all items, separating groups with an empty line. [default: separate]")
                .value_name("group-method")
                .min_values(0)
                .max_values(1)
                .conflicts_with_all(&[
                    options::REPEATED,
                    options::ALL_REPEATED,
                    options::UNIQUE,
                ]),
        )
        .arg(
            Arg::new(options::CHECK_CHARS)
                .short('w')
                .long(options::CHECK_CHARS)
                .help("compare no more than N characters in lines")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::COUNT)
                .short('c')
                .long(options::COUNT)
                .help("prefix lines by the number of occurrences"),
        )
        .arg(
            Arg::new(options::IGNORE_CASE)
                .short('i')
                .long(options::IGNORE_CASE)
                .help("ignore differences in case when comparing"),
        )
        .arg(
            Arg::new(options::REPEATED)
                .short('d')
                .long(options::REPEATED)
                .help("only print duplicate lines"),
        )
        .arg(
            Arg::new(options::SKIP_CHARS)
                .short('s')
                .long(options::SKIP_CHARS)
                .help("avoid comparing the first N characters")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::SKIP_FIELDS)
                .short('f')
                .long(options::SKIP_FIELDS)
                .help("avoid comparing the first N fields")
                .value_name("N"),
        )
        .arg(
            Arg::new(options::UNIQUE)
                .short('u')
                .long(options::UNIQUE)
                .help("only print unique lines"),
        )
        .arg(
            Arg::new(options::ZERO_TERMINATED)
                .short('z')
                .long(options::ZERO_TERMINATED)
                .help("end lines with 0 byte, not newline"),
        )
        .arg(
            Arg::new(ARG_FILES)
                .multiple_occurrences(true)
                .takes_value(true)
                .max_values(2),
        )
}

fn get_delimiter(matches: &ArgMatches) -> Delimiters {
    let value = matches
        .value_of(options::ALL_REPEATED)
        .or_else(|| matches.value_of(options::GROUP));
    if let Some(delimiter_arg) = value {
        Delimiters::from_str(delimiter_arg).unwrap() // All possible values for ALL_REPEATED are Delimiters (of type `&str`)
    } else if matches.is_present(options::GROUP) {
        Delimiters::Separate
    } else {
        Delimiters::None
    }
}

fn open_input_file(in_file_name: &str) -> UResult<BufReader<Box<dyn Read + 'static>>> {
    let in_file = if in_file_name == "-" {
        Box::new(stdin()) as Box<dyn Read>
    } else {
        let path = Path::new(in_file_name);
        let in_file = File::open(&path)
            .map_err_context(|| format!("Could not open {}", in_file_name.maybe_quote()))?;
        Box::new(in_file) as Box<dyn Read>
    };
    Ok(BufReader::new(in_file))
}

fn open_output_file(out_file_name: &str) -> UResult<BufWriter<Box<dyn Write + 'static>>> {
    let out_file = if out_file_name == "-" {
        Box::new(stdout()) as Box<dyn Write>
    } else {
        let path = Path::new(out_file_name);
        let out_file = File::create(&path)
            .map_err_context(|| format!("Could not create {}", out_file_name.maybe_quote()))?;
        Box::new(out_file) as Box<dyn Write>
    };
    Ok(BufWriter::new(out_file))
}
<|endoftext|>


<commit_before>use baseband::{Decoder, Decider};
use error::{P25Error, Result};
use nid;
use status::{StreamSymbol, StatusDeinterleaver};
use sync::{SyncCorrelator, SyncDetector};

use self::State::*;
use self::StateChange::*;

const PRIME_SAMPLES: u32 = 6000;

#[derive(Debug)]
pub enum ReceiverEvent {
    Symbol(StreamSymbol),
    NetworkID(nid::NetworkID),
}

#[derive(Copy, Clone)]
struct Receiver {
    recv: Decoder,
    status: StatusDeinterleaver,
}

impl Receiver {
    pub fn new(recv: Decoder) -> Receiver {
        Receiver {
            recv: recv,
            status: StatusDeinterleaver::new(),
        }
    }

    pub fn feed(&mut self, s: f32) -> Option<StreamSymbol> {
        let dibit = match self.recv.feed(s) {
            Some(dibit) => dibit,
            None => return None,
        };

        Some(self.status.feed(dibit))
    }
}

enum State {
    Prime(u32),
    Sync(SyncDetector),
    DecodeNID(Receiver, nid::NIDReceiver),
    DecodePacket(Receiver),
    FlushPads(Receiver),
}

enum StateChange {
    Change(State),
    Event(ReceiverEvent),
    EventChange(ReceiverEvent, State),
    Error(P25Error),
    NoChange,
}

impl State {
    pub fn prime() -> State { Prime(1) }
    pub fn sync() -> State { Sync(SyncDetector::new()) }
    pub fn decode_nid(decoder: Decoder) -> State {
        DecodeNID(Receiver::new(decoder), nid::NIDReceiver::new())
    }
    pub fn decode_packet(recv: Receiver) -> State { DecodePacket(recv) }
    pub fn flush_pads(recv: Receiver) -> State { FlushPads(recv) }
}

pub struct DataUnitReceiver {
    state: State,
    corr: SyncCorrelator,
}

impl DataUnitReceiver {
    pub fn new() -> DataUnitReceiver {
        DataUnitReceiver {
            state: State::prime(),
            corr: SyncCorrelator::new(),
        }
    }

    pub fn flush_pads(&mut self) {
        if let DecodePacket(recv) = self.state {
            self.state = State::flush_pads(recv);
        } else {
            panic!("not decoding a packet");
        }
    }

    pub fn resync(&mut self) { self.state = State::sync(); }

    fn handle(&mut self, s: f32) -> StateChange {
        let (power, thresh) = self.corr.feed(s);

        match self.state {
            Prime(t) => if t == PRIME_SAMPLES {
                Change(State::sync())
            } else {
                Change(Prime(t + 1))
            },
            Sync(ref mut sync) => if sync.feed(power, thresh) {
                let (p, m, n) = self.corr.thresholds();
                Change(State::decode_nid(Decoder::new(Decider::new(p, m, n))))
            } else {
                NoChange
            },
            DecodeNID(ref mut recv, ref mut nid) => {
                let dibit = match recv.feed(s) {
                    Some(StreamSymbol::Data(d)) => d,
                    Some(s) => return Event(ReceiverEvent::Symbol(s)),
                    None => return NoChange,
                };

                match nid.feed(dibit) {
                    Some(Ok(nid)) => EventChange(ReceiverEvent::NetworkID(nid),
                                                 State::decode_packet(*recv)),
                    Some(Err(e)) => Error(e),
                    None => NoChange,
                }
            },
            DecodePacket(ref mut recv) => match recv.feed(s) {
                Some(x) => Event(ReceiverEvent::Symbol(x)),
                None => NoChange,
            },
            FlushPads(ref mut recv) => match recv.feed(s) {
                Some(StreamSymbol::Status(_)) => Change(State::sync()),
                _ => NoChange,
            },
        }
    }

    pub fn feed(&mut self, s: f32) -> Option<Result<ReceiverEvent>> {
        match self.handle(s) {
            Change(state) => {
                self.state = state;
                None
            },
            Event(event) => Some(Ok(event)),
            EventChange(event, state) => {
                self.state = state;
                Some(Ok(event))
            },
            Error(err) => Some(Err(err)),
            NoChange => None,
        }
    }
}
<commit_msg>Ignore flush_pads call while resyncing<commit_after>use baseband::{Decoder, Decider};
use error::{P25Error, Result};
use nid;
use status::{StreamSymbol, StatusDeinterleaver};
use sync::{SyncCorrelator, SyncDetector};

use self::State::*;
use self::StateChange::*;

const PRIME_SAMPLES: u32 = 6000;

#[derive(Debug)]
pub enum ReceiverEvent {
    Symbol(StreamSymbol),
    NetworkID(nid::NetworkID),
}

#[derive(Copy, Clone)]
struct Receiver {
    recv: Decoder,
    status: StatusDeinterleaver,
}

impl Receiver {
    pub fn new(recv: Decoder) -> Receiver {
        Receiver {
            recv: recv,
            status: StatusDeinterleaver::new(),
        }
    }

    pub fn feed(&mut self, s: f32) -> Option<StreamSymbol> {
        let dibit = match self.recv.feed(s) {
            Some(dibit) => dibit,
            None => return None,
        };

        Some(self.status.feed(dibit))
    }
}

enum State {
    Prime(u32),
    Sync(SyncDetector),
    DecodeNID(Receiver, nid::NIDReceiver),
    DecodePacket(Receiver),
    FlushPads(Receiver),
}

enum StateChange {
    Change(State),
    Event(ReceiverEvent),
    EventChange(ReceiverEvent, State),
    Error(P25Error),
    NoChange,
}

impl State {
    pub fn prime() -> State { Prime(1) }
    pub fn sync() -> State { Sync(SyncDetector::new()) }
    pub fn decode_nid(decoder: Decoder) -> State {
        DecodeNID(Receiver::new(decoder), nid::NIDReceiver::new())
    }
    pub fn decode_packet(recv: Receiver) -> State { DecodePacket(recv) }
    pub fn flush_pads(recv: Receiver) -> State { FlushPads(recv) }
}

pub struct DataUnitReceiver {
    state: State,
    corr: SyncCorrelator,
}

impl DataUnitReceiver {
    pub fn new() -> DataUnitReceiver {
        DataUnitReceiver {
            state: State::prime(),
            corr: SyncCorrelator::new(),
        }
    }

    pub fn flush_pads(&mut self) {
        match self.state {
            DecodePacket(recv) => self.state = State::flush_pads(recv),
            Sync(_) => {},
            _ => panic!("not decoding a packet"),
        }
    }

    pub fn resync(&mut self) { self.state = State::sync(); }

    fn handle(&mut self, s: f32) -> StateChange {
        let (power, thresh) = self.corr.feed(s);

        match self.state {
            Prime(t) => if t == PRIME_SAMPLES {
                Change(State::sync())
            } else {
                Change(Prime(t + 1))
            },
            Sync(ref mut sync) => if sync.feed(power, thresh) {
                let (p, m, n) = self.corr.thresholds();
                Change(State::decode_nid(Decoder::new(Decider::new(p, m, n))))
            } else {
                NoChange
            },
            DecodeNID(ref mut recv, ref mut nid) => {
                let dibit = match recv.feed(s) {
                    Some(StreamSymbol::Data(d)) => d,
                    Some(s) => return Event(ReceiverEvent::Symbol(s)),
                    None => return NoChange,
                };

                match nid.feed(dibit) {
                    Some(Ok(nid)) => EventChange(ReceiverEvent::NetworkID(nid),
                                                 State::decode_packet(*recv)),
                    Some(Err(e)) => Error(e),
                    None => NoChange,
                }
            },
            DecodePacket(ref mut recv) => match recv.feed(s) {
                Some(x) => Event(ReceiverEvent::Symbol(x)),
                None => NoChange,
            },
            FlushPads(ref mut recv) => match recv.feed(s) {
                Some(StreamSymbol::Status(_)) => Change(State::sync()),
                _ => NoChange,
            },
        }
    }

    pub fn feed(&mut self, s: f32) -> Option<Result<ReceiverEvent>> {
        match self.handle(s) {
            Change(state) => {
                self.state = state;
                None
            },
            Event(event) => Some(Ok(event)),
            EventChange(event, state) => {
                self.state = state;
                Some(Ok(event))
            },
            Error(err) => Some(Err(err)),
            NoChange => None,
        }
    }
}
<|endoftext|>


<commit_before>use gateway::{ReconnectType, Shard, ShardAction};
use internal::prelude::*;
use internal::ws_impl::ReceiverExt;
use model::event::{Event, GatewayEvent};
use parking_lot::Mutex;
use serde::Deserialize;
use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
use std::sync::Arc;
use super::super::super::{EventHandler, dispatch};
use super::{ShardClientMessage, ShardId, ShardManagerMessage, ShardRunnerMessage};
use threadpool::ThreadPool;
use typemap::ShareMap;
use websocket::message::{CloseData, OwnedMessage};
use websocket::WebSocketError;

#[cfg(feature = "framework")]
use framework::Framework;
#[cfg(feature = "voice")]
use internal::ws_impl::SenderExt;

/// A runner for managing a [`Shard`] and its respective WebSocket client.
///
/// [`Shard`]: ../../../gateway/struct.Shard.html
pub struct ShardRunner<H: EventHandler + Send + Sync + 'static> {
    data: Arc<Mutex<ShareMap>>,
    event_handler: Arc<H>,
    #[cfg(feature = "framework")]
    framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
    manager_tx: Sender<ShardManagerMessage>,
    // channel to receive messages from the shard manager and dispatches
    runner_rx: Receiver<ShardClientMessage>,
    // channel to send messages to the shard runner from the shard manager
    runner_tx: Sender<ShardClientMessage>,
    shard: Shard,
    threadpool: ThreadPool,
}

impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
    /// Creates a new runner for a Shard.
    #[allow(too_many_arguments)]
    #[cfg(feature = "framework")]
    pub fn new(
        shard: Shard,
        manager_tx: Sender<ShardManagerMessage>,
        framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> Self {
        let (tx, rx) = mpsc::channel();

        Self {
            runner_rx: rx,
            runner_tx: tx,
            data,
            event_handler,
            framework,
            manager_tx,
            shard,
            threadpool,
        }
    }

    /// Creates a new runner for a Shard.
    #[cfg(not(feature = "framework"))]
    pub fn new(
        shard: Shard,
        manager_tx: Sender<ShardManagerMessage>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> Self {
        let (tx, rx) = mpsc::channel();

        Self {
            runner_rx: rx,
            runner_tx: tx,
            data,
            event_handler,
            manager_tx,
            shard,
            threadpool,
        }
    }

    /// Starts the runner's loop to receive events.
    ///
    /// This runs a loop that performs the following in each iteration:
    ///
    /// 1. checks the receiver for [`ShardRunnerMessage`]s, possibly from the
    /// [`ShardManager`], and if there is one, acts on it.
    ///
    /// 2. checks if a heartbeat should be sent to the discord Gateway, and if
    /// so, sends one.
    ///
    /// 3. attempts to retrieve a message from the WebSocket, processing it into
    /// a [`GatewayEvent`]. This will block for 100ms before assuming there is
    /// no message available.
    ///
    /// 4. Checks with the [`Shard`] to determine if the gateway event is
    /// specifying an action to take (e.g. resuming, reconnecting, heartbeating)
    /// and then performs that action, if any.
    ///
    /// 5. Dispatches the event via the Client.
    ///
    /// 6. Go back to 1.
    ///
    /// [`GatewayEvent`]: ../../../model/event/enum.GatewayEvent.html
    /// [`Shard`]: ../../../gateway/struct.Shard.html
    /// [`ShardManager`]: struct.ShardManager.html
    /// [`ShardRunnerMessage`]: enum.ShardRunnerMessage.html
    pub fn run(&mut self) -> Result<()> {
        debug!("[ShardRunner {:?}] Running", self.shard.shard_info());

        loop {
            if !self.recv()? {
                return Ok(());
            }

            // check heartbeat
            if let Err(why) = self.shard.check_heartbeat() {
                warn!(
                    "[ShardRunner {:?}] Error heartbeating: {:?}",
                    self.shard.shard_info(),
                    why,
                );
                debug!(
                    "[ShardRunner {:?}] Requesting restart",
                    self.shard.shard_info(),
                );

                return self.request_restart();
            }

            #[cfg(feature = "voice")]
            {
                for message in self.shard.cycle_voice_recv() {
                    if let Err(why) = self.shard.client.send_json(&message) {
                        println!("Err sending from voice over WS: {:?}", why);
                    }
                }
            }

            let (event, action, successful) = self.recv_event();

            match action {
                Some(ShardAction::Reconnect(ReconnectType::Reidentify)) => {
                    return self.request_restart()
                },
                Some(other) => {
                    let _ = self.action(&other);
                },
                None => {},
            }

            if let Some(event) = event {
                self.dispatch(event);
            }

            if !successful && !self.shard.stage().is_connecting() {
                return self.request_restart();
            }
        }
    }

    /// Clones the internal copy of the Sender to the shard runner.
    pub(super) fn runner_tx(&self) -> Sender<ShardClientMessage> {
        self.runner_tx.clone()
    }

    /// Takes an action that a [`Shard`] has determined should happen and then
    /// does it.
    ///
    /// For example, if the shard says that an Identify message needs to be
    /// sent, this will do that.
    ///
    /// # Errors
    ///
    /// Returns
    fn action(&mut self, action: &ShardAction) -> Result<()> {
        match *action {
            ShardAction::Reconnect(ReconnectType::Reidentify) => {
                self.request_restart()
            },
            ShardAction::Reconnect(ReconnectType::Resume) => {
                self.shard.resume()
            },
            ShardAction::Heartbeat => self.shard.heartbeat(),
            ShardAction::Identify => self.shard.identify(),
        }
    }

    // Checks if the ID received to shutdown is equivalent to the ID of the
    // shard this runner is responsible. If so, it shuts down the WebSocket
    // client.
    //
    // Returns whether the WebSocket client is still active.
    //
    // If true, the WebSocket client was _not_ shutdown. If false, it was.
    fn checked_shutdown(&mut self, id: ShardId) -> bool {
        // First verify the ID so we know for certain this runner is
        // to shutdown.
        if id.0 != self.shard.shard_info()[0] {
            // Not meant for this runner for some reason, don't
            // shutdown.
            return true;
        }

        let close_data = CloseData::new(1000, String::new());
        let msg = OwnedMessage::Close(Some(close_data));
        let _ = self.shard.client.send_message(&msg);

        false
    }

    #[inline]
    fn dispatch(&self, event: Event) {
        dispatch(
            event,
            #[cfg(feature = "framework")]
            &self.framework,
            &self.data,
            &self.event_handler,
            &self.runner_tx,
            &self.threadpool,
            self.shard.shard_info()[0],
        );
    }

    // Handles a received value over the shard runner rx channel.
    //
    // Returns a boolean on whether the shard runner can continue.
    //
    // This always returns true, except in the case that the shard manager asked
    // the runner to shutdown.
    fn handle_rx_value(&mut self, value: ShardClientMessage) -> bool {
        match value {
            ShardClientMessage::Manager(x) => match x {
                ShardManagerMessage::Restart(id) |
                ShardManagerMessage::Shutdown(id) => {
                    self.checked_shutdown(id)
                },
                ShardManagerMessage::ShutdownAll => {
                    // This variant should never be received.
                    warn!(
                        "[ShardRunner {:?}] Received a ShutdownAll?",
                        self.shard.shard_info(),
                    );

                    true
                },
                ShardManagerMessage::ShutdownInitiated => {
                    // nb: not sent here

                    true
                },
            }
            ShardClientMessage::Runner(x) => match x {
                ShardRunnerMessage::ChunkGuilds { guild_ids, limit, query } => {
                    self.shard.chunk_guilds(
                        guild_ids,
                        limit,
                        query.as_ref().map(String::as_str),
                    ).is_ok()
                },
                ShardRunnerMessage::Close(code, reason) => {
                    let reason = reason.unwrap_or_else(String::new);
                    let data = CloseData::new(code, reason);
                    let msg = OwnedMessage::Close(Some(data));

                    self.shard.client.send_message(&msg).is_ok()
                },
                ShardRunnerMessage::Message(msg) => {
                    self.shard.client.send_message(&msg).is_ok()
                },
                ShardRunnerMessage::SetGame(game) => {
                    // To avoid a clone of `game`, we do a little bit of
                    // trickery here:
                    //
                    // First, we obtain a reference to the current presence of
                    // the shard, and create a new presence tuple of the new
                    // game we received over the channel as well as the online
                    // status that the shard already had.
                    //
                    // We then (attempt to) send the websocket message with the
                    // status update, expressively returning:
                    //
                    // - whether the message successfully sent
                    // - the original game we received over the channel
                    self.shard.set_game(game);

                    self.shard.update_presence().is_ok()
                },
                ShardRunnerMessage::SetPresence(status, game) => {
                    self.shard.set_presence(status, game);

                    self.shard.update_presence().is_ok()
                },
                ShardRunnerMessage::SetStatus(status) => {
                    self.shard.set_status(status);

                    self.shard.update_presence().is_ok()
                },
            },
        }
    }

    // Receives values over the internal shard runner rx channel and handles
    // them.
    //
    // This will loop over values until there is no longer one.
    //
    // Requests a restart if the sending half of the channel disconnects. This
    // should _never_ happen, as the sending half is kept on the runner.

    // Returns whether the shard runner is in a state that can continue.
    fn recv(&mut self) -> Result<bool> {
        loop {
            match self.runner_rx.try_recv() {
                Ok(value) => {
                    if !self.handle_rx_value(value) {
                        return Ok(false);
                    }
                },
                Err(TryRecvError::Disconnected) => {
                    warn!(
                        "[ShardRunner {:?}] Sending half DC; restarting",
                        self.shard.shard_info(),
                    );

                    let _ = self.request_restart();

                    return Ok(false);
                },
                Err(TryRecvError::Empty) => break,
            }
        }

        // There are no longer any values available.

        Ok(true)
    }

    /// Returns a received event, as well as whether reading the potentially
    /// present event was successful.
    fn recv_event(&mut self) -> (Option<Event>, Option<ShardAction>, bool) {
        let gw_event = match self.shard.client.recv_json() {
            Ok(Some(value)) => {
                GatewayEvent::deserialize(value).map(Some).map_err(From::from)
            },
            Ok(None) => Ok(None),
            Err(Error::WebSocket(WebSocketError::IoError(_))) => {
                // Check that an amount of time at least double the
                // heartbeat_interval has passed.
                //
                // If not, continue on trying to receive messages.
                //
                // If it has, attempt to auto-reconnect.
                {
                    let last = self.shard.last_heartbeat_ack();
                    let interval = self.shard.heartbeat_interval();

                    if let (Some(last_heartbeat_ack), Some(interval)) = (last, interval) {
                        let seconds_passed = last_heartbeat_ack.elapsed().as_secs();
                        let interval_in_secs = interval / 1000;

                        if seconds_passed <= interval_in_secs * 2 {
                            return (None, None, true);
                        }
                    } else {
                        return (None, None, true);
                    }
                }

                debug!("Attempting to auto-reconnect");

                match self.shard.reconnection_type() {
                    ReconnectType::Reidentify => return (None, None, false),
                    ReconnectType::Resume => {
                        if let Err(why) = self.shard.resume() {
                            warn!("Failed to resume: {:?}", why);

                            return (None, None, false);
                        }
                    },
                }

                return (None, None, true);
            },
            Err(Error::WebSocket(WebSocketError::NoDataAvailable)) => {
                // This is hit when the websocket client dies this will be
                // hit every iteration.
                return (None, None, false);
            },
            Err(why) => Err(why),
        };

        let event = match gw_event {
            Ok(Some(event)) => Ok(event),
            Ok(None) => return (None, None, true),
            Err(why) => Err(why),
        };

        let action = match self.shard.handle_event(&event) {
            Ok(Some(action)) => Some(action),
            Ok(None) => None,
            Err(why) => {
                error!("Shard handler received err: {:?}", why);

                return (None, None, true);
            },
        };

        let event = match event {
            Ok(GatewayEvent::Dispatch(_, event)) => Some(event),
            _ => None,
        };

        (event, action, true)
    }

    fn request_restart(&self) -> Result<()> {
        debug!(
            "[ShardRunner {:?}] Requesting restart",
            self.shard.shard_info(),
        );
        let shard_id = ShardId(self.shard.shard_info()[0]);
        let msg = ShardManagerMessage::Restart(shard_id);
        let _ = self.manager_tx.send(msg);

        Ok(())
    }
}
<commit_msg>Add missing trailing comma in match<commit_after>use gateway::{ReconnectType, Shard, ShardAction};
use internal::prelude::*;
use internal::ws_impl::ReceiverExt;
use model::event::{Event, GatewayEvent};
use parking_lot::Mutex;
use serde::Deserialize;
use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
use std::sync::Arc;
use super::super::super::{EventHandler, dispatch};
use super::{ShardClientMessage, ShardId, ShardManagerMessage, ShardRunnerMessage};
use threadpool::ThreadPool;
use typemap::ShareMap;
use websocket::message::{CloseData, OwnedMessage};
use websocket::WebSocketError;

#[cfg(feature = "framework")]
use framework::Framework;
#[cfg(feature = "voice")]
use internal::ws_impl::SenderExt;

/// A runner for managing a [`Shard`] and its respective WebSocket client.
///
/// [`Shard`]: ../../../gateway/struct.Shard.html
pub struct ShardRunner<H: EventHandler + Send + Sync + 'static> {
    data: Arc<Mutex<ShareMap>>,
    event_handler: Arc<H>,
    #[cfg(feature = "framework")]
    framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
    manager_tx: Sender<ShardManagerMessage>,
    // channel to receive messages from the shard manager and dispatches
    runner_rx: Receiver<ShardClientMessage>,
    // channel to send messages to the shard runner from the shard manager
    runner_tx: Sender<ShardClientMessage>,
    shard: Shard,
    threadpool: ThreadPool,
}

impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
    /// Creates a new runner for a Shard.
    #[allow(too_many_arguments)]
    #[cfg(feature = "framework")]
    pub fn new(
        shard: Shard,
        manager_tx: Sender<ShardManagerMessage>,
        framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> Self {
        let (tx, rx) = mpsc::channel();

        Self {
            runner_rx: rx,
            runner_tx: tx,
            data,
            event_handler,
            framework,
            manager_tx,
            shard,
            threadpool,
        }
    }

    /// Creates a new runner for a Shard.
    #[cfg(not(feature = "framework"))]
    pub fn new(
        shard: Shard,
        manager_tx: Sender<ShardManagerMessage>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> Self {
        let (tx, rx) = mpsc::channel();

        Self {
            runner_rx: rx,
            runner_tx: tx,
            data,
            event_handler,
            manager_tx,
            shard,
            threadpool,
        }
    }

    /// Starts the runner's loop to receive events.
    ///
    /// This runs a loop that performs the following in each iteration:
    ///
    /// 1. checks the receiver for [`ShardRunnerMessage`]s, possibly from the
    /// [`ShardManager`], and if there is one, acts on it.
    ///
    /// 2. checks if a heartbeat should be sent to the discord Gateway, and if
    /// so, sends one.
    ///
    /// 3. attempts to retrieve a message from the WebSocket, processing it into
    /// a [`GatewayEvent`]. This will block for 100ms before assuming there is
    /// no message available.
    ///
    /// 4. Checks with the [`Shard`] to determine if the gateway event is
    /// specifying an action to take (e.g. resuming, reconnecting, heartbeating)
    /// and then performs that action, if any.
    ///
    /// 5. Dispatches the event via the Client.
    ///
    /// 6. Go back to 1.
    ///
    /// [`GatewayEvent`]: ../../../model/event/enum.GatewayEvent.html
    /// [`Shard`]: ../../../gateway/struct.Shard.html
    /// [`ShardManager`]: struct.ShardManager.html
    /// [`ShardRunnerMessage`]: enum.ShardRunnerMessage.html
    pub fn run(&mut self) -> Result<()> {
        debug!("[ShardRunner {:?}] Running", self.shard.shard_info());

        loop {
            if !self.recv()? {
                return Ok(());
            }

            // check heartbeat
            if let Err(why) = self.shard.check_heartbeat() {
                warn!(
                    "[ShardRunner {:?}] Error heartbeating: {:?}",
                    self.shard.shard_info(),
                    why,
                );
                debug!(
                    "[ShardRunner {:?}] Requesting restart",
                    self.shard.shard_info(),
                );

                return self.request_restart();
            }

            #[cfg(feature = "voice")]
            {
                for message in self.shard.cycle_voice_recv() {
                    if let Err(why) = self.shard.client.send_json(&message) {
                        println!("Err sending from voice over WS: {:?}", why);
                    }
                }
            }

            let (event, action, successful) = self.recv_event();

            match action {
                Some(ShardAction::Reconnect(ReconnectType::Reidentify)) => {
                    return self.request_restart()
                },
                Some(other) => {
                    let _ = self.action(&other);
                },
                None => {},
            }

            if let Some(event) = event {
                self.dispatch(event);
            }

            if !successful && !self.shard.stage().is_connecting() {
                return self.request_restart();
            }
        }
    }

    /// Clones the internal copy of the Sender to the shard runner.
    pub(super) fn runner_tx(&self) -> Sender<ShardClientMessage> {
        self.runner_tx.clone()
    }

    /// Takes an action that a [`Shard`] has determined should happen and then
    /// does it.
    ///
    /// For example, if the shard says that an Identify message needs to be
    /// sent, this will do that.
    ///
    /// # Errors
    ///
    /// Returns
    fn action(&mut self, action: &ShardAction) -> Result<()> {
        match *action {
            ShardAction::Reconnect(ReconnectType::Reidentify) => {
                self.request_restart()
            },
            ShardAction::Reconnect(ReconnectType::Resume) => {
                self.shard.resume()
            },
            ShardAction::Heartbeat => self.shard.heartbeat(),
            ShardAction::Identify => self.shard.identify(),
        }
    }

    // Checks if the ID received to shutdown is equivalent to the ID of the
    // shard this runner is responsible. If so, it shuts down the WebSocket
    // client.
    //
    // Returns whether the WebSocket client is still active.
    //
    // If true, the WebSocket client was _not_ shutdown. If false, it was.
    fn checked_shutdown(&mut self, id: ShardId) -> bool {
        // First verify the ID so we know for certain this runner is
        // to shutdown.
        if id.0 != self.shard.shard_info()[0] {
            // Not meant for this runner for some reason, don't
            // shutdown.
            return true;
        }

        let close_data = CloseData::new(1000, String::new());
        let msg = OwnedMessage::Close(Some(close_data));
        let _ = self.shard.client.send_message(&msg);

        false
    }

    #[inline]
    fn dispatch(&self, event: Event) {
        dispatch(
            event,
            #[cfg(feature = "framework")]
            &self.framework,
            &self.data,
            &self.event_handler,
            &self.runner_tx,
            &self.threadpool,
            self.shard.shard_info()[0],
        );
    }

    // Handles a received value over the shard runner rx channel.
    //
    // Returns a boolean on whether the shard runner can continue.
    //
    // This always returns true, except in the case that the shard manager asked
    // the runner to shutdown.
    fn handle_rx_value(&mut self, value: ShardClientMessage) -> bool {
        match value {
            ShardClientMessage::Manager(x) => match x {
                ShardManagerMessage::Restart(id) |
                ShardManagerMessage::Shutdown(id) => {
                    self.checked_shutdown(id)
                },
                ShardManagerMessage::ShutdownAll => {
                    // This variant should never be received.
                    warn!(
                        "[ShardRunner {:?}] Received a ShutdownAll?",
                        self.shard.shard_info(),
                    );

                    true
                },
                ShardManagerMessage::ShutdownInitiated => {
                    // nb: not sent here

                    true
                },
            },
            ShardClientMessage::Runner(x) => match x {
                ShardRunnerMessage::ChunkGuilds { guild_ids, limit, query } => {
                    self.shard.chunk_guilds(
                        guild_ids,
                        limit,
                        query.as_ref().map(String::as_str),
                    ).is_ok()
                },
                ShardRunnerMessage::Close(code, reason) => {
                    let reason = reason.unwrap_or_else(String::new);
                    let data = CloseData::new(code, reason);
                    let msg = OwnedMessage::Close(Some(data));

                    self.shard.client.send_message(&msg).is_ok()
                },
                ShardRunnerMessage::Message(msg) => {
                    self.shard.client.send_message(&msg).is_ok()
                },
                ShardRunnerMessage::SetGame(game) => {
                    // To avoid a clone of `game`, we do a little bit of
                    // trickery here:
                    //
                    // First, we obtain a reference to the current presence of
                    // the shard, and create a new presence tuple of the new
                    // game we received over the channel as well as the online
                    // status that the shard already had.
                    //
                    // We then (attempt to) send the websocket message with the
                    // status update, expressively returning:
                    //
                    // - whether the message successfully sent
                    // - the original game we received over the channel
                    self.shard.set_game(game);

                    self.shard.update_presence().is_ok()
                },
                ShardRunnerMessage::SetPresence(status, game) => {
                    self.shard.set_presence(status, game);

                    self.shard.update_presence().is_ok()
                },
                ShardRunnerMessage::SetStatus(status) => {
                    self.shard.set_status(status);

                    self.shard.update_presence().is_ok()
                },
            },
        }
    }

    // Receives values over the internal shard runner rx channel and handles
    // them.
    //
    // This will loop over values until there is no longer one.
    //
    // Requests a restart if the sending half of the channel disconnects. This
    // should _never_ happen, as the sending half is kept on the runner.

    // Returns whether the shard runner is in a state that can continue.
    fn recv(&mut self) -> Result<bool> {
        loop {
            match self.runner_rx.try_recv() {
                Ok(value) => {
                    if !self.handle_rx_value(value) {
                        return Ok(false);
                    }
                },
                Err(TryRecvError::Disconnected) => {
                    warn!(
                        "[ShardRunner {:?}] Sending half DC; restarting",
                        self.shard.shard_info(),
                    );

                    let _ = self.request_restart();

                    return Ok(false);
                },
                Err(TryRecvError::Empty) => break,
            }
        }

        // There are no longer any values available.

        Ok(true)
    }

    /// Returns a received event, as well as whether reading the potentially
    /// present event was successful.
    fn recv_event(&mut self) -> (Option<Event>, Option<ShardAction>, bool) {
        let gw_event = match self.shard.client.recv_json() {
            Ok(Some(value)) => {
                GatewayEvent::deserialize(value).map(Some).map_err(From::from)
            },
            Ok(None) => Ok(None),
            Err(Error::WebSocket(WebSocketError::IoError(_))) => {
                // Check that an amount of time at least double the
                // heartbeat_interval has passed.
                //
                // If not, continue on trying to receive messages.
                //
                // If it has, attempt to auto-reconnect.
                {
                    let last = self.shard.last_heartbeat_ack();
                    let interval = self.shard.heartbeat_interval();

                    if let (Some(last_heartbeat_ack), Some(interval)) = (last, interval) {
                        let seconds_passed = last_heartbeat_ack.elapsed().as_secs();
                        let interval_in_secs = interval / 1000;

                        if seconds_passed <= interval_in_secs * 2 {
                            return (None, None, true);
                        }
                    } else {
                        return (None, None, true);
                    }
                }

                debug!("Attempting to auto-reconnect");

                match self.shard.reconnection_type() {
                    ReconnectType::Reidentify => return (None, None, false),
                    ReconnectType::Resume => {
                        if let Err(why) = self.shard.resume() {
                            warn!("Failed to resume: {:?}", why);

                            return (None, None, false);
                        }
                    },
                }

                return (None, None, true);
            },
            Err(Error::WebSocket(WebSocketError::NoDataAvailable)) => {
                // This is hit when the websocket client dies this will be
                // hit every iteration.
                return (None, None, false);
            },
            Err(why) => Err(why),
        };

        let event = match gw_event {
            Ok(Some(event)) => Ok(event),
            Ok(None) => return (None, None, true),
            Err(why) => Err(why),
        };

        let action = match self.shard.handle_event(&event) {
            Ok(Some(action)) => Some(action),
            Ok(None) => None,
            Err(why) => {
                error!("Shard handler received err: {:?}", why);

                return (None, None, true);
            },
        };

        let event = match event {
            Ok(GatewayEvent::Dispatch(_, event)) => Some(event),
            _ => None,
        };

        (event, action, true)
    }

    fn request_restart(&self) -> Result<()> {
        debug!(
            "[ShardRunner {:?}] Requesting restart",
            self.shard.shard_info(),
        );
        let shard_id = ShardId(self.shard.shard_info()[0]);
        let msg = ShardManagerMessage::Restart(shard_id);
        let _ = self.manager_tx.send(msg);

        Ok(())
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use imgui_sys::{self, ImVec4};

use model::{Color, FieldCoord, HexCoord};
use vec2::Vec2;

const SQRT_3: f32 = 1.732_050_807_568_877_f32;

// #f3e4cf
const FIELD_WHITE: [f32; 4] = [0.9529, 0.8941, 0.8118, 1.0];
// #998578
const FIELD_BLACK: [f32; 4] = [0.6, 0.5216, 0.4706, 1.0];
// #ffffff
const PIECE_WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];
// #000000
const PIECE_OUTLINE: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
// #757575
const PIECE_BLACK: [f32; 4] = [0.4588, 0.4588, 0.4588, 1.0];

macro_rules! im_color {
    ($v:expr) => (imgui_sys::igColorConvertFloat4ToU32(ImVec4::from($v)))
}

pub fn draw_hex(coord: &HexCoord, origin: Vec2, size: f32) {
    for i in 0..6 {
        draw_field(&coord.to_field(i), origin, size);
    }
}

fn draw_field(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(FIELD_WHITE),
            Color::Black => im_color!(FIELD_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
    }
}

pub fn highlight_field(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            highlight,
        );
    }
}

pub fn highlight_field_dot(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let center = field_center(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddCircleFilled(
            draw_list,
            center.into(),
            size / (4.0 * SQRT_3),
            highlight,
            15,
        );
    }
}

pub fn draw_piece(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center = field_center(coord, origin, size);

    const SCALE: f32 = 0.7;

    let v1 = center + (v1 - center) * SCALE;
    let v2 = center + (v2 - center) * SCALE;
    let v3 = center + (v3 - center) * SCALE;

    // Linear equation derived by human testing and regression
    let outline_size = 0.032 * size - 0.335;

    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(PIECE_WHITE),
            Color::Black => im_color!(PIECE_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
        imgui_sys::ImDrawList_AddTriangle(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            im_color!(PIECE_OUTLINE),
            outline_size,
        );
    }
}

fn field_center(coord: &FieldCoord, origin: Vec2, size: f32) -> Vec2 {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center_x = (v1.x + v2.x + v3.x) / 3.0;
    let min_y = (v1.y).min(v2.y).min(v3.y);
    let center_y = match coord.color() {
        Color::White => min_y + size / SQRT_3,
        Color::Black => min_y + size / (2.0 * SQRT_3),
    };

    Vec2::new(center_x, center_y)
}

fn field_vertexes(coord: &FieldCoord, origin: Vec2, size: f32) -> (Vec2, Vec2, Vec2) {
    let center = hex_to_pixel(&coord.to_hex(), origin, size);
    let height = size * SQRT_3 / 2.0;

    let west = Vec2::new(-size, 0.0);
    let east = Vec2::new(size, 0.0);

    let northwest = Vec2::new(-size / 2.0, -height);
    let southwest = Vec2::new(-size / 2.0, height);
    let northeast = Vec2::new(size / 2.0, -height);
    let southeast = Vec2::new(size / 2.0, height);

    let v1;
    let v2;
    match coord.f() {
        0 => {
            v1 = center + northwest;
            v2 = center + northeast;
        }
        1 => {
            v1 = center + northeast;
            v2 = center + east;
        }
        2 => {
            v1 = center + east;
            v2 = center + southeast;
        }
        3 => {
            v1 = center + southeast;
            v2 = center + southwest;
        }
        4 => {
            v1 = center + southwest;
            v2 = center + west;
        }
        5 => {
            v1 = center + west;
            v2 = center + northwest;
        }
        _ => unreachable!(),
    };

    (center, v1, v2)
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#hex-to-pixel
fn hex_to_pixel(coord: &HexCoord, origin: Vec2, size: f32) -> Vec2 {
    let x = f32::from(coord.x());
    let y = f32::from(coord.y());

    let p = Vec2::new(size * (3.0 / 2.0) * x, size * -SQRT_3 * (x / 2.0 + y));

    origin + p
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#pixel-to-hex
pub fn pixel_to_field(p: Vec2, origin: Vec2, size: f32) -> Option<FieldCoord> {
    let v = p - origin;

    let q = v.x * (2.0 / 3.0) / size;
    let r = (-v.x - SQRT_3 * v.y) / (size * 3.0);

    if let Some(hex) = round_hex_coord(q, r) {
        /*
           To find the field, we subtract the converted hex coordinates (q, r) from the rounded hex
           coordinates (hex) to get the fractional part of the coordinates. Here is a diagram of a
           single hex, with the fractional coordinates of each of its vertexes in the hex coordinate
           system:

                (-1/3, 2/3) _______ (1/3, 1/3)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
            (-2/3, 1/3) (----(0,0)----) (2/3, -1/3)
                         \    / \    /
                          \  /   \  /
              (-1/3, -1/3) \/_____\/ (1/3, -2/3)

           Suppose our converted coordinates were (q, r) = (1.333, 0.583). Our rounded coordinates
           are hex = (1, 1), and so our fractional coordinates are (0.333, -0.417).

           We now define three linear equations that will split this hexagon into its six fields:

                            _______ y = x (/)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
                        (------X------) y = -x/2 (-)
                         \    / \    /
                          \  /   \  /
                           \/_____\/ y = -2x (\)

           Using these three equations as linear inequalities, we can check any pair of fractional
           coordinates and find which field it is in. (As a reminder, fields are numbered clockwise
           starting from 0 at the top.)

           For our example (0.333, -0.417), we see that:
               y + x/2 >= 0, so our field is above (-), and is either 5, 0, or 1.
               y +  2x >= 0, so our field is to the right of (\), and is either 0 or 1.
               y -   x <  0, so our field is to the right of (/), and is 1.
        */
        let x_diff = q - f32::from(hex.x());
        let y_diff = r - f32::from(hex.y());
        let mut i = 0;

        if y_diff + x_diff / 2.0 >= 0.0 {
            i |= 0b100;
        }
        if y_diff + x_diff * 2.0 >= 0.0 {
            i |= 0b010;
        }
        if y_diff - x_diff >= 0.0 {
            i |= 0b001;
        }

        // Using a lookup table because nested ifs are too confusing
        const INVALID: u8 = 6;
        let field_lookup = [3, 4, 2, INVALID, INVALID, 5, 1, 0];

        Some(hex.to_field(field_lookup[i]))
    } else {
        None
    }
}

// Algorithm from http://www.redblobgames.com/grids/hexagons/#rounding
fn round_hex_coord(x: f32, y: f32) -> Option<HexCoord> {
    let z = -x - y;

    let mut rx = x.round();
    let mut ry = y.round();
    let rz = z.round();

    let x_diff = (rx - x).abs();
    let y_diff = (ry - y).abs();
    let z_diff = (rz - z).abs();

    if x_diff > y_diff && x_diff > z_diff {
        rx = -ry - rz;
    } else if y_diff > z_diff {
        ry = -rx - rz;
    }

    let rx = rx as i8;
    let ry = ry as i8;

    HexCoord::try_new(rx, ry)
}
<commit_msg>Replace array lookup with match<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use imgui_sys::{self, ImVec4};

use model::{Color, FieldCoord, HexCoord};
use vec2::Vec2;

const SQRT_3: f32 = 1.732_050_807_568_877_f32;

// #f3e4cf
const FIELD_WHITE: [f32; 4] = [0.9529, 0.8941, 0.8118, 1.0];
// #998578
const FIELD_BLACK: [f32; 4] = [0.6, 0.5216, 0.4706, 1.0];
// #ffffff
const PIECE_WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];
// #000000
const PIECE_OUTLINE: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
// #757575
const PIECE_BLACK: [f32; 4] = [0.4588, 0.4588, 0.4588, 1.0];

macro_rules! im_color {
    ($v:expr) => (imgui_sys::igColorConvertFloat4ToU32(ImVec4::from($v)))
}

pub fn draw_hex(coord: &HexCoord, origin: Vec2, size: f32) {
    for i in 0..6 {
        draw_field(&coord.to_field(i), origin, size);
    }
}

fn draw_field(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(FIELD_WHITE),
            Color::Black => im_color!(FIELD_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
    }
}

pub fn highlight_field(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            highlight,
        );
    }
}

pub fn highlight_field_dot(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let center = field_center(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddCircleFilled(
            draw_list,
            center.into(),
            size / (4.0 * SQRT_3),
            highlight,
            15,
        );
    }
}

pub fn draw_piece(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center = field_center(coord, origin, size);

    const SCALE: f32 = 0.7;

    let v1 = center + (v1 - center) * SCALE;
    let v2 = center + (v2 - center) * SCALE;
    let v3 = center + (v3 - center) * SCALE;

    // Linear equation derived by human testing and regression
    let outline_size = 0.032 * size - 0.335;

    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(PIECE_WHITE),
            Color::Black => im_color!(PIECE_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
        imgui_sys::ImDrawList_AddTriangle(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            im_color!(PIECE_OUTLINE),
            outline_size,
        );
    }
}

fn field_center(coord: &FieldCoord, origin: Vec2, size: f32) -> Vec2 {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center_x = (v1.x + v2.x + v3.x) / 3.0;
    let min_y = (v1.y).min(v2.y).min(v3.y);
    let center_y = match coord.color() {
        Color::White => min_y + size / SQRT_3,
        Color::Black => min_y + size / (2.0 * SQRT_3),
    };

    Vec2::new(center_x, center_y)
}

fn field_vertexes(coord: &FieldCoord, origin: Vec2, size: f32) -> (Vec2, Vec2, Vec2) {
    let center = hex_to_pixel(&coord.to_hex(), origin, size);
    let height = size * SQRT_3 / 2.0;

    let west = Vec2::new(-size, 0.0);
    let east = Vec2::new(size, 0.0);

    let northwest = Vec2::new(-size / 2.0, -height);
    let southwest = Vec2::new(-size / 2.0, height);
    let northeast = Vec2::new(size / 2.0, -height);
    let southeast = Vec2::new(size / 2.0, height);

    let v1;
    let v2;
    match coord.f() {
        0 => {
            v1 = center + northwest;
            v2 = center + northeast;
        }
        1 => {
            v1 = center + northeast;
            v2 = center + east;
        }
        2 => {
            v1 = center + east;
            v2 = center + southeast;
        }
        3 => {
            v1 = center + southeast;
            v2 = center + southwest;
        }
        4 => {
            v1 = center + southwest;
            v2 = center + west;
        }
        5 => {
            v1 = center + west;
            v2 = center + northwest;
        }
        _ => unreachable!(),
    };

    (center, v1, v2)
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#hex-to-pixel
fn hex_to_pixel(coord: &HexCoord, origin: Vec2, size: f32) -> Vec2 {
    let x = f32::from(coord.x());
    let y = f32::from(coord.y());

    let p = Vec2::new(size * (3.0 / 2.0) * x, size * -SQRT_3 * (x / 2.0 + y));

    origin + p
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#pixel-to-hex
pub fn pixel_to_field(p: Vec2, origin: Vec2, size: f32) -> Option<FieldCoord> {
    let v = p - origin;

    let q = v.x * (2.0 / 3.0) / size;
    let r = (-v.x - SQRT_3 * v.y) / (size * 3.0);

    if let Some(hex) = round_hex_coord(q, r) {
        /*
           To find the field, we subtract the converted hex coordinates (q, r) from the rounded hex
           coordinates (hex) to get the fractional part of the coordinates. Here is a diagram of a
           single hex, with the fractional coordinates of each of its vertexes in the hex coordinate
           system:

                (-1/3, 2/3) _______ (1/3, 1/3)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
            (-2/3, 1/3) (----(0,0)----) (2/3, -1/3)
                         \    / \    /
                          \  /   \  /
              (-1/3, -1/3) \/_____\/ (1/3, -2/3)

           Suppose our converted coordinates were (q, r) = (1.333, 0.583). Our rounded coordinates
           are hex = (1, 1), and so our fractional coordinates are (0.333, -0.417).

           We now define three linear equations that will split this hexagon into its six fields:

                            _______ y = x (/)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
                        (------X------) y = -x/2 (-)
                         \    / \    /
                          \  /   \  /
                           \/_____\/ y = -2x (\)

           Using these three equations as linear inequalities, we can check any pair of fractional
           coordinates and find which field it is in. (As a reminder, fields are numbered clockwise
           starting from 0 at the top.)

           For our example (0.333, -0.417), we see that:
               y + x/2 >= 0, so our field is above (-), and is either 5, 0, or 1.
               y +  2x >= 0, so our field is to the right of (\), and is either 0 or 1.
               y -   x <  0, so our field is to the right of (/), and is 1.
        */
        let x_diff = q - f32::from(hex.x());
        let y_diff = r - f32::from(hex.y());
        let mut i = 0;

        if y_diff + x_diff / 2.0 >= 0.0 {
            i |= 0b100;
        }
        if y_diff + x_diff * 2.0 >= 0.0 {
            i |= 0b010;
        }
        if y_diff - x_diff >= 0.0 {
            i |= 0b001;
        }

        Some(hex.to_field(match i {
            0 => 3,
            1 => 4,
            2 => 2,
            5 => 5,
            6 => 1,
            7 => 0,
            _ => unreachable!(),
        }))
    } else {
        None
    }
}

// Algorithm from http://www.redblobgames.com/grids/hexagons/#rounding
fn round_hex_coord(x: f32, y: f32) -> Option<HexCoord> {
    let z = -x - y;

    let mut rx = x.round();
    let mut ry = y.round();
    let rz = z.round();

    let x_diff = (rx - x).abs();
    let y_diff = (ry - y).abs();
    let z_diff = (rz - z).abs();

    if x_diff > y_diff && x_diff > z_diff {
        rx = -ry - rz;
    } else if y_diff > z_diff {
        ry = -rx - rz;
    }

    let rx = rx as i8;
    let ry = ry as i8;

    HexCoord::try_new(rx, ry)
}
<|endoftext|>


<commit_before>use bits::Dibit;
use voice::consts;

pub fn descramble(dibits: &[Dibit; consts::FRAME_DIBITS], idx: usize) -> u32 {
    DESCRAMBLERS[idx].descramble(dibits)
}

const DESCRAMBLERS: [VoiceFrameDescrambler; 8] = [
    VoiceFrameDescrambler(&[
        ZigZag::hi(0, 23),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(69, 1),
        ZigZag::lo(0, 22),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(66, 2),
        ZigZag::hi(1, 21),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(1, 20),
        ZigZag::lo(64, 3),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(61, 4),
        ZigZag::hi(2, 11),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(35, 13),
        ZigZag::lo(2, 2),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(8, 15),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::hi(53, 7),
    ]),
];

struct VoiceFrameDescrambler(&'static [ZigZag]);

impl VoiceFrameDescrambler {
    pub fn descramble(&self, dibits: &[Dibit; consts::FRAME_DIBITS]) -> u32 {
        self.0.iter().fold(0, |word, zz| {
            zz.fold(word, |buf, (idx, hi)| {
                buf << 1 | if hi {
                    dibits[idx].hi() as u32
                } else {
                    dibits[idx].lo() as u32
                }
            })
        })
    }
}

#[derive(Clone, Copy)]
struct ZigZag {
    hi: bool,
    idx: usize,
    num: usize,
}

impl ZigZag {
    pub const fn hi(start: usize, num: usize) -> ZigZag {
        ZigZag {
            hi: true,
            idx: start,
            num: num,
        }
    }

    pub const fn lo(start: usize, num: usize) -> ZigZag {
        ZigZag {
            hi: false,
            idx: start,
            num: num,
        }
    }
}

impl Iterator for ZigZag {
    type Item = (usize, bool);

    fn next(&mut self) -> Option<Self::Item> {
        if self.num == 0 {
            return None;
        }

        let cur = (self.idx, self.hi);

        self.idx += 3;
        self.num -= 1;
        self.hi = !self.hi;

        Some(cur)
    }
}

#[cfg(test)]
mod test {
    use super::DESCRAMBLERS;
    use voice::consts;

    #[test]
    fn test_steps_exhaustive() {
        let mut visited = [0u32; consts::FRAME_DIBITS];

        for d in DESCRAMBLERS.iter() {
            for &zz in d.0.iter() {
                for (idx, _) in zz {
                    visited[idx] += 1;
                }
            }
        }

        for &v in visited.iter() {
            assert_eq!(v, 2);
        }
    }
}
<commit_msg>Fix ordering in voice descrambling for codeword 3<commit_after>use bits::Dibit;
use voice::consts;

pub fn descramble(dibits: &[Dibit; consts::FRAME_DIBITS], idx: usize) -> u32 {
    DESCRAMBLERS[idx].descramble(dibits)
}

const DESCRAMBLERS: [VoiceFrameDescrambler; 8] = [
    VoiceFrameDescrambler(&[
        ZigZag::hi(0, 23),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(69, 1),
        ZigZag::lo(0, 22),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(66, 2),
        ZigZag::hi(1, 21),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(64, 3),
        ZigZag::lo(1, 20),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(61, 4),
        ZigZag::hi(2, 11),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(35, 13),
        ZigZag::lo(2, 2),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::lo(8, 15),
    ]),
    VoiceFrameDescrambler(&[
        ZigZag::hi(53, 7),
    ]),
];

struct VoiceFrameDescrambler(&'static [ZigZag]);

impl VoiceFrameDescrambler {
    pub fn descramble(&self, dibits: &[Dibit; consts::FRAME_DIBITS]) -> u32 {
        self.0.iter().fold(0, |word, zz| {
            zz.fold(word, |buf, (idx, hi)| {
                buf << 1 | if hi {
                    dibits[idx].hi() as u32
                } else {
                    dibits[idx].lo() as u32
                }
            })
        })
    }
}

#[derive(Clone, Copy)]
struct ZigZag {
    hi: bool,
    idx: usize,
    num: usize,
}

impl ZigZag {
    pub const fn hi(start: usize, num: usize) -> ZigZag {
        ZigZag {
            hi: true,
            idx: start,
            num: num,
        }
    }

    pub const fn lo(start: usize, num: usize) -> ZigZag {
        ZigZag {
            hi: false,
            idx: start,
            num: num,
        }
    }
}

impl Iterator for ZigZag {
    type Item = (usize, bool);

    fn next(&mut self) -> Option<Self::Item> {
        if self.num == 0 {
            return None;
        }

        let cur = (self.idx, self.hi);

        self.idx += 3;
        self.num -= 1;
        self.hi = !self.hi;

        Some(cur)
    }
}

#[cfg(test)]
mod test {
    use super::DESCRAMBLERS;
    use voice::consts;

    #[test]
    fn test_steps_exhaustive() {
        let mut visited = [0u32; consts::FRAME_DIBITS];

        for d in DESCRAMBLERS.iter() {
            for &zz in d.0.iter() {
                for (idx, _) in zz {
                    visited[idx] += 1;
                }
            }
        }

        for &v in visited.iter() {
            assert_eq!(v, 2);
        }
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A registry to store metrics in

#![allow(missing_docs)]

use std::collections::HashMap;
use metrics::Metric;
use reporter::Reporter;

// TODO break out any notion of metrics. Instead we should have a notion of a collector.
// A collector should be able to insert metrics, and a registry should not.
pub trait Registry<'a>: Send + Sync {
    fn add_scheduled_reporter(&mut self, reporter: Box<Reporter>);
    fn get(&'a self, name: &'a str) -> &'a Metric;
    fn get_metrics_names(&self) -> Vec<&str>;
    fn insert(&mut self, name: &'a str, metric: Metric);
    fn labels(&self) -> HashMap<String, String>;
}

#[derive(Default)]
pub struct StdRegistry<'a> {
    metrics: HashMap<&'a str, Metric>,
    reporter: HashMap<&'a str, Box<Reporter>>,
    labels: HashMap<String, String>,
}

// Specific stuff for registry goes here
impl<'a> Registry<'a> for StdRegistry<'a> {
    fn add_scheduled_reporter(&mut self, reporter: Box<Reporter>) {
        let reporter_name = reporter.get_unique_reporter_name();
        self.reporter.insert(reporter_name, reporter);
    }

    fn get(&'a self, name: &'a str) -> &'a Metric {
        &self.metrics[name]
    }

    fn insert(&mut self, name: &'a str, metric: Metric) {
        self.metrics.insert(name, metric);
    }

    fn get_metrics_names(&self) -> Vec<&str> {
        self.metrics.keys().cloned().collect()
    }

    fn labels(&self) -> HashMap<String, String> {
        self.labels.clone()
    }
}

impl<'a> StdRegistry<'a> {
    pub fn new_with_labels(labels: HashMap<String, String>) -> Self {
        StdRegistry {
            metrics: HashMap::new(),
            reporter: HashMap::new(),
            labels: labels,
        }
    }

    pub fn new() -> Self {
        StdRegistry {
            metrics: HashMap::new(),
            reporter: HashMap::new(),
            labels: HashMap::new(),
        }
    }
}

#[cfg(test)]
mod test {
    use metrics::{Counter, Gauge, Meter, Metric, StdCounter, StdGauge, StdMeter};
    use registry::{Registry, StdRegistry};
    use histogram::Histogram;

    // TODO add labels tests

    #[test]
    fn meter() {
        let mut r = StdRegistry::new();
        let m = StdMeter::new();
        m.mark(100);
        r.insert("meter1", Metric::Meter(m.clone()));
    }

    #[test]
    fn gauge() {
        let mut r = StdRegistry::new();
        let g = StdGauge::new();
        g.set(3);
        r.insert("gauge1", Metric::Gauge(g.clone()));
    }

    #[test]
    fn counter() {
        let mut r = StdRegistry::new();
        let c = StdCounter::new();
        c.add(1);
        r.insert("counter1", Metric::Counter(c.clone()));
    }

    #[test]
    fn histogram() {
        let mut r = StdRegistry::new();
        let mut h = Histogram::configure()
            .max_value(100)
            .precision(1)
            .build()
            .unwrap();
        h.increment_by(1, 1).unwrap();
        r.insert("histogram", Metric::Histogram(h));
    }
}
<commit_msg>Verify that metrics can be updated after insertion into registry.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A registry to store metrics in

#![allow(missing_docs)]

use std::collections::HashMap;
use metrics::Metric;
use reporter::Reporter;

// TODO break out any notion of metrics. Instead we should have a notion of a collector.
// A collector should be able to insert metrics, and a registry should not.
pub trait Registry<'a>: Send + Sync {
    fn add_scheduled_reporter(&mut self, reporter: Box<Reporter>);
    fn get(&'a self, name: &'a str) -> &'a Metric;
    fn get_metrics_names(&self) -> Vec<&str>;
    fn insert(&mut self, name: &'a str, metric: Metric);
    fn labels(&self) -> HashMap<String, String>;
}

#[derive(Default)]
pub struct StdRegistry<'a> {
    metrics: HashMap<&'a str, Metric>,
    reporter: HashMap<&'a str, Box<Reporter>>,
    labels: HashMap<String, String>,
}

// Specific stuff for registry goes here
impl<'a> Registry<'a> for StdRegistry<'a> {
    fn add_scheduled_reporter(&mut self, reporter: Box<Reporter>) {
        let reporter_name = reporter.get_unique_reporter_name();
        self.reporter.insert(reporter_name, reporter);
    }

    fn get(&'a self, name: &'a str) -> &'a Metric {
        &self.metrics[name]
    }

    fn insert(&mut self, name: &'a str, metric: Metric) {
        self.metrics.insert(name, metric);
    }

    fn get_metrics_names(&self) -> Vec<&str> {
        self.metrics.keys().cloned().collect()
    }

    fn labels(&self) -> HashMap<String, String> {
        self.labels.clone()
    }
}

impl<'a> StdRegistry<'a> {
    pub fn new_with_labels(labels: HashMap<String, String>) -> Self {
        StdRegistry {
            metrics: HashMap::new(),
            reporter: HashMap::new(),
            labels: labels,
        }
    }

    pub fn new() -> Self {
        StdRegistry {
            metrics: HashMap::new(),
            reporter: HashMap::new(),
            labels: HashMap::new(),
        }
    }
}

#[cfg(test)]
mod test {
    use metrics::{Counter, Gauge, Meter, Metric, StdCounter, StdGauge, StdMeter};
    use registry::{Registry, StdRegistry};
    use histogram::Histogram;

    // TODO add labels tests

    #[test]
    fn meter() {
        let mut r = StdRegistry::new();
        let m = StdMeter::new();
        m.mark(100);
        r.insert("meter1", Metric::Meter(m.clone()));
        m.mark(200);
        if let Metric::Meter(ref meter) = *r.get("meter1") {
            assert_eq!(meter.snapshot().count, 300);
        } else {
            panic!("We didn't get a meter back.");
        }
    }

    #[test]
    fn gauge() {
        let mut r = StdRegistry::new();
        let g = StdGauge::new();
        g.set(3);
        r.insert("gauge1", Metric::Gauge(g.clone()));
        g.set(5);
        if let Metric::Gauge(ref gauge) = *r.get("gauge1") {
            assert_eq!(gauge.snapshot().value, 5);
        } else {
            panic!("We didn't get a gauge back.");
        }
    }

    #[test]
    fn counter() {
        let mut r = StdRegistry::new();
        let c = StdCounter::new();
        c.add(1);
        r.insert("counter1", Metric::Counter(c.clone()));
        c.inc();
        if let Metric::Counter(ref counter) = *r.get("counter1") {
            assert_eq!(counter.snapshot().value, 2);
        } else {
            panic!("We didn't get a counter back.");
        }
    }

    #[test]
    fn histogram() {
        let mut r = StdRegistry::new();
        let mut h = Histogram::configure()
            .max_value(100)
            .precision(1)
            .build()
            .unwrap();
        h.increment_by(1, 1).unwrap();
        r.insert("histogram", Metric::Histogram(h));
    }
}
<|endoftext|>