<commit_before>#![crate_name = "wc"]

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

#![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use std::str::from_utf8;
use std::io::{print, stdin, File, BufferedReader};
use StdResult = std::result::Result;
use getopts::Matches;

#[path = "../common/util.rs"]
mod util;

struct Result {
    filename: String,
    bytes: uint,
    chars: uint,
    lines: uint,
    words: uint,
    max_line_length: uint,
}

static NAME: &'static str = "wc";

pub fn uumain(args: Vec<String>) -> int {
    let program = args.get(0).clone();
    let opts = [
        getopts::optflag("c", "bytes", "print the byte counts"),
        getopts::optflag("m", "chars", "print the character counts"),
        getopts::optflag("l", "lines", "print the newline counts"),
        getopts::optflag("L", "max-line-length", "print the length of the longest line"),
        getopts::optflag("w", "words", "print the word counts"),
        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(f) => {
            crash!(1, "Invalid options\n{}", f)
        }
    };

    if matches.opt_present("help") {
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Print newline, word and byte counts for each FILE", opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }

    if matches.opt_present("version") {
        println!("wc 1.0.0");
        return 0;
    }

    let mut files = matches.free.clone();
    if files.is_empty() {
        files = vec!("-".to_string());
    }

    match wc(files, &matches) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e
    }

    0
}

static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;
static SPACE: u8 = ' ' as u8;
static TAB: u8 = '\t' as u8;
static SYN: u8 = 0x16 as u8;
static FF: u8 = 0x0C as u8;

fn is_word_seperator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

pub fn wc(files: Vec<String>, matches: &Matches) -> StdResult<(), int> {
    let mut total_line_count: uint = 0;
    let mut total_word_count: uint = 0;
    let mut total_char_count: uint = 0;
    let mut total_byte_count: uint = 0;
    let mut total_longest_line_length: uint = 0;

    let mut results = vec!();
    let mut max_str_len: uint = 0;

    for path in files.iter() {
        let mut reader = match open(path.to_string()) {
            Ok(f) => f,
            Err(e) => { return Err(e); }
        };

        let mut line_count: uint = 0;
        let mut word_count: uint = 0;
        let mut byte_count: uint = 0;
        let mut char_count: uint = 0;
        let mut current_char_count: uint = 0;
        let mut longest_line_length: uint = 0;

        loop {
            // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
            // hence the option wrapped in a result here
            match reader.read_until(LF) {
                Ok(raw_line) => {
                    // GNU 'wc' only counts lines that end in LF as lines
                    if raw_line.iter().last().unwrap() == &LF {
                        line_count += 1;
                    }

                    byte_count += raw_line.iter().len();

                    // try and convert the bytes to UTF-8 first
                    match from_utf8(raw_line.as_slice()) {
                        Some(line) => {
                            word_count += line.words().count();
                            current_char_count = line.chars().count();
                            char_count += current_char_count;
                        },
                        None => {
                            word_count += raw_line.as_slice().split(|&x| is_word_seperator(x)).count();
                            for byte in raw_line.iter() {
                                match byte.is_ascii() {
                                    true => {
                                        current_char_count += 1;
                                    }
                                    false => { }
                                }
                            }
                            char_count += current_char_count;
                        }
                    }

                    if current_char_count > longest_line_length {
                        // we subtract one here because `line.iter().len()` includes the LF
                        // matches GNU 'wc' behaviour
                        longest_line_length = current_char_count - 1;
                    }
                },
                _ => break
            }

        }

        results.push(Result {
            filename: path.to_string(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_str_len = total_byte_count.to_string().len();
    }

    for result in results.iter() {
        print_stats(result.filename.as_slice(), result.lines, result.words, result.chars, result.bytes, result.max_line_length, matches, max_str_len);
    }

    if files.len() > 1 {
        print_stats("total", total_line_count, total_word_count, total_char_count, total_byte_count, total_longest_line_length, matches, max_str_len);
    }

    Ok(())
}

fn print_stats(filename: &str, line_count: uint, word_count: uint, char_count: uint,
    byte_count: uint, longest_line_length: uint, matches: &Matches, max_str_len: uint) {
    if matches.opt_present("lines") {
        print!("{:1$}", line_count, max_str_len);
    }
    if matches.opt_present("words") {
        print!("{:1$}", word_count, max_str_len);
    }
    if matches.opt_present("bytes") {
        print!("{:1$}", byte_count, max_str_len + 1);
    }
    if matches.opt_present("chars") {
        print!("{:1$}", char_count, max_str_len);
    }
    if matches.opt_present("max-line-length") {
        print!("{:1$}", longest_line_length, max_str_len);
    }

    // defaults
    if !matches.opt_present("bytes")
        && !matches.opt_present("chars")
        && !matches.opt_present("lines")
        && !matches.opt_present("words")
        && !matches.opt_present("max-line-length") {
        print!("{:1$}", line_count, max_str_len);
        print!("{:1$}", word_count, max_str_len + 1);
        print!("{:1$}", byte_count, max_str_len + 1);
    }

    if filename != "-" {
        println!(" {}", filename.as_slice());
    }
    else {
        println!("");
    }
}

fn open(path: String) -> StdResult<BufferedReader<Box<Reader>>, int> {
    if "-" == path.as_slice() {
        let reader = box stdin() as Box<Reader>;
        return Ok(BufferedReader::new(reader));
    }

    match File::open(&std::path::Path::new(path.as_slice())) {
        Ok(fd) => {
            let reader = box fd as Box<Reader>;
            Ok(BufferedReader::new(reader))
        },
        Err(e) => {
            show_error!("wc: {0:s}: {1:s}", path, e.desc.to_string());
            Err(1)
        }
    }
}
<commit_msg>wc: stop buffering stdin twice<commit_after>#![crate_name = "wc"]

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

#![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use std::str::from_utf8;
use std::io::{print, stdin_raw, File, BufferedReader};
use StdResult = std::result::Result;
use getopts::Matches;

#[path = "../common/util.rs"]
mod util;

struct Result {
    filename: String,
    bytes: uint,
    chars: uint,
    lines: uint,
    words: uint,
    max_line_length: uint,
}

static NAME: &'static str = "wc";

pub fn uumain(args: Vec<String>) -> int {
    let program = args.get(0).clone();
    let opts = [
        getopts::optflag("c", "bytes", "print the byte counts"),
        getopts::optflag("m", "chars", "print the character counts"),
        getopts::optflag("l", "lines", "print the newline counts"),
        getopts::optflag("L", "max-line-length", "print the length of the longest line"),
        getopts::optflag("w", "words", "print the word counts"),
        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(f) => {
            crash!(1, "Invalid options\n{}", f)
        }
    };

    if matches.opt_present("help") {
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Print newline, word and byte counts for each FILE", opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }

    if matches.opt_present("version") {
        println!("wc 1.0.0");
        return 0;
    }

    let mut files = matches.free.clone();
    if files.is_empty() {
        files = vec!("-".to_string());
    }

    match wc(files, &matches) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e
    }

    0
}

static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;
static SPACE: u8 = ' ' as u8;
static TAB: u8 = '\t' as u8;
static SYN: u8 = 0x16 as u8;
static FF: u8 = 0x0C as u8;

#[inline(always)]
fn is_word_seperator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

pub fn wc(files: Vec<String>, matches: &Matches) -> StdResult<(), int> {
    let mut total_line_count: uint = 0;
    let mut total_word_count: uint = 0;
    let mut total_char_count: uint = 0;
    let mut total_byte_count: uint = 0;
    let mut total_longest_line_length: uint = 0;

    let mut results = vec!();
    let mut max_str_len: uint = 0;

    for path in files.iter() {
        let mut reader = match open(path.to_string()) {
            Ok(f) => f,
            Err(e) => { return Err(e); }
        };

        let mut line_count: uint = 0;
        let mut word_count: uint = 0;
        let mut byte_count: uint = 0;
        let mut char_count: uint = 0;
        let mut current_char_count: uint = 0;
        let mut longest_line_length: uint = 0;

        loop {
            // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
            // hence the option wrapped in a result here
            match reader.read_until(LF) {
                Ok(raw_line) => {
                    // GNU 'wc' only counts lines that end in LF as lines
                    if raw_line.iter().last().unwrap() == &LF {
                        line_count += 1;
                    }

                    byte_count += raw_line.iter().len();

                    // try and convert the bytes to UTF-8 first
                    match from_utf8(raw_line.as_slice()) {
                        Some(line) => {
                            word_count += line.words().count();
                            current_char_count = line.chars().count();
                            char_count += current_char_count;
                        },
                        None => {
                            word_count += raw_line.as_slice().split(|&x| is_word_seperator(x)).count();
                            for byte in raw_line.iter() {
                                match byte.is_ascii() {
                                    true => {
                                        current_char_count += 1;
                                    }
                                    false => { }
                                }
                            }
                            char_count += current_char_count;
                        }
                    }

                    if current_char_count > longest_line_length {
                        // we subtract one here because `line.iter().len()` includes the LF
                        // matches GNU 'wc' behaviour
                        longest_line_length = current_char_count - 1;
                    }
                },
                _ => break
            }

        }

        results.push(Result {
            filename: path.to_string(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_str_len = total_byte_count.to_string().len();
    }

    for result in results.iter() {
        print_stats(result.filename.as_slice(), result.lines, result.words, result.chars, result.bytes, result.max_line_length, matches, max_str_len);
    }

    if files.len() > 1 {
        print_stats("total", total_line_count, total_word_count, total_char_count, total_byte_count, total_longest_line_length, matches, max_str_len);
    }

    Ok(())
}

fn print_stats(filename: &str, line_count: uint, word_count: uint, char_count: uint,
    byte_count: uint, longest_line_length: uint, matches: &Matches, max_str_len: uint) {
    if matches.opt_present("lines") {
        print!("{:1$}", line_count, max_str_len);
    }
    if matches.opt_present("words") {
        print!("{:1$}", word_count, max_str_len);
    }
    if matches.opt_present("bytes") {
        print!("{:1$}", byte_count, max_str_len + 1);
    }
    if matches.opt_present("chars") {
        print!("{:1$}", char_count, max_str_len);
    }
    if matches.opt_present("max-line-length") {
        print!("{:1$}", longest_line_length, max_str_len);
    }

    // defaults
    if !matches.opt_present("bytes")
        && !matches.opt_present("chars")
        && !matches.opt_present("lines")
        && !matches.opt_present("words")
        && !matches.opt_present("max-line-length") {
        print!("{:1$}", line_count, max_str_len);
        print!("{:1$}", word_count, max_str_len + 1);
        print!("{:1$}", byte_count, max_str_len + 1);
    }

    if filename != "-" {
        println!(" {}", filename.as_slice());
    }
    else {
        println!("");
    }
}

fn open(path: String) -> StdResult<BufferedReader<Box<Reader>>, int> {
    if "-" == path.as_slice() {
        let reader = box stdin_raw() as Box<Reader>;
        return Ok(BufferedReader::new(reader));
    }

    match File::open(&std::path::Path::new(path.as_slice())) {
        Ok(fd) => {
            let reader = box fd as Box<Reader>;
            Ok(BufferedReader::new(reader))
        },
        Err(e) => {
            show_error!("wc: {0:s}: {1:s}", path, e.desc.to_string());
            Err(1)
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// chameneos

use std::option;
use std::os;
use std::strbuf::StrBuf;
use std::task;

fn print_complements() {
    let all = [Blue, Red, Yellow];
    for aa in all.iter() {
        for bb in all.iter() {
            println!("{} + {} -> {}", show_color(*aa), show_color(*bb),
                show_color(transform(*aa, *bb)));
        }
    }
}

enum color { Red, Yellow, Blue }

struct CreatureInfo {
    name: uint,
    color: color
}

fn show_color(cc: color) -> &'static str {
    match cc {
        Red    => "red",
        Yellow => "yellow",
        Blue   => "blue"
    }
}

fn show_color_list(set: Vec<color>) -> StrBuf {
    let mut out = StrBuf::new();
    for col in set.iter() {
        out.push_char(' ');
        out.push_str(show_color(*col));
    }
    out
}

fn show_digit(nn: uint) -> &'static str {
    match nn {
        0 => {"zero"}
        1 => {"one"}
        2 => {"two"}
        3 => {"three"}
        4 => {"four"}
        5 => {"five"}
        6 => {"six"}
        7 => {"seven"}
        8 => {"eight"}
        9 => {"nine"}
        _ => {fail!("expected digits from 0 to 9...")}
    }
}

fn show_number(nn: uint) -> StrBuf {
    let mut out = vec![];
    let mut num = nn;
    let mut dig;
    let mut len = 0;
    if num == 0 { out.push(show_digit(0)) };

    while num != 0 {
        dig = num % 10;
        num = num / 10;
        out.push(" ");
        let s = show_digit(dig);
        out.push(s);
        len += 1 + s.len();
    }
    len += 1;
    out.push(" ");

    let mut ret = StrBuf::with_capacity(len);
    for s in out.iter().rev() {
        ret.push_str(*s);
    }
    ret
}

fn transform(aa: color, bb: color) -> color {
    match (aa, bb) {
        (Red,    Red   ) => { Red    }
        (Red,    Yellow) => { Blue   }
        (Red,    Blue  ) => { Yellow }
        (Yellow, Red   ) => { Blue   }
        (Yellow, Yellow) => { Yellow }
        (Yellow, Blue  ) => { Red    }
        (Blue,   Red   ) => { Yellow }
        (Blue,   Yellow) => { Red    }
        (Blue,   Blue  ) => { Blue   }
    }
}

fn creature(
    name: uint,
    color: color,
    from_rendezvous: Receiver<Option<CreatureInfo>>,
    to_rendezvous: Sender<CreatureInfo>,
    to_rendezvous_log: Sender<~str>
) {
    let mut color = color;
    let mut creatures_met = 0;
    let mut evil_clones_met = 0;

    loop {
        // ask for a pairing
        to_rendezvous.send(CreatureInfo {name: name, color: color});
        let resp = from_rendezvous.recv();

        // log and change, or print and quit
        match resp {
            option::Some(other_creature) => {
                color = transform(color, other_creature.color);

                // track some statistics
                creatures_met += 1;
                if other_creature.name == name {
                   evil_clones_met += 1;
                }
            }
            option::None => {
                // log creatures met and evil clones of self
                let report = format!("{} {}",
                                     creatures_met, show_number(evil_clones_met).as_slice());
                to_rendezvous_log.send(report);
                break;
            }
        }
    }
}

fn rendezvous(nn: uint, set: Vec<color>) {

    // these ports will allow us to hear from the creatures
    let (to_rendezvous, from_creatures) = channel::<CreatureInfo>();
    let (to_rendezvous_log, from_creatures_log) = channel::<~str>();

    // these channels will be passed to the creatures so they can talk to us

    // these channels will allow us to talk to each creature by 'name'/index
    let mut to_creature: Vec<Sender<Option<CreatureInfo>>> =
        set.iter().enumerate().map(|(ii, col)| {
            // create each creature as a listener with a port, and
            // give us a channel to talk to each
            let ii = ii;
            let col = *col;
            let to_rendezvous = to_rendezvous.clone();
            let to_rendezvous_log = to_rendezvous_log.clone();
            let (to_creature, from_rendezvous) = channel();
            task::spawn(proc() {
                creature(ii,
                         col,
                         from_rendezvous,
                         to_rendezvous.clone(),
                         to_rendezvous_log.clone());
            });
            to_creature
        }).collect();

    let mut creatures_met = 0;

    // set up meetings...
    for _ in range(0, nn) {
        let mut fst_creature: CreatureInfo = from_creatures.recv();
        let mut snd_creature: CreatureInfo = from_creatures.recv();

        creatures_met += 2;

        to_creature.get_mut(fst_creature.name).send(Some(snd_creature));
        to_creature.get_mut(snd_creature.name).send(Some(fst_creature));
    }

    // tell each creature to stop
    for to_one in to_creature.iter() {
        to_one.send(None);
    }

    // save each creature's meeting stats
    let mut report = Vec::new();
    for _to_one in to_creature.iter() {
        report.push(from_creatures_log.recv());
    }

    // print each color in the set
    println!("{}", show_color_list(set));

    // print each creature's stats
    for rep in report.iter() {
        println!("{}", *rep);
    }

    // print the total number of creatures met
    println!("{}", show_number(creatures_met));
}

fn main() {
    let args = os::args();
    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_owned(), "200000".to_owned())
    } else if args.len() <= 1u {
        vec!("".to_owned(), "600".to_owned())
    } else {
        args.move_iter().collect()
    };

    let nn = from_str::<uint>(*args.get(1)).unwrap();

    print_complements();
    println!("");

    rendezvous(nn, vec!(Blue, Red, Yellow));
    println!("");

    rendezvous(nn,
        vec!(Blue, Red, Yellow, Red, Yellow, Blue, Red, Yellow, Red, Blue));
}
<commit_msg>fix and improve shootout-chameneos-redux<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(phase)]
#[phase(syntax)] extern crate green;

use std::strbuf::StrBuf;
use std::fmt;

green_start!(main)

fn print_complements() {
    let all = [Blue, Red, Yellow];
    for aa in all.iter() {
        for bb in all.iter() {
            println!("{} + {} -> {}", *aa, *bb, transform(*aa, *bb));
        }
    }
}

enum Color { Red, Yellow, Blue }
impl fmt::Show for Color {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let str = match *self {
            Red => "red",
            Yellow => "yellow",
            Blue => "blue",
        };
        f.buf.write(str.as_bytes())
    }
}

struct CreatureInfo {
    name: uint,
    color: Color
}

fn show_color_list(set: Vec<Color>) -> StrBuf {
    let mut out = StrBuf::new();
    for col in set.iter() {
        out.push_char(' ');
        out.push_str(col.to_str());
    }
    out
}

fn show_digit(nn: uint) -> &'static str {
    match nn {
        0 => {" zero"}
        1 => {" one"}
        2 => {" two"}
        3 => {" three"}
        4 => {" four"}
        5 => {" five"}
        6 => {" six"}
        7 => {" seven"}
        8 => {" eight"}
        9 => {" nine"}
        _ => {fail!("expected digits from 0 to 9...")}
    }
}

struct Number(uint);
impl fmt::Show for Number {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut out = vec![];
        let Number(mut num) = *self;
        if num == 0 { out.push(show_digit(0)) };

        while num != 0 {
            let dig = num % 10;
            num = num / 10;
            let s = show_digit(dig);
            out.push(s);
        }

        for s in out.iter().rev() {
            try!(f.buf.write(s.as_bytes()));
        }
        Ok(())
    }
}

fn transform(aa: Color, bb: Color) -> Color {
    match (aa, bb) {
        (Red,    Red   ) => { Red    }
        (Red,    Yellow) => { Blue   }
        (Red,    Blue  ) => { Yellow }
        (Yellow, Red   ) => { Blue   }
        (Yellow, Yellow) => { Yellow }
        (Yellow, Blue  ) => { Red    }
        (Blue,   Red   ) => { Yellow }
        (Blue,   Yellow) => { Red    }
        (Blue,   Blue  ) => { Blue   }
    }
}

fn creature(
    name: uint,
    mut color: Color,
    from_rendezvous: Receiver<CreatureInfo>,
    to_rendezvous: Sender<CreatureInfo>,
    to_rendezvous_log: Sender<~str>
) {
    let mut creatures_met = 0;
    let mut evil_clones_met = 0;
    let mut rendezvous = from_rendezvous.iter();

    loop {
        // ask for a pairing
        to_rendezvous.send(CreatureInfo {name: name, color: color});

        // log and change, or quit
        match rendezvous.next() {
            Some(other_creature) => {
                color = transform(color, other_creature.color);

                // track some statistics
                creatures_met += 1;
                if other_creature.name == name {
                   evil_clones_met += 1;
                }
            }
            None => break
        }
    }
    // log creatures met and evil clones of self
    let report = format!("{}{}", creatures_met, Number(evil_clones_met));
    to_rendezvous_log.send(report);
}

fn rendezvous(nn: uint, set: Vec<Color>) {
    // these ports will allow us to hear from the creatures
    let (to_rendezvous, from_creatures) = channel::<CreatureInfo>();

    // these channels will be passed to the creatures so they can talk to us
    let (to_rendezvous_log, from_creatures_log) = channel::<~str>();

    // these channels will allow us to talk to each creature by 'name'/index
    let mut to_creature: Vec<Sender<CreatureInfo>> =
        set.iter().enumerate().map(|(ii, &col)| {
            // create each creature as a listener with a port, and
            // give us a channel to talk to each
            let to_rendezvous = to_rendezvous.clone();
            let to_rendezvous_log = to_rendezvous_log.clone();
            let (to_creature, from_rendezvous) = channel();
            spawn(proc() {
                creature(ii,
                         col,
                         from_rendezvous,
                         to_rendezvous,
                         to_rendezvous_log);
            });
            to_creature
        }).collect();

    let mut creatures_met = 0;

    // set up meetings...
    for _ in range(0, nn) {
        let fst_creature = from_creatures.recv();
        let snd_creature = from_creatures.recv();

        creatures_met += 2;

        to_creature.get_mut(fst_creature.name).send(snd_creature);
        to_creature.get_mut(snd_creature.name).send(fst_creature);
    }

    // tell each creature to stop
    drop(to_creature);

    // print each color in the set
    println!("{}", show_color_list(set));

    // print each creature's stats
    drop(to_rendezvous_log);
    for rep in from_creatures_log.iter() {
        println!("{}", rep);
    }

    // print the total number of creatures met
    println!("{}\n", Number(creatures_met));
}

fn main() {
    let nn = if std::os::getenv("RUST_BENCH").is_some() {
        200000
    } else {
        std::os::args().get(1).and_then(|arg| from_str(*arg)).unwrap_or(600)
    };

    print_complements();
    println!("");

    rendezvous(nn, vec!(Blue, Red, Yellow));

    rendezvous(nn,
        vec!(Blue, Red, Yellow, Red, Yellow, Blue, Red, Yellow, Red, Blue));
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, TreeMap};
use std::hash::Hash;

///////////////////////////////////////////////////////////////////////////////

trait Deserialize<S, E> {
    fn deserialize(state: &mut S) -> Result<Self, E>;
}

///////////////////////////////////////////////////////////////////////////////

trait VisitorState<E> {
    fn syntax_error(&mut self) -> E;

    fn visit<
        V: Visitor<T, SeqState, MapState, Self, E>,
        T,
        SeqState,
        MapState,
    >(&mut self, visitor: &mut V) -> Result<T, E>;
}

trait Visitor<
    T,
    SeqState,
    MapState,
    S: VisitorState<E>,
    E,
> {
    fn visit_null(&mut self, state: &mut S) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_int(&mut self, state: &mut S, _v: int) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_string(&mut self, state: &mut S, _v: String) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_seq(&mut self, state: &mut S, _len: uint) -> Result<SeqState, E> {
        Err(state.syntax_error())
    }

    fn visit_seq_elt(&mut self, state: &mut S, _values: &mut SeqState) -> Result<(), E> {
        Err(state.syntax_error())
    }

    fn visit_seq_end(&mut self, state: &mut S, _values: SeqState) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_map(&mut self, state: &mut S, _len: uint) -> Result<MapState, E> {
        Err(state.syntax_error())
    }

    fn visit_map_elt(&mut self, state: &mut S, _values: &mut MapState) -> Result<(), E> {
        Err(state.syntax_error())
    }

    fn visit_map_end(&mut self, state: &mut S, _values: MapState) -> Result<T, E> {
        Err(state.syntax_error())
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for int {
    fn deserialize(state: &mut S) -> Result<int, E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<int, (), (), S, E> for Visitor {
            fn visit_int(&mut self, _state: &mut S, v: int) -> Result<int, E> {
                Ok(v)
            }
        }

        state.visit(&mut Visitor)
    }
}

impl<
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for String {
    fn deserialize(state: &mut S) -> Result<String, E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<String, (), (), S, E> for Visitor {
            fn visit_string(&mut self, _state: &mut S, v: String) -> Result<String, E> {
                Ok(v)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T: Deserialize<S, E>,
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for Vec<T> {
    fn deserialize(state: &mut S) -> Result<Vec<T>, E> {
        struct Visitor;

        impl<
            T: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<Vec<T>, Vec<T>, (), S, E> for Visitor {
            fn visit_seq(&mut self, _state: &mut S, len: uint) -> Result<Vec<T>, E> {
                Ok(Vec::with_capacity(len))
            }

            fn visit_seq_elt(&mut self, state: &mut S, values: &mut Vec<T>) -> Result<(), E> {
                let value = try!(Deserialize::deserialize(state));
                values.push(value);
                Ok(())
            }

            fn visit_seq_end(&mut self, _state: &mut S, values: Vec<T>) -> Result<Vec<T>, E> {
                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    S: VisitorState<E>,
    E
> Deserialize<S, E> for () {
    fn deserialize(state: &mut S) -> Result<(), E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<(), (), (), S, E> for Visitor {
            fn visit_null(&mut self, _state: &mut S) -> Result<(), E> {
                Ok(())
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T0: Deserialize<S, E>,
    T1: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for (T0, T1) {
    fn deserialize(state: &mut S) -> Result<(T0, T1), E> {
        struct Visitor {
            state: uint,
        }

        impl<
            T0: Deserialize<S, E>,
            T1: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<(T0, T1), (Option<T0>, Option<T1>), (), S, E> for Visitor {
            fn visit_seq(&mut self, _state: &mut S, _len: uint) -> Result<(Option<T0>, Option<T1>), E> {
                Ok((None, None))
            }

            fn visit_seq_elt(&mut self, state: &mut S, values: &mut (Option<T0>, Option<T1>)) -> Result<(), E> {
                match self.state {
                    0 => {
                        *values.mut0() = Some(try!(Deserialize::deserialize(state)));
                        self.state += 1;
                        Ok(())
                    }
                    1 => {
                        *values.mut1() = Some(try!(Deserialize::deserialize(state)));
                        self.state += 1;
                        Ok(())
                    }
                    _ => {
                        Err(state.syntax_error())
                    }
                }
            }

            fn visit_seq_end(&mut self, state: &mut S, values: (Option<T0>, Option<T1>)) -> Result<(T0, T1), E> {
                match values {
                    (Some(t0), Some(t1)) => Ok((t0, t1)),
                    _ => Err(state.syntax_error()),
                }
            }
        }

        state.visit(&mut Visitor { state: 0 })
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    K: Deserialize<S, E> + Eq + Hash,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for HashMap<K, V> {
    fn deserialize(state: &mut S) -> Result<HashMap<K, V>, E> {
        struct Visitor;

        impl<
            K: Deserialize<S, E> + Eq + Hash,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<HashMap<K, V>, (), HashMap<K, V>, S, E> for Visitor {
            fn visit_map(&mut self, _state: &mut S, len: uint) -> Result<HashMap<K, V>, E> {
                Ok(HashMap::with_capacity(len))
            }

            fn visit_map_elt(&mut self, state: &mut S, values: &mut HashMap<K, V>) -> Result<(), E> {
                let key = try!(Deserialize::deserialize(state));
                let value = try!(Deserialize::deserialize(state));
                values.insert(key, value);
                Ok(())
            }

            fn visit_map_end(&mut self, _state: &mut S, values: HashMap<K, V>) -> Result<HashMap<K, V>, E> {
                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

impl<
    K: Deserialize<S, E> + Eq + Ord,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for TreeMap<K, V> {
    fn deserialize(state: &mut S) -> Result<TreeMap<K, V>, E> {
        struct Visitor;

        impl<
            K: Deserialize<S, E> + Eq + Ord,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<TreeMap<K, V>, (), TreeMap<K, V>, S, E> for Visitor {
            fn visit_map(&mut self, _state: &mut S, _len: uint) -> Result<TreeMap<K, V>, E> {
                Ok(TreeMap::new())
            }

            fn visit_map_elt(&mut self, state: &mut S, values: &mut TreeMap<K, V>) -> Result<(), E> {
                let key = try!(Deserialize::deserialize(state));
                let value = try!(Deserialize::deserialize(state));
                values.insert(key, value);
                Ok(())
            }

            fn visit_map_end(&mut self, _state: &mut S, values: TreeMap<K, V>) -> Result<TreeMap<K, V>, E> {
                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

mod json {
    use std::collections::TreeMap;

    #[deriving(Show)]
    pub enum Value {
        Null,
        Bool(bool),
        Int(int),
        String(String),
        Vec(Vec<Value>),
        Map(TreeMap<String, Value>),
    }

    impl<
        S: super::VisitorState<E>,
        E,
    > super::Deserialize<S, E> for Value {
        fn deserialize(state: &mut S) -> Result<Value, E> {
            struct Visitor;

            impl<
                S: super::VisitorState<E>,
                E,
            > super::Visitor<Value, Vec<Value>, TreeMap<String, Value>, S, E> for Visitor {
                fn visit_null(&mut self, _state: &mut S) -> Result<Value, E> {
                    Ok(Null)
                }

                fn visit_int(&mut self, _state: &mut S, v: int) -> Result<Value, E> {
                    Ok(Int(v))
                }

                fn visit_string(&mut self, _state: &mut S, v: String) -> Result<Value, E> {
                    Ok(String(v))
                }

                fn visit_seq(&mut self, _state: &mut S, len: uint) -> Result<Vec<Value>, E> {
                    Ok(Vec::with_capacity(len))
                }

                fn visit_seq_elt(&mut self, state: &mut S, values: &mut Vec<Value>) -> Result<(), E> {
                    let value = try!(::Deserialize::deserialize(state));
                    values.push(value);
                    Ok(())
                }

                fn visit_seq_end(&mut self, _state: &mut S, values: Vec<Value>) -> Result<Value, E> {
                    Ok(Vec(values))
                }

                fn visit_map(&mut self, _state: &mut S, _len: uint) -> Result<TreeMap<String, Value>, E> {
                    Ok(TreeMap::new())
                }

                fn visit_map_elt(&mut self, state: &mut S, values: &mut TreeMap<String, Value>) -> Result<(), E> {
                    let key = try!(::Deserialize::deserialize(state));
                    let value = try!(::Deserialize::deserialize(state));
                    values.insert(key, value);
                    Ok(())
                }

                fn visit_map_end(&mut self, _state: &mut S, values: TreeMap<String, Value>) -> Result<Value, E> {
                    Ok(Map(values))
                }
            }

            state.visit(&mut Visitor)
        }
    }

}

///////////////////////////////////////////////////////////////////////////////

trait Deserializer<S, E> {
    fn deserialize<T: Deserialize<S, E>>(&mut self) -> Result<T, E>;
}

///////////////////////////////////////////////////////////////////////////////

enum Token {
    Null,
    Int(int),
    String(String),
    SeqStart(uint),
    MapStart(uint),
    End,
}

///////////////////////////////////////////////////////////////////////////////

struct MyDeserializerState<Iter> {
    tokens: Iter,
    peeked: Option<Token>,
}

impl<Iter: Iterator<Token>> MyDeserializerState<Iter> {
    fn new(tokens: Iter) -> MyDeserializerState<Iter> {
        MyDeserializerState {
            tokens: tokens,
            peeked: None,
        }
    }

    fn next(&mut self) -> Option<Token> {
        match self.peeked.take() {
            Some(token) => { return Some(token); }
            None => { }
        }

        self.tokens.next()
    }

    fn peek<'a>(&'a mut self) -> Option<&'a Token> {
        match self.peeked {
            Some(_) => { }
            None => { self.peeked = self.tokens.next(); }
        }

        self.peeked.as_ref()
    }
}

impl<
    'a,
    Iter: Iterator<Token>,
> VisitorState<
    (),
> for MyDeserializerState<Iter> {
    fn syntax_error(&mut self) -> () {
        ()
    }

    fn visit<
        V: Visitor<T, SeqState, MapState, MyDeserializerState<Iter>, ()>,
        T,
        SeqState,
        MapState,
    >(&mut self, visitor: &mut V) -> Result<T, ()> {
        match self.next() {
            Some(Null) => {
                visitor.visit_null(self)
            }
            Some(Int(v)) => {
                visitor.visit_int(self, v)
            }
            Some(String(v)) => {
                visitor.visit_string(self, v)
            }
            Some(SeqStart(len)) => {
                let mut state = try!(visitor.visit_seq(self, len));

                loop {
                    match self.peek() {
                        Some(&End) => {
                            self.next();
                            break;
                        }
                        Some(_) => {
                            try!(visitor.visit_seq_elt(self, &mut state));
                        }
                        None => {
                            return Err(());
                        }
                    }
                }

                visitor.visit_seq_end(self, state)
            }
            Some(MapStart(len)) => {
                let mut state = try!(visitor.visit_map(self, len));

                loop {
                    match self.peek() {
                        Some(&End) => {
                            self.next();
                            break;
                        }
                        Some(_) => {
                            try!(visitor.visit_map_elt(self, &mut state));
                        }
                        None => {
                            return Err(());
                        }
                    }
                }

                visitor.visit_map_end(self, state)
            }
            Some(End) => {
                Err(())
            }
            None => {
                Err(())
            }
        }

    }
}

///////////////////////////////////////////////////////////////////////////////

fn main() {
    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<Vec<int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<(int, int), ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<json::Value, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<HashMap<String, int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<json::Value, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);
}
<commit_msg>change visit_seq/map into an iterator-ish pattern<commit_after>use std::collections::{HashMap, TreeMap};
use std::hash::Hash;

///////////////////////////////////////////////////////////////////////////////

trait Deserialize<S, E> {
    fn deserialize(state: &mut S) -> Result<Self, E>;
}

///////////////////////////////////////////////////////////////////////////////

trait VisitorState<E> {
    fn syntax_error(&mut self) -> E;

    fn visit<
        V: Visitor<T, Self, E>,
        T: Deserialize<Self, E>,
    >(&mut self, visitor: &mut V) -> Result<T, E>;
}

trait Visitor<
    T,
    S: VisitorState<E>,
    E,
> {
    fn visit_null(&mut self, state: &mut S) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_int(&mut self, state: &mut S, _v: int) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_string(&mut self, state: &mut S, _v: String) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_seq<
        V: SeqVisitor<S, E>,
    >(&mut self, state: &mut S, _visitor: V) -> Result<T, E> {
        Err(state.syntax_error())
    }

    fn visit_map<
        V: MapVisitor<S, E>,
    >(&mut self, state: &mut S, _visitor: V) -> Result<T, E> {
        Err(state.syntax_error())
    }
}

trait SeqVisitor<S, E> {
    fn next<
        T: Deserialize<S, E>,
    >(&mut self, state: &mut S) -> Option<Result<T, E>>;

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        (0, None)
    }
}

trait MapVisitor<S, E> {
    fn next<
        K: Deserialize<S, E>,
        V: Deserialize<S, E>,
    >(&mut self, state: &mut S) -> Option<Result<(K, V), E>>;

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        (0, None)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for int {
    fn deserialize(state: &mut S) -> Result<int, E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<int, S, E> for Visitor {
            fn visit_int(&mut self, _state: &mut S, v: int) -> Result<int, E> {
                Ok(v)
            }
        }

        state.visit(&mut Visitor)
    }
}

impl<
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for String {
    fn deserialize(state: &mut S) -> Result<String, E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<String, S, E> for Visitor {
            fn visit_string(&mut self, _state: &mut S, v: String) -> Result<String, E> {
                Ok(v)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T: Deserialize<S, E>,
    S: VisitorState<E>,
    E,
> Deserialize<S, E> for Vec<T> {
    fn deserialize(state: &mut S) -> Result<Vec<T>, E> {
        struct Visitor;

        impl<
            T: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<Vec<T>, S, E> for Visitor {
            fn visit_seq<
                Visitor: SeqVisitor<S, E>,
            >(&mut self, state: &mut S, mut visitor: Visitor) -> Result<Vec<T>, E> {
                let (len, _) = visitor.size_hint();
                let mut values = Vec::with_capacity(len);

                loop {
                    match visitor.next(state) {
                        Some(Ok(value)) => {
                            values.push(value);
                        }
                        Some(Err(err)) => {
                            return Err(err);
                        }
                        None => {
                            break;
                        }
                    }
                }

                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    S: VisitorState<E>,
    E
> Deserialize<S, E> for () {
    fn deserialize(state: &mut S) -> Result<(), E> {
        struct Visitor;

        impl<
            S: VisitorState<E>,
            E,
        > ::Visitor<(), S, E> for Visitor {
            fn visit_null(&mut self, _state: &mut S) -> Result<(), E> {
                Ok(())
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T0: Deserialize<S, E>,
    T1: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for (T0, T1) {
    fn deserialize(state: &mut S) -> Result<(T0, T1), E> {
        struct Visitor;

        impl<
            T0: Deserialize<S, E>,
            T1: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<(T0, T1), S, E> for Visitor {
            fn visit_seq<
                Visitor: SeqVisitor<S, E>,
            >(&mut self, visitor_state: &mut S, mut visitor: Visitor) -> Result<(T0, T1), E> {
                let mut state = 0u;
                let mut t0 = None;
                let mut t1 = None;

                loop {
                    match state {
                        0 => {
                            t0 = match visitor.next(visitor_state) {
                                Some(Ok(v)) => Some(v),
                                Some(Err(err)) => { return Err(err); }
                                None => { return Err(visitor_state.syntax_error()); }
                            };
                            state += 1;
                        }
                        1 => {
                            t1 = match visitor.next(visitor_state) {
                                Some(Ok(v)) => Some(v),
                                Some(Err(err)) => { return Err(err); }
                                None => { return Err(visitor_state.syntax_error()); }
                            };
                            state += 1;
                        }
                        _ => {
                            match visitor.next(visitor_state) {
                                Some(Ok(())) => { return Err(visitor_state.syntax_error()); }
                                Some(Err(err)) => { return Err(err); }
                                None => { break; }
                            }
                        }
                    }
                }

                match (t0, t1) {
                    (Some(t0), Some(t1)) => Ok((t0, t1)),
                    _ => Err(visitor_state.syntax_error()),
                }
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    K: Deserialize<S, E> + Eq + Hash,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for HashMap<K, V> {
    fn deserialize(state: &mut S) -> Result<HashMap<K, V>, E> {
        struct Visitor;

        impl<
            K: Deserialize<S, E> + Eq + Hash,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<HashMap<K, V>, S, E> for Visitor {
            fn visit_map<
                Visitor: MapVisitor<S, E>,
            >(&mut self, state: &mut S, mut visitor: Visitor) -> Result<HashMap<K, V>, E> {
                let (len, _) = visitor.size_hint();
                let mut values = HashMap::with_capacity(len);

                loop {
                    match visitor.next(state) {
                        Some(Ok((key, value))) => {
                            values.insert(key, value);
                        }
                        Some(Err(err)) => {
                            return Err(err);
                        }
                        None => {
                            break;
                        }
                    }
                }

                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

impl<
    K: Deserialize<S, E> + Eq + Ord,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for TreeMap<K, V> {
    fn deserialize(state: &mut S) -> Result<TreeMap<K, V>, E> {
        struct Visitor;

        impl<
            K: Deserialize<S, E> + Eq + Ord,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E,
        > ::Visitor<TreeMap<K, V>, S, E> for Visitor {
            fn visit_map<
                Visitor: MapVisitor<S, E>,
            >(&mut self, state: &mut S, mut visitor: Visitor) -> Result<TreeMap<K, V>, E> {
                let mut values = TreeMap::new();

                loop {
                    match visitor.next(state) {
                        Some(Ok((key, value))) => {
                            values.insert(key, value);
                        }
                        Some(Err(err)) => {
                            return Err(err);
                        }
                        None => {
                            break;
                        }
                    }
                }

                Ok(values)
            }
        }

        state.visit(&mut Visitor)
    }
}

///////////////////////////////////////////////////////////////////////////////

mod json {
    use std::collections::TreeMap;

    #[deriving(Show)]
    pub enum Value {
        Null,
        Bool(bool),
        Int(int),
        String(String),
        Vec(Vec<Value>),
        Map(TreeMap<String, Value>),
    }

    impl<
        S: super::VisitorState<E>,
        E,
    > super::Deserialize<S, E> for Value {
        fn deserialize(state: &mut S) -> Result<Value, E> {
            struct Visitor;

            impl<
                S: super::VisitorState<E>,
                E,
            > super::Visitor<Value, S, E> for Visitor {
                fn visit_null(&mut self, _state: &mut S) -> Result<Value, E> {
                    Ok(Null)
                }

                fn visit_int(&mut self, _state: &mut S, v: int) -> Result<Value, E> {
                    Ok(Int(v))
                }

                fn visit_string(&mut self, _state: &mut S, v: String) -> Result<Value, E> {
                    Ok(String(v))
                }

                fn visit_seq<
                    Visitor: ::SeqVisitor<S, E>,
                >(&mut self, state: &mut S, mut visitor: Visitor) -> Result<Value, E> {
                    let (len, _) = visitor.size_hint();
                    let mut values = Vec::with_capacity(len);

                    loop {
                        match visitor.next(state) {
                            Some(Ok(value)) => {
                                values.push(value);
                            }
                            Some(Err(err)) => {
                                return Err(err);
                            }
                            None => {
                                break;
                            }
                        }
                    }

                    Ok(Vec(values))
                }

                fn visit_map<
                    Visitor: ::MapVisitor<S, E>,
                >(&mut self, state: &mut S, mut visitor: Visitor) -> Result<Value, E> {
                    let mut values = TreeMap::new();

                    loop {
                        match visitor.next(state) {
                            Some(Ok((key, value))) => {
                                values.insert(key, value);
                            }
                            Some(Err(err)) => {
                                return Err(err);
                            }
                            None => {
                                break;
                            }
                        }
                    }

                    Ok(Map(values))
                }
            }

            state.visit(&mut Visitor)
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

trait Deserializer<S, E> {
    fn deserialize<T: Deserialize<S, E>>(&mut self) -> Result<T, E>;
}

///////////////////////////////////////////////////////////////////////////////

enum Token {
    Null,
    Int(int),
    String(String),
    SeqStart(uint),
    MapStart(uint),
    End,
}

///////////////////////////////////////////////////////////////////////////////

struct MyDeserializerState<Iter> {
    tokens: Iter,
    peeked: Option<Token>,
}

impl<Iter: Iterator<Token>> MyDeserializerState<Iter> {
    fn new(tokens: Iter) -> MyDeserializerState<Iter> {
        MyDeserializerState {
            tokens: tokens,
            peeked: None,
        }
    }

    fn next(&mut self) -> Option<Token> {
        match self.peeked.take() {
            Some(token) => { return Some(token); }
            None => { }
        }

        self.tokens.next()
    }

    fn peek<'a>(&'a mut self) -> Option<&'a Token> {
        match self.peeked {
            Some(_) => { }
            None => { self.peeked = self.tokens.next(); }
        }

        self.peeked.as_ref()
    }
}

impl<
    Iter: Iterator<Token>,
> VisitorState<
    (),
> for MyDeserializerState<Iter> {
    fn syntax_error(&mut self) -> () {
        ()
    }

    fn visit<
        V: Visitor<T, MyDeserializerState<Iter>, ()>,
        T: Deserialize<MyDeserializerState<Iter>, ()>,
    >(&mut self, visitor: &mut V) -> Result<T, ()> {
        match self.next() {
            Some(Null) => {
                visitor.visit_null(self)
            }
            Some(Int(v)) => {
                visitor.visit_int(self, v)
            }
            Some(String(v)) => {
                visitor.visit_string(self, v)
            }
            Some(SeqStart(len)) => {
                visitor.visit_seq(self, MySeqVisitor { len: len })
            }
            Some(MapStart(len)) => {
                visitor.visit_map(self, MyMapVisitor { len: len })
            }
            Some(End) => {
                Err(())
            }
            None => {
                Err(())
            }
        }
    }
}

struct MySeqVisitor {
    len: uint,
}

impl<
    Iter: Iterator<Token>,
> SeqVisitor<MyDeserializerState<Iter>, ()> for MySeqVisitor {
    fn next<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
    >(&mut self, state: &mut MyDeserializerState<Iter>) -> Option<Result<T, ()>> {
        match state.peek() {
            Some(&End) => {
                state.next();
                None
            }
            Some(_) => {
                self.len -= 1;
                Some(Deserialize::deserialize(state))
            }
            None => {
                Some(Err(state.syntax_error()))
            }
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        (self.len, Some(self.len))
    }
}

struct MyMapVisitor {
    len: uint,
}

impl<
    Iter: Iterator<Token>,
> MapVisitor<MyDeserializerState<Iter>, ()> for MyMapVisitor {
    fn next<
        K: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Deserialize<MyDeserializerState<Iter>, ()>,
    >(&mut self, state: &mut MyDeserializerState<Iter>) -> Option<Result<(K, V), ()>> {
        match state.peek() {
            Some(&End) => {
                state.next();
                None
            }
            Some(_) => {
                self.len -= 1;

                let key = match Deserialize::deserialize(state) {
                    Ok(key) => key,
                    Err(err) => { return Some(Err(err)); }
                };

                let value = match Deserialize::deserialize(state) {
                    Ok(value) => value,
                    Err(err) => { return Some(Err(err)); }
                };

                Some(Ok((key, value)))
            }
            None => {
                Some(Err(state.syntax_error()))
            }
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        (self.len, Some(self.len))
    }
}

///////////////////////////////////////////////////////////////////////////////

fn main() {
    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<Vec<int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<(int, int), ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<json::Value, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<HashMap<String, int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    ////

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<json::Value, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);
}
<|endoftext|>


<commit_before>use std::option::Option;

type OptInt32 = Option<i32>;
type OptInt64 = Option<i64>;
type OptString = Option<String>;

pub enum CellType {
    Int32 (OptInt32),
    Int64 (OptInt64),
    Text  (OptString)
}

pub enum CellMetatype {
    tInt32
}

trait Metainf {
    fn name(&self) -> String;
}

impl Metainf for CellType {
    fn name(&self) -> String {
        return match *self {
            CellType::Int32(_) => "INT32".to_string(),
            CellType::Int64(_) => "INT64".to_string(),
            CellType::Text(_) => "TEXT".to_string()
        }
    }
}

trait Metaformat<T> {
    fn construct(type: CellMetatype, name: String) -> T;
}

pub struct CellFormat {
    name: String,
    type: CellMetatype
}

pub struct Cell {
    value: CellType
}

pub impl Metaformat<Cell> for Cell {
    fn construct(type: CellMetatype, name: String) -> Cell {
        return match type {
            CellMetatype::tInt32 => Cell { value: None, name: name }
        }
    }
}

pub struct Row {
    cells: Vec<CellType>
}

pub struct Table {
    rows: Vec<Row>
}

trait ISet {
    fn size(&self) -> usize;
    fn add(&mut self, row: Row);
    fn remove(&mut self, index: usize); // FIXME
}
<commit_msg>Fix build<commit_after>use std::option::Option;

type OptInt32 = Option<i32>;
type OptInt64 = Option<i64>;
type OptString = Option<String>;

pub enum CellType {
    Int32 (OptInt32),
    Int64 (OptInt64),
    Text  (OptString)
}

pub enum CellMetatype {
    tInt32
}

trait Metainf {
    fn name(&self) -> String;
}

impl Metainf for CellType {
    fn name(&self) -> String {
        return match *self {
            CellType::Int32(_) => "INT32".to_string(),
            CellType::Int64(_) => "INT64".to_string(),
            CellType::Text(_) => "TEXT".to_string()
        }
    }
}

trait Metaformat<T> {
    fn construct(type_: CellMetatype, name: String) -> T;
}

pub struct CellFormat {
    name: String,
    type: CellMetatype
}

pub struct Cell {
    value: CellType
}

pub impl Metaformat<Cell> for Cell {
    fn construct(type_: CellMetatype, name: String) -> Cell {
        return match type {
            CellMetatype::tInt32 => Cell { value: CellType::Int32(None), name: name }
        }
    }
}

pub struct Row {
    cells: Vec<CellType>
}

pub struct Table {
    rows: Vec<Row>
}

trait ISet {
    fn size(&self) -> usize;
    fn add(&mut self, row: Row);
    fn remove(&mut self, index: usize); // FIXME
}
<|endoftext|>


<commit_before>//! Reimplements core logic and types from `std::io` in an `alloc`-friendly
//! fashion.

use lib::*;

pub enum ErrorKind {
    InvalidData,
    WriteZero,
    Other,
    UnexpectedEof,
}

impl ErrorKind {
    #[inline]
    fn as_str(&self) -> &'static str {
        match self {
            ErrorKind::InvalidData => "invalid data",
            ErrorKind::WriteZero => "write zero",
            ErrorKind::Other => "other os error",
            ErrorKind::UnexpectedEof => "unexpected end of file",
        }
    }
}

pub struct Error {
    repr: Repr,
}

enum Repr {
    Simple(ErrorKind),
    Custom(ErrorKind, Box<dyn Display + Send + Sync>),
}

impl Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.repr {
            Repr::Custom(_, msg) => write!(fmt, "{}", msg),
            Repr::Simple(kind) => write!(fmt, "{}", kind.as_str()),
        }
    }
}

impl Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(self, fmt)
    }
}

impl serde::de::StdError for Error {}

impl From<ErrorKind> for Error {
    #[inline]
    fn from(kind: ErrorKind) -> Error {
        Error {
            repr: Repr::Simple(kind),
        }
    }
}

impl Error {
    pub(crate) fn new<E>(kind: ErrorKind, error: E) -> Error
    where
        E: Display + Send + Sync + 'static,
    {
        Error {
            repr: Repr::Custom(kind, Box::new(error)),
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

pub trait Write {
    fn write(&mut self, buf: &[u8]) -> Result<usize>;

    fn write_all(&mut self, mut buf: &[u8]) -> Result<()> {
        while !buf.is_empty() {
            match self.write(buf) {
                Ok(0) => {
                    return Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => buf = &buf[n..],
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }

    fn flush(&mut self) -> Result<()>;
}

impl<W: Write> Write for &mut W {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        (*self).write(buf)
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        (*self).write_all(buf)
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        (*self).flush()
    }
}

impl Write for Vec<u8> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.extend_from_slice(buf);
        Ok(buf.len())
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        self.extend_from_slice(buf);
        Ok(())
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

pub trait Read {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize>;
    fn bytes(self) -> Bytes<Self>
    where
        Self: Sized,
    {
        Bytes { inner: self }
    }
}

pub struct Bytes<R> {
    inner: R,
}

impl<R: Read> Iterator for Bytes<R> {
    type Item = Result<u8>;

    fn next(&mut self) -> Option<Result<u8>> {
        let mut byte = 0;
        match self.inner.read(slice::from_mut(&mut byte)) {
            Ok(0) => None,
            Ok(..) => Some(Ok(byte)),
            Err(e) => Some(Err(e)),
        }
    }
}
<commit_msg>Remove unused Simple io error variant<commit_after>//! Reimplements core logic and types from `std::io` in an `alloc`-friendly
//! fashion.

use lib::*;

pub enum ErrorKind {
    InvalidData,
    WriteZero,
    Other,
    UnexpectedEof,
}

pub struct Error {
    repr: Repr,
}

enum Repr {
    Custom(ErrorKind, Box<dyn Display + Send + Sync>),
}

impl Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.repr {
            Repr::Custom(_, msg) => write!(fmt, "{}", msg),
        }
    }
}

impl Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(self, fmt)
    }
}

impl serde::de::StdError for Error {}

impl Error {
    pub(crate) fn new<E>(kind: ErrorKind, error: E) -> Error
    where
        E: Display + Send + Sync + 'static,
    {
        Error {
            repr: Repr::Custom(kind, Box::new(error)),
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

pub trait Write {
    fn write(&mut self, buf: &[u8]) -> Result<usize>;

    fn write_all(&mut self, mut buf: &[u8]) -> Result<()> {
        while !buf.is_empty() {
            match self.write(buf) {
                Ok(0) => {
                    return Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => buf = &buf[n..],
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }

    fn flush(&mut self) -> Result<()>;
}

impl<W: Write> Write for &mut W {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        (*self).write(buf)
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        (*self).write_all(buf)
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        (*self).flush()
    }
}

impl Write for Vec<u8> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.extend_from_slice(buf);
        Ok(buf.len())
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        self.extend_from_slice(buf);
        Ok(())
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

pub trait Read {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize>;
    fn bytes(self) -> Bytes<Self>
    where
        Self: Sized,
    {
        Bytes { inner: self }
    }
}

pub struct Bytes<R> {
    inner: R,
}

impl<R: Read> Iterator for Bytes<R> {
    type Item = Result<u8>;

    fn next(&mut self) -> Option<Result<u8>> {
        let mut byte = 0;
        match self.inner.read(slice::from_mut(&mut byte)) {
            Ok(0) => None,
            Ok(..) => Some(Ok(byte)),
            Err(e) => Some(Err(e)),
        }
    }
}
<|endoftext|>


<commit_before>use amq_protocol::frame::{GenError, Offset, gen_frame, parse_frame};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;

use std::{
  io::{self, Read, Write},
  sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
  },
  thread::{self, Builder as ThreadBuilder, JoinHandle},
  time::{Duration, Instant},
};

use crate::{
  buffer::Buffer,
  connection::Connection,
  connection_status::ConnectionState,
  error::{Error, ErrorKind},
};

const SOCKET:   Token = Token(1);
const DATA:     Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
  handle: Arc<Mutex<Option<JoinHandle<Result<(), Error>>>>>,
}

impl Default for IoLoopHandle {
  fn default() -> Self {
    Self { handle: Arc::new(Mutex::new(None)) }
  }
}

impl IoLoopHandle {
  pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
    *self.handle.lock() = Some(handle);
  }

  pub(crate) fn wait(&self) -> Result<(), Error> {
    if let Some(handle) = self.handle.lock().take() {
      handle.join().expect("io loop")?
    }
    Ok(())
  }
}

#[derive(Debug, PartialEq)]
enum Status {
  Initial,
  Setup,
  Stop,
}

pub(crate) struct IoLoop<T> {
  connection:     Connection,
  socket:         T,
  status:         Status,
  poll:           Poll,
  registration:   Registration,
  set_readiness:  SetReadiness,
  hb_handle:      Option<JoinHandle<()>>,
  frame_size:     usize,
  receive_buffer: Buffer,
  send_buffer:    Buffer,
  can_write:      bool,
  can_read:       bool,
  has_data:       bool,
  send_heartbeat: Arc<AtomicBool>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
  pub(crate) fn new(connection: Connection, socket: T) -> Result<Self, Error> {
    let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
    let (registration, set_readiness) = Registration::new2();
    let inner = Self {
      connection,
      socket,
      status:         Status::Initial,
      poll:           Poll::new().map_err(ErrorKind::IOError)?,
      registration,
      set_readiness,
      hb_handle:      None,
      frame_size,
      receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      send_buffer:    Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      can_write:      false,
      can_read:       false,
      has_data:       false,
      send_heartbeat: Arc::new(AtomicBool::new(false)),
    };
    inner.poll.register(&inner.socket, SOCKET, Ready::readable() | Ready::writable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.connection, DATA, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.registration, CONTINUE, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    Ok(inner)
  }

  fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
    let connection    = self.connection.clone();
    let send_hartbeat = self.send_heartbeat.clone();
    let hb_handle = ThreadBuilder::new().name("heartbeat".to_owned()).spawn(move || {
      while connection.status().connected() {
        let start         = Instant::now();
        let mut remaining = interval;

        loop {
          thread::park_timeout(remaining);
          let elapsed = start.elapsed();
          if elapsed >= remaining {
            break
          }
          remaining -= interval - elapsed;
        }

        send_hartbeat.store(true, Ordering::Relaxed);
      }
    }).map_err(ErrorKind::IOError)?;
    self.hb_handle = Some(hb_handle);
    Ok(())
  }

  fn heartbeat(&mut self) -> Result<(), Error> {
    self.connection.send_heartbeat()?;
    self.send_heartbeat.store(false, Ordering::Relaxed);
    Ok(())
  }

  fn ensure_setup(&mut self) -> Result<(), Error> {
    if self.status != Status::Setup && self.connection.status().connected() {
      let frame_max = self.connection.configuration().frame_max() as usize;
      self.frame_size = std::cmp::max(self.frame_size, frame_max);
      self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
      self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
      let heartbeat = self.connection.configuration().heartbeat();
      if heartbeat != 0 {
        trace!("io_loop: start heartbeat");
        self.start_heartbeat(Duration::from_secs(heartbeat as u64))?;
        trace!("io_loop: heartbeat started");
      }
      self.status = Status::Setup;
    }
    Ok(())
  }

  fn wants_to_write(&self) -> bool {
    self.can_write && self.has_data
  }

  fn wants_to_read(&self) -> bool {
    self.can_read
  }

  fn should_continue(&self) -> bool {
    (self.status == Status::Initial || self.connection.status().connected() || self.connection.status().closing()) && self.status != Status::Stop && !self.connection.status().errored()
  }

  pub(crate) fn run(mut self) -> Result<(), Error> {
    self.connection.clone().set_io_loop(ThreadBuilder::new().name("io_loop".to_owned()).spawn(move || {
      let mut events = Events::with_capacity(1024);
      while self.should_continue() {
        self.do_run(&mut events)?;
      }
      if let Some(hb_handle) = self.hb_handle.take() {
        hb_handle.thread().unpark();
        hb_handle.join().expect("heartbeat loop failed");
      }
      Ok(())
    }).map_err(ErrorKind::IOError)?);
    Ok(())
  }

  fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
    // First, update our internal state
    trace!("io_loop run");
    self.ensure_setup()?;
    trace!("io_loop poll");
    self.poll.poll(events, None).map_err(ErrorKind::IOError)?;
    trace!("io_loop poll done");
    for event in events.iter() {
      match event.token() {
        SOCKET    => {
          if event.readiness().is_readable() {
            self.can_read = true;
          }
          if event.readiness().is_writable() {
            self.can_write = true;
          }
        },
        DATA      => self.has_data = true,
        _         => {},
      }
    }

    // Then, actually run an iteration of the event loop
    trace!("io_loop do_run; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
    loop {
      if self.send_heartbeat.load(Ordering::Relaxed) {
        self.heartbeat()?;
      }
      if self.wants_to_write() && !self.connection.status().blocked() {
        if let Err(e) = self.write_to_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_write = false,
            _ => {
              error!("error writing: {:?}", e);
              if let ConnectionState::SentProtocolHeader(wait_handle, ..) = self.connection.status().state() {
                wait_handle.error(ErrorKind::ConnectionRefused.into());
                self.status = Status::Stop;
              }
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.send_buffer.shift_unless_available(self.frame_size);
      }
      if self.connection.status().closed() {
        self.status = Status::Stop;
      }
      if self.should_continue() && self.wants_to_read() {
        if let Err(e) = self.read_from_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_read = false,
            _ => {
              error!("error reading: {:?}", e);
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.receive_buffer.shift_unless_available(self.frame_size);
      }
      if self.can_parse() {
        self.parse()?;
      }
      if !self.wants_to_read() || !self.wants_to_write() || self.status == Status::Stop || self.connection.status().errored() {
        if self.status != Status::Stop && (self.wants_to_read() || self.can_parse() || self.has_data) {
          trace!("io_loop send continue; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
          self.set_readiness.set_readiness(Ready::readable()).map_err(ErrorKind::IOError)?;
        }
        break;
      }
    }
    trace!("io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}", self.can_read, self.can_write, self.has_data, self.status);
    Ok(())
  }

  fn can_parse(&self) -> bool {
    self.receive_buffer.available_data() > 0
  }

  fn write_to_stream(&mut self) -> Result<(), Error> {
    self.serialize()?;

    self.socket.write(&self.send_buffer.data()).map(|sz| {
      trace!("wrote {} bytes", sz);
      self.send_buffer.consume(sz);
    }).map_err(|e| ErrorKind::IOError(e).into())
  }

  fn read_from_stream(&mut self) -> Result<(), Error> {
    match self.connection.status().state() {
      ConnectionState::Closed => Ok(()),
      ConnectionState::Error  => Err(ErrorKind::InvalidConnectionState(ConnectionState::Error).into()),
      _                       => self.socket.read(&mut self.receive_buffer.space()).map(|sz| {
        trace!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      }).map_err(|e| ErrorKind::IOError(e).into()),
    }
  }

  fn serialize(&mut self) -> Result<(), Error> {
    if let Some((send_id, next_msg)) = self.connection.next_frame() {
      trace!("will write to buffer: {:?}", next_msg);
      match gen_frame(self.send_buffer.space(), &next_msg).map(|tup| tup.0) {
        Ok(sz) => {
          self.send_buffer.fill(sz);
          self.connection.mark_sent(send_id);
          Ok(())
        },
        Err(e) => {
          match e {
            GenError::BufferTooSmall(_) => {
              // Requeue msg
              self.connection.requeue_frame(send_id, next_msg)?;
              self.send_buffer.shift();
              Ok(())
            },
            GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
              error!("error generating frame: {:?}", e);
              self.connection.set_error()?;
              Err(ErrorKind::SerialisationError(e).into())
            }
          }
        }
      }
    } else {
      self.has_data = false;
      Ok(())
    }
  }

  fn parse(&mut self) -> Result<(), Error> {
    match parse_frame(self.receive_buffer.data()) {
      Ok((i, f)) => {
        let consumed = self.receive_buffer.data().offset(i);
        self.receive_buffer.consume(consumed);

        if let Err(e) = self.connection.handle_frame(f) {
          self.connection.set_error()?;
          Err(e)
        } else {
          Ok(())
        }
      },
      Err(e) => {
        if e.is_incomplete() {
          self.receive_buffer.shift();
          Ok(())
        } else {
          error!("parse error: {:?}", e);
          self.connection.set_error()?;
          Err(ErrorKind::ParsingError(format!("{:?}", e)).into())
        }
      }
    }
  }
}
<commit_msg>io_loop: simplification<commit_after>use amq_protocol::frame::{GenError, Offset, gen_frame, parse_frame};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;

use std::{
  io::{self, Read, Write},
  sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
  },
  thread::{self, Builder as ThreadBuilder, JoinHandle},
  time::{Duration, Instant},
};

use crate::{
  buffer::Buffer,
  connection::Connection,
  connection_status::ConnectionState,
  error::{Error, ErrorKind},
};

const SOCKET:   Token = Token(1);
const DATA:     Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
  handle: Arc<Mutex<Option<JoinHandle<Result<(), Error>>>>>,
}

impl Default for IoLoopHandle {
  fn default() -> Self {
    Self { handle: Arc::new(Mutex::new(None)) }
  }
}

impl IoLoopHandle {
  pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
    *self.handle.lock() = Some(handle);
  }

  pub(crate) fn wait(&self) -> Result<(), Error> {
    if let Some(handle) = self.handle.lock().take() {
      handle.join().expect("io loop")?
    }
    Ok(())
  }
}

#[derive(Debug, PartialEq)]
enum Status {
  Initial,
  Setup,
  Stop,
}

pub(crate) struct IoLoop<T> {
  connection:     Connection,
  socket:         T,
  status:         Status,
  poll:           Poll,
  registration:   Registration,
  set_readiness:  SetReadiness,
  hb_handle:      Option<JoinHandle<()>>,
  frame_size:     usize,
  receive_buffer: Buffer,
  send_buffer:    Buffer,
  can_write:      bool,
  can_read:       bool,
  has_data:       bool,
  send_heartbeat: Arc<AtomicBool>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
  pub(crate) fn new(connection: Connection, socket: T) -> Result<Self, Error> {
    let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
    let (registration, set_readiness) = Registration::new2();
    let inner = Self {
      connection,
      socket,
      status:         Status::Initial,
      poll:           Poll::new().map_err(ErrorKind::IOError)?,
      registration,
      set_readiness,
      hb_handle:      None,
      frame_size,
      receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      send_buffer:    Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      can_write:      false,
      can_read:       false,
      has_data:       false,
      send_heartbeat: Arc::new(AtomicBool::new(false)),
    };
    inner.poll.register(&inner.socket, SOCKET, Ready::readable() | Ready::writable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.connection, DATA, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.registration, CONTINUE, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    Ok(inner)
  }

  fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
    let connection    = self.connection.clone();
    let send_hartbeat = self.send_heartbeat.clone();
    let hb_handle = ThreadBuilder::new().name("heartbeat".to_owned()).spawn(move || {
      while connection.status().connected() {
        let start         = Instant::now();
        let mut remaining = interval;

        loop {
          thread::park_timeout(remaining);
          let elapsed = start.elapsed();
          if elapsed >= remaining {
            break
          }
          remaining -= interval - elapsed;
        }

        send_hartbeat.store(true, Ordering::Relaxed);
      }
    }).map_err(ErrorKind::IOError)?;
    self.hb_handle = Some(hb_handle);
    Ok(())
  }

  fn heartbeat(&mut self) -> Result<(), Error> {
    self.connection.send_heartbeat()?;
    self.send_heartbeat.store(false, Ordering::Relaxed);
    Ok(())
  }

  fn ensure_setup(&mut self) -> Result<(), Error> {
    if self.status != Status::Setup && self.connection.status().connected() {
      let frame_max = self.connection.configuration().frame_max() as usize;
      self.frame_size = std::cmp::max(self.frame_size, frame_max);
      self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
      self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
      let heartbeat = self.connection.configuration().heartbeat();
      if heartbeat != 0 {
        trace!("io_loop: start heartbeat");
        self.start_heartbeat(Duration::from_secs(heartbeat as u64))?;
        trace!("io_loop: heartbeat started");
      }
      self.status = Status::Setup;
    }
    Ok(())
  }

  fn wants_to_write(&self) -> bool {
    self.can_write && self.has_data
  }

  fn wants_to_read(&self) -> bool {
    self.can_read
  }

  fn should_continue(&self) -> bool {
    let connection_status = self.connection.status();
    (self.status == Status::Initial || connection_status.connected() || connection_status.closing()) && self.status != Status::Stop && !connection_status.errored()
  }

  pub(crate) fn run(mut self) -> Result<(), Error> {
    self.connection.clone().set_io_loop(ThreadBuilder::new().name("io_loop".to_owned()).spawn(move || {
      let mut events = Events::with_capacity(1024);
      while self.should_continue() {
        self.do_run(&mut events)?;
      }
      if let Some(hb_handle) = self.hb_handle.take() {
        hb_handle.thread().unpark();
        hb_handle.join().expect("heartbeat loop failed");
      }
      Ok(())
    }).map_err(ErrorKind::IOError)?);
    Ok(())
  }

  fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
    // First, update our internal state
    trace!("io_loop run");
    self.ensure_setup()?;
    trace!("io_loop poll");
    self.poll.poll(events, None).map_err(ErrorKind::IOError)?;
    trace!("io_loop poll done");
    for event in events.iter() {
      match event.token() {
        SOCKET    => {
          if event.readiness().is_readable() {
            self.can_read = true;
          }
          if event.readiness().is_writable() {
            self.can_write = true;
          }
        },
        DATA      => self.has_data = true,
        _         => {},
      }
    }

    // Then, actually run an iteration of the event loop
    trace!("io_loop do_run; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
    loop {
      if self.send_heartbeat.load(Ordering::Relaxed) {
        self.heartbeat()?;
      }
      if self.wants_to_write() && !self.connection.status().blocked() {
        if let Err(e) = self.write_to_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_write = false,
            _ => {
              error!("error writing: {:?}", e);
              if let ConnectionState::SentProtocolHeader(wait_handle, ..) = self.connection.status().state() {
                wait_handle.error(ErrorKind::ConnectionRefused.into());
                self.status = Status::Stop;
              }
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.send_buffer.shift_unless_available(self.frame_size);
      }
      if self.connection.status().closed() {
        self.status = Status::Stop;
      }
      if self.should_continue() && self.wants_to_read() {
        if let Err(e) = self.read_from_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_read = false,
            _ => {
              error!("error reading: {:?}", e);
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.receive_buffer.shift_unless_available(self.frame_size);
      }
      if self.can_parse() {
        self.parse()?;
      }
      if !self.wants_to_read() || !self.wants_to_write() || self.status == Status::Stop || self.connection.status().errored() {
        if self.status != Status::Stop && (self.wants_to_read() || self.can_parse() || self.has_data) {
          trace!("io_loop send continue; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
          self.set_readiness.set_readiness(Ready::readable()).map_err(ErrorKind::IOError)?;
        }
        break;
      }
    }
    trace!("io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}", self.can_read, self.can_write, self.has_data, self.status);
    Ok(())
  }

  fn can_parse(&self) -> bool {
    self.receive_buffer.available_data() > 0
  }

  fn write_to_stream(&mut self) -> Result<(), Error> {
    self.serialize()?;

    self.socket.write(&self.send_buffer.data()).map(|sz| {
      trace!("wrote {} bytes", sz);
      self.send_buffer.consume(sz);
    }).map_err(|e| ErrorKind::IOError(e).into())
  }

  fn read_from_stream(&mut self) -> Result<(), Error> {
    match self.connection.status().state() {
      ConnectionState::Closed => Ok(()),
      ConnectionState::Error  => Err(ErrorKind::InvalidConnectionState(ConnectionState::Error).into()),
      _                       => self.socket.read(&mut self.receive_buffer.space()).map(|sz| {
        trace!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      }).map_err(|e| ErrorKind::IOError(e).into()),
    }
  }

  fn serialize(&mut self) -> Result<(), Error> {
    if let Some((send_id, next_msg)) = self.connection.next_frame() {
      trace!("will write to buffer: {:?}", next_msg);
      match gen_frame(self.send_buffer.space(), &next_msg).map(|tup| tup.0) {
        Ok(sz) => {
          self.send_buffer.fill(sz);
          self.connection.mark_sent(send_id);
          Ok(())
        },
        Err(e) => {
          match e {
            GenError::BufferTooSmall(_) => {
              // Requeue msg
              self.connection.requeue_frame(send_id, next_msg)?;
              self.send_buffer.shift();
              Ok(())
            },
            GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
              error!("error generating frame: {:?}", e);
              self.connection.set_error()?;
              Err(ErrorKind::SerialisationError(e).into())
            }
          }
        }
      }
    } else {
      self.has_data = false;
      Ok(())
    }
  }

  fn parse(&mut self) -> Result<(), Error> {
    match parse_frame(self.receive_buffer.data()) {
      Ok((i, f)) => {
        let consumed = self.receive_buffer.data().offset(i);
        self.receive_buffer.consume(consumed);

        if let Err(e) = self.connection.handle_frame(f) {
          self.connection.set_error()?;
          Err(e)
        } else {
          Ok(())
        }
      },
      Err(e) => {
        if e.is_incomplete() {
          self.receive_buffer.shift();
          Ok(())
        } else {
          error!("parse error: {:?}", e);
          self.connection.set_error()?;
          Err(ErrorKind::ParsingError(format!("{:?}", e)).into())
        }
      }
    }
  }
}
<|endoftext|>