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

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, false);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = false;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                // FIXME: Need to fix, there are more than just unknown keys.
                _ => {
                    wrong_key = true;
                }
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .eq(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self
            .upper_mark
            .saturating_add(self.content_rows.into())
            .min(self.line_count.saturating_sub(self.content_rows.into()));
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: bool) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: bool) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, true) => "[Press 'h' for instructions. (unimplemented)]".to_string(),
            (true, false) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, true) => format!("{}{}", status, BELL),
            (false, false) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<commit_msg>Now displays the unknown key entered<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, None);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = None;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char(k),
                    ..
                }) => wrong_key = Some(k),
                _ => continue,
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .eq(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self
            .upper_mark
            .saturating_add(self.content_rows.into())
            .min(self.line_count.saturating_sub(self.content_rows.into()));
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: Option<char>) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: Option<char>) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, Some(key)) => {
                format!(
                    "{} [Unknown key: '{}'. Press 'h' for instructions. (unimplemented)]",
                    status, key
                )
            }
            (true, None) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, Some(_)) => format!("{}{}", status, BELL),
            (false, None) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<|endoftext|>


<commit_before>use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};
use opus::{
    packet as opus_packet,
    Application as CodingMode,
    Channels,
    Decoder as OpusDecoder,
    Encoder as OpusEncoder,
};
use sodiumoxide::crypto::secretbox::{self, Key, Nonce};
use std::collections::HashMap;
use std::io::Write;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::mpsc::{self, Receiver as MpscReceiver, Sender as MpscSender};
use std::sync::{Arc, Mutex};
use std::thread::{self, Builder as ThreadBuilder, JoinHandle};
use std::time::Duration;
use super::audio::{AudioReceiver, AudioSource, AudioType, HEADER_LEN, SAMPLE_RATE};
use super::connection_info::ConnectionInfo;
use super::{payload, VoiceError, CRYPTO_MODE};
use websocket::client::Url as WebsocketUrl;
use websocket::sync::client::ClientBuilder;
use websocket::sync::stream::{AsTcpStream, TcpStream, TlsStream};
use websocket::sync::Client as WsClient;
use internal::prelude::*;
use internal::ws_impl::{ReceiverExt, SenderExt};
use internal::Timer;
use model::event::VoiceEvent;
use model::UserId;

type Client = WsClient<TlsStream<TcpStream>>;

enum ReceiverStatus {
    Udp(Vec<u8>),
    Websocket(VoiceEvent),
}

#[allow(dead_code)]
struct ThreadItems {
    rx: MpscReceiver<ReceiverStatus>,
    udp_close_sender: MpscSender<i32>,
    udp_thread: JoinHandle<()>,
    ws_close_sender: MpscSender<i32>,
    ws_thread: JoinHandle<()>,
}

#[allow(dead_code)]
pub struct Connection {
    audio_timer: Timer,
    client: Arc<Mutex<Client>>,
    decoder_map: HashMap<(u32, Channels), OpusDecoder>,
    destination: SocketAddr,
    encoder: OpusEncoder,
    encoder_stereo: bool,
    keepalive_timer: Timer,
    key: Key,
    sequence: u16,
    silence_frames: u8,
    speaking: bool,
    ssrc: u32,
    thread_items: ThreadItems,
    timestamp: u32,
    udp: UdpSocket,
    user_id: UserId,
}

impl Connection {
    pub fn new(mut info: ConnectionInfo) -> Result<Connection> {
        let url = generate_url(&mut info.endpoint)?;

        let mut client = ClientBuilder::from_url(&url).connect_secure(None)?;

        client.send_json(&payload::build_identify(&info))?;

        let hello = loop {
            match client.recv_json(VoiceEvent::decode)? {
                Some(VoiceEvent::Hello(received_hello)) => {
                    break received_hello;
                },
                Some(VoiceEvent::Heartbeat(_)) => continue,
                other => {
                    debug!("[Voice] Expected hello/heartbeat; got: {:?}", other);

                    return Err(Error::Voice(VoiceError::ExpectedHandshake));
                },
            }
        };

        if !has_valid_mode(&hello.modes) {
            return Err(Error::Voice(VoiceError::VoiceModeUnavailable));
        }

        let destination = (&info.endpoint[..], hello.port)
            .to_socket_addrs()?
            .next()
            .ok_or(Error::Voice(VoiceError::HostnameResolve))?;

        // Important to note here: the length of the packet can be of either 4
        // or 70 bytes. If it is 4 bytes, then we need to send a 70-byte packet
        // to determine the IP.
        //
        // Past the initial 4 bytes, the packet _must_ be completely empty data.
        //
        // The returned packet will be a null-terminated string of the IP, and
        // the port encoded in LE in the last two bytes of the packet.
        let udp = UdpSocket::bind("0.0.0.0:0")?;

        {
            let mut bytes = [0; 70];

            (&mut bytes[..]).write_u32::<BigEndian>(hello.ssrc)?;
            udp.send_to(&bytes, destination)?;

            let mut bytes = [0; 256];
            let (len, _addr) = udp.recv_from(&mut bytes)?;

            // Find the position in the bytes that contains the first byte of 0,
            // indicating the "end of the address".
            let index = bytes
                .iter()
                .skip(4)
                .position(|&x| x == 0)
                .ok_or(Error::Voice(VoiceError::FindingByte))?;

            let pos = 4 + index;
            let addr = String::from_utf8_lossy(&bytes[4..pos]);
            let port_pos = len - 2;
            let port = (&bytes[port_pos..]).read_u16::<LittleEndian>()?;

            client
                .send_json(&payload::build_select_protocol(addr, port))?;
        }

        let key = encryption_key(&mut client)?;

        let _ = client
            .stream_ref()
            .as_tcp()
            .set_read_timeout(Some(Duration::from_millis(25)));

        let mutexed_client = Arc::new(Mutex::new(client));
        let thread_items = start_threads(mutexed_client.clone(), &udp)?;

        info!("[Voice] Connected to: {}", info.endpoint);

        let encoder = OpusEncoder::new(SAMPLE_RATE, Channels::Mono, CodingMode::Audio)?;

        Ok(Connection {
            audio_timer: Timer::new(1000 * 60 * 4),
            client: mutexed_client,
            decoder_map: HashMap::new(),
            destination: destination,
            encoder: encoder,
            encoder_stereo: false,
            key: key,
            keepalive_timer: Timer::new(hello.heartbeat_interval),
            udp: udp,
            sequence: 0,
            silence_frames: 0,
            speaking: false,
            ssrc: hello.ssrc,
            thread_items: thread_items,
            timestamp: 0,
            user_id: info.user_id,
        })
    }

    #[allow(unused_variables)]
    pub fn cycle(&mut self,
                 source: &mut Option<Box<AudioSource>>,
                 receiver: &mut Option<Box<AudioReceiver>>,
                 audio_timer: &mut Timer)
                 -> Result<()> {
        let mut buffer = [0i16; 960 * 2];
        let mut packet = [0u8; 512];
        let mut nonce = secretbox::Nonce([0; 24]);

        if let Some(receiver) = receiver.as_mut() {
            while let Ok(status) = self.thread_items.rx.try_recv() {
                match status {
                    ReceiverStatus::Udp(packet) => {
                        let mut handle = &packet[2..];
                        let seq = handle.read_u16::<BigEndian>()?;
                        let timestamp = handle.read_u32::<BigEndian>()?;
                        let ssrc = handle.read_u32::<BigEndian>()?;

                        nonce.0[..HEADER_LEN]
                            .clone_from_slice(&packet[..HEADER_LEN]);

                        if let Ok(decrypted) =
                            secretbox::open(&packet[HEADER_LEN..], &nonce, &self.key) {
                            let channels = opus_packet::get_nb_channels(&decrypted)?;

                            let entry =
                                self.decoder_map.entry((ssrc, channels)).or_insert_with(
                                    || OpusDecoder::new(SAMPLE_RATE, channels).unwrap(),
                                );

                            let len = entry.decode(&decrypted, &mut buffer, false)?;

                            let is_stereo = channels == Channels::Stereo;

                            let b = if is_stereo { len * 2 } else { len };

                            receiver
                                .voice_packet(ssrc, seq, timestamp, is_stereo, &buffer[..b]);
                        }
                    },
                    ReceiverStatus::Websocket(VoiceEvent::Speaking(ev)) => {
                        receiver.speaking_update(ev.ssrc, ev.user_id.0, ev.speaking);
                    },
                    ReceiverStatus::Websocket(other) => {
                        info!("[Voice] Received other websocket data: {:?}", other);
                    },
                }
            }
        } else {
            loop {
                if self.thread_items.rx.try_recv().is_err() {
                    break;
                }
            }
        }

        // Send the voice websocket keepalive if it's time
        if self.keepalive_timer.check() {
            self.client
                .lock()
                .unwrap()
                .send_json(&payload::build_keepalive())?;
        }

        // Send UDP keepalive if it's time
        if self.audio_timer.check() {
            let mut bytes = [0; 4];
            (&mut bytes[..]).write_u32::<BigEndian>(self.ssrc)?;
            self.udp.send_to(&bytes, self.destination)?;
        }

        let mut opus_frame = Vec::new();

        let len = match source.as_mut() {
            Some(stream) => {
                let is_stereo = stream.is_stereo();

                if is_stereo != self.encoder_stereo {
                    let channels = if is_stereo {
                        Channels::Stereo
                    } else {
                        Channels::Mono
                    };
                    self.encoder = OpusEncoder::new(SAMPLE_RATE, channels, CodingMode::Audio)?;
                    self.encoder_stereo = is_stereo;
                }

                match stream.get_type() {
                    AudioType::Opus => match stream.read_opus_frame() {
                        Some(frame) => {
                            opus_frame = frame;
                            opus_frame.len()
                        },
                        None => 0,
                    },
                    AudioType::Pcm => {
                        let buffer_len = if is_stereo { 960 * 2 } else { 960 };

                        match stream.read_pcm_frame(&mut buffer[..buffer_len]) {
                            Some(len) => len,
                            None => 0,
                        }
                    },
                }
            },
            None => 0,
        };

        if len == 0 {
            self.set_speaking(false)?;

            if self.silence_frames > 0 {
                self.silence_frames -= 1;

                for value in &mut buffer[..] {
                    *value = 0;
                }
            } else {
                audio_timer.await();

                return Ok(());
            }
        } else {
            self.silence_frames = 5;

            for value in &mut buffer[len..] {
                *value = 0;
            }
        }

        self.set_speaking(true)?;

        let index = self.prep_packet(&mut packet, buffer, &opus_frame, nonce)?;
        audio_timer.await();

        self.udp.send_to(&packet[..index], self.destination)?;
        self.audio_timer.reset();

        Ok(())
    }

    fn prep_packet(&mut self,
                   packet: &mut [u8; 512],
                   buffer: [i16; 1920],
                   opus_frame: &[u8],
                   mut nonce: Nonce)
                   -> Result<usize> {
        {
            let mut cursor = &mut packet[..HEADER_LEN];
            cursor.write_all(&[0x80, 0x78])?;
            cursor.write_u16::<BigEndian>(self.sequence)?;
            cursor.write_u32::<BigEndian>(self.timestamp)?;
            cursor.write_u32::<BigEndian>(self.ssrc)?;
        }

        nonce.0[..HEADER_LEN]
            .clone_from_slice(&packet[..HEADER_LEN]);

        let sl_index = packet.len() - 16;
        let buffer_len = if self.encoder_stereo { 960 * 2 } else { 960 };

        let len = if opus_frame.is_empty() {
            self.encoder
                .encode(&buffer[..buffer_len], &mut packet[HEADER_LEN..sl_index])?
        } else {
            let len = opus_frame.len();
            packet[HEADER_LEN..HEADER_LEN + len]
                .clone_from_slice(opus_frame);
            len
        };

        let crypted = {
            let slice = &packet[HEADER_LEN..HEADER_LEN + len];
            secretbox::seal(slice, &nonce, &self.key)
        };
        let index = HEADER_LEN + crypted.len();
        packet[HEADER_LEN..index].clone_from_slice(&crypted);

        self.sequence = self.sequence.wrapping_add(1);
        self.timestamp = self.timestamp.wrapping_add(960);

        Ok(HEADER_LEN + crypted.len())
    }

    fn set_speaking(&mut self, speaking: bool) -> Result<()> {
        if self.speaking == speaking {
            return Ok(());
        }

        self.speaking = speaking;

        self.client
            .lock()
            .unwrap()
            .send_json(&payload::build_speaking(speaking))
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        let _ = self.thread_items.udp_close_sender.send(0);
        let _ = self.thread_items.ws_close_sender.send(0);

        info!("[Voice] Disconnected");
    }
}

fn generate_url(endpoint: &mut String) -> Result<WebsocketUrl> {
    if endpoint.ends_with(":80") {
        let len = endpoint.len();

        endpoint.truncate(len - 3);
    }

    WebsocketUrl::parse(&format!("wss://{}", endpoint))
        .or(Err(Error::Voice(VoiceError::EndpointUrl)))
}

#[inline]
fn encryption_key(client: &mut Client) -> Result<Key> {
    loop {
        match client.recv_json(VoiceEvent::decode)? {
            Some(VoiceEvent::Ready(ready)) => {
                if ready.mode != CRYPTO_MODE {
                    return Err(Error::Voice(VoiceError::VoiceModeInvalid));
                }

                return Key::from_slice(&ready.secret_key)
                    .ok_or(Error::Voice(VoiceError::KeyGen));
            },
            Some(VoiceEvent::Unknown(op, value)) => {
                debug!(
                    "[Voice] Expected ready for key; got: op{}/v{:?}",
                    op.num(),
                    value
                );
            },
            _ => {},
        }
    }
}

#[inline]
fn has_valid_mode<T, It> (modes: It) -> bool
where T: PartialEq<str>,
      It : IntoIterator<Item=T>
{
    modes.into_iter().any(|s| s == CRYPTO_MODE)
}

#[inline]
fn start_threads(client: Arc<Mutex<Client>>, udp: &UdpSocket) -> Result<ThreadItems> {
    let (udp_close_sender, udp_close_reader) = mpsc::channel();
    let (ws_close_sender, ws_close_reader) = mpsc::channel();

    let current_thread = thread::current();
    let thread_name = current_thread.name().unwrap_or("serenity voice");

    let (tx, rx) = mpsc::channel();
    let tx_clone = tx.clone();
    let udp_clone = udp.try_clone()?;

    let udp_thread = ThreadBuilder::new()
        .name(format!("{} UDP", thread_name))
        .spawn(move || {
            let _ = udp_clone.set_read_timeout(Some(Duration::from_millis(250)));

            let mut buffer = [0; 512];

            loop {
                if let Ok((len, _)) = udp_clone.recv_from(&mut buffer) {
                    let piece = buffer[..len].to_vec();
                    let send = tx.send(ReceiverStatus::Udp(piece));

                    if send.is_err() {
                        return;
                    }
                } else if udp_close_reader.try_recv().is_ok() {
                    return;
                }
            }
        })?;

    let ws_thread = ThreadBuilder::new()
        .name(format!("{} WS", thread_name))
        .spawn(move || loop {
            while let Ok(Some(msg)) = client.lock().unwrap().recv_json(VoiceEvent::decode) {
                if tx_clone.send(ReceiverStatus::Websocket(msg)).is_ok() {
                    return;
                }
            }

            if ws_close_reader.try_recv().is_ok() {
                return;
            }

            thread::sleep(Duration::from_millis(25));
        })?;

    Ok(ThreadItems {
        rx: rx,
        udp_close_sender: udp_close_sender,
        udp_thread: udp_thread,
        ws_close_sender: ws_close_sender,
        ws_thread: ws_thread,
    })
}
<commit_msg>Add the lifetime again<commit_after>use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};
use opus::{
    packet as opus_packet,
    Application as CodingMode,
    Channels,
    Decoder as OpusDecoder,
    Encoder as OpusEncoder,
};
use sodiumoxide::crypto::secretbox::{self, Key, Nonce};
use std::collections::HashMap;
use std::io::Write;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::mpsc::{self, Receiver as MpscReceiver, Sender as MpscSender};
use std::sync::{Arc, Mutex};
use std::thread::{self, Builder as ThreadBuilder, JoinHandle};
use std::time::Duration;
use super::audio::{AudioReceiver, AudioSource, AudioType, HEADER_LEN, SAMPLE_RATE};
use super::connection_info::ConnectionInfo;
use super::{payload, VoiceError, CRYPTO_MODE};
use websocket::client::Url as WebsocketUrl;
use websocket::sync::client::ClientBuilder;
use websocket::sync::stream::{AsTcpStream, TcpStream, TlsStream};
use websocket::sync::Client as WsClient;
use internal::prelude::*;
use internal::ws_impl::{ReceiverExt, SenderExt};
use internal::Timer;
use model::event::VoiceEvent;
use model::UserId;

type Client = WsClient<TlsStream<TcpStream>>;

enum ReceiverStatus {
    Udp(Vec<u8>),
    Websocket(VoiceEvent),
}

#[allow(dead_code)]
struct ThreadItems {
    rx: MpscReceiver<ReceiverStatus>,
    udp_close_sender: MpscSender<i32>,
    udp_thread: JoinHandle<()>,
    ws_close_sender: MpscSender<i32>,
    ws_thread: JoinHandle<()>,
}

#[allow(dead_code)]
pub struct Connection {
    audio_timer: Timer,
    client: Arc<Mutex<Client>>,
    decoder_map: HashMap<(u32, Channels), OpusDecoder>,
    destination: SocketAddr,
    encoder: OpusEncoder,
    encoder_stereo: bool,
    keepalive_timer: Timer,
    key: Key,
    sequence: u16,
    silence_frames: u8,
    speaking: bool,
    ssrc: u32,
    thread_items: ThreadItems,
    timestamp: u32,
    udp: UdpSocket,
    user_id: UserId,
}

impl Connection {
    pub fn new(mut info: ConnectionInfo) -> Result<Connection> {
        let url = generate_url(&mut info.endpoint)?;

        let mut client = ClientBuilder::from_url(&url).connect_secure(None)?;

        client.send_json(&payload::build_identify(&info))?;

        let hello = loop {
            match client.recv_json(VoiceEvent::decode)? {
                Some(VoiceEvent::Hello(received_hello)) => {
                    break received_hello;
                },
                Some(VoiceEvent::Heartbeat(_)) => continue,
                other => {
                    debug!("[Voice] Expected hello/heartbeat; got: {:?}", other);

                    return Err(Error::Voice(VoiceError::ExpectedHandshake));
                },
            }
        };

        if !has_valid_mode(&hello.modes) {
            return Err(Error::Voice(VoiceError::VoiceModeUnavailable));
        }

        let destination = (&info.endpoint[..], hello.port)
            .to_socket_addrs()?
            .next()
            .ok_or(Error::Voice(VoiceError::HostnameResolve))?;

        // Important to note here: the length of the packet can be of either 4
        // or 70 bytes. If it is 4 bytes, then we need to send a 70-byte packet
        // to determine the IP.
        //
        // Past the initial 4 bytes, the packet _must_ be completely empty data.
        //
        // The returned packet will be a null-terminated string of the IP, and
        // the port encoded in LE in the last two bytes of the packet.
        let udp = UdpSocket::bind("0.0.0.0:0")?;

        {
            let mut bytes = [0; 70];

            (&mut bytes[..]).write_u32::<BigEndian>(hello.ssrc)?;
            udp.send_to(&bytes, destination)?;

            let mut bytes = [0; 256];
            let (len, _addr) = udp.recv_from(&mut bytes)?;

            // Find the position in the bytes that contains the first byte of 0,
            // indicating the "end of the address".
            let index = bytes
                .iter()
                .skip(4)
                .position(|&x| x == 0)
                .ok_or(Error::Voice(VoiceError::FindingByte))?;

            let pos = 4 + index;
            let addr = String::from_utf8_lossy(&bytes[4..pos]);
            let port_pos = len - 2;
            let port = (&bytes[port_pos..]).read_u16::<LittleEndian>()?;

            client
                .send_json(&payload::build_select_protocol(addr, port))?;
        }

        let key = encryption_key(&mut client)?;

        let _ = client
            .stream_ref()
            .as_tcp()
            .set_read_timeout(Some(Duration::from_millis(25)));

        let mutexed_client = Arc::new(Mutex::new(client));
        let thread_items = start_threads(mutexed_client.clone(), &udp)?;

        info!("[Voice] Connected to: {}", info.endpoint);

        let encoder = OpusEncoder::new(SAMPLE_RATE, Channels::Mono, CodingMode::Audio)?;

        Ok(Connection {
            audio_timer: Timer::new(1000 * 60 * 4),
            client: mutexed_client,
            decoder_map: HashMap::new(),
            destination: destination,
            encoder: encoder,
            encoder_stereo: false,
            key: key,
            keepalive_timer: Timer::new(hello.heartbeat_interval),
            udp: udp,
            sequence: 0,
            silence_frames: 0,
            speaking: false,
            ssrc: hello.ssrc,
            thread_items: thread_items,
            timestamp: 0,
            user_id: info.user_id,
        })
    }

    #[allow(unused_variables)]
    pub fn cycle(&mut self,
                 source: &mut Option<Box<AudioSource>>,
                 receiver: &mut Option<Box<AudioReceiver>>,
                 audio_timer: &mut Timer)
                 -> Result<()> {
        let mut buffer = [0i16; 960 * 2];
        let mut packet = [0u8; 512];
        let mut nonce = secretbox::Nonce([0; 24]);

        if let Some(receiver) = receiver.as_mut() {
            while let Ok(status) = self.thread_items.rx.try_recv() {
                match status {
                    ReceiverStatus::Udp(packet) => {
                        let mut handle = &packet[2..];
                        let seq = handle.read_u16::<BigEndian>()?;
                        let timestamp = handle.read_u32::<BigEndian>()?;
                        let ssrc = handle.read_u32::<BigEndian>()?;

                        nonce.0[..HEADER_LEN]
                            .clone_from_slice(&packet[..HEADER_LEN]);

                        if let Ok(decrypted) =
                            secretbox::open(&packet[HEADER_LEN..], &nonce, &self.key) {
                            let channels = opus_packet::get_nb_channels(&decrypted)?;

                            let entry =
                                self.decoder_map.entry((ssrc, channels)).or_insert_with(
                                    || OpusDecoder::new(SAMPLE_RATE, channels).unwrap(),
                                );

                            let len = entry.decode(&decrypted, &mut buffer, false)?;

                            let is_stereo = channels == Channels::Stereo;

                            let b = if is_stereo { len * 2 } else { len };

                            receiver
                                .voice_packet(ssrc, seq, timestamp, is_stereo, &buffer[..b]);
                        }
                    },
                    ReceiverStatus::Websocket(VoiceEvent::Speaking(ev)) => {
                        receiver.speaking_update(ev.ssrc, ev.user_id.0, ev.speaking);
                    },
                    ReceiverStatus::Websocket(other) => {
                        info!("[Voice] Received other websocket data: {:?}", other);
                    },
                }
            }
        } else {
            loop {
                if self.thread_items.rx.try_recv().is_err() {
                    break;
                }
            }
        }

        // Send the voice websocket keepalive if it's time
        if self.keepalive_timer.check() {
            self.client
                .lock()
                .unwrap()
                .send_json(&payload::build_keepalive())?;
        }

        // Send UDP keepalive if it's time
        if self.audio_timer.check() {
            let mut bytes = [0; 4];
            (&mut bytes[..]).write_u32::<BigEndian>(self.ssrc)?;
            self.udp.send_to(&bytes, self.destination)?;
        }

        let mut opus_frame = Vec::new();

        let len = match source.as_mut() {
            Some(stream) => {
                let is_stereo = stream.is_stereo();

                if is_stereo != self.encoder_stereo {
                    let channels = if is_stereo {
                        Channels::Stereo
                    } else {
                        Channels::Mono
                    };
                    self.encoder = OpusEncoder::new(SAMPLE_RATE, channels, CodingMode::Audio)?;
                    self.encoder_stereo = is_stereo;
                }

                match stream.get_type() {
                    AudioType::Opus => match stream.read_opus_frame() {
                        Some(frame) => {
                            opus_frame = frame;
                            opus_frame.len()
                        },
                        None => 0,
                    },
                    AudioType::Pcm => {
                        let buffer_len = if is_stereo { 960 * 2 } else { 960 };

                        match stream.read_pcm_frame(&mut buffer[..buffer_len]) {
                            Some(len) => len,
                            None => 0,
                        }
                    },
                }
            },
            None => 0,
        };

        if len == 0 {
            self.set_speaking(false)?;

            if self.silence_frames > 0 {
                self.silence_frames -= 1;

                for value in &mut buffer[..] {
                    *value = 0;
                }
            } else {
                audio_timer.await();

                return Ok(());
            }
        } else {
            self.silence_frames = 5;

            for value in &mut buffer[len..] {
                *value = 0;
            }
        }

        self.set_speaking(true)?;

        let index = self.prep_packet(&mut packet, buffer, &opus_frame, nonce)?;
        audio_timer.await();

        self.udp.send_to(&packet[..index], self.destination)?;
        self.audio_timer.reset();

        Ok(())
    }

    fn prep_packet(&mut self,
                   packet: &mut [u8; 512],
                   buffer: [i16; 1920],
                   opus_frame: &[u8],
                   mut nonce: Nonce)
                   -> Result<usize> {
        {
            let mut cursor = &mut packet[..HEADER_LEN];
            cursor.write_all(&[0x80, 0x78])?;
            cursor.write_u16::<BigEndian>(self.sequence)?;
            cursor.write_u32::<BigEndian>(self.timestamp)?;
            cursor.write_u32::<BigEndian>(self.ssrc)?;
        }

        nonce.0[..HEADER_LEN]
            .clone_from_slice(&packet[..HEADER_LEN]);

        let sl_index = packet.len() - 16;
        let buffer_len = if self.encoder_stereo { 960 * 2 } else { 960 };

        let len = if opus_frame.is_empty() {
            self.encoder
                .encode(&buffer[..buffer_len], &mut packet[HEADER_LEN..sl_index])?
        } else {
            let len = opus_frame.len();
            packet[HEADER_LEN..HEADER_LEN + len]
                .clone_from_slice(opus_frame);
            len
        };

        let crypted = {
            let slice = &packet[HEADER_LEN..HEADER_LEN + len];
            secretbox::seal(slice, &nonce, &self.key)
        };
        let index = HEADER_LEN + crypted.len();
        packet[HEADER_LEN..index].clone_from_slice(&crypted);

        self.sequence = self.sequence.wrapping_add(1);
        self.timestamp = self.timestamp.wrapping_add(960);

        Ok(HEADER_LEN + crypted.len())
    }

    fn set_speaking(&mut self, speaking: bool) -> Result<()> {
        if self.speaking == speaking {
            return Ok(());
        }

        self.speaking = speaking;

        self.client
            .lock()
            .unwrap()
            .send_json(&payload::build_speaking(speaking))
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        let _ = self.thread_items.udp_close_sender.send(0);
        let _ = self.thread_items.ws_close_sender.send(0);

        info!("[Voice] Disconnected");
    }
}

fn generate_url(endpoint: &mut String) -> Result<WebsocketUrl> {
    if endpoint.ends_with(":80") {
        let len = endpoint.len();

        endpoint.truncate(len - 3);
    }

    WebsocketUrl::parse(&format!("wss://{}", endpoint))
        .or(Err(Error::Voice(VoiceError::EndpointUrl)))
}

#[inline]
fn encryption_key(client: &mut Client) -> Result<Key> {
    loop {
        match client.recv_json(VoiceEvent::decode)? {
            Some(VoiceEvent::Ready(ready)) => {
                if ready.mode != CRYPTO_MODE {
                    return Err(Error::Voice(VoiceError::VoiceModeInvalid));
                }

                return Key::from_slice(&ready.secret_key)
                    .ok_or(Error::Voice(VoiceError::KeyGen));
            },
            Some(VoiceEvent::Unknown(op, value)) => {
                debug!(
                    "[Voice] Expected ready for key; got: op{}/v{:?}",
                    op.num(),
                    value
                );
            },
            _ => {},
        }
    }
}

#[inline]
fn has_valid_mode<T, It> (modes: It) -> bool
where T: for<'a> PartialEq<&'a str>,
      It : IntoIterator<Item=T>
{
    modes.into_iter().any(|s| s == CRYPTO_MODE)
}

#[inline]
fn start_threads(client: Arc<Mutex<Client>>, udp: &UdpSocket) -> Result<ThreadItems> {
    let (udp_close_sender, udp_close_reader) = mpsc::channel();
    let (ws_close_sender, ws_close_reader) = mpsc::channel();

    let current_thread = thread::current();
    let thread_name = current_thread.name().unwrap_or("serenity voice");

    let (tx, rx) = mpsc::channel();
    let tx_clone = tx.clone();
    let udp_clone = udp.try_clone()?;

    let udp_thread = ThreadBuilder::new()
        .name(format!("{} UDP", thread_name))
        .spawn(move || {
            let _ = udp_clone.set_read_timeout(Some(Duration::from_millis(250)));

            let mut buffer = [0; 512];

            loop {
                if let Ok((len, _)) = udp_clone.recv_from(&mut buffer) {
                    let piece = buffer[..len].to_vec();
                    let send = tx.send(ReceiverStatus::Udp(piece));

                    if send.is_err() {
                        return;
                    }
                } else if udp_close_reader.try_recv().is_ok() {
                    return;
                }
            }
        })?;

    let ws_thread = ThreadBuilder::new()
        .name(format!("{} WS", thread_name))
        .spawn(move || loop {
            while let Ok(Some(msg)) = client.lock().unwrap().recv_json(VoiceEvent::decode) {
                if tx_clone.send(ReceiverStatus::Websocket(msg)).is_ok() {
                    return;
                }
            }

            if ws_close_reader.try_recv().is_ok() {
                return;
            }

            thread::sleep(Duration::from_millis(25));
        })?;

    Ok(ThreadItems {
        rx: rx,
        udp_close_sender: udp_close_sender,
        udp_thread: udp_thread,
        ws_close_sender: ws_close_sender,
        ws_thread: ws_thread,
    })
}
<|endoftext|>


<commit_before>#![crate_name = "toy_aa"]
#![crate_type = "dylib"]

use std::fmt::Show;
use std::mem::{replace, swap};

type Link<T> = Option<Box<T>>;

struct Node<K, V> {
    key: K,
    value: V,
    left: Link<Node<K, V>>,
    right: Link<Node<K, V>>,
    level: uint
}

pub struct Tree<K, V> {
    root: Link<Node<K, V>>,
}

impl<K: Ord, V> Tree<K, V> {
    pub fn new() -> Tree<K, V> {
        Tree { root: None }
    }

    fn is_bst(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_bst()
        }
    }

    fn is_aa(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_aa()
        }
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        match self.root {
            None => None,
            Some(ref r) => r.find(key)
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.root {
            None => {
                self.root = Some(box Node::new(key, value));
                None
            },
            Some(ref mut r) => r.insert(key, value)
        }
    }
}

impl<K: Ord, V> Node<K, V> {
    pub fn new(key: K, value: V) -> Node<K, V> {
        Node { key: key, value: value, left: None, right: None, level: 1 }
    }

    fn max(&self) -> &K {
        match self.right {
            None => &self.key,
            Some(ref n) => n.max(),
        }
    }

    fn min(&self) -> &K {
        match self.left {
            None => &self.key,
            Some(ref n) => n.min(),
        }
    }

    pub fn is_bst(&self) -> bool {
        self.left.as_ref().map_or(true, |n| n.is_bst() && *n.max() < self.key)
        && self.right.as_ref().map_or(true, |n| n.is_bst() && *n.min() > self.key)
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        let ch = match key.cmp(&self.key) {
            Equal => return Some(&self.value),
            Less => & self.left,
            Greater => & self.right,
        };

        match *ch {
            None => None,
            Some(ref b) => b.find(key),
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let inserted =
        {
            let ch = match key.cmp(&self.key) {
                Equal => {
                    self.key = key;
                    return Some(replace(&mut self.value, value))
                },
                Less => &mut self.left,
                Greater => &mut self.right,
            };

            match *ch {
                None => {
                    *ch = Some(box Node::new(key, value));
                    return None
                },
                Some(ref mut b) => {
                    b.insert(key, value)
                },
            }
        };

        skew(self);
        split(self);
        inserted

    }

    // To be an AA tree, it must be a binary search tree and, for all nodes n:
    //   - if n is missing a child, its level must be 1
    //   - the left child must have a level one less than n's level
    //   - the right child must have a level equal to or one less than n's level
    //   - the right child's right child must not have the same level as n's level
    fn is_aa(&self) -> bool {
        let lvl = self.level;

        self.is_bst()
            && !((self.left.is_none() || self.right.is_none()) && self.level != 1)
            && self.left.as_ref().map_or(true, |n| n.is_aa())
            && self.no_red_left_child()
            && self.right.as_ref().map_or(true,
                |n| n.is_aa() && (n.level == lvl || n.level + 1 == lvl)
                    && !(n.level == lvl && !n.no_red_right_child()))
    }

    fn no_red_left_child(&self) -> bool {
        match self.left {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    fn no_red_right_child(&self) -> bool {
        match self.right {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    /*

       In the below, 'o' represents black, '*' represents red.

       First, note that deleting a red leaf requires no rebalancing at all.

       ---

       There's only 1 possible AA tree of size 2:

       o
        \
         *

       Deleting the root or the leaf requires no rebalancing.

       ---

       For size 3 trees, we only have:

         o
        / \
       o   o

        - deleting left leaf means we have to decrease the level of root. this suffices to rebalance
        - deleting right leaf means we have to decrease level of root but then make left leaf the
          new root (and former root the child of new root)
        - deleting the root means right child is new root, left child becomes new root's right child

       ---

       For size 4 trees, we have either

         o       o
        / \     / \
       o   o   o   o
        \           \
         *           *

       There are 4 classes of deletion:

        - delete red leaf (see above)
        - delete red node's parent: make red a child of its former grandparent
        - delete root: in the first case, red needs to become root, so increase its level and make it
          the new root. in the second case, right child of root needs to become new root.
        - delete other child: in the first case, red needs to become root and the old root needs to
          be right child of new root. in the second case, decrease level of old root, increase level
          of right child and make old root the child of the right child (which is the new root)

       ---

       For size 5 trees, we have one of

         o        o
        / \      / \
       o   o    o   *
        \   \      / \
         *   *    o   o

       ===

       One fundamental operation we need is "promote the max from the left subtree (or min from right subtree!)
       this is used when we delete a root, for example. this seems essentially whenever we delete a node that
       has two children.

     */
}

// Remove red left child by rotating right
/*
     a      b
    /        \
   b    =>    a
    \        /
     c      c

  provided that a.level == b.level
*/
fn skew<K: Ord, V>(node: &mut Node<K, V>) {
    if node.left.is_some() && node.left.get_ref().level == node.level {
        let mut save = node.left.take_unwrap();
        swap(&mut node.left, &mut save.right); // save.right now None
        swap(node, &mut *save);
        node.right = Some(save);
    }
}

// Remove red right child of red right child by rotating left and increasing
// level of the parent
/*
    a            b
     \          / \
      b    =>  a   c
     / \        \
    d   c        d

  provided that a.level == c.level
*/
fn split<K: Ord, V>(node: &mut Node<K, V>) {
    if node.right.as_ref().map_or(false,
      |x| x.right.is_some() && x.right.get_ref().level == node.level) {
        let mut save = node.right.take_unwrap();
        swap(&mut node.right, &mut save.left); // save.left now None
        save.level += 1;
        swap(node, &mut *save);
        node.left = Some(save);
    }
}

fn print_node_depth<K: Show, V: Show>(node: &Link<Node<K,V>>, depth: uint) {
    let mut pre = "".to_string();
    if depth > 0 {
        for i in range(0, depth) {
            pre = pre + "   ";
        }
    }

    match *node {
        Some(ref n) => {
            println!("{}{}:{}*{}", pre, n.key, n.value, n.level);
            print_node_depth(&n.left, depth + 1);
            print_node_depth(&n.right, depth + 1);
        },
        None => println!("{}-", pre),
    }
}

pub fn print_tree<K: Show + Ord, V: Show>(tree: &Tree<K, V>) {
    print_node_depth(&tree.root, 0);
    println!("Is AA: {}", tree.is_aa());
    println!("------------");
}

mod test {
    use super::Tree;
    use std::rand;
    use std::rand::distributions::{IndependentSample, Range};

    #[test]
    fn test_find() {
        let mut t = Tree::new();
        assert_eq!(t.find(&1u), None);
        t.insert(1u, 'j');
        assert_eq!(t.find(&1u), Some(&'j'));

    }

    // testing whether we can find all the things we inserted
    #[test]
    fn test_insert() {
        let mut t: Tree<uint, u8> = Tree::new();
        for (i, c) in range(0u, 10).zip(range(b'a', b'z')) {
            t.insert(i, c);
        }

        for (ref i, ref c) in range(0u, 10).zip(range(b'a', b'z')) {
            assert_eq!(t.find(i), Some(c));
        }

        assert_eq!(t.find(&10u), None);
    }

    fn insert_n_check_aa(n: uint, between: Range<uint>, rng: &mut rand::TaskRng) {
        let mut t = Tree::new();

        for _ in range(0u, n) {
            let a = between.ind_sample(rng);
            println!("{}", a);
            t.insert(a, ());
        }

        assert!(t.is_aa());
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true
    #[test]
    fn test_insert_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 100_000);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true,
    // but this time some of the keys are repeated
    #[test]
    fn test_insert_dups_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 15);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }
}
<commit_msg>Turn skew and split into methods<commit_after>#![crate_name = "toy_aa"]
#![crate_type = "dylib"]

use std::fmt::Show;
use std::mem::{replace, swap};

type Link<T> = Option<Box<T>>;

struct Node<K, V> {
    key: K,
    value: V,
    left: Link<Node<K, V>>,
    right: Link<Node<K, V>>,
    level: uint
}

pub struct Tree<K, V> {
    root: Link<Node<K, V>>,
}

impl<K: Ord, V> Tree<K, V> {
    pub fn new() -> Tree<K, V> {
        Tree { root: None }
    }

    fn is_bst(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_bst()
        }
    }

    fn is_aa(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_aa()
        }
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        match self.root {
            None => None,
            Some(ref r) => r.find(key)
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.root {
            None => {
                self.root = Some(box Node::new(key, value));
                None
            },
            Some(ref mut r) => r.insert(key, value)
        }
    }
}

impl<K: Ord, V> Node<K, V> {
    pub fn new(key: K, value: V) -> Node<K, V> {
        Node { key: key, value: value, left: None, right: None, level: 1 }
    }

    fn max(&self) -> &K {
        match self.right {
            None => &self.key,
            Some(ref n) => n.max(),
        }
    }

    fn min(&self) -> &K {
        match self.left {
            None => &self.key,
            Some(ref n) => n.min(),
        }
    }

    pub fn is_bst(&self) -> bool {
        self.left.as_ref().map_or(true, |n| n.is_bst() && *n.max() < self.key)
        && self.right.as_ref().map_or(true, |n| n.is_bst() && *n.min() > self.key)
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        let ch = match key.cmp(&self.key) {
            Equal => return Some(&self.value),
            Less => & self.left,
            Greater => & self.right,
        };

        match *ch {
            None => None,
            Some(ref b) => b.find(key),
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let inserted =
        {
            let ch = match key.cmp(&self.key) {
                Equal => {
                    self.key = key;
                    return Some(replace(&mut self.value, value))
                },
                Less => &mut self.left,
                Greater => &mut self.right,
            };

            match *ch {
                None => {
                    *ch = Some(box Node::new(key, value));
                    return None
                },
                Some(ref mut b) => {
                    b.insert(key, value)
                },
            }
        };

        self.skew();
        self.split();
        inserted

    }

    // To be an AA tree, it must be a binary search tree and, for all nodes n:
    //   - if n is missing a child, its level must be 1
    //   - the left child must have a level one less than n's level
    //   - the right child must have a level equal to or one less than n's level
    //   - the right child's right child must not have the same level as n's level
    fn is_aa(&self) -> bool {
        let lvl = self.level;

        self.is_bst()
            && !((self.left.is_none() || self.right.is_none()) && self.level != 1)
            && self.left.as_ref().map_or(true, |n| n.is_aa())
            && self.no_red_left_child()
            && self.right.as_ref().map_or(true,
                |n| n.is_aa() && (n.level == lvl || n.level + 1 == lvl)
                    && !(n.level == lvl && !n.no_red_right_child()))
    }

    fn no_red_left_child(&self) -> bool {
        match self.left {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    fn no_red_right_child(&self) -> bool {
        match self.right {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    // Remove red left child by rotating right
    /*
         a      b
        /        \
       b    =>    a
        \        /
         c      c

      provided that a.level == b.level
    */
    fn skew(&mut self) {
        if self.left.is_some() && self.left.get_ref().level == self.level {
            let mut save = self.left.take_unwrap();
            swap(&mut self.left, &mut save.right); // save.right now None
            swap(self, &mut *save);
            self.right = Some(save);
        }
    }

    // Remove red right child of red right child by rotating left and increasing
    // level of the parent
    /*
        a            b
         \          / \
          b    =>  a   c
         / \        \
        d   c        d

      provided that a.level == c.level
    */
    fn split(&mut self) {
        if self.right.as_ref().map_or(false,
          |x| x.right.is_some() && x.right.get_ref().level == self.level) {
            let mut save = self.right.take_unwrap();
            swap(&mut self.right, &mut save.left); // save.left now None
            save.level += 1;
            swap(self, &mut *save);
            self.left = Some(save);
        }
    }

    /*

       In the below, 'o' represents black, '*' represents red.

       First, note that deleting a red leaf requires no rebalancing at all.

       ---

       There's only 1 possible AA tree of size 2:

       o
        \
         *

       Deleting the root or the leaf requires no rebalancing.

       ---

       For size 3 trees, we only have:

         o
        / \
       o   o

        - deleting left leaf means we have to decrease the level of root. this suffices to rebalance
        - deleting right leaf means we have to decrease level of root but then make left leaf the
          new root (and former root the child of new root)
        - deleting the root means right child is new root, left child becomes new root's right child

       ---

       For size 4 trees, we have either

         o       o
        / \     / \
       o   o   o   o
        \           \
         *           *

       There are 4 classes of deletion:

        - delete red leaf (see above)
        - delete red node's parent: make red a child of its former grandparent
        - delete root: in the first case, red needs to become root, so increase its level and make it
          the new root. in the second case, right child of root needs to become new root.
        - delete other child: in the first case, red needs to become root and the old root needs to
          be right child of new root. in the second case, decrease level of old root, increase level
          of right child and make old root the child of the right child (which is the new root)

       ---

       For size 5 trees, we have one of

         o        o
        / \      / \
       o   o    o   *
        \   \      / \
         *   *    o   o

       ===

       One fundamental operation we need is "promote the max from the left subtree (or min from right subtree!)
       this is used when we delete a root, for example. this seems essentially whenever we delete a node that
       has two children.

     */
}

fn print_node_depth<K: Show, V: Show>(node: &Link<Node<K,V>>, depth: uint) {
    let mut pre = "".to_string();
    if depth > 0 {
        for i in range(0, depth) {
            pre = pre + "   ";
        }
    }

    match *node {
        Some(ref n) => {
            println!("{}{}:{}*{}", pre, n.key, n.value, n.level);
            print_node_depth(&n.left, depth + 1);
            print_node_depth(&n.right, depth + 1);
        },
        None => println!("{}-", pre),
    }
}

pub fn print_tree<K: Show + Ord, V: Show>(tree: &Tree<K, V>) {
    print_node_depth(&tree.root, 0);
    println!("Is AA: {}", tree.is_aa());
    println!("------------");
}

mod test {
    use super::Tree;
    use std::rand;
    use std::rand::distributions::{IndependentSample, Range};

    #[test]
    fn test_find() {
        let mut t = Tree::new();
        assert_eq!(t.find(&1u), None);
        t.insert(1u, 'j');
        assert_eq!(t.find(&1u), Some(&'j'));

    }

    // testing whether we can find all the things we inserted
    #[test]
    fn test_insert() {
        let mut t: Tree<uint, u8> = Tree::new();
        for (i, c) in range(0u, 10).zip(range(b'a', b'z')) {
            t.insert(i, c);
        }

        for (ref i, ref c) in range(0u, 10).zip(range(b'a', b'z')) {
            assert_eq!(t.find(i), Some(c));
        }

        assert_eq!(t.find(&10u), None);
    }

    fn insert_n_check_aa(n: uint, between: Range<uint>, rng: &mut rand::TaskRng) {
        let mut t = Tree::new();

        for _ in range(0u, n) {
            let a = between.ind_sample(rng);
            println!("{}", a);
            t.insert(a, ());
        }

        assert!(t.is_aa());
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true
    #[test]
    fn test_insert_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 100_000);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true,
    // but this time some of the keys are repeated
    #[test]
    fn test_insert_dups_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 15);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }
}
<|endoftext|>


<commit_before>#![crate_name = "toy_aa"]
#![crate_type = "dylib"]

use std::fmt::{mod, Show};
use std::mem::{replace, swap};

type Link<T> = Option<Box<T>>;

struct Node<K, V> {
    key: K,
    value: V,
    left: Link<Node<K, V>>,
    right: Link<Node<K, V>>,
    level: uint
}

pub struct Tree<K, V> {
    root: Link<Node<K, V>>,
}

impl<K: Ord, V> Tree<K, V> {
    pub fn new() -> Tree<K, V> {
        Tree { root: None }
    }

    fn is_bst(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_bst()
        }
    }

    fn is_aa(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_aa()
        }
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        match self.root {
            None => None,
            Some(ref r) => r.find(key)
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.root {
            None => {
                self.root = Some(box Node::new(key, value));
                None
            },
            Some(ref mut r) => r.insert(key, value)
        }
    }
}

impl<K: Ord, V> Node<K, V> {
    pub fn new(key: K, value: V) -> Node<K, V> {
        Node { key: key, value: value, left: None, right: None, level: 1 }
    }

    fn max(&self) -> &K {
        match self.right {
            None => &self.key,
            Some(ref n) => n.max(),
        }
    }

    fn min(&self) -> &K {
        match self.left {
            None => &self.key,
            Some(ref n) => n.min(),
        }
    }

    pub fn is_bst(&self) -> bool {
        self.left.as_ref().map_or(true, |n| n.is_bst() && *n.max() < self.key)
        && self.right.as_ref().map_or(true, |n| n.is_bst() && *n.min() > self.key)
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        let ch = match key.cmp(&self.key) {
            Equal => return Some(&self.value),
            Less => & self.left,
            Greater => & self.right,
        };

        match *ch {
            None => None,
            Some(ref b) => b.find(key),
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let replaced = match key.cmp(&self.key) {
            Equal => {
                self.key = key;
                return Some(replace(&mut self.value, value))
            },
            Less =>
                match self.left {
                    None => {
                        self.left = Some(box Node::new(key, value));
                        None
                    },
                    Some(ref mut b) => b.insert(key, value),
                },
            Greater =>
                match self.right {
                    None => {
                        self.right = Some(box Node::new(key, value));
                        None
                    },
                    Some(ref mut b) => b.insert(key, value),
                },
        };

        self.skew();
        self.split();
        replaced
    }

    // To be an AA tree, it must be a binary search tree and, for all nodes n:
    //   - if n is missing a child, its level must be 1
    //   - the left child must have a level one less than n's level
    //   - the right child must have a level equal to or one less than n's level
    //   - the right child's right child must not have the same level as n's level
    fn is_aa(&self) -> bool {
        let lvl = self.level;

        self.is_bst()
            && !((self.left.is_none() || self.right.is_none()) && self.level != 1)
            && self.left.as_ref().map_or(true, |n| n.is_aa())
            && self.no_red_left_child()
            && self.right.as_ref().map_or(true,
                |n| n.is_aa() && (n.level == lvl || n.level + 1 == lvl)
                    && !(n.level == lvl && !n.no_red_right_child()))
    }

    fn no_red_left_child(&self) -> bool {
        match self.left {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    fn no_red_right_child(&self) -> bool {
        match self.right {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    // Remove red left child by rotating right
    /*
         a      b
        /        \
       b    =>    a
        \        /
         c      c

      provided that a.level == b.level
    */
    fn skew(&mut self) {
        if self.left.is_some() && self.left.get_ref().level == self.level {
            let mut save = self.left.take_unwrap();
            swap(&mut self.left, &mut save.right); // save.right now None
            swap(self, &mut *save);
            self.right = Some(save);
        }
    }

    // Remove red right child of red right child by rotating left and increasing
    // level of the parent
    /*
        a            b
         \          / \
          b    =>  a   c
         / \        \
        d   c        d

      provided that a.level == c.level
    */
    fn split(&mut self) {
        if self.right.as_ref().map_or(false,
          |x| x.right.is_some() && x.right.get_ref().level == self.level) {
            let mut save = self.right.take_unwrap();
            swap(&mut self.right, &mut save.left); // save.left now None
            save.level += 1;
            swap(self, &mut *save);
            self.left = Some(save);
        }
    }

    /*

       In the below, 'o' represents black, '*' represents red.

       First, note that deleting a red leaf requires no rebalancing at all.

       ---

       There's only 1 possible AA tree of size 2:

       o
        \
         *

       Deleting the root or the leaf requires no rebalancing.

       ---

       For size 3 trees, we only have:

         o
        / \
       o   o

        - deleting left leaf means we have to decrease the level of root. this suffices to rebalance
        - deleting right leaf means we have to decrease level of root but then make left leaf the
          new root (and former root the child of new root)
        - deleting the root means right child is new root, left child becomes new root's right child

       ---

       For size 4 trees, we have either

         o       o
        / \     / \
       o   o   o   o
        \           \
         *           *

       There are 4 classes of deletion:

        - delete red leaf (see above)
        - delete red node's parent: make red a child of its former grandparent
        - delete root: in the first case, red needs to become root, so increase its level and make it
          the new root. in the second case, right child of root needs to become new root.
        - delete other child: in the first case, red needs to become root and the old root needs to
          be right child of new root. in the second case, decrease level of old root, increase level
          of right child and make old root the child of the right child (which is the new root)

       ---

       For size 5 trees, we have one of

         o        o
        / \      / \
       o   o    o   *
        \   \      / \
         *   *    o   o

       ===

       One fundamental operation we need is "promote the max from the left subtree (or min from right subtree!)
       this is used when we delete a root, for example. this seems essentially whenever we delete a node that
       has two children.

     */
}

impl<K: Show, V: Show> Show for Node<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.key, self.value)
    }
}

fn print_node_depth<K: Show, V: Show>(node: &Link<Node<K,V>>, depth: uint) {
    let mut pre = "".to_string();
    if depth > 0 {
        for i in range(0, depth) {
            pre = pre + "   ";
        }
    }

    match *node {
        Some(ref n) => {
            println!("{}{}^{}", pre, n, n.level);
            print_node_depth(&n.left, depth + 1);
            print_node_depth(&n.right, depth + 1);
        },
        None => println!("{}-", pre),
    }
}

pub fn print_tree<K: Show + Ord, V: Show>(tree: &Tree<K, V>) {
    print_node_depth(&tree.root, 0);
    println!("Is AA: {}", tree.is_aa());
    println!("------------");
}

mod test {
    use super::Tree;
    use std::rand;
    use std::rand::distributions::{IndependentSample, Range};

    #[test]
    fn test_find() {
        let mut t = Tree::new();
        assert_eq!(t.find(&1u), None);
        t.insert(1u, 'j');
        assert_eq!(t.find(&1u), Some(&'j'));

    }

    // testing whether we can find all the things we inserted
    #[test]
    fn test_insert() {
        let mut t: Tree<uint, u8> = Tree::new();
        for (i, c) in range(0u, 10).zip(range(b'a', b'z')) {
            t.insert(i, c);
        }

        for (ref i, ref c) in range(0u, 10).zip(range(b'a', b'z')) {
            assert_eq!(t.find(i), Some(c));
        }

        assert_eq!(t.find(&10u), None);
    }

    fn insert_n_check_aa(n: uint, between: Range<uint>, rng: &mut rand::TaskRng) {
        let mut t = Tree::new();

        for _ in range(0u, n) {
            let a = between.ind_sample(rng);
            println!("{}", a);
            t.insert(a, ());
        }

        assert!(t.is_aa());
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true
    #[test]
    fn test_insert_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 100_000);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true,
    // but this time some of the keys are repeated
    #[test]
    fn test_insert_dups_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 15);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }
}
<commit_msg>Add cfg(test) around mod test so it only gets compiled for tests<commit_after>#![crate_name = "toy_aa"]
#![crate_type = "dylib"]

use std::fmt::{mod, Show};
use std::mem::{replace, swap};

type Link<T> = Option<Box<T>>;

struct Node<K, V> {
    key: K,
    value: V,
    left: Link<Node<K, V>>,
    right: Link<Node<K, V>>,
    level: uint
}

pub struct Tree<K, V> {
    root: Link<Node<K, V>>,
}

impl<K: Ord, V> Tree<K, V> {
    pub fn new() -> Tree<K, V> {
        Tree { root: None }
    }

    fn is_bst(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_bst()
        }
    }

    fn is_aa(&self) -> bool {
        match self.root {
            None => true,
            Some(ref r) => r.is_aa()
        }
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        match self.root {
            None => None,
            Some(ref r) => r.find(key)
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.root {
            None => {
                self.root = Some(box Node::new(key, value));
                None
            },
            Some(ref mut r) => r.insert(key, value)
        }
    }
}

impl<K: Ord, V> Node<K, V> {
    pub fn new(key: K, value: V) -> Node<K, V> {
        Node { key: key, value: value, left: None, right: None, level: 1 }
    }

    fn max(&self) -> &K {
        match self.right {
            None => &self.key,
            Some(ref n) => n.max(),
        }
    }

    fn min(&self) -> &K {
        match self.left {
            None => &self.key,
            Some(ref n) => n.min(),
        }
    }

    pub fn is_bst(&self) -> bool {
        self.left.as_ref().map_or(true, |n| n.is_bst() && *n.max() < self.key)
        && self.right.as_ref().map_or(true, |n| n.is_bst() && *n.min() > self.key)
    }

    pub fn find<'a>(&'a self, key: &K) -> Option<&'a V> {
        let ch = match key.cmp(&self.key) {
            Equal => return Some(&self.value),
            Less => & self.left,
            Greater => & self.right,
        };

        match *ch {
            None => None,
            Some(ref b) => b.find(key),
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let replaced = match key.cmp(&self.key) {
            Equal => {
                self.key = key;
                return Some(replace(&mut self.value, value))
            },
            Less =>
                match self.left {
                    None => {
                        self.left = Some(box Node::new(key, value));
                        None
                    },
                    Some(ref mut b) => b.insert(key, value),
                },
            Greater =>
                match self.right {
                    None => {
                        self.right = Some(box Node::new(key, value));
                        None
                    },
                    Some(ref mut b) => b.insert(key, value),
                },
        };

        self.skew();
        self.split();
        replaced
    }

    // To be an AA tree, it must be a binary search tree and, for all nodes n:
    //   - if n is missing a child, its level must be 1
    //   - the left child must have a level one less than n's level
    //   - the right child must have a level equal to or one less than n's level
    //   - the right child's right child must not have the same level as n's level
    fn is_aa(&self) -> bool {
        let lvl = self.level;

        self.is_bst()
            && !((self.left.is_none() || self.right.is_none()) && self.level != 1)
            && self.left.as_ref().map_or(true, |n| n.is_aa())
            && self.no_red_left_child()
            && self.right.as_ref().map_or(true,
                |n| n.is_aa() && (n.level == lvl || n.level + 1 == lvl)
                    && !(n.level == lvl && !n.no_red_right_child()))
    }

    fn no_red_left_child(&self) -> bool {
        match self.left {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    fn no_red_right_child(&self) -> bool {
        match self.right {
            None => true,
            Some(ref n) => n.level + 1 == self.level,
        }
    }

    // Remove red left child by rotating right
    /*
         a      b
        /        \
       b    =>    a
        \        /
         c      c

      provided that a.level == b.level
    */
    fn skew(&mut self) {
        if self.left.is_some() && self.left.get_ref().level == self.level {
            let mut save = self.left.take_unwrap();
            swap(&mut self.left, &mut save.right); // save.right now None
            swap(self, &mut *save);
            self.right = Some(save);
        }
    }

    // Remove red right child of red right child by rotating left and increasing
    // level of the parent
    /*
        a            b
         \          / \
          b    =>  a   c
         / \        \
        d   c        d

      provided that a.level == c.level
    */
    fn split(&mut self) {
        if self.right.as_ref().map_or(false,
          |x| x.right.is_some() && x.right.get_ref().level == self.level) {
            let mut save = self.right.take_unwrap();
            swap(&mut self.right, &mut save.left); // save.left now None
            save.level += 1;
            swap(self, &mut *save);
            self.left = Some(save);
        }
    }

    /*

       In the below, 'o' represents black, '*' represents red.

       First, note that deleting a red leaf requires no rebalancing at all.

       ---

       There's only 1 possible AA tree of size 2:

       o
        \
         *

       Deleting the root or the leaf requires no rebalancing.

       ---

       For size 3 trees, we only have:

         o
        / \
       o   o

        - deleting left leaf means we have to decrease the level of root. this suffices to rebalance
        - deleting right leaf means we have to decrease level of root but then make left leaf the
          new root (and former root the child of new root)
        - deleting the root means right child is new root, left child becomes new root's right child

       ---

       For size 4 trees, we have either

         o       o
        / \     / \
       o   o   o   o
        \           \
         *           *

       There are 4 classes of deletion:

        - delete red leaf (see above)
        - delete red node's parent: make red a child of its former grandparent
        - delete root: in the first case, red needs to become root, so increase its level and make it
          the new root. in the second case, right child of root needs to become new root.
        - delete other child: in the first case, red needs to become root and the old root needs to
          be right child of new root. in the second case, decrease level of old root, increase level
          of right child and make old root the child of the right child (which is the new root)

       ---

       For size 5 trees, we have one of

         o        o
        / \      / \
       o   o    o   *
        \   \      / \
         *   *    o   o

       ===

       One fundamental operation we need is "promote the max from the left subtree (or min from right subtree!)
       this is used when we delete a root, for example. this seems essentially whenever we delete a node that
       has two children.

     */
}

impl<K: Show, V: Show> Show for Node<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.key, self.value)
    }
}

fn print_node_depth<K: Show, V: Show>(node: &Link<Node<K,V>>, depth: uint) {
    let mut pre = "".to_string();
    if depth > 0 {
        for i in range(0, depth) {
            pre = pre + "   ";
        }
    }

    match *node {
        Some(ref n) => {
            println!("{}{}^{}", pre, n, n.level);
            print_node_depth(&n.left, depth + 1);
            print_node_depth(&n.right, depth + 1);
        },
        None => println!("{}-", pre),
    }
}

pub fn print_tree<K: Show + Ord, V: Show>(tree: &Tree<K, V>) {
    print_node_depth(&tree.root, 0);
    println!("Is AA: {}", tree.is_aa());
    println!("------------");
}

#[cfg(test)]
mod test {
    use super::Tree;
    use std::rand;
    use std::rand::distributions::{IndependentSample, Range};

    #[test]
    fn test_find() {
        let mut t = Tree::new();
        assert_eq!(t.find(&1u), None);
        t.insert(1u, 'j');
        assert_eq!(t.find(&1u), Some(&'j'));

    }

    // testing whether we can find all the things we inserted
    #[test]
    fn test_insert() {
        let mut t: Tree<uint, u8> = Tree::new();
        for (i, c) in range(0u, 10).zip(range(b'a', b'z')) {
            t.insert(i, c);
        }

        for (ref i, ref c) in range(0u, 10).zip(range(b'a', b'z')) {
            assert_eq!(t.find(i), Some(c));
        }

        assert_eq!(t.find(&10u), None);
    }

    fn insert_n_check_aa(n: uint, between: Range<uint>, rng: &mut rand::TaskRng) {
        let mut t = Tree::new();

        for _ in range(0u, n) {
            let a = between.ind_sample(rng);
            println!("{}", a);
            t.insert(a, ());
        }

        assert!(t.is_aa());
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true
    #[test]
    fn test_insert_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 100_000);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }

    // testing whether, after inserting 20 random keys, is_aa() returns true,
    // but this time some of the keys are repeated
    #[test]
    fn test_insert_dups_is_aa() {
        let mut rng = rand::task_rng();
        let between = Range::new(0u, 15);

        for _ in range(0u, 300) {
            insert_n_check_aa(20, between, &mut rng);
        }
    }
}
<|endoftext|>


<commit_before>//! INI configuration reader for wlroots-rs examples.
//! Modeled after the ini reader from wlroots.

extern crate ini;
extern crate wlroots;
extern crate wlroots_sys;

use wlroots_sys::wl_output_transform;

//use wlroots::Area;

pub struct OutputConfig {
    name: String,
    transform: wl_output_transform,
    x: i32,
    y: i32
}

pub struct DevConfig {
    name: String,
    mapped_output: String,
    //mapped_box: Area
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }
}
<commit_msg>Filled out the rest of the ex config structs<commit_after>//! INI configuration reader for wlroots-rs examples.
//! Modeled after the ini reader from wlroots.

extern crate ini;
extern crate wlroots;
extern crate wlroots_sys;

use std::path::PathBuf;

use wlroots_sys::wl_output_transform;

use wlroots::types::Area;

/// Main example configuration, holds the configuration sections for
/// outputs and devices.
pub struct ExampleConfig {
    pub config_path: PathBuf,
    pub mapped_output: String,
    pub mapped_box: Area,
    pub outputs: Vec<OutputConfig>,
    pub devs: Vec<DevConfig>
}

/// Configuration section for outputs.
pub struct OutputConfig {
    pub name: String,
    pub transform: wl_output_transform,
    pub x: i32,
    pub y: i32
}

/// Configuration section for devices.
pub struct DevConfig {
    pub name: String,
    pub mapped_output: String,
    pub mapped_box: Area
}

pub fn parse_args(args: &[&str]) -> Option<ExampleConfig> {
    unimplemented!()
}
<|endoftext|>