<commit_before>extern crate gimli;
use gimli::{AttributeValue, DebugAbbrev, DebugInfo, Endianity, EntriesCursor, LittleEndian};

#[cfg(test)]
fn assert_current_name<'input, 'abbrev, 'unit, Endian>(cursor: &mut EntriesCursor<'input,
                                                                           'abbrev,
                                                                           'unit,
                                                                           Endian>,
                                                       name: &'static str)
    where Endian: Endianity
{
    let entry = cursor.current()
        .expect("Should have an entry result")
        .expect("and it should be ok");

    let value = entry.attr_value(gimli::DW_AT_name)
        .expect("Should have found the name attribute");

    let mut with_null: Vec<u8> = name.as_bytes().into();
    with_null.push(0);

    assert_eq!(value, AttributeValue::String(&with_null));
}

#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
const ENTRIES_CURSOR_TESTS_ABBREV_BUF: [u8; 8] = [
    // Code
    0x01,

    // DW_TAG_subprogram
    0x2e,

    // DW_CHILDREN_yes
    0x01,

    // Begin attributes

        // Attribute name = DW_AT_name
        0x03,
        // Attribute form = DW_FORM_string
        0x08,

    // End attributes
    0x00,
    0x00,

    // Null terminator
    0x00
];

#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
const ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF: [u8; 71] = [
    // Compilation unit header

    // 32-bit unit length = 67
    0x43, 0x00, 0x00, 0x00,
    // Version 4
    0x04, 0x00,
    // debug_abbrev_offset
    0x00, 0x00, 0x00, 0x00,
    // Address size
    0x04,

    // DIEs

    // Abbreviation code
    0x01,
    // Attribute of form DW_FORM_string = "001\0"
    0x30, 0x30, 0x31, 0x00,

    // Children

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "002\0"
        0x30, 0x30, 0x32, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "003\0"
            0x30, 0x30, 0x33, 0x00,

            // Children

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "004\0"
        0x30, 0x30, 0x34, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "005\0"
            0x30, 0x30, 0x35, 0x00,

            // Children

            // End of children
            0x00,

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "006\0"
            0x30, 0x30, 0x36, 0x00,

            // Children

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "007\0"
        0x30, 0x30, 0x37, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "008\0"
            0x30, 0x30, 0x38, 0x00,

            // Children

                // Abbreviation code
                0x01,
                // Attribute of form DW_FORM_string = "009\0"
                0x30, 0x30, 0x39, 0x00,

                // Children

                // End of children
                0x00,

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "010\0"
        0x30, 0x31, 0x30, 0x00,

        // Children

        // End of children
        0x00,

    // End of children
    0x00
];

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_dfs() {
    let info_buf = &ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF;
    let debug_info = DebugInfo::<LittleEndian>::new(info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrevs_buf = &ENTRIES_CURSOR_TESTS_ABBREV_BUF;
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(abbrevs_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);
    assert_current_name(&mut cursor, "001");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "002");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "003");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), -1);
    assert_current_name(&mut cursor, "004");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "005");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 0);
    assert_current_name(&mut cursor, "006");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), -1);
    assert_current_name(&mut cursor, "007");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "008");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "009");

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), -2);
    assert_current_name(&mut cursor, "010");

    assert!(cursor.next_dfs().is_none());
    assert!(cursor.current().is_none());
}

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_sibling_no_sibling_ptr() {
    let info_buf = &ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF;
    let debug_info = DebugInfo::<LittleEndian>::new(info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrevs_buf = &ENTRIES_CURSOR_TESTS_ABBREV_BUF;
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(abbrevs_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);

    assert_current_name(&mut cursor, "001");

    // Down to the first child of the root entry.

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "002");

    // Now iterate all children of the root via `next_sibling`.

    cursor.next_sibling().expect("Should not be done with traversal");
    assert_current_name(&mut cursor, "004");

    cursor.next_sibling().expect("Should not be done with traversal");
    assert_current_name(&mut cursor, "007");

    cursor.next_sibling().expect("Should not be done with traversal");
    assert_current_name(&mut cursor, "010");

    // And now the cursor should be exhausted.

    assert!(cursor.next_sibling().is_none());
    assert!(cursor.current().is_none());
}

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_sibling_with_sibling_ptr() {
    let info_buf = [
        // Compilation unit header

        // 32-bit unit length = 56
        0x38, 0x00, 0x00, 0x00,
        // Version 4
        0x04, 0x00,
        // debug_abbrev_offset
        0x00, 0x00, 0x00, 0x00,
        // Address size
        0x04,

        // DIEs

        // Abbreviation code
        0x01,

        // DW_AT_name of form DW_FORM_string = "001\0"
        0x30, 0x30, 0x31, 0x00,
        // DW_AT_sibling of form DW_FORM_ref1
        0x00,

        // Children

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "002\0"
            0x30, 0x30, 0x32, 0x00,
            // Valid DW_AT_sibling pointer of form DW_FORM_ref1 = 31
            0x1f,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "003\0"
                0x30, 0x30, 0x33, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "004\0"
            0x30, 0x30, 0x34, 0x00,
            // Invalid DW_AT_sibling of form DW_FORM_ref1 = 255
            0xff,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "005\0"
                0x30, 0x30, 0x35, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "006\0"
            0x30, 0x30, 0x36, 0x00,
            // DW_AT_sibling of form DW_FORM_ref1
            0x00,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "007\0"
                0x30, 0x30, 0x37, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

        // End children
        0x00,
    ];

    let debug_info = DebugInfo::<LittleEndian>::new(&info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrev_buf = [
        // Code
        0x01,

        // DW_TAG_subprogram
        0x2e,

        // DW_CHILDREN_yes
        0x01,

        // Begin attributes

        // Attribute name = DW_AT_name
        0x03,
        // Attribute form = DW_FORM_string
        0x08,

        // Attribute name = DW_AT_sibling
        0x01,
        // Attribute form = DW_FORM_ref1
        0x11,

        // End attributes
        0x00,
        0x00,

        // Null terminator
        0x00
    ];

    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&abbrev_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);

    assert_current_name(&mut cursor, "001");

    // Down to the first child of the root.

    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), 1);
    assert_current_name(&mut cursor, "002");

    // Now iterate all children of the root via `next_sibling`.

    cursor.next_sibling().expect("Should handle valid DW_AT_sibling pointer");
    assert_current_name(&mut cursor, "004");

    cursor.next_sibling().expect("Should handle invalid DW_AT_sibling pointer");
    assert_current_name(&mut cursor, "006");

    // And now the cursor should be exhausted.

    assert!(cursor.next_sibling().is_none());
    assert!(cursor.current().is_none());
}

<commit_msg>Add assert_next_dfs/assert_next_sibling<commit_after>extern crate gimli;
use gimli::{AttributeValue, DebugAbbrev, DebugInfo, Endianity, EntriesCursor, LittleEndian};

#[cfg(test)]
fn assert_current_name<'input, 'abbrev, 'unit, Endian>(cursor: &mut EntriesCursor<'input,
                                                                           'abbrev,
                                                                           'unit,
                                                                           Endian>,
                                                       name: &'static str)
    where Endian: Endianity
{
    let entry = cursor.current()
        .expect("Should have an entry result")
        .expect("and it should be ok");

    let value = entry.attr_value(gimli::DW_AT_name)
        .expect("Should have found the name attribute");

    let mut with_null: Vec<u8> = name.as_bytes().into();
    with_null.push(0);

    assert_eq!(value, AttributeValue::String(&with_null));
}

#[cfg(test)]
fn assert_next_dfs<'input, 'abbrev, 'unit, Endian>(cursor: &mut EntriesCursor<'input,
                                                                           'abbrev,
                                                                           'unit,
                                                                           Endian>,
                                                   name: &'static str,
                                                   depth: isize)
    where Endian: Endianity
{
    assert_eq!(cursor.next_dfs().expect("Should not be done with traversal"), depth);
    assert_current_name(cursor, name);
}

#[cfg(test)]
fn assert_next_sibling<'input, 'abbrev, 'unit, Endian>(cursor: &mut EntriesCursor<'input,
                                                                           'abbrev,
                                                                           'unit,
                                                                           Endian>,
                                                       name: &'static str)
    where Endian: Endianity
{
    cursor.next_sibling().expect("Should not be done with traversal");
    assert_current_name(cursor, name);
}

#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
const ENTRIES_CURSOR_TESTS_ABBREV_BUF: [u8; 8] = [
    // Code
    0x01,

    // DW_TAG_subprogram
    0x2e,

    // DW_CHILDREN_yes
    0x01,

    // Begin attributes

        // Attribute name = DW_AT_name
        0x03,
        // Attribute form = DW_FORM_string
        0x08,

    // End attributes
    0x00,
    0x00,

    // Null terminator
    0x00
];

#[cfg(test)]
#[cfg_attr(rustfmt, rustfmt_skip)]
const ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF: [u8; 71] = [
    // Compilation unit header

    // 32-bit unit length = 67
    0x43, 0x00, 0x00, 0x00,
    // Version 4
    0x04, 0x00,
    // debug_abbrev_offset
    0x00, 0x00, 0x00, 0x00,
    // Address size
    0x04,

    // DIEs

    // Abbreviation code
    0x01,
    // Attribute of form DW_FORM_string = "001\0"
    0x30, 0x30, 0x31, 0x00,

    // Children

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "002\0"
        0x30, 0x30, 0x32, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "003\0"
            0x30, 0x30, 0x33, 0x00,

            // Children

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "004\0"
        0x30, 0x30, 0x34, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "005\0"
            0x30, 0x30, 0x35, 0x00,

            // Children

            // End of children
            0x00,

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "006\0"
            0x30, 0x30, 0x36, 0x00,

            // Children

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "007\0"
        0x30, 0x30, 0x37, 0x00,

        // Children

            // Abbreviation code
            0x01,
            // Attribute of form DW_FORM_string = "008\0"
            0x30, 0x30, 0x38, 0x00,

            // Children

                // Abbreviation code
                0x01,
                // Attribute of form DW_FORM_string = "009\0"
                0x30, 0x30, 0x39, 0x00,

                // Children

                // End of children
                0x00,

            // End of children
            0x00,

        // End of children
        0x00,

        // Abbreviation code
        0x01,
        // Attribute of form DW_FORM_string = "010\0"
        0x30, 0x31, 0x30, 0x00,

        // Children

        // End of children
        0x00,

    // End of children
    0x00
];

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_dfs() {
    let info_buf = &ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF;
    let debug_info = DebugInfo::<LittleEndian>::new(info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrevs_buf = &ENTRIES_CURSOR_TESTS_ABBREV_BUF;
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(abbrevs_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);

    assert_current_name(&mut cursor, "001");
    assert_next_dfs(&mut cursor, "002", 1);
    assert_next_dfs(&mut cursor, "003", 1);
    assert_next_dfs(&mut cursor, "004", -1);
    assert_next_dfs(&mut cursor, "005", 1);
    assert_next_dfs(&mut cursor, "006", 0);
    assert_next_dfs(&mut cursor, "007", -1);
    assert_next_dfs(&mut cursor, "008", 1);
    assert_next_dfs(&mut cursor, "009", 1);
    assert_next_dfs(&mut cursor, "010", -2);

    assert!(cursor.next_dfs().is_none());
    assert!(cursor.current().is_none());
}

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_sibling_no_sibling_ptr() {
    let info_buf = &ENTRIES_CURSOR_TESTS_DEBUG_INFO_BUF;
    let debug_info = DebugInfo::<LittleEndian>::new(info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrevs_buf = &ENTRIES_CURSOR_TESTS_ABBREV_BUF;
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(abbrevs_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);

    assert_current_name(&mut cursor, "001");

    // Down to the first child of the root entry.

    assert_next_dfs(&mut cursor, "002", 1);

    // Now iterate all children of the root via `next_sibling`.

    assert_next_sibling(&mut cursor, "004");
    assert_next_sibling(&mut cursor, "007");
    assert_next_sibling(&mut cursor, "010");

    // And now the cursor should be exhausted.

    assert!(cursor.next_sibling().is_none());
    assert!(cursor.current().is_none());
}

#[test]
#[cfg_attr(rustfmt, rustfmt_skip)]
fn test_cursor_next_sibling_with_sibling_ptr() {
    let info_buf = [
        // Compilation unit header

        // 32-bit unit length = 56
        0x38, 0x00, 0x00, 0x00,
        // Version 4
        0x04, 0x00,
        // debug_abbrev_offset
        0x00, 0x00, 0x00, 0x00,
        // Address size
        0x04,

        // DIEs

        // Abbreviation code
        0x01,

        // DW_AT_name of form DW_FORM_string = "001\0"
        0x30, 0x30, 0x31, 0x00,
        // DW_AT_sibling of form DW_FORM_ref1
        0x00,

        // Children

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "002\0"
            0x30, 0x30, 0x32, 0x00,
            // Valid DW_AT_sibling pointer of form DW_FORM_ref1 = 31
            0x1f,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "003\0"
                0x30, 0x30, 0x33, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "004\0"
            0x30, 0x30, 0x34, 0x00,
            // Invalid DW_AT_sibling of form DW_FORM_ref1 = 255
            0xff,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "005\0"
                0x30, 0x30, 0x35, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

            // Abbreviation code
            0x01,

            // DW_AT_name of form DW_FORM_string = "006\0"
            0x30, 0x30, 0x36, 0x00,
            // DW_AT_sibling of form DW_FORM_ref1
            0x00,

            // Children

                // Abbreviation code
                0x01,

                // DW_AT_name of form DW_FORM_string = "007\0"
                0x30, 0x30, 0x37, 0x00,
                // DW_AT_sibling of form DW_FORM_ref1
                0x00,

                // No children
                0x00,

            // End children
            0x00,

        // End children
        0x00,
    ];

    let debug_info = DebugInfo::<LittleEndian>::new(&info_buf);

    let unit = debug_info.units().next()
        .expect("should have a unit result")
        .expect("and it should be ok");

    let abbrev_buf = [
        // Code
        0x01,

        // DW_TAG_subprogram
        0x2e,

        // DW_CHILDREN_yes
        0x01,

        // Begin attributes

        // Attribute name = DW_AT_name
        0x03,
        // Attribute form = DW_FORM_string
        0x08,

        // Attribute name = DW_AT_sibling
        0x01,
        // Attribute form = DW_FORM_ref1
        0x11,

        // End attributes
        0x00,
        0x00,

        // Null terminator
        0x00
    ];

    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&abbrev_buf);

    let abbrevs = unit.abbreviations(debug_abbrev)
        .expect("Should parse abbreviations");

    let mut cursor = unit.entries(&abbrevs);

    assert_current_name(&mut cursor, "001");

    // Down to the first child of the root.

    assert_next_dfs(&mut cursor, "002", 1);

    // Now iterate all children of the root via `next_sibling`.

    assert_next_sibling(&mut cursor, "004");
    assert_next_sibling(&mut cursor, "006");

    // And now the cursor should be exhausted.

    assert!(cursor.next_sibling().is_none());
    assert!(cursor.current().is_none());
}

<|endoftext|>


<commit_before>use self::build_tables::build_tables;
use self::parse_grammar::parse_grammar;
use self::prepare_grammar::prepare_grammar;
use self::render::render_c_code;
use crate::error::{Error, Result};
use lazy_static::lazy_static;
use regex::{Regex, RegexBuilder};
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

mod build_tables;
mod grammars;
mod nfa;
mod npm_files;
mod parse_grammar;
mod prepare_grammar;
mod render;
mod rules;
mod tables;

lazy_static! {
    static ref JSON_COMMENT_REGEX: Regex = RegexBuilder::new("^\\s*//.*")
        .multi_line(true)
        .build()
        .unwrap();
}

pub fn generate_parser_in_directory(
    repo_path: &PathBuf,
    grammar_path: Option<&str>,
    minimize: bool,
    state_ids_to_log: Vec<usize>,
) -> Result<()> {
    let repo_src_path = repo_path.join("src");
    let repo_header_path = repo_src_path.join("tree_sitter");

    fs::create_dir_all(&repo_src_path)?;
    fs::create_dir_all(&repo_header_path)?;

    let grammar_json;
    match grammar_path {
        Some(path) => {
            grammar_json = load_grammar_file(path.as_ref())?;
        }
        None => {
            let grammar_js_path = grammar_path.map_or(repo_path.join("grammar.js"), |s| s.into());
            grammar_json = load_grammar_file(&grammar_js_path)?;
            fs::write(&repo_src_path.join("grammar.json"), &grammar_json)?;
        }
    }

    let (language_name, c_code) =
        generate_parser_for_grammar_with_opts(&grammar_json, minimize, state_ids_to_log)?;
    fs::write(&repo_src_path.join("parser.c"), c_code)
        .map_err(|e| format!("Failed to write parser.c: {}", e))?;
    fs::write(
        &repo_header_path.join("parser.h"),
        tree_sitter::PARSER_HEADER,
    )
    .map_err(|e| format!("Failed to write parser.h: {}", e))?;
    ensure_file(&repo_src_path.join("binding.cc"), || {
        npm_files::binding_cc(&language_name)
    })?;
    ensure_file(&repo_path.join("binding.gyp"), || {
        npm_files::binding_gyp(&language_name)
    })?;
    ensure_file(&repo_path.join("index.js"), || {
        npm_files::index_js(&language_name)
    })?;
    Ok(())
}

pub fn generate_parser_for_grammar(grammar_json: &str) -> Result<(String, String)> {
    let grammar_json = JSON_COMMENT_REGEX.replace_all(grammar_json, "\n");
    generate_parser_for_grammar_with_opts(&grammar_json, true, Vec::new())
}

fn generate_parser_for_grammar_with_opts(
    grammar_json: &str,
    minimize: bool,
    state_ids_to_log: Vec<usize>,
) -> Result<(String, String)> {
    let input_grammar = parse_grammar(grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let (parse_table, main_lex_table, keyword_lex_table, keyword_capture_token) = build_tables(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &inlines,
        minimize,
        state_ids_to_log,
    )?;
    let c_code = render_c_code(
        &input_grammar.name,
        parse_table,
        main_lex_table,
        keyword_lex_table,
        keyword_capture_token,
        syntax_grammar,
        lexical_grammar,
        simple_aliases,
    );
    Ok((input_grammar.name, c_code))
}

fn load_grammar_file(grammar_path: &Path) -> Result<String> {
    match grammar_path.extension().and_then(|e| e.to_str()) {
        Some("js") => Ok(load_js_grammar_file(grammar_path)?),
        Some("json") => Ok(fs::read_to_string(grammar_path)?),
        _ => Err(Error(format!(
            "Unknown grammar file extension: {:?}",
            grammar_path
        ))),
    }
}

fn load_js_grammar_file(grammar_path: &Path) -> Result<String> {
    let mut node_process = Command::new("node")
        .env("TREE_SITTER_GRAMMAR_PATH", grammar_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to run `node`");

    let mut node_stdin = node_process
        .stdin
        .take()
        .expect("Failed to open stdin for node");
    let javascript_code = include_bytes!("./dsl.js");
    node_stdin
        .write(javascript_code)
        .expect("Failed to write to node's stdin");
    drop(node_stdin);
    let output = node_process
        .wait_with_output()
        .expect("Failed to read output from node");
    match output.status.code() {
        None => panic!("Node process was killed"),
        Some(0) => {}
        Some(code) => return Err(Error(format!("Node process exited with status {}", code))),
    }

    Ok(String::from_utf8(output.stdout).expect("Got invalid UTF8 from node"))
}

fn ensure_file<T: AsRef<[u8]>>(path: &PathBuf, f: impl Fn() -> T) -> Result<()> {
    if path.exists() {
        Ok(())
    } else {
        fs::write(path, f().as_ref())
            .map_err(|e| Error(format!("Failed to write file {:?}: {}", path, e)))
    }
}
<commit_msg>cli: put a newline character at the end of grammar.json files<commit_after>use self::build_tables::build_tables;
use self::parse_grammar::parse_grammar;
use self::prepare_grammar::prepare_grammar;
use self::render::render_c_code;
use crate::error::{Error, Result};
use lazy_static::lazy_static;
use regex::{Regex, RegexBuilder};
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

mod build_tables;
mod grammars;
mod nfa;
mod npm_files;
mod parse_grammar;
mod prepare_grammar;
mod render;
mod rules;
mod tables;

lazy_static! {
    static ref JSON_COMMENT_REGEX: Regex = RegexBuilder::new("^\\s*//.*")
        .multi_line(true)
        .build()
        .unwrap();
}

pub fn generate_parser_in_directory(
    repo_path: &PathBuf,
    grammar_path: Option<&str>,
    minimize: bool,
    state_ids_to_log: Vec<usize>,
) -> Result<()> {
    let repo_src_path = repo_path.join("src");
    let repo_header_path = repo_src_path.join("tree_sitter");

    fs::create_dir_all(&repo_src_path)?;
    fs::create_dir_all(&repo_header_path)?;

    let grammar_json;
    match grammar_path {
        Some(path) => {
            grammar_json = load_grammar_file(path.as_ref())?;
        }
        None => {
            let grammar_js_path = grammar_path.map_or(repo_path.join("grammar.js"), |s| s.into());
            grammar_json = load_grammar_file(&grammar_js_path)?;
            fs::write(&repo_src_path.join("grammar.json"), &grammar_json)?;
        }
    }

    let (language_name, c_code) =
        generate_parser_for_grammar_with_opts(&grammar_json, minimize, state_ids_to_log)?;
    fs::write(&repo_src_path.join("parser.c"), c_code)
        .map_err(|e| format!("Failed to write parser.c: {}", e))?;
    fs::write(
        &repo_header_path.join("parser.h"),
        tree_sitter::PARSER_HEADER,
    )
    .map_err(|e| format!("Failed to write parser.h: {}", e))?;
    ensure_file(&repo_src_path.join("binding.cc"), || {
        npm_files::binding_cc(&language_name)
    })?;
    ensure_file(&repo_path.join("binding.gyp"), || {
        npm_files::binding_gyp(&language_name)
    })?;
    ensure_file(&repo_path.join("index.js"), || {
        npm_files::index_js(&language_name)
    })?;
    Ok(())
}

pub fn generate_parser_for_grammar(grammar_json: &str) -> Result<(String, String)> {
    let grammar_json = JSON_COMMENT_REGEX.replace_all(grammar_json, "\n");
    generate_parser_for_grammar_with_opts(&grammar_json, true, Vec::new())
}

fn generate_parser_for_grammar_with_opts(
    grammar_json: &str,
    minimize: bool,
    state_ids_to_log: Vec<usize>,
) -> Result<(String, String)> {
    let input_grammar = parse_grammar(grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let (parse_table, main_lex_table, keyword_lex_table, keyword_capture_token) = build_tables(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &inlines,
        minimize,
        state_ids_to_log,
    )?;
    let c_code = render_c_code(
        &input_grammar.name,
        parse_table,
        main_lex_table,
        keyword_lex_table,
        keyword_capture_token,
        syntax_grammar,
        lexical_grammar,
        simple_aliases,
    );
    Ok((input_grammar.name, c_code))
}

fn load_grammar_file(grammar_path: &Path) -> Result<String> {
    match grammar_path.extension().and_then(|e| e.to_str()) {
        Some("js") => Ok(load_js_grammar_file(grammar_path)?),
        Some("json") => Ok(fs::read_to_string(grammar_path)?),
        _ => Err(Error(format!(
            "Unknown grammar file extension: {:?}",
            grammar_path
        ))),
    }
}

fn load_js_grammar_file(grammar_path: &Path) -> Result<String> {
    let mut node_process = Command::new("node")
        .env("TREE_SITTER_GRAMMAR_PATH", grammar_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to run `node`");

    let mut node_stdin = node_process
        .stdin
        .take()
        .expect("Failed to open stdin for node");
    let javascript_code = include_bytes!("./dsl.js");
    node_stdin
        .write(javascript_code)
        .expect("Failed to write to node's stdin");
    drop(node_stdin);
    let output = node_process
        .wait_with_output()
        .expect("Failed to read output from node");
    match output.status.code() {
        None => panic!("Node process was killed"),
        Some(0) => {}
        Some(code) => return Err(Error(format!("Node process exited with status {}", code))),
    }

    let mut result = String::from_utf8(output.stdout).expect("Got invalid UTF8 from node");
    result.push('\n');
    Ok(result)
}

fn ensure_file<T: AsRef<[u8]>>(path: &PathBuf, f: impl Fn() -> T) -> Result<()> {
    if path.exists() {
        Ok(())
    } else {
        fs::write(path, f().as_ref())
            .map_err(|e| Error(format!("Failed to write file {:?}: {}", path, e)))
    }
}
<|endoftext|>


<commit_before>use domination::DominationOrd;
use std::cmp::Ordering;
use std::mem;
use std::collections::VecDeque;

struct Entry<'a, S, I = usize>
where
    S: 'a,
{
    /// A reference to the solution
    solution: &'a S,

    /// The index that `solution` has within the `solutions` array.
    index: I,

    /// By how many other solutions is this solution dominated
    domination_count: I,

    /// Which solutions we dominate
    dominated_solutions: VecDeque<I>,
}

pub struct NonDominatedSorter<'a, S, I = usize>
where
    S: 'a,
{
    entries: Vec<Entry<'a, S, I>>,
    current_front: Vec<(I, &'a S)>,
}

impl<'a, S> NonDominatedSorter<'a, S> {
    /// Perform a non-dominated sort of `solutions`.
    ///
    /// Each pareto front (the indices of the `solutions`) can be obtained by calling `next()`.

    pub fn new<D>(solutions: &'a [S], domination: &D) -> Self
    where
        D: DominationOrd<Solution = S>,
    {
        let mut current_front = Vec::new();

        let mut entries: Vec<_> = solutions
            .iter()
            .enumerate()
            .map(|(index, solution)| Entry {
                solution,
                index,
                domination_count: 0,
                dominated_solutions: VecDeque::new(),
            })
            .collect();

        for mid in 1..entries.len() + 1 {
            let (front_slice, tail_slice) = entries.split_at_mut(mid);
            debug_assert!(front_slice.len() > 0);
            let p = front_slice.last_mut().unwrap();
            for q in tail_slice.iter_mut() {
                match domination.domination_ord(p.solution, q.solution) {
                    Ordering::Less => {
                        // p dominates q
                        // Add `q` to the set of solutions dominated by `p`.
                        p.dominated_solutions.push_back(q.index);
                        // q is dominated by p
                        q.domination_count += 1;
                    }
                    Ordering::Greater => {
                        // p is dominated by q
                        // Add `p` to the set of solutions dominated by `q`.
                        q.dominated_solutions.push_back(p.index);
                        // q dominates p
                        // Increment domination counter of `p`.
                        p.domination_count += 1;
                    }
                    Ordering::Equal => {}
                }
            }
            if p.domination_count == 0 {
                // `p` belongs to the first front as it is not dominated by any
                // other solution.
                current_front.push((p.index, p.solution));
            }
        }

        Self {
            entries,
            current_front,
        }
    }

    /// Returns an array containing all pareto fronts.
    pub fn pareto_fronts(self) -> Vec<Vec<(usize, &'a S)>> {
        let mut fronts = Vec::new();
        for front in self {
            fronts.push(front);
        }
        return fronts;
    }

    /// Returns an array containing the first pareto fronts, until
    /// `max_solutions` have been found. Note that always the whole fronts are
    /// returned, i.e. the number of solutions returned may be higher.
    pub fn pareto_fronts_stop_at(self, max_solutions: usize) -> Vec<Vec<(usize, &'a S)>> {
        let mut found_solutions = 0;
        let mut fronts = Vec::new();

        for front in self {
            found_solutions += front.len();
            fronts.push(front);
            if found_solutions >= max_solutions {
                break;
            }
        }
        return fronts;
    }
}

/// Iterate over the pareto fronts. Each call to next() will yield the
/// next pareto front.
impl<'a, S> Iterator for NonDominatedSorter<'a, S> {
    type Item = Vec<(usize, &'a S)>;

    /// Return the next pareto front

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

        // Calculate the next front

        let mut next_front = Vec::new();
        for &(p_i, _) in self.current_front.iter() {
            // to calculate the next front, we have to remove the
            // solutions of the current front, and as such, decrease the
            // domination_count of they dominated_solutions. We can
            // destruct the dominated_solutions array here, as we will
            // no longer need it.
            // The only problem with poping off solutions off the end is
            // that we will populate the fronts in reverse order. For
            // that reason, we are using a VecDeque. This should give us
            // a stable sort.

            while let Some(q_i) = self.entries[p_i].dominated_solutions.pop_front() {
                let q = &mut self.entries[q_i];
                debug_assert!(q.domination_count > 0);
                q.domination_count -= 1;
                if q.domination_count == 0 {
                    // q is not dominated by any other solution. it belongs to the next front.
                    next_front.push((q_i, q.solution));
                }
            }
        }

        // and return the current front, swapping it with the next

        let current_front = mem::replace(&mut self.current_front, next_front);

        return Some(current_front);
    }
}

#[cfg(test)]
mod helper {
    use test_helper_domination::Tuple;
    pub fn get_solutions() -> Vec<Tuple> {
        vec![
            Tuple(1, 2),
            Tuple(1, 2),
            Tuple(2, 1),
            Tuple(1, 3),
            Tuple(0, 2),
        ]
    }
    pub fn keep_only_index(front: &[(usize, &Tuple)]) -> Vec<usize> {
        front.iter().map(|&(i, _)| i).collect()
    }

    pub fn assert_front_eq(expected: &[usize], front: &[(usize, &Tuple)]) {
        assert_eq!(expected.to_owned(), keep_only_index(front));
    }
}

#[test]
fn test_non_dominated_sort() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();
    let fronts = NonDominatedSorter::new(&solutions, &TupleDominationOrd).pareto_fronts();

    assert_eq!(3, fronts.len());
    helper::assert_front_eq(&[2, 4], &fronts[0]);
    helper::assert_front_eq(&[0, 1], &fronts[1]);
    helper::assert_front_eq(&[3], &fronts[2]);
}

#[test]
fn test_non_dominated_sort_stop_at() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();

    {
        let fronts =
            NonDominatedSorter::new(&solutions, &TupleDominationOrd).pareto_fronts_stop_at(2);
        assert_eq!(1, fronts.len());
        helper::assert_front_eq(&[2, 4], &fronts[0]);
    }
    {
        let fronts =
            NonDominatedSorter::new(&solutions, &TupleDominationOrd).pareto_fronts_stop_at(3);
        assert_eq!(2, fronts.len());
        helper::assert_front_eq(&[2, 4], &fronts[0]);
        helper::assert_front_eq(&[0, 1], &fronts[1]);
    }
}

#[test]
fn test_non_dominated_sort_iter() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();
    let mut fronts = NonDominatedSorter::new(&solutions, &TupleDominationOrd);
    let f = &|front: Vec<_>| helper::keep_only_index(&front);

    assert_eq!(Some(vec![2, 4]), fronts.next().map(f));
    assert_eq!(Some(vec![0, 1]), fronts.next().map(f));
    assert_eq!(Some(vec![3]), fronts.next().map(f));
    assert_eq!(None, fronts.next().map(f));
}
<commit_msg>Rename to NonDominatedSort<commit_after>use domination::DominationOrd;
use std::cmp::Ordering;
use std::mem;
use std::collections::VecDeque;

struct Entry<'a, S, I = usize>
where
    S: 'a,
{
    /// A reference to the solution
    solution: &'a S,

    /// The index that `solution` has within the `solutions` array.
    index: I,

    /// By how many other solutions is this solution dominated
    domination_count: I,

    /// Which solutions we dominate
    dominated_solutions: VecDeque<I>,
}

pub struct NonDominatedSort<'a, S, I = usize>
where
    S: 'a,
{
    entries: Vec<Entry<'a, S, I>>,
    current_front: Vec<(I, &'a S)>,
}

impl<'a, S> NonDominatedSort<'a, S> {
    /// Perform a non-dominated sort of `solutions`.
    ///
    /// Each pareto front (the indices of the `solutions`) can be obtained by calling `next()`.

    pub fn new<D>(solutions: &'a [S], domination: &D) -> Self
    where
        D: DominationOrd<Solution = S>,
    {
        let mut current_front = Vec::new();

        let mut entries: Vec<_> = solutions
            .iter()
            .enumerate()
            .map(|(index, solution)| Entry {
                solution,
                index,
                domination_count: 0,
                dominated_solutions: VecDeque::new(),
            })
            .collect();

        for mid in 1..entries.len() + 1 {
            let (front_slice, tail_slice) = entries.split_at_mut(mid);
            debug_assert!(front_slice.len() > 0);
            let p = front_slice.last_mut().unwrap();
            for q in tail_slice.iter_mut() {
                match domination.domination_ord(p.solution, q.solution) {
                    Ordering::Less => {
                        // p dominates q
                        // Add `q` to the set of solutions dominated by `p`.
                        p.dominated_solutions.push_back(q.index);
                        // q is dominated by p
                        q.domination_count += 1;
                    }
                    Ordering::Greater => {
                        // p is dominated by q
                        // Add `p` to the set of solutions dominated by `q`.
                        q.dominated_solutions.push_back(p.index);
                        // q dominates p
                        // Increment domination counter of `p`.
                        p.domination_count += 1;
                    }
                    Ordering::Equal => {}
                }
            }
            if p.domination_count == 0 {
                // `p` belongs to the first front as it is not dominated by any
                // other solution.
                current_front.push((p.index, p.solution));
            }
        }

        Self {
            entries,
            current_front,
        }
    }

    /// Returns an array containing all pareto fronts.
    pub fn pareto_fronts(self) -> Vec<Vec<(usize, &'a S)>> {
        let mut fronts = Vec::new();
        for front in self {
            fronts.push(front);
        }
        return fronts;
    }

    /// Returns an array containing the first pareto fronts, until
    /// `max_solutions` have been found. Note that always the whole fronts are
    /// returned, i.e. the number of solutions returned may be higher.
    pub fn pareto_fronts_stop_at(self, max_solutions: usize) -> Vec<Vec<(usize, &'a S)>> {
        let mut found_solutions = 0;
        let mut fronts = Vec::new();

        for front in self {
            found_solutions += front.len();
            fronts.push(front);
            if found_solutions >= max_solutions {
                break;
            }
        }
        return fronts;
    }
}

/// Iterate over the pareto fronts. Each call to next() will yield the
/// next pareto front.
impl<'a, S> Iterator for NonDominatedSort<'a, S> {
    type Item = Vec<(usize, &'a S)>;

    /// Return the next pareto front

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

        // Calculate the next front

        let mut next_front = Vec::new();
        for &(p_i, _) in self.current_front.iter() {
            // to calculate the next front, we have to remove the
            // solutions of the current front, and as such, decrease the
            // domination_count of they dominated_solutions. We can
            // destruct the dominated_solutions array here, as we will
            // no longer need it.
            // The only problem with poping off solutions off the end is
            // that we will populate the fronts in reverse order. For
            // that reason, we are using a VecDeque. This should give us
            // a stable sort.

            while let Some(q_i) = self.entries[p_i].dominated_solutions.pop_front() {
                let q = &mut self.entries[q_i];
                debug_assert!(q.domination_count > 0);
                q.domination_count -= 1;
                if q.domination_count == 0 {
                    // q is not dominated by any other solution. it belongs to the next front.
                    next_front.push((q_i, q.solution));
                }
            }
        }

        // and return the current front, swapping it with the next

        let current_front = mem::replace(&mut self.current_front, next_front);

        return Some(current_front);
    }
}

#[cfg(test)]
mod helper {
    use test_helper_domination::Tuple;
    pub fn get_solutions() -> Vec<Tuple> {
        vec![
            Tuple(1, 2),
            Tuple(1, 2),
            Tuple(2, 1),
            Tuple(1, 3),
            Tuple(0, 2),
        ]
    }
    pub fn keep_only_index(front: &[(usize, &Tuple)]) -> Vec<usize> {
        front.iter().map(|&(i, _)| i).collect()
    }

    pub fn assert_front_eq(expected: &[usize], front: &[(usize, &Tuple)]) {
        assert_eq!(expected.to_owned(), keep_only_index(front));
    }
}

#[test]
fn test_non_dominated_sort() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();
    let fronts = NonDominatedSort::new(&solutions, &TupleDominationOrd).pareto_fronts();

    assert_eq!(3, fronts.len());
    helper::assert_front_eq(&[2, 4], &fronts[0]);
    helper::assert_front_eq(&[0, 1], &fronts[1]);
    helper::assert_front_eq(&[3], &fronts[2]);
}

#[test]
fn test_non_dominated_sort_stop_at() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();

    {
        let fronts =
            NonDominatedSort::new(&solutions, &TupleDominationOrd).pareto_fronts_stop_at(2);
        assert_eq!(1, fronts.len());
        helper::assert_front_eq(&[2, 4], &fronts[0]);
    }
    {
        let fronts =
            NonDominatedSort::new(&solutions, &TupleDominationOrd).pareto_fronts_stop_at(3);
        assert_eq!(2, fronts.len());
        helper::assert_front_eq(&[2, 4], &fronts[0]);
        helper::assert_front_eq(&[0, 1], &fronts[1]);
    }
}

#[test]
fn test_non_dominated_sort_iter() {
    use test_helper_domination::TupleDominationOrd;
    let solutions = helper::get_solutions();
    let mut fronts = NonDominatedSort::new(&solutions, &TupleDominationOrd);
    let f = &|front: Vec<_>| helper::keep_only_index(&front);

    assert_eq!(Some(vec![2, 4]), fronts.next().map(f));
    assert_eq!(Some(vec![0, 1]), fronts.next().map(f));
    assert_eq!(Some(vec![3]), fronts.next().map(f));
    assert_eq!(None, fronts.next().map(f));
}
<|endoftext|>


<commit_before>use std::io::{self, Read};

use rand::{Rng, RngCore};
use std::{cmp, iter};

use super::decoder::{DecoderReader, BUF_SIZE};
use crate::encode::encode_engine_string;
use crate::engine::fast_portable::FastPortable;
use crate::engine::DEFAULT_ENGINE;
use crate::tests::{random_alphabet, random_config, random_engine};
use crate::{decode_engine_vec, DecodeError};

#[test]
fn simple() {
    let tests: &[(&[u8], &[u8])] = &[
        (&b"0"[..], &b"MA=="[..]),
        (b"01", b"MDE="),
        (b"012", b"MDEy"),
        (b"0123", b"MDEyMw=="),
        (b"01234", b"MDEyMzQ="),
        (b"012345", b"MDEyMzQ1"),
        (b"0123456", b"MDEyMzQ1Ng=="),
        (b"01234567", b"MDEyMzQ1Njc="),
        (b"012345678", b"MDEyMzQ1Njc4"),
        (b"0123456789", b"MDEyMzQ1Njc4OQ=="),
    ][..];

    for (text_expected, base64data) in tests.iter() {
        // Read n bytes at a time.
        for n in 1..base64data.len() + 1 {
            let mut wrapped_reader = io::Cursor::new(base64data);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &DEFAULT_ENGINE);

            // handle errors as you normally would
            let mut text_got = Vec::new();
            let mut buffer = vec![0u8; n];
            while let Ok(read) = decoder.read(&mut buffer[..]) {
                if read == 0 {
                    break;
                }
                text_got.extend_from_slice(&buffer[..read]);
            }

            assert_eq!(
                text_got,
                *text_expected,
                "\nGot: {}\nExpected: {}",
                String::from_utf8_lossy(&text_got[..]),
                String::from_utf8_lossy(text_expected)
            );
        }
    }
}

// Make sure we error out on trailing junk.
#[test]
fn trailing_junk() {
    let tests: &[&[u8]] = &[&b"MDEyMzQ1Njc4*!@#$%^&"[..], b"MDEyMzQ1Njc4OQ== "][..];

    for base64data in tests.iter() {
        // Read n bytes at a time.
        for n in 1..base64data.len() + 1 {
            let mut wrapped_reader = io::Cursor::new(base64data);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &DEFAULT_ENGINE);

            // handle errors as you normally would
            let mut buffer = vec![0u8; n];
            let mut saw_error = false;
            loop {
                match decoder.read(&mut buffer[..]) {
                    Err(_) => {
                        saw_error = true;
                        break;
                    }
                    Ok(read) if read == 0 => break,
                    Ok(_) => (),
                }
            }

            assert!(saw_error);
        }
    }
}

#[test]
fn handles_short_read_from_delegate() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        bytes.truncate(size);
        rng.fill_bytes(&mut bytes[..size]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);
        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut wrapped_reader = io::Cursor::new(b64.as_bytes());
        let mut short_reader = RandomShortRead {
            delegate: &mut wrapped_reader,
            rng: &mut rng,
        };

        let mut decoder = DecoderReader::from(&mut short_reader, &engine);

        let decoded_len = decoder.read_to_end(&mut decoded).unwrap();
        assert_eq!(size, decoded_len);
        assert_eq!(&bytes[..], &decoded[..]);
    }
}

#[test]
fn read_in_short_increments() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        // leave room to play around with larger buffers
        decoded.extend(iter::repeat(0).take(size * 3));

        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut wrapped_reader = io::Cursor::new(&b64[..]);
        let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

        consume_with_short_reads_and_validate(&mut rng, &bytes[..], &mut decoded, &mut decoder);
    }
}

#[test]
fn read_in_short_increments_with_short_delegate_reads() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        // leave room to play around with larger buffers
        decoded.extend(iter::repeat(0).take(size * 3));

        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut base_reader = io::Cursor::new(&b64[..]);
        let mut decoder = DecoderReader::from(&mut base_reader, &engine);
        let mut short_reader = RandomShortRead {
            delegate: &mut decoder,
            rng: &mut rand::thread_rng(),
        };

        consume_with_short_reads_and_validate(&mut rng, &bytes[..], &mut decoded, &mut short_reader)
    }
}

#[test]
fn reports_invalid_last_symbol_correctly() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut b64_bytes = Vec::new();
    let mut decoded = Vec::new();
    let mut bulk_decoded = Vec::new();

    for _ in 0..1_000 {
        bytes.clear();
        b64.clear();
        b64_bytes.clear();

        let size = rng.gen_range(1, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        decoded.extend(iter::repeat(0).take(size));
        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let config = random_config(&mut rng);
        let alphabet = random_alphabet(&mut rng);
        // changing padding will cause invalid padding errors when we twiddle the last byte
        let engine = FastPortable::from(alphabet, config.with_encode_padding(false));
        encode_engine_string(&bytes[..], &mut b64, &engine);
        b64_bytes.extend(b64.bytes());
        assert_eq!(b64_bytes.len(), b64.len());

        // change the last character to every possible symbol. Should behave the same as bulk
        // decoding whether invalid or valid.
        for &s1 in alphabet.symbols.iter() {
            decoded.clear();
            bulk_decoded.clear();

            // replace the last
            *b64_bytes.last_mut().unwrap() = s1;
            let bulk_res = decode_engine_vec(&b64_bytes[..], &mut bulk_decoded, &engine);

            let mut wrapped_reader = io::Cursor::new(&b64_bytes[..]);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

            let stream_res = decoder.read_to_end(&mut decoded).map(|_| ()).map_err(|e| {
                e.into_inner()
                    .and_then(|e| e.downcast::<DecodeError>().ok())
            });

            assert_eq!(bulk_res.map_err(|e| Some(Box::new(e))), stream_res);
        }
    }
}

#[test]
fn reports_invalid_byte_correctly() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(1, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        rng.fill_bytes(&mut bytes[..size]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);
        // replace one byte, somewhere, with '*', which is invalid
        let bad_byte_pos = rng.gen_range(0, &b64.len());
        let mut b64_bytes = b64.bytes().collect::<Vec<u8>>();
        b64_bytes[bad_byte_pos] = b'*';

        let mut wrapped_reader = io::Cursor::new(b64_bytes.clone());
        let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

        // some gymnastics to avoid double-moving the io::Error, which is not Copy
        let read_decode_err = decoder
            .read_to_end(&mut decoded)
            .map_err(|e| {
                let kind = e.kind();
                let inner = e
                    .into_inner()
                    .and_then(|e| e.downcast::<DecodeError>().ok());
                inner.map(|i| (*i, kind))
            })
            .err()
            .and_then(|o| o);

        let mut bulk_buf = Vec::new();
        let bulk_decode_err = decode_engine_vec(&b64_bytes[..], &mut bulk_buf, &engine).err();

        // it's tricky to predict where the invalid data's offset will be since if it's in the last
        // chunk it will be reported at the first padding location because it's treated as invalid
        // padding. So, we just check that it's the same as it is for decoding all at once.
        assert_eq!(
            bulk_decode_err.map(|e| (e, io::ErrorKind::InvalidData)),
            read_decode_err
        );
    }
}

fn consume_with_short_reads_and_validate<R: Read>(
    rng: &mut rand::rngs::ThreadRng,
    expected_bytes: &[u8],
    decoded: &mut Vec<u8>,
    short_reader: &mut R,
) {
    let mut total_read = 0_usize;
    loop {
        assert!(
            total_read <= expected_bytes.len(),
            "tr {} size {}",
            total_read,
            expected_bytes.len()
        );
        if total_read == expected_bytes.len() {
            assert_eq!(expected_bytes, &decoded[..total_read]);
            // should be done
            assert_eq!(0, short_reader.read(&mut decoded[..]).unwrap());
            // didn't write anything
            assert_eq!(expected_bytes, &decoded[..total_read]);

            break;
        }
        let decode_len = rng.gen_range(1, cmp::max(2, expected_bytes.len() * 2));

        let read = short_reader
            .read(&mut decoded[total_read..total_read + decode_len])
            .unwrap();
        total_read += read;
    }
}

/// Limits how many bytes a reader will provide in each read call.
/// Useful for shaking out code that may work fine only with typical input sources that always fill
/// the buffer.
struct RandomShortRead<'a, 'b, R: io::Read, N: rand::Rng> {
    delegate: &'b mut R,
    rng: &'a mut N,
}

impl<'a, 'b, R: io::Read, N: rand::Rng> io::Read for RandomShortRead<'a, 'b, R, N> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
        // avoid 0 since it means EOF for non-empty buffers
        let effective_len = cmp::min(self.rng.gen_range(1, 20), buf.len());

        self.delegate.read(&mut buf[..effective_len])
    }
}
<commit_msg>use pointer args<commit_after>use std::io::{self, Read};

use rand::{Rng, RngCore};
use std::{cmp, iter};

use super::decoder::{DecoderReader, BUF_SIZE};
use crate::encode::encode_engine_string;
use crate::engine::fast_portable::FastPortable;
use crate::engine::DEFAULT_ENGINE;
use crate::tests::{random_alphabet, random_config, random_engine};
use crate::{decode_engine_vec, DecodeError};

#[test]
fn simple() {
    let tests: &[(&[u8], &[u8])] = &[
        (&b"0"[..], &b"MA=="[..]),
        (b"01", b"MDE="),
        (b"012", b"MDEy"),
        (b"0123", b"MDEyMw=="),
        (b"01234", b"MDEyMzQ="),
        (b"012345", b"MDEyMzQ1"),
        (b"0123456", b"MDEyMzQ1Ng=="),
        (b"01234567", b"MDEyMzQ1Njc="),
        (b"012345678", b"MDEyMzQ1Njc4"),
        (b"0123456789", b"MDEyMzQ1Njc4OQ=="),
    ][..];

    for (text_expected, base64data) in tests.iter() {
        // Read n bytes at a time.
        for n in 1..base64data.len() + 1 {
            let mut wrapped_reader = io::Cursor::new(base64data);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &DEFAULT_ENGINE);

            // handle errors as you normally would
            let mut text_got = Vec::new();
            let mut buffer = vec![0u8; n];
            while let Ok(read) = decoder.read(&mut buffer[..]) {
                if read == 0 {
                    break;
                }
                text_got.extend_from_slice(&buffer[..read]);
            }

            assert_eq!(
                text_got,
                *text_expected,
                "\nGot: {}\nExpected: {}",
                String::from_utf8_lossy(&text_got[..]),
                String::from_utf8_lossy(text_expected)
            );
        }
    }
}

// Make sure we error out on trailing junk.
#[test]
fn trailing_junk() {
    let tests: &[&[u8]] = &[&b"MDEyMzQ1Njc4*!@#$%^&"[..], b"MDEyMzQ1Njc4OQ== "][..];

    for base64data in tests.iter() {
        // Read n bytes at a time.
        for n in 1..base64data.len() + 1 {
            let mut wrapped_reader = io::Cursor::new(base64data);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &DEFAULT_ENGINE);

            // handle errors as you normally would
            let mut buffer = vec![0u8; n];
            let mut saw_error = false;
            loop {
                match decoder.read(&mut buffer[..]) {
                    Err(_) => {
                        saw_error = true;
                        break;
                    }
                    Ok(read) if read == 0 => break,
                    Ok(_) => (),
                }
            }

            assert!(saw_error);
        }
    }
}

#[test]
fn handles_short_read_from_delegate() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        bytes.truncate(size);
        rng.fill_bytes(&mut bytes[..size]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);
        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut wrapped_reader = io::Cursor::new(b64.as_bytes());
        let mut short_reader = RandomShortRead {
            delegate: &mut wrapped_reader,
            rng: &mut rng,
        };

        let mut decoder = DecoderReader::from(&mut short_reader, &engine);

        let decoded_len = decoder.read_to_end(&mut decoded).unwrap();
        assert_eq!(size, decoded_len);
        assert_eq!(&bytes[..], &decoded[..]);
    }
}

#[test]
fn read_in_short_increments() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        // leave room to play around with larger buffers
        decoded.extend(iter::repeat(0).take(size * 3));

        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut wrapped_reader = io::Cursor::new(&b64[..]);
        let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

        consume_with_short_reads_and_validate(&mut rng, &bytes[..], &mut decoded, &mut decoder);
    }
}

#[test]
fn read_in_short_increments_with_short_delegate_reads() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(0, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        // leave room to play around with larger buffers
        decoded.extend(iter::repeat(0).take(size * 3));

        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);

        let mut base_reader = io::Cursor::new(&b64[..]);
        let mut decoder = DecoderReader::from(&mut base_reader, &engine);
        let mut short_reader = RandomShortRead {
            delegate: &mut decoder,
            rng: &mut rand::thread_rng(),
        };

        consume_with_short_reads_and_validate(&mut rng, &bytes[..], &mut decoded, &mut short_reader)
    }
}

#[test]
fn reports_invalid_last_symbol_correctly() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut b64_bytes = Vec::new();
    let mut decoded = Vec::new();
    let mut bulk_decoded = Vec::new();

    for _ in 0..1_000 {
        bytes.clear();
        b64.clear();
        b64_bytes.clear();

        let size = rng.gen_range(1, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        decoded.extend(iter::repeat(0).take(size));
        rng.fill_bytes(&mut bytes[..]);
        assert_eq!(size, bytes.len());

        let config = random_config(&mut rng);
        let alphabet = random_alphabet(&mut rng);
        // changing padding will cause invalid padding errors when we twiddle the last byte
        let engine = FastPortable::from(alphabet, config.with_encode_padding(false));
        encode_engine_string(&bytes[..], &mut b64, &engine);
        b64_bytes.extend(b64.bytes());
        assert_eq!(b64_bytes.len(), b64.len());

        // change the last character to every possible symbol. Should behave the same as bulk
        // decoding whether invalid or valid.
        for &s1 in alphabet.symbols.iter() {
            decoded.clear();
            bulk_decoded.clear();

            // replace the last
            *b64_bytes.last_mut().unwrap() = s1;
            let bulk_res = decode_engine_vec(&b64_bytes[..], &mut bulk_decoded, &engine);

            let mut wrapped_reader = io::Cursor::new(&b64_bytes[..]);
            let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

            let stream_res = decoder.read_to_end(&mut decoded).map(|_| ()).map_err(|e| {
                e.into_inner()
                    .and_then(|e| e.downcast::<DecodeError>().ok())
            });

            assert_eq!(bulk_res.map_err(|e| Some(Box::new(e))), stream_res);
        }
    }
}

#[test]
fn reports_invalid_byte_correctly() {
    let mut rng = rand::thread_rng();
    let mut bytes = Vec::new();
    let mut b64 = String::new();
    let mut decoded = Vec::new();

    for _ in 0..10_000 {
        bytes.clear();
        b64.clear();
        decoded.clear();

        let size = rng.gen_range(1, 10 * BUF_SIZE);
        bytes.extend(iter::repeat(0).take(size));
        rng.fill_bytes(&mut bytes[..size]);
        assert_eq!(size, bytes.len());

        let engine = random_engine(&mut rng);

        encode_engine_string(&bytes[..], &mut b64, &engine);
        // replace one byte, somewhere, with '*', which is invalid
        let bad_byte_pos = rng.gen_range(0, &b64.len());
        let mut b64_bytes = b64.bytes().collect::<Vec<u8>>();
        b64_bytes[bad_byte_pos] = b'*';

        let mut wrapped_reader = io::Cursor::new(b64_bytes.clone());
        let mut decoder = DecoderReader::from(&mut wrapped_reader, &engine);

        // some gymnastics to avoid double-moving the io::Error, which is not Copy
        let read_decode_err = decoder
            .read_to_end(&mut decoded)
            .map_err(|e| {
                let kind = e.kind();
                let inner = e
                    .into_inner()
                    .and_then(|e| e.downcast::<DecodeError>().ok());
                inner.map(|i| (*i, kind))
            })
            .err()
            .and_then(|o| o);

        let mut bulk_buf = Vec::new();
        let bulk_decode_err = decode_engine_vec(&b64_bytes[..], &mut bulk_buf, &engine).err();

        // it's tricky to predict where the invalid data's offset will be since if it's in the last
        // chunk it will be reported at the first padding location because it's treated as invalid
        // padding. So, we just check that it's the same as it is for decoding all at once.
        assert_eq!(
            bulk_decode_err.map(|e| (e, io::ErrorKind::InvalidData)),
            read_decode_err
        );
    }
}

fn consume_with_short_reads_and_validate<R: Read>(
    rng: &mut rand::rngs::ThreadRng,
    expected_bytes: &[u8],
    decoded: &mut [u8],
    short_reader: &mut R,
) {
    let mut total_read = 0_usize;
    loop {
        assert!(
            total_read <= expected_bytes.len(),
            "tr {} size {}",
            total_read,
            expected_bytes.len()
        );
        if total_read == expected_bytes.len() {
            assert_eq!(expected_bytes, &decoded[..total_read]);
            // should be done
            assert_eq!(0, short_reader.read(&mut *decoded).unwrap());
            // didn't write anything
            assert_eq!(expected_bytes, &decoded[..total_read]);

            break;
        }
        let decode_len = rng.gen_range(1, cmp::max(2, expected_bytes.len() * 2));

        let read = short_reader
            .read(&mut decoded[total_read..total_read + decode_len])
            .unwrap();
        total_read += read;
    }
}

/// Limits how many bytes a reader will provide in each read call.
/// Useful for shaking out code that may work fine only with typical input sources that always fill
/// the buffer.
struct RandomShortRead<'a, 'b, R: io::Read, N: rand::Rng> {
    delegate: &'b mut R,
    rng: &'a mut N,
}

impl<'a, 'b, R: io::Read, N: rand::Rng> io::Read for RandomShortRead<'a, 'b, R, N> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
        // avoid 0 since it means EOF for non-empty buffers
        let effective_len = cmp::min(self.rng.gen_range(1, 20), buf.len());

        self.delegate.read(&mut buf[..effective_len])
    }
}
<|endoftext|>


<commit_before>use {io, EventSet, PollOpt, Token};
use event::IoEvent;
use nix::sys::event::{EventFilter, EventFlag, FilterFlag, KEvent, kqueue, kevent};
use nix::sys::event::{EV_ADD, EV_CLEAR, EV_DELETE, EV_DISABLE, EV_ENABLE, EV_EOF, EV_ONESHOT};
use std::{fmt, slice};
use std::os::unix::io::RawFd;
use std::collections::HashMap;

#[derive(Debug)]
pub struct Selector {
    kq: RawFd,
    changes: Events
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        Ok(Selector {
            kq: try!(kqueue().map_err(super::from_nix_error)),
            changes: Events::new()
        })
    }

    pub fn select(&mut self, evts: &mut Events, timeout_ms: usize) -> io::Result<()> {
        let cnt = try!(kevent(self.kq, self.changes.as_slice(), evts.as_mut_slice(), timeout_ms)
                                  .map_err(super::from_nix_error));

        self.changes.sys_events.clear();

        unsafe {
            evts.sys_events.set_len(cnt);
        }

        evts.coalesce();

        Ok(())
    }

    pub fn register(&mut self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        self.ev_register(fd, token.as_usize(), EventFilter::EVFILT_READ, interests.contains(EventSet::readable()), opts);
        self.ev_register(fd, token.as_usize(), EventFilter::EVFILT_WRITE, interests.contains(EventSet::writable()), opts);

        self.flush_changes()
    }

    pub fn reregister(&mut self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&mut self, fd: RawFd) -> io::Result<()> {
        self.ev_push(fd, 0, EventFilter::EVFILT_READ, EV_DELETE);
        self.ev_push(fd, 0, EventFilter::EVFILT_WRITE, EV_DELETE);

        self.flush_changes()
    }

    fn ev_register(&mut self, fd: RawFd, token: usize, filter: EventFilter, enable: bool, opts: PollOpt) {
        let mut flags = EV_ADD;

        if enable {
            flags = flags | EV_ENABLE;
        } else {
            flags = flags | EV_DISABLE;
        }

        if opts.contains(PollOpt::edge()) {
            flags = flags | EV_CLEAR;
        }

        if opts.contains(PollOpt::oneshot()) {
            flags = flags | EV_ONESHOT;
        }

        self.ev_push(fd, token, filter, flags);
    }

    fn ev_push(&mut self, fd: RawFd, token: usize, filter: EventFilter, flags: EventFlag) {
        self.changes.sys_events.push(
            KEvent {
                ident: fd as ::libc::uintptr_t,
                filter: filter,
                flags: flags,
                fflags: FilterFlag::empty(),
                data: 0,
                udata: token
            });
    }

    fn flush_changes(&mut self) -> io::Result<()> {
        try!(kevent(self.kq, self.changes.as_slice(), &mut [], 0)
             .map_err(super::from_nix_error));

        self.changes.sys_events.clear();
        Ok(())
    }
}

pub struct Events {
    sys_events: Vec<KEvent>,
    events: Vec<IoEvent>,
    event_map: HashMap<Token, usize>,
}

impl Events {
    pub fn new() -> Events {
        Events {
            sys_events: Vec::with_capacity(1024),
            events: Vec::with_capacity(1024),
            event_map: HashMap::with_capacity(1024)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    pub fn get(&self, idx: usize) -> IoEvent {
        self.events[idx]
    }

    pub fn coalesce(&mut self) {
        self.events.clear();
        self.event_map.clear();

        for e in self.sys_events.iter() {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(IoEvent::new(EventSet::none(), token));

            }

            if e.filter == EventFilter::EVFILT_READ {
                self.events[idx].kind.insert(EventSet::readable());
            } else if e.filter == EventFilter::EVFILT_WRITE {
                self.events[idx].kind.insert(EventSet::writable());
            }

            if e.flags.contains(EV_EOF) {
                self.events[idx].kind.insert(EventSet::hup());

                // When the read end of the socket is closed, EV_EOF is set on
                // flags, and fflags contains the error if there is one.
                if !e.fflags.is_empty() {
                    self.events[idx].kind.insert(EventSet::error());
                }
            }
        }
    }

    fn as_slice(&self) -> &[KEvent] {
        unsafe {
            let ptr = (&self.sys_events[..]).as_ptr();
            slice::from_raw_parts(ptr, self.sys_events.len())
        }
    }

    fn as_mut_slice(&mut self) -> &mut [KEvent] {
        unsafe {
            let ptr = (&mut self.sys_events[..]).as_mut_ptr();
            slice::from_raw_parts_mut(ptr, self.sys_events.capacity())
        }
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Events {{ len: {} }}", self.sys_events.len())
    }
}
<commit_msg>Handle kevent's EV_ERROR case.<commit_after>use {io, EventSet, PollOpt, Token};
use event::IoEvent;
use nix::sys::event::{EventFilter, EventFlag, FilterFlag, KEvent, kqueue, kevent};
use nix::sys::event::{EV_ADD, EV_CLEAR, EV_DELETE, EV_DISABLE, EV_ENABLE, EV_EOF, EV_ERROR, EV_ONESHOT};
use std::{fmt, slice};
use std::os::unix::io::RawFd;
use std::collections::HashMap;

#[derive(Debug)]
pub struct Selector {
    kq: RawFd,
    changes: Events
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        Ok(Selector {
            kq: try!(kqueue().map_err(super::from_nix_error)),
            changes: Events::new()
        })
    }

    pub fn select(&mut self, evts: &mut Events, timeout_ms: usize) -> io::Result<()> {
        let cnt = try!(kevent(self.kq, self.changes.as_slice(), evts.as_mut_slice(), timeout_ms)
                                  .map_err(super::from_nix_error));

        self.changes.sys_events.clear();

        unsafe {
            evts.sys_events.set_len(cnt);
        }

        evts.coalesce();

        Ok(())
    }

    pub fn register(&mut self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        self.ev_register(fd, token.as_usize(), EventFilter::EVFILT_READ, interests.contains(EventSet::readable()), opts);
        self.ev_register(fd, token.as_usize(), EventFilter::EVFILT_WRITE, interests.contains(EventSet::writable()), opts);

        self.flush_changes()
    }

    pub fn reregister(&mut self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&mut self, fd: RawFd) -> io::Result<()> {
        self.ev_push(fd, 0, EventFilter::EVFILT_READ, EV_DELETE);
        self.ev_push(fd, 0, EventFilter::EVFILT_WRITE, EV_DELETE);

        self.flush_changes()
    }

    fn ev_register(&mut self, fd: RawFd, token: usize, filter: EventFilter, enable: bool, opts: PollOpt) {
        let mut flags = EV_ADD;

        if enable {
            flags = flags | EV_ENABLE;
        } else {
            flags = flags | EV_DISABLE;
        }

        if opts.contains(PollOpt::edge()) {
            flags = flags | EV_CLEAR;
        }

        if opts.contains(PollOpt::oneshot()) {
            flags = flags | EV_ONESHOT;
        }

        self.ev_push(fd, token, filter, flags);
    }

    fn ev_push(&mut self, fd: RawFd, token: usize, filter: EventFilter, flags: EventFlag) {
        self.changes.sys_events.push(
            KEvent {
                ident: fd as ::libc::uintptr_t,
                filter: filter,
                flags: flags,
                fflags: FilterFlag::empty(),
                data: 0,
                udata: token
            });
    }

    fn flush_changes(&mut self) -> io::Result<()> {
        try!(kevent(self.kq, self.changes.as_slice(), &mut [], 0)
             .map_err(super::from_nix_error));

        self.changes.sys_events.clear();
        Ok(())
    }
}

pub struct Events {
    sys_events: Vec<KEvent>,
    events: Vec<IoEvent>,
    event_map: HashMap<Token, usize>,
}

impl Events {
    pub fn new() -> Events {
        Events {
            sys_events: Vec::with_capacity(1024),
            events: Vec::with_capacity(1024),
            event_map: HashMap::with_capacity(1024)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    pub fn get(&self, idx: usize) -> IoEvent {
        self.events[idx]
    }

    pub fn coalesce(&mut self) {
        self.events.clear();
        self.event_map.clear();

        for e in self.sys_events.iter() {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(IoEvent::new(EventSet::none(), token));

            }

            if e.flags.contains(EV_ERROR) {
                self.events[idx].kind.insert(EventSet::error());
            }

            if e.filter == EventFilter::EVFILT_READ {
                self.events[idx].kind.insert(EventSet::readable());
            } else if e.filter == EventFilter::EVFILT_WRITE {
                self.events[idx].kind.insert(EventSet::writable());
            }

            if e.flags.contains(EV_EOF) {
                self.events[idx].kind.insert(EventSet::hup());

                // When the read end of the socket is closed, EV_EOF is set on
                // flags, and fflags contains the error if there is one.
                if !e.fflags.is_empty() {
                    self.events[idx].kind.insert(EventSet::error());
                }
            }
        }
    }

    fn as_slice(&self) -> &[KEvent] {
        unsafe {
            let ptr = (&self.sys_events[..]).as_ptr();
            slice::from_raw_parts(ptr, self.sys_events.len())
        }
    }

    fn as_mut_slice(&mut self) -> &mut [KEvent] {
        unsafe {
            let ptr = (&mut self.sys_events[..]).as_mut_ptr();
            slice::from_raw_parts_mut(ptr, self.sys_events.capacity())
        }
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Events {{ len: {} }}", self.sys_events.len())
    }
}
<|endoftext|>


<commit_before>use slack::{Event, RtmClient, EventHandler, Message};
use router::{Router,Rule};

pub struct TelescreenHandler {
    router: Router,
}

impl TelescreenHandler {
    pub fn new(router: Router) -> TelescreenHandler {
        TelescreenHandler { router: router }
    }

    pub fn send_message(&self, cli: &RtmClient, source_channel_id: &str, unwrapped_source_user_name: &str, source_text: &str) {
        let channel_name = self.get_channel_name_from_id(cli, source_channel_id);

        let unwrapped_channel_name = match channel_name {
            None => { warn!("No channel: {:?}", source_channel_id); return },
            Some(c) => c,
        };

        let rules: &Vec<Rule> = self.router.rules.as_ref();
        for rule in rules {
            if rule.regex.is_match(unwrapped_channel_name) {
                let dest_channel_id = cli.start_response().channels.as_ref()
                    .and_then(|channels| {
                        channels.iter().find(|chan| match chan.name {
                            None => false,
                            Some(ref name) => name == &(rule.destination),
                        })
                    })
                    .and_then(|chan| chan.id.as_ref());

                let dest_channel_id_unwrap = match dest_channel_id {
                    None => { warn!("No channel: {:?}", dest_channel_id); return },
                    Some(c) => c,
                };

                if unwrapped_channel_name != &(rule.destination) {
                    let message = format!("{:} [ <#{}> ]:\n{:}", unwrapped_source_user_name, source_channel_id, source_text);
                    info!("MESSAGE: {:?}", message);
                    let _ = cli.sender().send_message(&dest_channel_id_unwrap, &message);
                }
            }
        }
    }

    fn get_channel_name_from_id<'a>(&self, cli: &'a RtmClient, channel_id: &str) -> Option<&'a String> {
        cli.start_response().channels.as_ref()
            .and_then(|channels| {
                channels.iter().find(|chan| match chan.id {
                    None => false,
                    Some(ref id) => { id == channel_id },
                })
            })
            .and_then(|chan| chan.name.as_ref())
    }
}

#[allow(unused_variables)]
impl EventHandler for TelescreenHandler {
    fn on_event(&mut self, cli: &RtmClient, event: Event) {
        debug!("EVENT: {:?}", event);

        match event {
            Event::Message(event) => {
                match *event {
                    Message::Standard(message) => {
                        let source_user_id = match message.user {
                            None => { warn!("No user: {:?}", message.user); return },
                            Some(u) => u,
                        };
                        let source_channel_id = match message.channel {
                            None => { warn!("No channel: {:?}", message.channel); return },
                            Some(c) => c,
                        };
                        let source_text = match message.text {
                            None => { warn!("No text: {:?}", message.text); return },
                            Some(t) => t,
                        };

                        let source_user_name = cli.start_response()
                            .users
                            .as_ref()
                            .and_then(|users| {
                                users.iter().find(|user| match user.id {
                                    None => false,
                                    Some(ref name) => name.to_string() == source_user_id,
                                })
                            }).and_then(|user| user.name.as_ref());

                        let unwrapped_source_user_name = match source_user_name {
                            None => { warn!("No user: {:?}", source_user_name); return },
                            Some(t) => t,
                        };

                        self.send_message(cli, &source_channel_id, unwrapped_source_user_name, &source_text);
                    },
                    _ => { /* noop */ },
                }
            },
            _ => { /* noop */ },
        }
    }

    fn on_close(&mut self, cli: &RtmClient) {
        info!("Disconnected");
    }

    fn on_connect(&mut self, cli: &RtmClient) {
        info!("Connected");
    }
}
<commit_msg>Refactor sending message function of TelescreenHandler<commit_after>use slack::{Event, RtmClient, EventHandler, Message};
use router::{Router,Rule};

pub struct TelescreenHandler {
    router: Router,
}

impl TelescreenHandler {
    pub fn new(router: Router) -> TelescreenHandler {
        TelescreenHandler { router: router }
    }

    pub fn send_message(&self, cli: &RtmClient, source_channel_id: &str, source_user_id: &str, source_text: &str) {
        let channel_name = self.get_channel_name_from_id(cli, source_channel_id);
        let unwrapped_channel_name = match channel_name {
            None => { warn!("No channel: {:?}", source_channel_id); return },
            Some(c) => c,
        };

        let user_name = self.get_user_name_from_id(cli, source_user_id);
        let unwrapped_user_name = match user_name {
            None => { warn!("No user: {:?}", user_name); return },
            Some(t) => t,
        };

        let rules: &Vec<Rule> = self.router.rules.as_ref();
        for rule in rules {
            if rule.regex.is_match(unwrapped_channel_name) {
                let dest_channel_id = cli.start_response().channels.as_ref()
                    .and_then(|channels| {
                        channels.iter().find(|chan| match chan.name {
                            None => false,
                            Some(ref name) => name == &(rule.destination),
                        })
                    })
                    .and_then(|chan| chan.id.as_ref());

                let dest_channel_id_unwrap = match dest_channel_id {
                    None => { warn!("No channel: {:?}", dest_channel_id); return },
                    Some(c) => c,
                };

                if unwrapped_channel_name != &(rule.destination) {
                    let message = format!("@{:} [ <#{}> ]:\n{:}", unwrapped_user_name, source_channel_id, source_text);
                    info!("MESSAGE: {:?}", message);
                    let _ = cli.sender().send_message(&dest_channel_id_unwrap, &message);
                }
            }
        }
    }

    fn get_channel_name_from_id<'a>(&self, cli: &'a RtmClient, channel_id: &str) -> Option<&'a String> {
        cli.start_response().channels.as_ref()
            .and_then(|channels| {
                channels.iter().find(|chan| match chan.id {
                    None => false,
                    Some(ref id) => { id == channel_id },
                })
            })
            .and_then(|chan| chan.name.as_ref())
    }

    fn get_user_name_from_id<'a>(&self, cli: &'a RtmClient, user_id: &str) -> Option<&'a String> {
        cli.start_response().users.as_ref()
            .and_then(|users| {
                users.iter().find(|user| match user.id {
                    None => false,
                    Some(ref name) => { name.to_string() == user_id },
                })
            })
            .and_then(|user| user.name.as_ref())
    }
}

#[allow(unused_variables)]
impl EventHandler for TelescreenHandler {
    fn on_event(&mut self, cli: &RtmClient, event: Event) {
        debug!("EVENT: {:?}", event);

        match event {
            Event::Message(event) => {
                match *event {
                    Message::Standard(message) => {
                        let source_channel_id = match message.channel {
                            None => { warn!("No channel: {:?}", message.channel); return },
                            Some(c) => c,
                        };
                        let source_user_id = match message.user {
                            None => { warn!("No user: {:?}", message.user); return },
                            Some(u) => u,
                        };
                        let source_text = match message.text {
                            None => { warn!("No text: {:?}", message.text); return },
                            Some(t) => t,
                        };

                        self.send_message(cli, &source_channel_id, &source_user_id, &source_text);
                    },
                    _ => { /* noop */ },
                }
            },
            _ => { /* noop */ },
        }
    }

    fn on_close(&mut self, cli: &RtmClient) {
        info!("Disconnected");
    }

    fn on_connect(&mut self, cli: &RtmClient) {
        info!("Connected");
    }
}
<|endoftext|>