<commit_before>use std::collections::BTreeSet;
use std::io;
use std::process;

const NSQRT: usize = 3;
const N: usize = NSQRT * NSQRT;
const NSQ: usize = N*N;

fn error(msg: String) -> ! {
    println!("error: {}", msg);
    process::exit(1);
}

/// Return the 0-based row index of `cell`.
fn row(cell: usize) -> usize {
    cell / N
}

/// Return the 0-based row index of `cell`.
fn col(cell: usize) -> usize {
    cell % N
}

/// Return the 0-based index of the upper-left cell of `cell`'s group.
fn group(cell: usize) -> usize {
    let r = row(cell);
    let c = col(cell);
    (N * (r - r % NSQRT)) + (c - c % NSQRT)
}

/// Return the neighbors (indices) of `cell`:
/// - The cells on the same row;
/// - The cells on the same column;
/// - The same in the same group.
/// Note: `cell` is not a neighbor of itself.
fn neighbors_of(cell: usize) -> Vec<usize> {
    let mut all_neighbors: BTreeSet<usize> = BTreeSet::new();

    // Neighbors in row and column
    for i in 0..N {
        all_neighbors.insert((N * row(cell)) + i);
        all_neighbors.insert((N * i) + col(cell));
    }

    // Neighbors in group
    let leader = group(cell);
    for r in row(leader) .. row(leader) + NSQRT {
        for c in col(leader) .. col(leader) + NSQRT {
            all_neighbors.insert(N * r + c);
        }
    }

    all_neighbors.remove(&cell);
    return all_neighbors.into_iter().collect();
}

type CandidateSet = u32;
const EMPTY_SET: CandidateSet = 0;
const FULL_SET: CandidateSet = 0x1FF;

/// A sudoku board is represented by a vector of u32's.
struct SudokuBoard {
    cells: Vec<CandidateSet>,
    neighbors: Vec<Vec<usize>>,
}

impl SudokuBoard {
    /// Create a new sudoku board from a string.
    /// A non-zero digit stands for itself,
    /// a dot stands for a blank cell,
    /// anything else is an error.
    fn from_str(digits: &str) -> Self {
        if digits.len() != NSQ {
            error(format!("invalid puzzle length; expected {}, got {}",
                          NSQ, digits.len()));
        }
        let mut cells = Vec::with_capacity(NSQ);
        for d in digits.chars() {
            match d {
                '.' => {
                    cells.push(FULL_SET);
                }
                '1' ... '9' => {
                    let n = d.to_digit(10).unwrap() as usize;
                    cells.push(1 << (n - 1));
                }
                _ => { error(format!("invalid digit ({:?}) in string", d)); }
            }
        }

        let mut neighbors: Vec<Vec<usize>> = Vec::with_capacity(NSQ);
        for i in 0 .. NSQ {
            neighbors.push(neighbors_of(i));
        }

        return SudokuBoard { cells, neighbors };
    }

    /// A cell is solved if its set of candidates is a singleton.
    fn cell_solved(&self, cell: usize) -> bool {
        self.cells[cell].count_ones() == 1
    }

    /// A cell is solvable if it has at least one candidate.
    fn cell_solvable(&self, cell: usize) -> bool {
        self.cells[cell] != 0
    }

    /// The board is solved if all cells are solved.
    fn solved(&self) -> bool {
        (0 .. NSQ).all(|cell| self.cell_solved(cell))
    }

    /// The board is solvable is all cells are solvable.
    fn solvable(&self) -> bool {
        (0 .. NSQ).all(|cell| self.cell_solvable(cell))
    }

    /// The non-candidates of cells are the solved values in
    /// the cell's neighbors.
    fn non_candidates(&self, cell: usize) -> u32 {
        let mut set: u32 = EMPTY_SET;
        for &n in self.neighbors[cell].iter() {
            // Branchless version of:
            // `if self.cell_solved(n) { set |= self.cells[n]; }`
            set |= self.cells[n] * (self.cell_solved(n) as u32);
        }
        set
    }

    /// Remove non-candidates from the cells of the board
    /// until a fixed point is reached, i.e., no more non-
    /// candidates can be removed anymore.
    fn propagate(&self) -> Self {
        let mut output = SudokuBoard {
            cells: self.cells.clone(),
            neighbors: self.neighbors.clone()
        };
        loop {
            let mut candidates_changed = false;
            for i in 0 .. NSQ {
                let q = output.cells[i] & !output.non_candidates(i);
                candidates_changed = candidates_changed || (q != output.cells[i]);
                output.cells[i] = q;
            }
            if !candidates_changed {
                break;
            }
        }
        return output;
    }

    /// Find the index of the unsolved cell with the
    /// fewest number of candidates.  Helps to speed
    /// up the solving process by making the search tree
    /// narrower.
    fn most_promising(&self) -> Option<usize> {
        let mut min_len = N;
        let mut min_index = NSQ;

        for i in 0 .. NSQ {
            if self.cell_solved(i) {
                continue;
            }
            let len = self.cells[i].count_ones() as usize;
            if len < min_len {
                min_index = i;
                min_len = len;
            }
        }

        if min_index == NSQ {
            None
        } else {
            Some(min_index)
        }
    }

    /// Solve the Sudoku board:
    /// 1. Propagate the set constraints
    /// 2a. If the board is solved, terminate.
    /// 2b. If the board is unsolvable, backtrack.
    /// 3. Pick the most promising cell and brute-force it.
    fn solve(&self) -> Option<Self> {
        let mut newboard = self.propagate();

        if newboard.solved() { return Some(newboard); }

        if !newboard.solvable() { return None; }

        if let Some(cell) = newboard.most_promising() {
            let cell_candidates = newboard.cells[cell];

            for c in 0 .. N {
                if cell_candidates & (1 << c) == 0 {
                    continue;
                }

                newboard.cells[cell] = 1 << c;
                match newboard.solve() {
                    Some(solved_board) => { return Some(solved_board); }
                    None => { }
                }
            }
        }

        return None;
    }

    /// Convert the board to a linear textual representation.
    fn to_str(&self) -> String {
        let mut output = String::with_capacity(NSQ);
        for i in 0 .. NSQ {
            if self.cell_solved(i) {
                output.push_str(&format!("{}", get_singleton(self.cells[i])));
            } else {
                output.push('.');
            }
        }
        output
    }
}

fn get_singleton(mut s: CandidateSet) -> u32 {
    let mut i = 0;
    while s != 0 {
        i += 1;
        s >>= 1;
    }
    return i;
}

fn main() {
    let stdin = io::stdin();
    let mut buf = String::with_capacity(NSQ);

    loop {
        buf.clear();
        match stdin.read_line(&mut buf) {
            Err(e) => { error(format!("I/O error, {:?}", e)); }
            Ok(0) => { return; }
            Ok(_) => { /* pass through */ }
        }
        let sb = SudokuBoard::from_str(&buf.trim());
        match sb.solve() {
            Some(solution) => { println!("{}", solution.to_str()); }
            None => { println!("No solution"); }
        }
    }
}

#[test]
fn test_row_col() {
    assert_eq!(row(11), 1);
    assert_eq!(col(11), 2);
}

#[test]
fn test_group() {
    assert_eq!(group(0), 0);
    assert_eq!(group(1), 0);
    assert_eq!(group(2), 0);
    assert_eq!(group(9), 0);
    assert_eq!(group(10), 0);
    assert_eq!(group(11), 0);
    assert_eq!(group(18), 0);
    assert_eq!(group(19), 0);
    assert_eq!(group(20), 0);
    assert_eq!(group(60), 60);
    assert_eq!(group(61), 60);
    assert_eq!(group(62), 60);
    assert_eq!(group(69), 60);
    assert_eq!(group(70), 60);
    assert_eq!(group(71), 60);
    assert_eq!(group(78), 60);
    assert_eq!(group(79), 60);
    assert_eq!(group(80), 60);
}
<commit_msg>Minor cleanups<commit_after>use std::collections::BTreeSet;
use std::io;
use std::process;

// Sudoku board constants
const NSQRT: usize = 3;
const N: usize = NSQRT * NSQRT;
const NSQ: usize = N*N;

// Set constants
type CandidateSet = u32;
const EMPTY_SET: CandidateSet = 0;
const FULL_SET: CandidateSet = 0x1FF;

fn error(msg: String) -> ! {
    println!("error: {}", msg);
    process::exit(1);
}

/// Return the 0-based row index of `cell`.
fn row(cell: usize) -> usize {
    cell / N
}

/// Return the 0-based row index of `cell`.
fn col(cell: usize) -> usize {
    cell % N
}

/// Return the 0-based index of the upper-left cell of `cell`'s group.
fn group(cell: usize) -> usize {
    let r = row(cell);
    let c = col(cell);
    (N * (r - r % NSQRT)) + (c - c % NSQRT)
}

/// Return the neighbors (indices) of `cell`:
/// - The cells on the same row;
/// - The cells on the same column;
/// - The same in the same group.
/// Note: `cell` is not a neighbor of itself.
fn neighbors_of(cell: usize) -> Vec<usize> {
    let mut all_neighbors: BTreeSet<usize> = BTreeSet::new();

    // Neighbors in row and column
    for i in 0..N {
        all_neighbors.insert((N * row(cell)) + i);
        all_neighbors.insert((N * i) + col(cell));
    }

    // Neighbors in group
    let leader = group(cell);
    for r in row(leader) .. row(leader) + NSQRT {
        for c in col(leader) .. col(leader) + NSQRT {
            all_neighbors.insert(N * r + c);
        }
    }

    all_neighbors.remove(&cell);
    return all_neighbors.into_iter().collect();
}

/// A sudoku board is represented by a vector of u32's.
struct SudokuBoard {
    cells: Vec<CandidateSet>,
    neighbors: Vec<Vec<usize>>,
}

impl SudokuBoard {
    /// Create a new sudoku board from a string.
    /// A non-zero digit stands for itself,
    /// a dot stands for a blank cell,
    /// anything else is an error.
    fn from_str(digits: &str) -> Self {
        if digits.len() != NSQ {
            error(format!("invalid puzzle length; expected {}, got {}",
                          NSQ, digits.len()));
        }
        let mut cells = Vec::with_capacity(NSQ);
        for d in digits.chars() {
            match d {
                '.' => {
                    cells.push(FULL_SET);
                }
                '1' ... '9' => {
                    let n = d.to_digit(10).unwrap() as usize;
                    cells.push(1 << (n - 1));
                }
                _ => { error(format!("invalid digit ({:?}) in string", d)); }
            }
        }

        // Neighbor indices never change, so we compute them once,
        // and store them in the struct.
        let mut neighbors: Vec<Vec<usize>> = Vec::with_capacity(NSQ);
        for i in 0 .. NSQ {
            neighbors.push(neighbors_of(i));
        }

        return SudokuBoard { cells, neighbors };
    }

    /// A cell is solved if its set of candidates is a singleton.
    fn cell_solved(&self, cell: usize) -> bool {
        self.cells[cell].count_ones() == 1
    }

    /// A cell is solvable if it has at least one candidate.
    fn cell_solvable(&self, cell: usize) -> bool {
        self.cells[cell] != 0
    }

    /// The board is solved if all cells are solved.
    fn solved(&self) -> bool {
        (0 .. NSQ).all(|cell| self.cell_solved(cell))
    }

    /// The board is solvable is all cells are solvable.
    fn solvable(&self) -> bool {
        (0 .. NSQ).all(|cell| self.cell_solvable(cell))
    }

    /// The non-candidates of a cell are the solved values in
    /// the cell's neighbors.
    fn non_candidates(&self, cell: usize) -> u32 {
        let mut set: u32 = EMPTY_SET;
        for &n in self.neighbors[cell].iter() {
            set |= self.cells[n] * (self.cell_solved(n) as u32);
        }
        return set;
    }

    /// Remove non-candidates from the cells of the board
    /// until a fixed point is reached, i.e., no more non-
    /// candidates can be removed anymore.
    fn propagate(&self) -> Self {
        let mut output = SudokuBoard {
            cells: self.cells.clone(),
            neighbors: self.neighbors.clone()
        };
        loop {
            let mut candidates_changed = false;
            for i in 0 .. NSQ {
                let q = output.cells[i] & !output.non_candidates(i);
                candidates_changed = candidates_changed || (q != output.cells[i]);
                output.cells[i] = q;
            }
            if !candidates_changed {
                break;
            }
        }
        return output;
    }

    /// Find the index of the unsolved cell with the
    /// fewest number of candidates.  Helps to speed
    /// up the solving process by making the search tree
    /// narrower.
    fn most_promising(&self) -> Option<usize> {
        let mut min_len = N;
        let mut min_index = NSQ;

        for i in 0 .. NSQ {
            if self.cell_solved(i) {
                continue;
            }
            let len = self.cells[i].count_ones() as usize;
            if len < min_len {
                min_index = i;
                min_len = len;
            }
        }

        if min_index == NSQ {
            None
        } else {
            Some(min_index)
        }
    }

    /// Solve the Sudoku board:
    /// 1. Propagate the set constraints
    /// 2a. If the board is solved, terminate.
    /// 2b. If the board is unsolvable, backtrack.
    /// 3. Pick the most promising cell and brute-force it.
    fn solve(&self) -> Option<Self> {
        let mut newboard = self.propagate();

        if newboard.solved() { return Some(newboard); }

        if !newboard.solvable() { return None; }

        if let Some(cell) = newboard.most_promising() {
            let cell_candidates = newboard.cells[cell];

            for c in 0 .. N {
                if cell_candidates & (1 << c) == 0 {
                    continue;
                }

                newboard.cells[cell] = 1 << c;
                match newboard.solve() {
                    Some(solved_board) => { return Some(solved_board); }
                    None => { }
                }
            }
        }

        return None;
    }

    /// Convert the board to a linear textual representation.
    fn to_str(&self) -> String {
        let mut output = String::with_capacity(NSQ);
        for i in 0 .. NSQ {
            if self.cell_solved(i) {
                output.push_str(&format!("{}", set_to_num(self.cells[i])));
            } else {
                output.push('.');
            }
        }
        output
    }
}

fn set_to_num(mut s: CandidateSet) -> u32 {
    let mut i = 0;
    while s != 0 {
        i += 1;
        s >>= 1;
    }
    return i;
}

fn main() {
    let stdin = io::stdin();
    let mut buf = String::with_capacity(NSQ);

    loop {
        buf.clear();
        match stdin.read_line(&mut buf) {
            Err(e) => { error(format!("I/O error, {:?}", e)); }
            Ok(0) => { return; }
            Ok(_) => { /* pass through */ }
        }
        let sb = SudokuBoard::from_str(&buf.trim());
        match sb.solve() {
            Some(solution) => { println!("{}", solution.to_str()); }
            None => { println!("No solution"); }
        }
    }
}

#[test]
fn test_row_col() {
    assert_eq!(row(11), 1);
    assert_eq!(col(11), 2);
}

#[test]
fn test_group() {
    assert_eq!(group(0), 0);
    assert_eq!(group(1), 0);
    assert_eq!(group(2), 0);
    assert_eq!(group(9), 0);
    assert_eq!(group(10), 0);
    assert_eq!(group(11), 0);
    assert_eq!(group(18), 0);
    assert_eq!(group(19), 0);
    assert_eq!(group(20), 0);
    assert_eq!(group(60), 60);
    assert_eq!(group(61), 60);
    assert_eq!(group(62), 60);
    assert_eq!(group(69), 60);
    assert_eq!(group(70), 60);
    assert_eq!(group(71), 60);
    assert_eq!(group(78), 60);
    assert_eq!(group(79), 60);
    assert_eq!(group(80), 60);
}
<|endoftext|>


<commit_before>#![allow(dead_code, deprecated)]

// To use String::as_str
#![feature(convert)]

// Dependencies for peg
#![feature(collections, str_char)]
#![feature(plugin)]
#![plugin(peg_syntax_ext)]

peg_file! wacc_parse("wacc.peg");
peg_file! wacc2_parse("wacc2.peg");

extern crate getopts;

mod ast;

use std::io::Read;
use std::fs::File;
use std::env;
use getopts::Options;

fn print_code(s: &str) {
    let lines: Vec<&str> = s.split("\n").collect();
    let mut counter = 1;
    for l in lines {
        println!("{} {}", counter, l);
        counter += 1;
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn compile(input: &String, use_new_standard: bool) {
    println!("Input:");
    print_code(&input);
    println!("");

    // Parse program
    let mut program = if use_new_standard {
        match wacc2_parse::program(&input) {
            Ok(p) => p,
            Err(s) => { println!("Syntax Error: {}", s); return }
        }
    } else {
        match wacc_parse::program(&input) {
            Ok(p) => p,
            Err(s) => { println!("Syntax Error: {}", s); return }
        }
    };
    println!("AST:\n{:?}", program);

    // Semantic check and derive types
    match ast::semantic::check_program(&mut program) {
        Ok(_) => {},
        Err(s) => { println!("Semantic Error: {}", s); return }
    };

    // Pass AST to the code generator selected
}

fn main() {
    // Collect command line args
    let args: Vec<_> = env::args().collect();
    let program = args[0].clone();

    // Parse options
    let mut opts = Options::new();
    opts.optopt("", "std","language standard, either 'new' or 'old'", "(new|old)");
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    // Display help
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let use_new_standard = match matches.opt_default("std", "new") {
        Some(s) => match s.as_str() {
            "old" => false,
            _ => true
        },
        None => true
    };

    // Read files
    if !matches.free.is_empty() {
        let mut f = File::open(matches.free[0].clone()).unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();
        compile(&s, use_new_standard);
    } else {
        print_usage(&program, opts);
        return;
    }
}
<commit_msg>Communicate the proper error code on exit<commit_after>#![allow(dead_code, deprecated)]

// To use String::as_str
#![feature(convert)]

// To use std::env::set_exit_status
#![feature(exit_status)]

// Dependencies for peg
#![feature(collections, str_char)]
#![feature(plugin)]
#![plugin(peg_syntax_ext)]

peg_file! wacc_parse("wacc.peg");
peg_file! wacc2_parse("wacc2.peg");

extern crate getopts;

mod ast;

use std::io::Read;
use std::fs::File;
use std::env;
use getopts::Options;

#[derive(PartialEq)]
enum Stage
{
    AST
}

#[derive(PartialEq)]
enum CompileError
{
    Syntax,
    Semantic
}

fn print_code(s: &str) {
    let lines: Vec<&str> = s.split("\n").collect();
    let mut counter = 1;
    for l in lines {
        println!("{} {}", counter, l);
        counter += 1;
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options] <inputs>", program);
    print!("{}", opts.usage(&brief));
}

fn compile(input: &String, use_new_syntax: bool, run_to_stage: Stage) -> Result<(), CompileError> {
    println!("Input:");
    print_code(&input);
    println!("");

    // Parse program
    let mut program = if use_new_syntax {
        match wacc2_parse::program(&input) {
            Ok(p)  => p,
            Err(s) => { println!("Syntax Error: {}", s); return Err(CompileError::Syntax) }
        }
    } else {
        match wacc_parse::program(&input) {
            Ok(p)  => p,
            Err(s) => { println!("Syntax Error: {}", s); return Err(CompileError::Syntax) }
        }
    };

    // Semantic check and derive types
    match ast::semantic::check_program(&mut program) {
        Ok(_)  => {},
        Err(s) => { println!("Semantic Error: {}", s); return Err(CompileError::Semantic) }
    };

    // Should we stop here?
    if run_to_stage == Stage::AST {
        println!("AST:\n{:?}", program);
        return Ok(());
    }

    // Pass AST to the code generator selected

    // All OK
    Ok(())
}

fn main() {
    // Collect command line args
    let args: Vec<_> = env::args().collect();
    let program = args[0].clone();

    // Parse options
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optopt("", "syntax","language syntax to use - must be either 'wacc' or 'wacc2'", "<type>");
    opts.optflag("", "ast", "display the abstract syntax tree and exit");
    let matches = match opts.parse(&args[1..]) {
        Ok(m)  => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    // Display help
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    // Read options
    let use_new_syntax = match matches.opt_default("syntax", "wacc2") {
        Some(s) => match s.as_str() {
            "wacc" => false,
            _     => true
        },
        None => true
    };
    let stage = if matches.opt_present("ast") { Stage::AST } else { Stage::AST };

    // Compile files
    if !matches.free.is_empty() {
        let mut f = File::open(matches.free[0].clone()).unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();
        if let Err(t) = compile(&s, use_new_syntax, stage) {
            match t {
                CompileError::Syntax   => std::env::set_exit_status(100),
                CompileError::Semantic => std::env::set_exit_status(200)
            }
        }
    } else {
        print_usage(&program, opts);
        return;
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]
#![cfg_attr(nightly, feature(scoped))]

extern crate unicode_width;

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

use version::*;
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;

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 {})", VERSION, TIMESTAMP);
        } else {
            println!("heatseeker {} ({}) (built {})", VERSION, COMMIT, TIMESTAMP);
        }
        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);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

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

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

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selections());
}

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

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

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

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

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

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

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 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.index = 0;
    }

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

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

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.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 {
            self.recompute_matches();
            let selection = self.matches.get(self.index).unwrap_or(&"").to_string();
            self.selections.insert(selection);
        }
        for selection in self.selections.iter() {
            ret.push_str(selection);
            ret.push('\n');
        }
        ret
    }

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

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

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

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

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

fn print_matches(screen: &mut Screen, matches: &[&str], query: &str, index: usize, selections: &HashSet<String>) {
    #[cfg(windows)] const SELECTED: char = '*';
    #[cfg(not(windows))] const SELECTED: char = '✓';
    let mut i = 1;
    for choice in matches.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(' ');
            annotated_choice.push(SELECTED);
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == 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;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut 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 chars_to_draw = min(chars_in_choice, max_width as usize);
    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();
    loop {
        let mut s = String::new();
        stdin.read_line(&mut s).unwrap();
        if s.len() == 0 {
            break;
        } else {
            trim(&mut s);
            lines.push(s);
        }
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            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.slice_unchecked(a, b) }
    }
}

#[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");
}

#[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>Strip CRLF on Windows<commit_after>#![cfg_attr(test, allow(dead_code))]
#![cfg_attr(nightly, feature(scoped))]

extern crate unicode_width;

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

use version::*;
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;

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 {})", VERSION, TIMESTAMP);
        } else {
            println!("heatseeker {} ({}) (built {})", VERSION, COMMIT, TIMESTAMP);
        }
        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);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

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

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

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selections());
}

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

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

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

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

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

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

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 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.index = 0;
    }

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

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

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.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 {
            self.recompute_matches();
            let selection = self.matches.get(self.index).unwrap_or(&"").to_string();
            self.selections.insert(selection);
        }
        for selection in self.selections.iter() {
            ret.push_str(selection);
            ret.push('\n');
        }
        ret
    }

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

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

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

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

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

fn print_matches(screen: &mut Screen, matches: &[&str], query: &str, index: usize, selections: &HashSet<String>) {
    #[cfg(windows)] const SELECTED: char = '*';
    #[cfg(not(windows))] const SELECTED: char = '✓';
    let mut i = 1;
    for choice in matches.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(' ');
            annotated_choice.push(SELECTED);
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == 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;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut 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 chars_to_draw = min(chars_in_choice, max_width as usize);
    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();
    loop {
        let mut s = String::new();
        stdin.read_line(&mut s).unwrap();
        if s.len() == 0 {
            break;
        } else {
            trim(&mut s);
            lines.push(s);
        }
    }

    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.slice_unchecked(a, b) }
    }
}

#[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>//! # Triple buffering
//!
//! In this crate, we attempt to implement a triple buffering mechanism,
//! suitable for emulating shared memory cells in a thread-safe and
//! non-blocking fashion between one single writer and one single reader.

use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// ## Indexing
///
/// These types are used to index into triple buffers
/// TODO: Switch to i8 / AtomicI8 once the later is stable
type TripleBufferIndex = usize;
type AtomicTripleBufferIndex = AtomicUsize;

/// ## Shared state
///
/// In a triple buffering communication protocol, the reader and writer share
/// the following storage:
///
/// - Three memory buffers suitable for storing the data at hand
/// - One index pointing to the buffer which no one is currently using
/// - One index pointing to the most recently updated buffer
///
#[derive(Debug)]
struct TripleBufferSharedState<T: Clone + PartialEq> {
    /// Three buffers of data
    buffers: [UnsafeCell<T>; 3],

    /// Index of the back-buffer (buffer which no one is currently accessing)
    back_idx: AtomicTripleBufferIndex,

    /// Index of the most recently updated buffer
    last_idx: AtomicTripleBufferIndex,
}
//
impl<T: Clone + PartialEq> Clone for TripleBufferSharedState<T> {
    fn clone(&self) -> Self {
        // The use of UnsafeCell makes buffers somwhat cumbersome to clone...
        let clone_buffer = | i: TripleBufferIndex | {
            UnsafeCell::new(
                unsafe { (*self.buffers[i].get()).clone() }
            )
        };

        // ...so better define a shortcut before we move on to actual cloning
        TripleBufferSharedState {
            buffers: [
                clone_buffer(0),
                clone_buffer(1),
                clone_buffer(2),
            ],
            back_idx: AtomicTripleBufferIndex::new(
                self.back_idx.load(Ordering::Relaxed)
            ),
            last_idx: AtomicTripleBufferIndex::new(
                self.last_idx.load(Ordering::Relaxed)
            ),
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBufferSharedState<T> {
    fn eq(&self, other: &Self) -> bool {
        // The use of UnsafeCell makes buffers somwhat cumbersome to compare...
        let buffer_eq = | i: TripleBufferIndex | -> bool {
            unsafe { *self.buffers[i].get() == *other.buffers[i].get() }
        };

        // ...and atomics aren't much better...
        let atomics_eq = | x: &AtomicTripleBufferIndex,
                           y: &AtomicTripleBufferIndex | -> bool {
            x.load(Ordering::Relaxed) == y.load(Ordering::Relaxed)
        };

        // ...so better define some shortcuts before getting started:
        buffer_eq(0) && buffer_eq(1) && buffer_eq(2) &&
        atomics_eq(&self.back_idx, &other.back_idx) &&
        atomics_eq(&self.last_idx, &other.last_idx)
    }
}
//
unsafe impl<T: Clone + PartialEq> Sync for TripleBufferSharedState<T> {}

/// ## Producing data ("input")
///
/// The producer of data has exclusive write access to the so-called "write
/// buffer". To send an update, it writes it into the back buffer, then swaps
/// the write buffer with the back buffer, making the data readable.
///
#[derive(Debug)]
struct TripleBufferInput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    write_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferInput<T> {
    /// Write a new value into the triple buffer
    fn write(&mut self, value: T) {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Determine which shared buffer is the write buffer
        let active_idx = self.write_idx;

        // Move the input value into the write buffer
        let write_ptr = shared_state.buffers[active_idx].get();
        unsafe { *write_ptr = value; }

        // Swap the back buffer and the write buffer. This makes the new data
        // block available to the reader and gives us a new buffer to write to.
        self.write_idx = shared_state.back_idx.swap(
            active_idx,
            Ordering::Release  // Send previous updates to the reader
        );

        // Notify the reader that we submitted an update
        shared_state.last_idx.store(active_idx, Ordering::Relaxed);
    }
}

/// ## Consuming data ("output")
///
/// The consumer of data has exclusive read access to the so-called "read
/// buffer". To access the latest data, it swaps the read buffer with the
/// back buffer if the later contains more recent data, then reads from
/// the (possibly new) read buffer.
///
#[derive(Debug)]
struct TripleBufferOutput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    read_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferOutput<T> {
    /// Access the latest fully written value from the triple buffer
    fn read(&mut self) -> &T {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Check if the writer has submitted an update
        let last_idx = shared_state.last_idx.load(Ordering::Relaxed);

        // If an update is pending in the back buffer, make it our read buffer
        if self.read_idx != last_idx {
            // Swap the back buffer and the read buffer. We get exclusive read
            // access to the data, and the writer gets a new back buffer.
            self.read_idx = shared_state.back_idx.swap(
                self.read_idx,
                Ordering::Acquire  // Synchronize with the writer's updates
            );
        }

        // Access data from the current read buffer
        let read_ptr = shared_state.buffers[self.read_idx].get();
        unsafe { &*read_ptr }
    }
}

/// ## Triple buffer
///
/// A triple buffer is created as the combination of an input and an output.
///
#[derive(Debug)]
struct TripleBuffer<T: Clone + PartialEq> {
    input: TripleBufferInput<T>,
    output: TripleBufferOutput<T>,
}
//
impl<T: Clone + PartialEq> TripleBuffer<T> {
    /// Construct a triple buffer with a certain initial value
    fn new(initial: T) -> Self {
        // Start with the shared state...
        let shared_state = Arc::new(
            TripleBufferSharedState {
                buffers: [
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial)
                ],
                back_idx: AtomicTripleBufferIndex::new(0),
                last_idx: AtomicTripleBufferIndex::new(2),
            }
        );

        // ...then construct the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: 1,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: 2,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> Clone for TripleBuffer<T> {
    fn clone(&self) -> Self {
        // Clone the shared state...
        let shared_state = Arc::new(
            (*self.input.shared).clone()
        );

        // ...then the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: self.input.write_idx,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: self.output.read_idx,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBuffer<T> {
    fn eq(&self, other: &Self) -> bool {
        (*self.input.shared == *other.input.shared) &&
        (self.input.write_idx == other.input.write_idx) &&
        (self.output.read_idx == other.output.read_idx)
    }
}

/// ## Testing
///
/// Unit tests and benchmarks for triple buffers
///
#[cfg(test)]
mod tests {
    /// Test that triple buffers are properly initialized
    #[test]
    fn test_init() {
        // Let's create a triple buffer
        let buf = ::TripleBuffer::new(42);

        // Access the shared state
        let ref buf_shared = *buf.input.shared;
        let back_idx = buf_shared.back_idx.load(::Ordering::Relaxed);
        let last_idx = buf_shared.last_idx.load(::Ordering::Relaxed);

        // Write-/read-/back-buffer indexes must be in range
        assert!(index_in_range(buf.input.write_idx));
        assert!(index_in_range(buf.output.read_idx));
        assert!(index_in_range(back_idx));

        // Write-/read-/back-buffer indexes must be distinct
        assert!(buf.input.write_idx != buf.output.read_idx);
        assert!(buf.input.write_idx != back_idx);
        assert!(buf.output.read_idx != back_idx);

        // Read buffer must be properly initialized
        let read_ptr = buf_shared.buffers[buf.output.read_idx].get();
        assert!(unsafe { *read_ptr } == 42);

        // Last-written index must initially point to the read buffer
        assert_eq!(last_idx, buf.output.read_idx);
    }

    /// Test that (sequentially) writing to a triple buffer works
    #[test]
    fn test_seq_write() {
        // Let's create a triple buffer
        let mut buf = ::TripleBuffer::new(false);

        // Back up the initial buffer state
        let old_buf = buf.clone();
        let ref old_shared = old_buf.input.shared;

        // Perform a write
        buf.input.write(true);

        // Check new implementation state
        {
            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the former write buffer to have received the new value
            let old_write_idx = old_buf.input.write_idx;
            let write_buf_ptr = expected_shared.buffers[old_write_idx].get();
            unsafe { *write_buf_ptr = true; }

            // We expect the former write buffer to be the new back buffer
            expected_shared.back_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the last updated index to point to former write buffer
            expected_shared.last_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the old back buffer to become the new write buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.input.write_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }
    }

    /// Check that (sequentially) reading from a triple buffer works
    #[test]
    fn test_seq_read() {
        // Let's create a triple buffer and write into it
        let mut buf = ::TripleBuffer::new(1.0);
        buf.input.write(4.2);

        // Test readout from dirty (freshly written) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();
            let ref old_shared = old_buf.input.shared;

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the new back index to point to the former read buffer
            expected_shared.back_idx.store(
                old_buf.output.read_idx,
                ::Ordering::Relaxed
            );

            // We expect the new read index to point to the former back buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.output.read_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }

        // Test readout from clean (unchanged) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Buffer state should be unchanged
            assert_eq!(buf, old_buf);
        }
    }

    // Check that concurrent reads and writes work
    #[test]
    fn test_conc_access() {
        // We will stress the infrastructure by performing this many writes
        // as a reader continuously reads the latest value
        let test_write_count = 1000;

        // This is the buffer that our reader and writer will share
        let mut buf = ::TripleBuffer::new(0u64);

        // Extract the input stage so that we can send it to the writer
        let mut buf_input = buf.input;

        // The writer continuously increments the buffered value, with some
        // rate limiting to ensure that the reader can see the updates
        let writer = ::thread::spawn(move || {
            // TODO: Use a barrier here instead
            for value in 1 .. test_write_count+1 {
                buf_input.write(value);
                ::thread::yield_now();
                ::thread::sleep(::Duration::from_millis(1));
            }
        });

        // The reader continuously checks the buffered value, and should see
        // every update without any incoherent value
        let mut last_value = 0u64;
        while last_value < test_write_count {
            let new_value = *buf.output.read();
            assert!(
                (new_value >= last_value) && (new_value-last_value <= 1)
            );
            last_value = new_value;
        }

        // Wait for the writer to finish
        writer.join().unwrap();
    }

    // TODO: Benchmark the performance of concurrent reads and writes

    /// Range check for triple buffer indexes
    #[allow(unused_comparisons)]
    fn index_in_range(idx: ::TripleBufferIndex) -> bool {
        (idx >= 0) & (idx <= 2)
    }
}

fn main() {
    println!("Hello, world!");

    let mut buf = TripleBuffer::new("Hello!");
    println!("Initial value is {}", buf.output.read());

    buf.input.write("World!");
    println!("New value is {}", buf.output.read())
}
<commit_msg>Changed last_idx ordering to Acq/Rel<commit_after>//! # Triple buffering
//!
//! In this crate, we attempt to implement a triple buffering mechanism,
//! suitable for emulating shared memory cells in a thread-safe and
//! non-blocking fashion between one single writer and one single reader.

use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// ## Indexing
///
/// These types are used to index into triple buffers
/// TODO: Switch to i8 / AtomicI8 once the later is stable
type TripleBufferIndex = usize;
type AtomicTripleBufferIndex = AtomicUsize;

/// ## Shared state
///
/// In a triple buffering communication protocol, the reader and writer share
/// the following storage:
///
/// - Three memory buffers suitable for storing the data at hand
/// - One index pointing to the buffer which no one is currently using
/// - One index pointing to the most recently updated buffer
///
#[derive(Debug)]
struct TripleBufferSharedState<T: Clone + PartialEq> {
    /// Three buffers of data
    buffers: [UnsafeCell<T>; 3],

    /// Index of the back-buffer (buffer which no one is currently accessing)
    back_idx: AtomicTripleBufferIndex,

    /// Index of the most recently updated buffer
    last_idx: AtomicTripleBufferIndex,
}
//
impl<T: Clone + PartialEq> Clone for TripleBufferSharedState<T> {
    fn clone(&self) -> Self {
        // The use of UnsafeCell makes buffers somwhat cumbersome to clone...
        let clone_buffer = | i: TripleBufferIndex | {
            UnsafeCell::new(
                unsafe { (*self.buffers[i].get()).clone() }
            )
        };

        // ...so better define a shortcut before we move on to actual cloning
        TripleBufferSharedState {
            buffers: [
                clone_buffer(0),
                clone_buffer(1),
                clone_buffer(2),
            ],
            back_idx: AtomicTripleBufferIndex::new(
                self.back_idx.load(Ordering::Relaxed)
            ),
            last_idx: AtomicTripleBufferIndex::new(
                self.last_idx.load(Ordering::Relaxed)
            ),
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBufferSharedState<T> {
    fn eq(&self, other: &Self) -> bool {
        // The use of UnsafeCell makes buffers somwhat cumbersome to compare...
        let buffer_eq = | i: TripleBufferIndex | -> bool {
            unsafe { *self.buffers[i].get() == *other.buffers[i].get() }
        };

        // ...and atomics aren't much better...
        let atomics_eq = | x: &AtomicTripleBufferIndex,
                           y: &AtomicTripleBufferIndex | -> bool {
            x.load(Ordering::Relaxed) == y.load(Ordering::Relaxed)
        };

        // ...so better define some shortcuts before getting started:
        buffer_eq(0) && buffer_eq(1) && buffer_eq(2) &&
        atomics_eq(&self.back_idx, &other.back_idx) &&
        atomics_eq(&self.last_idx, &other.last_idx)
    }
}
//
unsafe impl<T: Clone + PartialEq> Sync for TripleBufferSharedState<T> {}

/// ## Producing data ("input")
///
/// The producer of data has exclusive write access to the so-called "write
/// buffer". To send an update, it writes it into the back buffer, then swaps
/// the write buffer with the back buffer, making the data readable.
///
#[derive(Debug)]
struct TripleBufferInput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    write_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferInput<T> {
    /// Write a new value into the triple buffer
    fn write(&mut self, value: T) {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Determine which shared buffer is the write buffer
        let active_idx = self.write_idx;

        // Move the input value into the write buffer
        let write_ptr = shared_state.buffers[active_idx].get();
        unsafe { *write_ptr = value; }

        // Swap the back buffer and the write buffer. This makes the new data
        // block available to the reader and gives us a new buffer to write to.
        self.write_idx = shared_state.back_idx.swap(
            active_idx,
            Ordering::Release  // Need this for new buffer state to be visible
        );

        // Notify the reader that we submitted an update
        shared_state.last_idx.store(
            active_idx,
            Ordering::Release  // Need this for back_idx to be visible
        );
    }
}

/// ## Consuming data ("output")
///
/// The consumer of data has exclusive read access to the so-called "read
/// buffer". To access the latest data, it swaps the read buffer with the
/// back buffer if the later contains more recent data, then reads from
/// the (possibly new) read buffer.
///
#[derive(Debug)]
struct TripleBufferOutput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    read_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferOutput<T> {
    /// Access the latest fully written value from the triple buffer
    fn read(&mut self) -> &T {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Check if the writer has submitted an update
        let last_idx = shared_state.last_idx.load(Ordering::Acquire);

        // If an update is pending in the back buffer, make it our read buffer
        if self.read_idx != last_idx {
            // Swap the back buffer and the read buffer. We get exclusive read
            // access to the data, and the writer gets a new back buffer.
            self.read_idx = shared_state.back_idx.swap(
                self.read_idx,
                Ordering::Acquire
            );
        }

        // Access data from the current read buffer
        let read_ptr = shared_state.buffers[self.read_idx].get();
        unsafe { &*read_ptr }
    }
}

/// ## Triple buffer
///
/// A triple buffer is created as the combination of an input and an output.
///
#[derive(Debug)]
struct TripleBuffer<T: Clone + PartialEq> {
    input: TripleBufferInput<T>,
    output: TripleBufferOutput<T>,
}
//
impl<T: Clone + PartialEq> TripleBuffer<T> {
    /// Construct a triple buffer with a certain initial value
    fn new(initial: T) -> Self {
        // Start with the shared state...
        let shared_state = Arc::new(
            TripleBufferSharedState {
                buffers: [
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial)
                ],
                back_idx: AtomicTripleBufferIndex::new(0),
                last_idx: AtomicTripleBufferIndex::new(2),
            }
        );

        // ...then construct the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: 1,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: 2,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> Clone for TripleBuffer<T> {
    fn clone(&self) -> Self {
        // Clone the shared state...
        let shared_state = Arc::new(
            (*self.input.shared).clone()
        );

        // ...then the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: self.input.write_idx,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: self.output.read_idx,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBuffer<T> {
    fn eq(&self, other: &Self) -> bool {
        (*self.input.shared == *other.input.shared) &&
        (self.input.write_idx == other.input.write_idx) &&
        (self.output.read_idx == other.output.read_idx)
    }
}

/// ## Testing
///
/// Unit tests and benchmarks for triple buffers
///
#[cfg(test)]
mod tests {
    /// Test that triple buffers are properly initialized
    #[test]
    fn test_init() {
        // Let's create a triple buffer
        let buf = ::TripleBuffer::new(42);

        // Access the shared state
        let ref buf_shared = *buf.input.shared;
        let back_idx = buf_shared.back_idx.load(::Ordering::Relaxed);
        let last_idx = buf_shared.last_idx.load(::Ordering::Relaxed);

        // Write-/read-/back-buffer indexes must be in range
        assert!(index_in_range(buf.input.write_idx));
        assert!(index_in_range(buf.output.read_idx));
        assert!(index_in_range(back_idx));

        // Write-/read-/back-buffer indexes must be distinct
        assert!(buf.input.write_idx != buf.output.read_idx);
        assert!(buf.input.write_idx != back_idx);
        assert!(buf.output.read_idx != back_idx);

        // Read buffer must be properly initialized
        let read_ptr = buf_shared.buffers[buf.output.read_idx].get();
        assert!(unsafe { *read_ptr } == 42);

        // Last-written index must initially point to the read buffer
        assert_eq!(last_idx, buf.output.read_idx);
    }

    /// Test that (sequentially) writing to a triple buffer works
    #[test]
    fn test_seq_write() {
        // Let's create a triple buffer
        let mut buf = ::TripleBuffer::new(false);

        // Back up the initial buffer state
        let old_buf = buf.clone();
        let ref old_shared = old_buf.input.shared;

        // Perform a write
        buf.input.write(true);

        // Check new implementation state
        {
            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the former write buffer to have received the new value
            let old_write_idx = old_buf.input.write_idx;
            let write_buf_ptr = expected_shared.buffers[old_write_idx].get();
            unsafe { *write_buf_ptr = true; }

            // We expect the former write buffer to be the new back buffer
            expected_shared.back_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the last updated index to point to former write buffer
            expected_shared.last_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the old back buffer to become the new write buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.input.write_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }
    }

    /// Check that (sequentially) reading from a triple buffer works
    #[test]
    fn test_seq_read() {
        // Let's create a triple buffer and write into it
        let mut buf = ::TripleBuffer::new(1.0);
        buf.input.write(4.2);

        // Test readout from dirty (freshly written) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();
            let ref old_shared = old_buf.input.shared;

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the new back index to point to the former read buffer
            expected_shared.back_idx.store(
                old_buf.output.read_idx,
                ::Ordering::Relaxed
            );

            // We expect the new read index to point to the former back buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.output.read_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }

        // Test readout from clean (unchanged) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Buffer state should be unchanged
            assert_eq!(buf, old_buf);
        }
    }

    // Check that concurrent reads and writes work
    #[test]
    fn test_conc_access() {
        // We will stress the infrastructure by performing this many writes
        // as a reader continuously reads the latest value
        let test_write_count = 1000;

        // This is the buffer that our reader and writer will share
        let mut buf = ::TripleBuffer::new(0u64);

        // Extract the input stage so that we can send it to the writer
        let mut buf_input = buf.input;

        // The writer continuously increments the buffered value, with some
        // rate limiting to ensure that the reader can see the updates
        let writer = ::thread::spawn(move || {
            // TODO: Use a barrier here instead
            for value in 1 .. test_write_count+1 {
                buf_input.write(value);
                ::thread::yield_now();
                ::thread::sleep(::Duration::from_millis(1));
            }
        });

        // The reader continuously checks the buffered value, and should see
        // every update without any incoherent value
        let mut last_value = 0u64;
        while last_value < test_write_count {
            let new_value = *buf.output.read();
            assert!(
                (new_value >= last_value) && (new_value-last_value <= 1)
            );
            last_value = new_value;
        }

        // Wait for the writer to finish
        writer.join().unwrap();
    }

    // TODO: Benchmark the performance of concurrent reads and writes

    /// Range check for triple buffer indexes
    #[allow(unused_comparisons)]
    fn index_in_range(idx: ::TripleBufferIndex) -> bool {
        (idx >= 0) & (idx <= 2)
    }
}

fn main() {
    println!("Hello, world!");

    let mut buf = TripleBuffer::new("Hello!");
    println!("Initial value is {}", buf.output.read());

    buf.input.write("World!");
    println!("New value is {}", buf.output.read())
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate time;
extern crate image;
extern crate num_cpus;
extern crate threadpool;
extern crate rustc_serialize;

use clap::{Arg, App};
use time::PreciseTime;
use rustc_serialize::json;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::thread::{self, JoinHandle};

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Size {
    width: f32,
    height: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Point {
    x: f32,
    y: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Rect {
    origin: Point,
    size: Size,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Config {
    ppu: u32,
    limit: u32,
    color_steps: f32,
    color_components: u8,
    color_palette: Vec<Vec<f32>>,
    window: Rect,
}

fn idx2point(idx: u32, width: u32) -> Point {
    let x = idx % width;
    let y = idx / width;
    Point { x: x as f32, y: y as f32}
}

fn point2idx(p: Point, width: u32) -> u32 {
    width * (p.y as u32) + (p.x as u32)
}

fn mandelbrot(cx: f32, cy: f32, limit: u32) -> u32 {
    let mut x = cx;
    let mut y = cy;
    let mut count = 0;
    while count < limit {
        let xy = x * y;
        let xx = x * x;
        let yy = y * y;
        let sum = xx + yy;
        if sum > 4.0 {
            break
        }
        count += 1;
        x = xx - yy + cx;
        y = xy * 2.0 + cy;
    }
    count as u32
}

fn gen_mandelbrot(size: &Size, config: &Config) -> Vec<u32> {
    let window = &config.window;
    let limit = config.limit;

    let thread_count = (num_cpus::get() as f32 * 1.0).floor() as usize;

    let data_size = size.width as u32 * size.height as u32;
    let mut data: Vec<u32> = Vec::with_capacity(data_size as usize);

    let mut guards: Vec<JoinHandle<Vec<u32>>> = vec![];

    let thread_work = (data_size as f32 / thread_count as f32).ceil() as u32;
    let mut thread_start = 0;
    let mut thread_end = thread_start + thread_work;

    println!("Data Size: {}\nThread Work: {}", data_size, thread_work);

    for t in 0..thread_count {
        let t_size = size.clone();
        let t_window = window.clone();
        let t_limit = limit;

        if thread_end > data_size {
            thread_end = data_size;
        }

        if t == thread_count - 1 {
            thread_end = data_size;
        }

        let guard = thread::spawn(move || {
            println!("Starting thread [{}] working on data {} to {}", t, thread_start, thread_end);

            let thread_size = thread_end - thread_start;
            let mut thread_data = Vec::with_capacity(thread_size as usize);

            for i in thread_start..thread_end {
                let p = idx2point(i, t_size.width as u32);

                let px: f32 = p.x / t_size.width;
                let py: f32 = p.y / t_size.height;

                let cx = t_window.origin.x + px * t_window.size.width;
                let cy = (t_window.origin.y + t_window.size.height) - py * t_window.size.height;

                let c = mandelbrot(cx, cy, t_limit);

                thread_data.push(c);
            }

            thread_data
        });

        guards.push(guard);

        thread_start = thread_end;
        thread_end = thread_end + thread_work;
    }

    for g in guards {
        data.extend(g.join().unwrap().into_iter());
    }

    data
}

fn rbg_from_palette(palette: &Vec<Vec<f32>>, idx: usize) -> (f32, f32, f32) {
    let color = &palette[idx];
    (color[0], color[1], color[2])
}

fn color_for_val_with_config(val: u32, config: &Config) -> (u8, u8, u8) {
    let (r, g, b);

    let limit = config.limit;
    let steps = config.color_steps;
    let palette = &config.color_palette;

    if val == limit as u32 {
        r = 0;
        g = 0;
        b = 0;
    } else {
        let val = (val as f32 % steps) * (palette.len() as f32) / steps;
        let left = val as usize % palette.len();
        let right = (left + 1) % palette.len();

        let p = val - left as f32;
        let (r1, g1, b1) = rbg_from_palette(palette, left);
        let (r2, g2, b2) = rbg_from_palette(palette, right);
        r = (r1 + (r2 - r1) * p) as u8;
        g = (g1 + (g2 - g1) * p) as u8;
        b = (b1 + (b2 - b1) * p) as u8;
    }
    (r, g, b)
}

fn validate_config(conf: &Config) {

    // Check if limit is 'to large'
    if conf.limit > 10000 {
        panic!("Config Error: limit is over 10,000");
    }

    // Check if color component count is valid
    match conf.color_components {
        3 => println!("Using RBG colors"),
        _ => panic!("Unsuported color component count"),
    };

    // Check colors in palette match component count
    for v in &conf.color_palette {
        if v.len() != conf.color_components as usize {
            panic!("Config Error: Color {:?} does not match color component count", v);
        }
    }
}

fn main() {
    let start = PreciseTime::now();

    let args = App::new("Mandelbrot Generator")
        .version(&crate_version!()[..])
        .author("DJ Edmonson <djedmonson@gmail.com>")
        .about("Generates a mandelbrot image")
        .arg(Arg::with_name("CONFIG")
             .long("config")
             .help("Config JSON file to use. Output will be at <input_file_path>.png")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("output-palette")
             .long("output-palette")
             .help("Generate image with 100px squares of the provided colors in order. Outputs to <input_file_path>-palette.png"))
        .get_matches();

    let config_file_path = Path::new(args.value_of("CONFIG").unwrap());
    println!("Getting config from {}", config_file_path.display());

    let mut config_file = match File::open(&config_file_path) {
        Err(why) => panic!("Could not open config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(f) => f,
    };

    let mut config_json = String::new();
    match config_file.read_to_string(&mut config_json) {
        Err(why) => panic!("Could not read config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(_) => println!("Read config file"),
    };

    let config: Config = match json::decode(&config_json) {
        Err(why) => panic!("Error parsing config JSON: {}",
                           why),
        Ok(conf) => conf,
    };

    validate_config(&config);

    println!("Bootstrap time:\n{}", start.to(PreciseTime::now()));

    if args.is_present("output-palette") {
        let root = match config_file_path.file_stem() {
            None => unreachable!(),
            Some(r) => r.to_str().unwrap().to_string(),
        };

        let output_path_string = root + "-palette.png";
        let output_path = Path::new(&output_path_string);
        println!("Generating color palette image at {}", output_path.display());

        let width = config.color_palette.len() * 100;
        let height = 100;

        let mut imgbuf = image::ImageBuffer::new(width as u32, height as u32);

        for (x, _, pixel) in imgbuf.enumerate_pixels_mut() {
            let color_idx = (x as f32 / width as f32) * config.color_palette.len() as f32;

            let (r, b, g) = rbg_from_palette(&config.color_palette, color_idx as usize);

            *pixel = image::Rgb([r as u8, b as u8, g as u8]);
        }

        let _ = imgbuf.save(output_path);
    } else {
        let output_path = config_file_path.with_extension("ppm");

        let mut output_file =  match File::create(&output_path) {
            Err(why) => panic!("Could not open output file at {}: {}",
                               output_path.display(),
                               Error::description(&why)),
            Ok(f) => f,
        };

        let img_width = config.ppu as f32 * config.window.size.width;
        let img_height = config.ppu as f32 * config.window.size.height;

        println!("Generating image at {} with size {}x{}", output_path.display(), img_width, img_height);

        let size = Size {width: img_width, height: img_height};

        let render_start = PreciseTime::now();
        let mut phase_start = PreciseTime::now();

        let imgdata = gen_mandelbrot(&size, &config);

        println!("Generation Duration:\n{}", phase_start.to(PreciseTime::now()));

        phase_start = PreciseTime::now();

        writeln!(output_file, "P6").unwrap();
        writeln!(output_file, "{} {}", img_width as usize, img_height as usize).unwrap();
        writeln!(output_file, "255").unwrap();

        let mut linebuf = vec![0; img_width as usize * 3];

        for y in 0 .. img_height as usize {
            for x in 0 .. img_width as usize {
                let idx = point2idx(Point{ x: x as f32, y: y as f32}, size.width as u32) as usize;

                let it = imgdata[idx];

                let (r, g, b) = color_for_val_with_config(it, &config);

                let offset = 3 * x;

                linebuf[offset]     = r;
                linebuf[offset + 1] = g;
                linebuf[offset + 2] = b
            }
            output_file.write(&linebuf).unwrap();
        }
        println!("Render Duration:\n{}", phase_start.to(PreciseTime::now()));
        println!("Total Render Duration:\n{}", render_start.to(PreciseTime::now()));
    }

    println!("Total time:\n{}", start.to(PreciseTime::now()));
}

<commit_msg>Change to reasonable thread count and move to f64 for work calc<commit_after>#[macro_use]
extern crate clap;
extern crate time;
extern crate image;
extern crate num_cpus;
extern crate threadpool;
extern crate rustc_serialize;

use clap::{Arg, App};
use time::PreciseTime;
use rustc_serialize::json;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::thread::{self, JoinHandle};

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Size {
    width: f32,
    height: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Point {
    x: f32,
    y: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable, Clone)]
struct Rect {
    origin: Point,
    size: Size,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Config {
    ppu: u32,
    limit: u32,
    color_steps: f32,
    color_components: u8,
    color_palette: Vec<Vec<f32>>,
    window: Rect,
}

fn idx2point(idx: u32, width: u32) -> Point {
    let x = idx % width;
    let y = idx / width;
    Point { x: x as f32, y: y as f32}
}

fn point2idx(p: Point, width: u32) -> u32 {
    width * (p.y as u32) + (p.x as u32)
}

fn mandelbrot(cx: f32, cy: f32, limit: u32) -> u32 {
    let mut x = cx;
    let mut y = cy;
    let mut count = 0;
    while count < limit {
        let xy = x * y;
        let xx = x * x;
        let yy = y * y;
        let sum = xx + yy;
        if sum > 4.0 {
            break
        }
        count += 1;
        x = xx - yy + cx;
        y = xy * 2.0 + cy;
    }
    count as u32
}

fn gen_mandelbrot(size: &Size, config: &Config) -> Vec<u32> {
    let window = &config.window;
    let limit = config.limit;

    // Seems like a good number on most devices
    let thread_count = num_cpus::get() * 3;

    let data_size = size.width as u32 * size.height as u32;
    let mut data: Vec<u32> = Vec::with_capacity(data_size as usize);

    let mut guards: Vec<JoinHandle<Vec<u32>>> = vec![];

    let thread_work = (data_size as f64 / thread_count as f64).ceil() as u32;
    let mut thread_start = 0;
    let mut thread_end = thread_start + thread_work;

    println!("Data Size: {}\nThread Work: {}", data_size, thread_work);

    for t in 0..thread_count {
        let t_size = size.clone();
        let t_window = window.clone();
        let t_limit = limit;

        if thread_end > data_size {
            thread_end = data_size;
        }

        if t == thread_count - 1 {
            thread_end = data_size;
        }

        let guard = thread::spawn(move || {
            println!("Starting thread [{}] working on data {} to {}", t, thread_start, thread_end);

            let thread_size = thread_end - thread_start;
            let mut thread_data = Vec::with_capacity(thread_size as usize);

            for i in thread_start..thread_end {
                let p = idx2point(i, t_size.width as u32);

                let px: f32 = p.x / t_size.width;
                let py: f32 = p.y / t_size.height;

                let cx = t_window.origin.x + px * t_window.size.width;
                let cy = (t_window.origin.y + t_window.size.height) - py * t_window.size.height;

                let c = mandelbrot(cx, cy, t_limit);

                thread_data.push(c);
            }

            thread_data
        });

        guards.push(guard);

        thread_start = thread_end;
        thread_end = thread_end + thread_work;
    }

    for g in guards {
        data.extend(g.join().unwrap().into_iter());
    }

    data
}

fn rbg_from_palette(palette: &Vec<Vec<f32>>, idx: usize) -> (f32, f32, f32) {
    let color = &palette[idx];
    (color[0], color[1], color[2])
}

fn color_for_val_with_config(val: u32, config: &Config) -> (u8, u8, u8) {
    let (r, g, b);

    let limit = config.limit;
    let steps = config.color_steps;
    let palette = &config.color_palette;

    if val == limit as u32 {
        r = 0;
        g = 0;
        b = 0;
    } else {
        let val = (val as f32 % steps) * (palette.len() as f32) / steps;
        let left = val as usize % palette.len();
        let right = (left + 1) % palette.len();

        let p = val - left as f32;
        let (r1, g1, b1) = rbg_from_palette(palette, left);
        let (r2, g2, b2) = rbg_from_palette(palette, right);
        r = (r1 + (r2 - r1) * p) as u8;
        g = (g1 + (g2 - g1) * p) as u8;
        b = (b1 + (b2 - b1) * p) as u8;
    }
    (r, g, b)
}

fn validate_config(conf: &Config) {

    // Check if limit is 'to large'
    if conf.limit > 10000 {
        panic!("Config Error: limit is over 10,000");
    }

    // Check if color component count is valid
    match conf.color_components {
        3 => println!("Using RBG colors"),
        _ => panic!("Unsuported color component count"),
    };

    // Check colors in palette match component count
    for v in &conf.color_palette {
        if v.len() != conf.color_components as usize {
            panic!("Config Error: Color {:?} does not match color component count", v);
        }
    }
}

fn main() {
    let start = PreciseTime::now();

    let args = App::new("Mandelbrot Generator")
        .version(&crate_version!()[..])
        .author("DJ Edmonson <djedmonson@gmail.com>")
        .about("Generates a mandelbrot image")
        .arg(Arg::with_name("CONFIG")
             .long("config")
             .help("Config JSON file to use. Output will be at <input_file_path>.png")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("output-palette")
             .long("output-palette")
             .help("Generate image with 100px squares of the provided colors in order. Outputs to <input_file_path>-palette.png"))
        .get_matches();

    let config_file_path = Path::new(args.value_of("CONFIG").unwrap());
    println!("Getting config from {}", config_file_path.display());

    let mut config_file = match File::open(&config_file_path) {
        Err(why) => panic!("Could not open config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(f) => f,
    };

    let mut config_json = String::new();
    match config_file.read_to_string(&mut config_json) {
        Err(why) => panic!("Could not read config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(_) => println!("Read config file"),
    };

    let config: Config = match json::decode(&config_json) {
        Err(why) => panic!("Error parsing config JSON: {}",
                           why),
        Ok(conf) => conf,
    };

    validate_config(&config);

    println!("Bootstrap time:\n{}", start.to(PreciseTime::now()));

    if args.is_present("output-palette") {
        let root = match config_file_path.file_stem() {
            None => unreachable!(),
            Some(r) => r.to_str().unwrap().to_string(),
        };

        let output_path_string = root + "-palette.png";
        let output_path = Path::new(&output_path_string);
        println!("Generating color palette image at {}", output_path.display());

        let width = config.color_palette.len() * 100;
        let height = 100;

        let mut imgbuf = image::ImageBuffer::new(width as u32, height as u32);

        for (x, _, pixel) in imgbuf.enumerate_pixels_mut() {
            let color_idx = (x as f32 / width as f32) * config.color_palette.len() as f32;

            let (r, b, g) = rbg_from_palette(&config.color_palette, color_idx as usize);

            *pixel = image::Rgb([r as u8, b as u8, g as u8]);
        }

        let _ = imgbuf.save(output_path);
    } else {
        let output_path = config_file_path.with_extension("ppm");

        let mut output_file =  match File::create(&output_path) {
            Err(why) => panic!("Could not open output file at {}: {}",
                               output_path.display(),
                               Error::description(&why)),
            Ok(f) => f,
        };

        let img_width = config.ppu as f32 * config.window.size.width;
        let img_height = config.ppu as f32 * config.window.size.height;

        println!("Generating image at {} with size {}x{}", output_path.display(), img_width, img_height);

        let size = Size {width: img_width, height: img_height};

        let render_start = PreciseTime::now();
        let mut phase_start = PreciseTime::now();

        let imgdata = gen_mandelbrot(&size, &config);

        println!("Generation Duration:\n{}", phase_start.to(PreciseTime::now()));

        phase_start = PreciseTime::now();

        writeln!(output_file, "P6").unwrap();
        writeln!(output_file, "{} {}", img_width as usize, img_height as usize).unwrap();
        writeln!(output_file, "255").unwrap();

        let mut linebuf = vec![0; img_width as usize * 3];

        for y in 0 .. img_height as usize {
            for x in 0 .. img_width as usize {
                let idx = point2idx(Point{ x: x as f32, y: y as f32}, size.width as u32) as usize;

                let it = imgdata[idx];

                let (r, g, b) = color_for_val_with_config(it, &config);

                let offset = 3 * x;

                linebuf[offset]     = r;
                linebuf[offset + 1] = g;
                linebuf[offset + 2] = b
            }
            output_file.write(&linebuf).unwrap();
        }
        println!("Render Duration:\n{}", phase_start.to(PreciseTime::now()));
        println!("Total Render Duration:\n{}", render_start.to(PreciseTime::now()));
    }

    println!("Total time:\n{}", start.to(PreciseTime::now()));
}

<|endoftext|>