<commit_before>#[macro_use]
extern crate lazy_static;

use backtrace::Backtrace;
use constants::*;
use git_version::git_version;
use isatty::{stdin_isatty, stdout_isatty};
use refiner::Refiner;
use regex::Regex;
use std::env;
use std::fs;
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::panic;
use std::process::exit;
use std::process::{Command, Stdio};
use std::str;

mod constants;
mod refiner;
mod tokenizer;

const HELP_TEXT: &str = r#"
Usage: diff ... | riff

Colors diff and highlights what parts of changed lines have changed.

Git integration:
    git config --global pager.diff riff
    git config --global pager.show riff
    git config --global interactive.filter riff
"#;

const HELP_TEXT_FOOTER: &str = r#"
Report issues at <https://github.com/walles/riff>.
"#;

const CRASH_FOOTER: &str = r#"
Please copy all of the above up to the --- RIFF CRASHED --- marker and report it at one of:
* <https://github.com/walles/riff/issues> (preferred)
* <johan.walles@gmail.com>
"#;

const HUNK_HEADER: &str = "\x1b[36m"; // Cyan
const PAGER_FORKBOMB_STOP: &str = "_RIFF_IGNORE_PAGER";

const GIT_VERSION: &str = git_version!();

lazy_static! {
    static ref STATIC_HEADERS: Vec<(Regex, &'static str)> = vec![
        (Regex::new("^diff ").unwrap(), BOLD),
        (Regex::new("^index ").unwrap(), BOLD),
        (Regex::new("^--- ").unwrap(), BOLD),
        (Regex::new("^\\+\\+\\+ ").unwrap(), BOLD),
        (Regex::new("^@@ ").unwrap(), HUNK_HEADER),
    ];
    static ref ANSI_COLOR_REGEX: Regex = Regex::new("\x1b[^m]*m").unwrap();
}

enum LastLineKind {
    Initial,
    Add,
    Remove,
}

#[must_use]
fn get_fixed_highlight(line: &str) -> &str {
    for static_header in STATIC_HEADERS.iter() {
        let re = &static_header.0;
        if re.is_match(line) {
            return static_header.1;
        }
    }

    return "";
}

fn print(stream: &mut BufWriter<&mut dyn Write>, text: &str) {
    stream.write(text.as_bytes()).unwrap();
}

fn println(stream: &mut BufWriter<&mut dyn Write>, text: &str) {
    print(stream, &text);
    print(stream, "\n");
}

fn highlight_diff(input: &mut dyn io::Read, output: &mut dyn io::Write) {
    let mut adds = String::new();
    let mut removes = String::new();
    let input = BufReader::new(input);
    let output = &mut BufWriter::new(output);
    let mut last_line_kind = LastLineKind::Initial;
    for line in input.lines() {
        let line = line.unwrap();

        // Strip out incoming ANSI formatting. This enables us to highlight
        // already-colored input.
        let line = ANSI_COLOR_REGEX.replace_all(&line, "");

        let fixed_highlight = get_fixed_highlight(&line);
        if !fixed_highlight.is_empty() {
            // Drain outstanding adds / removes
            for line in Refiner::create(&adds, &removes).format() {
                println(output, &line);
            }
            adds.clear();
            removes.clear();

            print(output, fixed_highlight);
            print(output, &line);
            println(output, NORMAL);
            continue;
        }

        // Collect adds / removes
        if !line.is_empty() && line.chars().next().unwrap() == '+' {
            adds.push_str(&line[1..]);
            adds.push('\n');
            last_line_kind = LastLineKind::Add;
            continue;
        } else if !line.is_empty() && line.chars().next().unwrap() == '-' {
            removes.push_str(&line[1..]);
            removes.push('\n');
            last_line_kind = LastLineKind::Remove;
            continue;
        }

        if line == NO_EOF_NEWLINE_MARKER {
            match last_line_kind {
                LastLineKind::Add => {
                    assert!(adds.pop().unwrap() == '\n');
                    continue;
                }
                LastLineKind::Remove => {
                    assert!(removes.pop().unwrap() == '\n');
                    continue;
                }
                LastLineKind::Initial => {
                    // This block intentionally left blank
                }
            }
        }

        last_line_kind = LastLineKind::Initial;

        // Drain outstanding adds / removes
        for line in Refiner::create(&adds, &removes).format() {
            println(output, &line);
        }
        adds.clear();
        removes.clear();

        // Print current line
        if line == NO_EOF_NEWLINE_MARKER {
            print(output, NO_EOF_NEWLINE_COLOR);
            print(output, &line);
            println(output, NORMAL);
        } else {
            println(output, &line);
        }
    }
    for line in Refiner::create(&adds, &removes).format() {
        println(output, &line);
    }
}

/// Try paging using the named pager (`$PATH` will be searched).
///
/// Returns `true` if the pager was found, `false` otherwise.
#[must_use]
fn try_pager(pager_name: &str) -> bool {
    let mut command = Command::new(pager_name);

    if env::var(PAGER_FORKBOMB_STOP).is_ok() {
        // Try preventing fork bombing if $PAGER is set to riff
        return false;
    }
    command.env(PAGER_FORKBOMB_STOP, "1");

    if !env::var("LESS").is_ok() {
        // Set by git when paging
        command.env("LESS", "FRX");
    }

    if !env::var("LV").is_ok() {
        // Set by git when paging
        command.env("LV", "-c");
    }

    command.stdin(Stdio::piped());

    match command.spawn() {
        Ok(mut pager) => {
            let pager_stdin = pager.stdin.as_mut().unwrap();
            highlight_diff(&mut io::stdin().lock(), pager_stdin);

            // FIXME: Report pager exit status if non-zero, together with
            // contents of pager stderr as well if possible.
            pager.wait().expect("Waiting for pager failed");

            return true;
        }
        Err(_) => {
            return false;
        }
    }
}

/// If `option` is found in `argv`, all instances of `option` will be removed
/// from `argv`.
///
/// Returns `true` if `option` was found and consumed, false otherwise.
#[must_use]
fn consume(option: &str, argv: &mut Vec<String>) -> bool {
    if !argv.contains(&option.to_string()) {
        // Not found
        return false;
    }

    argv.retain(|x| x != option);
    return true;
}

fn print_help(output: &mut dyn io::Write) {
    output.write(HELP_TEXT.trim().as_bytes()).unwrap();
    output.write(b"\n").unwrap();
    output.write(b"\n").unwrap();

    let self_path = env::args().into_iter().next().unwrap();
    let self_path = fs::canonicalize(self_path).unwrap();
    let self_path = self_path.as_path().to_str().unwrap();

    // FIXME: Do this only if we aren't already in the $PATH
    output
        .write("Installing riff in the $PATH:\n".as_bytes())
        .unwrap();
    output
        .write(&format!("  * sudo cp {} /usr/local/bin\n", self_path).as_bytes())
        .unwrap();
    output.write(b"\n").unwrap();

    output.write(HELP_TEXT_FOOTER.trim().as_bytes()).unwrap();
    output.write(b"\n").unwrap();
}

fn panic_handler(panic_info: &panic::PanicInfo) {
    let stderr: &mut dyn Write = &mut io::stderr();
    let stderr = &mut BufWriter::new(stderr);
    println(stderr, "\n\n------------ RIFF CRASHED -------------------");

    // Panic message
    if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
        println(stderr, &format!("Panic message: <{:?}>", s));
        println(stderr, "");
    }

    // Backtrace
    // FIXME: Ditch the Backtrace-internal frames from this backtrace
    // FIXME: Ditch the panic internal frames at the end of the backtrace
    println(stderr, &format!("{:?}", Backtrace::new()));

    println(stderr, &format!("Riff version: {}", GIT_VERSION));

    println(stderr, CRASH_FOOTER);
}

fn main() {
    panic::set_hook(Box::new(|panic_info: &panic::PanicInfo| {
        panic_handler(panic_info);
    }));

    let mut args: Vec<String> = env::args().collect();
    if consume("--help", &mut args) {
        print_help(&mut io::stdout());
        return;
    }

    if consume("--version", &mut args) {
        println!("riff {}", GIT_VERSION);
        println!("");
        println!("Source code available at <https://github.com/walles/riff>.");
        return;
    }

    if stdin_isatty() {
        eprintln!("ERROR: Expected input from a pipe");
        eprintln!("");
        print_help(&mut io::stderr());
        exit(1);
    }

    if !stdout_isatty() {
        // We're being piped, just do stdin -> stdout
        highlight_diff(&mut io::stdin().lock(), &mut io::stdout());
        return;
    }

    let pager_env_var = env::var("PAGER");
    if pager_env_var.is_ok() {
        let pager_value = pager_env_var.unwrap();
        if try_pager(&pager_value) {
            return;
        }

        // FIXME: Print warning at the end if $PAGER was set to something that
        // doesn't exist.
    }

    if try_pager("moar") {
        return;
    }

    if try_pager("less") {
        return;
    }

    // No pager found, wth?
    highlight_diff(&mut io::stdin().lock(), &mut io::stdout());
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    fn remove(text: &str) -> String {
        return format!("{}{}{}", REMOVE, text, NORMAL);
    }

    fn add(text: &str) -> String {
        return format!("{}{}{}", ADD, text, NORMAL);
    }

    #[test]
    fn test_remove_trailing_newline() {
        let mut input = "-hej\n\
            +hej\n\
            \\ No newline at end of file\n\
            "
        .as_bytes();

        let expected = format!(
            "{}\n{}\n{}\n",
            remove(&format!("-hej{}⏎", INVERSE_VIDEO)),
            add("+hej"),
            format!(
                "{}\\ No newline at end of file{}",
                NO_EOF_NEWLINE_COLOR, NORMAL
            )
        );

        let mut output: Vec<u8> = Vec::new();
        highlight_diff(&mut input, &mut output);
        assert_eq!(std::str::from_utf8(&output).unwrap(), expected);
    }

    #[test]
    fn test_trailing_newline_context() {
        let mut input = "+bepa\n apa\n\\ No newline at end of file\n".as_bytes();

        let expected = format!(
            "{}\n{}\n{}\n",
            add("+bepa"),
            " apa",
            format!(
                "{}\\ No newline at end of file{}",
                NO_EOF_NEWLINE_COLOR, NORMAL
            )
        );

        let mut output: Vec<u8> = Vec::new();
        highlight_diff(&mut input, &mut output);
        assert_eq!(std::str::from_utf8(&output).unwrap(), expected);
    }

    #[test]
    fn test_testdata_examples() {
        // Example value: `/Users/johan/src/riff/target/debug/deps/riff-7a8916c06b0d3d6c`
        let exe_path = std::env::current_exe().unwrap();

        // Example value: `/Users/johan/src/riff`
        let project_path = exe_path
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .parent()
            .unwrap();

        // Example value: `/Users/johan/src/riff/testdata`
        let testdata_path = &project_path.join("testdata");
        assert!(testdata_path.is_dir());

        // Iterate all files in there
        for diff in fs::read_dir(testdata_path).unwrap() {
            let diff = diff.unwrap();
            let diff = diff.path();
            let diff = diff.as_path();
            if !diff.is_file() {
                continue;
            }

            if diff.extension().unwrap() != "diff" {
                continue;
            }

            println!("Evaluating example file <{}>...", diff.to_str().unwrap());

            // Run highlighting on the file into a memory buffer
            let mut actual_result: Vec<u8> = Vec::new();
            highlight_diff(&mut fs::File::open(diff).unwrap(), &mut actual_result);
            let actual_result = str::from_utf8(&actual_result).unwrap();

            // Load the corresponding .riff-output file into a string
            let basename = diff.file_stem().unwrap().to_str().unwrap();
            let expected_path = format!(
                "{}/{}.riff-output",
                testdata_path.to_str().unwrap(),
                basename
            );
            let expected_result = fs::read_to_string(expected_path).unwrap();

            // FIXME: Assert that the highlighting output matches the contents of .riff-output
            assert_eq!(actual_result, expected_result);
        }
    }
}
<commit_msg>Accept -h as help, and document --help and --version<commit_after>#[macro_use]
extern crate lazy_static;

use backtrace::Backtrace;
use constants::*;
use git_version::git_version;
use isatty::{stdin_isatty, stdout_isatty};
use refiner::Refiner;
use regex::Regex;
use std::env;
use std::fs;
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::panic;
use std::process::exit;
use std::process::{Command, Stdio};
use std::str;

mod constants;
mod refiner;
mod tokenizer;

const HELP_TEXT: &str = r#"
Usage: diff ... | riff

Colors diff and highlights what parts of changed lines have changed.

Git integration:
    git config --global pager.diff riff
    git config --global pager.show riff
    git config --global interactive.filter riff

Options:
    --help:    Print this text
    --version: Print version number
"#;

const HELP_TEXT_FOOTER: &str = r#"
Report issues at <https://github.com/walles/riff>.
"#;

const CRASH_FOOTER: &str = r#"
Please copy all of the above up to the --- RIFF CRASHED --- marker and report it at one of:
* <https://github.com/walles/riff/issues> (preferred)
* <johan.walles@gmail.com>
"#;

const HUNK_HEADER: &str = "\x1b[36m"; // Cyan
const PAGER_FORKBOMB_STOP: &str = "_RIFF_IGNORE_PAGER";

const GIT_VERSION: &str = git_version!();

lazy_static! {
    static ref STATIC_HEADERS: Vec<(Regex, &'static str)> = vec![
        (Regex::new("^diff ").unwrap(), BOLD),
        (Regex::new("^index ").unwrap(), BOLD),
        (Regex::new("^--- ").unwrap(), BOLD),
        (Regex::new("^\\+\\+\\+ ").unwrap(), BOLD),
        (Regex::new("^@@ ").unwrap(), HUNK_HEADER),
    ];
    static ref ANSI_COLOR_REGEX: Regex = Regex::new("\x1b[^m]*m").unwrap();
}

enum LastLineKind {
    Initial,
    Add,
    Remove,
}

#[must_use]
fn get_fixed_highlight(line: &str) -> &str {
    for static_header in STATIC_HEADERS.iter() {
        let re = &static_header.0;
        if re.is_match(line) {
            return static_header.1;
        }
    }

    return "";
}

fn print(stream: &mut BufWriter<&mut dyn Write>, text: &str) {
    stream.write(text.as_bytes()).unwrap();
}

fn println(stream: &mut BufWriter<&mut dyn Write>, text: &str) {
    print(stream, &text);
    print(stream, "\n");
}

fn highlight_diff(input: &mut dyn io::Read, output: &mut dyn io::Write) {
    let mut adds = String::new();
    let mut removes = String::new();
    let input = BufReader::new(input);
    let output = &mut BufWriter::new(output);
    let mut last_line_kind = LastLineKind::Initial;
    for line in input.lines() {
        let line = line.unwrap();

        // Strip out incoming ANSI formatting. This enables us to highlight
        // already-colored input.
        let line = ANSI_COLOR_REGEX.replace_all(&line, "");

        let fixed_highlight = get_fixed_highlight(&line);
        if !fixed_highlight.is_empty() {
            // Drain outstanding adds / removes
            for line in Refiner::create(&adds, &removes).format() {
                println(output, &line);
            }
            adds.clear();
            removes.clear();

            print(output, fixed_highlight);
            print(output, &line);
            println(output, NORMAL);
            continue;
        }

        // Collect adds / removes
        if !line.is_empty() && line.chars().next().unwrap() == '+' {
            adds.push_str(&line[1..]);
            adds.push('\n');
            last_line_kind = LastLineKind::Add;
            continue;
        } else if !line.is_empty() && line.chars().next().unwrap() == '-' {
            removes.push_str(&line[1..]);
            removes.push('\n');
            last_line_kind = LastLineKind::Remove;
            continue;
        }

        if line == NO_EOF_NEWLINE_MARKER {
            match last_line_kind {
                LastLineKind::Add => {
                    assert!(adds.pop().unwrap() == '\n');
                    continue;
                }
                LastLineKind::Remove => {
                    assert!(removes.pop().unwrap() == '\n');
                    continue;
                }
                LastLineKind::Initial => {
                    // This block intentionally left blank
                }
            }
        }

        last_line_kind = LastLineKind::Initial;

        // Drain outstanding adds / removes
        for line in Refiner::create(&adds, &removes).format() {
            println(output, &line);
        }
        adds.clear();
        removes.clear();

        // Print current line
        if line == NO_EOF_NEWLINE_MARKER {
            print(output, NO_EOF_NEWLINE_COLOR);
            print(output, &line);
            println(output, NORMAL);
        } else {
            println(output, &line);
        }
    }
    for line in Refiner::create(&adds, &removes).format() {
        println(output, &line);
    }
}

/// Try paging using the named pager (`$PATH` will be searched).
///
/// Returns `true` if the pager was found, `false` otherwise.
#[must_use]
fn try_pager(pager_name: &str) -> bool {
    let mut command = Command::new(pager_name);

    if env::var(PAGER_FORKBOMB_STOP).is_ok() {
        // Try preventing fork bombing if $PAGER is set to riff
        return false;
    }
    command.env(PAGER_FORKBOMB_STOP, "1");

    if !env::var("LESS").is_ok() {
        // Set by git when paging
        command.env("LESS", "FRX");
    }

    if !env::var("LV").is_ok() {
        // Set by git when paging
        command.env("LV", "-c");
    }

    command.stdin(Stdio::piped());

    match command.spawn() {
        Ok(mut pager) => {
            let pager_stdin = pager.stdin.as_mut().unwrap();
            highlight_diff(&mut io::stdin().lock(), pager_stdin);

            // FIXME: Report pager exit status if non-zero, together with
            // contents of pager stderr as well if possible.
            pager.wait().expect("Waiting for pager failed");

            return true;
        }
        Err(_) => {
            return false;
        }
    }
}

/// If `option` is found in `argv`, all instances of `option` will be removed
/// from `argv`.
///
/// Returns `true` if `option` was found and consumed, false otherwise.
#[must_use]
fn consume(option: &str, argv: &mut Vec<String>) -> bool {
    if !argv.contains(&option.to_string()) {
        // Not found
        return false;
    }

    argv.retain(|x| x != option);
    return true;
}

fn print_help(output: &mut dyn io::Write) {
    output.write(HELP_TEXT.trim().as_bytes()).unwrap();
    output.write(b"\n").unwrap();
    output.write(b"\n").unwrap();

    let self_path = env::args().into_iter().next().unwrap();
    let self_path = fs::canonicalize(self_path).unwrap();
    let self_path = self_path.as_path().to_str().unwrap();

    // FIXME: Do this only if we aren't already in the $PATH
    output
        .write("Installing riff in the $PATH:\n".as_bytes())
        .unwrap();
    output
        .write(&format!("  * sudo cp {} /usr/local/bin\n", self_path).as_bytes())
        .unwrap();
    output.write(b"\n").unwrap();

    output.write(HELP_TEXT_FOOTER.trim().as_bytes()).unwrap();
    output.write(b"\n").unwrap();
}

fn panic_handler(panic_info: &panic::PanicInfo) {
    let stderr: &mut dyn Write = &mut io::stderr();
    let stderr = &mut BufWriter::new(stderr);
    println(stderr, "\n\n------------ RIFF CRASHED -------------------");

    // Panic message
    if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
        println(stderr, &format!("Panic message: <{:?}>", s));
        println(stderr, "");
    }

    // Backtrace
    // FIXME: Ditch the Backtrace-internal frames from this backtrace
    // FIXME: Ditch the panic internal frames at the end of the backtrace
    println(stderr, &format!("{:?}", Backtrace::new()));

    println(stderr, &format!("Riff version: {}", GIT_VERSION));

    println(stderr, CRASH_FOOTER);
}

fn main() {
    panic::set_hook(Box::new(|panic_info: &panic::PanicInfo| {
        panic_handler(panic_info);
    }));

    let mut args: Vec<String> = env::args().collect();
    if consume("--help", &mut args) || consume("-h", &mut args) {
        print_help(&mut io::stdout());
        return;
    }

    if consume("--version", &mut args) {
        println!("riff {}", GIT_VERSION);
        println!("");
        println!("Source code available at <https://github.com/walles/riff>.");
        return;
    }

    if stdin_isatty() {
        eprintln!("ERROR: Expected input from a pipe");
        eprintln!("");
        print_help(&mut io::stderr());
        exit(1);
    }

    if !stdout_isatty() {
        // We're being piped, just do stdin -> stdout
        highlight_diff(&mut io::stdin().lock(), &mut io::stdout());
        return;
    }

    let pager_env_var = env::var("PAGER");
    if pager_env_var.is_ok() {
        let pager_value = pager_env_var.unwrap();
        if try_pager(&pager_value) {
            return;
        }

        // FIXME: Print warning at the end if $PAGER was set to something that
        // doesn't exist.
    }

    if try_pager("moar") {
        return;
    }

    if try_pager("less") {
        return;
    }

    // No pager found, wth?
    highlight_diff(&mut io::stdin().lock(), &mut io::stdout());
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    fn remove(text: &str) -> String {
        return format!("{}{}{}", REMOVE, text, NORMAL);
    }

    fn add(text: &str) -> String {
        return format!("{}{}{}", ADD, text, NORMAL);
    }

    #[test]
    fn test_remove_trailing_newline() {
        let mut input = "-hej\n\
            +hej\n\
            \\ No newline at end of file\n\
            "
        .as_bytes();

        let expected = format!(
            "{}\n{}\n{}\n",
            remove(&format!("-hej{}⏎", INVERSE_VIDEO)),
            add("+hej"),
            format!(
                "{}\\ No newline at end of file{}",
                NO_EOF_NEWLINE_COLOR, NORMAL
            )
        );

        let mut output: Vec<u8> = Vec::new();
        highlight_diff(&mut input, &mut output);
        assert_eq!(std::str::from_utf8(&output).unwrap(), expected);
    }

    #[test]
    fn test_trailing_newline_context() {
        let mut input = "+bepa\n apa\n\\ No newline at end of file\n".as_bytes();

        let expected = format!(
            "{}\n{}\n{}\n",
            add("+bepa"),
            " apa",
            format!(
                "{}\\ No newline at end of file{}",
                NO_EOF_NEWLINE_COLOR, NORMAL
            )
        );

        let mut output: Vec<u8> = Vec::new();
        highlight_diff(&mut input, &mut output);
        assert_eq!(std::str::from_utf8(&output).unwrap(), expected);
    }

    #[test]
    fn test_testdata_examples() {
        // Example value: `/Users/johan/src/riff/target/debug/deps/riff-7a8916c06b0d3d6c`
        let exe_path = std::env::current_exe().unwrap();

        // Example value: `/Users/johan/src/riff`
        let project_path = exe_path
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .parent()
            .unwrap();

        // Example value: `/Users/johan/src/riff/testdata`
        let testdata_path = &project_path.join("testdata");
        assert!(testdata_path.is_dir());

        // Iterate all files in there
        for diff in fs::read_dir(testdata_path).unwrap() {
            let diff = diff.unwrap();
            let diff = diff.path();
            let diff = diff.as_path();
            if !diff.is_file() {
                continue;
            }

            if diff.extension().unwrap() != "diff" {
                continue;
            }

            println!("Evaluating example file <{}>...", diff.to_str().unwrap());

            // Run highlighting on the file into a memory buffer
            let mut actual_result: Vec<u8> = Vec::new();
            highlight_diff(&mut fs::File::open(diff).unwrap(), &mut actual_result);
            let actual_result = str::from_utf8(&actual_result).unwrap();

            // Load the corresponding .riff-output file into a string
            let basename = diff.file_stem().unwrap().to_str().unwrap();
            let expected_path = format!(
                "{}/{}.riff-output",
                testdata_path.to_str().unwrap(),
                basename
            );
            let expected_result = fs::read_to_string(expected_path).unwrap();

            // FIXME: Assert that the highlighting output matches the contents of .riff-output
            assert_eq!(actual_result, expected_result);
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{Read, Write, Result};

mod buffer;
#[cfg(test)]
mod parser_tests;
#[cfg(test)]
mod interp_tests;

#[derive(Debug, Clone, PartialEq)]
enum Instruction {
    MoveLeft(usize),
    MoveRight(usize),
    Inc(u8),
    Dec(u8),
    Input,
    Output,
    LoopEntry(usize),
    LoopExit(usize),
}

use Instruction::*;

#[derive(Debug, PartialEq)]
struct Program {
    instructions: Vec<Instruction>,
}

impl Program {
    pub fn new(text: String) -> Program {
        // open brackets while parsing
        let mut bracket_stack = Vec::new();
        // idx to patch with index of LoopEntry to patch in
        let mut relocs = Vec::new();
        // "lex" token stream - aka skip whitespace
        let token = text.chars()
            .filter(move |c| {
                        *c == '>' || *c == '<' || *c == '+' || *c == '-' || *c == '.' ||
                        *c == ',' || *c == '[' || *c == ']'
                    });
        let mut instructions = token
            .enumerate()
            .map(|(idx, c)| {
                match c {
                    '<' => MoveLeft(1),
                    '>' => MoveRight(1),
                    '+' => Inc(1),
                    '-' => Dec(1),
                    '.' => Output,
                    ',' => Input,
                    '[' => {
                        bracket_stack.push(idx);
                        // value must be patched later
                        LoopEntry(std::usize::MAX)
                    }
                    ']' => {
                        if let Some(loop_entry) = bracket_stack.pop() {
                            relocs.push((loop_entry, idx));
                            LoopExit(loop_entry)
                        } else {
                            panic!("Unbalanced {:?} at pc={:}", c, idx);
                        }
                    }
                    c => panic!("Unknown instruction {:?} at pc={:}", c, idx),
                }
            })
            .collect::<Vec<_>>();
        if let Some(unbalanced_idx) = bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
        Program { instructions: instructions }
    }

    pub fn interp(&self, input: &mut Read, output: &mut Write) {
        let mut memory = vec![0u8; 30000];

        let mut pc: usize = 0;
        let mut dataptr: usize = 0;

        fn get_char(input: &mut Read) -> u8 {
            let mut buf = [0u8; 1];
            match input.read(&mut buf) {
                Err(_) => panic!("Cannot read from stdin"),
                Ok(_) => buf[0],
            }
        }

        while pc < self.instructions.len() {
            match self.instructions[pc] {
                MoveLeft(offset) => dataptr -= offset,
                MoveRight(offset) => dataptr += offset,
                Inc(increment) => memory[dataptr] = memory[dataptr].wrapping_add(increment),
                Dec(decrement) => memory[dataptr] = memory[dataptr].wrapping_sub(decrement),
                Output => write!(output, "{:}", memory[dataptr] as char).expect("Output error"),
                Input => memory[dataptr] = get_char(input),
                LoopEntry(target) => {
                    if 0 == memory[dataptr] {
                        pc = target;
                    }
                }
                LoopExit(target) => {
                    if 0 != memory[dataptr] {
                        pc = target;
                    }
                }
                // ref c => panic!("Unknown instruction {:?} at pc={:}", c, pc),
            }
            pc += 1;
        }
    }
}

fn load_program(fname: String) -> Result<Program> {
    let mut file = File::open(fname)?;
    let mut contents: Vec<u8> = Vec::new();
    file.read_to_end(&mut contents)?;
    let text = String::from_utf8(contents).unwrap();

    Ok(Program::new(text))
}

fn main() {
    match load_program("examples/mandelbrot.bf".to_string()) {
        Ok(ref mut p) => {
            //println!("{:?}", p);
            p.interp(&mut std::io::stdin(), &mut std::io::stdout());
        }
        Err(err) => panic!("Cannot read file because {:?}", err),
    }
}
<commit_msg>started with peephole optimizations<commit_after>use std::fs::File;
use std::io::{Read, Write, Result};

mod buffer;
#[cfg(test)]
mod parser_tests;
#[cfg(test)]
mod interp_tests;

#[derive(Debug, Clone, Copy, PartialEq)]
enum Instruction {
    MoveLeft(usize),
    MoveRight(usize),
    Inc(u8),
    Dec(u8),
    Input,
    Output,
    LoopEntry(usize),
    LoopExit(usize),
}

use Instruction::*;

#[derive(Debug, PartialEq)]
struct Program {
    instructions: Vec<Instruction>,
}

impl Program {
    pub fn new(text: String) -> Program {
        // open brackets while parsing
        let mut bracket_stack = Vec::new();
        // idx to patch with index of LoopEntry to patch in
        let mut relocs = Vec::new();
        // "lex" token stream - aka skip whitespace
        let token = text.chars()
            .filter(move |c| {
                        *c == '>' || *c == '<' || *c == '+' || *c == '-' || *c == '.' ||
                        *c == ',' || *c == '[' || *c == ']'
                    });
        let mut instructions = token
            .enumerate()
            .map(|(idx, c)| {
                match c {
                    '<' => MoveLeft(1),
                    '>' => MoveRight(1),
                    '+' => Inc(1),
                    '-' => Dec(1),
                    '.' => Output,
                    ',' => Input,
                    '[' => {
                        bracket_stack.push(idx);
                        // value must be patched later
                        LoopEntry(std::usize::MAX)
                    }
                    ']' => {
                        if let Some(loop_entry) = bracket_stack.pop() {
                            relocs.push((loop_entry, idx));
                            LoopExit(loop_entry)
                        } else {
                            panic!("Unbalanced {:?} at pc={:}", c, idx);
                        }
                    }
                    c => panic!("Unknown instruction {:?} at pc={:}", c, idx),
                }
            })
            .collect::<Vec<_>>();
        if let Some(unbalanced_idx) = bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
        Program { instructions: instructions }
    }

    pub fn optimize(&self) -> Program {
        // optimized program
        let mut instructions:Vec<Instruction> = Vec::new();
        // open brackets while parsing
        let mut bracket_stack = Vec::new();
        // idx to patch with index of LoopEntry to patch in
        let mut relocs = Vec::new();

        for (idx, instr) in self.instructions.iter().enumerate() {
            match *instr {
                MoveLeft(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveLeft(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveLeft(old_offset + offset);
                    }
                    else {
                        instructions.push(*instr);
                    }
                }
                MoveRight(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveRight(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveRight(old_offset + offset);
                    }
                    else {
                        instructions.push(*instr);
                    }
                }
                Inc(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Inc(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Inc(old_val.wrapping_add(val));
                    }
                    else {
                        instructions.push(*instr);
                    }
                }
                Dec(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Dec(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Dec(old_val.wrapping_add(val));
                    }
                    else {
                        instructions.push(*instr);
                    }
                }
                LoopEntry(_) => {
                    bracket_stack.push(instructions.len());
                    instructions.push(LoopEntry(std::usize::MAX));
                }
                LoopExit(_) => {
                    if let Some(loop_entry) = bracket_stack.pop() {
                        relocs.push((loop_entry, instructions.len()));
                        instructions.push(LoopExit(loop_entry));
                    } else {
                        panic!("Unbalanced {:?} at pc={:}", instr, idx);
                    }
                }
                instr => instructions.push(instr),
            }
        }
        if let Some(unbalanced_idx) = bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
        Program { instructions: instructions }
    }

    pub fn interp(&self, input: &mut Read, output: &mut Write) {
        let mut memory = vec![0u8; 30000];

        let mut pc: usize = 0;
        let mut dataptr: usize = 0;

        fn get_char(input: &mut Read) -> u8 {
            let mut buf = [0u8; 1];
            match input.read(&mut buf) {
                Err(_) => panic!("Cannot read from stdin"),
                Ok(_) => buf[0],
            }
        }

        while pc < self.instructions.len() {
            match self.instructions[pc] {
                MoveLeft(offset) => dataptr -= offset,
                MoveRight(offset) => dataptr += offset,
                Inc(increment) => memory[dataptr] = memory[dataptr].wrapping_add(increment),
                Dec(decrement) => memory[dataptr] = memory[dataptr].wrapping_sub(decrement),
                Output => write!(output, "{:}", memory[dataptr] as char).expect("Output error"),
                Input => memory[dataptr] = get_char(input),
                LoopEntry(target) => {
                    if 0 == memory[dataptr] {
                        pc = target;
                    }
                }
                LoopExit(target) => {
                    if 0 != memory[dataptr] {
                        pc = target;
                    }
                }
                // ref c => panic!("Unknown instruction {:?} at pc={:}", c, pc),
            }
            pc += 1;
        }
    }
}

fn load_program(fname: String) -> Result<Program> {
    let mut file = File::open(fname)?;
    let mut contents: Vec<u8> = Vec::new();
    file.read_to_end(&mut contents)?;
    let text = String::from_utf8(contents).unwrap();

    Ok(Program::new(text))
}

fn main() {
    match load_program("examples/mandelbrot.bf".to_string()) {
        Ok(ref mut p) => {
            let p = p.optimize();
            //println!("{:?}", p);
            p.interp(&mut std::io::stdin(), &mut std::io::stdout());
        }
        Err(err) => panic!("Cannot read file because {:?}", err),
    }
}
<|endoftext|>


<commit_before>//
// main.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 13 Mar 2017 21:03:41 +0100 (CET)
//

extern crate rand;
extern crate cursive;

#[cfg(feature = "go-slow")]
use std::thread;
#[cfg(feature = "go-slow")]
use std::time;

mod process;
mod scheduler;
mod fcfs;
mod round_robin;
mod sjf;
mod srtf;
mod tui;

use scheduler::*;
use tui::*;

/// Length of a single system clock tick (frequency of actual, real time)
#[cfg(feature = "go-slow")]
const CLOCK_HZ: u64 = 1000;
/// Switch context every <value> clock ticks
const SYSTEM_HZ: usize = 8;

/// Spawn this many processes in each scenario
const SPAWNED_PROCESS_NUM: usize = 32;
/// Number of scenarios
const SCENARIOS_NUM: usize = 4;

struct SimulationResult {
  average_waiting_time: f64,
  average_turnaround_time: f64,
  context_switch_num: usize,
}

fn main() {
  let mut tui = tui::new();

  let scenarios = {
    (0..SCENARIOS_NUM).map(|_| (0..SPAWNED_PROCESS_NUM).map(|_| {
      let upper_limit = rand::random::<usize>() % (SYSTEM_HZ * 2) + 1;
      let burst_time = rand::random::<usize>() % upper_limit + SYSTEM_HZ;

      burst_time + 1 // + 1 so there's no chance of having a process with 0 burst time
    }).collect::<Vec<usize>>()).collect::<Vec<Vec<usize>>>()
  };

  for (idx, ref scenario) in scenarios.iter().enumerate() {
    tui.add_scenario(format!("{}", idx), scenario);
  }

  macro_rules! run_simulation_suite {
    ($scheduler:ident) => ({
      let mut average_waiting_time = 0f64;
      let mut average_turnaround_time = 0f64;
      let mut total_context_switches = 0usize;
      let mut average_context_switches = 0f64;

      for (run, scenario) in scenarios.iter().enumerate() {
        tui.set_header(format!("Algorithm: {}, run #{}", stringify!($scheduler), run));

        let scheduler = $scheduler::new();
        let result = run_simulation(&mut tui, scheduler, scenario.clone());

        average_waiting_time -= average_waiting_time / (run + 1) as f64;
        average_waiting_time += result.average_waiting_time / (run + 1) as f64;
        total_context_switches += result.context_switch_num;

        average_turnaround_time -= average_turnaround_time / (run + 1) as f64;
        average_turnaround_time += result.average_turnaround_time / (run + 1) as f64;

        average_context_switches -= average_context_switches / (run + 1) as f64;
        average_context_switches += result.context_switch_num as f64 / (run + 1) as f64;
      }

      tui.add_result(format!("{}: Average waiting time: {:02.2}", stringify!($scheduler), average_waiting_time));
      tui.add_result(format!("{}: Average turnaround time: {:02.2}", stringify!($scheduler), average_turnaround_time));
      tui.add_result(format!("{}: Total # of context switches: {:02}", stringify!($scheduler), total_context_switches));
      tui.add_result(format!("{}: Average # of context switches: {:02.2}", stringify!($scheduler), average_context_switches));
    })
  }

  run_simulation_suite!(fcfs);
  run_simulation_suite!(round_robin);
  run_simulation_suite!(sjf);
  run_simulation_suite!(srtf);

  tui.update();

  loop {}
}

fn run_simulation<S>(mut tui: &mut Tui, mut scheduler: S, mut process_list: Vec<usize>) -> SimulationResult
where S: Scheduler {
  let mut clock_tick = 0;
  let mut process_spawner = process::new_spawner();
  let mut average_waiting_time = 0f64;
  let mut average_turnaround_time = 0f64;
  let mut num_of_spawned_procs = 1;

  tui.debug(format!("{:05}: Starting simulation using {}", clock_tick, scheduler.name()));

  // vectors only support the pop (from the back) operation so
  // reverse it (so we pop from the front)
  process_list.reverse();

  // Add the equivalent of the 'init' process
  scheduler.add_process(process_spawner.spawn(process_list.pop().unwrap(), clock_tick));

  loop {
    // update the process view
    scheduler.list_processes(&mut tui);

    if cfg!(feature = "go-slow") || clock_tick % 2 == 0 {
      // update all the views
      tui.update();
    }

    if !scheduler.has_processes() {
      // if there's no more processes - end the simulation
      // if a process can only be created by another process - there's
      // nothing else to do
      break;
    }

    // simulate executing the current process
    tui.debug(format!("{:05}: Executing {}", clock_tick, scheduler.current_proc().unwrap().name));
    scheduler.current_proc_mut().unwrap().record_execution();

    // Simulate other processes creating new processes
    // Assume that you can't create new processes, if there's none already
    if clock_tick % SYSTEM_HZ == 0 {
      if let Some(burst_time) = process_list.pop() {
        let new_proc = process_spawner.spawn(burst_time, clock_tick);
        tui.debug(format!("{:05}: Spawning {}", clock_tick, new_proc.name));
        scheduler.add_process(new_proc);
        num_of_spawned_procs += 1;
      }
    }

    if scheduler.current_proc().unwrap().done_executing() {
      tui.debug(format!("{:05}: Killing process {}", clock_tick, scheduler.current_proc().unwrap().name));
      scheduler.kill_current_proc();
    }

    if scheduler.has_processes() {
      if clock_tick % SYSTEM_HZ == SYSTEM_HZ - 1 {
        let prev_proc_name = scheduler.current_proc().unwrap().name.clone();
        scheduler.schedule();
        tui.debug(format!("{:05}: Switching context: {} -> {}", clock_tick, prev_proc_name, scheduler.current_proc().unwrap().name));
        tui.debug(format!("{:05}: {} was waiting {} clock ticks", clock_tick, scheduler.current_proc().unwrap().name, scheduler.current_proc().unwrap().waiting_time));

        average_waiting_time -= average_waiting_time / num_of_spawned_procs as f64;
        average_waiting_time += scheduler.current_proc().unwrap().waiting_time as f64 / num_of_spawned_procs as f64;

        average_turnaround_time -= average_turnaround_time / num_of_spawned_procs as f64;
        average_turnaround_time += (clock_tick as f64 - scheduler.current_proc().unwrap().arrival_time as f64) / num_of_spawned_procs as f64;
      }

      // increase the waiting time for every process
      scheduler.increase_waiting_times();
    }

    #[cfg(feature = "go-slow")]
    thread::sleep(time::Duration::from_millis(1000 / CLOCK_HZ));

    clock_tick += 1;
  }

  SimulationResult {
    average_waiting_time: average_waiting_time,
    average_turnaround_time: average_turnaround_time,
    context_switch_num: scheduler.context_switch_num(),
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<commit_msg>Show which algorithm had the lowest / shortest average times<commit_after>//
// main.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 13 Mar 2017 21:03:41 +0100 (CET)
//

extern crate rand;
extern crate cursive;

#[cfg(feature = "go-slow")]
use std::thread;
#[cfg(feature = "go-slow")]
use std::time;

mod process;
mod scheduler;
mod fcfs;
mod round_robin;
mod sjf;
mod srtf;
mod tui;

use scheduler::*;
use tui::*;

/// Length of a single system clock tick (frequency of actual, real time)
#[cfg(feature = "go-slow")]
const CLOCK_HZ: u64 = 1000;
/// Switch context every <value> clock ticks
const SYSTEM_HZ: usize = 8;

/// Spawn this many processes in each scenario
const SPAWNED_PROCESS_NUM: usize = 32;
/// Number of scenarios
const SCENARIOS_NUM: usize = 4;

struct SimulationResult {
  average_waiting_time: f64,
  average_turnaround_time: f64,
  context_switch_num: usize,
}

fn main() {
  let mut tui = tui::new();

  let mut shortest_average_waiting_time = (String::new(), std::f64::MAX);
  let mut shortest_average_turnaround_time = (String::new(), std::f64::MAX);
  let mut lowest_context_switch_num = (String::new(), std::usize::MAX);
  let mut lowest_average_context_switch_num = (String::new(), std::f64::MAX);

  let scenarios = {
    (0..SCENARIOS_NUM).map(|_| (0..SPAWNED_PROCESS_NUM).map(|_| {
      let upper_limit = rand::random::<usize>() % (SYSTEM_HZ * 2) + 1;
      let burst_time = rand::random::<usize>() % upper_limit + SYSTEM_HZ;

      burst_time + 1 // + 1 so there's no chance of having a process with 0 burst time
    }).collect::<Vec<usize>>()).collect::<Vec<Vec<usize>>>()
  };

  for (idx, ref scenario) in scenarios.iter().enumerate() {
    tui.add_scenario(format!("{}", idx), scenario);
  }

  macro_rules! run_simulation_suite {
    ($scheduler:ident) => ({
      let mut average_waiting_time = 0f64;
      let mut average_turnaround_time = 0f64;
      let mut total_context_switches = 0usize;
      let mut average_context_switches = 0f64;

      for (run, scenario) in scenarios.iter().enumerate() {
        tui.set_header(format!("Algorithm: {}, run #{}", stringify!($scheduler), run));

        let scheduler = $scheduler::new();
        let result = run_simulation(&mut tui, scheduler, scenario.clone());

        average_waiting_time -= average_waiting_time / (run + 1) as f64;
        average_waiting_time += result.average_waiting_time / (run + 1) as f64;
        total_context_switches += result.context_switch_num;

        average_turnaround_time -= average_turnaround_time / (run + 1) as f64;
        average_turnaround_time += result.average_turnaround_time / (run + 1) as f64;

        average_context_switches -= average_context_switches / (run + 1) as f64;
        average_context_switches += result.context_switch_num as f64 / (run + 1) as f64;
      }

      tui.add_result(format!("{}: Average waiting time: {:02.2}", stringify!($scheduler), average_waiting_time));
      tui.add_result(format!("{}: Average turnaround time: {:02.2}", stringify!($scheduler), average_turnaround_time));
      tui.add_result(format!("{}: Total # of context switches: {:02}", stringify!($scheduler), total_context_switches));
      tui.add_result(format!("{}: Average # of context switches: {:02.2}", stringify!($scheduler), average_context_switches));

      if average_waiting_time < shortest_average_waiting_time.1 {
        shortest_average_waiting_time = (String::from(stringify!($scheduler)), average_waiting_time);
      }

      if average_turnaround_time < shortest_average_turnaround_time.1 {
        shortest_average_turnaround_time = (String::from(stringify!($scheduler)), average_turnaround_time);
      }

      if total_context_switches < lowest_context_switch_num.1 {
        lowest_context_switch_num = (String::from(stringify!($scheduler)), total_context_switches);
      }

      if average_context_switches < lowest_average_context_switch_num.1 {
        lowest_average_context_switch_num = (String::from(stringify!($scheduler)), average_context_switches);
      }
    })
  }

  run_simulation_suite!(fcfs);
  run_simulation_suite!(round_robin);
  run_simulation_suite!(sjf);
  run_simulation_suite!(srtf);

  tui.add_result(format!(""));
  tui.add_result(format!("-- Summary --"));
  tui.add_result(format!("Shortest average waiting time: {}", shortest_average_waiting_time.0));
  tui.add_result(format!("Shortest average turnaround time: {}", shortest_average_turnaround_time.0));
  tui.add_result(format!("Lowest context switch #: {}", lowest_context_switch_num.0));
  tui.add_result(format!("Lowest average context switch #: {}", lowest_average_context_switch_num.0));

  tui.update();

  loop {}
}

fn run_simulation<S>(mut tui: &mut Tui, mut scheduler: S, mut process_list: Vec<usize>) -> SimulationResult
where S: Scheduler {
  let mut clock_tick = 0;
  let mut process_spawner = process::new_spawner();
  let mut average_waiting_time = 0f64;
  let mut average_turnaround_time = 0f64;
  let mut num_of_spawned_procs = 1;

  tui.debug(format!("{:05}: Starting simulation using {}", clock_tick, scheduler.name()));

  // vectors only support the pop (from the back) operation so
  // reverse it (so we pop from the front)
  process_list.reverse();

  // Add the equivalent of the 'init' process
  scheduler.add_process(process_spawner.spawn(process_list.pop().unwrap(), clock_tick));

  loop {
    // update the process view
    scheduler.list_processes(&mut tui);

    if cfg!(feature = "go-slow") || clock_tick % 2 == 0 {
      // update all the views
      tui.update();
    }

    if !scheduler.has_processes() {
      // if there's no more processes - end the simulation
      // if a process can only be created by another process - there's
      // nothing else to do
      break;
    }

    // simulate executing the current process
    tui.debug(format!("{:05}: Executing {}", clock_tick, scheduler.current_proc().unwrap().name));
    scheduler.current_proc_mut().unwrap().record_execution();

    // Simulate other processes creating new processes
    // Assume that you can't create new processes, if there's none already
    if clock_tick % SYSTEM_HZ == 0 {
      if let Some(burst_time) = process_list.pop() {
        let new_proc = process_spawner.spawn(burst_time, clock_tick);
        tui.debug(format!("{:05}: Spawning {}", clock_tick, new_proc.name));
        scheduler.add_process(new_proc);
        num_of_spawned_procs += 1;
      }
    }

    if scheduler.current_proc().unwrap().done_executing() {
      tui.debug(format!("{:05}: Killing process {}", clock_tick, scheduler.current_proc().unwrap().name));
      scheduler.kill_current_proc();
    }

    if scheduler.has_processes() {
      if clock_tick % SYSTEM_HZ == SYSTEM_HZ - 1 {
        let prev_proc_name = scheduler.current_proc().unwrap().name.clone();
        scheduler.schedule();
        tui.debug(format!("{:05}: Switching context: {} -> {}", clock_tick, prev_proc_name, scheduler.current_proc().unwrap().name));
        tui.debug(format!("{:05}: {} was waiting {} clock ticks", clock_tick, scheduler.current_proc().unwrap().name, scheduler.current_proc().unwrap().waiting_time));

        average_waiting_time -= average_waiting_time / num_of_spawned_procs as f64;
        average_waiting_time += scheduler.current_proc().unwrap().waiting_time as f64 / num_of_spawned_procs as f64;

        average_turnaround_time -= average_turnaround_time / num_of_spawned_procs as f64;
        average_turnaround_time += (clock_tick as f64 - scheduler.current_proc().unwrap().arrival_time as f64) / num_of_spawned_procs as f64;
      }

      // increase the waiting time for every process
      scheduler.increase_waiting_times();
    }

    #[cfg(feature = "go-slow")]
    thread::sleep(time::Duration::from_millis(1000 / CLOCK_HZ));

    clock_tick += 1;
  }

  SimulationResult {
    average_waiting_time: average_waiting_time,
    average_turnaround_time: average_turnaround_time,
    context_switch_num: scheduler.context_switch_num(),
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<|endoftext|>


<commit_before>#![feature(collections, core, io, os, path, rustc_private, std_misc, env, test)]

#[macro_use] extern crate log;

extern crate syntax;
extern crate collections;
extern crate core;

#[cfg(not(test))]
use racer::Match;
#[cfg(not(test))]
use racer::util::getline;
#[cfg(not(test))]
use racer::nameres::{do_file_search, do_external_search};
#[cfg(not(test))]
use racer::scopes;

pub mod racer;

#[cfg(not(test))]
fn match_with_snippet_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords_from_file(&m.filepath, m.point).unwrap();
    if m.matchstr == "" {
        panic!("MATCHSTR is empty - waddup?");
    }

    let snippet = racer::snippets::snippet_for_match(&m);
    println!("MATCH {};{};{};{};{};{:?};{}", m.matchstr,
                                    snippet,
                                    linenum.to_string(),
                                    charnum.to_string(),
                                    m.filepath.as_str().unwrap(),
                                    m.mtype,
                                    m.contextstr,
             );
}

#[cfg(not(test))]
fn match_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords_from_file(&m.filepath, m.point).unwrap();
    if m.matchstr == "" {
        panic!("MATCHSTR is empty - waddup?");
    }
    println!("MATCH {},{},{},{},{:?},{}", m.matchstr,
                                    linenum.to_string(),
                                    charnum.to_string(),
                                    m.filepath.as_str().unwrap(),
                                    m.mtype,
                                    m.contextstr
             );
}

#[cfg(not(test))]
fn complete(match_found : &Fn(Match)) {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 3 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    match args[2].parse::<usize>() {
        Ok(linenum) => {
            // input: linenum, colnum, fname
            if args.len() < 5 {
                println!("Provide more arguments!");
                print_usage();
                std::env::set_exit_status(1);
                return;
            }
            let charnum = args[3].parse::<usize>().unwrap();
            let fname = &args[4][];
            let fpath = Path::new(fname);
            let src = racer::load_file(&fpath);
            let line = &*getline(&fpath, linenum);
            let (start, pos) = racer::util::expand_ident(line, charnum);
            println!("PREFIX {},{},{}", start, pos, &line[start..pos]);

            let point = scopes::coords_to_point(&*src, linenum, charnum);
            for m in racer::complete_from_file(&*src, &fpath, point) {
                match_found(m);
            }
        }
        Err(_) => {
            // input: a command line string passed in
            let arg = &args[2][];
            let it = arg.split_str("::");
            let p : Vec<&str> = it.collect();

            for m in do_file_search(p[0], &Path::new(".")) {
                if p.len() == 1 {
                    match_found(m);
                } else {
                    for m in do_external_search(&p[1..], &m.filepath, m.point, racer::SearchType::StartsWith, racer::Namespace::BothNamespaces) {
                        match_found(m);
                    }
                }
            }
        }
    }
}

#[cfg(not(test))]
fn prefix() {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 5 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    let linenum = args[2].parse::<usize>().unwrap();
    let charnum = args[3].parse::<usize>().unwrap();
    let fname = &args[4][];

    // print the start, end, and the identifier prefix being matched
    let path = Path::new(fname);
    let line = &*getline(&path, linenum);
    let (start, pos) = racer::util::expand_ident(line, charnum);
    println!("PREFIX {},{},{}", start, pos, &line[start..pos]);
}

#[cfg(not(test))]
fn find_definition() {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 5 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    let linenum = args[2].parse::<usize>().unwrap();
    let charnum = args[3].parse::<usize>().unwrap();
    let fname = &args[4][];
    let fpath = Path::new(fname);
    let src = racer::load_file(&fpath);
    let pos = scopes::coords_to_point(&*src, linenum, charnum);

    racer::find_definition(&*src, &fpath, pos).map(match_fn);
}

#[cfg(not(test))]
fn print_usage() {
    let program = std::env::args().next().unwrap().into_string().unwrap().clone();
    println!("usage: {} complete linenum charnum fname", program);
    println!("or:    {} find-definition linenum charnum fname", program);
    println!("or:    {} complete fullyqualifiedname   (e.g. std::io::)",program);
    println!("or:    {} prefix linenum charnum fname",program);
    println!("or replace complete with complete-with-snippet for more detailed completions.");
}

#[cfg(not(test))]
fn main() {
    if std::env::var_string("RUST_SRC_PATH").is_err() {
        println!("RUST_SRC_PATH environment variable must be set");
        std::env::set_exit_status(1);
        return;
    }

    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();

    if args.len() == 1 {
        print_usage();
        std::env::set_exit_status(1);
        return;
    }

    let command = &args[1][];
    match command {
        "prefix" => prefix(),
        "complete" => complete(&match_fn),
        "complete-with-snippet" => complete(&match_with_snippet_fn),
        "find-definition" => find_definition(),
        "help" => print_usage(),
        _ => {
            println!("Sorry, I didn't understand command {}", command );
            print_usage();
            std::env::set_exit_status(1);
            return;
        }
    }
}
<commit_msg>Work around rust-lang/rust#22008 (fixes #139)<commit_after>#![feature(collections, core, io, os, path, rustc_private, std_misc, env, test)]

#[macro_use] extern crate log;

extern crate syntax;
extern crate collections;
extern crate core;

#[cfg(not(test))]
use racer::Match;
#[cfg(not(test))]
use racer::util::getline;
#[cfg(not(test))]
use racer::nameres::{do_file_search, do_external_search};
#[cfg(not(test))]
use racer::scopes;

pub mod racer;

#[cfg(not(test))]
fn match_with_snippet_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords_from_file(&m.filepath, m.point).unwrap();
    if m.matchstr == "" {
        panic!("MATCHSTR is empty - waddup?");
    }

    let snippet = racer::snippets::snippet_for_match(&m);
    println!("MATCH {};{};{};{};{};{:?};{}", m.matchstr,
                                    snippet,
                                    linenum.to_string(),
                                    charnum.to_string(),
                                    m.filepath.as_str().unwrap(),
                                    m.mtype,
                                    m.contextstr,
             );
}

#[cfg(not(test))]
fn match_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords_from_file(&m.filepath, m.point).unwrap();
    if m.matchstr == "" {
        panic!("MATCHSTR is empty - waddup?");
    }
    println!("MATCH {},{},{},{},{:?},{}", m.matchstr,
                                    linenum.to_string(),
                                    charnum.to_string(),
                                    m.filepath.as_str().unwrap(),
                                    m.mtype,
                                    m.contextstr
             );
}

#[cfg(not(test))]
fn complete(match_found : &Fn(Match)) {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 3 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    match args[2].parse::<usize>() {
        Ok(linenum) => {
            // input: linenum, colnum, fname
            if args.len() < 5 {
                println!("Provide more arguments!");
                print_usage();
                std::env::set_exit_status(1);
                return;
            }
            let charnum = args[3].parse::<usize>().unwrap();
            let fname = &args[4][];
            let fpath = Path::new(fname);
            let src = racer::load_file(&fpath);
            let line = &*getline(&fpath, linenum);
            let (start, pos) = racer::util::expand_ident(line, charnum);
            println!("PREFIX {},{},{}", start, pos, &line[start..pos]);

            let point = scopes::coords_to_point(&*src, linenum, charnum);
            for m in racer::complete_from_file(&*src, &fpath, point) {
                match_found(m);
            }
        }
        Err(_) => {
            // input: a command line string passed in
            let arg = &args[2][];
            let it = arg.split_str("::");
            let p : Vec<&str> = it.collect();

            for m in do_file_search(p[0], &Path::new(".")) {
                if p.len() == 1 {
                    match_found(m);
                } else {
                    for m in do_external_search(&p[1..], &m.filepath, m.point, racer::SearchType::StartsWith, racer::Namespace::BothNamespaces) {
                        match_found(m);
                    }
                }
            }
        }
    }
}

#[cfg(not(test))]
fn prefix() {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 5 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    let linenum = args[2].parse::<usize>().unwrap();
    let charnum = args[3].parse::<usize>().unwrap();
    let fname = &args[4][];

    // print the start, end, and the identifier prefix being matched
    let path = Path::new(fname);
    let line = &*getline(&path, linenum);
    let (start, pos) = racer::util::expand_ident(line, charnum);
    println!("PREFIX {},{},{}", start, pos, &line[start..pos]);
}

#[cfg(not(test))]
fn find_definition() {
    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();
    if args.len() < 5 {
        println!("Provide more arguments!");
        print_usage();
        std::env::set_exit_status(1);
        return;
    }
    let linenum = args[2].parse::<usize>().unwrap();
    let charnum = args[3].parse::<usize>().unwrap();
    let fname = &args[4][];
    let fpath = Path::new(fname);
    let src = racer::load_file(&fpath);
    let pos = scopes::coords_to_point(&*src, linenum, charnum);

    racer::find_definition(&*src, &fpath, pos).map(match_fn);
}

#[cfg(not(test))]
fn print_usage() {
    let program = std::env::args().next().unwrap().into_string().unwrap().clone();
    println!("usage: {} complete linenum charnum fname", program);
    println!("or:    {} find-definition linenum charnum fname", program);
    println!("or:    {} complete fullyqualifiedname   (e.g. std::io::)",program);
    println!("or:    {} prefix linenum charnum fname",program);
    println!("or replace complete with complete-with-snippet for more detailed completions.");
}

#[cfg(not(test))]
fn main() {
    if std::env::var_string("RUST_SRC_PATH").is_err() {
        println!("RUST_SRC_PATH environment variable must be set");
        std::env::set_exit_status(1);
        return;
    }

    let args = std::env::args().map(|a| a.into_string().unwrap()).collect::<Vec<_>>();

    if args.len() == 1 {
        print_usage();
        std::env::set_exit_status(1);
        return;
    }

    let command = &args[1][];
    match command {
        "prefix" => prefix(),
        "complete" => complete(&match_fn),
        "complete-with-snippet" => complete(&match_with_snippet_fn),
        "find-definition" => find_definition(),
        "help" => print_usage(),
        cmd => {
            println!("Sorry, I didn't understand command {}", cmd);
            print_usage();
            std::env::set_exit_status(1);
            return;
        }
    }
}
<|endoftext|>


<commit_before>extern crate mio;

#[macro_use]
extern crate log;

extern crate env_logger;

use std::io;
//use std::io::{Error, ErrorKind};
use std::net::SocketAddr;
use mio::*;
use mio::tcp::{TcpListener, TcpStream, Shutdown};
use mio::util::Slab;
use mio::buf::RingBuf;

const ENDPOINT: &'static str = "127.0.0.1:6666";
const DESTINATION: &'static str = "127.0.0.1:80";

//const BUF_SIZE: usize = 524288; //131072;
//const BUF_SIZE: usize = 8388608;
const BUF_SIZE: usize = 1048576;

const INVALID: Token = Token(0);
const ACCEPTOR: Token = Token(1);
const FLOW: Token = Token(2);

const OUTMASK: usize = 2147483648; //2 ** 31
const NOTMASK: usize = !OUTMASK;

struct Conn {
    // socket
    sock: TcpStream,

    // token used to register the inbound with the event loop
    token: Token,

    // set of events we are interested in
    interest: EventSet,

    // whether the remote peer has half-closed the socket
    // i.e. reading makes no more sense
    dead: bool,

    // whether the socket can accept writes right now
    writable: bool,

    // octets in transit
    buf: RingBuf,
}

/// A stateful wrapper around a non-blocking stream. This connection is not
/// the SERVER connection. This connection represents the client connections
/// _accepted_ by the SERVER connection.
struct Flow {

    // inbound connection
    inb: Conn,

    // outbound connection
    out: Conn,
}

struct Nexus {
    // main server socket (that accepts inbound connections)
    acceptor: TcpListener,

    // token of the acceptor
    token: Token,

    // a list of all inbound and outbound connections
    conns: Slab<Flow>,
}

impl Conn {
    fn new(token: Token, sock: TcpStream) -> Conn {
        Conn {
            sock: sock,
            token: token,
            interest: EventSet::all(),
            dead: false,
            writable: false,
            buf: RingBuf::new(BUF_SIZE),
        }
    }

    /// Re-register flow interest in read events with the event_loop.
    #[inline]
    fn reregister(&mut self, event_loop: &mut EventLoop<Nexus>)
                  -> io::Result<()> {

        debug!("Reregistering {:?} for {:?}", self.token, self.interest);
        event_loop.reregister(
            &self.sock, self.token, self.interest, PollOpt::edge()).or_else(|e| {
                error!("Failed to reregister {:?} for {:?} due to {:?}", self.token, self.interest, e);
                Err(e)
            })
    }
}

impl Flow {

    fn new(insock: TcpStream, ousock: TcpStream, token: Token) -> Flow {
        Flow {
            inb: Conn::new(token, insock),
            out: Conn::new(Token(OUTMASK + token.as_usize()), ousock),
        }
    }

    /// Register flow interest in read events with the event_loop.
    ///
    /// This will let our connection accept reads starting next event loop tick.
    fn register(&mut self, event_loop: &mut EventLoop<Nexus>)
                -> io::Result<()> {

        event_loop.register_opt(
            &self.inb.sock,
            self.inb.token,
            self.inb.interest,
            PollOpt::edge()).and_then(|_|
        event_loop.register_opt(
            &self.out.sock,
            self.out.token,
            self.out.interest,
            PollOpt::edge())).
        or_else(|e| {
            error!("Failed to register {:?} due to {:?}", self.inb.token, e);
            Err(e)
        })
    }

    /// Handle flow read event from event loop.
    ///
    #[inline]
    fn read(&mut self, inbo: bool, event_loop: &mut EventLoop<Nexus>) -> bool {

        let conn;
        let peer;
        if inbo {
            conn = &mut self.inb;
            peer = &mut self.out;
        } else {
            conn = &mut self.out;
            peer = &mut self.inb;
        }

        read1(conn, peer, event_loop)
    }

    #[inline]
    fn write(&mut self, inbo: bool, event_loop: &mut EventLoop<Nexus>) -> bool {
        let conn;
        let peer;
        if inbo {
            conn = &mut self.inb;
            peer = &mut self.out;
        } else {
            conn = &mut self.out;
            peer = &mut self.inb;
        }

        conn.writable = true;

        write1(conn, peer, event_loop)
    }
}

fn read1(conn: &mut Conn, peer: &mut Conn, event_loop: &mut EventLoop<Nexus>) -> bool {
    loop {
        match conn.sock.try_read_buf(&mut conn.buf) {
            Ok(Some(0)) => {
                debug!("Successful read of zero bytes (i.e. EOF) from token {:?}", conn.token);
                conn.dead = true;
                conn.sock.shutdown(Shutdown::Write);
                if peer.dead == true {
                    return false;
                } else {
                    return true;
                }
            },
            Ok(Some(n)) => {
                debug!("Successfully read {} bytes from {:?}, buf size: {}",
                       n, conn.token, <RingBuf as Buf>::remaining(&conn.buf));
                if peer.writable {
                    write1(peer, conn, event_loop);
                }
                if conn.buf.is_full() {
                    warn!("Suspending reads due to buffer full for {:?}", conn.token);
                    conn.interest = conn.interest & !EventSet::readable();
                    conn.reregister(event_loop);
                    return true;
                }
            },
            Ok(None) => {
                debug!("Done reading for token {:?}", conn.token);
                return true;
            },
            Err(e) => {
                warn!("Read failure {:?} on token {:?}", e, conn.token);
                conn.dead = true;
                return true;
            }
        }
    }
}

fn write1(conn: &mut Conn, peer: &mut Conn, event_loop: &mut EventLoop<Nexus>) -> bool {
    if peer.buf.is_empty() {
        if conn.dead && conn.buf.is_empty() {
            return false;
        }
    }

    loop {

        if peer.buf.is_empty() {
            if !peer.interest.is_readable() {
                peer.interest = peer.interest | EventSet::readable();
                peer.reregister(event_loop);
            }
            return true;
        }

        match conn.sock.try_write_buf(&mut peer.buf) {
            Ok(Some(n)) => {
                debug!("Successfully wrote {} bytes to token {:?}", n, conn.token);
                if n == 0 {
                    if !peer.interest.is_readable() {
                        peer.interest = peer.interest | EventSet::readable();
                        peer.reregister(event_loop);
                    }
                    //conn.writable = false;
                    return true;
                }
            },
            Ok(None) => {
                debug!("Write not accepted");
                conn.writable = false;
                return true;
            },
            Err(e) => {
                warn!("Write failure {:?} on token {:?}", e, conn.token);
                conn.writable = false;
                return false;
            }
        }
    }
}

impl Nexus {
    fn new(acceptor: TcpListener) -> Nexus {
        Nexus {
            acceptor: acceptor,

            // I don't use Token(0) because kqueue will send stuff to Token(0)
            // by default causing really strange behavior. This way, if I see
            // something as Token(0), I know there are kqueue shenanigans
            // going on.
            token: ACCEPTOR,

            // SERVER is Token(1), so start after that
            // we can deal with a max of 126 connections
            conns: Slab::new_starting_at(FLOW, 128)
        }
    }

    /// Register Server with the event loop.
    ///
    /// This keeps the registration details neatly tucked away inside of our implementation.
    fn register(&mut self, event_loop: &mut EventLoop<Nexus>) -> io::Result<()> {
        event_loop.register_opt(
            &self.acceptor,
            self.token,
            EventSet::readable(),
            PollOpt::edge(), // | PollOpt::oneshot()
        ).or_else(|e| {
            error!("Failed to register server {:?}, {:?}", self.token, e);
            Err(e)
        })
    }

    /// Accept a _new_ client connection.
    ///
    /// The server will keep track of the new connection and forward any events from the event loop
    /// to this connection.
    fn accept(&mut self, event_loop: &mut EventLoop<Nexus>) {
        debug!("tnexus accepting new socket(s)");

        loop {
        // Log an error if there is no socket, but otherwise move on so we do not tear down the
        // entire server.
        let inbound = match self.acceptor.accept() {
            Ok(s) => {
                match s {
                    Some(sock) => sock,
                    None => {
                        error!("Failed to accept new socket");
                        //self.reregister(event_loop);
                        return;
                    }
                }
            },
            Err(e) => {
                error!("Failed to accept new socket, {:?}", e);
                //self.reregister(event_loop);
                return;
            }
        };

        let dst_addr: SocketAddr = DESTINATION.parse()
            .ok().expect("Failed to parse destination endpoint");

        let outbound = TcpStream::connect(&dst_addr).ok().expect("outbound failure not handled yet");

        match self.conns.insert_with(|token| {
            debug!("Inserting {:?} into slab", token);
            Flow::new(inbound, outbound, token)
        }) {
            Some(token) => {
                match self.find_connection_by_token(token).register(event_loop) {
                    Ok(_) => {
                        debug!("Registered inbound token {:?}", token);

                    },
                    Err(e) => {
                        error!("Failed to register inbound {:?} connection with event loop, {:?}", token, e);
                        self.conns.remove(token);
                    }
                }
            },
            None => {
                // If we fail to insert, the socks will go out of scope and be dropped.
                error!("Failed to insert into slab");
            }
        };

        }
        // We are using edge-triggered polling. Even our SERVER token needs to reregister.
        // self.reregister(event_loop);
    }

    fn stop_flow(&mut self, token: Token) {
        debug!("Stopping flow {:?}", token);
        self.conns.remove(token);
    }

    /// Find a connection in the slab using the given token.
    fn find_connection_by_token<'a>(&'a mut self, token: Token) -> &'a mut Flow {
        &mut self.conns[token]
    }

    /*
    fn with_flow<F, R>(&mut self, token: Token, mut closure: F) -> Option<R>
        where F: FnMut(&mut Flow) -> R {

        match self.conns.get_mut(token) {
            None => {
                warn!("Token not found: {:?}", token);
                None
            },
            Some(flow) => {
                Some(closure(flow))
            }
        }
    }
    */
}

impl Handler for Nexus {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Nexus>, _token: Token, events: EventSet) {
        assert!(_token != INVALID, "[BUG]: Received event for Token(0)");

        let tokval = _token.as_usize();
        let inb = tokval & OUTMASK == 0;
        let token = Token(tokval & NOTMASK);

        debug!("READY {:X} {} {:?} {:?}", tokval, inb, token, events);

        //if events.is_error() || events.is_hup() {
        if events.is_error() {
            warn!("Final event {:?} for token {}", events, tokval);
            //self.disable_flow(token, event_loop);
            return;
        }

        // We never expect a write event for our `Server` token . A write event for any other token
        // should be handed off to that connection.
        if events.is_writable() {
            debug!("Write event for {:X}", tokval);
            assert!(self.token != token, "Received writable event for Server");

            let write_result = match self.conns.get_mut(token) {
                None => {
                    warn!("Unknown token {:X}", tokval);
                    return;
                },
                Some(flow) => {
                    flow.write(inb, event_loop)
                }
            };

            if !write_result {
                self.stop_flow(token);
                return;
            }
        }

        // A read event for our `Server` token means we are establishing a new connection. A read
        // event for any other token should be handed off to that connection.
        if events.is_readable() {
            if ACCEPTOR == token {
                self.accept(event_loop);
            } else {
                if !self.find_connection_by_token(token).read(inb, event_loop) {
                    self.stop_flow(token);
                }
            }
        }
    }
}

fn main() {
    env_logger::init().ok().expect("Failed to init logger");

    info!("Starting tnexus...");

    let endpoint_addr: SocketAddr = ENDPOINT.parse()
        .ok().expect("Failed to parse server enpoint");

    // Setup the acceptor socket
    let acceptor = TcpListener::bind(&endpoint_addr)
        .ok().expect("Failed to bind server endpoint");

    // Create an event loop
    let mut event_loop = EventLoop::new()
        .ok().expect("Could not initialize MIO event loop");

    let mut nexus = Nexus::new(acceptor);

    // Start listening for incoming connections
    nexus.register(&mut event_loop)
        .ok().expect("Failed to register acceptor in event loop");

    // Start handling events
    event_loop.run(&mut nexus)
        .ok().expect("Failed to start event loop");

    info!("Over.");
}
<commit_msg>Various minor cleanups.<commit_after>extern crate mio;

#[macro_use]
extern crate log;

extern crate env_logger;

use std::io;
use std::net::SocketAddr;
use mio::*;
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use mio::buf::RingBuf;

const ENDPOINT: &'static str = "127.0.0.1:6666";
const DESTINATION: &'static str = "127.0.0.1:80";

//const BUF_SIZE: usize = 524288; //131072;
//const BUF_SIZE: usize = 8388608;
const BUF_SIZE: usize = 1048576;

const INVALID: Token = Token(0);
const ACCEPTOR: Token = Token(1);
const FLOW: Token = Token(2);

const OUTMASK: usize = 2147483648; //2 ** 31
const NOTMASK: usize = !OUTMASK;

struct Conn {
    // socket
    sock: TcpStream,

    // token used to register the inbound with the event loop
    token: Token,

    // set of events we are interested in
    interest: EventSet,

    // whether the remote peer has half-closed the socket
    // i.e. reading makes no more sense
    dead: bool,

    // whether the socket can accept writes right now
    writable: bool,

    // octets in transit
    buf: RingBuf,
}

/// A stateful wrapper around a non-blocking stream. This connection is not
/// the SERVER connection. This connection represents the client connections
/// _accepted_ by the SERVER connection.
struct Flow {

    // inbound connection
    inb: Conn,

    // outbound connection
    out: Conn,
}

struct Nexus {
    // main server socket (that accepts inbound connections)
    acceptor: TcpListener,

    // token of the acceptor
    token: Token,

    // a list of all inbound and outbound connections
    conns: Slab<Flow>,
}

impl Conn {
    fn new(token: Token, sock: TcpStream) -> Conn {
        Conn {
            sock: sock,
            token: token,
            interest: EventSet::all(),
            dead: false,
            writable: false,
            buf: RingBuf::new(BUF_SIZE),
        }
    }

    /// Re-register flow interest in read events with the event_loop.
    #[inline]
    fn reregister(&mut self, event_loop: &mut EventLoop<Nexus>)
                  -> io::Result<()> {

        debug!("Reregistering {:?} for {:?}", self.token, self.interest);
        event_loop.reregister(
            &self.sock, self.token, self.interest, PollOpt::edge()).or_else(|e| {
                error!("Failed to reregister {:?} for {:?} due to {:?}", self.token, self.interest, e);
                Err(e)
            })
    }
}

impl Flow {

    fn new(insock: TcpStream, ousock: TcpStream, token: Token) -> Flow {
        Flow {
            inb: Conn::new(token, insock),
            out: Conn::new(Token(OUTMASK + token.as_usize()), ousock),
        }
    }

    /// Register flow interest in read events with the event_loop.
    ///
    /// This will let our connection accept reads starting next event loop tick.
    fn register(&mut self, event_loop: &mut EventLoop<Nexus>)
                -> io::Result<()> {

        event_loop.register_opt(
            &self.inb.sock,
            self.inb.token,
            self.inb.interest,
            PollOpt::edge()).and_then(|_|
        event_loop.register_opt(
            &self.out.sock,
            self.out.token,
            self.out.interest,
            PollOpt::edge())).
        or_else(|e| {
            error!("Failed to register {:?} due to {:?}", self.inb.token, e);
            Err(e)
        })
    }

    /// Handle flow read event from event loop.
    ///
    #[inline]
    fn read(&mut self, inbo: bool, event_loop: &mut EventLoop<Nexus>) -> bool {

        let conn;
        let peer;
        if inbo {
            conn = &mut self.inb;
            peer = &mut self.out;
        } else {
            conn = &mut self.out;
            peer = &mut self.inb;
        }

        read1(conn, peer, event_loop)
    }

    #[inline]
    fn write(&mut self, inbo: bool, event_loop: &mut EventLoop<Nexus>) -> bool {
        let conn;
        let peer;
        if inbo {
            conn = &mut self.inb;
            peer = &mut self.out;
        } else {
            conn = &mut self.out;
            peer = &mut self.inb;
        }

        conn.writable = true;

        write1(conn, peer, event_loop)
    }
}

fn read1(conn: &mut Conn, peer: &mut Conn, event_loop: &mut EventLoop<Nexus>) -> bool {
    loop {
        match conn.sock.try_read_buf(&mut conn.buf) {
            Ok(Some(0)) => {
                debug!("Successful read of zero bytes (i.e. EOF) from token {:?}", conn.token);
                conn.dead = true;
                if peer.dead == true && conn.buf.is_empty() && peer.buf.is_empty() {
                	//we're all done:
                	debug!("All done {:?}", conn.token);
                    return false;
                } else {
                    debug!("Suspending reads due to EOF for {:?}", conn.token);
                    conn.interest = conn.interest & !EventSet::readable();
                    conn.reregister(event_loop);
                    return true;
                }
            },
            Ok(Some(n)) => {
                debug!("Successfully read {} bytes from {:?}, buf size: {}",
                       n, conn.token, <RingBuf as Buf>::remaining(&conn.buf));
                if peer.writable {
                    write1(peer, conn, event_loop);
                }
                if conn.buf.is_full() {
                    warn!("Suspending reads due to buffer full for {:?}", conn.token);
                    conn.interest = conn.interest & !EventSet::readable();
                    conn.reregister(event_loop);
                    return true;
                }
            },
            Ok(None) => {
                debug!("Done reading for token {:?}", conn.token);
                return true;
            },
            Err(e) => {
                warn!("Read failure {:?} on token {:?}", e, conn.token);
                conn.dead = true;
                conn.interest = conn.interest & !EventSet::readable();
                conn.reregister(event_loop);
	            return true;
            }
        }
    }
}

fn write1(conn: &mut Conn, peer: &mut Conn, event_loop: &mut EventLoop<Nexus>) -> bool {
    if peer.buf.is_empty() {
        if conn.dead && conn.buf.is_empty() {
            return false;
        }
    }

    loop {

        if peer.buf.is_empty() {
            if !peer.interest.is_readable() {
                peer.interest = peer.interest | EventSet::readable();
                peer.reregister(event_loop);
            }
            return true;
        }

        match conn.sock.try_write_buf(&mut peer.buf) {
            Ok(Some(n)) => {
                debug!("Successfully wrote {} bytes to token {:?}", n, conn.token);
            },
            Ok(None) => {
                debug!("Write not accepted");
                conn.writable = false;
                return true;
            },
            Err(e) => {
                warn!("Write failure {:?} on token {:?}", e, conn.token);
                conn.writable = false;
                return false;
            }
        }
    }
}

impl Nexus {
    fn new(acceptor: TcpListener) -> Nexus {
        Nexus {
            acceptor: acceptor,

            // I don't use Token(0) because kqueue will send stuff to Token(0)
            // by default causing really strange behavior. This way, if I see
            // something as Token(0), I know there are kqueue shenanigans
            // going on.
            token: ACCEPTOR,

            // SERVER is Token(1), so start after that
            // we can deal with a max of 126 connections
            conns: Slab::new_starting_at(FLOW, 128)
        }
    }

    /// Register Server with the event loop.
    ///
    /// This keeps the registration details neatly tucked away inside of our implementation.
    fn register(&mut self, event_loop: &mut EventLoop<Nexus>) -> io::Result<()> {
        event_loop.register_opt(
            &self.acceptor,
            self.token,
            EventSet::readable(),
            PollOpt::edge(), // | PollOpt::oneshot()
        ).or_else(|e| {
            error!("Failed to register server {:?}, {:?}", self.token, e);
            Err(e)
        })
    }

    /// Accept a _new_ client connection.
    ///
    /// The server will keep track of the new connection and forward any events from the event loop
    /// to this connection.
    fn accept(&mut self, event_loop: &mut EventLoop<Nexus>) {
        debug!("tnexus accepting new socket(s)");

        loop {
        // Log an error if there is no socket, but otherwise move on so we do not tear down the
        // entire server.
        let inbound = match self.acceptor.accept() {
            Ok(s) => {
                match s {
                    Some(sock) => sock,
                    None => {
                        debug!("No more new sockets");
                        return;
                    }
                }
            },
            Err(e) => {
                error!("Failed to accept new socket, {:?}", e);
                return;
            }
        };

        let dst_addr: SocketAddr = DESTINATION.parse()
            .ok().expect("Failed to parse destination endpoint");

        let outbound = TcpStream::connect(&dst_addr).ok().expect("outbound failure not handled yet");

        match self.conns.insert_with(|token| {
            debug!("Inserting {:?} into slab", token);
            Flow::new(inbound, outbound, token)
        }) {
            Some(token) => {
                match self.find_connection_by_token(token).register(event_loop) {
                    Ok(_) => {
                        debug!("Registered inbound token {:?}", token);

                    },
                    Err(e) => {
                        error!("Failed to register inbound {:?} connection with event loop, {:?}", token, e);
                        self.conns.remove(token);
                    }
                }
            },
            None => {
                // If we fail to insert, the socks will go out of scope and be dropped.
                error!("Failed to insert into slab");
            }
        };

        }
    }

    fn stop_flow(&mut self, token: Token) {
        debug!("Stopping flow {:?}", token);
        self.conns.remove(token);
    }

    /// Find a connection in the slab using the given token.
    fn find_connection_by_token<'a>(&'a mut self, token: Token) -> &'a mut Flow {
        &mut self.conns[token]
    }

    /*
    fn with_flow<F, R>(&mut self, token: Token, mut closure: F) -> Option<R>
        where F: FnMut(&mut Flow) -> R {

        match self.conns.get_mut(token) {
            None => {
                warn!("Token not found: {:?}", token);
                None
            },
            Some(flow) => {
                Some(closure(flow))
            }
        }
    }
    */
}

impl Handler for Nexus {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Nexus>, _token: Token, events: EventSet) {
        assert!(_token != INVALID, "[BUG]: Received event for Token(0)");

        let tokval = _token.as_usize();
        let inb = tokval & OUTMASK == 0;
        let token = Token(tokval & NOTMASK);

        debug!("READY {:X} {} {:?} {:?}", tokval, inb, token, events);

        if events.is_error() {
            warn!("Final event {:?} for token {:X}", events, tokval);
        }

        // We never expect a write event for our `Server` token . A write event for any other token
        // should be handed off to that connection.
        if events.is_writable() {
            debug!("Write event for {:X}", tokval);
            assert!(self.token != token, "Received writable event for Server");

            let write_result = match self.conns.get_mut(token) {
                None => {
                    warn!("Unknown token {:X}", tokval);
                    return;
                },
                Some(flow) => {
                    flow.write(inb, event_loop)
                }
            };

            if !write_result {
                self.stop_flow(token);
                return;
            }
        }

        // A read event for our `Server` token means we are establishing a new connection. A read
        // event for any other token should be handed off to that connection.
        if events.is_readable() {
            if ACCEPTOR == token {
                self.accept(event_loop);
            } else {
                if !self.find_connection_by_token(token).read(inb, event_loop) {
                    self.stop_flow(token);
                }
            }
        }
    }
}

fn main() {
    env_logger::init().ok().expect("Failed to init logger");

    info!("Starting tnexus...");

    let endpoint_addr: SocketAddr = ENDPOINT.parse()
        .ok().expect("Failed to parse server enpoint");

    // Setup the acceptor socket
    let acceptor = TcpListener::bind(&endpoint_addr)
        .ok().expect("Failed to bind server endpoint");

    // Create an event loop
    let mut event_loop = EventLoop::new()
        .ok().expect("Could not initialize MIO event loop");

    let mut nexus = Nexus::new(acceptor);

    // Start listening for incoming connections
    nexus.register(&mut event_loop)
        .ok().expect("Failed to register acceptor in event loop");

    // Start handling events
    event_loop.run(&mut nexus)
        .ok().expect("Failed to start event loop");

    info!("Over.");
}
<|endoftext|>