<commit_before>use std::cmp::max;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use grep::matcher::{Captures, Matcher};
use grep::regex::{RegexMatcher, RegexMatcherBuilder};
use grep::searcher::{
    BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkContextKind, SinkMatch,
};
use ignore::Walk;
use structopt::StructOpt;
use text_io::read;

#[derive(Debug, StructOpt)]
#[structopt(name = "fnr")]
/// Look for things, optionally replace them.
struct Opts {
    /// Run search in case insensitive mode
    #[structopt(short, long)]
    insensitive: bool,

    /// Print out replacements without actually performing them
    #[structopt(short, long)]
    dry_run: bool,

    #[structopt(short = "A", long, default_value = "0")]
    after: usize,
    #[structopt(short = "B", long, default_value = "0")]
    before: usize,
    #[structopt(short = "C", long, default_value = "0")]
    context: usize,

    #[structopt(short, long)]
    prompt: bool,

    /// What to search for.
    #[structopt(name = "FIND")]
    find: String,

    /// What to replace it with.
    #[structopt(name = "REPLACE")]
    replace: String,

    /// Locations to search. Current directory if not specified.
    #[structopt(name = "PATH", parse(from_os_str))]
    paths: Vec<PathBuf>,
}

#[derive(Debug, Clone)]
struct Line(u64, String);

#[derive(Debug)]
struct SearchMatch {
    line: Line,
    context_pre: Vec<Line>,
    context_post: Vec<Line>,
}

impl SearchMatch {
    fn as_change(&self, line: String) -> Change {
        return Change {
            line_number: self.line.0,
            new_line: line,
        };
    }

    fn display_change(&self, replacement_line: &str) {
        for line in &self.context_pre {
            print!("    {}: {}", line.0, line.1);
        }

        // TODO: Multiple matches on same line
        // TODO: Disable colors when not atty
        // TODO: Align line number with padding.
        print!("\x1B[31m-   {}: {}\x1B[0m", self.line.0, self.line.1);
        print!("\x1B[32m+   {}: {}\x1B[0m", self.line.0, replacement_line);

        for line in &self.context_post {
            print!("    {}: {}", line.0, line.1);
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum MatchState {
    Init,
    Before,
    Match,
    After,
}

#[derive(Debug)]
struct SearchMatchCollector {
    state: MatchState,

    cur_context_pre: Vec<Line>,
    cur_context_post: Vec<Line>,
    cur_match_line: Option<Line>,

    matches: Vec<SearchMatch>,
}

impl SearchMatchCollector {
    fn new() -> SearchMatchCollector {
        return SearchMatchCollector {
            state: MatchState::Init,
            cur_match_line: None,
            cur_context_pre: Vec::new(),
            cur_context_post: Vec::new(),
            matches: Vec::new(),
        };
    }

    fn maybe_emit(&mut self) {
        if let Some(line) = &self.cur_match_line {
            // TODO: clones are wasteful here.
            let m = SearchMatch {
                line: line.clone(),
                context_pre: self.cur_context_pre.clone(),
                context_post: self.cur_context_post.clone(),
            };

            self.matches.push(m);
            self.reset();
        }
    }

    fn reset(&mut self) {
        self.state = MatchState::Init;

        self.cur_match_line = None;
        self.cur_context_pre.clear();
        self.cur_context_post.clear();
    }

    #[inline]
    fn transition(&mut self, next: MatchState) {
        match (self.state, next) {
            // Beginning a new match or ending a previous one
            (MatchState::Match, MatchState::Before)       // Have before context, no after context
            | (MatchState::Match, MatchState::Match)      // No before context, no after context
            | (MatchState::After, MatchState::Before)     // Have before context, have after context
            | (MatchState::After, MatchState::Match) => { // Have after context, no before context
                self.maybe_emit();
            }

            (_prev, next) => {
                self.state = next;
            }
        }
    }

    fn collect(&mut self) -> &Vec<SearchMatch> {
        self.maybe_emit();
        return &self.matches;
    }
}

impl Drop for SearchMatchCollector {
    fn drop(&mut self) {
        self.maybe_emit();
    }
}

impl Sink for SearchMatchCollector {
    type Error = std::io::Error;

    fn matched(
        &mut self,
        _searcher: &Searcher,
        mat: &SinkMatch<'_>,
    ) -> Result<bool, std::io::Error> {
        self.transition(MatchState::Match);

        // TODO: handle errors
        let line = Line(
            mat.line_number().expect("grab line number"),
            String::from_utf8_lossy(mat.bytes()).to_string(),
        );

        self.cur_match_line = Some(line);

        Ok(true)
    }

    fn context(
        &mut self,
        _searcher: &Searcher,
        ctx: &SinkContext<'_>,
    ) -> Result<bool, std::io::Error> {
        // TODO: handle errors.
        // TODO: can we avoid the additional .to_string() by storing a COW instead?
        let line = Line(
            ctx.line_number().expect("grab line number"),
            String::from_utf8_lossy(ctx.bytes()).to_string(),
        );

        match ctx.kind() {
            &SinkContextKind::Before => {
                self.transition(MatchState::Before);
                self.cur_context_pre.push(line);
            }
            &SinkContextKind::After => {
                self.transition(MatchState::After);
                self.cur_context_post.push(line);
            }
            &SinkContextKind::Other => {}
        }

        Ok(true)
    }
}

trait Replacer {
    fn replace(&self, input: &str) -> Option<String>;
}

struct RegexReplacer<'a> {
    matcher: &'a RegexMatcher,
    template: &'a str,
}

impl<'a> Replacer for RegexReplacer<'a> {
    fn replace(&self, input: &str) -> Option<String> {
        let mut caps = self.matcher.new_captures().unwrap();
        let mut dst = vec![];

        self.matcher
            .replace_with_captures(input.as_bytes(), &mut caps, &mut dst, |caps, dst| {
                caps.interpolate(
                    |name| self.matcher.capture_index(name),
                    input.as_bytes(),
                    self.template.as_bytes(),
                    dst,
                );
                true
            })
            .unwrap();

        Some(String::from_utf8_lossy(&dst).to_string())
    }
}

#[derive(Debug)]
struct Change {
    line_number: u64,
    // TODO: use slice here
    new_line: String,
}

struct SearchMatchProcessor<'a> {
    replacer: &'a dyn Replacer,
    replacement_decider: ReplacementDecider,

    total_matches: usize,
    total_replacements: usize,
}

impl<'a> SearchMatchProcessor<'a> {
    fn new(
        replacer: &'a dyn Replacer,
        replacement_decider: ReplacementDecider,
    ) -> SearchMatchProcessor<'a> {
        return SearchMatchProcessor {
            replacer,
            replacement_decider,

            total_matches: 0,
            total_replacements: 0,
        };
    }

    fn handle_path(
        &mut self,
        path: &Path,
        matches: &Vec<SearchMatch>,
    ) -> Result<(), std::io::Error> {
        if matches.is_empty() {
            return Ok(());
        }
        println!(
            "--- {}: {} match{}",
            path.display(),
            matches.len(),
            if matches.len() == 1 { "" } else { "es" }
        );
        self.total_matches += matches.len();
        self.replacement_decider.reset_local_decision();

        let mut change_list = vec![];
        for m in matches {
            let replacement = if let Some(r) = self.replacer.replace(&m.line.1) {
                r
            } else {
                println!("TODO: replacer failed on line");
                continue;
            };

            let decision = self.replacement_decider.decide(m, &replacement);
            let change = match decision {
                ReplacementDecision::Accept => m.as_change(replacement),
                ReplacementDecision::Ignore => continue,
                ReplacementDecision::Edit => {
                    let mut line = read_input("Replace with [^D to skip] ")?;
                    if line == "" {
                        println!("... skipped ...");
                        continue;
                    }

                    line.push('\n');
                    m.display_change(&line);
                    println!("--");
                    m.as_change(line)
                }
            };

            change_list.push(change);
            self.total_replacements += 1;
        }

        // TODO: For --prompt, confirm before applying here
        if !change_list.is_empty() {
            self.apply_changes(path, &mut change_list)?;
        }

        Ok(())
    }

    fn apply_changes(&self, path: &Path, mut changes: &[Change]) -> Result<(), std::io::Error> {
        /*println!(
            "{}: {} change{}",
            path.display(),
            changes.len(),
            if changes.len() != 1 { "s" } else { "" }
        );*/
        let dst_path = path.with_extension("~");
        let src = File::open(path)?;
        let dst = File::create(&dst_path)?;

        let mut reader = BufReader::new(src);
        let mut writer = BufWriter::new(dst);

        let mut line_num = 0;
        loop {
            let mut line = String::new();
            let bytes_read = reader.read_line(&mut line)?;
            // EOF reached
            if bytes_read == 0 {
                break;
            }

            line_num += 1;

            if !changes.is_empty() && changes[0].line_number == line_num {
                writer.write(changes[0].new_line.as_bytes())?;
                changes = &changes[1..];
            } else {
                writer.write(line.as_bytes())?;
            }
        }

        std::fs::rename(dst_path, path)?;
        return Ok(());
    }

    fn finalize(&self) {
        println!(
            "All done. Replaced {} of {} matches",
            self.total_replacements, self.total_matches
        );
    }
}

#[derive(Debug, Copy, Clone)]
enum ReplacementDecision {
    Accept,
    Ignore,
    Edit,
}

fn read_input(prompt: &str) -> Result<String, std::io::Error> {
    print!("{}", prompt);
    std::io::stdout().flush()?;

    Ok(read!("{}\n"))
}

struct ReplacementDecider {
    should_prompt: bool,
    global_decision: Option<ReplacementDecision>,
    local_decision: Option<ReplacementDecision>,
}

impl ReplacementDecider {
    fn with_prompt() -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: true,
            global_decision: None,
            local_decision: None,
        };
    }

    fn constantly(decision: ReplacementDecision) -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: false,
            global_decision: Some(decision),
            local_decision: None,
        };
    }

    fn reset_local_decision(&mut self) {
        self.local_decision = None;
    }

    fn decide(&mut self, match_: &SearchMatch, replacement: &str) -> ReplacementDecision {
        if let Some(decision) = self.global_decision {
            return decision;
        } else if let Some(decision) = self.local_decision {
            return decision;
        }

        if !self.should_prompt {
            panic!("invalid state: no decision, but should not prompt");
        }

        // TODO: Should we display this always?
        match_.display_change(replacement);
        return self.prompt_for_decision();
    }

    fn prompt_for_decision(&mut self) -> ReplacementDecision {
        loop {
            let line = read_input("Stage this replacement [y,n,q,a,e,d,?] ").unwrap();

            return match line.as_str() {
                "y" => ReplacementDecision::Accept,
                "n" => ReplacementDecision::Ignore,
                "q" => {
                    self.global_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Ignore
                }
                "a" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "d" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "e" => {
                    self.local_decision = Some(ReplacementDecision::Edit);
                    ReplacementDecision::Edit
                }

                "?" | _ => {
                    println!(
                        "\x1B[31m
Y - replace this line
n - do not replace this line
q - quit; do not replace this line or any remaining ones
a - replace this line and all remaining ones in this file
d - do not replace this line nor any remaining ones in this file
e - edit this replacement
? - show help
\x1B[0m"
                    );
                    continue;
                }
            };
        }
    }
}

fn main() {
    let opts: Opts = Opts::from_args();
    println!("Parsed opts: {:?}", opts);

    let matcher = RegexMatcherBuilder::new()
        .case_insensitive(opts.insensitive)
        .build(&opts.find)
        .expect("bad pattern");

    let mut searcher = SearcherBuilder::new()
        .binary_detection(BinaryDetection::quit(0x00))
        .line_number(true)
        .before_context(max(opts.context, opts.before))
        .after_context(max(opts.context, opts.after))
        .build();

    // TODO: Confirm that template does not reference more capture groups than exist.
    let replacer = RegexReplacer {
        matcher: &matcher,
        template: &opts.replace,
    };

    if opts.dry_run && opts.prompt {
        println!("WARN: --prompt does not make sense with --dry-run, skipping");
    }

    let replacement_decider = if opts.dry_run {
        ReplacementDecider::constantly(ReplacementDecision::Ignore)
    } else if opts.prompt {
        ReplacementDecider::with_prompt()
    } else {
        ReplacementDecider::constantly(ReplacementDecision::Accept)
    };

    let mut proc = SearchMatchProcessor::new(&replacer, replacement_decider);

    let paths = if opts.paths.is_empty() {
        vec![Path::new("./").to_owned()]
    } else {
        opts.paths
    };

    for path in paths {
        // TODO: Support files as well as directories here
        for dir_entry in Walk::new(path) {
            let dir_entry = dir_entry.unwrap();
            let path = dir_entry.path();
            if !path.is_file() {
                continue;
            }

            let mut sink = SearchMatchCollector::new();

            searcher
                .search_path(&matcher, dir_entry.path(), &mut sink)
                .expect("search failed");

            let matches = sink.collect();
            proc.handle_path(path, matches).expect("handle path");
        }
    }

    if opts.dry_run {
        println!("--dry-run enabled, no changes applied.");
    }

    proc.finalize();
}
<commit_msg>Fix prompt decision mapping<commit_after>use std::cmp::max;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use grep::matcher::{Captures, Matcher};
use grep::regex::{RegexMatcher, RegexMatcherBuilder};
use grep::searcher::{
    BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkContextKind, SinkMatch,
};
use ignore::Walk;
use structopt::StructOpt;
use text_io::read;

#[derive(Debug, StructOpt)]
#[structopt(name = "fnr")]
/// Look for things, optionally replace them.
struct Opts {
    /// Run search in case insensitive mode
    #[structopt(short, long)]
    insensitive: bool,

    /// Print out replacements without actually performing them
    #[structopt(short, long)]
    dry_run: bool,

    #[structopt(short = "A", long, default_value = "0")]
    after: usize,
    #[structopt(short = "B", long, default_value = "0")]
    before: usize,
    #[structopt(short = "C", long, default_value = "0")]
    context: usize,

    #[structopt(short, long)]
    prompt: bool,

    /// What to search for.
    #[structopt(name = "FIND")]
    find: String,

    /// What to replace it with.
    #[structopt(name = "REPLACE")]
    replace: String,

    /// Locations to search. Current directory if not specified.
    #[structopt(name = "PATH", parse(from_os_str))]
    paths: Vec<PathBuf>,
}

#[derive(Debug, Clone)]
struct Line(u64, String);

#[derive(Debug)]
struct SearchMatch {
    line: Line,
    context_pre: Vec<Line>,
    context_post: Vec<Line>,
}

impl SearchMatch {
    fn as_change(&self, line: String) -> Change {
        return Change {
            line_number: self.line.0,
            new_line: line,
        };
    }

    fn display_change(&self, replacement_line: &str) {
        for line in &self.context_pre {
            print!("    {}: {}", line.0, line.1);
        }

        // TODO: Multiple matches on same line
        // TODO: Disable colors when not atty
        // TODO: Align line number with padding.
        print!("\x1B[31m-   {}: {}\x1B[0m", self.line.0, self.line.1);
        print!("\x1B[32m+   {}: {}\x1B[0m", self.line.0, replacement_line);

        for line in &self.context_post {
            print!("    {}: {}", line.0, line.1);
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum MatchState {
    Init,
    Before,
    Match,
    After,
}

#[derive(Debug)]
struct SearchMatchCollector {
    state: MatchState,

    cur_context_pre: Vec<Line>,
    cur_context_post: Vec<Line>,
    cur_match_line: Option<Line>,

    matches: Vec<SearchMatch>,
}

impl SearchMatchCollector {
    fn new() -> SearchMatchCollector {
        return SearchMatchCollector {
            state: MatchState::Init,
            cur_match_line: None,
            cur_context_pre: Vec::new(),
            cur_context_post: Vec::new(),
            matches: Vec::new(),
        };
    }

    fn maybe_emit(&mut self) {
        if let Some(line) = &self.cur_match_line {
            // TODO: clones are wasteful here.
            let m = SearchMatch {
                line: line.clone(),
                context_pre: self.cur_context_pre.clone(),
                context_post: self.cur_context_post.clone(),
            };

            self.matches.push(m);
            self.reset();
        }
    }

    fn reset(&mut self) {
        self.state = MatchState::Init;

        self.cur_match_line = None;
        self.cur_context_pre.clear();
        self.cur_context_post.clear();
    }

    #[inline]
    fn transition(&mut self, next: MatchState) {
        match (self.state, next) {
            // Beginning a new match or ending a previous one
            (MatchState::Match, MatchState::Before)       // Have before context, no after context
            | (MatchState::Match, MatchState::Match)      // No before context, no after context
            | (MatchState::After, MatchState::Before)     // Have before context, have after context
            | (MatchState::After, MatchState::Match) => { // Have after context, no before context
                self.maybe_emit();
            }

            (_prev, next) => {
                self.state = next;
            }
        }
    }

    fn collect(&mut self) -> &Vec<SearchMatch> {
        self.maybe_emit();
        return &self.matches;
    }
}

impl Drop for SearchMatchCollector {
    fn drop(&mut self) {
        self.maybe_emit();
    }
}

impl Sink for SearchMatchCollector {
    type Error = std::io::Error;

    fn matched(
        &mut self,
        _searcher: &Searcher,
        mat: &SinkMatch<'_>,
    ) -> Result<bool, std::io::Error> {
        self.transition(MatchState::Match);

        // TODO: handle errors
        let line = Line(
            mat.line_number().expect("grab line number"),
            String::from_utf8_lossy(mat.bytes()).to_string(),
        );

        self.cur_match_line = Some(line);

        Ok(true)
    }

    fn context(
        &mut self,
        _searcher: &Searcher,
        ctx: &SinkContext<'_>,
    ) -> Result<bool, std::io::Error> {
        // TODO: handle errors.
        // TODO: can we avoid the additional .to_string() by storing a COW instead?
        let line = Line(
            ctx.line_number().expect("grab line number"),
            String::from_utf8_lossy(ctx.bytes()).to_string(),
        );

        match ctx.kind() {
            &SinkContextKind::Before => {
                self.transition(MatchState::Before);
                self.cur_context_pre.push(line);
            }
            &SinkContextKind::After => {
                self.transition(MatchState::After);
                self.cur_context_post.push(line);
            }
            &SinkContextKind::Other => {}
        }

        Ok(true)
    }
}

trait Replacer {
    fn replace(&self, input: &str) -> Option<String>;
}

struct RegexReplacer<'a> {
    matcher: &'a RegexMatcher,
    template: &'a str,
}

impl<'a> Replacer for RegexReplacer<'a> {
    fn replace(&self, input: &str) -> Option<String> {
        let mut caps = self.matcher.new_captures().unwrap();
        let mut dst = vec![];

        self.matcher
            .replace_with_captures(input.as_bytes(), &mut caps, &mut dst, |caps, dst| {
                caps.interpolate(
                    |name| self.matcher.capture_index(name),
                    input.as_bytes(),
                    self.template.as_bytes(),
                    dst,
                );
                true
            })
            .unwrap();

        Some(String::from_utf8_lossy(&dst).to_string())
    }
}

#[derive(Debug)]
struct Change {
    line_number: u64,
    // TODO: use slice here
    new_line: String,
}

struct SearchMatchProcessor<'a> {
    replacer: &'a dyn Replacer,
    replacement_decider: ReplacementDecider,

    total_matches: usize,
    total_replacements: usize,
}

impl<'a> SearchMatchProcessor<'a> {
    fn new(
        replacer: &'a dyn Replacer,
        replacement_decider: ReplacementDecider,
    ) -> SearchMatchProcessor<'a> {
        return SearchMatchProcessor {
            replacer,
            replacement_decider,

            total_matches: 0,
            total_replacements: 0,
        };
    }

    fn handle_path(
        &mut self,
        path: &Path,
        matches: &Vec<SearchMatch>,
    ) -> Result<(), std::io::Error> {
        if matches.is_empty() {
            return Ok(());
        }
        println!(
            "--- {}: {} match{}",
            path.display(),
            matches.len(),
            if matches.len() == 1 { "" } else { "es" }
        );
        self.total_matches += matches.len();
        self.replacement_decider.reset_local_decision();

        let mut change_list = vec![];
        for m in matches {
            let replacement = if let Some(r) = self.replacer.replace(&m.line.1) {
                r
            } else {
                println!("TODO: replacer failed on line");
                continue;
            };

            let decision = self.replacement_decider.decide(m, &replacement);
            let change = match decision {
                ReplacementDecision::Accept => m.as_change(replacement),
                ReplacementDecision::Ignore => continue,
                ReplacementDecision::Edit => {
                    let mut line = read_input("Replace with [^D to skip] ")?;
                    if line == "" {
                        println!("... skipped ...");
                        continue;
                    }

                    line.push('\n');
                    m.display_change(&line);
                    println!("--");
                    m.as_change(line)
                }
            };

            change_list.push(change);
            self.total_replacements += 1;
        }

        // TODO: For --prompt, confirm before applying here
        if !change_list.is_empty() {
            self.apply_changes(path, &mut change_list)?;
        }

        Ok(())
    }

    fn apply_changes(&self, path: &Path, mut changes: &[Change]) -> Result<(), std::io::Error> {
        /*println!(
            "{}: {} change{}",
            path.display(),
            changes.len(),
            if changes.len() != 1 { "s" } else { "" }
        );*/
        let dst_path = path.with_extension("~");
        let src = File::open(path)?;
        let dst = File::create(&dst_path)?;

        let mut reader = BufReader::new(src);
        let mut writer = BufWriter::new(dst);

        let mut line_num = 0;
        loop {
            let mut line = String::new();
            let bytes_read = reader.read_line(&mut line)?;
            // EOF reached
            if bytes_read == 0 {
                break;
            }

            line_num += 1;

            if !changes.is_empty() && changes[0].line_number == line_num {
                writer.write(changes[0].new_line.as_bytes())?;
                changes = &changes[1..];
            } else {
                writer.write(line.as_bytes())?;
            }
        }

        std::fs::rename(dst_path, path)?;
        return Ok(());
    }

    fn finalize(&self) {
        println!(
            "All done. Replaced {} of {} matches",
            self.total_replacements, self.total_matches
        );
    }
}

#[derive(Debug, Copy, Clone)]
enum ReplacementDecision {
    Accept,
    Ignore,
    Edit,
}

fn read_input(prompt: &str) -> Result<String, std::io::Error> {
    print!("{}", prompt);
    std::io::stdout().flush()?;

    Ok(read!("{}\n"))
}

struct ReplacementDecider {
    should_prompt: bool,
    global_decision: Option<ReplacementDecision>,
    local_decision: Option<ReplacementDecision>,
}

impl ReplacementDecider {
    fn with_prompt() -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: true,
            global_decision: None,
            local_decision: None,
        };
    }

    fn constantly(decision: ReplacementDecision) -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: false,
            global_decision: Some(decision),
            local_decision: None,
        };
    }

    fn reset_local_decision(&mut self) {
        self.local_decision = None;
    }

    fn decide(&mut self, match_: &SearchMatch, replacement: &str) -> ReplacementDecision {
        if let Some(decision) = self.global_decision {
            return decision;
        } else if let Some(decision) = self.local_decision {
            return decision;
        }

        if !self.should_prompt {
            panic!("invalid state: no decision, but should not prompt");
        }

        // TODO: Should we display this always?
        match_.display_change(replacement);
        return self.prompt_for_decision();
    }

    fn prompt_for_decision(&mut self) -> ReplacementDecision {
        loop {
            let line = read_input("Stage this replacement [y,n,q,a,e,d,?] ").unwrap();

            return match line.as_str() {
                "y" => ReplacementDecision::Accept,
                "n" => ReplacementDecision::Ignore,
                "q" => {
                    self.global_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Ignore
                }
                "a" => {
                    self.local_decision = Some(ReplacementDecision::Accept);
                    ReplacementDecision::Accept
                }
                "d" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Ignore
                }
                "e" => {
                    ReplacementDecision::Edit
                }

                "?" | _ => {
                    println!(
                        "\x1B[31m
Y - replace this line
n - do not replace this line
q - quit; do not replace this line or any remaining ones
a - replace this line and all remaining ones in this file
d - do not replace this line nor any remaining ones in this file
e - edit this replacement
? - show help
\x1B[0m"
                    );
                    continue;
                }
            };
        }
    }
}

fn main() {
    // TODO: Paths from stdin
    let opts: Opts = Opts::from_args();
    println!("Parsed opts: {:?}", opts);

    let matcher = RegexMatcherBuilder::new()
        .case_insensitive(opts.insensitive)
        .build(&opts.find)
        .expect("bad pattern");

    let mut searcher = SearcherBuilder::new()
        .binary_detection(BinaryDetection::quit(0x00))
        .line_number(true)
        .before_context(max(opts.context, opts.before))
        .after_context(max(opts.context, opts.after))
        .build();

    // TODO: Confirm that template does not reference more capture groups than exist.
    let replacer = RegexReplacer {
        matcher: &matcher,
        template: &opts.replace,
    };

    if opts.dry_run && opts.prompt {
        println!("WARN: --prompt does not make sense with --dry-run, skipping");
    }

    let replacement_decider = if opts.dry_run {
        ReplacementDecider::constantly(ReplacementDecision::Ignore)
    } else if opts.prompt {
        ReplacementDecider::with_prompt()
    } else {
        ReplacementDecider::constantly(ReplacementDecision::Accept)
    };

    let mut proc = SearchMatchProcessor::new(&replacer, replacement_decider);

    let paths = if opts.paths.is_empty() {
        vec![Path::new("./").to_owned()]
    } else {
        opts.paths
    };

    for path in paths {
        // TODO: Support files as well as directories here
        for dir_entry in Walk::new(path) {
            let dir_entry = dir_entry.unwrap();
            let path = dir_entry.path();
            if !path.is_file() {
                continue;
            }

            let mut sink = SearchMatchCollector::new();

            searcher
                .search_path(&matcher, dir_entry.path(), &mut sink)
                .expect("search failed");

            let matches = sink.collect();
            proc.handle_path(path, matches).expect("handle path");
        }
    }

    if opts.dry_run {
        println!("--dry-run enabled, no changes applied.");
    }

    proc.finalize();
}
<|endoftext|>


<commit_before>#![feature(phase)]
#![allow(uppercase_variables)]
#![allow(unused_imports)]

extern crate test;
extern crate ndarray;

use ndarray::{Array, S, Si};
use ndarray::{arr0, arr1, arr2};
use ndarray::d2;

use test::black_box;

#[bench]
fn bench_std_add(bench: &mut test::Bencher)
{
    let mut a = arr2::<f32>([[1., 2., 2.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [5., 6., 6.]]);
    let b = a.clone();
    bench.iter(|| a.iadd(&b));
}

#[bench]
fn bench_std_iter_1d(bench: &mut test::Bencher)
{
    let a = arr1::<f32>([1., 2., 2.,
                         3., 4., 4.,
                         3., 4., 4.,
                         3., 4., 4.,
                         5., 6., 6.]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_std_iter_2d(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_iter_diag(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    bench.iter(|| for elt in a.diag_iter() { black_box(elt) })
}

#[bench]
fn bench_row_iter(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    let it = a.row_iter(17);
    bench.iter(|| for elt in it.clone() { black_box(elt) })
}

#[bench]
fn bench_col_iter(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    bench.iter(|| for elt in a.col_iter(17) { black_box(elt) })
}

#[bench]
fn bench_mat_mul(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    let mut at = a.clone();
    at.swap_axes(0, 1);
    bench.iter(|| at.mat_mul(&a));
}

#[bench]
fn lst_squares(bench: &mut test::Bencher)
{
    let xs = arr2::<f32>([[ 2.,  3.],
                          [-2., -1.],
                          [ 1.,  5.],
                          [-1.,  2.]]);
    let b = arr1([1., -1., 2., 1.]);
    bench.iter(|| ndarray::linalg::least_squares(&xs, &b));
}
<commit_msg>benches: Add iterator over .raw_data() for comparison<commit_after>#![feature(phase)]
#![allow(uppercase_variables)]
#![allow(unused_imports)]

extern crate test;
extern crate ndarray;

use ndarray::{Array, S, Si};
use ndarray::{arr0, arr1, arr2};
use ndarray::d2;

use test::black_box;

#[bench]
fn bench_std_add(bench: &mut test::Bencher)
{
    let mut a = arr2::<f32>([[1., 2., 2.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [5., 6., 6.]]);
    let b = a.clone();
    bench.iter(|| a.iadd(&b));
}

#[bench]
fn bench_std_iter_1d(bench: &mut test::Bencher)
{
    let a = arr1::<f32>([1., 2., 2.,
                         3., 4., 4.,
                         3., 4., 4.,
                         3., 4., 4.,
                         5., 6., 6.]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_std_iter_1d_raw(bench: &mut test::Bencher)
{
    let a = arr1::<f32>([1., 2., 2.,
                         3., 4., 4.,
                         3., 4., 4.,
                         3., 4., 4.,
                         5., 6., 6.]);
    bench.iter(|| for elt in a.raw_data().iter() { black_box(elt) })
}

#[bench]
fn bench_std_iter_2d(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_iter_diag(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    bench.iter(|| for elt in a.diag_iter() { black_box(elt) })
}

#[bench]
fn bench_row_iter(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    let it = a.row_iter(17);
    bench.iter(|| for elt in it.clone() { black_box(elt) })
}

#[bench]
fn bench_col_iter(bench: &mut test::Bencher)
{
    let a = Array::<f32, _>::zeros(d2(1024, 1024));
    bench.iter(|| for elt in a.col_iter(17) { black_box(elt) })
}

#[bench]
fn bench_mat_mul(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    let mut at = a.clone();
    at.swap_axes(0, 1);
    bench.iter(|| at.mat_mul(&a));
}

#[bench]
fn lst_squares(bench: &mut test::Bencher)
{
    let xs = arr2::<f32>([[ 2.,  3.],
                          [-2., -1.],
                          [ 1.,  5.],
                          [-1.,  2.]]);
    let b = arr1([1., -1., 2., 1.]);
    bench.iter(|| ndarray::linalg::least_squares(&xs, &b));
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate test;

use std::collections::VecDeque;

#[macro_use]
extern crate strider;
use strider::SliceRing;
use strider::SliceRingImpl;

#[bench]
fn bench_empty_to_probe_for_release_mode(b: &mut test::Bencher) {
    b.iter(|| 1)
}

macro_rules! bench_push_many_back {
    ($bencher:expr, $new:expr) => {{
        let vec = (0..1000).collect::<Vec<i32>>();
        $bencher.iter(|| {
            let mut testable = $new;
            testable.push_many_back(&vec[..]);
            testable
        });
    }}
}
#[bench]
fn bench_push_many_back_deque(b: &mut test::Bencher) {
    bench_push_many_back!(b, VecDeque::new());
}
#[bench]
fn bench_push_many_back_deque_with_capacity(b: &mut test::Bencher) {
    bench_push_many_back!(b, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_push_many_back_optimized(b: &mut test::Bencher) {
    bench_push_many_back!(b, SliceRingImpl::new());
}
#[bench]
fn bench_push_many_back_optimized_with_capacity(b: &mut test::Bencher) {
    bench_push_many_back!(b, SliceRingImpl::with_capacity(1000));
}

macro_rules! bench_drop_many_front {
    ($bencher:expr, $new:expr) => {{
        let vec = (0..1000).collect::<Vec<i32>>();
        $bencher.iter(|| {
            let mut testable = $new;
            testable.push_many_back(&vec[..]);
            testable.drop_many_front(1000)
        });
    }}
}

#[bench]
fn bench_drop_many_front_deque(b: &mut test::Bencher) {
    bench_drop_many_front!(b, VecDeque::new());
}
#[bench]
fn bench_drop_many_front_deque_with_capacity(b: &mut test::Bencher) {
    bench_drop_many_front!(b, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_drop_many_front_optimized(b: &mut test::Bencher) {
    bench_drop_many_front!(b, SliceRingImpl::new());
}
#[bench]
fn bench_drop_many_front_optimized_with_capacity(b: &mut test::Bencher) {
    bench_drop_many_front!(b, SliceRingImpl::with_capacity(1000));
}

macro_rules! bench_read_many_front {
    ($bencher:expr, $new:expr) => {{
        let mut output: Vec<i32> = std::iter::repeat(0).take(1000).collect();
        let vec = (0..1000).collect::<Vec<i32>>();
        let mut testable = $new;
        testable.push_many_back(&vec[..]);
        $bencher.iter(|| {
            testable.read_many_front(&mut output[..])
        });
    }}
}
#[bench]
fn bench_read_many_front_deque(b: &mut test::Bencher) {
    bench_read_many_front!(b, VecDeque::new());
}
#[bench]
fn bench_read_many_front_deque_with_capacity(b: &mut test::Bencher) {
    bench_read_many_front!(b, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_read_many_front_optimized(b: &mut test::Bencher) {
    bench_read_many_front!(b, SliceRingImpl::new());
}
#[bench]
fn bench_read_many_front_optimized_with_capacity(b: &mut test::Bencher) {
    bench_read_many_front!(b, SliceRingImpl::with_capacity(1000));
}

#[bench]
fn bench_test_slice_ring_deque(b: &mut test::Bencher) {
    b.iter(|| {
        test_slice_ring!(VecDeque::<i32>::new());
    });
}
#[bench]
fn bench_test_slice_ring_optimized(b: &mut test::Bencher) {
    b.iter(|| {
        test_slice_ring!(SliceRingImpl::<i32>::new())
    });
}
#[macro_export]
macro_rules! bench_slice_ring_windowing {
    ($new:expr) => {{
        let window_size = 4096;
        let step_size = 512;
        let mut testable = $new;
        let input: Vec<i32> = std::iter::repeat(0).take(6000).collect();
        let mut output: Vec<i32> = std::iter::repeat(0).take(window_size).collect();
        for _ in 0..100 {
            testable.push_many_back(&input[..]);
            while testable.len() >= window_size {
                testable.read_many_front(&mut output[..]);
                testable.drop_many_front(step_size);
            }
        }
        testable
    }};
}
#[bench]
fn bench_slice_ring_windowing_deque(b: &mut test::Bencher) {
    b.iter(|| {
        bench_slice_ring_windowing!(VecDeque::<i32>::new());
    });
}
#[bench]
fn bench_slice_ring_windowing_optimized(b: &mut test::Bencher) {
    b.iter(|| {
        bench_slice_ring_windowing!(SliceRingImpl::<i32>::new())
    });
}

#[macro_export]
macro_rules! bench_slice_ring_one_minute_audio {
    ($bencher:expr, $new:expr) => {{
        // we move this outside the bencher
        // to measure only the time taken by the
        // slice ring operations.
        // keep all allocations out of the benchmark iteration.
        let samples_per_second = 44100;
        let seconds = 60;
        let sample_count: usize = seconds * samples_per_second;
        // println!("sample_count = {}", sample_count);
        let samples = (0..(sample_count as i32)).collect::<Vec<i32>>();
        let input_slice_size = 10000;
        let window_size = 1024;
        let step_size = 512;

        // determine the output_size by
        // simulating just the indexes without data
        let mut output_size = 0;
        let mut input_index = 0;
        let mut output_index = 0;
        let mut ring_size = 0;
        while input_index < samples.len() {
            let input_range_start = input_index;
            let input_range_end = std::cmp::min(sample_count, input_index + input_slice_size);
            // push_many_back
            ring_size += input_range_end - input_range_start;
            while ring_size >= window_size {
                let output_range_start = output_index;
                let output_range_end = output_index + window_size;
                output_size += output_range_end - output_range_start;
                ring_size -= step_size;
                output_index += window_size;
            }
            input_index += input_slice_size;
        }
        // println!("output_size = {}", output_size);
        let mut output = std::iter::repeat(0).take(output_size).collect::<Vec<i32>>();
        // println!("output.len() = {}", output.len());
        let mut expected_output = Vec::<i32>::new();
        let mut expected_output_index = 0;
        while expected_output_index < sample_count {
            let range_start = expected_output_index as i32;
            let range_end = (expected_output_index + window_size) as i32;
            let mut range_vec = (range_start..range_end).collect::<Vec<i32>>();
            expected_output.append(&mut range_vec);
            expected_output_index += step_size;
        }
        // println!("expected_output.len() = {}", expected_output.len());
        $bencher.iter(|| {
            let mut testable = $new;
            let mut input_index = 0;
            let mut output_index = 0;
            while input_index < samples.len() {
                let input_range_start = input_index;
                let input_range_end = std::cmp::min(sample_count, input_index + input_slice_size);
                let input_slice =
                    &samples[input_range_start..input_range_end];
                testable.push_many_back(input_slice);
                while testable.len() >= window_size {
                    let output_range_start = output_index;
                    let output_range_end = output_index + window_size;
                    let output_slice =
                        &mut output[output_range_start..output_range_end];
                    testable.read_many_front(output_slice);
                    testable.drop_many_front(step_size);
                    output_index += window_size;
                }
                input_index += input_slice_size;
            }
            // println!("input_index = {}", input_index);
            // println!("output_index = {}", output_index);
            testable
        });
        for i in 0..output.len() {
            assert!(output[i] == expected_output[i], "{} != {} at {}", output[i], expected_output[i], i);
        }
        // assert_eq!(output.len(), expected_output.len());
    }};
}
#[bench]
fn bench_slice_ring_one_minute_audio_deque(b: &mut test::Bencher) {
    bench_slice_ring_one_minute_audio!(b, VecDeque::<i32>::new());
}
#[bench]
fn bench_slice_ring_one_minute_audio_optimized(b: &mut test::Bencher) {
    bench_slice_ring_one_minute_audio!(b, SliceRingImpl::<i32>::new());
}
<commit_msg>b -> bencher<commit_after>#![feature(test)]
extern crate test;

use std::collections::VecDeque;

#[macro_use]
extern crate strider;
use strider::SliceRing;
use strider::SliceRingImpl;

#[bench]
fn bench_empty_to_probe_for_release_mode(bencher: &mut test::Bencher) {
    bencher.iter(|| 1)
}

macro_rules! bench_push_many_back {
    ($bencher:expr, $new:expr) => {{
        let vec = (0..1000).collect::<Vec<i32>>();
        $bencher.iter(|| {
            let mut testable = $new;
            testable.push_many_back(&vec[..]);
            testable
        });
    }}
}
#[bench]
fn bench_push_many_back_deque(bencher: &mut test::Bencher) {
    bench_push_many_back!(bencher, VecDeque::new());
}
#[bench]
fn bench_push_many_back_deque_with_capacity(bencher: &mut test::Bencher) {
    bench_push_many_back!(bencher, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_push_many_back_optimized(bencher: &mut test::Bencher) {
    bench_push_many_back!(bencher, SliceRingImpl::new());
}
#[bench]
fn bench_push_many_back_optimized_with_capacity(bencher: &mut test::Bencher) {
    bench_push_many_back!(bencher, SliceRingImpl::with_capacity(1000));
}

macro_rules! bench_drop_many_front {
    ($bencher:expr, $new:expr) => {{
        let vec = (0..1000).collect::<Vec<i32>>();
        $bencher.iter(|| {
            let mut testable = $new;
            testable.push_many_back(&vec[..]);
            testable.drop_many_front(1000)
        });
    }}
}

#[bench]
fn bench_drop_many_front_deque(bencher: &mut test::Bencher) {
    bench_drop_many_front!(bencher, VecDeque::new());
}
#[bench]
fn bench_drop_many_front_deque_with_capacity(bencher: &mut test::Bencher) {
    bench_drop_many_front!(bencher, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_drop_many_front_optimized(bencher: &mut test::Bencher) {
    bench_drop_many_front!(bencher, SliceRingImpl::new());
}
#[bench]
fn bench_drop_many_front_optimized_with_capacity(bencher: &mut test::Bencher) {
    bench_drop_many_front!(bencher, SliceRingImpl::with_capacity(1000));
}

macro_rules! bench_read_many_front {
    ($bencher:expr, $new:expr) => {{
        let mut output: Vec<i32> = std::iter::repeat(0).take(1000).collect();
        let vec = (0..1000).collect::<Vec<i32>>();
        let mut testable = $new;
        testable.push_many_back(&vec[..]);
        $bencher.iter(|| {
            testable.read_many_front(&mut output[..])
        });
    }}
}
#[bench]
fn bench_read_many_front_deque(bencher: &mut test::Bencher) {
    bench_read_many_front!(bencher, VecDeque::new());
}
#[bench]
fn bench_read_many_front_deque_with_capacity(bencher: &mut test::Bencher) {
    bench_read_many_front!(bencher, VecDeque::with_capacity(1000));
}
#[bench]
fn bench_read_many_front_optimized(bencher: &mut test::Bencher) {
    bench_read_many_front!(bencher, SliceRingImpl::new());
}
#[bench]
fn bench_read_many_front_optimized_with_capacity(bencher: &mut test::Bencher) {
    bench_read_many_front!(bencher, SliceRingImpl::with_capacity(1000));
}

#[bench]
fn bench_test_slice_ring_deque(bencher: &mut test::Bencher) {
    bencher.iter(|| {
        test_slice_ring!(VecDeque::<i32>::new());
    });
}
#[bench]
fn bench_test_slice_ring_optimized(bencher: &mut test::Bencher) {
    bencher.iter(|| {
        test_slice_ring!(SliceRingImpl::<i32>::new())
    });
}
#[macro_export]
macro_rules! bench_slice_ring_windowing {
    ($new:expr) => {{
        let window_size = 4096;
        let step_size = 512;
        let mut testable = $new;
        let input: Vec<i32> = std::iter::repeat(0).take(6000).collect();
        let mut output: Vec<i32> = std::iter::repeat(0).take(window_size).collect();
        for _ in 0..100 {
            testable.push_many_back(&input[..]);
            while testable.len() >= window_size {
                testable.read_many_front(&mut output[..]);
                testable.drop_many_front(step_size);
            }
        }
        testable
    }};
}
#[bench]
fn bench_slice_ring_windowing_deque(bencher: &mut test::Bencher) {
    bencher.iter(|| {
        bench_slice_ring_windowing!(VecDeque::<i32>::new());
    });
}
#[bench]
fn bench_slice_ring_windowing_optimized(bencher: &mut test::Bencher) {
    bencher.iter(|| {
        bench_slice_ring_windowing!(SliceRingImpl::<i32>::new())
    });
}

#[macro_export]
macro_rules! bench_slice_ring_one_minute_audio {
    ($bencher:expr, $new:expr) => {{
        // we move this outside the bencher
        // to measure only the time taken by the
        // slice ring operations.
        // keep all allocations out of the benchmark iteration.
        let samples_per_second = 44100;
        let seconds = 60;
        let sample_count: usize = seconds * samples_per_second;
        // println!("sample_count = {}", sample_count);
        let samples = (0..(sample_count as i32)).collect::<Vec<i32>>();
        let input_slice_size = 10000;
        let window_size = 1024;
        let step_size = 512;

        // determine the output_size by
        // simulating just the indexes without data
        let mut output_size = 0;
        let mut input_index = 0;
        let mut output_index = 0;
        let mut ring_size = 0;
        while input_index < samples.len() {
            let input_range_start = input_index;
            let input_range_end = std::cmp::min(sample_count, input_index + input_slice_size);
            // push_many_back
            ring_size += input_range_end - input_range_start;
            while ring_size >= window_size {
                let output_range_start = output_index;
                let output_range_end = output_index + window_size;
                output_size += output_range_end - output_range_start;
                ring_size -= step_size;
                output_index += window_size;
            }
            input_index += input_slice_size;
        }
        // println!("output_size = {}", output_size);
        let mut output = std::iter::repeat(0).take(output_size).collect::<Vec<i32>>();
        // println!("output.len() = {}", output.len());
        let mut expected_output = Vec::<i32>::new();
        let mut expected_output_index = 0;
        while expected_output_index < sample_count {
            let range_start = expected_output_index as i32;
            let range_end = (expected_output_index + window_size) as i32;
            let mut range_vec = (range_start..range_end).collect::<Vec<i32>>();
            expected_output.append(&mut range_vec);
            expected_output_index += step_size;
        }
        // println!("expected_output.len() = {}", expected_output.len());
        $bencher.iter(|| {
            let mut testable = $new;
            let mut input_index = 0;
            let mut output_index = 0;
            while input_index < samples.len() {
                let input_range_start = input_index;
                let input_range_end = std::cmp::min(sample_count, input_index + input_slice_size);
                let input_slice =
                    &samples[input_range_start..input_range_end];
                testable.push_many_back(input_slice);
                while testable.len() >= window_size {
                    let output_range_start = output_index;
                    let output_range_end = output_index + window_size;
                    let output_slice =
                        &mut output[output_range_start..output_range_end];
                    testable.read_many_front(output_slice);
                    testable.drop_many_front(step_size);
                    output_index += window_size;
                }
                input_index += input_slice_size;
            }
            // println!("input_index = {}", input_index);
            // println!("output_index = {}", output_index);
            testable
        });
        for i in 0..output.len() {
            assert!(output[i] == expected_output[i], "{} != {} at {}", output[i], expected_output[i], i);
        }
        // assert_eq!(output.len(), expected_output.len());
    }};
}
#[bench]
fn bench_slice_ring_one_minute_audio_deque(bencher: &mut test::Bencher) {
    bench_slice_ring_one_minute_audio!(bencher, VecDeque::<i32>::new());
}
#[bench]
fn bench_slice_ring_one_minute_audio_optimized(bencher: &mut test::Bencher) {
    bench_slice_ring_one_minute_audio!(bencher, SliceRingImpl::<i32>::new());
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate mtl;

extern crate test;
extern crate rand;

use test::Bencher;
use rand::{thread_rng, Rng};

use mtl::Dense2D;
use mtl::Matrix;

const SIZE: usize = 20;

#[bench]
fn access_via_double_index(b: &mut Bencher) {
    let mut mv = Vec::new();
    for _ in 0 .. SIZE {
        mv.push(thread_rng().gen_iter().take(SIZE).collect::<Vec<f64>>());
    }
    let m = Dense2D::from_vec(mv);
    let mut v = 0f64;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = m[i][j];
            }
        }
    });
}

 #[bench]
fn access_via_tuple_index(b: &mut Bencher) {
    let mut mv = Vec::new();
    for _ in 0 .. SIZE {
        mv.push(thread_rng().gen_iter().take(SIZE).collect::<Vec<f64>>());
    }
    let m = Dense2D::from_vec(mv);
    let mut v = 0f64;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = m[(i,j)];
            }
        }
    });
}

#[bench]
fn access_via_matrix_get(b: &mut Bencher) {
    let mut mv = Vec::new();
    for _ in 0 .. SIZE {
        mv.push(thread_rng().gen_iter().take(SIZE).collect::<Vec<f64>>());
    }
    let m = Dense2D::from_vec(mv);
    let mut v = 0f64;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = *m.get(i,j).unwrap();
            }
        }
    });
}

#[bench]
fn access_via_matrix_get_unchecked(b: &mut Bencher) {
    let mut mv = Vec::new();
    for _ in 0 .. SIZE {
        mv.push(thread_rng().gen_iter().take(SIZE).collect::<Vec<f64>>());
    }
    let m = Dense2D::from_vec(mv);
    let mut v = 0f64;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = unsafe { *m.get_unchecked(i, j) };
            }
        }
    });
}
<commit_msg>add sparse matrix benchmarks<commit_after>#![feature(test)]

extern crate mtl;

extern crate test;
extern crate rand;

use test::Bencher;
use rand::{thread_rng, Rng};

// use mtl::Dense2D;
use mtl::matrix::Matrix;
use mtl::matrix::sparse::SparseYale;

const SIZE: usize = 500;
const SPARSY: u32 = 20;

#[bench]
fn access_via_double_index(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    // for _ in 0 .. SIZE {
    //     mv.push(thread_rng().gen_iter().take(SIZE).collect::<Vec<i32>>());
    // }
    let m = Matrix::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = m[i][j];
            }
        }
    });
}

 #[bench]
fn access_via_tuple_index(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = Matrix::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = m[(i,j)];
            }
        }
    });
}

#[bench]
fn access_via_matrix_get(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = Matrix::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = *m.get(i,j).unwrap();
            }
        }
    });
}

#[bench]
fn access_via_matrix_get_unchecked(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = Matrix::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = unsafe { *m.get_unchecked(i, j) };
            }
        }
    });
}

#[bench]
fn sparse_matrix_access_via_tuple_index(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = SparseYale::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = m[(i,j)];
            }
        }
    });
}

#[bench]
fn sparse_matrix_access_via_get(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = SparseYale::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                v = *m.get((i,j)).unwrap();
            }
        }
    });
}

#[bench]
fn sparse_matrix_access_via_get_unchecked(b: &mut Bencher) {
    let mut mv = Vec::new();
    let mut rng = thread_rng();
    for _ in 0 .. SIZE {
        let mut row = Vec::new();
        for _ in 0 .. SIZE {
            if rng.gen_weighted_bool(SPARSY) {
                row.push(rng.gen::<i32>());
            } else {
                row.push(0);
            }
        }
        mv.push(row);
    }
    let m = SparseYale::from_vec(mv);
    let mut v = 0i32;
    b.iter(|| {
        for i in 0 .. m.shape().0 {
            for j in 0 .. m.shape().1 {
                unsafe {
                    v = *m.get_unchecked((i,j));
                }
            }
        }
    });
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate bencher;
#[macro_use]
extern crate soa_derive;

use bencher::Bencher;

#[derive(StructOfArray)]
pub struct Small {
    x: f64,
    y: f64,
    z: f64,
}

impl Small {
    fn new() -> Small {
        Small {
            x: 1.0,
            y: 0.2,
            z: -2.3,
        }
    }

    fn aos_vec(size: usize) -> Vec<Small> {
        let mut vec = Vec::new();
        for _ in 0..size {
            vec.push(Small::new())
        }
        return vec;
    }

    fn soa_vec(size: usize) -> SmallVec {
        let mut vec = SmallVec::new();
        for _ in 0..size {
            vec.push(Small::new())
        }
        return vec;
    }
}

#[derive(StructOfArray)]
pub struct Big {
    position: (f64, f64, f64),
    velocity: (f64, f64, f64),
    id: usize,
    name: String,
    userdata: String
}

impl Big {
    fn new() -> Big {
        Big {
            position: (1.0, 0.2, -2.3),
            velocity: (1.0, 0.2, -2.3),
            id: 67,
            name: "foo".into(),
            userdata: "bar".into()
        }
    }

    fn aos_vec(size: usize) -> Vec<Big> {
        let mut vec = Vec::new();
        for _ in 0..size {
            vec.push(Big::new())
        }
        return vec;
    }

    fn soa_vec(size: usize) -> BigVec {
        let mut vec = BigVec::new();
        for _ in 0..size {
            vec.push(Big::new())
        }
        return vec;
    }
}

fn aos_small_push(bencher: &mut Bencher) {
    let mut vec = Vec::new();
    bencher.iter(||{
        vec.push(Small::new())
    })
}

fn soa_small_push(bencher: &mut Bencher) {
    let mut vec = SmallVec::new();
    bencher.iter(||{
        vec.push(Small::new())
    })
}

fn aos_big_push(bencher: &mut Bencher) {
    let mut vec = Vec::new();
    bencher.iter(||{
        vec.push(Big::new())
    })
}

fn soa_big_push(bencher: &mut Bencher) {
    let mut vec = BigVec::new();
    bencher.iter(||{
        vec.push(Big::new())
    })
}

fn aos_small_do_work_10000(bencher: &mut Bencher) {
    let vec = Small::aos_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.x + v.y;
        }
        s
    })
}

fn soa_small_do_work_10000(bencher: &mut Bencher) {
    let vec = Small::soa_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for (x, y) in vec.x.iter().zip(&vec.y) {
            s += x + y;
        }
        s
    })
}

fn aos_big_do_work_1000(bencher: &mut Bencher) {
    let vec = Big::aos_vec(1000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.position.0 + v.velocity.0 * 0.1;
        }
        s
    })
}

fn aos_big_do_work_10000(bencher: &mut Bencher) {
    let vec = Big::aos_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.position.0 + v.velocity.0 * 0.1;
        }
        s
    })
}

fn soa_big_do_work_1000(bencher: &mut Bencher) {
    let vec = Big::soa_vec(1000);
    bencher.iter(||{
        let mut s = 0.0;
        for (position, velocity) in vec.position.iter().zip(&vec.velocity) {
            s += position.0 + velocity.0;
        }
        s
    })
}

fn soa_big_do_work_10000(bencher: &mut Bencher) {
    let vec = Big::soa_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for (position, velocity) in vec.position.iter().zip(&vec.velocity) {
            s += position.0 + velocity.0;
        }
        s
    })
}

benchmark_group!(aos,
    aos_small_push, aos_big_push, aos_small_do_work_10000, aos_big_do_work_1000,
    aos_big_do_work_10000
);
benchmark_group!(soa,
    soa_small_push, soa_big_push, soa_small_do_work_10000, soa_big_do_work_1000,
    soa_big_do_work_10000
);
benchmark_main!(soa, aos);
<commit_msg>Update benchmarks<commit_after>use soa_derive::StructOfArray;

use bencher::{Bencher, benchmark_group, benchmark_main};

#[derive(StructOfArray)]
pub struct Small {
    x: f64,
    y: f64,
    z: f64,
}

impl Small {
    fn new() -> Small {
        Small {
            x: 1.0,
            y: 0.2,
            z: -2.3,
        }
    }

    fn aos_vec(size: usize) -> Vec<Small> {
        let mut vec = Vec::new();
        for _ in 0..size {
            vec.push(Small::new())
        }
        return vec;
    }

    fn soa_vec(size: usize) -> SmallVec {
        let mut vec = SmallVec::new();
        for _ in 0..size {
            vec.push(Small::new())
        }
        return vec;
    }
}

#[derive(StructOfArray)]
pub struct Big {
    position: (f64, f64, f64),
    velocity: (f64, f64, f64),
    id: usize,
    name: String,
    userdata: String
}

impl Big {
    fn new() -> Big {
        Big {
            position: (1.0, 0.2, -2.3),
            velocity: (1.0, 0.2, -2.3),
            id: 67,
            name: "foo".into(),
            userdata: "bar".into()
        }
    }

    fn aos_vec(size: usize) -> Vec<Big> {
        let mut vec = Vec::new();
        for _ in 0..size {
            vec.push(Big::new())
        }
        return vec;
    }

    fn soa_vec(size: usize) -> BigVec {
        let mut vec = BigVec::new();
        for _ in 0..size {
            vec.push(Big::new())
        }
        return vec;
    }
}

fn aos_small_push(bencher: &mut Bencher) {
    let mut vec = Vec::new();
    bencher.iter(||{
        vec.push(Small::new())
    })
}

fn soa_small_push(bencher: &mut Bencher) {
    let mut vec = SmallVec::new();
    bencher.iter(||{
        vec.push(Small::new())
    })
}

fn aos_big_push(bencher: &mut Bencher) {
    let mut vec = Vec::new();
    bencher.iter(||{
        vec.push(Big::new())
    })
}

fn soa_big_push(bencher: &mut Bencher) {
    let mut vec = BigVec::new();
    bencher.iter(||{
        vec.push(Big::new())
    })
}

fn aos_small_do_work_10000(bencher: &mut Bencher) {
    let vec = Small::aos_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.x + v.y;
        }
        s
    })
}

fn soa_small_do_work_10000(bencher: &mut Bencher) {
    let vec = Small::soa_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for (x, y) in vec.x.iter().zip(&vec.y) {
            s += x + y;
        }
        s
    })
}

fn aos_big_do_work_1000(bencher: &mut Bencher) {
    let vec = Big::aos_vec(1000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.position.0 + v.velocity.0 * 0.1;
        }
        s
    })
}

fn aos_big_do_work_10000(bencher: &mut Bencher) {
    let vec = Big::aos_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for v in &vec {
            s += v.position.0 + v.velocity.0 * 0.1;
        }
        s
    })
}

fn soa_big_do_work_1000(bencher: &mut Bencher) {
    let vec = Big::soa_vec(1000);
    bencher.iter(||{
        let mut s = 0.0;
        for (position, velocity) in vec.position.iter().zip(&vec.velocity) {
            s += position.0 + velocity.0;
        }
        s
    })
}

fn soa_big_do_work_10000(bencher: &mut Bencher) {
    let vec = Big::soa_vec(10000);
    bencher.iter(||{
        let mut s = 0.0;
        for (position, velocity) in vec.position.iter().zip(&vec.velocity) {
            s += position.0 + velocity.0;
        }
        s
    })
}

benchmark_group!(aos,
    aos_small_push, aos_big_push, aos_small_do_work_10000, aos_big_do_work_1000,
    aos_big_do_work_10000
);
benchmark_group!(soa,
    soa_small_push, soa_big_push, soa_small_do_work_10000, soa_big_do_work_1000,
    soa_big_do_work_10000
);
benchmark_main!(soa, aos);
<|endoftext|>


<commit_before>//! Generates random trees and tree-like plants.

use prop::plant::{Branch, ControlPoint};
use math::{Basis3, Deg, Euler, InnerSpace, Point3f, Rotation, Vector1, Vector3f};
use rand::{Rand, Rng};
use rand::distributions::range::SampleRange;
use rand::distributions::{self, IndependentSample};
use std::ops::Range;

/// Parameters for the tree generator.
#[derive(Debug)]
struct Preset {
    name: &'static str,
    /// Diameter of the first branch we create (the trunk).
    trunk_diameter: Range<f32>,
    /// Trunk height. Note that branches going upward can increase plant height
    /// beyond this.
    trunk_height: Range<f32>,
    /// Trunk diameter at `trunk_height`. Should be smaller than
    /// `trunk_diameter`.
    trunk_diameter_top: Range<f32>,
    /// Trunk height at which we start creating branches.
    min_branch_height: Range<f32>,
    /// Range of subbranch diameters as a factor of the parent branch.
    branch_diameter_factor: Range<f32>,
    /// Range of subbranch angles in degrees.
    branch_angle_deg: Range<f32>,
    /// Factor by which to reduce segment diameter between consecutive points,
    /// sampled per branch.
    branch_diam_reduction: Range<f32>,
    /// Range of angles to use for rotation of new segments.
    ///
    /// The higher the angle, the more "twisted" branches appear.
    branch_segment_angle: Range<f32>,
    /// Range of segment counts for branches.
    ///
    /// Together with `branch_segment_length`, this defines the overall branch
    /// length.
    branch_segment_count: Range<u32>,
    /// Range of segment lengths to use for branches.
    ///
    /// Each segment will have a random length in this range.
    branch_segment_length: Range<f32>, // FIXME unused
}

static PRESETS: &'static [Preset] = &[Preset {
                                          name: "'Regular' Tree",
                                          trunk_diameter: 0.3..0.5,
                                          trunk_height: 3.0..6.0,
                                          trunk_diameter_top: 0.2..0.4,
                                          min_branch_height: 0.4..0.6,
                                          branch_diameter_factor: 0.3..0.5,
                                          branch_angle_deg: 70.0..110.0,
                                          branch_diam_reduction: 0.75..0.85,
                                          branch_segment_angle: 2.0..10.0,
                                          branch_segment_count: 5..20,
                                          branch_segment_length: 0.05..0.15,
                                      }];

pub struct TreeGen {
    preset: &'static Preset,
    /// Buffer for branches, filled as they're created.
    branches: Vec<Branch>,
}

impl TreeGen {
    /// Generates a new branch at the given start point.
    ///
    /// # Parameters
    ///
    /// * `rng`: The RNG to use
    /// * `start`: The point at which the branch should start
    /// * `dir`: Direction in which the branch should grow
    /// * `depth`: Recursion depth, used to limit the maximum branch depth
    /// * `parent_diam`: Diameter of the node of the parent branch where this
    ///   branch starts
    fn create_branch<R: Rng>(&mut self,
                             rng: &mut R,
                             start: Point3f,
                             dir: Vector3f,
                             depth: u32,
                             parent_diam: f32) {
        if depth > 5 {
            // Limit recursion
            return;
        }

        // Current normalized growing direction, variated slightly as segments are
        // generated
        let mut dir = dir.normalize();

        // Determine starting diameter of the new branch
        let mut diam = range_sample(&self.preset.branch_diameter_factor, rng) * parent_diam;
        // Determine how much segment diameter is reduced
        let diam_factor = range_sample(&self.preset.branch_diam_reduction, rng) * parent_diam;
        // How many segments should this branch get?
        let segment_count = range_sample(&self.preset.branch_segment_count, rng);

        let mut points = vec![ControlPoint {
                                  point: start,
                                  diameter: diam,
                              }];

        {
            let mut last = start;
            // Helper for adding a new point to this branch, which possibly grows a new
            // branch.
            let mut add_point = |dist, diam| {
                // First, get a random angle by which to variate this segment.
                let x_angle = range_sample(&self.preset.branch_segment_angle, rng);
                let y_angle = range_sample(&self.preset.branch_segment_angle, rng);

                let rotation = Basis3::from(Euler {
                    x: Deg::new(x_angle),
                    y: Deg::new(y_angle),
                    z: Deg::new(0.0),
                });
                dir = rotation.rotate_vector(dir);

                // We need to add a point with a distance of `dist` from `last`.
                let point = last + dir * dist;
                last = point;

                // FIXME Make branch spawn chance configurable
                // 1/20 chance to spawn a branch at any point
                if rng.gen_weighted_bool(20) {
                    // Build a vector for the branch direction (Z is up)
                    let dir = self.gen_branch_direction(rng, dir);
                    self.create_branch(rng, point, dir, depth + 1, diam);
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            // In a loop, get the length of the next segment from the current diameter.
            for _ in 0..segment_count {
                let length = segment_dist(diam);
                diam *= diam_factor;
                add_point(length, diam);
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });
    }

    /// Given the growing direction of the parent branch, calculates a growing
    /// direction to use for a new child branch.
    fn gen_branch_direction<R: Rng>(&self, rng: &mut R, parent_dir: Vector3f) -> Vector3f {
        // `branch_angle_deg` specifies the angle range in degrees
        let mut x_angle = range_sample(&self.preset.branch_angle_deg, rng);
        let mut y_angle = range_sample(&self.preset.branch_angle_deg, rng);

        // Invert sign with 50%, to mirror the specified range to the other side
        x_angle = if rng.gen() { -x_angle } else { x_angle };
        y_angle = if rng.gen() { -y_angle } else { y_angle };

        // Rotate the growing direction of the parent branch
        let rotation = Basis3::from(Euler {
            x: Deg::new(x_angle),
            y: Deg::new(y_angle),
            z: Deg::new(0.0),
        });
        rotation.rotate_vector(parent_dir)
    }

    fn create_trunk<R: Rng>(&mut self, rng: &mut R) {
        let trunk_diameter = range_sample(&self.preset.trunk_diameter, rng);
        let trunk_height = range_sample(&self.preset.trunk_height, rng);
        let trunk_diameter_top = range_sample(&self.preset.trunk_diameter_top, rng);
        let min_branch_height = range_sample(&self.preset.min_branch_height, rng) * trunk_height;

        debug!("trunk diam {} to {}, height {}, branch start at {}",
               trunk_diameter,
               trunk_diameter_top,
               trunk_height,
               min_branch_height);

        let mut points = Vec::new();

        {
            let mut add_point = |height, diam| {
                let point = Point3f::new(0.0, 0.0, height);
                if height >= min_branch_height {
                    // FIXME Make branch spawn chance configurable
                    // 1/3 chance to spawn a branch at any point
                    if rng.gen_weighted_bool(3) {
                        // Build a vector for the branch direction (Z is up)
                        let dir = self.gen_branch_direction(rng, Vector3f::new(0.0, 0.0, 1.0));
                        self.create_branch(rng, point, dir, 1, diam);
                    }
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            let diam_start = Vector1::new(trunk_diameter);
            let diam_end = Vector1::new(trunk_diameter_top);

            // Split trunk in segments
            // FIXME Vary the segment direction like we do for normal branches
            // FIXME Make segment count depend on the trunk height
            const SEGMENT_COUNT: u32 = 10;
            for i in 0..SEGMENT_COUNT + 1 {
                let height = i as f32 * trunk_height / SEGMENT_COUNT as f32;
                let height_frac = height / trunk_height;
                let diam = diam_start.lerp(diam_end, height_frac);

                add_point(height, diam.x);
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });

        debug!("generated tree with {} branches", self.branches.len());
    }

    /// Generates a random tree according to the stored parameters.
    ///
    /// The tree is returned as a list of branches for now.
    pub fn generate<R: Rng>(mut self, rng: &mut R) -> Vec<Branch> {
        // Recursively create the tree and put all branches in a buffer.
        self.create_trunk(rng);
        self.branches
    }
}

impl Rand for TreeGen {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        // Select a random preset that we'll use
        let preset = rng.choose(PRESETS).unwrap().clone();

        TreeGen {
            preset: preset,
            branches: Vec::new(),
        }
    }
}

/// Samples a random element from a range.
fn range_sample<T: SampleRange + PartialOrd + Copy, R: Rng>(range: &Range<T>, rng: &mut R) -> T {
    // Build a `rand` crate Range. We use `std`s Range for the cool `a..b` syntax ;)
    distributions::Range::new(range.start, range.end).ind_sample(rng)
}

/// Approximation of real-world distance of branch segments, depending on the
/// starting branch diameter.
fn segment_dist(diameter: f32) -> f32 {
    diameter * 11.25
}
<commit_msg>tweak treegen some more<commit_after>//! Generates random trees and tree-like plants.

use prop::plant::{Branch, ControlPoint};
use math::{ApproxEq, Basis3, Deg, Euler, InnerSpace, Point3f, Rad, Rotation, Rotation3, Vector1,
           Vector3f};
use rand::{Rand, Rng};
use rand::distributions::range::SampleRange;
use rand::distributions::{self, IndependentSample};
use std::ops::Range;

/// Parameters for the tree generator.
#[derive(Debug)]
struct Preset {
    name: &'static str,
    /// Diameter of the first branch we create (the trunk).
    trunk_diameter: Range<f32>,
    /// Trunk height. Note that branches going upward can increase plant height
    /// beyond this.
    trunk_height: Range<f32>,
    /// Trunk diameter at `trunk_height`. Should be smaller than
    /// `trunk_diameter`.
    trunk_diameter_top: Range<f32>,
    /// Trunk height at which we start creating branches.
    min_branch_height: Range<f32>,
    /// Range of subbranch diameters as a factor of the parent branch.
    branch_diameter_factor: Range<f32>,
    /// Range of subbranch angles in degrees.
    branch_angle_deg: Range<f32>,
    /// Factor by which to reduce segment diameter between consecutive points,
    /// sampled per branch.
    branch_diam_reduction: Range<f32>,
    /// Range of angles to use for rotation of new segments.
    ///
    /// The higher the angle, the more "twisted" branches appear.
    branch_segment_angle: Range<f32>,
    /// Range of segment counts for branches.
    ///
    /// Together with `branch_segment_length`, this defines the overall branch
    /// length.
    branch_segment_count: Range<u32>,
    /// Range of segment lengths to use for branches.
    ///
    /// Each segment will have a random length in this range.
    branch_segment_length: Range<f32>, // FIXME unused
}

static PRESETS: &'static [Preset] = &[Preset {
                                          name: "'Regular' Tree",
                                          trunk_diameter: 0.3..0.5,
                                          trunk_height: 3.0..6.0,
                                          trunk_diameter_top: 0.2..0.4,
                                          min_branch_height: 0.4..0.6,
                                          branch_diameter_factor: 0.3..0.5,
                                          branch_angle_deg: 70.0..110.0,
                                          branch_diam_reduction: 0.75..0.85,
                                          branch_segment_angle: 5.0..15.0,
                                          branch_segment_count: 3..10,
                                          branch_segment_length: 0.30..0.40,
                                      }];

pub struct TreeGen {
    preset: &'static Preset,
    /// Buffer for branches, filled as they're created.
    branches: Vec<Branch>,
}

impl TreeGen {
    /// Generates a new branch at the given start point.
    ///
    /// # Parameters
    ///
    /// * `rng`: The RNG to use
    /// * `start`: The point at which the branch should start
    /// * `dir`: Direction in which the branch should grow
    /// * `depth`: Recursion depth, used to limit the maximum branch depth
    /// * `parent_diam`: Diameter of the node of the parent branch where this
    ///   branch starts
    fn create_branch<R: Rng>(&mut self,
                             rng: &mut R,
                             start: Point3f,
                             dir: Vector3f,
                             depth: u32,
                             parent_diam: f32) {
        if depth > 5 {
            // Limit recursion
            return;
        }

        // Current normalized growing direction, variated slightly as segments are
        // generated
        let mut dir = dir.normalize();

        // Determine starting diameter of the new branch
        let mut diam = range_sample(&self.preset.branch_diameter_factor, rng) * parent_diam;
        // Determine how much segment diameter is reduced
        let diam_factor = range_sample(&self.preset.branch_diam_reduction, rng) * parent_diam;
        // How many segments should this branch get?
        let segment_count = range_sample(&self.preset.branch_segment_count, rng);

        let mut points = vec![ControlPoint {
                                  point: start,
                                  diameter: diam,
                              }];

        {
            let mut last = start;
            // Helper for adding a new point to this branch, which possibly grows a new
            // branch.
            let mut add_point = |dist, diam| {
                // First, get a random angle by which to variate this segment.
                let mut x_angle = range_sample(&self.preset.branch_segment_angle, rng);
                let mut y_angle = range_sample(&self.preset.branch_segment_angle, rng);

                // Invert sign with 50%, to mirror the specified range to the other side
                x_angle = if rng.gen() { -x_angle } else { x_angle };
                y_angle = if rng.gen() { -y_angle } else { y_angle };

                let rotation = Basis3::from(Euler {
                    x: Deg::new(x_angle),
                    y: Deg::new(y_angle),
                    z: Deg::new(0.0),
                });
                dir = rotation.rotate_vector(dir);

                // We need to add a point with a distance of `dist` from `last`.
                let point = last + dir * dist;
                last = point;

                // FIXME Make branch spawn chance configurable
                if rng.gen_weighted_bool(depth * 3) {
                    // Build a vector for the branch direction (Z is up)
                    let dir = self.gen_branch_direction(rng, dir);
                    self.create_branch(rng, point, dir, depth + 1, diam);
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            // In a loop, get the length of the next segment from the current diameter.
            for _ in 0..segment_count {
                let length = segment_dist(diam);
                diam *= diam_factor;

                add_point(length, diam);

                if diam < 0.001 {
                    // Bail out at 1mm
                    break;
                }
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });
    }

    /// Given the growing direction of the parent branch, calculates a growing
    /// direction to use for a new child branch.
    fn gen_branch_direction<R: Rng>(&self, rng: &mut R, parent_dir: Vector3f) -> Vector3f {
        // `branch_angle_deg` specifies the angle range in degrees
        let angle = range_sample(&self.preset.branch_angle_deg, rng);

        // AWESOME hack: Generate *any* vector that is not parallel to `parent_dir`,
        // calculate the cross product between it and `parent_dir`.
        // This gets us a vector perpendicular to `parent_dir`, which we can rotate
        // `parent_dir` around to tilt it by `angle` degrees.
        let mut rand_vec = Vector3f::rand(rng).normalize();
        while rand_vec.angle(parent_dir).approx_eq(&Rad::new(0.0)) {
            // They're parallel, so generate a new vector. This is probably unnecessary,
            // but someone
            // who *actually* knows their math instead of faking it should check.
            rand_vec = Vector3f::rand(rng);
        }

        // Create vector around we'll tilt `parent_dir`
        let rot_vec = parent_dir.cross(rand_vec);

        // Tilt the growing direction of the parent branch by about 90° to get the
        // direction of the
        let tilt_rotation = Basis3::from_axis_angle(rot_vec, Deg::new(angle).into());

        let spin_angle = range_sample(&(0.0..360.0), rng);

        // Then rotate this vector randomly (0-360°) around the parent branch
        let around_we_go = Basis3::from_axis_angle(parent_dir, Deg::new(spin_angle).into());

        around_we_go.rotate_vector(tilt_rotation.rotate_vector(parent_dir))
    }

    fn create_trunk<R: Rng>(&mut self, rng: &mut R) {
        let trunk_diameter = range_sample(&self.preset.trunk_diameter, rng);
        let trunk_height = range_sample(&self.preset.trunk_height, rng);
        let trunk_diameter_top = range_sample(&self.preset.trunk_diameter_top, rng);
        let min_branch_height = range_sample(&self.preset.min_branch_height, rng) * trunk_height;

        debug!("trunk diam {} to {}, height {}, branch start at {}",
               trunk_diameter,
               trunk_diameter_top,
               trunk_height,
               min_branch_height);

        let mut points = Vec::new();

        {
            let mut add_point = |height, diam| {
                let point = Point3f::new(0.0, 0.0, height);
                if height >= min_branch_height {
                    // FIXME Make branch spawn chance configurable
                    let branches = &[0, 0, 1, 1, 1, 2];
                    for _ in 0..*rng.choose(branches).unwrap() {
                        // Build a vector for the branch direction (Z is up)
                        let dir = self.gen_branch_direction(rng, Vector3f::new(0.0, 0.0, 1.0));
                        self.create_branch(rng, point, dir, 1, diam);
                    }
                }

                points.push(ControlPoint {
                    point: point,
                    diameter: diam,
                });
            };

            let diam_start = Vector1::new(trunk_diameter);
            let diam_end = Vector1::new(trunk_diameter_top);

            // Split trunk in segments
            // FIXME Vary the segment direction like we do for normal branches
            // FIXME Make segment count depend on the trunk height
            const SEGMENT_COUNT: u32 = 10;
            for i in 0..SEGMENT_COUNT + 1 {
                let height = i as f32 * trunk_height / SEGMENT_COUNT as f32;
                let height_frac = height / trunk_height;
                let diam = diam_start.lerp(diam_end, height_frac);

                add_point(height, diam.x);
            }
        }

        assert!(points.len() >= 2,
                "should've generated at least 2 points :(");
        self.branches.push(Branch {
            points: points,
            // FIXME Fixed color for now, we should use a configurable random color (or at least
            // make it brown).
            color: Vector3f::new(0.0, 1.0, 0.0),
        });

        debug!("generated tree with {} branches", self.branches.len());
    }

    /// Generates a random tree according to the stored parameters.
    ///
    /// The tree is returned as a list of branches for now.
    pub fn generate<R: Rng>(mut self, rng: &mut R) -> Vec<Branch> {
        // Recursively create the tree and put all branches in a buffer.
        self.create_trunk(rng);
        self.branches
    }
}

impl Rand for TreeGen {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        // Select a random preset that we'll use
        let preset = rng.choose(PRESETS).unwrap().clone();

        TreeGen {
            preset: preset,
            branches: Vec::new(),
        }
    }
}

/// Samples a random element from a range.
fn range_sample<T: SampleRange + PartialOrd + Copy, R: Rng>(range: &Range<T>, rng: &mut R) -> T {
    // Build a `rand` crate Range. We use `std`s Range for the cool `a..b` syntax ;)
    distributions::Range::new(range.start, range.end).ind_sample(rng)
}

/// Approximation of real-world distance of branch segments, depending on the
/// starting branch diameter.
fn segment_dist(diameter: f32) -> f32 {
    diameter * 11.25
}
<|endoftext|>