<commit_before>extern crate unicode_width;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::collections::HashSet;
use std::process;
use std::io::{stdin, BufRead};
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;
use unicode_width::UnicodeWidthStr;

#[cfg(windows)] pub const NEWLINE: &'static str = "\r\n";
#[cfg(not(windows))] pub const NEWLINE: &'static str = "\n";

const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/version.txt"));
const TIMESTAMP: &'static str = include_str!(concat!(env!("OUT_DIR"), "/timestamp.txt"));
const TARGET: &'static str = include_str!(concat!(env!("OUT_DIR"), "/target.txt"));
const COMMIT: &'static str = include_str!(concat!(env!("OUT_DIR"), "/commit.txt"));

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            process::exit(1);
        },
    };

    if args.help { return; }

    if args.version {
        if COMMIT == "" {
            println!("heatseeker {} (built {} for {})", VERSION, TIMESTAMP, TARGET);
        } else {
            println!("heatseeker {}-{} (built {} for {})", VERSION, COMMIT, TIMESTAMP, TARGET);
        }
        return;
    }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search, args.filter_only);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        let desired_rows = if args.full_screen { 999 } else { 20 };
        let selections = event_loop(desired_rows, &choices, &initial_search, args.filter_only);
        print!("{}", selections);
    }
}

fn event_loop(desired_rows: u16, choices: &[&str], initial_search: &str, filter_only: bool) -> String {
    let mut search = Search::new(choices, initial_search.to_string(), filter_only);
    let mut screen = screen::new(desired_rows);

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(screen.as_mut(), &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in &keys {
            handle_key(&mut search, key, screen.visible_choices());
        }
    }

    screen.blank_screen();
    search.get_selections()
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace |
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') |
        Control('g') => search.cancel(),
        Control('t') => { search.toggle_selection(); search.down(visible_choices); },
        Control('p') |
        Up |
        ShiftTab => search.up(visible_choices),
        Control('n') |
        Down |
        Tab => search.down(visible_choices),
        Home => search.home(),
        End => search.end(visible_choices),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    scroll_offset: usize,
    cursor_index: usize,
    state: SearchState,
    selections: HashSet<String>,
    filter_only: bool,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String, filter_only: bool) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            scroll_offset: 0,
            cursor_index: 0,
            state: InProgress,
            selections: HashSet::new(),
            filter_only: filter_only,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.scroll_offset == 0;
        if self.cursor_index == 0 {
            if should_wrap {
                if match_count > visible_choices as usize {
                    self.scroll_offset = match_count - visible_choices as usize;
                    self.cursor_index = visible_choices as usize - 1;
                } else {
                    self.cursor_index = limit;
                }
            } else {
                self.scroll_offset -= 1;
            }
        } else {
            self.cursor_index -= 1;
        }
    }

    fn down(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.cursor_index + self.scroll_offset == match_count - 1;
        if self.cursor_index == limit {
            if should_wrap {
                self.cursor_index = 0;
                self.scroll_offset = 0;
            } else {
                self.scroll_offset += 1;
            }
        } else {
            self.cursor_index += 1;
        }
    }

    fn home(&mut self) {
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn end(&mut self, visible_choices: u16) {
        self.home();
        self.up(visible_choices);
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query, self.filter_only);
            self.stale = false;
        }
    }

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
        if self.selections.contains(&selection) {
            self.selections.remove(&selection);
        } else {
            self.selections.insert(selection);
        }
    }

    fn get_selections(&mut self) -> String {
        let mut ret = String::new();
        if self.state != Canceled {
            for selection in &self.selections {
                ret.push_str(selection);
                ret.push_str(NEWLINE);
            }
            if ret.is_empty() {
                self.recompute_matches();
                let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
                ret.push_str(&selection);
            }
        }
        ret
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut dyn Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices){}", search.query, search.matches.len(), search.choices.len(), NEWLINE));

    print_matches(screen, &search.matches, &search.query, search.scroll_offset, search.cursor_index, &search.selections);

    let query_str: &str = &search.query;
    screen.move_cursor_to_prompt_line(2 + UnicodeWidthStr::width(query_str) as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut dyn Screen, matches: &[&str], query: &str, scroll_offset: usize, cursor_index: usize, selections: &HashSet<String>) {
    let mut i = 1;
    for choice in matches[scroll_offset..].iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width();
        let mut annotated_choice = choice.to_string();
        if selections.contains(&annotated_choice) {
            annotated_choice.push(' ');
            if cfg!(windows) {
                annotated_choice.push('√');
            } else {
                annotated_choice.push('✓');
            }
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == cursor_index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else if highlight {
                screen.write_red(s);
            } else { screen.write(s); }
        });
        if i >= screen.visible_choices() as usize {
            return;
        }
        screen.write(NEWLINE);
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut dyn FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().count();
    let mut chars_to_draw = min(chars_in_choice, max_width as usize);
    while UnicodeWidthStr::width(slice_chars(choice, 0, chars_to_draw)) > max_width as usize {
        chars_to_draw -= 1;
    }
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(slice_chars(choice, last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(slice_chars(choice, idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(slice_chars(choice, last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let stdin = stdin();
    let mut lines = Vec::new();

    let mut stdin = stdin.lock();
    let mut first_error = None;
    let mut suppressed = 0;
    loop {
        let mut s = String::new();
        match stdin.read_line(&mut s) {
            Ok(_) => {
                if s.is_empty() { break; }
                trim(&mut s);
                lines.push(s);
            },
            Err(e) => {
                if first_error.is_some() {
                    suppressed = suppressed + 1;
                } else {
                    first_error = Some(e);
                }
            }
        }
    }
    if first_error.is_some() {
        eprintln!("Warning: Failed to parse one or more lines (\"{}\"); {} additional error(s) suppressed", first_error.unwrap(), suppressed);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' && x != '\r' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

// This is technically a standard library function, but it's plastered with
// stability warnings and therefore only available on the nightly channel.
fn slice_chars(s: &str, begin: usize, end: usize) -> &str {
    assert!(begin <= end);
    let mut count = 0;
    let mut begin_byte = None;
    let mut end_byte = None;

    // This could be even more efficient by not decoding,
    // only finding the char boundaries
    for (idx, _) in s.char_indices() {
        if count == begin { begin_byte = Some(idx); }
        if count == end { end_byte = Some(idx); break; }
        count += 1;
    }
    if begin_byte.is_none() && count == begin { begin_byte = Some(s.len()) }
    if end_byte.is_none() && count == end { end_byte = Some(s.len()) }

    match (begin_byte, end_byte) {
        (None, _) => panic!("slice_chars: `begin` is beyond end of string"),
        (_, None) => panic!("slice_chars: `end` is beyond end of string"),
        (Some(a), Some(b)) => unsafe { s.get_unchecked(a..b) }
    }
}

#[cfg(test)]
mod tests {
    use super::{trim, delete_last_word};

    #[test]
    fn trim_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            trim(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("\n", "");
        should_become("\n\n", "");
        should_become("asdf", "asdf");
        should_become("asdf\n", "asdf");
        should_become("asdf\nasdf\n", "asdf\nasdf");
        should_become("asdf\r\n", "asdf");
    }

    #[test]
    fn delete_word_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            delete_last_word(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("a", "");
        should_become("asdf", "");
        should_become("asdf asdf asdf", "asdf asdf ");
        should_become("asdf asdf asdf ", "asdf asdf ");
        should_become("asdf asdf asdf  ", "asdf asdf ");
    }
}
<commit_msg>Implement PgUp/PgDown<commit_after>extern crate unicode_width;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::collections::HashSet;
use std::process;
use std::io::{stdin, BufRead};
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;
use unicode_width::UnicodeWidthStr;

#[cfg(windows)] pub const NEWLINE: &'static str = "\r\n";
#[cfg(not(windows))] pub const NEWLINE: &'static str = "\n";

const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/version.txt"));
const TIMESTAMP: &'static str = include_str!(concat!(env!("OUT_DIR"), "/timestamp.txt"));
const TARGET: &'static str = include_str!(concat!(env!("OUT_DIR"), "/target.txt"));
const COMMIT: &'static str = include_str!(concat!(env!("OUT_DIR"), "/commit.txt"));

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            process::exit(1);
        },
    };

    if args.help { return; }

    if args.version {
        if COMMIT == "" {
            println!("heatseeker {} (built {} for {})", VERSION, TIMESTAMP, TARGET);
        } else {
            println!("heatseeker {}-{} (built {} for {})", VERSION, COMMIT, TIMESTAMP, TARGET);
        }
        return;
    }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search, args.filter_only);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        let desired_rows = if args.full_screen { 999 } else { 20 };
        let selections = event_loop(desired_rows, &choices, &initial_search, args.filter_only);
        print!("{}", selections);
    }
}

fn event_loop(desired_rows: u16, choices: &[&str], initial_search: &str, filter_only: bool) -> String {
    let mut search = Search::new(choices, initial_search.to_string(), filter_only);
    let mut screen = screen::new(desired_rows);

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(screen.as_mut(), &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in &keys {
            handle_key(&mut search, key, screen.visible_choices());
        }
    }

    screen.blank_screen();
    search.get_selections()
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace |
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') |
        Control('g') => search.cancel(),
        Control('t') => { search.toggle_selection(); search.down(visible_choices); },
        Control('p') |
        Up |
        ShiftTab => search.up(visible_choices),
        Control('n') |
        Down |
        Tab => search.down(visible_choices),
        Home => search.home(),
        End => search.end(visible_choices),
        Enter => search.done(),
        PgUp => search.pgup(visible_choices),
        PgDown => search.pgdown(visible_choices),
        _ => {},
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    scroll_offset: usize,
    cursor_index: usize,
    state: SearchState,
    selections: HashSet<String>,
    filter_only: bool,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String, filter_only: bool) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            scroll_offset: 0,
            cursor_index: 0,
            state: InProgress,
            selections: HashSet::new(),
            filter_only: filter_only,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.scroll_offset == 0;
        if self.cursor_index == 0 {
            if should_wrap {
                if match_count > visible_choices as usize {
                    self.scroll_offset = match_count - visible_choices as usize;
                    self.cursor_index = visible_choices as usize - 1;
                } else {
                    self.cursor_index = limit;
                }
            } else {
                self.scroll_offset -= 1;
            }
        } else {
            self.cursor_index -= 1;
        }
    }

    fn down(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.cursor_index + self.scroll_offset == match_count - 1;
        if self.cursor_index == limit {
            if should_wrap {
                self.cursor_index = 0;
                self.scroll_offset = 0;
            } else {
                self.scroll_offset += 1;
            }
        } else {
            self.cursor_index += 1;
        }
    }

    fn home(&mut self) {
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn end(&mut self, visible_choices: u16) {
        self.home();
        self.up(visible_choices);
    }

    fn pgup(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset == 0 && self.cursor_index == 0 {
                return;
            }
            self.up(visible_choices);
        }
    }

    fn pgdown(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset + self.cursor_index == self.matches.len() - 1 {
                return;
            }
            self.down(visible_choices);
        }
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query, self.filter_only);
            self.stale = false;
        }
    }

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
        if self.selections.contains(&selection) {
            self.selections.remove(&selection);
        } else {
            self.selections.insert(selection);
        }
    }

    fn get_selections(&mut self) -> String {
        let mut ret = String::new();
        if self.state != Canceled {
            for selection in &self.selections {
                ret.push_str(selection);
                ret.push_str(NEWLINE);
            }
            if ret.is_empty() {
                self.recompute_matches();
                let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
                ret.push_str(&selection);
            }
        }
        ret
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut dyn Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices){}", search.query, search.matches.len(), search.choices.len(), NEWLINE));

    print_matches(screen, &search.matches, &search.query, search.scroll_offset, search.cursor_index, &search.selections);

    let query_str: &str = &search.query;
    screen.move_cursor_to_prompt_line(2 + UnicodeWidthStr::width(query_str) as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut dyn Screen, matches: &[&str], query: &str, scroll_offset: usize, cursor_index: usize, selections: &HashSet<String>) {
    let mut i = 1;
    for choice in matches[scroll_offset..].iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width();
        let mut annotated_choice = choice.to_string();
        if selections.contains(&annotated_choice) {
            annotated_choice.push(' ');
            if cfg!(windows) {
                annotated_choice.push('√');
            } else {
                annotated_choice.push('✓');
            }
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == cursor_index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else if highlight {
                screen.write_red(s);
            } else { screen.write(s); }
        });
        if i >= screen.visible_choices() as usize {
            return;
        }
        screen.write(NEWLINE);
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut dyn FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().count();
    let mut chars_to_draw = min(chars_in_choice, max_width as usize);
    while UnicodeWidthStr::width(slice_chars(choice, 0, chars_to_draw)) > max_width as usize {
        chars_to_draw -= 1;
    }
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(slice_chars(choice, last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(slice_chars(choice, idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(slice_chars(choice, last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let stdin = stdin();
    let mut lines = Vec::new();

    let mut stdin = stdin.lock();
    let mut first_error = None;
    let mut suppressed = 0;
    loop {
        let mut s = String::new();
        match stdin.read_line(&mut s) {
            Ok(_) => {
                if s.is_empty() { break; }
                trim(&mut s);
                lines.push(s);
            },
            Err(e) => {
                if first_error.is_some() {
                    suppressed = suppressed + 1;
                } else {
                    first_error = Some(e);
                }
            }
        }
    }
    if first_error.is_some() {
        eprintln!("Warning: Failed to parse one or more lines (\"{}\"); {} additional error(s) suppressed", first_error.unwrap(), suppressed);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' && x != '\r' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

// This is technically a standard library function, but it's plastered with
// stability warnings and therefore only available on the nightly channel.
fn slice_chars(s: &str, begin: usize, end: usize) -> &str {
    assert!(begin <= end);
    let mut count = 0;
    let mut begin_byte = None;
    let mut end_byte = None;

    // This could be even more efficient by not decoding,
    // only finding the char boundaries
    for (idx, _) in s.char_indices() {
        if count == begin { begin_byte = Some(idx); }
        if count == end { end_byte = Some(idx); break; }
        count += 1;
    }
    if begin_byte.is_none() && count == begin { begin_byte = Some(s.len()) }
    if end_byte.is_none() && count == end { end_byte = Some(s.len()) }

    match (begin_byte, end_byte) {
        (None, _) => panic!("slice_chars: `begin` is beyond end of string"),
        (_, None) => panic!("slice_chars: `end` is beyond end of string"),
        (Some(a), Some(b)) => unsafe { s.get_unchecked(a..b) }
    }
}

#[cfg(test)]
mod tests {
    use super::{trim, delete_last_word};

    #[test]
    fn trim_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            trim(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("\n", "");
        should_become("\n\n", "");
        should_become("asdf", "asdf");
        should_become("asdf\n", "asdf");
        should_become("asdf\nasdf\n", "asdf\nasdf");
        should_become("asdf\r\n", "asdf");
    }

    #[test]
    fn delete_word_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            delete_last_word(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("a", "");
        should_become("asdf", "");
        should_become("asdf asdf asdf", "asdf asdf ");
        should_become("asdf asdf asdf ", "asdf asdf ");
        should_become("asdf asdf asdf  ", "asdf asdf ");
    }
}
<|endoftext|>


<commit_before>extern crate rand;

#[macro_use]
extern crate lazy_static;

use rand::{XorShiftRng, Rng, thread_rng, SeedableRng};
use std::fmt;
use std::fmt::{Formatter, Debug};
use std::sync::Mutex;

trait Fitness: Ord {
    fn max() -> Self;
}

trait Specimen: Sized {
    type F: Fitness;

    fn fitness(&self) -> Self::F;
    fn reevaluate(&mut self, rng: &mut XorShiftRng);
    fn mutate(generation: &mut[Self], rng: &mut XorShiftRng);
    fn breed(parents: &[Self], rng: &mut XorShiftRng) -> Self;
    fn filter_strongest(species: &mut Vec<Self>);

    fn initial() -> Vec<Self>;
}

fn genetic<S: Specimen>(threshold: S::F, max_iters: usize, nparents: usize, nchildren: usize) -> S {
    let mut species: Vec<S> = S::initial();
    let mut best = (S::F::max(), 0);
    for _ in 0..max_iters {
        species = next_gen(species, nparents, nchildren);

        for (i, specimen) in species.iter().enumerate() {
            let fitness = specimen.fitness();
            if fitness <= best.0 {
                best = (fitness, i)
            }
        }

        if best.0 <= threshold {
            println!("found!");
            break;
        }
    }

    species.swap_remove(best.1)
}

#[inline(never)]
fn next_gen<S: Specimen>(mut species: Vec<S>, nparents: usize, nchildren: usize) -> Vec<S> {
    let mut rng = RNG.lock().unwrap();

    // 1. filter strongest
    S::filter_strongest(&mut species);

    // 2. bear children
    let mut children = vec![];
    let mut family = vec![];
    for _ in 0..nchildren {
        make_family(&mut species, nparents, &mut family, &mut rng);
        let child = S::breed(&family, &mut rng);
        children.push(child);
        species.extend(family.drain(..));
    }

    // 3. mutate
    S::mutate(&mut children, &mut rng);
    for child in children.iter_mut() {
        child.reevaluate(&mut rng);
    }

    if KILL_PARENTS {
        children
    } else {
        species.extend(children);
        species
    }
}

fn make_family<S: Specimen>(species: &mut Vec<S>, nparents: usize, family: &mut Vec<S>, rng: &mut XorShiftRng) {
    assert!(family.is_empty());
    for _ in 0..nparents {
        let next = rng.gen_range(0, species.len());
        let s = species.swap_remove(next);
        family.push(s);
    }
}

//---------------------------------------------------

// BEST FOR SIZE=1000
//const ELITE: usize = 3;
//const NCHILDREN: usize = 65;
//const KILL_PARENTS: bool = false;
//const SIZE: usize = 1000;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;

//const ELITE: usize = 4;
//const NCHILDREN: usize = 130;
//const KILL_PARENTS: bool = true;
const ELITE: usize = 2;
//const NCHILDREN: usize = 65;
const NCHILDREN: usize = 2;
const KILL_PARENTS: bool = false;
const SIZE: usize = 100000;
const MUTATION_PROBABILITY: f32 = 0.01f32;
const MAX_ITERS: usize = 50000;
const NPARENTS: usize = 2;

//
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 10000;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 20;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 1000;

lazy_static! {
    static ref RNG: Mutex<XorShiftRng> = Mutex::new(XorShiftRng::new_unseeded());
}

struct Board {
    fitness: usize,
    queens: Vec<usize>,
}

impl Board {
    fn new(queens: Vec<usize>) -> Board {
        Board { fitness: ::std::usize::MAX, queens: queens }
    }

    #[inline(never)]
    fn do_mutate(&mut self, rng: &mut XorShiftRng) {
        let i = rng.gen_range(0, SIZE);
        let mut j = rng.gen_range(0, SIZE-1);

        if j >= i {
            j += 1;
        }

        self.queens.swap(i, j);
    }

    fn breed_pmx_norng(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        let mut child = Board::new(parents[0].queens.clone());

        // A very important assert! Without it, the program runs twice slower! o_O
        assert!(parents.len() == NPARENTS);

        {
            let cq: &mut [usize] = &mut child.queens;
            let mut inverse = vec![0; SIZE];
            for i in 0..SIZE {
                let y = cq[i];
                inverse[y] = i;
            }

            let mut x = rng.gen_range(0, SIZE);
            for _ in 0..SIZE/2 {
                x += 1;
                if x == SIZE {
                    x = 0;
                }

                let parent = 1;
                let mother_y = cq[x];
                let father_y = parents[parent].queens[x];
                let father_x = inverse[father_y];
                cq[father_x] = mother_y;
                cq[x] = father_y;
                inverse[mother_y] = father_x;
                if NPARENTS > 2 {
                    // Optimization: this is not required when NPARENTS==2
                    inverse[father_y] = x;
                }
            }
        }

        child
    }

    // The algorithm is described in "Genetic Algorithm Solution of the TSP Avoiding Special Crossover and Mutation"
    // http://user.ceng.metu.edu.tr/~ucoluk/research/publications/tspnew.pdf
    #[inline(never)]
    fn breed_pmx(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        // A very important assert! Without it, the program runs twice slower! o_O
        assert!(parents.len() == NPARENTS);

        let mut child = Board::new(parents[0].queens.clone());
        {
            let mut inverse = vec![0; SIZE];
            let cq = &mut child.queens;
            for i in 0..SIZE {
                let y = cq[i];
                inverse[y] = i;
            }

            for x in 0..SIZE {
                let parent = rng.gen_range(0, parents.len());
                if parent == 1 {
                    let mother_y = cq[x];
                    let father_y = parents[parent].queens[x];
                    let father_x = inverse[father_y];
                    cq[father_x] = mother_y;
                    cq[x] = father_y;
                    inverse[mother_y] = father_x;
                    if NPARENTS > 2 { // Optimization: this is not required when NPARENTS==2
                        inverse[father_y] = x;
                    }
                }
            }
        }

        child
    }

//    fn permutation_valid(p: &[usize]) -> bool {
//        let mut used = vec![0; p.len()];
//        for &i in p.iter() {
//            used[i] = 1;
//        }
//        for &u in used.iter() {
//            if u != 1 {
//                return false;
//            }
//        }
//
//        true
//    }
}

impl Fitness for usize {
    fn max() -> Self {
        ::std::usize::MAX
    }
}

impl Specimen for Board {
    type F = usize;

    #[inline(never)]
    fn reevaluate(&mut self, rng: &mut XorShiftRng) {
        let q: &mut [usize] = &mut self.queens;
        let mut diag_counts = vec![0; 4*SIZE-2];
        let (diag_counts1, diag_counts2) = diag_counts.split_at_mut(2*SIZE-1);
        for x in 0..SIZE {
            let d1 = SIZE-1 + x - q[x];
            let d2 = x + q[x];
            diag_counts1[d1] += 1;
            diag_counts2[d2] += 1;
        }

        let mut f = 0;
        for x in 0..SIZE {
            let d1 = SIZE-1 + x - q[x];
            let d2 = x + q[x];
            f += diag_counts1[d1] + diag_counts2[d2];
        }

        self.fitness = f;

        // perform a conflict minimization step
        let i = rng.gen_range(0, SIZE);
        let qi = q[i];
        let di1o = SIZE-1 + i - q[i];
        let di2o = i + q[i];
        let ci_old = diag_counts1[di1o] + diag_counts2[di2o];
        let mut max_deltaf = 0;
        let mut max_deltaf_j = 0;
        for j in 0..SIZE {
            if i == j {
                continue;
            }

            let qj = q[j];

            let dj1o = SIZE-1 + j - qj;
            let dj2o = j + qj;
            let di1n = SIZE-1 + j - qi;
            let di2n = j + qi;
            let dj1n = SIZE-1 + i - qj;
            let dj2n = i + qj;

            let fold;
            let fnew;
            if di1o == dj1o {
                assert!(di2n == dj2n);
                let c1o = (diag_counts1[di1o]-1)*4; // neg
                let c2n = (diag_counts2[di2n]+1)*4;

                let c2o = diag_counts2[di2o] + diag_counts2[dj2o];
                let c1n = diag_counts1[di1n] + diag_counts1[dj1n];

                fold = c1o + (c2o - 1)*2;
                fnew = c2n + (c1n + 1)*2;
            } else if di2o == dj2o {
                assert!(di1n == dj1n);
                let c2o = (diag_counts2[di2o]-1)*4; // neg
                let c1n = (diag_counts1[di1n]+1)*4;

                let c1o = diag_counts1[di1o] + diag_counts1[dj1o]; // neg
                let c2n = diag_counts2[di2n] + diag_counts2[dj2n];

                fold = c2o + (c1o - 1)*2;
                fnew = c1n + (c2n + 1)*2;
            } else {
                let cj_old = diag_counts1[dj1o] + diag_counts2[dj2o];
                let ci_new = diag_counts1[di1n] + diag_counts2[di2n] + 2;
                let cj_new = diag_counts1[dj1n] + diag_counts2[dj2n] + 2;

                fold = (ci_old + cj_old)*2;
                fnew = (ci_new + cj_new)*2;
            }

            if fold > max_deltaf + fnew {
                max_deltaf = fold - fnew;
                max_deltaf_j = j;
            }
        }

        if max_deltaf != 0 {
//            println!("queens_o={:?}, fitness_o={}, case={}", q, self.fitness, case);
            q.swap(i, max_deltaf_j);
            self.fitness -= max_deltaf;

//            // TEST DEBUG
//            let mut diag_counts = vec![0; 4*SIZE-2];
//            let (diag_counts1, diag_counts2) = diag_counts.split_at_mut(2*SIZE-1);
//            for x in 0..SIZE {
//                let d1 = SIZE-1 + x - q[x];
//                let d2 = x + q[x];
//                diag_counts1[d1] += 1;
//                diag_counts2[d2] += 1;
//            }
//
//            let mut f = 0;
//            for x in 0..SIZE {
//                let d1 = SIZE-1 + x - q[x];
//                let d2 = x + q[x];
//                f += diag_counts1[d1] + diag_counts2[d2];
//            }
//
//            assert!(f == self.fitness, "max_deltaf={}, f={}, fitness={}, queens={:?}", max_deltaf, f, self.fitness, q);
        }
    }

    fn fitness(&self) -> usize {
        self.fitness
    }

    #[inline(never)]
    fn mutate(new_gen: &mut [Self], rng: &mut XorShiftRng) {
        for board in new_gen.iter_mut() {
            if rng.next_f32() < MUTATION_PROBABILITY {
                board.do_mutate(rng);
            }
        }
    }

    #[inline(never)]
    fn breed(parents: &[Self], rng: &mut XorShiftRng) -> Self {
//        Self::breed_pmx(parents, rng)
        Self::breed_pmx_norng(parents, rng)
    }

    #[inline(never)]
    fn filter_strongest(species: &mut Vec<Self>) {
        if species.len() <= ELITE {
            return;
        }

        species.sort_by(|s, t| s.fitness.partial_cmp(&t.fitness).unwrap());
        species.truncate(ELITE);
    }

    #[inline(never)]
    fn initial() -> Vec<Self> {
        let mut boards = Vec::with_capacity(NCHILDREN);

        let mut ys = Vec::with_capacity(SIZE);

        let mut rng = RNG.lock().unwrap();
        for _ in 0..NCHILDREN {
            let mut board = Board::new(Vec::with_capacity(SIZE));

            for y in 0..SIZE {
                ys.push(y);
            }

            for _ in 0..SIZE {
                let z = rng.gen_range(0, ys.len());
                let y = ys.swap_remove(z);
                board.queens.push(y);
            }

            board.reevaluate(&mut rng);
            boards.push(board);
        }
        boards
    }
}

impl Debug for Board {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "fitness={}, queens=[{:?}]", self.fitness, self.queens)
    }
}

fn genetic_queens() {
    {
        let mut rng = RNG.lock().unwrap();
        let mut trng = thread_rng();
        rng.reseed([trng.next_u32(), trng.next_u32(), trng.next_u32(), trng.next_u32()]);
    }

    let best = genetic::<Board>(2*SIZE, MAX_ITERS, NPARENTS, NCHILDREN);
//    println!("best: {:?}", &best);
    println!("best: {:?}", best.fitness());
}

fn main() {
    for _ in 0..10 {
        genetic_queens()
    }
}

<commit_msg>slight optimization and readability improvement<commit_after>extern crate rand;

#[macro_use]
extern crate lazy_static;

use rand::{XorShiftRng, Rng, thread_rng, SeedableRng};
use std::fmt;
use std::fmt::{Formatter, Debug};
use std::sync::Mutex;
use std::cmp;

trait Fitness: Ord {
    fn max() -> Self;
}

trait Specimen: Sized {
    type F: Fitness;

    fn fitness(&self) -> Self::F;
    fn reevaluate(&mut self, rng: &mut XorShiftRng);
    fn mutate(generation: &mut[Self], rng: &mut XorShiftRng);
    fn breed(parents: &[Self], rng: &mut XorShiftRng) -> Self;
    fn filter_strongest(species: &mut Vec<Self>);

    fn initial() -> Vec<Self>;
}

fn genetic<S: Specimen>(threshold: S::F, max_iters: usize, nparents: usize, nchildren: usize) -> S {
    let mut species: Vec<S> = S::initial();
    let mut best = (S::F::max(), 0);
    for _ in 0..max_iters {
        species = next_gen(species, nparents, nchildren);

        for (i, specimen) in species.iter().enumerate() {
            let fitness = specimen.fitness();
            if fitness <= best.0 {
                best = (fitness, i)
            }
        }

        if best.0 <= threshold {
            println!("found!");
            break;
        }
    }

    species.swap_remove(best.1)
}

#[inline(never)]
fn next_gen<S: Specimen>(mut species: Vec<S>, nparents: usize, nchildren: usize) -> Vec<S> {
    let mut rng = RNG.lock().unwrap();

    // 1. filter strongest
    S::filter_strongest(&mut species);

    // 2. bear children
    let mut children = vec![];
    let mut family = vec![];
    for _ in 0..nchildren {
        make_family(&mut species, nparents, &mut family, &mut rng);
        let child = S::breed(&family, &mut rng);
        children.push(child);
        species.extend(family.drain(..));
    }

    // 3. mutate
    S::mutate(&mut children, &mut rng);
    for child in children.iter_mut() {
        child.reevaluate(&mut rng);
    }

    if KILL_PARENTS {
        children
    } else {
        species.extend(children);
        species
    }
}

fn make_family<S: Specimen>(species: &mut Vec<S>, nparents: usize, family: &mut Vec<S>, rng: &mut XorShiftRng) {
    assert!(family.is_empty());
    for _ in 0..nparents {
        let next = rng.gen_range(0, species.len());
        let s = species.swap_remove(next);
        family.push(s);
    }
}

//---------------------------------------------------

// BEST FOR SIZE=1000
//const ELITE: usize = 3;
//const NCHILDREN: usize = 65;
//const KILL_PARENTS: bool = false;
//const SIZE: usize = 1000;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;

//const ELITE: usize = 4;
//const NCHILDREN: usize = 130;
//const KILL_PARENTS: bool = true;
const ELITE: usize = 2;
//const NCHILDREN: usize = 65;
const NCHILDREN: usize = 1;
const KILL_PARENTS: bool = false;
const SIZE: usize = 100000;
const MUTATION_PROBABILITY: f32 = 0.01f32;
const MAX_ITERS: usize = 50000;
const NPARENTS: usize = 2;

//
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 10000;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 20;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 1000;

lazy_static! {
    static ref RNG: Mutex<XorShiftRng> = Mutex::new(XorShiftRng::new_unseeded());
}

struct Board {
    fitness: usize,
    queens: Vec<usize>,
}

impl Board {
    fn new(queens: Vec<usize>) -> Board {
        Board { fitness: ::std::usize::MAX, queens: queens }
    }

    #[inline(never)]
    fn do_mutate(&mut self, rng: &mut XorShiftRng) {
        let i = rng.gen_range(0, SIZE);
        let mut j = rng.gen_range(0, SIZE-1);

        if j >= i {
            j += 1;
        }

        self.queens.swap(i, j);
    }

    fn breed_pmx_norng(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        let mut child = Board::new(parents[0].queens.clone());

        // A very important assert! Without it, the program runs twice slower! o_O
        assert!(parents.len() == NPARENTS);

        {
            let cq: &mut [usize] = &mut child.queens;
            let mut inverse = vec![0; SIZE];
            for i in 0..SIZE {
                let y = cq[i];
                inverse[y] = i;
            }

            let mut x = rng.gen_range(0, SIZE);
            for _ in 0..SIZE/2 {
                x += 1;
                if x == SIZE {
                    x = 0;
                }

                let parent = 1;
                let mother_y = cq[x];
                let father_y = parents[parent].queens[x];
                let father_x = inverse[father_y];
                cq[father_x] = mother_y;
                cq[x] = father_y;
                inverse[mother_y] = father_x;
                if NPARENTS > 2 {
                    // Optimization: this is not required when NPARENTS==2
                    inverse[father_y] = x;
                }
            }
        }

        child
    }

    // The algorithm is described in "Genetic Algorithm Solution of the TSP Avoiding Special Crossover and Mutation"
    // http://user.ceng.metu.edu.tr/~ucoluk/research/publications/tspnew.pdf
    #[inline(never)]
    fn breed_pmx(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        // A very important assert! Without it, the program runs twice slower! o_O
        assert!(parents.len() == NPARENTS);

        let mut child = Board::new(parents[0].queens.clone());
        {
            let mut inverse = vec![0; SIZE];
            let cq = &mut child.queens;
            for i in 0..SIZE {
                let y = cq[i];
                inverse[y] = i;
            }

            for x in 0..SIZE {
                let parent = rng.gen_range(0, parents.len());
                if parent == 1 {
                    let mother_y = cq[x];
                    let father_y = parents[parent].queens[x];
                    let father_x = inverse[father_y];
                    cq[father_x] = mother_y;
                    cq[x] = father_y;
                    inverse[mother_y] = father_x;
                    if NPARENTS > 2 { // Optimization: this is not required when NPARENTS==2
                        inverse[father_y] = x;
                    }
                }
            }
        }

        child
    }

//    fn permutation_valid(p: &[usize]) -> bool {
//        let mut used = vec![0; p.len()];
//        for &i in p.iter() {
//            used[i] = 1;
//        }
//        for &u in used.iter() {
//            if u != 1 {
//                return false;
//            }
//        }
//
//        true
//    }
}

impl Fitness for usize {
    fn max() -> Self {
        ::std::usize::MAX
    }
}

impl Specimen for Board {
    type F = usize;

    #[inline(never)]
    fn reevaluate(&mut self, rng: &mut XorShiftRng) {
        let q: &mut [usize] = &mut self.queens;
        let mut diag_counts = vec![0; 4*SIZE-2];
        let (diag_counts1, diag_counts2) = diag_counts.split_at_mut(2*SIZE-1);
        for x in 0..SIZE {
            let d1 = SIZE-1 + x - q[x];
            let d2 = x + q[x];
            diag_counts1[d1] += 1;
            diag_counts2[d2] += 1;
        }

        let mut f = 0;
        for x in 0..SIZE {
            let d1 = SIZE-1 + x - q[x];
            let d2 = x + q[x];
            f += diag_counts1[d1] + diag_counts2[d2];
        }

        self.fitness = f;

        // perform a conflict minimization step
        let i = rng.gen_range(0, SIZE);
        let qi = q[i];
        let (di1o, di2o) = (SIZE-1 + i - qi, i + qi);
        let (ci1o, ci2o) = (diag_counts1[di1o], diag_counts2[di2o]);
        let ci_old = ci1o + ci2o;
        let mut max_deltaf = 0;
        let mut max_deltaf_j = 0;
        for j in 0..SIZE {
            if i == j {
                continue;
            }

            let qj = q[j];

            let (dj1o, dj2o) = (SIZE-1 + j - qj, j + qj);
            let (cj1o, cj2o) = (diag_counts1[dj1o], diag_counts2[dj2o]);
            let cj_old = cj1o + cj2o;

            if ci_old==2 && cj_old==2 {
                continue;
            }

            let (di1n, di2n) = (SIZE-1 + j - qi, j + qi);
            let (ci1n, ci2n) = (diag_counts1[di1n], diag_counts2[di2n]);

            let (dj1n, dj2n) = (SIZE-1 + i - qj, i + qj);
            let (cj1n, cj2n) = (diag_counts1[dj1n], diag_counts2[dj2n]);

            let fold;
            let fnew;
            if di1o == dj1o {
                assert!(di2n == dj2n);
                let c1o = (ci1o-1)*4; // neg
                let c2n = (ci2n+1)*4;

                let c2o = ci2o + cj2o;
                let c1n = ci1n + cj1n;

                fold = c1o + (c2o - 1)*2;
                fnew = c2n + (c1n + 1)*2;
            } else if di2o == dj2o {
                assert!(di1n == dj1n);
                let c2o = (ci2o-1)*4; // neg
                let c1n = (ci1n+1)*4;

                let c1o = ci1o + cj1o; // neg
                let c2n = ci2n + cj2n;

                fold = c2o + (c1o - 1)*2;
                fnew = c1n + (c2n + 1)*2;
            } else {
                let ci_new = ci1n + ci2n + 2;
                let cj_new = cj1n + cj2n + 2;

                fold = (ci_old + cj_old)*2;
                fnew = (ci_new + cj_new)*2;
            }

            if fold > max_deltaf + fnew {
                max_deltaf = fold - fnew;
                max_deltaf_j = j;
            }
        }

        if max_deltaf != 0 {
//            println!("queens_o={:?}, fitness_o={}, case={}", q, self.fitness, case);
            q.swap(i, max_deltaf_j);
            self.fitness -= max_deltaf;

//            // TEST DEBUG
//            let mut diag_counts = vec![0; 4*SIZE-2];
//            let (diag_counts1, diag_counts2) = diag_counts.split_at_mut(2*SIZE-1);
//            for x in 0..SIZE {
//                let d1 = SIZE-1 + x - q[x];
//                let d2 = x + q[x];
//                diag_counts1[d1] += 1;
//                diag_counts2[d2] += 1;
//            }
//
//            let mut f = 0;
//            for x in 0..SIZE {
//                let d1 = SIZE-1 + x - q[x];
//                let d2 = x + q[x];
//                f += diag_counts1[d1] + diag_counts2[d2];
//            }
//
//            assert!(f == self.fitness, "max_deltaf={}, f={}, fitness={}, queens={:?}", max_deltaf, f, self.fitness, q);
        }
    }

    fn fitness(&self) -> usize {
        self.fitness
    }

    #[inline(never)]
    fn mutate(new_gen: &mut [Self], rng: &mut XorShiftRng) {
        for board in new_gen.iter_mut() {
            if rng.next_f32() < MUTATION_PROBABILITY {
                board.do_mutate(rng);
            }
        }
    }

    #[inline(never)]
    fn breed(parents: &[Self], rng: &mut XorShiftRng) -> Self {
//        Self::breed_pmx(parents, rng)
        Self::breed_pmx_norng(parents, rng)
    }

    #[inline(never)]
    fn filter_strongest(species: &mut Vec<Self>) {
        if species.len() <= ELITE {
            return;
        }

        species.sort_by(|s, t| s.fitness.partial_cmp(&t.fitness).unwrap());
        species.truncate(ELITE);
    }

    #[inline(never)]
    fn initial() -> Vec<Self> {
        let population = cmp::max(NCHILDREN, NPARENTS);
        let mut boards = Vec::with_capacity(population);

        let mut ys = Vec::with_capacity(SIZE);

        let mut rng = RNG.lock().unwrap();
        for _ in 0..population {
            let mut board = Board::new(Vec::with_capacity(SIZE));

            for y in 0..SIZE {
                ys.push(y);
            }

            for _ in 0..SIZE {
                let z = rng.gen_range(0, ys.len());
                let y = ys.swap_remove(z);
                board.queens.push(y);
            }

            board.reevaluate(&mut rng);
            boards.push(board);
        }
        boards
    }
}

impl Debug for Board {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "fitness={}, queens=[{:?}]", self.fitness, self.queens)
    }
}

fn genetic_queens() {
    {
        let mut rng = RNG.lock().unwrap();
        let mut trng = thread_rng();
        rng.reseed([trng.next_u32(), trng.next_u32(), trng.next_u32(), trng.next_u32()]);
    }

    let best = genetic::<Board>(2*SIZE, MAX_ITERS, NPARENTS, NCHILDREN);
//    println!("best: {:?}", &best);
    println!("best: {:?}", best.fitness());
}

fn main() {
    for _ in 0..10 {
        genetic_queens()
    }
}

<|endoftext|>


<commit_before>#![feature(globs)]
#![allow(dead_code)]

extern crate libc;
extern crate collections;
use libc::funcs::posix88::unistd::*;
use std::os;
use std::mem;
use posix::*;
use std::c_str::CString;
use std::io::extensions::*;

mod ptrace;
mod posix;
mod syscall;

fn main () {
    let env = "/usr/bin/env";
    let mut c = vec!["env".to_string()];
    if os::args().tail().len() != 0 {
        c.push_all(os::args().tail());
    } else {
        c.push_all(["echo".to_string(), "No argumets".to_string()]);
    }
    unsafe {
        let result = fork_wrapper();
        match result {
            ForkChild =>  {
                let mut vec = Vec::with_capacity(c.len());
                let mut ptr_vec = Vec::with_capacity(c.len());

                for v in c.iter()
                {
                    let c_str = v.as_slice().to_c_str();
                    ptr_vec.push(c_str.as_ptr());
                    vec.push(c_str);
                }

                let t = ptr_vec.as_mut_ptr();
                ptrace::trace_me();
                execv(env.to_c_str().as_ptr(), t);
            }
            ForkFailure(_) => {
                println!("An error occurred: {}", result.get_error_as_string());
            }
            ForkParent(child_pid) => {
                let result = ptrace::init_trace(child_pid);

                if result.is_error() {
                    println!("An error in init_ptrace occurred: {}", result.get_error_as_string());
                }
                loop {
                    let result = posix::waitpid_wrapper(child_pid, 0);
                    match result {
                        WaitPidFailure(_) => {
                            println!("An error in wait occurred: {}", result.get_error_as_string());
                            break;
                        },
                        WaitPidSuccess(pid, status) => {
                            if ((status >> 8) & (0x80 | posix::SIGTRAP)) != 0 {
                                match ptrace::get_registers(pid) {
                                    Ok(ptrace::UserRegs { orig_rax: syscall_no, rdi: rdi, .. }) => {
                                            //println!("Syscall {} with {}", syscall::get_syscall_name_as_str(syscall_no as uint), handle_syscall_arguments(pid, rdi));
                                        if syscall_no == 2 {
                                            println!("{}", handle_syscall_arguments(pid, rdi));
                                        }
                                        ptrace::syscall(pid); //need this to get only at end
                                        posix::waitpid_wrapper(pid, 0); // and this too
                                    },
                                    Err(errno) => {
                                        println!("Error in ptrace {}", get_strerror(errno)); break;
                                    },
                                }
                            }
                            let result = ptrace::resume_trace(child_pid);
                            if result.is_error() {
                                println!("An error in ptrace_resume occurred: {}", result.get_error_as_string());
                                break;
                            }
                        },
                    }
                }

            }
        }
    }
}

fn pstrdup(word: u64) -> String {
    unsafe {
        let mut bytes    = Vec::new();
        //let mut mut_addr = addr as libc::uint64_t;

        let mut i = 0;

        // XXX I'm not using a for loop because of a bug in Rust
        while i < mem::size_of::<libc::uint64_t>() {
            // XXX byte order
            let lsb = (word >> (i * 8)) & 0xFF;
            if lsb == 0 {
                break;
            }
            bytes.push(lsb as u8);
            i += 1;
        }
        //mut_addr += mem::size_of::<libc::uint64_t>() as libc::uint64_t;
        let result = std::str::raw::from_utf8(bytes.as_slice()).to_string();
        result
            //TODO notes touch shows no open, > shows open
        //println!("Syscall {} with {}", syscall::get_syscall_name_as_str(syscall_no as uint), CString::new(rdi as *const i8, false).as_str()); //TODO
    }
}

fn get_program_args(pid: int, addr: *mut libc::c_void) -> String {
    unsafe {
    let mut mut_addr = addr as libc::uint64_t;
    let mut bytes = Vec::new();

    'j: loop {
        match ptrace::peektext(pid, mut_addr as *mut libc::c_void) {
            Err(_) | Ok(0) => break,
            Ok(word)       => {
                println!("{}", bytes);
                u64_to_be_bytes(word, 8, |v| bytes.push_all(v) ); // {
                //for b in
                    //if b == 0 {
                        //break 'j;
                    //}
                    //bytes.push(b);
                //}
            }
        }

    }
    let result = std::str::raw::from_utf8(bytes.as_slice()).to_string();
    result
}
}

fn handle_syscall_arguments(pid: int, argv_ptr: libc::uint64_t) -> String {
    let argv = get_program_args(pid, argv_ptr as *mut libc::c_void);
    argv
}
<commit_msg>add get args<commit_after>#![feature(globs)]
#![allow(dead_code)]

extern crate libc;
extern crate collections;
use libc::funcs::posix88::unistd::*;
use std::os;
use std::mem;
use posix::*;
use std::io::extensions::*;

mod ptrace;
mod posix;
mod syscall;

fn main () {
    let env = "/usr/bin/env";
    let mut c = vec!["env".to_string()];
    if os::args().tail().len() != 0 {
        c.push_all(os::args().tail());
    } else {
        c.push_all(["echo".to_string(), "No argumets".to_string()]);
    }
    unsafe {
        let result = fork_wrapper();
        match result {
            ForkChild =>  {
                let mut vec = Vec::with_capacity(c.len());
                let mut ptr_vec = Vec::with_capacity(c.len());

                for v in c.iter()
                {
                    let c_str = v.as_slice().to_c_str();
                    ptr_vec.push(c_str.as_ptr());
                    vec.push(c_str);
                }

                let t = ptr_vec.as_mut_ptr();
                ptrace::trace_me();
                execv(env.to_c_str().as_ptr(), t);
            }
            ForkFailure(_) => {
                println!("An error occurred: {}", result.get_error_as_string());
            }
            ForkParent(child_pid) => {
                let result = ptrace::init_trace(child_pid);

                if result.is_error() {
                    println!("An error in init_ptrace occurred: {}", result.get_error_as_string());
                }
                loop {
                    let result = posix::waitpid_wrapper(child_pid, 0);
                    match result {
                        WaitPidFailure(_) => {
                            println!("An error in wait occurred: {}", result.get_error_as_string());
                            break;
                        },
                        WaitPidSuccess(pid, status) => {
                            if ((status >> 8) & (0x80 | posix::SIGTRAP)) != 0 {
                                match ptrace::get_registers(pid) {
                                    Ok(ptrace::UserRegs { orig_rax: syscall_no, rdi: rdi, .. }) => {
                                        if syscall_no == 2 {
                                            println!("{}", handle_syscall_arguments(pid, rdi));
                                        }
                                        ptrace::syscall(pid); //need this to get only at end
                                        posix::waitpid_wrapper(pid, 0); // and this too
                                    },
                                    Err(errno) => {
                                        println!("Error in ptrace {}", get_strerror(errno)); break;
                                    },
                                }
                            }
                            let result = ptrace::resume_trace(child_pid);
                            if result.is_error() {
                                println!("An error in ptrace_resume occurred: {}", result.get_error_as_string());
                                break;
                            }
                        },
                    }
                }

            }
        }
    }
}

fn get_program_args(pid: int, addr: *mut libc::c_void) -> String {
    unsafe {
        let mut mut_addr = addr as libc::uint64_t;
        let mut bytes = Vec::new();
        'j: loop {
            match ptrace::peektext(pid, mut_addr as *mut libc::c_void) {
                Err(_) | Ok(0) => break 'j,
                Ok(word)       => {
                    u64_to_le_bytes(word, 8, |v| bytes.push_all(v));
                    for &b in bytes.iter() {
                        if b == 0u8 {
                            break 'j;
                        }
                    }
                }
            }
            mut_addr += mem::size_of::<libc::uint64_t>() as libc::uint64_t;
        }
        std::str::raw::from_utf8(bytes.as_slice()).to_string()
    }
}

fn handle_syscall_arguments(pid: int, argv_ptr: libc::uint64_t) -> String {
    get_program_args(pid, argv_ptr as *mut libc::c_void)
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate flate2;

mod padding;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::u8;
use crypto::digest::Digest;
use crypto::sha1;

const SHA1_BYTE_LENGTH: usize = 20;

struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

struct SearchParams<'a> {
    current_message: &'a str,
    desired_prefix: HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    let search_params = SearchParams {
        current_message,
        desired_prefix,
        counter_range: ops::Range { start: 0, end: std::u64::MAX },
        extension_word_length: 8
    };

    match search_for_match(&search_params) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn search_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = &params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}
<commit_msg>Move logic for iteration parameters into a separate function<commit_after>extern crate crypto;
extern crate flate2;

mod padding;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::u8;
use crypto::digest::Digest;
use crypto::sha1;

const SHA1_BYTE_LENGTH: usize = 20;

struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

struct SearchParams<'a> {
    current_message: &'a str,
    desired_prefix: HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    match find_match(current_message, desired_prefix) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn find_match(current_message: &str, desired_prefix: HashPrefix) -> Option<HashMatch> {
    let search_params = &SearchParams {
        current_message,
        desired_prefix,
        counter_range: ops::Range { start: 0, end: std::u64::MAX },
        extension_word_length: 8
    };

    iterate_for_match(search_params)
}

fn iterate_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = &params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}
<|endoftext|>


<commit_before>// Author: Alex Chernyakhovsky (achernya@mit.edu)

// TODO(achernya): Per the documentation on doc.rust-lang.org, std::io
// is not yet ready, so use old_io until 1.0-final.
#![feature(old_io)]

// std::env contains a lot of nice functions that otherwise would
// require std::os to use; std::os has lots of deprecated functions.
#![feature(env)]

// TODO(achernya): Remove this feature when std::env moves over to
// std::path.
#![feature(old_path)]

use std::env;
use std::old_io as io;

// println_stderr is like println, but to stderr.
macro_rules! println_stderr(
    ($($arg:tt)*) => (
        match writeln!(&mut io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
    )
);

// ShellCommand is a trait that defines a runnable POSIX shell
// command. An implementation is an abstract representation of shell
// commands such as simple invocations, invocations with redirection,
// and even shell pipelines.
trait ShellCommand {
    fn run(&self);
}

fn shell_loop() {
    let mut stdin = io::stdin();
    loop {
        print!("$ ");
        let line = stdin.read_line();
        match line {
            Ok(expr) => handle_command(expr),
            Err(_) => break,
        }
    }
}

fn handle_command(user_expr: String) {
    // Clean up the string by removing the newline at the end
    let expr = &user_expr.trim_matches('\n');
    let components: Vec<&str> = expr.split(' ').collect();
    if builtins(components) {
        return;
    }
}

fn builtins(command: Vec<&str>) -> bool {
    match command[0] {
        "cd" => cd(command),
        "pwd" => pwd(),
        _ => return false,
    }
    return true;
}

fn cd(command: Vec<&str>) {
    // cd is the "change directory" command. It can take either 0 or 1
    // arguments. If given no arguments, then the $HOME directory is
    // chosen.
    let dir: Option<Path> = match command.len() {
        0 => panic!("invalid cd invocation"),
        1 => env::home_dir(),
        _ => Some(Path::new(command[1]))
    };
    if dir.is_none() {
        println_stderr!("cd: no directory to change to");
        return;
    }
    let directory = dir.unwrap();
    let result = env::set_current_dir(&directory);
    match result {
        Err(err) => {
            println_stderr!("cd: {}: {}", directory.display(), err);
        },
        _ => (),
    }
}

fn pwd() {
    let p = env::current_dir().unwrap_or(Path::new("/"));
    println!("{}", p.display());
}

fn main() {
    // TODO(achernya): is there any initialization we want to do
    // before we enter the shell loop?
    shell_loop();
}
<commit_msg>Clean up some style issues found while reading Rust Book<commit_after>// Author: Alex Chernyakhovsky (achernya@mit.edu)

// TODO(achernya): Per the documentation on doc.rust-lang.org, std::io
// is not yet ready, so use old_io until 1.0-final.
#![feature(old_io)]

// std::env contains a lot of nice functions that otherwise would
// require std::os to use; std::os has lots of deprecated functions.
#![feature(env)]

// TODO(achernya): Remove this feature when std::env moves over to
// std::path.
#![feature(old_path)]

use std::env;
use std::old_io as io;

// println_stderr is like println, but to stderr.
macro_rules! println_stderr(
    ($($arg:tt)*) => (
        match writeln!(&mut io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
    )
);

// ShellCommand is a trait that defines a runnable POSIX shell
// command. An implementation is an abstract representation of shell
// commands such as simple invocations, invocations with redirection,
// and even shell pipelines.
trait ShellCommand {
    fn run(&self);
}

fn shell_loop() {
    let mut stdin = io::stdin();
    loop {
        print!("$ ");
        let line = stdin.read_line();
        match line {
            Ok(expr) => handle_command(&expr),
            Err(_) => break,
        }
    }
}

fn handle_command(user_expr: &str) {
    // Clean up the string by removing the newline at the end
    let expr = user_expr.trim_matches('\n');
    let components: Vec<&str> = expr.split(' ').collect();
    if builtins(&components) {
        return;
    }
}

fn builtins(command: &Vec<&str>) -> bool {
    match command[0] {
        "cd" => cd(command),
        "pwd" => pwd(),
        _ => return false,
    }
    true
}

fn cd(command: &Vec<&str>) {
    // cd is the "change directory" command. It can take either 0 or 1
    // arguments. If given no arguments, then the $HOME directory is
    // chosen.
    let dir: Option<Path> = match command.len() {
        0 => panic!("invalid cd invocation"),
        1 => env::home_dir(),
        _ => Some(Path::new(command[1]))
    };
    if dir.is_none() {
        println_stderr!("cd: no directory to change to");
        return;
    }
    let directory = dir.unwrap();
    let result = env::set_current_dir(&directory);
    match result {
        Err(err) => {
            println_stderr!("cd: {}: {}", directory.display(), err);
        },
        _ => {},
    }
}

fn pwd() {
    let p = env::current_dir().unwrap_or(Path::new("/"));
    println!("{}", p.display());
}

fn main() {
    // TODO(achernya): is there any initialization we want to do
    // before we enter the shell loop?
    shell_loop();
}
<|endoftext|>


<commit_before>// Copyright © 2016 Bart Massey
// [This program is licensed under the "MIT License"]
// Please see the file COPYING in the source
// distribution of this software for license terms.

// Rocket pastebin demo

#![feature(plugin)]
#![feature(plugin, custom_derive)]
#![feature(proc_macro)]
#![plugin(rocket_codegen)]

static UPLOAD: &'static str = "upload";
static NID: usize = 8;

extern crate rocket;
extern crate rocket_contrib;
extern crate rand;
extern crate serde;
#[macro_use] extern crate serde_derive;

use std::io;
use std::io::{Read, BufRead, Write};
use std::path::Path;
use std::collections::HashMap;
use std::fs::{File, create_dir, metadata, remove_file};

use rocket::Data;
use rocket::response::NamedFile;
use rocket_contrib::Template;
use rocket::request::Form;
use rocket::response::Redirect;

mod paste_id;
use paste_id::PasteID;

#[derive(Serialize)]
struct PasteInfo {
    paste_id: String,
    snippet: String,
    ellipsis: bool
}

#[get("/")]
fn index() -> io::Result<Template> {
    let dirents = std::fs::read_dir(UPLOAD)
                  .expect("could not open upload directory");
    let pastes = dirents.map(|d| {
        let paste_id = d.unwrap().file_name().into_string().unwrap();
        let paste_filename = Path::new(UPLOAD).join(&paste_id);
        let paste_file = NamedFile::open(&paste_filename)
                         .expect("could not open paste file");
        let mut paste_reader = io::BufReader::new(paste_file);
        let mut paste_line = String::new();
        let paste_len = paste_reader.read_line(&mut paste_line)
                        .expect("could not read first line of paste");
        let (ellipsis, snippet) =
            if paste_len < 24 {
                (false, paste_line)
            } else {
                let result = (&paste_line)[0..24].to_string();
                (true, result)
            };
        PasteInfo { paste_id: paste_id, snippet: snippet, ellipsis: ellipsis }
    }).collect::<Vec<PasteInfo>>();
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("pastes", &pastes);
    Ok(Template::render("index", &context))
}

#[post("/", data = "<paste>")]
fn upload(paste: Data) -> io::Result<String> {
    let paste_id = PasteID::new(NID);
    let path = Path::new(UPLOAD).join(&paste_id.to_string());
    paste.stream_to_file(path)?;
    Ok(format!("/{}\n", paste_id))
}

fn open_paste(id: &PasteID) -> io::Result<NamedFile> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    NamedFile::open(&filename)
}

fn create_paste(id: &PasteID) -> io::Result<File> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    File::create(&filename)
}

#[get("/<id>")]
fn retrieve(id: PasteID) -> Option<NamedFile> {
    open_paste(&id).ok()
}

#[get("/new")]
fn edit_new() -> io::Result<Redirect> {
    let paste_id = PasteID::new(NID);
    let f = try!(create_paste(&paste_id));
    drop(f);
    Ok(Redirect::to(&format!("/edit/{}", paste_id.to_string())))
}

#[get("/edit/<id>")]
fn make_edit(id: PasteID) -> io::Result<Template> {
    let paste = open_paste(&id)?;
    let mut paste_reader = io::BufReader::new(paste);
    let mut paste_contents = String::new();
    paste_reader.read_to_string(&mut paste_contents)?;
    let paste_id = id.to_string();
    let paste_rows = format!("{}", 100);
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("contents", &paste_contents);
    context.insert("paste_id", &paste_id);
    context.insert("paste_rows", &paste_rows);
    Ok(Template::render("edit", &context))
}

#[derive(FromForm)]
struct EditForm {
    paste: String
}

#[post("/edit/<id>", data = "<form>")]
fn accept_edit(id: PasteID, form: Form<EditForm>) -> io::Result<Redirect> {
    let edit = form.get();
    let paste = create_paste(&id)?;
    let mut paste_writer = io::BufWriter::new(paste);
    paste_writer.write_all(edit.paste.as_bytes())?;
    paste_writer.flush()?;
    Ok(Redirect::to("/"))
}

#[get("/delete/<id>")]
fn delete(id: PasteID) -> io::Result<Redirect> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    try!(remove_file(&filename));
    Ok(Redirect::to("/"))
}

fn ok_dir<P: AsRef<Path>>(path: P) -> bool {
    match metadata(path.as_ref()) {
        Ok(md) => md.is_dir(),
        _ => false
    }
}

fn check_upload_dir() {
    if !ok_dir(UPLOAD) {
        create_dir(UPLOAD).expect("could not create upload directory");
    }
}

fn main() {
    check_upload_dir();
    let rocket = rocket::ignite();
    let rocket = rocket.mount("/", routes![index, upload, edit_new, retrieve,
                                           make_edit, accept_edit, delete]);
    rocket.launch()
}

#[cfg(test)]
mod test {
    extern crate scraper;

    use super::rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{Status, Method};
    use self::scraper::{Html, Selector};

    #[test]
    fn index() {
        let rocket = rocket::ignite().mount("/", routes![super::index]);

        let mut req = MockRequest::new(Method::Get, "/");
        let mut response = req.dispatch_with(&rocket);

        assert_eq!(response.status(), Status::Ok);

        let body_str = response.body().and_then(|b| b.into_string()).expect("Missing response body!");
        assert_has_text(&body_str, "Pastebin!");
        assert_has_selector(&body_str, "ul");
        assert_no_selector(&body_str, "ul li");
    }

    fn assert_has_text(body: &str, expected: &str) {
        let body_html = Html::parse_fragment(body);
        let selector = Selector::parse("*").unwrap();
        let nodes = body_html.select(&selector);
        let texts = nodes.map(|node| node.text().collect::<Vec<_>>().concat());
        let actual = texts.collect::<Vec<_>>().concat();
        assert!(&actual.contains(expected));
    }

    fn assert_has_selector(body: &str, selector_text: &str) {
        let body_html = Html::parse_fragment(body);
        let selector = Selector::parse(selector_text).expect("Invalid selector!");
        let count = body_html.select(&selector).count();
        assert!(count > 0);
    }

    fn assert_no_selector(body: &str, selector_text: &str) {
        let body_html = Html::parse_fragment(body);
        let selector = Selector::parse(selector_text).expect("Invalid selector!");
        let count = body_html.select(&selector).count();
        assert_eq!(0, count);
    }
}
<commit_msg>refactored test assertions some<commit_after>// Copyright © 2016 Bart Massey
// [This program is licensed under the "MIT License"]
// Please see the file COPYING in the source
// distribution of this software for license terms.

// Rocket pastebin demo

#![feature(plugin)]
#![feature(plugin, custom_derive)]
#![feature(proc_macro)]
#![plugin(rocket_codegen)]

static UPLOAD: &'static str = "upload";
static NID: usize = 8;

extern crate rocket;
extern crate rocket_contrib;
extern crate rand;
extern crate serde;
#[macro_use] extern crate serde_derive;

use std::io;
use std::io::{Read, BufRead, Write};
use std::path::Path;
use std::collections::HashMap;
use std::fs::{File, create_dir, metadata, remove_file};

use rocket::Data;
use rocket::response::NamedFile;
use rocket_contrib::Template;
use rocket::request::Form;
use rocket::response::Redirect;

mod paste_id;
use paste_id::PasteID;

#[derive(Serialize)]
struct PasteInfo {
    paste_id: String,
    snippet: String,
    ellipsis: bool
}

#[get("/")]
fn index() -> io::Result<Template> {
    let dirents = std::fs::read_dir(UPLOAD)
                  .expect("could not open upload directory");
    let pastes = dirents.map(|d| {
        let paste_id = d.unwrap().file_name().into_string().unwrap();
        let paste_filename = Path::new(UPLOAD).join(&paste_id);
        let paste_file = NamedFile::open(&paste_filename)
                         .expect("could not open paste file");
        let mut paste_reader = io::BufReader::new(paste_file);
        let mut paste_line = String::new();
        let paste_len = paste_reader.read_line(&mut paste_line)
                        .expect("could not read first line of paste");
        let (ellipsis, snippet) =
            if paste_len < 24 {
                (false, paste_line)
            } else {
                let result = (&paste_line)[0..24].to_string();
                (true, result)
            };
        PasteInfo { paste_id: paste_id, snippet: snippet, ellipsis: ellipsis }
    }).collect::<Vec<PasteInfo>>();
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("pastes", &pastes);
    Ok(Template::render("index", &context))
}

#[post("/", data = "<paste>")]
fn upload(paste: Data) -> io::Result<String> {
    let paste_id = PasteID::new(NID);
    let path = Path::new(UPLOAD).join(&paste_id.to_string());
    paste.stream_to_file(path)?;
    Ok(format!("/{}\n", paste_id))
}

fn open_paste(id: &PasteID) -> io::Result<NamedFile> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    NamedFile::open(&filename)
}

fn create_paste(id: &PasteID) -> io::Result<File> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    File::create(&filename)
}

#[get("/<id>")]
fn retrieve(id: PasteID) -> Option<NamedFile> {
    open_paste(&id).ok()
}

#[get("/new")]
fn edit_new() -> io::Result<Redirect> {
    let paste_id = PasteID::new(NID);
    let f = try!(create_paste(&paste_id));
    drop(f);
    Ok(Redirect::to(&format!("/edit/{}", paste_id.to_string())))
}

#[get("/edit/<id>")]
fn make_edit(id: PasteID) -> io::Result<Template> {
    let paste = open_paste(&id)?;
    let mut paste_reader = io::BufReader::new(paste);
    let mut paste_contents = String::new();
    paste_reader.read_to_string(&mut paste_contents)?;
    let paste_id = id.to_string();
    let paste_rows = format!("{}", 100);
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("contents", &paste_contents);
    context.insert("paste_id", &paste_id);
    context.insert("paste_rows", &paste_rows);
    Ok(Template::render("edit", &context))
}

#[derive(FromForm)]
struct EditForm {
    paste: String
}

#[post("/edit/<id>", data = "<form>")]
fn accept_edit(id: PasteID, form: Form<EditForm>) -> io::Result<Redirect> {
    let edit = form.get();
    let paste = create_paste(&id)?;
    let mut paste_writer = io::BufWriter::new(paste);
    paste_writer.write_all(edit.paste.as_bytes())?;
    paste_writer.flush()?;
    Ok(Redirect::to("/"))
}

#[get("/delete/<id>")]
fn delete(id: PasteID) -> io::Result<Redirect> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    try!(remove_file(&filename));
    Ok(Redirect::to("/"))
}

fn ok_dir<P: AsRef<Path>>(path: P) -> bool {
    match metadata(path.as_ref()) {
        Ok(md) => md.is_dir(),
        _ => false
    }
}

fn check_upload_dir() {
    if !ok_dir(UPLOAD) {
        create_dir(UPLOAD).expect("could not create upload directory");
    }
}

fn main() {
    check_upload_dir();
    let rocket = rocket::ignite();
    let rocket = rocket.mount("/", routes![index, upload, edit_new, retrieve,
                                           make_edit, accept_edit, delete]);
    rocket.launch()
}

#[cfg(test)]
mod test {
    extern crate scraper;

    use super::rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{Status, Method};
    use self::scraper::{Html, Selector};

    #[test]
    fn index() {
        let rocket = rocket::ignite().mount("/", routes![super::index]);

        let mut req = MockRequest::new(Method::Get, "/");
        let mut response = req.dispatch_with(&rocket);

        assert_eq!(response.status(), Status::Ok);

        let body_str = response.body()
                       .and_then(|b| b.into_string())
                       .expect("Missing response body!");
        let body_html = Html::parse_fragment(&body_str);
        assert_has_text(&body_html, "Pastebin!");
        assert_has_selector(&body_html, "ul");
        assert_no_selector(&body_html, "ul li");
    }

    fn assert_has_text(body: &Html, expected: &str) {
        let selector = Selector::parse("*").unwrap();
        let nodes = body.select(&selector);
        let texts = nodes.map(|node| node.text().collect::<Vec<_>>().concat());
        let actual = texts.collect::<Vec<_>>().concat();
        assert!(&actual.contains(expected));
    }

    fn count_selectors(body: &Html, selector_text: &str) -> usize {
        let selector = Selector::parse(selector_text)
                       .expect("Invalid selector!");
        body.select(&selector).count()
    }

    fn assert_has_selector(body: &Html, selector_text: &str) {
        let count = count_selectors(body, selector_text);
        assert!(count > 0);
    }

    fn assert_no_selector(body: &Html, selector_text: &str) {
        let count = count_selectors(body, selector_text);
        assert_eq!(0, count);
    }
}
<|endoftext|>


<commit_before>extern crate config;
extern crate clap;
extern crate hyper;
extern crate rustc_serialize;
extern crate url;
#[macro_use]
extern crate error_chain;
#[cfg(feature = "clipboard_support")]
extern crate clipboard;

mod bins;

use bins::error::*;
use bins::Bins;
use bins::arguments;
use bins::configuration::{BinsConfiguration, Configurable};
use std::io::Write;
#[cfg(feature = "clipboard_support")] use clipboard::ClipboardContext;

macro_rules! println_stderr {
  ($fmt:expr) => { { writeln!(std::io::stderr(), $fmt).expect("error writing to stderr"); } };
  ($fmt:expr, $($arg:tt)*) => { { writeln!(std::io::stderr(), $fmt, $($arg)*).expect("error writing to stderr"); } };
}

macro_rules! or_exit {
  ($expr: expr) => { match $expr { Ok(x) => x, Err(e) => { for err in e.iter() { println_stderr!("{}", err); } return 1; } } };
}

fn make_bins() -> Result<Bins> {
  let configuration = BinsConfiguration::new();
  let config = try!(configuration.parse_config());
  let arguments = arguments::get_arguments(&config);
  Ok(Bins::new(config, arguments))
}

#[cfg(feature = "clipboard_support")]
fn copy_to_clipboard(ref data: &String) -> Result<()> {
  let mut clipboard = try!(ClipboardContext::new().map_err(|e| e.to_string()));
  clipboard.set_contents(data.clone().to_owned()).map_err(|e| e.to_string().into())
}

#[cfg(not(feature = "clipboard_support"))]
fn copy_to_clipboard(_: &String) -> Result<()> {
  Ok(())
}

fn inner() -> i32 {
  let bins = or_exit!(make_bins());
  let output = or_exit!(bins.get_output());
  if bins.arguments.copy {
    or_exit!(copy_to_clipboard(&output));
  }
  println!("{}", output);
  0
}

fn main() {
  let exit_code = inner();
  std::process::exit(exit_code);
}
<commit_msg>refactor: make changes according to clippy suggestions<commit_after>extern crate config;
extern crate clap;
extern crate hyper;
extern crate rustc_serialize;
extern crate url;
#[macro_use]
extern crate error_chain;
#[cfg(feature = "clipboard_support")]
extern crate clipboard;

mod bins;

use bins::error::*;
use bins::Bins;
use bins::arguments;
use bins::configuration::{BinsConfiguration, Configurable};
use std::io::Write;
#[cfg(feature = "clipboard_support")] use clipboard::ClipboardContext;

macro_rules! println_stderr {
  ($fmt:expr) => { { writeln!(std::io::stderr(), $fmt).expect("error writing to stderr"); } };
  ($fmt:expr, $($arg:tt)*) => { { writeln!(std::io::stderr(), $fmt, $($arg)*).expect("error writing to stderr"); } };
}

macro_rules! or_exit {
  ($expr: expr) => { match $expr { Ok(x) => x, Err(e) => { for err in e.iter() { println_stderr!("{}", err); } return 1; } } };
}

fn make_bins() -> Result<Bins> {
  let configuration = BinsConfiguration::new();
  let config = try!(configuration.parse_config());
  let arguments = arguments::get_arguments(&config);
  Ok(Bins::new(config, arguments))
}

#[cfg(feature = "clipboard_support")]
fn copy_to_clipboard(data: &str) -> Result<()> {
  let mut clipboard = try!(ClipboardContext::new().map_err(|e| e.to_string()));
  clipboard.set_contents((*data).to_owned()).map_err(|e| e.to_string().into())
}

#[cfg(not(feature = "clipboard_support"))]
fn copy_to_clipboard(_: &String) -> Result<()> {
  Ok(())
}

fn inner() -> i32 {
  let bins = or_exit!(make_bins());
  let output = or_exit!(bins.get_output());
  if bins.arguments.copy {
    or_exit!(copy_to_clipboard(&output));
  }
  println!("{}", output);
  0
}

fn main() {
  let exit_code = inner();
  std::process::exit(exit_code);
}
<|endoftext|>