<commit_before>/// Given a list of entries `items` and the query string, filter out the
/// matched entries using fuzzy search algorithm.

extern crate num_cpus;

use std::sync::{Arc, RwLock, Mutex};

use std::sync::mpsc::channel;
use std::collections::HashMap;
use event::Event;
use item::{Item, MatchedItem, MatchedRange};
use util::eventbox::EventBox;
use score;
use orderedvec::OrderedVec;
use std::cmp::min;
use std::thread;
use std::time::Instant;
use getopts;
use regex::Regex;

const MATCHER_CHUNK_SIZE: usize = 100;
const PROCESS_START_UPDATE_DURATION: u32 = 200; // milliseconds
const PROCESS_UPDATE_DURATION: u64 = 100; // milliseconds
const RESULT_UPDATE_DURATION: u64 = 300; // milliseconds

#[derive(Clone, Copy)]
enum Algorithm {
    FUZZY,
    REGEX,
}

pub struct Matcher {
    pub eb_req: Arc<EventBox<Event>>,       // event box that recieve requests
    eb_notify: Arc<EventBox<Event>>,    // event box that send out notification
    items: Arc<RwLock<Vec<Item>>>,
    new_items: Arc<RwLock<Vec<Item>>>,
    query: Query,
    cache: HashMap<String, MatcherCache>,
    partitions: usize,
    rank_criterion: Arc<Vec<RankCriteria>>,
    is_interactive: bool,
    algorithm: Algorithm,
}

impl<'a> Matcher {
    pub fn new(items: Arc<RwLock<Vec<Item>>>,
               new_items: Arc<RwLock<Vec<Item>>>,
               eb_notify: Arc<EventBox<Event>>) -> Self {

        let mut cache = HashMap::new();
        cache.entry("".to_string()).or_insert(MatcherCache::new());

        Matcher {
            eb_req: Arc::new(EventBox::new()),
            eb_notify: eb_notify,
            items: items,
            new_items: new_items,
            query: Query::new(""),
            cache: cache,
            partitions: num_cpus::get(),
            rank_criterion: Arc::new(vec![RankCriteria::Score,
                                          RankCriteria::Index,
                                          RankCriteria::Begin,
                                          RankCriteria::End]),
            is_interactive: false,
            algorithm: Algorithm::FUZZY,
        }
    }

    pub fn parse_options(&mut self, options: &getopts::Matches) {
        if let Some(tie_breaker) = options.opt_str("t") {
            let mut vec = Vec::new();
            for criteria in tie_breaker.split(',') {
                if let Some(c) = parse_criteria(criteria) {
                    vec.push(c);
                }
            }
            self.rank_criterion = Arc::new(vec);
        }

        if options.opt_present("i") {
            self.is_interactive = true;
        }

        if options.opt_present("regex") {
            self.algorithm = Algorithm::REGEX;
        }
    }

    pub fn process(&mut self) {
        let ref mut cache = self.cache.get_mut(&self.query.get()).unwrap();

        let query = Arc::new(self.query.clone());
        let (tx, rx) = channel();
        let mut guards = vec![];

        let start_pos = Arc::new(Mutex::new(cache.item_pos));

        for _ in 0..self.partitions {
            let items = self.items.clone();
            let start_pos = start_pos.clone();
            let query = query.clone();
            let tx = tx.clone();
            let eb_req = self.eb_req.clone();
            let criterion = self.rank_criterion.clone();
            let is_interactive = self.is_interactive;
            let algorithm = self.algorithm;

            let guard = thread::spawn(move || {
                let items = items.read().unwrap();
                loop {
                    let (start, end) = { // to release the start_pos lock as soon as possible
                        let mut start_idx = start_pos.lock().unwrap();
                        if *start_idx >= items.len() {
                            break;
                        }

                        let start = *start_idx;
                        let end = min(start + MATCHER_CHUNK_SIZE, items.len());
                        *start_idx = end;
                        (start, end)
                    };

                    for i in start..end {
                        let ref item = items[i];
                        if let Some(matched) = match_item(i, &item, &query, &criterion, algorithm) {
                            let _ = tx.send(Some(matched));
                        }
                    }

                    if eb_req.peek(Event::EvMatcherResetQuery) {
                        break;
                    }
                }
                let _ = tx.send(None); // to notify match process end
            });
            guards.push(guard);
        };

        let items_len = self.items.read().unwrap().len();
        let mut matched_items = cache.matched_items.write().unwrap();
        let timer = Instant::now();
        while let Ok(result) = rx.recv() {
            if let Some(matched) = result {
                matched_items.push(matched);
            }

            let start_idx = {*start_pos.lock().unwrap()};

            // update process
            let time = timer.elapsed();
            let mills = (time.as_secs()*1000) as u32 + time.subsec_nanos()/1000/1000;
            if mills > PROCESS_START_UPDATE_DURATION {
                self.eb_notify.set_throttle(Event::EvMatcherUpdateProcess, Box::new(((start_idx+1) *100/(items_len+1)) as u64), PROCESS_UPDATE_DURATION);
            }

            if start_idx >= items_len {
                break;
            }

            if self.eb_req.peek(Event::EvMatcherResetQuery) {
                break;
            }
        }

        // wait for all threads to exit
        for guard in guards {
            let _ = guard.join();
        }

        // consume remaining results.
        while let Ok(result) = rx.try_recv() {
            if let Some(matched) = result {
                matched_items.push(matched);
            }
        }
        cache.item_pos = *start_pos.lock().unwrap();
        self.eb_notify.set(Event::EvMatcherUpdateProcess, Box::new(100 as u64));
    }

    fn reset_query(&mut self, query: &str) {
        self.query = Query::new(query);
        if self.is_interactive {
            self.new_items.write().unwrap().clear();
            self.cache.remove(&query.to_string());
        }
        self.cache.entry(query.to_string()).or_insert(MatcherCache::new());
    }

    pub fn run(&mut self) {
        loop {
            for (e, val) in self.eb_req.wait() {
                match e {
                    Event::EvMatcherNewItem => {}
                    Event::EvMatcherResetQuery => {
                        self.reset_query(&val.downcast::<String>().unwrap());
                        if self.is_interactive {
                            self.eb_notify.set(Event::EvMatcherSync, Box::new(true));
                            let _ = self.eb_req.wait_for(Event::EvModelAck);
                        }
                    }
                    _ => {}
                }
            }

            // insert new items
            {
                let mut buffer = self.new_items.write().unwrap();
                if buffer.len() > 0 {
                    let mut items = self.items.write().unwrap();
                    items.append(&mut buffer);
                }
            }

            self.process();
            if !self.eb_req.peek(Event::EvMatcherResetQuery) {
                let matched_items = self.cache.get_mut(&self.query.get()).unwrap().matched_items.clone();
                if self.is_interactive {
                    self.eb_notify.set_debounce(Event::EvMatcherEnd, Box::new(matched_items), RESULT_UPDATE_DURATION);
                } else {
                    self.eb_notify.set(Event::EvMatcherEnd, Box::new(matched_items));
                }
            }
        }
    }
}

fn match_item(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria],
              algorithm: Algorithm) -> Option<MatchedItem> {
    match algorithm {
        Algorithm::FUZZY => match_item_fuzzy(index, item, query, criterion),
        Algorithm::REGEX => match_item_regex(index, item, query, criterion),
    }
}

fn match_item_fuzzy(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria]) -> Option<MatchedItem> {
    let matched_result = score::fuzzy_match(item.get_lower_chars(), query.get_chars(), query.get_lower_chars());

    if matched_result == None {
        return None;
    }

    let (score, matched_range) = matched_result.unwrap();

    let begin = *matched_range.get(0).unwrap_or(&0) as i64;
    let end = *matched_range.last().unwrap_or(&0) as i64;
    let rank = build_rank(criterion, -score, index as i64, begin, end);

    let mut item = MatchedItem::new(index);
    item.set_matched_range(MatchedRange::Chars(matched_range));
    item.set_rank(rank);
    Some(item)
}

fn match_item_regex(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria]) -> Option<MatchedItem> {
    let matched_result = if query.empty() {
        Some((0, 0))
    } else {
        score::regex_match(item.get_text(), query.get_regex())
    };

    if matched_result == None {
        return None;
    }

    let (begin, end) = matched_result.unwrap();
    let score = end - begin;
    let rank = build_rank(criterion, score as i64, index as i64, begin as i64, end as i64);

    let mut item = MatchedItem::new(index);
    item.set_matched_range(MatchedRange::Range(begin, end));
    item.set_rank(rank);
    Some(item)
}

fn build_rank(criterion: &[RankCriteria], score: i64, index: i64, begin: i64, end: i64) -> [i64; 4] {
    let mut rank = [0; 4];
    for (idx, criteria) in criterion.iter().enumerate().take(4) {
        rank[idx] = match *criteria {
            RankCriteria::Score    => score,
            RankCriteria::Index    => index,
            RankCriteria::Begin    => begin,
            RankCriteria::End      => end,
            RankCriteria::NegScore => -score,
            RankCriteria::NegIndex => -index,
            RankCriteria::NegBegin => -begin,
            RankCriteria::NegEnd   => -end,
        }
    }
    rank
}

struct MatcherCache {
    pub matched_items: Arc<RwLock<OrderedVec<MatchedItem>>>,
    pub item_pos: usize,
}

impl MatcherCache {
    pub fn new() -> Self {
        MatcherCache {
            item_pos: 0,
            matched_items: Arc::new(RwLock::new(OrderedVec::new())),
        }
    }
}

pub enum RankCriteria {
    Score,
    Index,
    Begin,
    End,
    NegScore,
    NegIndex,
    NegBegin,
    NegEnd,
}

pub fn parse_criteria(text: &str) -> Option<RankCriteria> {
    match text.to_lowercase().as_ref() {
        "score"  => Some(RankCriteria::Score),
        "index"  => Some(RankCriteria::Index),
        "begin"  => Some(RankCriteria::Begin),
        "end"    => Some(RankCriteria::End),
        "-score" => Some(RankCriteria::NegScore),
        "-index" => Some(RankCriteria::NegIndex),
        "-begin" => Some(RankCriteria::NegBegin),
        "-end"   => Some(RankCriteria::NegEnd),
        _ => None,
    }
}

// cache for lowercases and others.
#[derive(Clone)]
struct Query {
    query: String,
    query_chars: Vec<char>,
    query_lower_chars: Vec<char>,
    regex: Option<Regex>,
}

impl Query {
    pub fn new(query: &str) -> Self {
        Query {
            query: query.to_string(),
            query_chars: query.chars().collect(),
            query_lower_chars: query.to_lowercase().chars().collect(),
            regex: Regex::new(query).ok(),
        }
    }

    pub fn get(&self) -> String {
        self.query.clone()
    }

    pub fn get_chars(&self) -> &[char] {
        &self.query_chars
    }

    pub fn get_lower_chars(&self) -> &[char] {
        &self.query_lower_chars
    }

    pub fn get_regex(&self) -> &Option<Regex> {
        &self.regex
    }

    pub fn empty(&self) -> bool {
        &self.query == ""
    }
}
<commit_msg>[matcher] skip matching process for interactive mode<commit_after>/// Given a list of entries `items` and the query string, filter out the
/// matched entries using fuzzy search algorithm.

extern crate num_cpus;

use std::sync::{Arc, RwLock, Mutex};

use std::sync::mpsc::channel;
use std::collections::HashMap;
use event::Event;
use item::{Item, MatchedItem, MatchedRange};
use util::eventbox::EventBox;
use score;
use orderedvec::OrderedVec;
use std::cmp::min;
use std::thread;
use std::time::Instant;
use getopts;
use regex::Regex;

const MATCHER_CHUNK_SIZE: usize = 100;
const PROCESS_START_UPDATE_DURATION: u32 = 200; // milliseconds
const PROCESS_UPDATE_DURATION: u64 = 100; // milliseconds
const RESULT_UPDATE_DURATION: u64 = 300; // milliseconds

#[derive(Clone, Copy)]
enum Algorithm {
    FUZZY,
    REGEX,
}

pub struct Matcher {
    pub eb_req: Arc<EventBox<Event>>,       // event box that recieve requests
    eb_notify: Arc<EventBox<Event>>,    // event box that send out notification
    items: Arc<RwLock<Vec<Item>>>,
    new_items: Arc<RwLock<Vec<Item>>>,
    query: Query,
    cache: HashMap<String, MatcherCache>,
    partitions: usize,
    rank_criterion: Arc<Vec<RankCriteria>>,
    is_interactive: bool,
    algorithm: Algorithm,
}

impl<'a> Matcher {
    pub fn new(items: Arc<RwLock<Vec<Item>>>,
               new_items: Arc<RwLock<Vec<Item>>>,
               eb_notify: Arc<EventBox<Event>>) -> Self {

        let mut cache = HashMap::new();
        cache.entry("".to_string()).or_insert(MatcherCache::new());

        Matcher {
            eb_req: Arc::new(EventBox::new()),
            eb_notify: eb_notify,
            items: items,
            new_items: new_items,
            query: Query::new(""),
            cache: cache,
            partitions: num_cpus::get(),
            rank_criterion: Arc::new(vec![RankCriteria::Score,
                                          RankCriteria::Index,
                                          RankCriteria::Begin,
                                          RankCriteria::End]),
            is_interactive: false,
            algorithm: Algorithm::FUZZY,
        }
    }

    pub fn parse_options(&mut self, options: &getopts::Matches) {
        if let Some(tie_breaker) = options.opt_str("t") {
            let mut vec = Vec::new();
            for criteria in tie_breaker.split(',') {
                if let Some(c) = parse_criteria(criteria) {
                    vec.push(c);
                }
            }
            self.rank_criterion = Arc::new(vec);
        }

        if options.opt_present("i") {
            self.is_interactive = true;
        }

        if options.opt_present("regex") {
            self.algorithm = Algorithm::REGEX;
        }
    }

    pub fn process(&mut self) {
        let ref mut cache = self.cache.get_mut(&self.query.get()).unwrap();

        let query = Arc::new(self.query.clone());
        let (tx, rx) = channel();
        let mut guards = vec![];

        let start_pos = Arc::new(Mutex::new(cache.item_pos));

        for _ in 0..self.partitions {
            let items = self.items.clone();
            let start_pos = start_pos.clone();
            let query = query.clone();
            let tx = tx.clone();
            let eb_req = self.eb_req.clone();
            let criterion = self.rank_criterion.clone();
            let algorithm = self.algorithm;

            let guard = thread::spawn(move || {
                let items = items.read().unwrap();
                loop {
                    let (start, end) = { // to release the start_pos lock as soon as possible
                        let mut start_idx = start_pos.lock().unwrap();
                        if *start_idx >= items.len() {
                            break;
                        }

                        let start = *start_idx;
                        let end = min(start + MATCHER_CHUNK_SIZE, items.len());
                        *start_idx = end;
                        (start, end)
                    };

                    for i in start..end {
                        let ref item = items[i];
                        if let Some(matched) = match_item(i, &item, &query, &criterion, algorithm) {
                            let _ = tx.send(Some(matched));
                        }
                    }

                    if eb_req.peek(Event::EvMatcherResetQuery) {
                        break;
                    }
                }
                let _ = tx.send(None); // to notify match process end
            });
            guards.push(guard);
        };

        let items_len = self.items.read().unwrap().len();
        let mut matched_items = cache.matched_items.write().unwrap();
        let timer = Instant::now();
        while let Ok(result) = rx.recv() {
            if let Some(matched) = result {
                matched_items.push(matched);
            }

            let start_idx = {*start_pos.lock().unwrap()};

            // update process
            let time = timer.elapsed();
            let mills = (time.as_secs()*1000) as u32 + time.subsec_nanos()/1000/1000;
            if mills > PROCESS_START_UPDATE_DURATION {
                self.eb_notify.set_throttle(Event::EvMatcherUpdateProcess, Box::new(((start_idx+1) *100/(items_len+1)) as u64), PROCESS_UPDATE_DURATION);
            }

            if start_idx >= items_len {
                break;
            }

            if self.eb_req.peek(Event::EvMatcherResetQuery) {
                break;
            }
        }

        // wait for all threads to exit
        for guard in guards {
            let _ = guard.join();
        }

        // consume remaining results.
        while let Ok(result) = rx.try_recv() {
            if let Some(matched) = result {
                matched_items.push(matched);
            }
        }
        cache.item_pos = *start_pos.lock().unwrap();
        self.eb_notify.set(Event::EvMatcherUpdateProcess, Box::new(100 as u64));
    }

    pub fn process_interactive(&mut self) {
        let ref mut cache = self.cache.get_mut(&self.query.get()).unwrap();
        let mut matched_items = cache.matched_items.write().unwrap();
        let items_len = self.items.read().unwrap().len();
        let timer = Instant::now();
        let start_pos = cache.item_pos;
        let mut last_pos = start_pos;

        for index in start_pos..items_len {
            let mut matched = MatchedItem::new(index);
            matched.set_matched_range(MatchedRange::Range(0, 0));
            matched.set_rank(build_rank(&self.rank_criterion, index as i64, index as i64, 0, 0));

            matched_items.push(matched);

            if self.eb_req.peek(Event::EvMatcherResetQuery) {
                break;
            }

            last_pos = index+1;

            // update process
            let time = timer.elapsed();
            let mills = (time.as_secs()*1000) as u32 + time.subsec_nanos()/1000/1000;
            if mills > PROCESS_START_UPDATE_DURATION {
                self.eb_notify.set_throttle(Event::EvMatcherUpdateProcess, Box::new(((index+1) *100/(items_len+1)) as u64), PROCESS_UPDATE_DURATION);
            }
        }
        cache.item_pos = last_pos;
        self.eb_notify.set(Event::EvMatcherUpdateProcess, Box::new(100 as u64));
    }

    fn reset_query(&mut self, query: &str) {
        self.query = Query::new(query);
        if self.is_interactive {
            self.new_items.write().unwrap().clear();
            self.cache.remove(&query.to_string());
        }
        self.cache.entry(query.to_string()).or_insert(MatcherCache::new());
    }

    pub fn run(&mut self) {
        loop {
            for (e, val) in self.eb_req.wait() {
                match e {
                    Event::EvMatcherNewItem => {}
                    Event::EvMatcherResetQuery => {
                        self.reset_query(&val.downcast::<String>().unwrap());
                        if self.is_interactive {
                            self.eb_notify.set(Event::EvMatcherSync, Box::new(true));
                            let _ = self.eb_req.wait_for(Event::EvModelAck);
                        }
                    }
                    _ => {}
                }
            }

            // insert new items
            {
                let mut buffer = self.new_items.write().unwrap();
                if buffer.len() > 0 {
                    let mut items = self.items.write().unwrap();
                    items.append(&mut buffer);
                }
            }

            if self.is_interactive {
                self.process_interactive();
            } else {
                self.process();
            }

            if !self.eb_req.peek(Event::EvMatcherResetQuery) {
                let matched_items = self.cache.get_mut(&self.query.get()).unwrap().matched_items.clone();
                if self.is_interactive {
                    self.eb_notify.set_debounce(Event::EvMatcherEnd, Box::new(matched_items), RESULT_UPDATE_DURATION);
                } else {
                    self.eb_notify.set(Event::EvMatcherEnd, Box::new(matched_items));
                }
            }
        }
    }
}

fn match_item(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria],
              algorithm: Algorithm) -> Option<MatchedItem> {
    match algorithm {
        Algorithm::FUZZY => match_item_fuzzy(index, item, query, criterion),
        Algorithm::REGEX => match_item_regex(index, item, query, criterion),
    }
}

fn match_item_fuzzy(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria]) -> Option<MatchedItem> {
    let matched_result = score::fuzzy_match(item.get_lower_chars(), query.get_chars(), query.get_lower_chars());

    if matched_result == None {
        return None;
    }

    let (score, matched_range) = matched_result.unwrap();

    let begin = *matched_range.get(0).unwrap_or(&0) as i64;
    let end = *matched_range.last().unwrap_or(&0) as i64;
    let rank = build_rank(criterion, -score, index as i64, begin, end);

    let mut item = MatchedItem::new(index);
    item.set_matched_range(MatchedRange::Chars(matched_range));
    item.set_rank(rank);
    Some(item)
}

fn match_item_regex(index: usize, item: &Item, query: &Query, criterion: &[RankCriteria]) -> Option<MatchedItem> {
    let matched_result = if query.empty() {
        Some((0, 0))
    } else {
        score::regex_match(item.get_text(), query.get_regex())
    };

    if matched_result == None {
        return None;
    }

    let (begin, end) = matched_result.unwrap();
    let score = end - begin;
    let rank = build_rank(criterion, score as i64, index as i64, begin as i64, end as i64);

    let mut item = MatchedItem::new(index);
    item.set_matched_range(MatchedRange::Range(begin, end));
    item.set_rank(rank);
    Some(item)
}

fn build_rank(criterion: &[RankCriteria], score: i64, index: i64, begin: i64, end: i64) -> [i64; 4] {
    let mut rank = [0; 4];
    for (idx, criteria) in criterion.iter().enumerate().take(4) {
        rank[idx] = match *criteria {
            RankCriteria::Score    => score,
            RankCriteria::Index    => index,
            RankCriteria::Begin    => begin,
            RankCriteria::End      => end,
            RankCriteria::NegScore => -score,
            RankCriteria::NegIndex => -index,
            RankCriteria::NegBegin => -begin,
            RankCriteria::NegEnd   => -end,
        }
    }
    rank
}

struct MatcherCache {
    pub matched_items: Arc<RwLock<OrderedVec<MatchedItem>>>,
    pub item_pos: usize,
}

impl MatcherCache {
    pub fn new() -> Self {
        MatcherCache {
            item_pos: 0,
            matched_items: Arc::new(RwLock::new(OrderedVec::new())),
        }
    }
}

pub enum RankCriteria {
    Score,
    Index,
    Begin,
    End,
    NegScore,
    NegIndex,
    NegBegin,
    NegEnd,
}

pub fn parse_criteria(text: &str) -> Option<RankCriteria> {
    match text.to_lowercase().as_ref() {
        "score"  => Some(RankCriteria::Score),
        "index"  => Some(RankCriteria::Index),
        "begin"  => Some(RankCriteria::Begin),
        "end"    => Some(RankCriteria::End),
        "-score" => Some(RankCriteria::NegScore),
        "-index" => Some(RankCriteria::NegIndex),
        "-begin" => Some(RankCriteria::NegBegin),
        "-end"   => Some(RankCriteria::NegEnd),
        _ => None,
    }
}

// cache for lowercases and others.
#[derive(Clone)]
struct Query {
    query: String,
    query_chars: Vec<char>,
    query_lower_chars: Vec<char>,
    regex: Option<Regex>,
}

impl Query {
    pub fn new(query: &str) -> Self {
        Query {
            query: query.to_string(),
            query_chars: query.chars().collect(),
            query_lower_chars: query.to_lowercase().chars().collect(),
            regex: Regex::new(query).ok(),
        }
    }

    pub fn get(&self) -> String {
        self.query.clone()
    }

    pub fn get_chars(&self) -> &[char] {
        &self.query_chars
    }

    pub fn get_lower_chars(&self) -> &[char] {
        &self.query_lower_chars
    }

    pub fn get_regex(&self) -> &Option<Regex> {
        &self.regex
    }

    pub fn empty(&self) -> bool {
        &self.query == ""
    }
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::File;
use std::path;

mod mbc0;
mod mbc1;
mod mbc3;
mod mbc5;

pub trait MBC {
    fn readrom(&self, a: u16) -> u8;
    fn readram(&self, a: u16) -> u8;
    fn writerom(&mut self, a: u16, v: u8);
    fn writeram(&mut self, a: u16, v: u8);
}

pub fn get_mbc(file: path::PathBuf) -> ::StrResult<Box<MBC+'static>> {
    let mut data = vec![];
    try!(File::open(&file).and_then(|mut f| f.read_to_end(&mut data)).map_err(|_| "Could not read ROM"));
    if data.len() < 0x150 { return Err("Rom size to small"); }
    try!(check_checksum(&data));
    match data[0x147] {
        0x00 => mbc0::MBC0::new(data).map(|v| Box::new(v) as Box<MBC>),
        0x01 ... 0x03 => mbc1::MBC1::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        0x0F ... 0x13 => mbc3::MBC3::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        0x19 ... 0x1E => mbc5::MBC5::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        _ => { Err("Unsupported MBC type") },
    }
}

fn ram_size(v: u8) -> usize {
    match v {
        1 => 0x800,
        2 => 0x2000,
        3 => 0x8000,
        4 => 0x20000,
        _ => 0,
    }
}

fn check_checksum(data: &Vec<u8>) -> ::StrResult<()> {
    let mut value: u8 = 0;
    for i in (0x134usize .. 0x14D) {
        value = value.wrapping_sub(data[i]).wrapping_sub(1);
    }
    match data[0x14D] == value
    {
        true => Ok(()),
        false => Err("Cartridge checksum is invalid"),
    }
}
<commit_msg>Add test to mbc checksum<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::path;

mod mbc0;
mod mbc1;
mod mbc3;
mod mbc5;

pub trait MBC {
    fn readrom(&self, a: u16) -> u8;
    fn readram(&self, a: u16) -> u8;
    fn writerom(&mut self, a: u16, v: u8);
    fn writeram(&mut self, a: u16, v: u8);
}

pub fn get_mbc(file: path::PathBuf) -> ::StrResult<Box<MBC+'static>> {
    let mut data = vec![];
    try!(File::open(&file).and_then(|mut f| f.read_to_end(&mut data)).map_err(|_| "Could not read ROM"));
    if data.len() < 0x150 { return Err("Rom size to small"); }
    try!(check_checksum(&data));
    match data[0x147] {
        0x00 => mbc0::MBC0::new(data).map(|v| Box::new(v) as Box<MBC>),
        0x01 ... 0x03 => mbc1::MBC1::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        0x0F ... 0x13 => mbc3::MBC3::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        0x19 ... 0x1E => mbc5::MBC5::new(data, file).map(|v| Box::new(v) as Box<MBC>),
        _ => { Err("Unsupported MBC type") },
    }
}

fn ram_size(v: u8) -> usize {
    match v {
        1 => 0x800,
        2 => 0x2000,
        3 => 0x8000,
        4 => 0x20000,
        _ => 0,
    }
}

fn check_checksum(data: &[u8]) -> ::StrResult<()> {
    let mut value: u8 = 0;
    for i in (0x134 .. 0x14D) {
        value = value.wrapping_sub(data[i]).wrapping_sub(1);
    }
    match data[0x14D] == value
    {
        true => Ok(()),
        false => Err("Cartridge checksum is invalid"),
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn checksum_zero() {
        let mut data = vec![0; 0x150];
        data[0x14D] = -(0x14D_i32 - 0x134_i32) as u8;
        super::check_checksum(&data).unwrap();
    }

    #[test]
    fn checksum_ones() {
        let mut data = vec![1; 0x150];
        data[0x14D] = (-(0x14D_i32 - 0x134_i32) * 2) as u8;
        super::check_checksum(&data).unwrap();
    }
}
<|endoftext|>


<commit_before>use ansi_term::{ANSIString, Colour, Style};
use config::Config;

/// Merges in the default values for the program
pub fn merge_defaults(c: &mut Config) {
    c.set_default("global.modules", vec!["directory", "git", "prompt"]).unwrap();
    c.set_default("global.foreground", "bright_white").unwrap();
    c.set_default("global.background", "blue").unwrap();
    c.set_default("global.style", "default").unwrap();
    c.set_default("global.separator", "").unwrap();
    c.set_default("global.padding_left", " ").unwrap();
    c.set_default("global.padding_right", " ").unwrap();

    c.set_default("modules.directory.max_depth", 4).unwrap();

    c.set_default("modules.exit_code.bg_success", "green").unwrap();
    c.set_default("modules.exit_code.bg_error", "red").unwrap();

    c.set_default("modules.git.symbol_insertion", "+").unwrap();
    c.set_default("modules.git.symbol_deletion", "-").unwrap();
    c.set_default("modules.git.symbol_push", "⇡").unwrap();
    c.set_default("modules.git.symbol_pull", "⇣").unwrap();
    c.set_default("modules.git.show_changed", true).unwrap();
    c.set_default("modules.git.show_diff_stats", false).unwrap();
    c.set_default("modules.git.show_unpushed", true).unwrap();

    c.set_default("modules.prompt.output", "$").unwrap();
    c.set_default("modules.prompt.bg_success", "green").unwrap();
    c.set_default("modules.prompt.bg_error", "red").unwrap();
}

/// Generic formatting method
pub fn format_module<'a>(c: &mut Config,
                         name: &'a str,
                         output: Option<String>,
                         last_successful: Option<&str>)
                         -> (Option<&'a str>, Option<ANSIString<'static>>) {
    // Formatting was not successful if there was nothing to format
    if c.get_str(&format!("modules.{}.output", name)).is_none() && output.is_none() {
        return (None, None);
    }

    // Get config options
    let fg = c.get_str(&format!("modules.{}.foreground", name))
        .unwrap_or_else(|| c.get_str("global.foreground").unwrap_or_default());
    let fg = string_to_colour(fg);

    let bg = c.get_str(&format!("modules.{}.background", name))
        .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
    let bg = string_to_colour(bg);

    let padding_left = c.get_str(&format!("modules.{}.padding_left", name))
        .unwrap_or_else(|| c.get_str("global.padding_left").unwrap_or_default());
    let padding_right = c.get_str(&format!("modules.{}.padding_right", name))
        .unwrap_or_else(|| c.get_str("global.padding_right").unwrap_or_default());
    let separator = c.get_str(&format!("modules.{}.separator", name))
        .unwrap_or_else(|| c.get_str("global.separator").unwrap_or_default());

    let main_style = c.get_str(&format!("modules.{}.style", name))
        .unwrap_or_else(|| c.get_str("global.style").unwrap_or_default());
    let main_style = string_to_style(main_style).on(bg).fg(fg);

    // Override the output if there is "output" in this module's config
    let output = if let Some(s) = c.get_str(&format!("modules.{}.output", name)) {
        s
    } else {
        output.unwrap()
    };

    // Format the main content
    let mut content = format!("\\[{}\\]{}{}{}\\[{}\\]",
                              main_style.prefix(),
                              padding_left,
                              output,
                              padding_right,
                              main_style.suffix());

    // Format the separator
    let main_style = main_style.on(fg).fg(bg);
    if let Some(name) = last_successful {
        // There is a visible module that comes after this one
        let next_bg = c.get_str(&format!("modules.{}.background", name))
            .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
        let next_bg = string_to_colour(next_bg);

        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          main_style.on(next_bg).prefix(),
                          separator,
                          main_style.on(next_bg).suffix());
    } else {
        // This is the final module
        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          bg.prefix(),
                          separator,
                          bg.suffix());
    }

    (Some(name), Some(ANSIString::from(content)))
}

// Converts a string (from the config file) to a Colour
// See: https://upload.wikimedia.org/wikipedia/commons/1/15/Xterm_256color_chart.svg
fn string_to_colour(s: String) -> Colour {
    if let Ok(i) = s.parse::<u8>() {
        Colour::Fixed(i)
    } else {
        let s = s.to_lowercase();
        match s.as_ref() {
            "black" => Colour::Fixed(0),
            "bright_black" => Colour::Fixed(8),
            "red" => Colour::Fixed(1),
            "bright_red" => Colour::Fixed(9),
            "green" => Colour::Fixed(2),
            "bright_green" => Colour::Fixed(10),
            "yellow" => Colour::Fixed(3),
            "bright_yellow" => Colour::Fixed(11),
            "blue" => Colour::Fixed(4),
            "bright_blue" => Colour::Fixed(12),
            "purple" => Colour::Fixed(5),
            "bright_purple" => Colour::Fixed(13),
            "cyan" => Colour::Fixed(6),
            "bright_cyan" => Colour::Fixed(14),
            "white" => Colour::Fixed(7),
            "bright_white" => Colour::Fixed(15),
            _ => panic!("Invalid color option: {} in config file!", s),
        }
    }
}

fn string_to_style(s: String) -> Style {
    let s = s.to_lowercase();
    match s.as_ref() {
        "default" | "" | "normal" => Style::new(),
        "bold" => Style::new().bold(),
        "dimmed" => Style::new().dimmed(),
        "italic" => Style::new().italic(),
        "underline" => Style::new().underline(),
        "blink" => Style::new().blink(),
        "reverse" => Style::new().reverse(),
        "hidden" => Style::new().hidden(),
        "strikethrough" => Style::new().strikethrough(),
        _ => panic!("Unknown style property: {} in config file!", s),
    }
}

pub fn format_module_prompt<'a>(c: &mut Config,
                                last_successful: Option<&'a str>,
                                exit_code: &str)
                                -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.prompt.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.prompt.bg_error").unwrap_or_default(),
    };
    c.set("modules.prompt.background", bg).unwrap();

    let output = c.get_str("modules.prompt.output").unwrap_or_default();

    format_module(c, "prompt", Some(output), last_successful)
}

pub fn format_module_exit_code<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>,
                                   exit_code: &str)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.exit_code.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.exit_code.bg_error").unwrap_or_default(),
    };
    c.set("modules.exit_code.background", bg).unwrap();

    format_module(c, "exit_code", Some(exit_code.to_string()), last_successful)
}

pub fn format_module_directory<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use std::env;
    use std::path::PathBuf;

    let home = env::var("HOME").unwrap();
    let cwd = env::current_dir().unwrap();

    // Convert "/home/user/directory" to "~/directory"
    let mut shortened_cwd: PathBuf;
    if let Ok(stripped_cwd) = cwd.strip_prefix(&home) {
        shortened_cwd = PathBuf::from("~").join(stripped_cwd);
    } else {
        shortened_cwd = env::current_dir().unwrap();
    }

    let depth = shortened_cwd.components().count();

    // Max number of directories we want to see
    let max_depth = c.get_int("modules.directory.max_depth").unwrap_or_default() as usize;

    if depth > max_depth {
        let comp_iter = shortened_cwd.clone();
        let comp_iter = comp_iter.components();

        shortened_cwd = PathBuf::new();
        for (i, component) in comp_iter.enumerate() {
            if i < (max_depth / 2) || i >= (depth - (max_depth / 2)) {
                shortened_cwd.push(component.as_os_str());
            } else if i == (max_depth / 2) {
                shortened_cwd.push("...");
            }
        }
    }

    format_module(c,
                  "directory",
                  Some(format!("{}", shortened_cwd.display())),
                  last_successful)
}

pub fn format_module_git<'a>(c: &mut Config,
                             last_successful: Option<&'a str>)
                             -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use git2::{Branch, Repository};
    use std::env;

    let mut output = String::new();

    if let Ok(repo) = Repository::discover(env::current_dir().unwrap_or_default()) {
        let local = repo.head().unwrap();

        // Output current branch name
        output.push_str(local.shorthand().unwrap());

        // Show local changes
        let show_diffs = c.get_bool("modules.git.show_diff_stats").unwrap_or_default();
        let show_changed = c.get_bool("modules.git.show_changed").unwrap_or_default();
        if show_changed {
            let diff_stats = repo.diff_index_to_workdir(None, None).unwrap();
            let diff_stats = diff_stats.stats().unwrap();

            if diff_stats.files_changed() > 0 {
                let symbol_insertion = c.get_str("modules.git.symbol_insertion")
                    .unwrap_or_default();

                if show_diffs {
                    // Show actual number of insertions or deletions
                    let symbol_deletion = c.get_str("modules.git.symbol_deletion")
                        .unwrap_or_default();
                    output.push_str(&format!(" ({}{}, {}{})",
                                             symbol_deletion,
                                             diff_stats.deletions(),
                                             symbol_insertion,
                                             diff_stats.insertions()));
                } else {
                    // Do not show insertions nor deletions, just say
                    // something was changed
                    output.push_str(&format!(" {}", symbol_insertion));
                }
            }
        }

        let local = Branch::wrap(local);

        // Show unpushed/unpulled commits
        let show_unpushed = c.get_bool("modules.git.show_unpushed").unwrap_or_default();
        if show_unpushed {
            if let Ok(upstream) = local.upstream() {
                let local_ref = local.get();
                let upstream_ref = upstream.get();

                if let Ok((ahead, behind)) =
                    repo.graph_ahead_behind(local_ref.target().unwrap(),
                                            upstream_ref.target().unwrap()) {

                    if ahead != 0 {
                        let symbol_push = c.get_str("modules.git.symbol_push").unwrap_or_default();
                        output.push_str(&format!(" {}{}", symbol_push, ahead));
                    }

                    if behind != 0 {
                        let symbol_pull = c.get_str("modules.git.symbol_pull").unwrap_or_default();
                        output.push_str(&format!(" {}{}", symbol_pull, behind));
                    }
                }
            }
        }
    }

    if output.is_empty() {
        (None, None)
    } else {
        format_module(c, "git", Some(output), last_successful)
    }
}

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

    // Note: ansi_term has already tested this kind of thing, this is
    // just to ensure that passing a string and matching on it works
    // like I expect (sanity-check)

    #[test]
    fn test_colour_and_style() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            color: &'static str,
            style: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         color: "blue",
                         style: "bold",
                         expected: &format!("\x1B[1;38;5;4m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         color: "bright_white",
                         style: "",
                         expected: &format!("\x1B[38;5;15m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let style = string_to_style(test.style.to_string());
            let color = string_to_colour(test.color.to_string());
            let result = format!("{}", style.fg(color).paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    fn test_string_to_style() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "bold",
                         expected: &format!("\x1B[1m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "italic",
                         expected: &format!("\x1B[3m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_style(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_style_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "bold" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_style(test.input.to_string());
        }
    }

    #[test]
    fn test_string_to_colour() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "green",
                         expected: &format!("\x1B[38;5;2m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "bright_green",
                         expected: &format!("\x1B[38;5;10m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_colour(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_colour_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "green" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_colour(test.input.to_string());
        }
    }
}
<commit_msg>Fix compilation error with env::current_dir()<commit_after>use ansi_term::{ANSIString, Colour, Style};
use config::Config;

/// Merges in the default values for the program
pub fn merge_defaults(c: &mut Config) {
    c.set_default("global.modules", vec!["directory", "git", "prompt"]).unwrap();
    c.set_default("global.foreground", "bright_white").unwrap();
    c.set_default("global.background", "blue").unwrap();
    c.set_default("global.style", "default").unwrap();
    c.set_default("global.separator", "").unwrap();
    c.set_default("global.padding_left", " ").unwrap();
    c.set_default("global.padding_right", " ").unwrap();

    c.set_default("modules.directory.max_depth", 4).unwrap();

    c.set_default("modules.exit_code.bg_success", "green").unwrap();
    c.set_default("modules.exit_code.bg_error", "red").unwrap();

    c.set_default("modules.git.symbol_insertion", "+").unwrap();
    c.set_default("modules.git.symbol_deletion", "-").unwrap();
    c.set_default("modules.git.symbol_push", "⇡").unwrap();
    c.set_default("modules.git.symbol_pull", "⇣").unwrap();
    c.set_default("modules.git.show_changed", true).unwrap();
    c.set_default("modules.git.show_diff_stats", false).unwrap();
    c.set_default("modules.git.show_unpushed", true).unwrap();

    c.set_default("modules.prompt.output", "$").unwrap();
    c.set_default("modules.prompt.bg_success", "green").unwrap();
    c.set_default("modules.prompt.bg_error", "red").unwrap();
}

/// Generic formatting method
pub fn format_module<'a>(c: &mut Config,
                         name: &'a str,
                         output: Option<String>,
                         last_successful: Option<&str>)
                         -> (Option<&'a str>, Option<ANSIString<'static>>) {
    // Formatting was not successful if there was nothing to format
    if c.get_str(&format!("modules.{}.output", name)).is_none() && output.is_none() {
        return (None, None);
    }

    // Get config options
    let fg = c.get_str(&format!("modules.{}.foreground", name))
        .unwrap_or_else(|| c.get_str("global.foreground").unwrap_or_default());
    let fg = string_to_colour(fg);

    let bg = c.get_str(&format!("modules.{}.background", name))
        .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
    let bg = string_to_colour(bg);

    let padding_left = c.get_str(&format!("modules.{}.padding_left", name))
        .unwrap_or_else(|| c.get_str("global.padding_left").unwrap_or_default());
    let padding_right = c.get_str(&format!("modules.{}.padding_right", name))
        .unwrap_or_else(|| c.get_str("global.padding_right").unwrap_or_default());
    let separator = c.get_str(&format!("modules.{}.separator", name))
        .unwrap_or_else(|| c.get_str("global.separator").unwrap_or_default());

    let main_style = c.get_str(&format!("modules.{}.style", name))
        .unwrap_or_else(|| c.get_str("global.style").unwrap_or_default());
    let main_style = string_to_style(main_style).on(bg).fg(fg);

    // Override the output if there is "output" in this module's config
    let output = if let Some(s) = c.get_str(&format!("modules.{}.output", name)) {
        s
    } else {
        output.unwrap()
    };

    // Format the main content
    let mut content = format!("\\[{}\\]{}{}{}\\[{}\\]",
                              main_style.prefix(),
                              padding_left,
                              output,
                              padding_right,
                              main_style.suffix());

    // Format the separator
    let main_style = main_style.on(fg).fg(bg);
    if let Some(name) = last_successful {
        // There is a visible module that comes after this one
        let next_bg = c.get_str(&format!("modules.{}.background", name))
            .unwrap_or_else(|| c.get_str("global.background").unwrap_or_default());
        let next_bg = string_to_colour(next_bg);

        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          main_style.on(next_bg).prefix(),
                          separator,
                          main_style.on(next_bg).suffix());
    } else {
        // This is the final module
        content = format!("{}\\[{}\\]{}\\[{}\\]",
                          content,
                          bg.prefix(),
                          separator,
                          bg.suffix());
    }

    (Some(name), Some(ANSIString::from(content)))
}

// Converts a string (from the config file) to a Colour
// See: https://upload.wikimedia.org/wikipedia/commons/1/15/Xterm_256color_chart.svg
fn string_to_colour(s: String) -> Colour {
    if let Ok(i) = s.parse::<u8>() {
        Colour::Fixed(i)
    } else {
        let s = s.to_lowercase();
        match s.as_ref() {
            "black" => Colour::Fixed(0),
            "bright_black" => Colour::Fixed(8),
            "red" => Colour::Fixed(1),
            "bright_red" => Colour::Fixed(9),
            "green" => Colour::Fixed(2),
            "bright_green" => Colour::Fixed(10),
            "yellow" => Colour::Fixed(3),
            "bright_yellow" => Colour::Fixed(11),
            "blue" => Colour::Fixed(4),
            "bright_blue" => Colour::Fixed(12),
            "purple" => Colour::Fixed(5),
            "bright_purple" => Colour::Fixed(13),
            "cyan" => Colour::Fixed(6),
            "bright_cyan" => Colour::Fixed(14),
            "white" => Colour::Fixed(7),
            "bright_white" => Colour::Fixed(15),
            _ => panic!("Invalid color option: {} in config file!", s),
        }
    }
}

fn string_to_style(s: String) -> Style {
    let s = s.to_lowercase();
    match s.as_ref() {
        "default" | "" | "normal" => Style::new(),
        "bold" => Style::new().bold(),
        "dimmed" => Style::new().dimmed(),
        "italic" => Style::new().italic(),
        "underline" => Style::new().underline(),
        "blink" => Style::new().blink(),
        "reverse" => Style::new().reverse(),
        "hidden" => Style::new().hidden(),
        "strikethrough" => Style::new().strikethrough(),
        _ => panic!("Unknown style property: {} in config file!", s),
    }
}

pub fn format_module_prompt<'a>(c: &mut Config,
                                last_successful: Option<&'a str>,
                                exit_code: &str)
                                -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.prompt.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.prompt.bg_error").unwrap_or_default(),
    };
    c.set("modules.prompt.background", bg).unwrap();

    let output = c.get_str("modules.prompt.output").unwrap_or_default();

    format_module(c, "prompt", Some(output), last_successful)
}

pub fn format_module_exit_code<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>,
                                   exit_code: &str)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    let bg = match exit_code.as_ref() {
        "0" => c.get_str("modules.exit_code.bg_success").unwrap_or_default(),
        _ => c.get_str("modules.exit_code.bg_error").unwrap_or_default(),
    };
    c.set("modules.exit_code.background", bg).unwrap();

    format_module(c, "exit_code", Some(exit_code.to_string()), last_successful)
}

pub fn format_module_directory<'a>(c: &mut Config,
                                   last_successful: Option<&'a str>)
                                   -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use std::env;
    use std::path::PathBuf;

    let home = env::var("HOME").unwrap();
    let cwd = env::current_dir().unwrap();

    // Convert "/home/user/directory" to "~/directory"
    let mut shortened_cwd: PathBuf;
    if let Ok(stripped_cwd) = cwd.strip_prefix(&home) {
        shortened_cwd = PathBuf::from("~").join(stripped_cwd);
    } else {
        shortened_cwd = env::current_dir().unwrap();
    }

    let depth = shortened_cwd.components().count();

    // Max number of directories we want to see
    let max_depth = c.get_int("modules.directory.max_depth").unwrap_or_default() as usize;

    if depth > max_depth {
        let comp_iter = shortened_cwd.clone();
        let comp_iter = comp_iter.components();

        shortened_cwd = PathBuf::new();
        for (i, component) in comp_iter.enumerate() {
            if i < (max_depth / 2) || i >= (depth - (max_depth / 2)) {
                shortened_cwd.push(component.as_os_str());
            } else if i == (max_depth / 2) {
                shortened_cwd.push("...");
            }
        }
    }

    format_module(c,
                  "directory",
                  Some(format!("{}", shortened_cwd.display())),
                  last_successful)
}

pub fn format_module_git<'a>(c: &mut Config,
                             last_successful: Option<&'a str>)
                             -> (Option<&'a str>, Option<ANSIString<'static>>) {
    use git2::{Branch, Repository};
    use std::env;

    let mut output = String::new();

    if let Ok(repo) = Repository::discover(env::current_dir().unwrap()) {
        let local = repo.head().unwrap();

        // Output current branch name
        output.push_str(local.shorthand().unwrap());

        // Show local changes
        let show_diffs = c.get_bool("modules.git.show_diff_stats").unwrap_or_default();
        let show_changed = c.get_bool("modules.git.show_changed").unwrap_or_default();
        if show_changed {
            let diff_stats = repo.diff_index_to_workdir(None, None).unwrap();
            let diff_stats = diff_stats.stats().unwrap();

            if diff_stats.files_changed() > 0 {
                let symbol_insertion = c.get_str("modules.git.symbol_insertion")
                    .unwrap_or_default();

                if show_diffs {
                    // Show actual number of insertions or deletions
                    let symbol_deletion = c.get_str("modules.git.symbol_deletion")
                        .unwrap_or_default();
                    output.push_str(&format!(" ({}{}, {}{})",
                                             symbol_deletion,
                                             diff_stats.deletions(),
                                             symbol_insertion,
                                             diff_stats.insertions()));
                } else {
                    // Do not show insertions nor deletions, just say
                    // something was changed
                    output.push_str(&format!(" {}", symbol_insertion));
                }
            }
        }

        let local = Branch::wrap(local);

        // Show unpushed/unpulled commits
        let show_unpushed = c.get_bool("modules.git.show_unpushed").unwrap_or_default();
        if show_unpushed {
            if let Ok(upstream) = local.upstream() {
                let local_ref = local.get();
                let upstream_ref = upstream.get();

                if let Ok((ahead, behind)) =
                    repo.graph_ahead_behind(local_ref.target().unwrap(),
                                            upstream_ref.target().unwrap()) {

                    if ahead != 0 {
                        let symbol_push = c.get_str("modules.git.symbol_push").unwrap_or_default();
                        output.push_str(&format!(" {}{}", symbol_push, ahead));
                    }

                    if behind != 0 {
                        let symbol_pull = c.get_str("modules.git.symbol_pull").unwrap_or_default();
                        output.push_str(&format!(" {}{}", symbol_pull, behind));
                    }
                }
            }
        }
    }

    if output.is_empty() {
        (None, None)
    } else {
        format_module(c, "git", Some(output), last_successful)
    }
}

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

    // Note: ansi_term has already tested this kind of thing, this is
    // just to ensure that passing a string and matching on it works
    // like I expect (sanity-check)

    #[test]
    fn test_colour_and_style() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            color: &'static str,
            style: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         color: "blue",
                         style: "bold",
                         expected: &format!("\x1B[1;38;5;4m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         color: "bright_white",
                         style: "",
                         expected: &format!("\x1B[38;5;15m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let style = string_to_style(test.style.to_string());
            let color = string_to_colour(test.color.to_string());
            let result = format!("{}", style.fg(color).paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    fn test_string_to_style() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "bold",
                         expected: &format!("\x1B[1m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "italic",
                         expected: &format!("\x1B[3m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_style(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_style_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "bold" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_style(test.input.to_string());
        }
    }

    #[test]
    fn test_string_to_colour() {
        const CONTENT: &'static str = "hi";

        struct Test<'a> {
            input: &'static str,
            expected: &'a str,
        }

        let tests = [Test {
                         input: "green",
                         expected: &format!("\x1B[38;5;2m{}\x1B[0m", CONTENT),
                     },
                     Test {
                         input: "bright_green",
                         expected: &format!("\x1B[38;5;10m{}\x1B[0m", CONTENT),
                     }];

        for test in &tests {
            let result = string_to_colour(test.input.to_string());
            let result = format!("{}", result.paint(CONTENT));

            assert_eq!(test.expected, result);
        }
    }

    #[test]
    #[should_panic]
    fn test_string_to_colour_invalid_input() {
        struct Test {
            input: &'static str,
        }

        let tests = [Test { input: "green" }, Test { input: "invalid" }];

        for test in &tests {
            string_to_colour(test.input.to_string());
        }
    }
}
<|endoftext|>


<commit_before>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use transaction::{TransactionId, TransactionIdGenerator};

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Vec<u8>)
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>)
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.get_address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![]
        }
    }

    pub fn run(self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.get_address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.get_address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            let mut origin = self.routing_table
                                .find_node(&origin_address)
                                .expect("Got find node response from unknown node");
                            origin.received_response(transaction_id);
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        let mut origin = self.routing_table
                            .find_node(&origin_address)
                            .expect("Got packet response from unknown node");
                        origin.received_response(transaction_id);
                    }
                    Response::Ping => {
                        let mut origin = self.routing_table
                            .find_node(&origin_address)
                            .expect("Got ping response from unknown node");
                        origin.received_response(transaction_id);
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.get_address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) => {
                self.network.send_packet(recipient, payload, event_loop);
            }
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect())).unwrap();
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<commit_msg>Always bind udp server to 0.0.0.0<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use transaction::{TransactionId, TransactionIdGenerator};

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Vec<u8>)
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>)
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.get_address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![]
        }
    }

    pub fn run(self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.get_address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.get_address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            let mut origin = self.routing_table
                                .find_node(&origin_address)
                                .expect("Got find node response from unknown node");
                            origin.received_response(transaction_id);
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        let mut origin = self.routing_table
                            .find_node(&origin_address)
                            .expect("Got packet response from unknown node");
                        origin.received_response(transaction_id);
                    }
                    Response::Ping => {
                        let mut origin = self.routing_table
                            .find_node(&origin_address)
                            .expect("Got ping response from unknown node");
                        origin.received_response(transaction_id);
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.get_address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) => {
                self.network.send_packet(recipient, payload, event_loop);
            }
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect())).unwrap();
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<|endoftext|>


<commit_before>//! Write benchmark information to the output repository

use errors::{Result, ResultExt};

use std::fs::{self, File, OpenOptions, read_dir};
use std::io::{Read, Write, BufRead, BufReader};
use std::path::{PathBuf};
use std::process::Command;
use std::str;
use std::collections::HashSet;

use serde_json;
use rustc_perf_collector::CommitData;
use rust_sysroot::git::Commit as GitCommit;
use execute::Benchmark;

pub struct Repo {
    path: PathBuf,
    retries: Vec<String>
}

impl Repo {
    fn git(&self, args: &[&str]) -> Result<()> {
        let mut command = Command::new("git");
        command.current_dir(&self.path);
        info!("git {:?}", args);
        command.args(args);
        let status = command.status().chain_err(|| format!("could not spawn git {:?}", args))?;
        if !status.success() {
            bail!("command `git {:?}` failed in `{}`", args, self.path.display());
        }
        Ok(())
    }

    pub fn open(path: PathBuf) -> Result<Self> {
        let mut result = Repo { path: path, retries: vec![] };

        if !result.retries_file().exists() {
            // try not to nuke random repositories.
            bail!("`{}` file not present", result.retries_file().display());
        }

        result.git(&["fetch"])?;
        result.git(&["reset", "--hard", "@{upstream}"])?;

        fs::create_dir_all(result.times()).chain_err(|| "can't create `times/`")?;
        result.load_retries()?;

        Ok(result)
    }

    pub fn success(&self, data: &CommitData) -> Result<()> {
        self.add_commit_data(data)?;
        self.commit_and_push(&format!("{} - success", data.commit.sha))?;
        Ok(())
    }

    pub fn find_missing_commits<'a>(
        &self,
        commits: &'a [GitCommit],
        benchmarks: &[Benchmark],
        triple: &str,
    ) -> Result<Vec<&'a GitCommit>> {
        let mut have = HashSet::new();
        let path = self.times();
        for entry in read_dir(path)? {
            let entry = entry?;
            let filename = entry.file_name().to_string_lossy().to_string();
            let sha = &filename[filename.find("00:00").unwrap() + 6..filename.find("-x86").unwrap()];
            have.insert(sha.to_string());
        }

        let file = OpenOptions::new()
            .write(true).read(true).create(true).open(self.broken_commits_file())?;
        let file = BufReader::new(file);
        for line in file.lines() {
            let line = line?;
            let sha = &line[..line.find(":").unwrap()];
            have.insert(sha.to_string());
        }

        let missing = commits.iter().filter(|c| {
            !have.contains(&c.sha) || {
                self.load_commit_data(c, triple).ok().map(|data| {
                    benchmarks.iter()
                        .all(|b| data.benchmarks.keys().find(|k| **k == b.name).is_some())
                }).unwrap_or(true)
            }
        }).collect::<Vec<_>>();

        Ok(missing)
    }

    fn commit_and_push(&self, message: &str) -> Result<()> {
        self.write_retries()?;
        self.git(&["add", "retries", "times"])?;
        self.git(&["commit", "-m", message])?;
        self.git(&["push"])?;
        Ok(())
    }

    pub fn load_commit_data(&self, commit: &GitCommit, triple: &str) -> Result<CommitData> {
        let filepath = self.times()
            .join(format!("{}-{}-{}.json", commit.date.to_rfc3339(), commit.sha, triple));
        info!("loading file {}", filepath.display());
        let mut file = File::open(&filepath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        let data = serde_json::from_str(&contents)?;
        Ok(data)
    }

    fn add_commit_data(&self, data: &CommitData) -> Result<()> {
        let commit = &data.commit;
        let filepath = self.times()
            .join(format!("{}-{}-{}.json", commit.date.to_rfc3339(), commit.sha, data.triple));
        info!("creating file {}", filepath.display());
        let mut file = File::create(&filepath)?;
        serde_json::to_writer(&mut file, &data)?;
        Ok(())
    }

    fn load_retries(&mut self) -> Result<()> {
        let mut retries = OpenOptions::new().read(true).write(true).create(true).open(self.retries_file())
            .chain_err(|| format!("can't create `{}`", self.retries_file().display()))?;
        let mut retries_s = String::new();
        retries.read_to_string(&mut retries_s)?;
        self.retries = retries_s.split('\n')
            .map(|line| line.trim())
            .filter(|line| !line.is_empty())
            .map(|line| {
                if line.len() == 40 {
                    Ok(line.to_owned())
                } else {
                    bail!("bad retry hash `{}`", line)
                }
            }).collect::<Result<_>>()?;
        info!("loaded retries: {:?}", self.retries);
        Ok(())
    }

    fn write_retries(&self) -> Result<()> {
        info!("writing retries");
        let mut retries = OpenOptions::new().write(true).truncate(true)
            .open(self.retries_file())
            .chain_err(|| "can't create `retries`")?;
        for retry in self.retries.iter() {
            writeln!(retries, "{}", retry)?;
        }
        Ok(())
    }

    pub fn next_retry(&mut self) -> Option<String> {
        if self.retries.len() == 0 {
            None
        } else {
            Some(self.retries.remove(0))
        }
    }

    fn broken_commits_file(&self) -> PathBuf {
        self.path.join("broken-commits-log")
    }

    fn retries_file(&self) -> PathBuf {
        self.path.join("retries")
    }

    pub fn times(&self) -> PathBuf {
        self.path.join("times")
    }
}
<commit_msg>Correct logic for missing benchmarks.<commit_after>//! Write benchmark information to the output repository

use errors::{Result, ResultExt};

use std::fs::{self, File, OpenOptions, read_dir};
use std::io::{Read, Write, BufRead, BufReader};
use std::path::{PathBuf};
use std::process::Command;
use std::str;
use std::collections::HashSet;

use serde_json;
use rustc_perf_collector::CommitData;
use rust_sysroot::git::Commit as GitCommit;
use execute::Benchmark;

pub struct Repo {
    path: PathBuf,
    retries: Vec<String>
}

impl Repo {
    fn git(&self, args: &[&str]) -> Result<()> {
        let mut command = Command::new("git");
        command.current_dir(&self.path);
        info!("git {:?}", args);
        command.args(args);
        let status = command.status().chain_err(|| format!("could not spawn git {:?}", args))?;
        if !status.success() {
            bail!("command `git {:?}` failed in `{}`", args, self.path.display());
        }
        Ok(())
    }

    pub fn open(path: PathBuf) -> Result<Self> {
        let mut result = Repo { path: path, retries: vec![] };

        if !result.retries_file().exists() {
            // try not to nuke random repositories.
            bail!("`{}` file not present", result.retries_file().display());
        }

        result.git(&["fetch"])?;
        result.git(&["reset", "--hard", "@{upstream}"])?;

        fs::create_dir_all(result.times()).chain_err(|| "can't create `times/`")?;
        result.load_retries()?;

        Ok(result)
    }

    pub fn success(&self, data: &CommitData) -> Result<()> {
        self.add_commit_data(data)?;
        self.commit_and_push(&format!("{} - success", data.commit.sha))?;
        Ok(())
    }

    pub fn find_missing_commits<'a>(
        &self,
        commits: &'a [GitCommit],
        benchmarks: &[Benchmark],
        triple: &str,
    ) -> Result<Vec<&'a GitCommit>> {
        let mut have = HashSet::new();
        let path = self.times();
        for entry in read_dir(path)? {
            let entry = entry?;
            let filename = entry.file_name().to_string_lossy().to_string();
            let sha = &filename[filename.find("00:00").unwrap() + 6..filename.find("-x86").unwrap()];
            have.insert(sha.to_string());
        }

        let file = OpenOptions::new()
            .write(true).read(true).create(true).open(self.broken_commits_file())?;
        let file = BufReader::new(file);
        for line in file.lines() {
            let line = line?;
            let sha = &line[..line.find(":").unwrap()];
            have.insert(sha.to_string());
        }

        let missing = commits.iter().filter(|c| {
            !have.contains(&c.sha) || {
                self.load_commit_data(c, triple).ok().map(|data| {
                    benchmarks.iter().any(|b| data.benchmarks.keys().find(|k| **k == b.name).is_none())
                }).unwrap_or(true)
            }
        }).collect::<Vec<_>>();

        Ok(missing)
    }

    fn commit_and_push(&self, message: &str) -> Result<()> {
        self.write_retries()?;
        self.git(&["add", "retries", "times"])?;
        self.git(&["commit", "-m", message])?;
        self.git(&["push"])?;
        Ok(())
    }

    pub fn load_commit_data(&self, commit: &GitCommit, triple: &str) -> Result<CommitData> {
        let filepath = self.times()
            .join(format!("{}-{}-{}.json", commit.date.to_rfc3339(), commit.sha, triple));
        info!("loading file {}", filepath.display());
        let mut file = File::open(&filepath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        let data = serde_json::from_str(&contents)?;
        Ok(data)
    }

    fn add_commit_data(&self, data: &CommitData) -> Result<()> {
        let commit = &data.commit;
        let filepath = self.times()
            .join(format!("{}-{}-{}.json", commit.date.to_rfc3339(), commit.sha, data.triple));
        info!("creating file {}", filepath.display());
        let mut file = File::create(&filepath)?;
        serde_json::to_writer(&mut file, &data)?;
        Ok(())
    }

    fn load_retries(&mut self) -> Result<()> {
        let mut retries = OpenOptions::new().read(true).write(true).create(true).open(self.retries_file())
            .chain_err(|| format!("can't create `{}`", self.retries_file().display()))?;
        let mut retries_s = String::new();
        retries.read_to_string(&mut retries_s)?;
        self.retries = retries_s.split('\n')
            .map(|line| line.trim())
            .filter(|line| !line.is_empty())
            .map(|line| {
                if line.len() == 40 {
                    Ok(line.to_owned())
                } else {
                    bail!("bad retry hash `{}`", line)
                }
            }).collect::<Result<_>>()?;
        info!("loaded retries: {:?}", self.retries);
        Ok(())
    }

    fn write_retries(&self) -> Result<()> {
        info!("writing retries");
        let mut retries = OpenOptions::new().write(true).truncate(true)
            .open(self.retries_file())
            .chain_err(|| "can't create `retries`")?;
        for retry in self.retries.iter() {
            writeln!(retries, "{}", retry)?;
        }
        Ok(())
    }

    pub fn next_retry(&mut self) -> Option<String> {
        if self.retries.len() == 0 {
            None
        } else {
            Some(self.retries.remove(0))
        }
    }

    fn broken_commits_file(&self) -> PathBuf {
        self.path.join("broken-commits-log")
    }

    fn retries_file(&self) -> PathBuf {
        self.path.join("retries")
    }

    pub fn times(&self) -> PathBuf {
        self.path.join("times")
    }
}
<|endoftext|>