<commit_before>// Implements http://rosettacode.org/wiki/Anagrams

extern crate collections;

use collections::HashMap;
use std::str;
use std::io::{File, BufferedReader};
use std::cmp::max;

fn sort_string(string: &str) -> ~str {
	let mut chars: ~[char] = string.chars().collect();
	chars.sort();
	str::from_chars(chars)
}

// Returns groups of anagrams (each group consists of a vector containing the words)
fn get_anagrams<T: Reader>(mut reader: BufferedReader<T>) -> HashMap<~str, Vec<~str>> {
    let mut map: HashMap<~str, Vec<~str>> = HashMap::new();

    // Make groups of words according to the letters they contain
    // i.e. evil, live would be in the same group because they share the same letters
	for line in reader.lines() {
		let s = line.unwrap().trim().to_owned();
		map.insert_or_update_with(sort_string(s.clone()), vec!(s.clone()),
				   |_k, v| v.push(s.clone())
				);
	}

    map
}

// Returns the groups of anagrams that contain the most words in them
fn get_biggest_groups_anagrams(groups: &HashMap<~str, Vec<~str>>) -> HashMap<~str, Vec<~str>> {
    let max_length = groups.iter().fold(0, |s, (_, v)| max(s, v.len()));

    groups.iter().filter(|&(_, v)| v.len() == max_length).map(|(x, y)| (x.clone(), y.clone())).collect()
}

fn main () {
	let path = Path::new("resources/unixdict.txt");
	let reader = BufferedReader::new(File::open(&path));

	let map = get_anagrams(reader);
	let biggest_groups = get_biggest_groups_anagrams(&map);

	for (_, v) in biggest_groups.iter() {
        for s in v.iter() {
            print!("{} ", *s)
        }
        println!("")
	}
}
<commit_msg>Added test to anagram.rs and refactored a bit<commit_after>// Implements http://rosettacode.org/wiki/Anagrams

extern crate collections;

use collections::{HashMap, HashSet};
use std::str;
use std::io::{File, BufferedReader};
use std::cmp::max;

fn sort_string(string: &str) -> ~str {
	let mut chars: ~[char] = string.chars().collect();
	chars.sort();
	str::from_chars(chars)
}

// Returns groups of anagrams (each group consists of a set containing the words)
fn get_anagrams<T: Buffer>(mut reader: T) -> HashMap<~str, HashSet<~str>> {
    let mut groups: HashMap<~str, HashSet<~str>> = HashMap::new();

    // Make groups of words according to the letters they contain
	for line in reader.lines() {
		let s = line.unwrap().trim().to_owned();
		groups.insert_or_update_with(
            sort_string(s.clone()),                                       // The key
            { let mut set = HashSet::new(); set.insert(s.clone()); set }, // The value
			|_, group| { group.insert(s.clone()); }                       // The closure to update the value
		);
	}

    groups
}

// Returns the groups of anagrams that contain the most words in them
fn get_biggest_groups(groups: &HashMap<~str, HashSet<~str>>) -> HashMap<~str, HashSet<~str>> {
    let max_length = groups.iter().fold(0, |current_max, (_, group)| max(current_max, group.len()));
    groups.iter().filter(|&(_, group)| group.len() == max_length).map(|(x, y)| (x.clone(), y.clone())).collect()
}

#[cfg(not(test))]
fn main () {
	let path = Path::new("resources/unixdict.txt");
	let reader = BufferedReader::new(File::open(&path));

	let anagram_groups = get_anagrams(reader);
	let biggest_groups = get_biggest_groups(&anagram_groups);

    // Print the words in the biggest groups of anagrams
	for (_, group) in biggest_groups.iter() {
        for word in group.iter() {
            print!("{} ", *word)
        }
        println!("")
	}
}

#[test]
fn basic_test() {
    // Groups of anagrams
    let group1: HashSet<~str> = vec!("lane".to_owned(), "neal".to_owned(), "lean".to_owned()).move_iter().collect();
    let group2: HashSet<~str> = vec!("angel".to_owned(), "angle".to_owned(), "galen".to_owned()).move_iter().collect();
    let group3: HashSet<~str> = vec!("glare".to_owned(), "large".to_owned()).move_iter().collect();

    // Prepare the input for the program
    // We will get a string like "lane\nneal\nlean\nangel\nangle\ngalen\nglare\nlarge"
    let mut word_iter = group1.iter().chain(group2.iter().chain(group3.iter()));
    let mut words = StrBuf::new();

    words.push_str(word_iter.next().unwrap().as_slice());
    for word in word_iter {
        words.push_str("\n");
        words.push_str(word.as_slice());
    }

    // Here begins the real testing
    let all_groups = get_anagrams(std::io::MemReader::new(words.to_str().bytes().collect()));
    let biggest_groups = get_biggest_groups(&all_groups);

    // Groups 1, 2 and 3 are contained in "all_groups"
    assert!(all_groups.iter().any(|(_, group)| *group == group1));
    assert!(all_groups.iter().any(|(_, group)| *group == group2));
    assert!(all_groups.iter().any(|(_, group)| *group == group3));

    // Groups 1 and 2 are contained in "biggest_groups". Group 3 is not.
    assert!(biggest_groups.iter().any(|(_, group)| *group == group1));
    assert!(biggest_groups.iter().any(|(_, group)| *group == group2));
    assert!(biggest_groups.iter().all(|(_, group)| *group != group3));
}<|endoftext|>


<commit_before>pub fn borrow(num: &str, digit: char) -> Result<String, &'static str> {
    match digit {
        'V' => {
            let mut parts: Vec<&str> = num.rsplitn(2, "X").collect();
            parts.reverse();
            Ok(parts.join("VV"))
        },
        'I' => {
            Ok(num.replace("X", "VV").replace("V", "IIIII"))
        },
        _ => {
            Ok(num.to_string())
        }
    }
}

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

    #[test]
    fn cannot_borrow_i_from_ii() {
        assert_eq!("II", borrow("II", 'I').unwrap());
    }

    #[test]
    fn borrow_i_from_v() {
        assert_eq!("IIIII", borrow("V", 'I').unwrap());
    }

    #[test]
    fn borrow_v_from_x() {
        assert_eq!("VV", borrow("X", 'V').unwrap());
    }

    #[test]
    fn borrow_v_from_cx() {
        assert_eq!("CVV", borrow("CX", 'V').unwrap());
    }

    #[test]
    fn borrow_v_from_xx_only_borrows_once() {
        assert_eq!("XVV", borrow("XX", 'V').unwrap());
    }
}
<commit_msg>Failing transitive borrow test for Rust<commit_after>pub fn borrow(num: &str, digit: char) -> Result<String, &'static str> {
    match digit {
        'V' => {
            let mut parts: Vec<&str> = num.rsplitn(2, "X").collect();
            parts.reverse();
            Ok(parts.join("VV"))
        },
        'I' => {
            Ok(num.replace("X", "VV").replace("V", "IIIII"))
        },
        _ => {
            Ok(num.to_string())
        }
    }
}

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

    #[test]
    fn cannot_borrow_i_from_ii() {
        assert_eq!("II", borrow("II", 'I').unwrap());
    }

    #[test]
    fn borrow_i_from_v() {
        assert_eq!("IIIII", borrow("V", 'I').unwrap());
    }

    #[test]
    fn borrow_v_from_x() {
        assert_eq!("VV", borrow("X", 'V').unwrap());
    }

    #[test]
    fn borrow_v_from_cx() {
        assert_eq!("CVV", borrow("CX", 'V').unwrap());
    }

    #[test]
    fn borrow_i_from_x_borrows_transitively() {
        assert_eq!("VIIIII", borrow("X", 'I').unwrap());
    }

    #[test]
    fn borrow_v_from_xx_only_borrows_once() {
        assert_eq!("XVV", borrow("XX", 'V').unwrap());
    }
}
<|endoftext|>


<commit_before>/* rustdoc: rust -> markdown translator
 * Copyright 2011 Google Inc.
 */

use std;
use rustc;

import option;
import option::{some, none};
import rustc::driver::diagnostic;
import rustc::syntax::ast;
import rustc::syntax::codemap;
import rustc::syntax::parse::parser;
import rustc::syntax::print::pprust;
import rustc::syntax::visit;
import std::io;
import io::writer_util;
import std::map;

type rustdoc = {
    ps: pprust::ps,
    w: io::writer
};

type fndoc = {
    brief: str,
    desc: option::t<str>,
    return: option::t<str>,
    args: map::hashmap<str, str>
};

#[doc(
  brief = "Documents a single function.",
  args(rd = "Rustdoc context",
       ident = "Identifier for this function",
       doc = "Function docs extracted from attributes",
       _fn = "AST object representing this function")
)]
fn doc_fn(rd: rustdoc, ident: str, doc: fndoc, decl: ast::fn_decl) {
    rd.w.write_line("## Function `" + ident + "`");
    rd.w.write_line(doc.brief);
    alt doc.desc {
        some(_d) {
            rd.w.write_line("");
            rd.w.write_line(_d);
            rd.w.write_line("");
        }
        none. { }
    }
    for arg: ast::arg in decl.inputs {
        rd.w.write_str("### Argument `" + arg.ident + "`: ");
        rd.w.write_line("`" + pprust::ty_to_str(arg.ty) + "`");
        alt doc.args.find(arg.ident) {
            some(_d) {
                rd.w.write_line(_d);
            }
            none. { }
        };
    }
    rd.w.write_line("### Returns `" + pprust::ty_to_str(decl.output) + "`");
    alt doc.return {
        some(_r) { rd.w.write_line(_r); }
        none. { }
    }
}

#[doc(
  brief = "Parses function docs from a complex #[doc] attribute.",
  desc = "Supported attributes:

* `brief`: Brief description
* `desc`: Long description
* `return`: Description of return value
* `args`: List of argname = argdesc pairs
",
  args(items = "Doc attribute contents"),
  return = "Parsed function docs."
)]
fn parse_compound_fndoc(items: [@ast::meta_item]) -> fndoc {
    let brief = none;
    let desc = none;
    let return = none;
    let argdocs = map::new_str_hash::<str>();
    let argdocsfound = none;
    for item: @ast::meta_item in items {
        alt item.node {
            ast::meta_name_value("brief", {node: ast::lit_str(value),
                                           span: _}) {
                brief = some(value);
            }
            ast::meta_name_value("desc", {node: ast::lit_str(value),
                                              span: _}) {
                desc = some(value);
            }
            ast::meta_name_value("return", {node: ast::lit_str(value),
                                            span: _}) {
                return = some(value);
            }
            ast::meta_list("args", args) {
                argdocsfound = some(args);
            }
            _ { }
        }
    }

    alt argdocsfound {
        none. { }
        some(ds) {
            for d: @ast::meta_item in ds {
                alt d.node {
                    ast::meta_name_value(key, {node: ast::lit_str(value),
                                               span: _}) {
                        argdocs.insert(key, value);
                    }
                }
            }
        }
    }

    let _brief = alt brief {
        some(_b) { _b }
        none. { "_undocumented_" }
    };

    { brief: _brief, desc: desc, return: return, args: argdocs }
}

#[doc(
  brief = "Documents a single crate item.",
  args(rd = "Rustdoc context",
       item = "AST item to document")
)]
fn doc_item(rd: rustdoc, item: @ast::item) {
    let _fndoc = none;
    let noargdocs = map::new_str_hash::<str>();
    for attr: ast::attribute in item.attrs {
        alt attr.node.value.node {
            ast::meta_name_value(
                "doc", {node: ast::lit_str(value), span: _}) {
                _fndoc = some({ brief: value,
                                desc: none,
                                return: none,
                                args: noargdocs });
            }
            ast::meta_list("doc", docs) {
                _fndoc = some(parse_compound_fndoc(docs));
            }
        }
    }

    let _fndoc0 = alt _fndoc {
        some(_d) { _d }
        none. { { brief: "_undocumented_",
                 desc: none,
                 return: none,
                 args: noargdocs } }
    };

    alt item.node {
        ast::item_const(ty, expr) { }
        ast::item_fn(decl, _, _) {
            doc_fn(rd, item.ident, _fndoc0, decl);
        }
        ast::item_mod(_mod) { }
        ast::item_ty(ty, typarams) { }
        ast::item_tag(variant, typarams) { }
        ast::item_res(_, _, _, _, _) { }
    };
}

#[doc(
  brief = "Generate a crate document header.",
  args(rd = "Rustdoc context",
       name = "Crate name")
)]
fn doc_header(rd: rustdoc, name: str) {
    rd.w.write_line("# Crate " + name);
}

#[doc(
  brief = "Main function.",
  desc = "Command-line arguments:

*  argv[1]: crate file name",
  args(argv = "Command-line arguments.")
)]

fn main(argv: [str]) {

    let w = io::stdout();

    if vec::len(argv) != 2u {
        w.write_str(#fmt("usage: %s <input>\n", argv[0]));
        ret;
    }

    let cm = codemap::new_codemap();
    let sess = @{
        cm: cm,
        mutable next_id: 0,
        diagnostic: diagnostic::mk_handler(cm, none)
    };
    let rd = { ps: pprust::rust_printer(w), w: w };
    doc_header(rd, argv[1]);
    let p = parser::parse_crate_from_source_file(argv[1], [], sess);
    let v = visit::mk_simple_visitor(@{
        visit_item: bind doc_item(rd, _)
        with *visit::default_simple_visitor()});
    visit::visit_crate(*p, (), v);
}
<commit_msg>rustdoc: Move AST parsing into the parse module<commit_after>/* rustdoc: rust -> markdown translator
 * Copyright 2011 Google Inc.
 */

use std;
use rustc;

import option;
import option::{some, none};
import rustc::driver::diagnostic;
import rustc::syntax::ast;
import rustc::syntax::codemap;
import rustc::syntax::parse::parser;
import rustc::syntax::print::pprust;
import rustc::syntax::visit;
import std::io;
import io::writer_util;
import std::map;

type rustdoc = {
    ps: pprust::ps,
    w: io::writer
};

type fndoc = {
    brief: str,
    desc: option::t<str>,
    return: option::t<str>,
    args: map::hashmap<str, str>
};

#[doc(
  brief = "Documents a single function.",
  args(rd = "Rustdoc context",
       ident = "Identifier for this function",
       doc = "Function docs extracted from attributes",
       _fn = "AST object representing this function")
)]
fn doc_fn(rd: rustdoc, ident: str, doc: fndoc, decl: ast::fn_decl) {
    rd.w.write_line("## Function `" + ident + "`");
    rd.w.write_line(doc.brief);
    alt doc.desc {
        some(_d) {
            rd.w.write_line("");
            rd.w.write_line(_d);
            rd.w.write_line("");
        }
        none. { }
    }
    for arg: ast::arg in decl.inputs {
        rd.w.write_str("### Argument `" + arg.ident + "`: ");
        rd.w.write_line("`" + pprust::ty_to_str(arg.ty) + "`");
        alt doc.args.find(arg.ident) {
            some(_d) {
                rd.w.write_line(_d);
            }
            none. { }
        };
    }
    rd.w.write_line("### Returns `" + pprust::ty_to_str(decl.output) + "`");
    alt doc.return {
        some(_r) { rd.w.write_line(_r); }
        none. { }
    }
}

#[doc(
  brief = "Parses function docs from a complex #[doc] attribute.",
  desc = "Supported attributes:

* `brief`: Brief description
* `desc`: Long description
* `return`: Description of return value
* `args`: List of argname = argdesc pairs
",
  args(items = "Doc attribute contents"),
  return = "Parsed function docs."
)]
fn parse_compound_fndoc(items: [@ast::meta_item]) -> fndoc {
    let brief = none;
    let desc = none;
    let return = none;
    let argdocs = map::new_str_hash::<str>();
    let argdocsfound = none;
    for item: @ast::meta_item in items {
        alt item.node {
            ast::meta_name_value("brief", {node: ast::lit_str(value),
                                           span: _}) {
                brief = some(value);
            }
            ast::meta_name_value("desc", {node: ast::lit_str(value),
                                              span: _}) {
                desc = some(value);
            }
            ast::meta_name_value("return", {node: ast::lit_str(value),
                                            span: _}) {
                return = some(value);
            }
            ast::meta_list("args", args) {
                argdocsfound = some(args);
            }
            _ { }
        }
    }

    alt argdocsfound {
        none. { }
        some(ds) {
            for d: @ast::meta_item in ds {
                alt d.node {
                    ast::meta_name_value(key, {node: ast::lit_str(value),
                                               span: _}) {
                        argdocs.insert(key, value);
                    }
                }
            }
        }
    }

    let _brief = alt brief {
        some(_b) { _b }
        none. { "_undocumented_" }
    };

    { brief: _brief, desc: desc, return: return, args: argdocs }
}

#[doc(
  brief = "Documents a single crate item.",
  args(rd = "Rustdoc context",
       item = "AST item to document")
)]
fn doc_item(rd: rustdoc, item: @ast::item) {
    let _fndoc = none;
    let noargdocs = map::new_str_hash::<str>();
    for attr: ast::attribute in item.attrs {
        alt attr.node.value.node {
            ast::meta_name_value(
                "doc", {node: ast::lit_str(value), span: _}) {
                _fndoc = some({ brief: value,
                                desc: none,
                                return: none,
                                args: noargdocs });
            }
            ast::meta_list("doc", docs) {
                _fndoc = some(parse_compound_fndoc(docs));
            }
        }
    }

    let _fndoc0 = alt _fndoc {
        some(_d) { _d }
        none. { { brief: "_undocumented_",
                 desc: none,
                 return: none,
                 args: noargdocs } }
    };

    alt item.node {
        ast::item_const(ty, expr) { }
        ast::item_fn(decl, _, _) {
            doc_fn(rd, item.ident, _fndoc0, decl);
        }
        ast::item_mod(_mod) { }
        ast::item_ty(ty, typarams) { }
        ast::item_tag(variant, typarams) { }
        ast::item_res(_, _, _, _, _) { }
    };
}

#[doc(
  brief = "Generate a crate document header.",
  args(rd = "Rustdoc context",
       name = "Crate name")
)]
fn doc_header(rd: rustdoc, name: str) {
    rd.w.write_line("# Crate " + name);
}

mod parse {
    fn from_file(file: str) -> @ast::crate {
        let cm = codemap::new_codemap();
        let sess = @{
            cm: cm,
            mutable next_id: 0,
            diagnostic: diagnostic::mk_handler(cm, none)
        };
        parser::parse_crate_from_source_file(file, [], sess)
    }
}

#[doc(
  brief = "Main function.",
  desc = "Command-line arguments:

*  argv[1]: crate file name",
  args(argv = "Command-line arguments.")
)]
fn main(argv: [str]) {

    if vec::len(argv) != 2u {
        io::println(#fmt("usage: %s <input>", argv[0]));
        ret;
    }

    let crate = parse::from_file(argv[1]);

    let w = io::stdout();
    let rd = { ps: pprust::rust_printer(w), w: w };
    doc_header(rd, argv[1]);

    let v = visit::mk_simple_visitor(@{
        visit_item: bind doc_item(rd, _)
        with *visit::default_simple_visitor()});
    visit::visit_crate(*crate, (), v);
}
<|endoftext|>


<commit_before>use core::pin::Pin;
use futures_core::future::{TryFuture};
use futures_core::stream::{Stream, TryStream};
use futures_core::task::{Context, Poll};
use futures_sink::Sink;
use pin_utils::{unsafe_pinned, unsafe_unpinned};

/// Stream for the [`try_filter_map`](super::TryStreamExt::try_filter_map)
/// method.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct TryFilterMap<St, Fut, F> {
    stream: St,
    f: F,
    pending: Option<Fut>,
}

impl<St, Fut, F> Unpin for TryFilterMap<St, Fut, F>
    where St: Unpin, Fut: Unpin,
{}

impl<St, Fut, F> TryFilterMap<St, Fut, F> {
    unsafe_pinned!(stream: St);
    unsafe_unpinned!(f: F);
    unsafe_pinned!(pending: Option<Fut>);

    pub(super) fn new(stream: St, f: F) -> Self {
        TryFilterMap { stream, f, pending: None }
    }

    /// Acquires a reference to the underlying stream that this combinator is
    /// pulling from.
    pub fn get_ref(&self) -> &St {
        &self.stream
    }

    /// Acquires a mutable reference to the underlying stream that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_mut(&mut self) -> &mut St {
        &mut self.stream
    }

    /// Acquires a pinned mutable reference to the underlying stream that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut St> {
        self.stream()
    }

    /// Consumes this combinator, returning the underlying stream.
    ///
    /// Note that this may discard intermediate state of this combinator, so
    /// care should be taken to avoid losing resources when this is called.
    pub fn into_inner(self) -> St {
        self.stream
    }
}

impl<St, Fut, F, T> Stream for TryFilterMap<St, Fut, F>
    where St: TryStream,
          Fut: TryFuture<Ok = Option<T>, Error = St::Error>,
          F: FnMut(St::Ok) -> Fut,
{
    type Item = Result<T, St::Error>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<T, St::Error>>> {
        loop {
            if self.pending.is_none() {
                let item = match ready!(self.as_mut().stream().try_poll_next(cx)) {
                    Some(Ok(x)) => x,
                    Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                    None => return Poll::Ready(None),
                };
                let fut = (self.as_mut().f())(item);
                self.as_mut().pending().set(Some(fut));
            }

            let result = ready!(self.as_mut().pending().as_pin_mut().unwrap().try_poll(cx));
            self.as_mut().pending().set(None);
            match result {
                Ok(Some(x)) => return Poll::Ready(Some(Ok(x))),
                Err(e) => return Poll::Ready(Some(Err(e))),
                Ok(None) => {},
            }
        }
    }
}

// Forwarding impl of Sink from the underlying stream
impl<S, Fut, F, T, Item> Sink<Item> for TryFilterMap<S, Fut, F>
    where S: TryStream + Sink<Item>,
          Fut: TryFuture<Ok = Option<T>, Error = S::Error>,
          F: FnMut(S::Ok) -> Fut,
{
    type SinkError = S::SinkError;

    delegate_sink!(stream, Item);
}
<commit_msg>Add FusedStream implementation for TryFilterMap<commit_after>use core::pin::Pin;
use futures_core::future::{TryFuture};
use futures_core::stream::{Stream, TryStream, FusedStream};
use futures_core::task::{Context, Poll};
use futures_sink::Sink;
use pin_utils::{unsafe_pinned, unsafe_unpinned};

/// Stream for the [`try_filter_map`](super::TryStreamExt::try_filter_map)
/// method.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct TryFilterMap<St, Fut, F> {
    stream: St,
    f: F,
    pending: Option<Fut>,
}

impl<St, Fut, F> Unpin for TryFilterMap<St, Fut, F>
    where St: Unpin, Fut: Unpin,
{}

impl<St, Fut, F> TryFilterMap<St, Fut, F> {
    unsafe_pinned!(stream: St);
    unsafe_unpinned!(f: F);
    unsafe_pinned!(pending: Option<Fut>);

    pub(super) fn new(stream: St, f: F) -> Self {
        TryFilterMap { stream, f, pending: None }
    }

    /// Acquires a reference to the underlying stream that this combinator is
    /// pulling from.
    pub fn get_ref(&self) -> &St {
        &self.stream
    }

    /// Acquires a mutable reference to the underlying stream that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_mut(&mut self) -> &mut St {
        &mut self.stream
    }

    /// Acquires a pinned mutable reference to the underlying stream that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut St> {
        self.stream()
    }

    /// Consumes this combinator, returning the underlying stream.
    ///
    /// Note that this may discard intermediate state of this combinator, so
    /// care should be taken to avoid losing resources when this is called.
    pub fn into_inner(self) -> St {
        self.stream
    }
}

impl<St, Fut, F, T> FusedStream for TryFilterMap<St, Fut, F>
    where St: TryStream + FusedStream,
          Fut: TryFuture<Ok = Option<T>, Error = St::Error>,
          F: FnMut(St::Ok) -> Fut,
{
    fn is_terminated(&self) -> bool {
        self.pending.is_none() && self.stream.is_terminated()
    }
}

impl<St, Fut, F, T> Stream for TryFilterMap<St, Fut, F>
    where St: TryStream,
          Fut: TryFuture<Ok = Option<T>, Error = St::Error>,
          F: FnMut(St::Ok) -> Fut,
{
    type Item = Result<T, St::Error>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<T, St::Error>>> {
        loop {
            if self.pending.is_none() {
                let item = match ready!(self.as_mut().stream().try_poll_next(cx)) {
                    Some(Ok(x)) => x,
                    Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                    None => return Poll::Ready(None),
                };
                let fut = (self.as_mut().f())(item);
                self.as_mut().pending().set(Some(fut));
            }

            let result = ready!(self.as_mut().pending().as_pin_mut().unwrap().try_poll(cx));
            self.as_mut().pending().set(None);
            match result {
                Ok(Some(x)) => return Poll::Ready(Some(Ok(x))),
                Err(e) => return Poll::Ready(Some(Err(e))),
                Ok(None) => {},
            }
        }
    }
}

// Forwarding impl of Sink from the underlying stream
impl<S, Fut, F, T, Item> Sink<Item> for TryFilterMap<S, Fut, F>
    where S: TryStream + Sink<Item>,
          Fut: TryFuture<Ok = Option<T>, Error = S::Error>,
          F: FnMut(S::Ok) -> Fut,
{
    type SinkError = S::SinkError;

    delegate_sink!(stream, Item);
}
<|endoftext|>


<commit_before>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::usize;
use std::time::Duration;
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{Builder, JoinHandle};
use std::marker::PhantomData;
use std::boxed::FnBox;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering};
use std::fmt::Write;

pub const DEFAULT_TASKS_PER_TICK: usize = 10000;
const DEFAULT_QUEUE_CAPACITY: usize = 1000;
const DEFAULT_THREAD_COUNT: usize = 1;
const NAP_SECS: u64 = 1;
const QUEUE_MAX_CAPACITY: usize = 8 * DEFAULT_QUEUE_CAPACITY;

pub trait Context: Send {
    fn on_task_started(&mut self) {}
    fn on_task_finished(&mut self) {}
    fn on_tick(&mut self) {}
}

#[derive(Default)]
pub struct DefaultContext;

impl Context for DefaultContext {}

pub trait ContextFactory<Ctx: Context> {
    fn create(&self) -> Ctx;
}

pub struct DefaultContextFactory;

impl<C: Context + Default> ContextFactory<C> for DefaultContextFactory {
    fn create(&self) -> C {
        C::default()
    }
}

pub struct Task<C> {
    task: Box<FnBox(&mut C) + Send>,
}

impl<C: Context> Task<C> {
    fn new<F>(job: F) -> Task<C>
    where
        for<'r> F: FnOnce(&'r mut C) + Send + 'static,
    {
        Task {
            task: Box::new(job),
        }
    }
}

// First in first out queue.
pub struct FifoQueue<C> {
    queue: VecDeque<Task<C>>,
}

impl<C: Context> FifoQueue<C> {
    fn new() -> FifoQueue<C> {
        FifoQueue {
            queue: VecDeque::with_capacity(DEFAULT_QUEUE_CAPACITY),
        }
    }

    fn push(&mut self, task: Task<C>) {
        self.queue.push_back(task);
    }

    fn pop(&mut self) -> Option<Task<C>> {
        let task = self.queue.pop_front();

        if self.queue.is_empty() && self.queue.capacity() > QUEUE_MAX_CAPACITY {
            self.queue = VecDeque::with_capacity(DEFAULT_QUEUE_CAPACITY);
        }

        task
    }
}

pub struct ThreadPoolBuilder<C, F> {
    name: String,
    thread_count: usize,
    tasks_per_tick: usize,
    factory: F,
    _ctx: PhantomData<C>,
}

impl<C: Context + Default + 'static> ThreadPoolBuilder<C, DefaultContextFactory> {
    pub fn with_default_factory(name: String) -> ThreadPoolBuilder<C, DefaultContextFactory> {
        ThreadPoolBuilder::new(name, DefaultContextFactory)
    }
}

impl<C: Context + 'static, F: ContextFactory<C>> ThreadPoolBuilder<C, F> {
    pub fn new(name: String, factory: F) -> ThreadPoolBuilder<C, F> {
        ThreadPoolBuilder {
            name: name,
            thread_count: DEFAULT_THREAD_COUNT,
            tasks_per_tick: DEFAULT_TASKS_PER_TICK,
            factory: factory,
            _ctx: PhantomData,
        }
    }

    pub fn thread_count(mut self, count: usize) -> ThreadPoolBuilder<C, F> {
        self.thread_count = count;
        self
    }

    pub fn tasks_per_tick(mut self, count: usize) -> ThreadPoolBuilder<C, F> {
        self.tasks_per_tick = count;
        self
    }

    pub fn build(self) -> ThreadPool<C> {
        ThreadPool::new(
            self.name,
            self.thread_count,
            self.tasks_per_tick,
            self.factory,
        )
    }
}

/// `ThreadPool` is used to execute tasks in parallel.
/// Each task would be pushed into the pool, and when a thread
/// is ready to process a task, it will get a task from the pool
/// according to the `ScheduleQueue` provided in initialization.
pub struct ThreadPool<Ctx> {
    stop_flag: Arc<AtomicBool>,
    task_pool: Arc<(Mutex<FifoQueue<Ctx>>, Condvar)>,
    threads: Vec<JoinHandle<()>>,
    task_count: Arc<AtomicUsize>,
}

impl<Ctx> ThreadPool<Ctx>
where
    Ctx: Context + 'static,
{
    fn new<C: ContextFactory<Ctx>>(
        name: String,
        num_threads: usize,
        tasks_per_tick: usize,
        f: C,
    ) -> ThreadPool<Ctx> {
        assert!(num_threads >= 1);
        let task_pool = Arc::new((Mutex::new(FifoQueue::new()), Condvar::new()));
        let mut threads = Vec::with_capacity(num_threads);
        let task_count = Arc::new(AtomicUsize::new(0));
        let stop_flag = Arc::new(AtomicBool::new(false));
        // Threadpool threads
        for _ in 0..num_threads {
            let tasks = task_pool.clone();
            let task_num = task_count.clone();
            let ctx = f.create();
            let stop = stop_flag.clone();
            let thread = Builder::new()
                .name(name.clone())
                .spawn(move || {
                    let mut worker = Worker::new(tasks, task_num, tasks_per_tick, stop, ctx);
                    worker.run();
                })
                .unwrap();
            threads.push(thread);
        }

        ThreadPool {
            task_pool: task_pool,
            stop_flag: stop_flag,
            threads: threads,
            task_count: task_count,
        }
    }

    pub fn execute<F>(&self, job: F)
    where
        F: FnOnce(&mut Ctx) + Send + 'static,
        Ctx: Context,
    {
        if self.stop_flag.load(AtomicOrdering::SeqCst) {
            return;
        }
        let task = Task::new(job);
        let &(ref lock, ref cvar) = &*self.task_pool;
        {
            let mut queue = lock.lock().unwrap();
            queue.push(task);
            cvar.notify_one();
        }
        self.task_count.fetch_add(1, AtomicOrdering::SeqCst);
    }

    #[inline]
    pub fn get_task_count(&self) -> usize {
        self.task_count.load(AtomicOrdering::SeqCst)
    }

    pub fn stop(&mut self) -> Result<(), String> {
        self.stop_flag.store(true, AtomicOrdering::SeqCst);
        let &(_, ref cvar) = &*self.task_pool;
        cvar.notify_all();
        let mut err_msg = String::new();
        for t in self.threads.drain(..) {
            if let Err(e) = t.join() {
                write!(&mut err_msg, "Failed to join thread with err: {:?};", e).unwrap();
            }
        }
        if !err_msg.is_empty() {
            return Err(err_msg);
        }
        Ok(())
    }
}

// Each thread has a worker.
struct Worker<C> {
    stop_flag: Arc<AtomicBool>,
    task_queue: Arc<(Mutex<FifoQueue<C>>, Condvar)>,
    task_count: Arc<AtomicUsize>,
    tasks_per_tick: usize,
    task_counter: usize,
    ctx: C,
}

impl<C> Worker<C>
where
    C: Context,
{
    fn new(
        task_queue: Arc<(Mutex<FifoQueue<C>>, Condvar)>,
        task_count: Arc<AtomicUsize>,
        tasks_per_tick: usize,
        stop_flag: Arc<AtomicBool>,
        ctx: C,
    ) -> Worker<C> {
        Worker {
            stop_flag: stop_flag,
            task_queue: task_queue,
            task_count: task_count,
            tasks_per_tick: tasks_per_tick,
            task_counter: 0,
            ctx: ctx,
        }
    }

    fn get_task_timeout(&mut self, timeout: Option<Duration>) -> Option<Task<C>> {
        let &(ref lock, ref cvar) = &*self.task_queue;
        let mut task_queue = lock.lock().unwrap();

        if let Some(task) = task_queue.pop() {
            return Some(task);
        }
        let mut q = match timeout {
            Some(t) => cvar.wait_timeout(task_queue, t).unwrap().0,
            None => cvar.wait(task_queue).unwrap(),
        };
        q.pop()
    }

    fn run(&mut self) {
        let mut timeout = Some(Duration::from_secs(NAP_SECS));
        while !self.stop_flag.load(AtomicOrdering::SeqCst) {
            if let Some(t) = self.get_task_timeout(timeout) {
                self.ctx.on_task_started();
                (t.task).call_once((&mut self.ctx,));
                self.ctx.on_task_finished();
                self.task_count.fetch_sub(1, AtomicOrdering::SeqCst);
                self.task_counter += 1;
                if self.task_counter == self.tasks_per_tick {
                    self.task_counter = 0;
                    self.ctx.on_tick();
                }
                timeout = Some(Duration::from_secs(NAP_SECS));
            } else {
                self.task_counter = 0;
                self.ctx.on_tick();
                timeout = None;
            }
        }
    }
}

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

    use std::time::Duration;
    use std::sync::mpsc::{channel, Sender};
    use std::sync::{Arc, Mutex};
    use std::sync::atomic::{AtomicIsize, Ordering};

    #[test]
    fn test_get_task_count() {
        let name = thd_name!("test_get_task_count");
        let mut task_pool = ThreadPoolBuilder::with_default_factory(name).build();
        let (tx, rx) = channel();
        let (ftx, frx) = channel();
        let receiver = Arc::new(Mutex::new(rx));
        let timeout = Duration::from_secs(2);
        let group_num = 4;
        let mut task_num = 0;
        for gid in 0..group_num {
            let rxer = receiver.clone();
            let ftx = ftx.clone();
            task_pool.execute(move |_: &mut DefaultContext| {
                let rx = rxer.lock().unwrap();
                let id = rx.recv_timeout(timeout).unwrap();
                assert_eq!(id, gid);
                ftx.send(true).unwrap();
            });
            task_num += 1;
            assert_eq!(task_pool.get_task_count(), task_num);
        }

        for gid in 0..group_num {
            tx.send(gid).unwrap();
            frx.recv_timeout(timeout).unwrap();
            let left_num = task_pool.get_task_count();
            // current task may be still running.
            assert!(
                left_num == task_num || left_num == task_num - 1,
                format!("left_num {},task_num {}", left_num, task_num)
            );
            task_num -= 1;
        }
        task_pool.stop().unwrap();
    }

    #[test]
    fn test_task_context() {
        struct TestContext {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        unsafe impl Send for TestContext {}

        impl Context for TestContext {
            fn on_task_started(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
            }
            fn on_task_finished(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
                self.tx.send(()).unwrap();
            }
            fn on_tick(&mut self) {}
        }

        struct TestContextFactory {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        impl ContextFactory<TestContext> for TestContextFactory {
            fn create(&self) -> TestContext {
                TestContext {
                    counter: self.counter.clone(),
                    tx: self.tx.clone(),
                }
            }
        }

        let (tx, rx) = channel();

        let f = TestContextFactory {
            counter: Arc::new(AtomicIsize::new(0)),
            tx: tx,
        };
        let ctx = f.create();
        let name = thd_name!("test_tasks_with_contexts");
        let mut task_pool = ThreadPoolBuilder::new(name, f).thread_count(5).build();

        for _ in 0..10 {
            task_pool.execute(move |_: &mut TestContext| {});
        }
        for _ in 0..10 {
            rx.recv_timeout(Duration::from_millis(20)).unwrap();
        }
        task_pool.stop().unwrap();
        assert_eq!(ctx.counter.load(Ordering::SeqCst), 20);
    }

    #[test]
    fn test_task_tick() {
        struct TestContext {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        unsafe impl Send for TestContext {}

        impl Context for TestContext {
            fn on_task_started(&mut self) {}
            fn on_task_finished(&mut self) {}
            fn on_tick(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
                let _ = self.tx.send(());
            }
        }

        struct TestContextFactory {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        impl ContextFactory<TestContext> for TestContextFactory {
            fn create(&self) -> TestContext {
                TestContext {
                    counter: self.counter.clone(),
                    tx: self.tx.clone(),
                }
            }
        }

        let (tx, rx) = channel();

        let f = TestContextFactory {
            counter: Arc::new(AtomicIsize::new(0)),
            tx: tx,
        };
        let ctx = f.create();
        let name = thd_name!("test_tasks_tick");
        let mut task_pool = ThreadPoolBuilder::new(name, f)
            .thread_count(5)
            .tasks_per_tick(1)
            .build();

        for _ in 0..10 {
            task_pool.execute(move |_: &mut TestContext| {});
        }
        for _ in 0..10 {
            rx.recv_timeout(Duration::from_millis(20)).unwrap();
        }
        task_pool.stop().unwrap();
        // `on_tick` may be called even if there is no task.
        assert!(ctx.counter.load(Ordering::SeqCst) >= 10);
    }
}
<commit_msg>threadpool: fix race (#2266)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::usize;
use std::time::Duration;
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{Builder, JoinHandle};
use std::marker::PhantomData;
use std::boxed::FnBox;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::fmt::Write;

pub const DEFAULT_TASKS_PER_TICK: usize = 10000;
const DEFAULT_QUEUE_CAPACITY: usize = 1000;
const DEFAULT_THREAD_COUNT: usize = 1;
const NAP_SECS: u64 = 1;
const QUEUE_MAX_CAPACITY: usize = 8 * DEFAULT_QUEUE_CAPACITY;

pub trait Context: Send {
    fn on_task_started(&mut self) {}
    fn on_task_finished(&mut self) {}
    fn on_tick(&mut self) {}
}

#[derive(Default)]
pub struct DefaultContext;

impl Context for DefaultContext {}

pub trait ContextFactory<Ctx: Context> {
    fn create(&self) -> Ctx;
}

pub struct DefaultContextFactory;

impl<C: Context + Default> ContextFactory<C> for DefaultContextFactory {
    fn create(&self) -> C {
        C::default()
    }
}

pub struct Task<C> {
    task: Box<FnBox(&mut C) + Send>,
}

impl<C: Context> Task<C> {
    fn new<F>(job: F) -> Task<C>
    where
        for<'r> F: FnOnce(&'r mut C) + Send + 'static,
    {
        Task {
            task: Box::new(job),
        }
    }
}

// First in first out queue.
pub struct FifoQueue<C> {
    queue: VecDeque<Task<C>>,
}

impl<C: Context> FifoQueue<C> {
    fn new() -> FifoQueue<C> {
        FifoQueue {
            queue: VecDeque::with_capacity(DEFAULT_QUEUE_CAPACITY),
        }
    }

    fn push(&mut self, task: Task<C>) {
        self.queue.push_back(task);
    }

    fn pop(&mut self) -> Option<Task<C>> {
        let task = self.queue.pop_front();

        if self.queue.is_empty() && self.queue.capacity() > QUEUE_MAX_CAPACITY {
            self.queue = VecDeque::with_capacity(DEFAULT_QUEUE_CAPACITY);
        }

        task
    }
}

pub struct ThreadPoolBuilder<C, F> {
    name: String,
    thread_count: usize,
    tasks_per_tick: usize,
    factory: F,
    _ctx: PhantomData<C>,
}

impl<C: Context + Default + 'static> ThreadPoolBuilder<C, DefaultContextFactory> {
    pub fn with_default_factory(name: String) -> ThreadPoolBuilder<C, DefaultContextFactory> {
        ThreadPoolBuilder::new(name, DefaultContextFactory)
    }
}

impl<C: Context + 'static, F: ContextFactory<C>> ThreadPoolBuilder<C, F> {
    pub fn new(name: String, factory: F) -> ThreadPoolBuilder<C, F> {
        ThreadPoolBuilder {
            name: name,
            thread_count: DEFAULT_THREAD_COUNT,
            tasks_per_tick: DEFAULT_TASKS_PER_TICK,
            factory: factory,
            _ctx: PhantomData,
        }
    }

    pub fn thread_count(mut self, count: usize) -> ThreadPoolBuilder<C, F> {
        self.thread_count = count;
        self
    }

    pub fn tasks_per_tick(mut self, count: usize) -> ThreadPoolBuilder<C, F> {
        self.tasks_per_tick = count;
        self
    }

    pub fn build(self) -> ThreadPool<C> {
        ThreadPool::new(
            self.name,
            self.thread_count,
            self.tasks_per_tick,
            self.factory,
        )
    }
}

struct ScheduleState<Ctx> {
    queue: FifoQueue<Ctx>,
    stopped: bool,
}

/// `ThreadPool` is used to execute tasks in parallel.
/// Each task would be pushed into the pool, and when a thread
/// is ready to process a task, it will get a task from the pool
/// according to the `ScheduleQueue` provided in initialization.
pub struct ThreadPool<Ctx> {
    state: Arc<(Mutex<ScheduleState<Ctx>>, Condvar)>,
    threads: Vec<JoinHandle<()>>,
    task_count: Arc<AtomicUsize>,
}

impl<Ctx> ThreadPool<Ctx>
where
    Ctx: Context + 'static,
{
    fn new<C: ContextFactory<Ctx>>(
        name: String,
        num_threads: usize,
        tasks_per_tick: usize,
        f: C,
    ) -> ThreadPool<Ctx> {
        assert!(num_threads >= 1);
        let state = ScheduleState {
            queue: FifoQueue::new(),
            stopped: false,
        };
        let state = Arc::new((Mutex::new(state), Condvar::new()));
        let mut threads = Vec::with_capacity(num_threads);
        let task_count = Arc::new(AtomicUsize::new(0));
        // Threadpool threads
        for _ in 0..num_threads {
            let state = state.clone();
            let task_num = task_count.clone();
            let ctx = f.create();
            let thread = Builder::new()
                .name(name.clone())
                .spawn(move || {
                    let mut worker = Worker::new(state, task_num, tasks_per_tick, ctx);
                    worker.run();
                })
                .unwrap();
            threads.push(thread);
        }

        ThreadPool {
            state: state,
            threads: threads,
            task_count: task_count,
        }
    }

    pub fn execute<F>(&self, job: F)
    where
        F: FnOnce(&mut Ctx) + Send + 'static,
        Ctx: Context,
    {
        let task = Task::new(job);
        let &(ref lock, ref cvar) = &*self.state;
        {
            let mut state = lock.lock().unwrap();
            if state.stopped {
                return;
            }
            state.queue.push(task);
            cvar.notify_one();
        }
        self.task_count.fetch_add(1, AtomicOrdering::SeqCst);
    }

    #[inline]
    pub fn get_task_count(&self) -> usize {
        self.task_count.load(AtomicOrdering::SeqCst)
    }

    pub fn stop(&mut self) -> Result<(), String> {
        let &(ref lock, ref cvar) = &*self.state;
        {
            let mut state = lock.lock().unwrap();
            state.stopped = true;
            cvar.notify_all();
        }
        let mut err_msg = String::new();
        for t in self.threads.drain(..) {
            if let Err(e) = t.join() {
                write!(&mut err_msg, "Failed to join thread with err: {:?};", e).unwrap();
            }
        }
        if !err_msg.is_empty() {
            return Err(err_msg);
        }
        Ok(())
    }
}

// Each thread has a worker.
struct Worker<C> {
    state: Arc<(Mutex<ScheduleState<C>>, Condvar)>,
    task_count: Arc<AtomicUsize>,
    tasks_per_tick: usize,
    task_counter: usize,
    ctx: C,
}

impl<C> Worker<C>
where
    C: Context,
{
    fn new(
        state: Arc<(Mutex<ScheduleState<C>>, Condvar)>,
        task_count: Arc<AtomicUsize>,
        tasks_per_tick: usize,
        ctx: C,
    ) -> Worker<C> {
        Worker {
            state: state,
            task_count: task_count,
            tasks_per_tick: tasks_per_tick,
            task_counter: 0,
            ctx: ctx,
        }
    }

    fn next_task(&mut self) -> Option<Task<C>> {
        let &(ref lock, ref cvar) = &*self.state;
        let mut state = lock.lock().unwrap();
        let mut timeout = Some(Duration::from_secs(NAP_SECS));
        loop {
            if state.stopped {
                return None;
            }
            match state.queue.pop() {
                Some(t) => {
                    self.task_counter += 1;
                    return Some(t);
                }
                None => {
                    state = match timeout {
                        Some(t) => cvar.wait_timeout(state, t).unwrap().0,
                        None => {
                            self.task_counter = 0;
                            self.ctx.on_tick();
                            cvar.wait(state).unwrap()
                        }
                    };
                    timeout = None;
                }
            }
        }
    }

    fn run(&mut self) {
        loop {
            let task = match self.next_task() {
                None => return,
                Some(t) => t,
            };

            self.ctx.on_task_started();
            (task.task).call_once((&mut self.ctx,));
            self.ctx.on_task_finished();
            self.task_count.fetch_sub(1, AtomicOrdering::SeqCst);
            if self.task_counter == self.tasks_per_tick {
                self.task_counter = 0;
                self.ctx.on_tick();
            }
        }
    }
}

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

    use std::time::Duration;
    use std::sync::mpsc::{channel, Sender};
    use std::sync::{Arc, Mutex};
    use std::sync::atomic::{AtomicIsize, Ordering};

    #[test]
    fn test_get_task_count() {
        let name = thd_name!("test_get_task_count");
        let mut task_pool = ThreadPoolBuilder::with_default_factory(name).build();
        let (tx, rx) = channel();
        let (ftx, frx) = channel();
        let receiver = Arc::new(Mutex::new(rx));
        let timeout = Duration::from_secs(2);
        let group_num = 4;
        let mut task_num = 0;
        for gid in 0..group_num {
            let rxer = receiver.clone();
            let ftx = ftx.clone();
            task_pool.execute(move |_: &mut DefaultContext| {
                let rx = rxer.lock().unwrap();
                let id = rx.recv_timeout(timeout).unwrap();
                assert_eq!(id, gid);
                ftx.send(true).unwrap();
            });
            task_num += 1;
            assert_eq!(task_pool.get_task_count(), task_num);
        }

        for gid in 0..group_num {
            tx.send(gid).unwrap();
            frx.recv_timeout(timeout).unwrap();
            let left_num = task_pool.get_task_count();
            // current task may be still running.
            assert!(
                left_num == task_num || left_num == task_num - 1,
                format!("left_num {},task_num {}", left_num, task_num)
            );
            task_num -= 1;
        }
        task_pool.stop().unwrap();
    }

    #[test]
    fn test_task_context() {
        struct TestContext {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        unsafe impl Send for TestContext {}

        impl Context for TestContext {
            fn on_task_started(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
            }
            fn on_task_finished(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
                self.tx.send(()).unwrap();
            }
            fn on_tick(&mut self) {}
        }

        struct TestContextFactory {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        impl ContextFactory<TestContext> for TestContextFactory {
            fn create(&self) -> TestContext {
                TestContext {
                    counter: self.counter.clone(),
                    tx: self.tx.clone(),
                }
            }
        }

        let (tx, rx) = channel();

        let f = TestContextFactory {
            counter: Arc::new(AtomicIsize::new(0)),
            tx: tx,
        };
        let ctx = f.create();
        let name = thd_name!("test_tasks_with_contexts");
        let mut task_pool = ThreadPoolBuilder::new(name, f).thread_count(5).build();

        for _ in 0..10 {
            task_pool.execute(move |_: &mut TestContext| {});
        }
        for _ in 0..10 {
            rx.recv_timeout(Duration::from_millis(20)).unwrap();
        }
        task_pool.stop().unwrap();
        assert_eq!(ctx.counter.load(Ordering::SeqCst), 20);
    }

    #[test]
    fn test_task_tick() {
        struct TestContext {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        unsafe impl Send for TestContext {}

        impl Context for TestContext {
            fn on_task_started(&mut self) {}
            fn on_task_finished(&mut self) {}
            fn on_tick(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
                let _ = self.tx.send(());
            }
        }

        struct TestContextFactory {
            counter: Arc<AtomicIsize>,
            tx: Sender<()>,
        }

        impl ContextFactory<TestContext> for TestContextFactory {
            fn create(&self) -> TestContext {
                TestContext {
                    counter: self.counter.clone(),
                    tx: self.tx.clone(),
                }
            }
        }

        let (tx, rx) = channel();

        let f = TestContextFactory {
            counter: Arc::new(AtomicIsize::new(0)),
            tx: tx,
        };
        let ctx = f.create();
        let name = thd_name!("test_tasks_tick");
        let mut task_pool = ThreadPoolBuilder::new(name, f)
            .thread_count(5)
            .tasks_per_tick(1)
            .build();

        for _ in 0..10 {
            task_pool.execute(move |_: &mut TestContext| {});
        }
        for _ in 0..10 {
            rx.recv_timeout(Duration::from_millis(20)).unwrap();
        }
        task_pool.stop().unwrap();
        // `on_tick` may be called even if there is no task.
        assert!(ctx.counter.load(Ordering::SeqCst) >= 10);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate mdbook;
#[macro_use]
extern crate clap;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    try!(book.init());

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        try!(book.copy_theme());
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.build());

    if args.is_present("open") {
        open(book.get_dest().join("index.html"));
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    if args.is_present("open") {
        try!(book.build());
        open(book.get_dest().join("index.html"));
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => {},
        }
        println!("");
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();
    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");
    let interface = args.value_of("interface").unwrap_or("localhost");
    let public_address = args.value_of("address").unwrap_or(interface);
    let open_browser = args.is_present("open");

    let address = format!("{}:{}", interface, port);
    let ws_address = format!("{}:{}", interface, ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#, public_address, ws_port, RELOAD_COMMAND).to_owned());

    try!(book.build());

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| {
        |_| {
            Ok(())
        }
    }).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || {
        ws_server.listen(&*ws_address).unwrap();
    });

    println!("\nServing on {}", address);

    if open_browser {
        open(format!("http://{}", address));
    }

    trigger_on_change(&mut book, move |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
        }
        println!("");
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.test());

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        }
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_src(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
        ::std::process::exit(0);
    };

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if let Err(_) = watcher.watch(book.get_root().join("book.json"), NonRecursive) {
        // do nothing if book.json is not found
    }
    if let Err(_) = watcher.watch(book.get_root().join("book.toml"), NonRecursive) {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => match event {
                NoticeWrite(path) |
                NoticeRemove(path) |
                Create(path) |
                Write(path) |
                Remove(path) |
                Rename(_, path) => {
                    closure(&path, book);
                }
                _ => {}
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<commit_msg>Fix spelling of omitted<commit_after>#[macro_use]
extern crate mdbook;
#[macro_use]
extern crate clap;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    try!(book.init());

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        try!(book.copy_theme());
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.build());

    if args.is_present("open") {
        open(book.get_dest().join("index.html"));
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    if args.is_present("open") {
        try!(book.build());
        open(book.get_dest().join("index.html"));
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => {},
        }
        println!("");
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();
    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");
    let interface = args.value_of("interface").unwrap_or("localhost");
    let public_address = args.value_of("address").unwrap_or(interface);
    let open_browser = args.is_present("open");

    let address = format!("{}:{}", interface, port);
    let ws_address = format!("{}:{}", interface, ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#, public_address, ws_port, RELOAD_COMMAND).to_owned());

    try!(book.build());

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| {
        |_| {
            Ok(())
        }
    }).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || {
        ws_server.listen(&*ws_address).unwrap();
    });

    println!("\nServing on {}", address);

    if open_browser {
        open(format!("http://{}", address));
    }

    trigger_on_change(&mut book, move |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
        }
        println!("");
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.test());

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        }
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_src(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
        ::std::process::exit(0);
    };

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if let Err(_) = watcher.watch(book.get_root().join("book.json"), NonRecursive) {
        // do nothing if book.json is not found
    }
    if let Err(_) = watcher.watch(book.get_root().join("book.toml"), NonRecursive) {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => match event {
                NoticeWrite(path) |
                NoticeRemove(path) |
                Create(path) |
                Write(path) |
                Remove(path) |
                Rename(_, path) => {
                    closure(&path, book);
                }
                _ => {}
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<|endoftext|>


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

#![feature(io, std_misc, core, collections, os)]

#[macro_use]
extern crate log;

extern crate getopts;

extern crate sdl2;

extern crate rboy;

use rboy::device::Device;
use std::time::Duration;
use std::sync::{Arc,RwLock};
use std::sync::mpsc::{Sender,Receiver};
use std::sync::mpsc::TryRecvError::{Disconnected,Empty};
use std::old_io::timer;

static SCALE: usize = 2;
static EXITCODE_INCORRECTOPTIONS: isize = 1;
static EXITCODE_CPULOADFAILS: isize = 2;

#[cfg(not(test))]
fn set_exit_status(exitcode: isize)
{
	std::os::set_exit_status(exitcode);
}

fn main() {
	let args = std::os::args();

	let mut opts = getopts::Options::new();
	opts.optflag("s", "serial", "Output serial to stdout");
	opts.optflag("c", "classic", "Force Classic mode");

	let matches = match opts.parse(args.tail()) {
		Ok(m) => { m }
		Err(f) => { println!("{}", f); return }
	};

	let filename = if !matches.free.is_empty() {
		matches.free[0].clone()
	} else {
		let mut info_start = args[0].clone();
		info_start.push_str(" <filename>");
		println!("{}", opts.usage(info_start.as_slice()));
		set_exit_status(EXITCODE_INCORRECTOPTIONS);
		return;
	};

	sdl2::init(sdl2::INIT_VIDEO);
	let window = match sdl2::video::Window::new("RBoy - A gameboy in Rust",
												sdl2::video::WindowPos::PosUndefined,
												sdl2::video::WindowPos::PosUndefined,
												160*SCALE as i32,
												144*SCALE as i32,
												sdl2::video::WindowFlags::empty()) {
		Ok(window) => window,
		Err(err) => panic!("failed to open window: {}", err),
	};
	let renderer = match sdl2::render::Renderer::from_window(window, sdl2::render::RenderDriverIndex::Auto, sdl2::render::ACCELERATED) {
		Ok(screen) => screen,
		Err(err) => panic!("failed to open screen: {}", err),
	};

	let (sdl_tx, cpu_rx) = std::sync::mpsc::channel();
	let (cpu_tx, sdl_rx) = std::sync::mpsc::channel();
	let rawscreen = ::std::iter::repeat(0u8).take(160*144*3).collect();
	let arc = Arc::new(RwLock::new(rawscreen));
	let arc2 = arc.clone();

	let cpuloop_thread = std::thread::Thread::scoped(move|| cpuloop(&cpu_tx, &cpu_rx, arc2, filename.as_slice(), &matches));

	let mut timer = timer::Timer::new().unwrap();
	let periodic = timer.periodic(Duration::milliseconds(8));

	'main : loop {
		let _ = periodic.recv();
		match sdl_rx.try_recv() {
			Err(Disconnected) => { break 'main },
			Ok(_) => recalculate_screen(&renderer, &arc),
			Err(Empty) => {},
		}
		'event : loop {
			match sdl2::event::poll_event() {
				sdl2::event::Event::Quit(_) => break 'main,
				sdl2::event::Event::None => break 'event,
				sdl2::event::Event::KeyDown(_, _, sdl2::keycode::KeyCode::Escape, _, _, _)
					=> break 'main,
				sdl2::event::Event::KeyDown(_, _, sdl2::keycode::KeyCode::LShift, _, _, _)
					=> { let _ = sdl_tx.send(GBEvent::SpeedUp); },
				sdl2::event::Event::KeyUp(_, _, sdl2::keycode::KeyCode::LShift, _, _, _)
					=> { let _ = sdl_tx.send(GBEvent::SlowDown); },
				sdl2::event::Event::KeyDown(_, _, sdlkey, _, _, _) => {
					match sdl_to_keypad(sdlkey) {
						Some(key) =>  { let _ = sdl_tx.send(GBEvent::KeyDown(key)); },
						None => {},
					}
				},
				sdl2::event::Event::KeyUp(_, _, sdlkey, _, _, _) => {
					match sdl_to_keypad(sdlkey) {
						Some(key) => { let _ = sdl_tx.send(GBEvent::KeyUp(key)); },
						None => {},
					}
				},
				_ => {}
			}
		}
	}

	drop(sdl_tx); // Disconnect such that the cpuloop will exit
	let _ = cpuloop_thread.join();
}

fn sdl_to_keypad(key: sdl2::keycode::KeyCode) -> Option<rboy::KeypadKey> {
	match key {
		sdl2::keycode::KeyCode::Z => Some(rboy::KeypadKey::A),
		sdl2::keycode::KeyCode::X => Some(rboy::KeypadKey::B),
		sdl2::keycode::KeyCode::Up => Some(rboy::KeypadKey::Up),
		sdl2::keycode::KeyCode::Down => Some(rboy::KeypadKey::Down),
		sdl2::keycode::KeyCode::Left => Some(rboy::KeypadKey::Left),
		sdl2::keycode::KeyCode::Right => Some(rboy::KeypadKey::Right),
		sdl2::keycode::KeyCode::Space => Some(rboy::KeypadKey::Select),
		sdl2::keycode::KeyCode::Return => Some(rboy::KeypadKey::Start),
		_ => None,
	}
}

fn recalculate_screen(screen: &sdl2::render::Renderer, arc: &Arc<RwLock<Vec<u8>>>) {
	let mut drawer = screen.drawer();
	drawer.set_draw_color(sdl2::pixels::Color::RGB(0xFF, 0xFF, 0xFF));
	drawer.clear();

	let data =  arc.read().unwrap().clone();
	for y in range(0, 144) {
		for x in range(0, 160) {
			drawer.set_draw_color(sdl2::pixels::Color::RGB(data[y*160*3 + x*3 + 0],
															data[y*160*3 + x*3 + 1],
															data[y*160*3 + x*3 + 2]));
			drawer.draw_rect(&sdl2::rect::Rect::new((x*SCALE) as i32, (y*SCALE) as i32, SCALE as i32, SCALE as i32));
		}
	}
	drawer.present();
}

enum GBEvent {
	KeyUp(rboy::KeypadKey),
	KeyDown(rboy::KeypadKey),
	SpeedUp,
	SlowDown,
}

fn cpuloop(cpu_tx: &Sender<u32>, cpu_rx: &Receiver<GBEvent>, arc: Arc<RwLock<Vec<u8>>>, filename: &str, matches: &getopts::Matches) {
	let opt_c = match matches.opt_present("classic") {
		true => Device::new(filename),
		false => Device::new_cgb(filename),
	};
	let mut c = match opt_c
	{
		Some(cpu) => { cpu },
		None => { error!("Could not get a valid gameboy"); set_exit_status(EXITCODE_CPULOADFAILS); return; },
	};
	c.set_stdout(matches.opt_present("serial"));

	let mut timer = timer::Timer::new().unwrap();
	let periodic = timer.periodic(Duration::milliseconds(8));
	let mut limit_speed = true;

	let waitticks = (4194.304f32 * 4.0) as u32;

	let mut ticks = 0;
	'cpuloop: loop {
		while ticks < waitticks {
			ticks += c.do_cycle();
			if c.check_and_reset_gpu_updated() {
				let mut data = arc.write().unwrap();
				let gpu_data = c.get_gpu_data();
				for i in range(0, data.len()) { data[i] = gpu_data[i]; }
				if cpu_tx.send(0).is_err() { break 'cpuloop };
			}
		}
		ticks -= waitticks;
		if limit_speed { let _ = periodic.recv(); }

		match cpu_rx.try_recv() {
			Ok(event) => match event {
				GBEvent::KeyUp(key) => c.keyup(key),
				GBEvent::KeyDown(key) => c.keydown(key),
				GBEvent::SpeedUp => limit_speed = false,
				GBEvent::SlowDown => limit_speed = true,
			},
			Err(Empty) => {},
			Err(Disconnected) => { break },
		};
	}
}
<commit_msg>Update to new sdl2 event enum<commit_after>#![crate_name = "rboy"]

#![feature(io, std_misc, core, collections, os)]

#[macro_use]
extern crate log;

extern crate getopts;

extern crate sdl2;

extern crate rboy;

use rboy::device::Device;
use std::time::Duration;
use std::sync::{Arc,RwLock};
use std::sync::mpsc::{Sender,Receiver};
use std::sync::mpsc::TryRecvError::{Disconnected,Empty};
use std::old_io::timer;

static SCALE: usize = 2;
static EXITCODE_INCORRECTOPTIONS: isize = 1;
static EXITCODE_CPULOADFAILS: isize = 2;

#[cfg(not(test))]
fn set_exit_status(exitcode: isize)
{
	std::os::set_exit_status(exitcode);
}

fn main() {
	let args = std::os::args();

	let mut opts = getopts::Options::new();
	opts.optflag("s", "serial", "Output serial to stdout");
	opts.optflag("c", "classic", "Force Classic mode");

	let matches = match opts.parse(args.tail()) {
		Ok(m) => { m }
		Err(f) => { println!("{}", f); return }
	};

	let filename = if !matches.free.is_empty() {
		matches.free[0].clone()
	} else {
		let mut info_start = args[0].clone();
		info_start.push_str(" <filename>");
		println!("{}", opts.usage(info_start.as_slice()));
		set_exit_status(EXITCODE_INCORRECTOPTIONS);
		return;
	};

	sdl2::init(sdl2::INIT_VIDEO);
	let window = match sdl2::video::Window::new("RBoy - A gameboy in Rust",
												sdl2::video::WindowPos::PosUndefined,
												sdl2::video::WindowPos::PosUndefined,
												160*SCALE as i32,
												144*SCALE as i32,
												sdl2::video::WindowFlags::empty()) {
		Ok(window) => window,
		Err(err) => panic!("failed to open window: {}", err),
	};
	let renderer = match sdl2::render::Renderer::from_window(window, sdl2::render::RenderDriverIndex::Auto, sdl2::render::ACCELERATED) {
		Ok(screen) => screen,
		Err(err) => panic!("failed to open screen: {}", err),
	};

	let (sdl_tx, cpu_rx) = std::sync::mpsc::channel();
	let (cpu_tx, sdl_rx) = std::sync::mpsc::channel();
	let rawscreen = ::std::iter::repeat(0u8).take(160*144*3).collect();
	let arc = Arc::new(RwLock::new(rawscreen));
	let arc2 = arc.clone();

	let cpuloop_thread = std::thread::Thread::scoped(move|| cpuloop(&cpu_tx, &cpu_rx, arc2, filename.as_slice(), &matches));

	let mut timer = timer::Timer::new().unwrap();
	let periodic = timer.periodic(Duration::milliseconds(8));

	'main : loop {
		let _ = periodic.recv();
		match sdl_rx.try_recv() {
			Err(Disconnected) => { break 'main },
			Ok(_) => recalculate_screen(&renderer, &arc),
			Err(Empty) => {},
		}
		'event : loop {
			match sdl2::event::poll_event() {
				sdl2::event::Event::Quit { .. } => break 'main,
				sdl2::event::Event::None => break 'event,
				sdl2::event::Event::KeyDown { keycode: sdl2::keycode::KeyCode::Escape, .. }
					=> break 'main,
				sdl2::event::Event::KeyDown { keycode: sdl2::keycode::KeyCode::LShift, .. }
					=> { let _ = sdl_tx.send(GBEvent::SpeedUp); },
				sdl2::event::Event::KeyUp { keycode: sdl2::keycode::KeyCode::LShift, .. }
					=> { let _ = sdl_tx.send(GBEvent::SlowDown); },
				sdl2::event::Event::KeyDown { keycode: sdlkey, .. } => {
					match sdl_to_keypad(sdlkey) {
						Some(key) =>  { let _ = sdl_tx.send(GBEvent::KeyDown(key)); },
						None => {},
					}
				},
				sdl2::event::Event::KeyUp { keycode: sdlkey, .. } => {
					match sdl_to_keypad(sdlkey) {
						Some(key) => { let _ = sdl_tx.send(GBEvent::KeyUp(key)); },
						None => {},
					}
				},
				_ => {}
			}
		}
	}

	drop(sdl_tx); // Disconnect such that the cpuloop will exit
	let _ = cpuloop_thread.join();
}

fn sdl_to_keypad(key: sdl2::keycode::KeyCode) -> Option<rboy::KeypadKey> {
	match key {
		sdl2::keycode::KeyCode::Z => Some(rboy::KeypadKey::A),
		sdl2::keycode::KeyCode::X => Some(rboy::KeypadKey::B),
		sdl2::keycode::KeyCode::Up => Some(rboy::KeypadKey::Up),
		sdl2::keycode::KeyCode::Down => Some(rboy::KeypadKey::Down),
		sdl2::keycode::KeyCode::Left => Some(rboy::KeypadKey::Left),
		sdl2::keycode::KeyCode::Right => Some(rboy::KeypadKey::Right),
		sdl2::keycode::KeyCode::Space => Some(rboy::KeypadKey::Select),
		sdl2::keycode::KeyCode::Return => Some(rboy::KeypadKey::Start),
		_ => None,
	}
}

fn recalculate_screen(screen: &sdl2::render::Renderer, arc: &Arc<RwLock<Vec<u8>>>) {
	let mut drawer = screen.drawer();
	drawer.set_draw_color(sdl2::pixels::Color::RGB(0xFF, 0xFF, 0xFF));
	drawer.clear();

	let data =  arc.read().unwrap().clone();
	for y in range(0, 144) {
		for x in range(0, 160) {
			drawer.set_draw_color(sdl2::pixels::Color::RGB(data[y*160*3 + x*3 + 0],
															data[y*160*3 + x*3 + 1],
															data[y*160*3 + x*3 + 2]));
			drawer.draw_rect(&sdl2::rect::Rect::new((x*SCALE) as i32, (y*SCALE) as i32, SCALE as i32, SCALE as i32));
		}
	}
	drawer.present();
}

enum GBEvent {
	KeyUp(rboy::KeypadKey),
	KeyDown(rboy::KeypadKey),
	SpeedUp,
	SlowDown,
}

fn cpuloop(cpu_tx: &Sender<u32>, cpu_rx: &Receiver<GBEvent>, arc: Arc<RwLock<Vec<u8>>>, filename: &str, matches: &getopts::Matches) {
	let opt_c = match matches.opt_present("classic") {
		true => Device::new(filename),
		false => Device::new_cgb(filename),
	};
	let mut c = match opt_c
	{
		Some(cpu) => { cpu },
		None => { error!("Could not get a valid gameboy"); set_exit_status(EXITCODE_CPULOADFAILS); return; },
	};
	c.set_stdout(matches.opt_present("serial"));

	let mut timer = timer::Timer::new().unwrap();
	let periodic = timer.periodic(Duration::milliseconds(8));
	let mut limit_speed = true;

	let waitticks = (4194.304f32 * 4.0) as u32;

	let mut ticks = 0;
	'cpuloop: loop {
		while ticks < waitticks {
			ticks += c.do_cycle();
			if c.check_and_reset_gpu_updated() {
				let mut data = arc.write().unwrap();
				let gpu_data = c.get_gpu_data();
				for i in range(0, data.len()) { data[i] = gpu_data[i]; }
				if cpu_tx.send(0).is_err() { break 'cpuloop };
			}
		}
		ticks -= waitticks;
		if limit_speed { let _ = periodic.recv(); }

		match cpu_rx.try_recv() {
			Ok(event) => match event {
				GBEvent::KeyUp(key) => c.keyup(key),
				GBEvent::KeyDown(key) => c.keydown(key),
				GBEvent::SpeedUp => limit_speed = false,
				GBEvent::SlowDown => limit_speed = true,
			},
			Err(Empty) => {},
			Err(Disconnected) => { break },
		};
	}
}
<|endoftext|>