<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::{io, os, str, run, task};
use extra::{tempfile, json};
use extra::json::ToJson;

use super::*;

fn write_whole_file(path: &Path, data: &str) {
    match io::file_writer(path, [io::Create]) {
        Ok(writer) => writer.write_str(data),
        Err(message) => fail!(message),
    }
}

fn almost_equals(a: &json::Json, b: &json::Json) -> bool {
    match (a, b) {
        (&json::Number(a), &json::Number(b)) => (a - b).abs() < 1e-6,
        (&json::String(ref a), &json::String(ref b)) => a == b,
        (&json::Boolean(a), &json::Boolean(b)) => a == b,
        (&json::List(ref a), &json::List(ref b))
            => a.iter().zip(b.iter()).all(|(ref a, ref b)| almost_equals(*a, *b)),
        (&json::Object(_), &json::Object(_)) => fail!(~"Not implemented"),
        (&json::Null, &json::Null) => true,
        _ => false,
    }
}

fn assert_json_eq(results: json::Json, expected: json::Json, message: ~str) {
    if !almost_equals(&results, &expected) {
        let temp = tempfile::mkdtemp(&os::tmpdir(), "rust-cssparser-tests").unwrap();
        let temp_ = temp.clone();
        let results = json::to_pretty_str(&results) + "\n";
        let expected = json::to_pretty_str(&expected) + "\n";
        do task::try {
            let result_path = temp.push("results.json");
            let expected_path = temp.push("expected.json");
            write_whole_file(&result_path, results);
            write_whole_file(&expected_path, expected);
            run::process_status("colordiff", [~"-u1000", result_path.to_str(),
                                              expected_path.to_str()]);
        };
        os::remove_dir_recursive(&temp_);

        fail!(message)
    }
}

fn run_json_tests<T: ToJson>(json_data: &str, parse: &fn (input: ~str) -> T) {
    let items = match json::from_str(json_data) {
        Ok(json::List(items)) => items,
        _ => fail!("Invalid JSON")
    };
    assert!(items.len() % 2 == 0);
    let mut input: Option<~str> = None;
    for item in items.consume_iter() {
        match (&input, item) {
            (&None, json::String(string)) => input = Some(string),
            (&Some(_), expected) => {
                let input = input.take_unwrap();
                let result = parse(input.to_owned()).to_json();
                assert_json_eq(result, expected, input);
            },
            _ => fail!("Unexpected JSON")
        };
    }
}

#[test]
fn component_value_list() {
    do run_json_tests(include_str!("css-parsing-tests/component_value_list.json")) |input| {
        tokenize(input).transform(|(c, _)| c).to_owned_vec()
    }
}

#[test]
fn one_component_value() {
    do run_json_tests(include_str!("css-parsing-tests/one_component_value.json")) |input| {
        parse_one_component_value(tokenize(input))
    }
}

#[test]
fn declaration_list() {
    do run_json_tests(include_str!("css-parsing-tests/declaration_list.json")) |input| {
        parse_declaration_list(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn one_declaration() {
    do run_json_tests(include_str!("css-parsing-tests/one_declaration.json")) |input| {
        parse_one_declaration(tokenize(input))
    }
}

#[test]
fn rule_list() {
    do run_json_tests(include_str!("css-parsing-tests/rule_list.json")) |input| {
        parse_rule_list(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn stylesheet() {
    do run_json_tests(include_str!("css-parsing-tests/stylesheet.json")) |input| {
        parse_stylesheet_rules(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn one_rule() {
    do run_json_tests(include_str!("css-parsing-tests/one_rule.json")) |input| {
        parse_one_rule(tokenize(input))
    }
}

fn run_color_tests(json_data: &str, to_json: &fn(result: Option<Color>) -> json::Json) {
    do run_json_tests(json_data) |input| {
        match parse_one_component_value(tokenize(input)) {
            Ok(component_value) => to_json(parse_color(&component_value)),
            Err(_reason) => json::Null,
        }
    }
}

#[test]
fn color3() {
    run_color_tests(include_str!("css-parsing-tests/color3.json"), |c| c.to_json())
}

#[test]
fn color3_hsl() {
    run_color_tests(include_str!("css-parsing-tests/color3_hsl.json"), |c| c.to_json())
}

/// color3_keywords.json is different: R, G and B are in 0..255 rather than 0..1
#[test]
fn color3_keywords() {
    do run_color_tests(include_str!("css-parsing-tests/color3_keywords.json")) |c| {
        match c {
            Some(RGBA(r, g, b, a)) => (~[r * 255., g * 255., b * 255., a]).to_json(),
            Some(CurrentColor) => json::String(~"currentColor"),
            None => json::Null,
        }
    }
}

impl ToJson for Result<Rule, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<DeclarationListItem, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<Declaration, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<ComponentValue, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for SyntaxError {
    fn to_json(&self) -> json::Json {
        json::List(~[json::String(~"error"), json::String(match self.reason {
            ErrEmptyInput => ~"empty",
            ErrExtraInput => ~"extra-input",
            _ => ~"invalid",
        })])
    }
}

impl ToJson for Color {
    fn to_json(&self) -> json::Json {
        match *self {
            RGBA(r, g, b, a) => (~[r, g, b, a]).to_json(),
            CurrentColor => json::String(~"currentColor"),
        }
    }
}

impl ToJson for Rule {
    fn to_json(&self) -> json::Json {
        match *self {
            QualifiedRule(ref rule) => rule.to_json(),
            AtRule(ref rule) => rule.to_json(),
        }
    }
}

impl ToJson for DeclarationListItem {
    fn to_json(&self) -> json::Json {
        match *self {
            Declaration(ref declaration) => declaration.to_json(),
            Decl_AtRule(ref at_rule) => at_rule.to_json(),
        }
    }
}

fn list_to_json(list: &~[(ComponentValue, SourceLocation)]) -> ~[json::Json] {
    list.map(|tuple| {
        match *tuple {
            (ref c, _) => c.to_json()
        }
    })
}

impl ToJson for AtRule {
    fn to_json(&self) -> json::Json {
        match *self {
            AtRule{name: ref name, prelude: ref prelude, block: ref block, _}
            => json::List(~[json::String(~"at-rule"), name.to_json(),
                            prelude.to_json(), block.map(list_to_json).to_json()])
        }
    }
}

impl ToJson for QualifiedRule {
    fn to_json(&self) -> json::Json {
        match *self {
            QualifiedRule{prelude: ref prelude, block: ref block, _}
            => json::List(~[json::String(~"qualified rule"),
                            prelude.to_json(), json::List(list_to_json(block))])
        }
    }
}

impl ToJson for Declaration {
    fn to_json(&self) -> json::Json {
        match *self {
            Declaration{name: ref name, value: ref value, important: ref important, _}
            =>  json::List(~[json::String(~"declaration"), name.to_json(),
                             value.to_json(), important.to_json()])
        }
    }
}

impl ToJson for ComponentValue {
    fn to_json(&self) -> json::Json {
        use JList = extra::json::List;
        use JString = extra::json::String;

        fn numeric(value: &NumericValue) -> ~[json::Json] {
            match *value {
                NumericValue{representation: ref r, value: ref v, int_value: ref i}
                => ~[r.to_json(), v.to_json(),
                     JString(match *i { Some(_) => ~"integer", _ => ~"number" })]
            }
        }

        match *self {
            Ident(ref value) => JList(~[JString(~"ident"), value.to_json()]),
            AtKeyword(ref value) => JList(~[JString(~"at-keyword"), value.to_json()]),
            Hash(ref value) => JList(~[JString(~"hash"), value.to_json(),
                                       JString(~"unrestricted")]),
            IDHash(ref value) => JList(~[JString(~"hash"), value.to_json(), JString(~"id")]),
            String(ref value) => JList(~[JString(~"string"), value.to_json()]),
            URL(ref value) => JList(~[JString(~"url"), value.to_json()]),
            Delim('\\') => JString(~"\\"),
            Delim(value) => JString(str::from_char(value)),

            Number(ref value) => JList(~[JString(~"number")] + numeric(value)),
            Percentage(ref value) => JList(~[JString(~"percentage")] + numeric(value)),
            Dimension(ref value, ref unit)
            => JList(~[JString(~"dimension")] + numeric(value) + ~[unit.to_json()]),

            // TODO:
            UnicodeRange(_start, _end) => fail!(),
            EmptyUnicodeRange => fail!(),

            WhiteSpace => JString(~" "),
            Colon => JString(~":"),
            Semicolon => JString(~";"),
            Comma => JString(~","),
            IncludeMath => JString(~"~="),
            DashMatch => JString(~"|="),
            PrefixMatch => JString(~"^="),
            SuffixMatch => JString(~"$="),
            SubstringMatch => JString(~"*="),
            Column => JString(~"||"),
            CDO => JString(~"<!--"),
            CDC => JString(~"-->"),

            Function(ref name, ref arguments)
            => JList(~[JString(~"function"), name.to_json()] + arguments.map(|a| a.to_json())),
            ParenthesisBlock(ref content)
            => JList(~[JString(~"()")] + content.map(|c| c.to_json())),
            SquareBracketBlock(ref content)
            => JList(~[JString(~"[]")] + content.map(|c| c.to_json())),
            CurlyBracketBlock(ref content)
            => JList(~[JString(~"{}")] + list_to_json(content)),

            BadURL => JList(~[JString(~"error"), JString(~"bad-url")]),
            BadString => JList(~[JString(~"error"), JString(~"bad-string")]),
            CloseParenthesis => JList(~[JString(~"error"), JString(~")")]),
            CloseSquareBracket => JList(~[JString(~"error"), JString(~"]")]),
            CloseCurlyBracket => JList(~[JString(~"error"), JString(~"}")]),
        }
    }
}
<commit_msg>Typo fix.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::{io, os, str, run, task};
use extra::{tempfile, json};
use extra::json::ToJson;

use super::*;

fn write_whole_file(path: &Path, data: &str) {
    match io::file_writer(path, [io::Create]) {
        Ok(writer) => writer.write_str(data),
        Err(message) => fail!(message),
    }
}

fn almost_equals(a: &json::Json, b: &json::Json) -> bool {
    match (a, b) {
        (&json::Number(a), &json::Number(b)) => (a - b).abs() < 1e-6,
        (&json::String(ref a), &json::String(ref b)) => a == b,
        (&json::Boolean(a), &json::Boolean(b)) => a == b,
        (&json::List(ref a), &json::List(ref b))
            => a.iter().zip(b.iter()).all(|(ref a, ref b)| almost_equals(*a, *b)),
        (&json::Object(_), &json::Object(_)) => fail!(~"Not implemented"),
        (&json::Null, &json::Null) => true,
        _ => false,
    }
}

fn assert_json_eq(results: json::Json, expected: json::Json, message: ~str) {
    if !almost_equals(&results, &expected) {
        let temp = tempfile::mkdtemp(&os::tmpdir(), "rust-cssparser-tests").unwrap();
        let temp_ = temp.clone();
        let results = json::to_pretty_str(&results) + "\n";
        let expected = json::to_pretty_str(&expected) + "\n";
        do task::try {
            let result_path = temp.push("results.json");
            let expected_path = temp.push("expected.json");
            write_whole_file(&result_path, results);
            write_whole_file(&expected_path, expected);
            run::process_status("colordiff", [~"-u1000", result_path.to_str(),
                                              expected_path.to_str()]);
        };
        os::remove_dir_recursive(&temp_);

        fail!(message)
    }
}

fn run_json_tests<T: ToJson>(json_data: &str, parse: &fn (input: ~str) -> T) {
    let items = match json::from_str(json_data) {
        Ok(json::List(items)) => items,
        _ => fail!("Invalid JSON")
    };
    assert!(items.len() % 2 == 0);
    let mut input: Option<~str> = None;
    for item in items.consume_iter() {
        match (&input, item) {
            (&None, json::String(string)) => input = Some(string),
            (&Some(_), expected) => {
                let input = input.take_unwrap();
                let result = parse(input.to_owned()).to_json();
                assert_json_eq(result, expected, input);
            },
            _ => fail!("Unexpected JSON")
        };
    }
}

#[test]
fn component_value_list() {
    do run_json_tests(include_str!("css-parsing-tests/component_value_list.json")) |input| {
        tokenize(input).transform(|(c, _)| c).to_owned_vec()
    }
}

#[test]
fn one_component_value() {
    do run_json_tests(include_str!("css-parsing-tests/one_component_value.json")) |input| {
        parse_one_component_value(tokenize(input))
    }
}

#[test]
fn declaration_list() {
    do run_json_tests(include_str!("css-parsing-tests/declaration_list.json")) |input| {
        parse_declaration_list(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn one_declaration() {
    do run_json_tests(include_str!("css-parsing-tests/one_declaration.json")) |input| {
        parse_one_declaration(tokenize(input))
    }
}

#[test]
fn rule_list() {
    do run_json_tests(include_str!("css-parsing-tests/rule_list.json")) |input| {
        parse_rule_list(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn stylesheet() {
    do run_json_tests(include_str!("css-parsing-tests/stylesheet.json")) |input| {
        parse_stylesheet_rules(tokenize(input)).to_owned_vec()
    }
}

#[test]
fn one_rule() {
    do run_json_tests(include_str!("css-parsing-tests/one_rule.json")) |input| {
        parse_one_rule(tokenize(input))
    }
}

fn run_color_tests(json_data: &str, to_json: &fn(result: Option<Color>) -> json::Json) {
    do run_json_tests(json_data) |input| {
        match parse_one_component_value(tokenize(input)) {
            Ok(component_value) => to_json(parse_color(&component_value)),
            Err(_reason) => json::Null,
        }
    }
}

#[test]
fn color3() {
    run_color_tests(include_str!("css-parsing-tests/color3.json"), |c| c.to_json())
}

#[test]
fn color3_hsl() {
    run_color_tests(include_str!("css-parsing-tests/color3_hsl.json"), |c| c.to_json())
}

/// color3_keywords.json is different: R, G and B are in 0..255 rather than 0..1
#[test]
fn color3_keywords() {
    do run_color_tests(include_str!("css-parsing-tests/color3_keywords.json")) |c| {
        match c {
            Some(RGBA(r, g, b, a)) => (~[r * 255., g * 255., b * 255., a]).to_json(),
            Some(CurrentColor) => json::String(~"currentColor"),
            None => json::Null,
        }
    }
}

impl ToJson for Result<Rule, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<DeclarationListItem, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<Declaration, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for Result<ComponentValue, SyntaxError> {
    fn to_json(&self) -> json::Json {
        match *self {
            Ok(ref a) => a.to_json(),
            Err(ref b) => b.to_json(),
        }
    }
}

impl ToJson for SyntaxError {
    fn to_json(&self) -> json::Json {
        json::List(~[json::String(~"error"), json::String(match self.reason {
            ErrEmptyInput => ~"empty",
            ErrExtraInput => ~"extra-input",
            _ => ~"invalid",
        })])
    }
}

impl ToJson for Color {
    fn to_json(&self) -> json::Json {
        match *self {
            RGBA(r, g, b, a) => (~[r, g, b, a]).to_json(),
            CurrentColor => json::String(~"currentColor"),
        }
    }
}

impl ToJson for Rule {
    fn to_json(&self) -> json::Json {
        match *self {
            QualifiedRule(ref rule) => rule.to_json(),
            AtRule(ref rule) => rule.to_json(),
        }
    }
}

impl ToJson for DeclarationListItem {
    fn to_json(&self) -> json::Json {
        match *self {
            Declaration(ref declaration) => declaration.to_json(),
            Decl_AtRule(ref at_rule) => at_rule.to_json(),
        }
    }
}

fn list_to_json(list: &~[(ComponentValue, SourceLocation)]) -> ~[json::Json] {
    list.map(|tuple| {
        match *tuple {
            (ref c, _) => c.to_json()
        }
    })
}

impl ToJson for AtRule {
    fn to_json(&self) -> json::Json {
        match *self {
            AtRule{name: ref name, prelude: ref prelude, block: ref block, _}
            => json::List(~[json::String(~"at-rule"), name.to_json(),
                            prelude.to_json(), block.map(list_to_json).to_json()])
        }
    }
}

impl ToJson for QualifiedRule {
    fn to_json(&self) -> json::Json {
        match *self {
            QualifiedRule{prelude: ref prelude, block: ref block, _}
            => json::List(~[json::String(~"qualified rule"),
                            prelude.to_json(), json::List(list_to_json(block))])
        }
    }
}

impl ToJson for Declaration {
    fn to_json(&self) -> json::Json {
        match *self {
            Declaration{name: ref name, value: ref value, important: ref important, _}
            =>  json::List(~[json::String(~"declaration"), name.to_json(),
                             value.to_json(), important.to_json()])
        }
    }
}

impl ToJson for ComponentValue {
    fn to_json(&self) -> json::Json {
        use JList = extra::json::List;
        use JString = extra::json::String;

        fn numeric(value: &NumericValue) -> ~[json::Json] {
            match *value {
                NumericValue{representation: ref r, value: ref v, int_value: ref i}
                => ~[r.to_json(), v.to_json(),
                     JString(match *i { Some(_) => ~"integer", _ => ~"number" })]
            }
        }

        match *self {
            Ident(ref value) => JList(~[JString(~"ident"), value.to_json()]),
            AtKeyword(ref value) => JList(~[JString(~"at-keyword"), value.to_json()]),
            Hash(ref value) => JList(~[JString(~"hash"), value.to_json(),
                                       JString(~"unrestricted")]),
            IDHash(ref value) => JList(~[JString(~"hash"), value.to_json(), JString(~"id")]),
            String(ref value) => JList(~[JString(~"string"), value.to_json()]),
            URL(ref value) => JList(~[JString(~"url"), value.to_json()]),
            Delim('\\') => JString(~"\\"),
            Delim(value) => JString(str::from_char(value)),

            Number(ref value) => JList(~[JString(~"number")] + numeric(value)),
            Percentage(ref value) => JList(~[JString(~"percentage")] + numeric(value)),
            Dimension(ref value, ref unit)
            => JList(~[JString(~"dimension")] + numeric(value) + ~[unit.to_json()]),

            // TODO:
            UnicodeRange(_start, _end) => fail!(),
            EmptyUnicodeRange => fail!(),

            WhiteSpace => JString(~" "),
            Colon => JString(~":"),
            Semicolon => JString(~";"),
            Comma => JString(~","),
            IncludeMatch => JString(~"~="),
            DashMatch => JString(~"|="),
            PrefixMatch => JString(~"^="),
            SuffixMatch => JString(~"$="),
            SubstringMatch => JString(~"*="),
            Column => JString(~"||"),
            CDO => JString(~"<!--"),
            CDC => JString(~"-->"),

            Function(ref name, ref arguments)
            => JList(~[JString(~"function"), name.to_json()] + arguments.map(|a| a.to_json())),
            ParenthesisBlock(ref content)
            => JList(~[JString(~"()")] + content.map(|c| c.to_json())),
            SquareBracketBlock(ref content)
            => JList(~[JString(~"[]")] + content.map(|c| c.to_json())),
            CurlyBracketBlock(ref content)
            => JList(~[JString(~"{}")] + list_to_json(content)),

            BadURL => JList(~[JString(~"error"), JString(~"bad-url")]),
            BadString => JList(~[JString(~"error"), JString(~"bad-string")]),
            CloseParenthesis => JList(~[JString(~"error"), JString(~")")]),
            CloseSquareBracket => JList(~[JString(~"error"), JString(~"]")]),
            CloseCurlyBracket => JList(~[JString(~"error"), JString(~"}")]),
        }
    }
}
<|endoftext|>


<commit_before>#![crate_id(name="tr", vers="1.0.0", author="Michael Gehring")]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate collections;
extern crate getopts;

use collections::hashmap::{HashMap, HashSet};
use getopts::OptGroup;
use std::char::from_u32;
use std::io::print;
use std::io::stdio::{stdin,stdout};
use std::iter::FromIterator;
use std::os;
use std::vec::Vec;

static NAME : &'static str = "tr";
static VERSION : &'static str = "1.0.0";

fn unescape_char(c: char) -> char {
    match c {
        'a' => 0x07u8 as char,
        'b' => 0x08u8 as char,
        'f' => 0x0cu8 as char,
        'v' => 0x0bu8 as char,
        'n' => '\n',
        'r' => '\r',
        't' => '\t',
        _ => c,
    }
}

fn unescape(v: Vec<char>) -> Vec<char> {
    let mut out = Vec::new();
    let mut input = v.as_slice();
    loop {
        input = match input {
            ['\\', e, ..rest] => {
                out.push(unescape_char(e));
                rest
            }
            [c, ..rest] => {
                out.push(c);
                rest
            }
            [] => break
        }
    }
    out
}

fn expand_range(from: char, to: char) -> Vec<char> {
    range(from as u32, to as u32 + 1).map(|c| from_u32(c).unwrap()).collect()
}

fn expand_set(s: &str) -> Vec<char> {
    let mut set = Vec::<char>::new();
    let unesc = unescape(FromIterator::from_iter(s.chars()));
    let mut input = unesc.as_slice();

    loop {
        input = match input {
            [f, '-', t, ..rest] => {
                set.push_all(expand_range(f, t).as_slice());
                rest
            }
            [c, ..rest] => {
                set.push(c);
                rest
            }
            [] => break
        };
    }
    set
}

fn delete(set: Vec<char>) {
    let mut hset = HashSet::new();
    let mut out = stdout();

    for &c in set.iter() {
        hset.insert(c);
    }

    for c in stdin().chars() {
        match c {
            Ok(c) if !hset.contains(&c) => out.write_char(c).unwrap(),
            Ok(_) => (),
            Err(err) => fail!("{}", err),
        };
    }
}

fn tr(set1: &[char], set2: &[char]) {
    let mut map = HashMap::new();
    let mut out = stdout();

    for i in range(0, set1.len()) {
        if i >= set2.len() {
            map.insert(set1[i], set2[set2.len()-1]);
        } else {
            map.insert(set1[i], set2[i]);
        }
    }

    for c in stdin().chars() {
        match c {
            Ok(inc) => {
                let trc = match map.find(&inc) {
                    Some(t) => *t,
                    None => inc,
                };
                out.write_char(trc).unwrap();
            }
            Err(err) => {
                fail!("{}", err);
            }
        }
    }
}

fn usage(opts: &[OptGroup]) {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] SET1 [SET2]", NAME);
        println!("");
        print(getopts::usage("Translate or delete characters.", opts).as_slice());
}

pub fn main() {
    let args: Vec<StrBuf> = os::args().iter().map(|x| x.to_strbuf()).collect();
    let opts = [
        getopts::optflag("d", "delete", "delete characters in SET1"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => fail!("{}", err.to_err_msg()),
    };

    if matches.opt_present("help") {
        usage(opts);
        return
    }

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return;
    }

    if matches.free.len() == 0 {
        usage(opts);
        os::set_exit_status(1);
        return
    }

    let dflag = matches.opt_present("d");
    let sets = matches.free;

    if dflag {
        let set1 = expand_set(sets.get(0).as_slice());
        delete(set1);
    } else {
        let set1 = expand_set(sets.get(0).as_slice());
        let set2 = expand_set(sets.get(1).as_slice());
        tr(set1.as_slice(), set2.as_slice());
    }
}
<commit_msg>tr: use SmallIntMap instead of HashMap<commit_after>#![crate_id(name="tr", vers="1.0.0", author="Michael Gehring")]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate collections;
extern crate getopts;

use collections::hashmap::HashSet;
use collections::smallintmap::SmallIntMap;
use getopts::OptGroup;
use std::char::from_u32;
use std::io::print;
use std::io::stdio::{stdin,stdout};
use std::iter::FromIterator;
use std::os;
use std::vec::Vec;

static NAME : &'static str = "tr";
static VERSION : &'static str = "1.0.0";

fn unescape_char(c: char) -> char {
    match c {
        'a' => 0x07u8 as char,
        'b' => 0x08u8 as char,
        'f' => 0x0cu8 as char,
        'v' => 0x0bu8 as char,
        'n' => '\n',
        'r' => '\r',
        't' => '\t',
        _ => c,
    }
}

fn unescape(v: Vec<char>) -> Vec<char> {
    let mut out = Vec::new();
    let mut input = v.as_slice();
    loop {
        input = match input {
            ['\\', e, ..rest] => {
                out.push(unescape_char(e));
                rest
            }
            [c, ..rest] => {
                out.push(c);
                rest
            }
            [] => break
        }
    }
    out
}

fn expand_range(from: char, to: char) -> Vec<char> {
    range(from as u32, to as u32 + 1).map(|c| from_u32(c).unwrap()).collect()
}

fn expand_set(s: &str) -> Vec<char> {
    let mut set = Vec::<char>::new();
    let unesc = unescape(FromIterator::from_iter(s.chars()));
    let mut input = unesc.as_slice();

    loop {
        input = match input {
            [f, '-', t, ..rest] => {
                set.push_all(expand_range(f, t).as_slice());
                rest
            }
            [c, ..rest] => {
                set.push(c);
                rest
            }
            [] => break
        };
    }
    set
}

fn delete(set: Vec<char>) {
    let mut hset = HashSet::new();
    let mut out = stdout();

    for &c in set.iter() {
        hset.insert(c);
    }

    for c in stdin().chars() {
        match c {
            Ok(c) if !hset.contains(&c) => out.write_char(c).unwrap(),
            Ok(_) => (),
            Err(err) => fail!("{}", err),
        };
    }
}

fn tr(set1: &[char], set2: &[char]) {
    let mut map = SmallIntMap::<char>::new();
    let mut out = stdout();

    for i in range(0, set1.len()) {
        if i >= set2.len() {
            map.insert(set1[i] as uint, set2[set2.len()-1]);
        } else {
            map.insert(set1[i] as uint, set2[i]);
        }
    }

    for c in stdin().chars() {
        match c {
            Ok(inc) => {
                let trc = match map.find(&(inc as uint)) {
                    Some(t) => *t,
                    None => inc,
                };
                out.write_char(trc).unwrap();
            }
            Err(err) => {
                fail!("{}", err);
            }
        }
    }
}

fn usage(opts: &[OptGroup]) {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] SET1 [SET2]", NAME);
        println!("");
        print(getopts::usage("Translate or delete characters.", opts).as_slice());
}

pub fn main() {
    let args: Vec<StrBuf> = os::args().iter().map(|x| x.to_strbuf()).collect();
    let opts = [
        getopts::optflag("d", "delete", "delete characters in SET1"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => fail!("{}", err.to_err_msg()),
    };

    if matches.opt_present("help") {
        usage(opts);
        return
    }

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return;
    }

    if matches.free.len() == 0 {
        usage(opts);
        os::set_exit_status(1);
        return
    }

    let dflag = matches.opt_present("d");
    let sets = matches.free;

    if dflag {
        let set1 = expand_set(sets.get(0).as_slice());
        delete(set1);
    } else {
        let set1 = expand_set(sets.get(0).as_slice());
        let set2 = expand_set(sets.get(1).as_slice());
        tr(set1.as_slice(), set2.as_slice());
    }
}
<|endoftext|>


<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json::{self, Value};
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::error::Error as StdError;

#[derive(Debug)]
pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for Value {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        // Arrays are a special case: since ReQL commands are sent as arrays
        // We have to wrap them using MAKE_ARRAY
        let val = wrap_arrays(self.clone());
        match serde_json::to_string(&val) {
            Ok(cmd) => Ok((Some(cmd), None)),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl<T> IntoCommandArg for (T, Value)
        where T: IntoCommandArg
{
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        if !self.1.is_object() {
            let msg = String::from("Only objects are allowed as function options. You should use `r.object()` to pass optional arguments in your functions.");
            return Err(From::from(DriverError::Other(msg)));
        }
        let arg = try!(self.0.to_arg());
        let opt = try!(self.1.to_arg());
        Ok((arg.0, opt.0))
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match self.0 {
            Ok(ref cmd) => Ok((Some(cmd.to_string()), None)),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

macro_rules! define {
    (impl IntoCommandArg for $T:ty) => {
        impl IntoCommandArg for $T {
            fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
                Ok((Some(self.to_string()), None))
            }
        }
    }
}

define!{ impl IntoCommandArg for bool }
define!{ impl IntoCommandArg for char }
define!{ impl IntoCommandArg for u8 }
define!{ impl IntoCommandArg for u16 }
define!{ impl IntoCommandArg for u32 }
define!{ impl IntoCommandArg for u64 }
define!{ impl IntoCommandArg for usize }
define!{ impl IntoCommandArg for i8 }
define!{ impl IntoCommandArg for i16 }
define!{ impl IntoCommandArg for i32 }
define!{ impl IntoCommandArg for i64 }
define!{ impl IntoCommandArg for isize }
define!{ impl IntoCommandArg for f32 }
define!{ impl IntoCommandArg for f64 }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        match e.to_arg() {
            Ok(arg) => if let Some(arg) = arg.0 {
                RootCommand(Ok(arg))
            } else {
                RootCommand(Ok(String::new()))
            },
            Err(e) => RootCommand(Err(e)),
        }
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn table_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_drop(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

macro_rules! command {
    ($name:ident, $cmd:ident) => {
        pub fn $name<T>(self, arg: T) -> RootCommand
            where T: IntoCommandArg
            {
                let commands = match self.0 {
                    Ok(t) => t,
                    Err(e) => return RootCommand(Err(e)),
                };
                RootCommand(Command::wrap(tt::$cmd,
                                          Some(arg),
                                          Some(&commands)))
            }
    }
}

impl RootCommand {
    command!(table_create, TABLE_CREATE);
    command!(table_drop, TABLE_DROP);
    command!(table, TABLE);
    command!(filter, FILTER);
    command!(insert, INSERT);

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE,
                                  None as Option<&str>,
                                  Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                input: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
        {
            let mut cmds = format!("[{},", command.value());
            let mut args = String::new();
            if let Some(commands) = commands {
                args.push_str(commands);
                if input.is_some() {
                    args.push_str(",");
                }
            }
            let mut arguments: Option<String> = None;
            let mut options: Option<String> = None;
            if let Some(input) = input {
                let (args, opts) = try!(input.to_arg());
                arguments = args;
                options = opts;
            }
            if let Some(arguments) = arguments {
                args.push_str(&arguments);
            }
            cmds.push_str(&format!("[{}]", args));
            if let Some(options) = options {
                cmds.push_str(&format!(",{}", options));
            }
            cmds.push(']');
            Ok(cmds)
        }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
        -> String {
            let mut qry = format!("[{}", query_type.value());
            if let Some(query) = query {
                qry.push_str(&format!(",{}", query));
            }
            if let Some(options) = options {
                qry.push_str(&format!(",{}", options));
            }
            qry.push_str("]");
            qry
        }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}

fn wrap_arrays(mut val: Value) -> Value {
    if val.is_array() {
        let mut array = Vec::with_capacity(2);
        array.push(Value::I64(tt::MAKE_ARRAY.value() as i64));
        if let Value::Array(vec) = val {
            let mut new_val = Vec::with_capacity(vec.len());
            for v in vec.into_iter() {
                if v.is_array() || v.is_object() {
                    new_val.push(wrap_arrays(v));
                } else {
                    new_val.push(v)
                }
            }
            val = Value::Array(new_val);
        }
        array.push(val);
        val = Value::Array(array);
    } else if val.is_object() {
        if let Value::Object(mut obj) = val {
            for (k, v) in obj.clone() {
                obj.insert(k, wrap_arrays(v));
            }
            val = Value::Object(obj);
        }
    }
    val
}
<commit_msg>Add more commands<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json::{self, Value};
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::error::Error as StdError;

#[derive(Debug)]
pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for Value {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        // Arrays are a special case: since ReQL commands are sent as arrays
        // We have to wrap them using MAKE_ARRAY
        let val = wrap_arrays(self.clone());
        match serde_json::to_string(&val) {
            Ok(cmd) => Ok((Some(cmd), None)),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl<T> IntoCommandArg for (T, Value)
        where T: IntoCommandArg
{
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        if !self.1.is_object() {
            let msg = String::from("Only objects are allowed as function options. You should use `r.object()` to pass optional arguments in your functions.");
            return Err(From::from(DriverError::Other(msg)));
        }
        let arg = try!(self.0.to_arg());
        let opt = try!(self.1.to_arg());
        Ok((arg.0, opt.0))
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match self.0 {
            Ok(ref cmd) => Ok((Some(cmd.to_string()), None)),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

macro_rules! define {
    (impl IntoCommandArg for $T:ty) => {
        impl IntoCommandArg for $T {
            fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
                Ok((Some(self.to_string()), None))
            }
        }
    }
}

define!{ impl IntoCommandArg for bool }
define!{ impl IntoCommandArg for char }
define!{ impl IntoCommandArg for u8 }
define!{ impl IntoCommandArg for u16 }
define!{ impl IntoCommandArg for u32 }
define!{ impl IntoCommandArg for u64 }
define!{ impl IntoCommandArg for usize }
define!{ impl IntoCommandArg for i8 }
define!{ impl IntoCommandArg for i16 }
define!{ impl IntoCommandArg for i32 }
define!{ impl IntoCommandArg for i64 }
define!{ impl IntoCommandArg for isize }
define!{ impl IntoCommandArg for f32 }
define!{ impl IntoCommandArg for f64 }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        match e.to_arg() {
            Ok(arg) => if let Some(arg) = arg.0 {
                RootCommand(Ok(arg))
            } else {
                RootCommand(Ok(String::new()))
            },
            Err(e) => RootCommand(Err(e)),
        }
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn table_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_drop(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

macro_rules! command {
    ($name:ident, $cmd:ident) => {
        pub fn $name<T>(self, arg: T) -> RootCommand
            where T: IntoCommandArg
            {
                let commands = match self.0 {
                    Ok(t) => t,
                    Err(e) => return RootCommand(Err(e)),
                };
                RootCommand(Command::wrap(tt::$cmd,
                                          Some(arg),
                                          Some(&commands)))
            }
    }
}

impl RootCommand {
    command!(table_create, TABLE_CREATE);
    command!(table_drop, TABLE_DROP);
    command!(table, TABLE);
    command!(index_drop, INDEX_DROP);
    command!(index_create, INDEX_CREATE);
    command!(replace, REPLACE);
    command!(update, UPDATE);
    command!(order_by, ORDER_BY);
    command!(without, WITHOUT);
    command!(contains, CONTAINS);
    command!(limit, LIMIT);
    command!(get, GET);
    command!(get_all, GET_ALL);
    command!(filter, FILTER);
    command!(insert, INSERT);

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE,
                                  None as Option<&str>,
                                  Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                input: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
        {
            let mut cmds = format!("[{},", command.value());
            let mut args = String::new();
            if let Some(commands) = commands {
                args.push_str(commands);
                if input.is_some() {
                    args.push_str(",");
                }
            }
            let mut arguments: Option<String> = None;
            let mut options: Option<String> = None;
            if let Some(input) = input {
                let (args, opts) = try!(input.to_arg());
                arguments = args;
                options = opts;
            }
            if let Some(arguments) = arguments {
                args.push_str(&arguments);
            }
            cmds.push_str(&format!("[{}]", args));
            if let Some(options) = options {
                cmds.push_str(&format!(",{}", options));
            }
            cmds.push(']');
            Ok(cmds)
        }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
        -> String {
            let mut qry = format!("[{}", query_type.value());
            if let Some(query) = query {
                qry.push_str(&format!(",{}", query));
            }
            if let Some(options) = options {
                qry.push_str(&format!(",{}", options));
            }
            qry.push_str("]");
            qry
        }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}

fn wrap_arrays(mut val: Value) -> Value {
    if val.is_array() {
        let mut array = Vec::with_capacity(2);
        array.push(Value::I64(tt::MAKE_ARRAY.value() as i64));
        if let Value::Array(vec) = val {
            let mut new_val = Vec::with_capacity(vec.len());
            for v in vec.into_iter() {
                if v.is_array() || v.is_object() {
                    new_val.push(wrap_arrays(v));
                } else {
                    new_val.push(v)
                }
            }
            val = Value::Array(new_val);
        }
        array.push(val);
        val = Value::Array(array);
    } else if val.is_object() {
        if let Value::Object(mut obj) = val {
            for (k, v) in obj.clone() {
                obj.insert(k, wrap_arrays(v));
            }
            val = Value::Object(obj);
        }
    }
    val
}
<|endoftext|>


<commit_before>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();

        println!("num_pieces = {}", num_pieces);
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.as_bytes().iter().cloned());
        message.extend((&[0; 8]).iter().cloned());
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        try!(self.stream.write_all(&message.encode()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        println!("Sending request for {}", piece);
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        println!("Request size: {}", request_size);
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request(usize, usize, usize),
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let offset = convert_big_endian_to_integer(&body[8..12]);
                Message::Request(index, begin, offset)
            },
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn encode(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, begin, amount) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(convert_usize_to_bytes(amount).iter().cloned());
            },
            Message::Piece(index, begin, data) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = convert_usize_to_bytes(payload.len());
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref begin, ref offset) => write!(f, "Request({}, {}, {})", index, begin, offset),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Cleanup.<commit_after>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("Recieved: {:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.as_bytes().iter().cloned());
        message.extend((&[0; 8]).iter().cloned());
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.encode()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            Ok(Message::new(&message[0], &message[1..]))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<usize> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i)
            }
        }
        println!("Done downloading file!");
        None
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        let num_pieces = self.downloaded.len();
        let request_size = if piece == num_pieces - 1 {
            self.metainfo.info.length as usize - (self.metainfo.info.piece_length as usize * (num_pieces - 1))
        } else {
            BLOCK_SIZE
        };
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request(usize, usize, usize),
    Piece(usize, usize, Block),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let offset = convert_big_endian_to_integer(&body[8..12]);
                Message::Request(index, begin, offset)
            },
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn encode(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, begin, amount) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(convert_usize_to_bytes(amount).iter().cloned());
            },
            Message::Piece(index, begin, data) => {
                payload.push(6);
                payload.extend(convert_usize_to_bytes(index).iter().cloned());
                payload.extend(convert_usize_to_bytes(begin).iter().cloned());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = convert_usize_to_bytes(payload.len());
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref begin, ref offset) => write!(f, "Request({}, {}, {})", index, begin, offset),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::collections::hash_map::{self, Entry};
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::IntoIterator;

/// Ordered hashmap built on top of std::collections::HashMap
#[derive(Debug)]
pub struct OrderedHashMap<K, V> {
    #[doc(hidden)]
    base: HashMap<K, V>,
    keys: Vec<K>,
}

impl<K, V> OrderedHashMap<K, V>
where
    K: Eq + Hash + Clone,
{
    /// Creates an empty `OrderedHashMap`.
    ///
    /// The hash map is initially created with a capacity of 0, so it will not allocate until it
    /// is first inserted into.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use ordered_hashmap::OrderedHashMap;
    /// let mut map: OrderedHashMap<&str, i32> = HashMap::new();
    /// ```
    pub fn new() -> OrderedHashMap<K, V> {
        OrderedHashMap { base: HashMap::<K, V>::new(), keys: Vec::<K>::new() }
    }

    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Hash + Eq,
    {
        self.base.get(k)
    }

    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Hash + Eq,
    {
        self.base.get_mut(k)
    }

    pub fn remove(&mut self, k: &K) -> Option<V> {
        match self.keys.iter().position(|x| x == k) {
            Some(index) => {
                self.keys.swap_remove(index);
                self.base.remove(k)
            }
            None => None,
        }
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        if !self.base.contains_key(&k) {
            self.keys.push(k.clone());
        }
        self.base.insert(k, v)
    }

    pub fn iter(&self) -> Iter<'_, K, V> {
        Iter { base: &self.base, keys_iterator: self.keys.iter() }
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        self.base.iter_mut()
    }

    pub fn keys(&self) -> std::slice::Iter<K> {
        self.keys.iter()
    }

    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
        if !self.base.contains_key(&key) {
            self.keys.push(key.clone());
        }
        self.base.entry(key)
    }
}

impl<K, V> Default for OrderedHashMap<K, V>
where
    K: Eq + Hash + Clone,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<'a, K, V> IntoIterator for &'a OrderedHashMap<K, V>
where
    K: Eq + Hash,
{
    type Item = (&'a K, &'a V);
    type IntoIter = Iter<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        Iter { base: &self.base, keys_iterator: self.keys.iter() }
    }
}

impl<K, V> IntoIterator for OrderedHashMap<K, V>
where
    K: Eq + Hash,
{
    type Item = (K, V);
    type IntoIter = IntoIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter { base: self.base, keys_iterator: self.keys.into_iter() }
    }
}

pub struct Iter<'a, K, V> {
    #[doc(hidden)]
    base: &'a HashMap<K, V>,
    keys_iterator: std::slice::Iter<'a, K>,
}

impl<'a, K, V> Iterator for Iter<'a, K, V>
where
    K: Eq + Hash,
{
    type Item = (&'a K, &'a V);
    fn next(&mut self) -> Option<Self::Item> {
        match self.keys_iterator.next() {
            Some(k) => self.base.get_key_value(&k),
            None => None,
        }
    }
}

pub struct IntoIter<K, V> {
    #[doc(hidden)]
    base: HashMap<K, V>,
    keys_iterator: std::vec::IntoIter<K>,
}

impl<K, V> Iterator for IntoIter<K, V>
where
    K: Eq + Hash,
{
    type Item = (K, V);
    fn next(&mut self) -> Option<Self::Item> {
        match self.keys_iterator.next() {
            Some(k) => self.base.remove_entry(&k),
            None => None,
        }
    }
}

pub type IterMut<'a, K, V> = hash_map::IterMut<'a, K, V>;

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

    #[test]
    fn get() {
        let mut map = OrderedHashMap::new();
        map.insert("a", 1);
        assert_eq!(map.get("a"), Some(&1));
        assert_eq!(map.get("b"), None);
    }
}
<commit_msg>more docs for ordered hashmap<commit_after>//! Ordered Hashmap
//!
//! Needs to iterate over items in predictable order
//! e.g. for save ini sections and items in the same order as loaded or added

use std::borrow::Borrow;
use std::collections::hash_map::{self, Entry};
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::IntoIterator;

/// Ordered hashmap built on top of std::collections::HashMap
/// Keys are stored in the field `keys` in the order they were added
#[derive(Debug)]
pub struct OrderedHashMap<K, V> {
    #[doc(hidden)]
    base: HashMap<K, V>,
    keys: Vec<K>,
}

impl<K, V> OrderedHashMap<K, V>
where
    K: Eq + Hash + Clone,
{
    /// Creates an empty `OrderedHashMap`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map: OrderedHashMap<&str, i32> = HashMap::new();
    /// ```
    pub fn new() -> OrderedHashMap<K, V> {
        OrderedHashMap { base: HashMap::<K, V>::new(), keys: Vec::<K>::new() }
    }

    /// Returns a reference to the value corresponding to the key.
    ///
    /// The key may be any borrowed form of the map's key type, but
    /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
    /// the key type.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// map.insert(1, "a");
    /// assert_eq!(map.get(&1), Some(&"a"));
    /// assert_eq!(map.get(&2), None);
    /// ```
    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq,
    {
        self.base.get(k)
    }

    /// Returns a mutable reference to the value corresponding to the key.
    ///
    /// The key may be any borrowed form of the map's key type, but
    /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
    /// the key type.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use ordered_hashmap::OrderedHashMap;
    ///
    /// let mut map = OrderedHashMap::new();
    /// map.insert(1, "a");
    /// if let Some(x) = map.get_mut(&1) {
    ///     *x = "b";
    /// }
    /// assert_eq!(map[&1], "b");
    /// ```
    pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq,
    {
        self.base.get_mut(k)
    }

    /// Inserts a key-value pair into the map.
    ///
    /// If the map did not have this key present, [`None`] is returned.
    ///
    /// If the map did have this key present, the value is updated, and the old
    /// value is returned. The key is not updated, though; this matters for
    /// types that can be `==` without being identical.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// assert_eq!(map.insert(37, "a"), None);
    /// assert_eq!(map.is_empty(), false);
    ///
    /// map.insert(37, "b");
    /// assert_eq!(map.insert(37, "c"), Some("b"));
    /// assert_eq!(map[&37], "c");
    /// ```
    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        if !self.base.contains_key(&k) {
            self.keys.push(k.clone());
        }
        self.base.insert(k, v)
    }

    /// Removes a key from the map, returning the value at the key if the key
    /// was previously in the map.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// map.insert(1, "a");
    /// assert_eq!(map.remove(&1), Some("a"));
    /// assert_eq!(map.remove(&1), None);
    /// ```
    pub fn remove(&mut self, k: &K) -> Option<V> {
        match self.keys.iter().position(|x| x == k) {
            Some(index) => {
                self.keys.swap_remove(index);
                self.base.remove(k)
            }
            None => None,
        }
    }

    /// An iterator visiting all key-value pairs in the order they were added.
    /// The iterator element type is `(&'a K, &'a V)`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// map.insert("a", 1);
    /// map.insert("b", 2);
    /// map.insert("c", 3);
    ///
    /// for (key, val) in map.iter() {
    ///     println!("key: {} val: {}", key, val);
    /// }
    /// ```
    pub fn iter(&self) -> Iter<'_, K, V> {
        Iter { base: &self.base, keys_iterator: self.keys.iter() }
    }

    /// An iterator visiting all key-value pairs in the order they were added,
    /// with mutable references to the values.
    /// The iterator element type is `(&'a K, &'a mut V)`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// map.insert("a", 1);
    /// map.insert("b", 2);
    /// map.insert("c", 3);
    ///
    /// // Update all values
    /// for (_, val) in map.iter_mut() {
    ///     *val *= 2;
    /// }
    ///
    /// for (key, val) in &map {
    ///     println!("key: {} val: {}", key, val);
    /// }
    /// ```
    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        self.base.iter_mut()
    }

    /// An iterator visiting all keys in the order they were added.
    /// The iterator element type is `&'a K`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut map = OrderedHashMap::new();
    /// map.insert("a", 1);
    /// map.insert("b", 2);
    /// map.insert("c", 3);
    ///
    /// for key in map.keys() {
    ///     println!("{}", key);
    /// }
    /// ```
    pub fn keys(&self) -> std::slice::Iter<K> {
        self.keys.iter()
    }

    /// Gets the given key's corresponding entry in the map for in-place manipulation.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut letters = OrderedHashMap::new();
    ///
    /// for ch in "a short treatise on fungi".chars() {
    ///     let counter = letters.entry(ch).or_insert(0);
    ///     *counter += 1;
    /// }
    ///
    /// assert_eq!(letters[&'s'], 2);
    /// assert_eq!(letters[&'t'], 3);
    /// assert_eq!(letters[&'u'], 1);
    /// assert_eq!(letters.get(&'y'), None);
    /// ```
    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
        if !self.base.contains_key(&key) {
            self.keys.push(key.clone());
        }
        self.base.entry(key)
    }
}

impl<K, V> Default for OrderedHashMap<K, V>
where
    K: Eq + Hash + Clone,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<'a, K, V> IntoIterator for &'a OrderedHashMap<K, V>
where
    K: Eq + Hash,
{
    type Item = (&'a K, &'a V);
    type IntoIter = Iter<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        Iter { base: &self.base, keys_iterator: self.keys.iter() }
    }
}

impl<K, V> IntoIterator for OrderedHashMap<K, V>
where
    K: Eq + Hash,
{
    type Item = (K, V);
    type IntoIter = IntoIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter { base: self.base, keys_iterator: self.keys.into_iter() }
    }
}

pub struct Iter<'a, K, V> {
    #[doc(hidden)]
    base: &'a HashMap<K, V>,
    keys_iterator: std::slice::Iter<'a, K>,
}

impl<'a, K, V> Iterator for Iter<'a, K, V>
where
    K: Eq + Hash,
{
    type Item = (&'a K, &'a V);
    fn next(&mut self) -> Option<Self::Item> {
        match self.keys_iterator.next() {
            Some(k) => self.base.get_key_value(&k),
            None => None,
        }
    }
}

pub struct IntoIter<K, V> {
    #[doc(hidden)]
    base: HashMap<K, V>,
    keys_iterator: std::vec::IntoIter<K>,
}

impl<K, V> Iterator for IntoIter<K, V>
where
    K: Eq + Hash,
{
    type Item = (K, V);
    fn next(&mut self) -> Option<Self::Item> {
        match self.keys_iterator.next() {
            Some(k) => self.base.remove_entry(&k),
            None => None,
        }
    }
}

pub type IterMut<'a, K, V> = hash_map::IterMut<'a, K, V>;

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

    #[test]
    fn get() {
        let mut map = OrderedHashMap::new();
        map.insert("a", 1);
        assert_eq!(map.get("a"), Some(&1));
        assert_eq!(map.get("b"), None);
    }
}
<|endoftext|>


<commit_before>const CHUNKSIZE: usize = 128;
const HASHSIZE: usize = 32;

pub trait TreeHash {
    fn tree_hash(&self) -> Vec<u8>;
}

/// Returns a 32 byte hash of 'list' - a vector of byte vectors.
/// Note that this will consume 'list'.
pub fn merkle_hash(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    // flatten list
    let data = &mut list_to_blob(list);

    // get data_len as bytes. It will hashed will the merkle root
    let dlen = data.len() as u64;
    let data_len_bytes = &mut dlen.tree_hash();
    data_len_bytes.resize(32, 0);

    // merklize
    let mut mhash = hash_level(data, CHUNKSIZE);
    while mhash.len() > HASHSIZE {
        mhash = hash_level(&mut mhash, HASHSIZE);
    }

    mhash.append(data_len_bytes);
    mhash.as_slice().tree_hash()
}

/// Takes a flat vector of bytes. It then hashes 'chunk_size * 2' slices into
/// a byte vector of hashes, divisible by HASHSIZE
fn hash_level(data: &mut Vec<u8>, chunk_size: usize) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    for two_chunks in data.chunks(chunk_size * 2) {
        if two_chunks.len() == chunk_size && data.len() > chunk_size {
            // if there is only one chunk here, hash it with a zero-byte
            // CHUNKSIZE vector
            let mut c = two_chunks.to_vec();
            c.append(&mut vec![0; CHUNKSIZE]);
            result.append(&mut c.as_slice().tree_hash());
        } else {
            result.append(&mut two_chunks.tree_hash());
        }
    }

    result
}

fn list_to_blob(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    if list[0].len().is_power_of_two() == false {
        for x in list.iter_mut() {
            extend_to_power_of_2(x);
        }
    }

    let mut data_len = list[0].len() * list.len();

    // do we need padding?
    let extend_by = if data_len % CHUNKSIZE > 0 {
        CHUNKSIZE - (data_len % CHUNKSIZE)
    } else {
        0
    };

    println!("data_len {}, extend_by {}", data_len, extend_by);

    // allocate buffer and append each list element (flatten the vec of vecs)
    data_len += extend_by;
    let mut data: Vec<u8> = Vec::with_capacity(data_len);
    for x in list.iter_mut() {
        data.append(x);
    }

    // add padding
    let mut i = 0;
    while i < extend_by {
        data.push(0);
        i += 1;
    }

    data
}

/// Extends data length to a power of 2 by minimally right-zero-padding
fn extend_to_power_of_2(data: &mut Vec<u8>) {
    let len = data.len();
    let new_len = len.next_power_of_two();
    if new_len > len {
        data.resize(new_len, 0);
    }
}

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

    #[test]
    fn test_merkle_hash() {
        let data1 = vec![1; 100];
        let data2 = vec![2; 100];
        let data3 = vec![3; 100];
        let mut list = vec![data1, data2, data3];
        let result = merkle_hash(&mut list);

        //note: should test againt a known test hash value
        assert_eq!(HASHSIZE, result.len());
        println!("merkle_hash: {:?}", result);
    }

}
<commit_msg>removed a debugging println statement<commit_after>const CHUNKSIZE: usize = 128;
const HASHSIZE: usize = 32;

pub trait TreeHash {
    fn tree_hash(&self) -> Vec<u8>;
}

/// Returns a 32 byte hash of 'list' - a vector of byte vectors.
/// Note that this will consume 'list'.
pub fn merkle_hash(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    // flatten list
    let data = &mut list_to_blob(list);

    // get data_len as bytes. It will hashed will the merkle root
    let dlen = data.len() as u64;
    let data_len_bytes = &mut dlen.tree_hash();
    data_len_bytes.resize(32, 0);

    // merklize
    let mut mhash = hash_level(data, CHUNKSIZE);
    while mhash.len() > HASHSIZE {
        mhash = hash_level(&mut mhash, HASHSIZE);
    }

    mhash.append(data_len_bytes);
    mhash.as_slice().tree_hash()
}

/// Takes a flat vector of bytes. It then hashes 'chunk_size * 2' slices into
/// a byte vector of hashes, divisible by HASHSIZE
fn hash_level(data: &mut Vec<u8>, chunk_size: usize) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    for two_chunks in data.chunks(chunk_size * 2) {
        if two_chunks.len() == chunk_size && data.len() > chunk_size {
            // if there is only one chunk here, hash it with a zero-byte
            // CHUNKSIZE vector
            let mut c = two_chunks.to_vec();
            c.append(&mut vec![0; CHUNKSIZE]);
            result.append(&mut c.as_slice().tree_hash());
        } else {
            result.append(&mut two_chunks.tree_hash());
        }
    }

    result
}

fn list_to_blob(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    if list[0].len().is_power_of_two() == false {
        for x in list.iter_mut() {
            extend_to_power_of_2(x);
        }
    }

    let mut data_len = list[0].len() * list.len();

    // do we need padding?
    let extend_by = if data_len % CHUNKSIZE > 0 {
        CHUNKSIZE - (data_len % CHUNKSIZE)
    } else {
        0
    };

    // allocate buffer and append each list element (flatten the vec of vecs)
    data_len += extend_by;
    let mut data: Vec<u8> = Vec::with_capacity(data_len);
    for x in list.iter_mut() {
        data.append(x);
    }

    // add padding
    let mut i = 0;
    while i < extend_by {
        data.push(0);
        i += 1;
    }

    data
}

/// Extends data length to a power of 2 by minimally right-zero-padding
fn extend_to_power_of_2(data: &mut Vec<u8>) {
    let len = data.len();
    let new_len = len.next_power_of_two();
    if new_len > len {
        data.resize(new_len, 0);
    }
}

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

    #[test]
    fn test_merkle_hash() {
        let data1 = vec![1; 100];
        let data2 = vec![2; 100];
        let data3 = vec![3; 100];
        let mut list = vec![data1, data2, data3];
        let result = merkle_hash(&mut list);

        //note: should test againt a known test hash value
        assert_eq!(HASHSIZE, result.len());
        println!("merkle_hash: {:?}", result);
    }

}
<|endoftext|>