<commit_before>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate rand;

pub mod basetypes;
pub mod chess_move;
pub mod position;
pub mod search;
pub mod uci;
pub mod tt;

use std::process::exit;
use std::sync::Arc;
use std::time::SystemTime;
use std::collections::VecDeque;
use basetypes::*;
use tt::*;
use position::Position;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use search::{TimeManagement, Variation, SearchStatus, MultipvSearch};

/// The version of the program.
pub const VERSION: &'static str = "0.1";

// The name of the program.
const NAME: &'static str = "Alcibiades";

// The author of the program.
const AUTHOR: &'static str = "Evgeni Pandurski";

// The starting position in Forsyth–Edwards notation (FEN).
const STARTING_POSITION: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";

// Represents a condition for terminating the search.
enum PlayWhen {
    TimeManagement(TimeManagement), // Stop when `TimeManagement` says.
    MoveTime(u64), // Stop after the given number of milliseconds.
    Nodes(NodeCount), // Stop when the given number of nodes has been searched.
    Depth(u8), // Stop when the given search depth has been reached.
    Never, // An infinite search.
}

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    position: Position,
    current_depth: u8,

    // `Engine` will hand over the real work to `MultipvSearch`.
    search: MultipvSearch,

    // Tells the engine when it must stop thinking and play the best
    // move it has found.
    play_when: PlayWhen,

    // Tells the engine how many best lines to calculate and send to
    // the GUI (the first move in each best line is different). This
    // is the so called "MultiPV" mode.
    variations_count: usize,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // `true` if the engine is thinking in pondering mode.
    is_pondering: bool,

    // Marks the last time when a message was sent by the engine to
    // the GUI, or if no message has been sent yet -- the time when
    // the search stared.
    silent_since: SystemTime,

    // A queue for the messages send by the engine to the GUI.
    queue: VecDeque<EngineReply>,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);

        Engine {
            tt: tt.clone(),
            position: Position::from_fen(STARTING_POSITION).ok().unwrap(),
            current_depth: 0,
            search: MultipvSearch::new(tt),
            play_when: PlayWhen::Never,
            variations_count: 1,
            pondering_is_allowed: false,
            is_pondering: false,
            silent_since: SystemTime::now(),
            queue: VecDeque::new(),
        }
    }

    // A helper method. It it adds a progress report message to
    // `self.queue`.
    fn queue_progress_report(&mut self) {
        let &SearchStatus { depth, searched_nodes, nps, .. } = self.search.status();
        self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It it adds a message containing the current
    // (multi)PV to `self.queue`.
    fn queue_pv(&mut self) {
        let &SearchStatus { depth, ref variations, searched_nodes, duration_millis, nps, .. } =
            self.search.status();
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let bound_suffix = match bound {
                BOUND_EXACT => "",
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                _ => panic!("unexpected bound type"),
            };
            let mut moves_string = String::new();
            for m in moves {
                moves_string.push_str(&m.notation());
                moves_string.push(' ');
            }
            self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("multipv".to_string(), format!("{}", i + 1)),
            ("score".to_string(), format!("cp {}{}", value, bound_suffix)),
            ("time".to_string(), format!("{}", duration_millis)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
            ("pv".to_string(), moves_string),
        ]));
        }
        self.silent_since = SystemTime::now();
    }

    // A helper method. It it adds a message containing the current
    // best move to `self.queue`.
    fn queue_best_move(&mut self) {
        let &SearchStatus { ref variations, .. } = self.search.status();
        self.queue.push_back(EngineReply::BestMove {
            best_move: variations[0].moves.get(0).map_or("0000".to_string(), |m| m.notation()),
            ponder_move: variations[0].moves.get(1).map(|m| m.notation()),
        });
        self.silent_since = SystemTime::now();
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            "MultiPV" => {
                self.variations_count = match value.parse::<usize>().unwrap_or(0) {
                    0 => 1,
                    n if n > 500 => 500,
                    n => n,
                };
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        self.tt.clear();
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if let Ok(p) = Position::from_history(fen, moves) {
            self.position = p;
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        // NOTE: We ignore the "mate" parameter.

        self.search.stop();
        self.tt.new_search();
        self.current_depth = 0;
        self.is_pondering = ponder;
        self.play_when = if infinite {
            PlayWhen::Never
        } else if movetime.is_some() {
            PlayWhen::MoveTime(movetime.unwrap())
        } else if nodes.is_some() {
            PlayWhen::Nodes(nodes.unwrap())
        } else if depth.is_some() {
            PlayWhen::Depth(depth.unwrap() as u8)
        } else {
            PlayWhen::TimeManagement(TimeManagement::new(self.position.board().to_move(),
                                                         self.pondering_is_allowed,
                                                         wtime,
                                                         btime,
                                                         winc,
                                                         binc,
                                                         movestogo))
        };
        self.silent_since = SystemTime::now();
        self.search.start(&self.position, searchmoves, self.variations_count);
    }

    fn ponder_hit(&mut self) {
        if self.search.status().done {
            self.queue_best_move();
        } else {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        self.search.stop();
        self.queue_best_move();
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        if !self.search.status().done {
            // Update the search status.
            self.search.update_status();
            let &SearchStatus { done, depth, searched_nodes, duration_millis, .. } = self.search
                                                                                .status();

            // Send the (multi)PV when changed.
            if depth > self.current_depth {
                self.current_depth = depth;
                self.queue_pv();
            }

            // Send periodic progress reports.
            if self.silent_since.elapsed().unwrap().as_secs() > 10 {
                self.queue_progress_report();
            }

            // Check if we must play now.
            if !self.is_pondering &&
               match self.play_when {
                PlayWhen::TimeManagement(ref tm) => done || tm.must_play(self.search.status()),
                PlayWhen::MoveTime(t) => done || duration_millis >= t,
                PlayWhen::Nodes(n) => done || searched_nodes >= n,
                PlayWhen::Depth(d) => done || depth >= d,
                PlayWhen::Never => false,
            } {
                self.stop();
            }
        }

        self.queue.pop_front()
    }

    fn exit(&mut self) {
        self.search.exit();
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("{} {}", NAME, VERSION)
    }

    fn author(&self) -> String {
        AUTHOR.to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 500, default: 1 }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

fn main() {
    if let Ok(mut uci_loop) = uci::Server::wait_for_hanshake(EngineFactory) {
        match uci_loop.serve() {
            Ok(_) => {
                exit(0);
            }
            Err(_) => {
                exit(1);
            }
        }
    }
    exit(2);
}
<commit_msg>minor beautification<commit_after>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate rand;

pub mod basetypes;
pub mod chess_move;
pub mod position;
pub mod search;
pub mod uci;
pub mod tt;

use std::process::exit;
use std::sync::Arc;
use std::time::SystemTime;
use std::collections::VecDeque;
use basetypes::NodeCount;
use tt::{TranspositionTable, BOUND_EXACT, BOUND_UPPER, BOUND_LOWER};
use position::Position;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use search::{TimeManagement, Variation, SearchStatus, MultipvSearch};

/// The version of the program.
pub const VERSION: &'static str = "0.1";

// The name of the program.
const NAME: &'static str = "Alcibiades";

// The author of the program.
const AUTHOR: &'static str = "Evgeni Pandurski";

// The starting position in Forsyth–Edwards notation (FEN).
const STARTING_POSITION: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";

// Represents a condition for terminating the search.
enum PlayWhen {
    TimeManagement(TimeManagement), // Stop when `TimeManagement` says.
    MoveTime(u64), // Stop after the given number of milliseconds.
    Nodes(NodeCount), // Stop when the given number of nodes has been searched.
    Depth(u8), // Stop when the given search depth has been reached.
    Never, // An infinite search.
}

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    position: Position,
    current_depth: u8,

    // `Engine` will hand over the real work to `MultipvSearch`.
    search: MultipvSearch,

    // Tells the engine when it must stop thinking and play the best
    // move it has found.
    play_when: PlayWhen,

    // Tells the engine how many best lines to calculate and send to
    // the GUI (the first move in each best line is different). This
    // is the so called "MultiPV" mode.
    variations_count: usize,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // `true` if the engine is thinking in pondering mode.
    is_pondering: bool,

    // Marks the last time when a message was sent by the engine to
    // the GUI, or if no message has been sent yet -- the time when
    // the search stared.
    silent_since: SystemTime,

    // A queue for the messages send by the engine to the GUI.
    queue: VecDeque<EngineReply>,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);

        Engine {
            tt: tt.clone(),
            position: Position::from_fen(STARTING_POSITION).ok().unwrap(),
            current_depth: 0,
            search: MultipvSearch::new(tt),
            play_when: PlayWhen::Never,
            variations_count: 1,
            pondering_is_allowed: false,
            is_pondering: false,
            silent_since: SystemTime::now(),
            queue: VecDeque::new(),
        }
    }

    // A helper method. It it adds a progress report message to
    // `self.queue`.
    fn queue_progress_report(&mut self) {
        let &SearchStatus { depth, searched_nodes, nps, .. } = self.search.status();
        self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It it adds a message containing the current
    // (multi)PV to `self.queue`.
    fn queue_pv(&mut self) {
        let &SearchStatus { depth, ref variations, searched_nodes, duration_millis, nps, .. } =
            self.search.status();
        for (i, &Variation { ref moves, value, bound }) in variations.iter().enumerate() {
            let bound_suffix = match bound {
                BOUND_EXACT => "",
                BOUND_UPPER => " upperbound",
                BOUND_LOWER => " lowerbound",
                _ => panic!("unexpected bound type"),
            };
            let mut moves_string = String::new();
            for m in moves {
                moves_string.push_str(&m.notation());
                moves_string.push(' ');
            }
            self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("multipv".to_string(), format!("{}", i + 1)),
            ("score".to_string(), format!("cp {}{}", value, bound_suffix)),
            ("time".to_string(), format!("{}", duration_millis)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
            ("pv".to_string(), moves_string),
        ]));
        }
        self.silent_since = SystemTime::now();
    }

    // A helper method. It it adds a message containing the current
    // best move to `self.queue`.
    fn queue_best_move(&mut self) {
        let &SearchStatus { ref variations, .. } = self.search.status();
        self.queue.push_back(EngineReply::BestMove {
            best_move: variations[0].moves.get(0).map_or("0000".to_string(), |m| m.notation()),
            ponder_move: variations[0].moves.get(1).map(|m| m.notation()),
        });
        self.silent_since = SystemTime::now();
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            "MultiPV" => {
                self.variations_count = match value.parse::<usize>().unwrap_or(0) {
                    0 => 1,
                    n if n > 500 => 500,
                    n => n,
                };
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        self.tt.clear();
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if let Ok(p) = Position::from_history(fen, moves) {
            self.position = p;
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        // NOTE: We ignore the "mate" parameter.

        self.search.stop();
        self.tt.new_search();
        self.current_depth = 0;
        self.is_pondering = ponder;
        self.play_when = if infinite {
            PlayWhen::Never
        } else if movetime.is_some() {
            PlayWhen::MoveTime(movetime.unwrap())
        } else if nodes.is_some() {
            PlayWhen::Nodes(nodes.unwrap())
        } else if depth.is_some() {
            PlayWhen::Depth(depth.unwrap() as u8)
        } else {
            PlayWhen::TimeManagement(TimeManagement::new(self.position.board().to_move(),
                                                         self.pondering_is_allowed,
                                                         wtime,
                                                         btime,
                                                         winc,
                                                         binc,
                                                         movestogo))
        };
        self.silent_since = SystemTime::now();
        self.search.start(&self.position, searchmoves, self.variations_count);
    }

    fn ponder_hit(&mut self) {
        if self.search.status().done {
            self.queue_best_move();
        } else {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        self.search.stop();
        self.queue_best_move();
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        if !self.search.status().done {
            // Update the search status.
            self.search.update_status();
            let &SearchStatus { done, depth, searched_nodes, duration_millis, .. } = self.search
                                                                                .status();

            // Send the (multi)PV when changed.
            if depth > self.current_depth {
                self.current_depth = depth;
                self.queue_pv();
            }

            // Send periodic progress reports.
            if self.silent_since.elapsed().unwrap().as_secs() > 10 {
                self.queue_progress_report();
            }

            // Check if we must play now.
            if !self.is_pondering &&
               match self.play_when {
                PlayWhen::TimeManagement(ref tm) => done || tm.must_play(self.search.status()),
                PlayWhen::MoveTime(t) => done || duration_millis >= t,
                PlayWhen::Nodes(n) => done || searched_nodes >= n,
                PlayWhen::Depth(d) => done || depth >= d,
                PlayWhen::Never => false,
            } {
                self.stop();
            }
        }

        self.queue.pop_front()
    }

    fn exit(&mut self) {
        self.search.exit();
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("{} {}", NAME, VERSION)
    }

    fn author(&self) -> String {
        AUTHOR.to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 500, default: 1 }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

fn main() {
    if let Ok(mut uci_loop) = uci::Server::wait_for_hanshake(EngineFactory) {
        match uci_loop.serve() {
            Ok(_) => {
                exit(0);
            }
            Err(_) => {
                exit(1);
            }
        }
    }
    exit(2);
}
<|endoftext|>


<commit_before>use std::env;
use std::io::prelude::*;
use std::fs::File;
use std::path::PathBuf;
extern crate uuid;
use uuid::Uuid;
extern crate getopts;
use getopts::Options;
use std::io;

fn print_usage(program: &str, opts: Options) {
    let path = PathBuf::from(program);
    let command = path.file_name().unwrap().to_string_lossy();

    let brief = format!("Usage: {} FILE [options]", command);
    let info = "Safely rewrite contents of FILE with stdin, even\nwhere FILE is being read by \
                upstream command";
    print!("{}", opts.usage(&[&brief, info].join("\n")));
}

fn redirect_to_file(outfile: &str) -> Result<(), io::Error> {
    let mut tempfile = env::temp_dir();
    tempfile.push(Uuid::new_v4().hyphenated().to_string());
    // println!("{}", tempfile.display());

    {
        let mut buffer = [0; 512];
        let mut stdin = io::stdin();
        let mut f = File::create(&tempfile).expect("Failed to create temporary file!");

        loop {
            let read_bytes = stdin.read(&mut buffer).expect("Error reading from stdin!");
            if read_bytes == 0 {
                break;
            }

            let write_bytes = f.write(&buffer[0..read_bytes])
                .expect("Failed to write to temporary file!");

            assert!(write_bytes == read_bytes);
        }
    }

    match std::fs::rename(&tempfile, &outfile) {
        Ok(m) => m,
        _ => {
            // fs::rename does not support cross-device linking
            // copy and delete instead
            std::fs::copy(&tempfile, &outfile).expect("Failed to create output file!");
            std::fs::remove_file(&tempfile)
                .unwrap_or_else(|e| print!("Failed to delete temporary file!\n{:?}", e));
        }
    };

    return Ok(());
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optflag("h", "help", "prints this help info");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let infile = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    assert!(redirect_to_file(&infile).is_ok());
}
<commit_msg>Write temporary file to the same directory as outfile<commit_after>use std::env;
use std::io::prelude::*;
use std::fs::File;
use std::path::PathBuf;
extern crate uuid;
use uuid::Uuid;
extern crate getopts;
use getopts::Options;
use std::io;

fn print_usage(program: &str, opts: Options) {
    let path = PathBuf::from(program);
    let command = path.file_name().unwrap().to_string_lossy();

    let brief = format!("Usage: {} FILE [options]", command);
    let info = "Safely rewrite contents of FILE with stdin, even\nwhere FILE is being read by \
                upstream command";
    print!("{}", opts.usage(&[&brief, info].join("\n")));
}

fn redirect_to_file(outfile: &str) -> Result<(), io::Error> {
    // create the temporary file in the same directory as outfile
    // this lets us guarantee a rename (instead of a move) on completion
    let mut tempfile = PathBuf::from(outfile);
    tempfile.pop(); //now refers to parent, which might be nothing
    tempfile.push(Uuid::new_v4().hyphenated().to_string());
    // println!("{}", tempfile.display());

    {
        let mut buffer = [0; 512];
        let mut stdin = io::stdin();
        let mut f = File::create(&tempfile).expect("Failed to create temporary file!");

        loop {
            let read_bytes = stdin.read(&mut buffer).expect("Error reading from stdin!");
            if read_bytes == 0 {
                break;
            }

            let write_bytes = f.write(&buffer[0..read_bytes])
                .expect("Failed to write to temporary file!");

            assert!(write_bytes == read_bytes);
        }
    }

    match std::fs::rename(&tempfile, &outfile) {
        Ok(m) => m,
        _ => {
            // fs::rename does not support cross-device linking
            // copy and delete instead
            std::fs::copy(&tempfile, &outfile).expect("Failed to create output file!");
            std::fs::remove_file(&tempfile)
                .unwrap_or_else(|e| print!("Failed to delete temporary file!\n{:?}", e));
        }
    };

    return Ok(());
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optflag("h", "help", "prints this help info");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let infile = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    assert!(redirect_to_file(&infile).is_ok());
}
<|endoftext|>


<commit_before>// TODO: Fix unwraps

extern crate git2;
extern crate term_painter;

use std::error::Error;
use std::io;
use std::env;
use std::path::{Path, PathBuf};

use term_painter::{ToStyle, Color};

fn main() {
    let working_dir = match env::current_dir() {
        Ok(path) => path,
        Err(err) => {
            println!("Error getting working directory: {}", err.description());
            return;
        }
    };

    println!("{:?}", walk_dirs(&working_dir));
}

fn walk_dirs(path: &Path) -> io::Result<()> {
    let mut pending: Vec<PathBuf> = Vec::new();

    loop {
        let current_dir = if pending.len() == 0 {
            path.to_path_buf()
        } else {
            pending.pop().unwrap()
        };

        let read_result = current_dir.read_dir();

        if read_result.is_ok() {
            let non_repo_iter = read_result.unwrap()
                .filter(|x| x.is_ok())
                .map(|x| x.unwrap())
                .filter(|x| match x.file_type() {
                    Ok(t) => t.is_dir(),
                    Err(_) => false,
                })
                .filter(|x| match x.path().file_name() {
                    Some(name) => {
                        match name.to_str() {
                            Some(name_str) => {
                                !name_str.starts_with(".") && !name_str.starts_with("$")
                            }
                            None => false,
                        }
                    }
                    None => false,
                })
                .filter(|x| match git_changes(&x.path()) {
                    Err(GitError::OpenRepo) => true, // Only return folders that arent repos
                    _ => false,
                })
                .map(|x| x.path().to_path_buf());

            for path in non_repo_iter {
                pending.push(path);
            }
        }

        if pending.len() == 0 {
            break;
        }
    }

    Ok(())
}

enum GitError {
    OpenRepo,
    Status,
}

// This can become its own iterator
fn git_changes(path: &Path) -> Result<(), GitError> {
    let repo = git2::Repository::open(path).map_err(|_| GitError::OpenRepo)?;

    if repo.is_bare() {
        return Ok(());
    }

    let mut opts = git2::StatusOptions::new();

    opts.include_ignored(false)
        .include_untracked(true)
        .recurse_untracked_dirs(true)
        .include_unreadable_as_untracked(true)
        .disable_pathspec_match(true)
        .exclude_submodules(true);

    let statuses = repo.statuses(Some(&mut opts)).map_err(|_| GitError::Status)?;
    let mut statuses_iter = statuses.iter()
        .filter(|x| {
            if x.status() != git2::STATUS_WT_DELETED {
                return true;
            }

            // For some reason, some files with the deleted
            // status actually still exist, so ignore these.
            match x.path() {
                None => false,
                Some(p) => {
                    let mut del_path = path.to_path_buf();
                    del_path.push(p);
                    !del_path.exists()
                }
            }
        })
        .peekable();

    if statuses_iter.peek().is_none() {
        return Ok(());
    }

    println!("{}", path.to_str().unwrap());

    for entry in statuses_iter {
        let pre = match entry.status() {
            git2::STATUS_WT_DELETED => Color::BrightRed.paint("    Deleted"),
            git2::STATUS_WT_MODIFIED => Color::BrightCyan.paint("   Modified"),
            git2::STATUS_WT_NEW => Color::BrightGreen.paint("        New"),
            git2::STATUS_WT_RENAMED => Color::BrightCyan.paint("    Renamed"),
            git2::STATUS_WT_TYPECHANGE => Color::BrightCyan.paint("Typechanged"),
            _ => Color::BrightMagenta.paint("    Unknown"),
        };

        println!("  {} {}", pre, entry.path().unwrap());
    }

    Ok(())
}
<commit_msg>Only print out error messages, nothing on Ok<commit_after>// TODO: Fix unwraps

extern crate git2;
extern crate term_painter;

use std::error::Error;
use std::io;
use std::env;
use std::path::{Path, PathBuf};

use term_painter::{ToStyle, Color};

fn main() {
    let working_dir = match env::current_dir() {
        Ok(path) => path,
        Err(err) => {
            println!("Error getting working directory: {}", err.description());
            return;
        }
    };

    match walk_dirs(&working_dir) {
        Err(e) => println!("{:?}", e),
        _ => {}
    }
}

fn walk_dirs(path: &Path) -> io::Result<()> {
    let mut pending: Vec<PathBuf> = Vec::new();

    loop {
        let current_dir = if pending.len() == 0 {
            path.to_path_buf()
        } else {
            pending.pop().unwrap()
        };

        let read_result = current_dir.read_dir();

        if read_result.is_ok() {
            let non_repo_iter = read_result.unwrap()
                .filter(|x| x.is_ok())
                .map(|x| x.unwrap())
                .filter(|x| match x.file_type() {
                    Ok(t) => t.is_dir(),
                    Err(_) => false,
                })
                .filter(|x| match x.path().file_name() {
                    Some(name) => {
                        match name.to_str() {
                            Some(name_str) => {
                                !name_str.starts_with(".") && !name_str.starts_with("$")
                            }
                            None => false,
                        }
                    }
                    None => false,
                })
                .filter(|x| match git_changes(&x.path()) {
                    Err(GitError::OpenRepo) => true, // Only return folders that arent repos
                    _ => false,
                })
                .map(|x| x.path().to_path_buf());

            for path in non_repo_iter {
                pending.push(path);
            }
        }

        if pending.len() == 0 {
            break;
        }
    }

    Ok(())
}

enum GitError {
    OpenRepo,
    Status,
}

// This can become its own iterator
fn git_changes(path: &Path) -> Result<(), GitError> {
    let repo = git2::Repository::open(path).map_err(|_| GitError::OpenRepo)?;

    if repo.is_bare() {
        return Ok(());
    }

    let mut opts = git2::StatusOptions::new();

    opts.include_ignored(false)
        .include_untracked(true)
        .recurse_untracked_dirs(true)
        .include_unreadable_as_untracked(true)
        .disable_pathspec_match(true)
        .exclude_submodules(true);

    let statuses = repo.statuses(Some(&mut opts)).map_err(|_| GitError::Status)?;
    let mut statuses_iter = statuses.iter()
        .filter(|x| {
            if x.status() != git2::STATUS_WT_DELETED {
                return true;
            }

            // For some reason, some files with the deleted
            // status actually still exist, so ignore these.
            match x.path() {
                None => false,
                Some(p) => {
                    let mut del_path = path.to_path_buf();
                    del_path.push(p);
                    !del_path.exists()
                }
            }
        })
        .peekable();

    if statuses_iter.peek().is_none() {
        return Ok(());
    }

    println!("{}", path.to_str().unwrap());

    for entry in statuses_iter {
        let pre = match entry.status() {
            git2::STATUS_WT_DELETED => Color::BrightRed.paint("    Deleted"),
            git2::STATUS_WT_MODIFIED => Color::BrightCyan.paint("   Modified"),
            git2::STATUS_WT_NEW => Color::BrightGreen.paint("        New"),
            git2::STATUS_WT_RENAMED => Color::BrightCyan.paint("    Renamed"),
            git2::STATUS_WT_TYPECHANGE => Color::BrightCyan.paint("Typechanged"),
            _ => Color::BrightMagenta.paint("    Unknown"),
        };

        println!("  {} {}", pre, entry.path().unwrap());
    }

    Ok(())
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

mod logger;
mod toml_file;
mod git;
mod changelog;
mod commit_analyzer;
mod cargo;
mod error;
mod github;
mod config;
mod utils;

extern crate rustc_serialize;
extern crate toml;
extern crate regex;
extern crate semver;
extern crate docopt;
extern crate git2;
extern crate clog;
extern crate hyper;
extern crate hubcaps;
extern crate url;
extern crate travis_after_all;
extern crate env_logger;

use docopt::Docopt;
use commit_analyzer::CommitType;
use config::ConfigBuilder;
use std::process;
use semver::Version;
use std::{env,fs};
use std::path::Path;
use std::error::Error;
use std::thread;
use std::time::Duration;
use travis_after_all::Build;
use utils::user_repo_from_url;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const USERAGENT: &'static str = concat!("semantic-rs/", env!("CARGO_PKG_VERSION"));
const USAGE: &'static str = "
semantic.rs 🚀

Usage:
  semantic-rs [options]
  semantic-rs --version

Options:
  -h --help              Show this screen.
  --version              Show version.
  -p PATH, --path=PATH   Specifies the repository path. [default: .]
  -w W, --write=W        Write changes to files (default: yes if CI is set, otherwise no).
  -r R, --release=R      Create release on GitHub and publish on crates.io (only in write mode) [default: yes].
  -b B, --branch=B       The branch on which releases should happen. [default: master].
";

const COMMITTER_ERROR_MESSAGE: &'static str = r"
A release commit needs a committer name and email address.
We tried fetching it from different locations, but couldn't find one.

Committer information is taken from the following environment variables, if set:

GIT_COMMITTER_NAME
GIT_COMMITTER_EMAIL

If none is set the normal git config is tried in the following order:

Local repository config
User config
Global config";

macro_rules! print_exit {
    ($fmt:expr) => {{
        logger::stderr($fmt);
        process::exit(1);
    }};
    ($fmt:expr, $($arg:tt)*) => {{
        logger::stderr(format!($fmt, $($arg)*));
        process::exit(1);
    }};
}

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_path: String,
    flag_write: Option<String>,
    flag_version: bool,
    flag_release: String,
    flag_branch: String,
}

fn string_to_bool(answer: &str) -> bool {
    match &answer.to_lowercase()[..] {
        "yes" | "true" | "1" => true,
        _ => false
    }
}

fn version_bump(version: &Version, bump: CommitType) -> Option<Version> {
    let mut version = version.clone();
    match bump {
        CommitType::Unknown => return None,
        CommitType::Patch => version.increment_patch(),
        CommitType::Minor => version.increment_minor(),
        CommitType::Major => version.increment_major(),
    }

    Some(version)
}

fn ci_env_set() -> bool {
    env::var("CI").is_ok()
}

fn current_branch(repo: &git2::Repository) -> Option<String> {
    if let Ok(branch) = env::var("TRAVIS_BRANCH") {
        return Some(branch)
    }

    let head = repo.head().expect("No HEAD found for repository");

    if head.is_branch() {
        let short = head.shorthand().expect("No branch name found");
        return Some(short.into());
    }

    None
}

fn is_release_branch(current: &str, release: &str) -> bool {
    if let Ok(pr) = env::var("TRAVIS_PULL_REQUEST") {
        if pr != "false" {
            return false;
        }
    }

    current == release
}

fn push_to_github(config: &config::Config, tag_name: &str) {
    logger::stdout("Pushing new commit and tag");
    git::push(&config, &tag_name)
        .unwrap_or_else(|err| print_exit!("Failed to push git: {:?}", err));

    logger::stdout("Waiting a tiny bit, so GitHub can store the git tag");
    thread::sleep(Duration::from_secs(1));
}

fn release_on_github(config: &config::Config, tag_message: &str, tag_name: &str) {
    if github::can_release(&config) {
        logger::stdout("Creating GitHub release");
        github::release(&config, &tag_name, &tag_message)
            .unwrap_or_else(|err| print_exit!("Failed to create GitHub release: {:?}", err));
    } else {
        logger::stdout("Project not hosted on GitHub. Skipping release step");
    }
}

fn release_on_cratesio(config: &config::Config) {
    logger::stdout("Publishing crate on crates.io");
    if !cargo::publish(&config.repository_path, &config.cargo_token.as_ref().unwrap()) {
        print_exit!("Failed to publish on crates.io");
    }
}

fn generate_changelog(repository_path: &str, version: &Version, new_version: &String) -> String {
    logger::stdout(format!("New version would be: {}", new_version));
    logger::stdout("Would write the following Changelog:");
    match changelog::generate(repository_path, &version.to_string(), new_version) {
        Ok(log) => log,
        Err(err) => {
            logger::stderr(format!("Generating Changelog failed: {:?}", err));
            process::exit(1)
        }
    }
}

fn write_changelog(repository_path: &str, version: &Version, new_version: &str) {
    logger::stdout("Writing Changelog");
    changelog::write(repository_path, &version.to_string(), &new_version)
        .unwrap_or_else(|err| print_exit!("Writing Changelog failed: {:?}", err));
}

fn print_changelog(changelog: &str) {
    logger::stdout("====================================");
    logger::stdout(changelog);
    logger::stdout("====================================");
    logger::stdout("Would create annotated git tag");
}

fn package_crate(config: &config::Config, repository_path: &str, new_version: &str) {
    if config.release_mode {
        logger::stdout("Updating lockfile");
        if !cargo::update_lockfile(repository_path) {
            print_exit!("`cargo fetch` failed. See above for the cargo error message.");
        }
    }

    git::commit_files(&config, &new_version)
        .unwrap_or_else(|err| print_exit!("Committing files failed: {:?}", err));

    logger::stdout("Package crate");
    if !cargo::package(repository_path) {
        print_exit!("`cargo package` failed. See above for the cargo error message.");
    }
}

fn get_repo(repository_path: &str) -> git2::Repository {
    match git2::Repository::open(repository_path) {
        Ok(repo) => repo,
        Err(e) => {
            logger::stderr(format!("Could not open the git repository: {:?}", e));
            process::exit(1);
        }
    }
}

fn get_repository_path(args: &Args) -> String {
    let path = Path::new(&args.flag_path);
    let path = fs::canonicalize(path)
        .unwrap_or_else(|_| print_exit!("Path does not exist or a component is
                                                            not a directory"));
    let repo_path = path.to_str().unwrap_or_else(|| print_exit!("Path is not valid unicode"));
    repo_path.to_string()
}

fn get_signature<'a>(repository_path: String) -> git2::Signature<'a> {
    let repo = get_repo(&repository_path);
    let signature = match git::get_signature(&repo) {
        Ok(sig) => sig,
            Err(e) => {
                logger::stderr(format!("Failed to get the committer's name and email address: {}", e.description()));
                logger::stderr(COMMITTER_ERROR_MESSAGE);
                process::exit(1);
            }
    };

    signature.to_owned()
}

fn get_user_and_repo(repository_path: &str) -> (Option<String>, Option<String>) {
    let repo = get_repo(repository_path);
    let remote_or_none = repo.find_remote("origin");
    match remote_or_none {
        Ok(remote) => {
            let url = remote.url().expect("Remote URL is not valid UTF-8").to_owned();
            let (user, repo_name) = user_repo_from_url(&url)
                .unwrap_or_else(|e| print_exit!("Could not extract user and repository name from URL: {:?}", e));

            (Some(user), Some(repo_name))

            //We'll do that separately
            // let cargo_token = env::var("CARGO_TOKEN")
            //     .unwrap_or_else(|err| print_exit!("CARGO_TOKEN not set: {:?}", err));

            // config_builder.cargo_token(cargo_token);
        },
        Err(err) => {
            logger::warn(format!("Could not determine the origin remote url: {:?}", err));
            logger::warn("semantic-rs can't push changes or create a release on GitHub");
            (None, None)
        }
    }
}

fn get_github_token(repository_path: &str) -> Option<String> {
    let repo = get_repo(repository_path);
    let remote_or_none = repo.find_remote("origin");
    match remote_or_none {
        Ok(remote) => {
            let url = remote.url().expect("Remote URL is not valid UTF-8").to_owned();
            if github::is_github_url(&url) {
                match env::var("GH_TOKEN") {
                    Ok(token) => Some(token),
                    Err(err) => {
                        logger::warn(format!("GH_TOKEN not set: {:?}", err));
                        None
                    }
                }
            } else {
                None
            }
        },
        Err(_) => None
    }
}

fn main() {
    env_logger::init().expect("Can't instantiate env logger");
    println!("semantic.rs 🚀");

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let mut config_builder = ConfigBuilder::new();

    if args.flag_version {
        println!("semantic.rs 🚀 -- v{}", VERSION);
        process::exit(0);
    }

    // If write mode is requested OR denied,
    // adhere to the user's wish,
    // otherwise we decide based on whether we are running in CI.
    let write_mode = match args.flag_write {
        None => ci_env_set(),
        Some(ref flag) => string_to_bool(flag)
    };

    // We can only release, if we are allowed to write
    let release_mode = write_mode && string_to_bool(&args.flag_release);
    let repository_path = get_repository_path(&args);

    config_builder.write(write_mode);
    config_builder.release(release_mode);
    config_builder.branch(args.flag_branch.clone());
    config_builder.repository_path(repository_path.clone());
    config_builder.signature(get_signature(repository_path.clone()));
    let (user, repo) = get_user_and_repo(&repository_path);
    if user.is_some() {
        config_builder.user(user.unwrap());
    }
    if repo.is_some() {
        config_builder.user(repo.unwrap());
    }
    let gh_token = get_github_token(&repository_path);
    if gh_token.is_some() {
        config_builder.gh_token(gh_token.unwrap());
    }
}
<commit_msg>refactor: Set cargo token<commit_after>#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

mod logger;
mod toml_file;
mod git;
mod changelog;
mod commit_analyzer;
mod cargo;
mod error;
mod github;
mod config;
mod utils;

extern crate rustc_serialize;
extern crate toml;
extern crate regex;
extern crate semver;
extern crate docopt;
extern crate git2;
extern crate clog;
extern crate hyper;
extern crate hubcaps;
extern crate url;
extern crate travis_after_all;
extern crate env_logger;

use docopt::Docopt;
use commit_analyzer::CommitType;
use config::ConfigBuilder;
use std::process;
use semver::Version;
use std::{env,fs};
use std::path::Path;
use std::error::Error;
use std::thread;
use std::time::Duration;
use travis_after_all::Build;
use utils::user_repo_from_url;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const USERAGENT: &'static str = concat!("semantic-rs/", env!("CARGO_PKG_VERSION"));
const USAGE: &'static str = "
semantic.rs 🚀

Usage:
  semantic-rs [options]
  semantic-rs --version

Options:
  -h --help              Show this screen.
  --version              Show version.
  -p PATH, --path=PATH   Specifies the repository path. [default: .]
  -w W, --write=W        Write changes to files (default: yes if CI is set, otherwise no).
  -r R, --release=R      Create release on GitHub and publish on crates.io (only in write mode) [default: yes].
  -b B, --branch=B       The branch on which releases should happen. [default: master].
";

const COMMITTER_ERROR_MESSAGE: &'static str = r"
A release commit needs a committer name and email address.
We tried fetching it from different locations, but couldn't find one.

Committer information is taken from the following environment variables, if set:

GIT_COMMITTER_NAME
GIT_COMMITTER_EMAIL

If none is set the normal git config is tried in the following order:

Local repository config
User config
Global config";

macro_rules! print_exit {
    ($fmt:expr) => {{
        logger::stderr($fmt);
        process::exit(1);
    }};
    ($fmt:expr, $($arg:tt)*) => {{
        logger::stderr(format!($fmt, $($arg)*));
        process::exit(1);
    }};
}

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_path: String,
    flag_write: Option<String>,
    flag_version: bool,
    flag_release: String,
    flag_branch: String,
}

fn string_to_bool(answer: &str) -> bool {
    match &answer.to_lowercase()[..] {
        "yes" | "true" | "1" => true,
        _ => false
    }
}

fn version_bump(version: &Version, bump: CommitType) -> Option<Version> {
    let mut version = version.clone();
    match bump {
        CommitType::Unknown => return None,
        CommitType::Patch => version.increment_patch(),
        CommitType::Minor => version.increment_minor(),
        CommitType::Major => version.increment_major(),
    }

    Some(version)
}

fn ci_env_set() -> bool {
    env::var("CI").is_ok()
}

fn current_branch(repo: &git2::Repository) -> Option<String> {
    if let Ok(branch) = env::var("TRAVIS_BRANCH") {
        return Some(branch)
    }

    let head = repo.head().expect("No HEAD found for repository");

    if head.is_branch() {
        let short = head.shorthand().expect("No branch name found");
        return Some(short.into());
    }

    None
}

fn is_release_branch(current: &str, release: &str) -> bool {
    if let Ok(pr) = env::var("TRAVIS_PULL_REQUEST") {
        if pr != "false" {
            return false;
        }
    }

    current == release
}

fn push_to_github(config: &config::Config, tag_name: &str) {
    logger::stdout("Pushing new commit and tag");
    git::push(&config, &tag_name)
        .unwrap_or_else(|err| print_exit!("Failed to push git: {:?}", err));

    logger::stdout("Waiting a tiny bit, so GitHub can store the git tag");
    thread::sleep(Duration::from_secs(1));
}

fn release_on_github(config: &config::Config, tag_message: &str, tag_name: &str) {
    if github::can_release(&config) {
        logger::stdout("Creating GitHub release");
        github::release(&config, &tag_name, &tag_message)
            .unwrap_or_else(|err| print_exit!("Failed to create GitHub release: {:?}", err));
    } else {
        logger::stdout("Project not hosted on GitHub. Skipping release step");
    }
}

fn release_on_cratesio(config: &config::Config) {
    logger::stdout("Publishing crate on crates.io");
    if !cargo::publish(&config.repository_path, &config.cargo_token.as_ref().unwrap()) {
        print_exit!("Failed to publish on crates.io");
    }
}

fn generate_changelog(repository_path: &str, version: &Version, new_version: &String) -> String {
    logger::stdout(format!("New version would be: {}", new_version));
    logger::stdout("Would write the following Changelog:");
    match changelog::generate(repository_path, &version.to_string(), new_version) {
        Ok(log) => log,
        Err(err) => {
            logger::stderr(format!("Generating Changelog failed: {:?}", err));
            process::exit(1)
        }
    }
}

fn write_changelog(repository_path: &str, version: &Version, new_version: &str) {
    logger::stdout("Writing Changelog");
    changelog::write(repository_path, &version.to_string(), &new_version)
        .unwrap_or_else(|err| print_exit!("Writing Changelog failed: {:?}", err));
}

fn print_changelog(changelog: &str) {
    logger::stdout("====================================");
    logger::stdout(changelog);
    logger::stdout("====================================");
    logger::stdout("Would create annotated git tag");
}

fn package_crate(config: &config::Config, repository_path: &str, new_version: &str) {
    if config.release_mode {
        logger::stdout("Updating lockfile");
        if !cargo::update_lockfile(repository_path) {
            print_exit!("`cargo fetch` failed. See above for the cargo error message.");
        }
    }

    git::commit_files(&config, &new_version)
        .unwrap_or_else(|err| print_exit!("Committing files failed: {:?}", err));

    logger::stdout("Package crate");
    if !cargo::package(repository_path) {
        print_exit!("`cargo package` failed. See above for the cargo error message.");
    }
}

fn get_repo(repository_path: &str) -> git2::Repository {
    match git2::Repository::open(repository_path) {
        Ok(repo) => repo,
        Err(e) => {
            logger::stderr(format!("Could not open the git repository: {:?}", e));
            process::exit(1);
        }
    }
}

fn get_repository_path(args: &Args) -> String {
    let path = Path::new(&args.flag_path);
    let path = fs::canonicalize(path)
        .unwrap_or_else(|_| print_exit!("Path does not exist or a component is
                                                            not a directory"));
    let repo_path = path.to_str().unwrap_or_else(|| print_exit!("Path is not valid unicode"));
    repo_path.to_string()
}

fn get_signature<'a>(repository_path: String) -> git2::Signature<'a> {
    let repo = get_repo(&repository_path);
    let signature = match git::get_signature(&repo) {
        Ok(sig) => sig,
            Err(e) => {
                logger::stderr(format!("Failed to get the committer's name and email address: {}", e.description()));
                logger::stderr(COMMITTER_ERROR_MESSAGE);
                process::exit(1);
            }
    };

    signature.to_owned()
}

fn get_user_and_repo(repository_path: &str) -> (Option<String>, Option<String>) {
    let repo = get_repo(repository_path);
    let remote_or_none = repo.find_remote("origin");
    match remote_or_none {
        Ok(remote) => {
            let url = remote.url().expect("Remote URL is not valid UTF-8").to_owned();
            let (user, repo_name) = user_repo_from_url(&url)
                .unwrap_or_else(|e| print_exit!("Could not extract user and repository name from URL: {:?}", e));

            (Some(user), Some(repo_name))
        },
        Err(err) => {
            logger::warn(format!("Could not determine the origin remote url: {:?}", err));
            logger::warn("semantic-rs can't push changes or create a release on GitHub");
            (None, None)
        }
    }
}

fn get_github_token(repository_path: &str) -> Option<String> {
    let repo = get_repo(repository_path);
    let remote_or_none = repo.find_remote("origin");
    match remote_or_none {
        Ok(remote) => {
            let url = remote.url().expect("Remote URL is not valid UTF-8").to_owned();
            if github::is_github_url(&url) {
                match env::var("GH_TOKEN") {
                    Ok(token) => Some(token),
                    Err(err) => {
                        logger::warn(format!("GH_TOKEN not set: {:?}", err));
                        None
                    }
                }
            } else {
                None
            }
        },
        Err(_) => None
    }
}

fn get_cargo_token() -> String {
    let cargo_token = env::var("CARGO_TOKEN")
        .unwrap_or_else(|err| print_exit!("CARGO_TOKEN not set: {:?}", err));
    cargo_token
}

fn main() {
    env_logger::init().expect("Can't instantiate env logger");
    println!("semantic.rs 🚀");

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let mut config_builder = ConfigBuilder::new();

    if args.flag_version {
        println!("semantic.rs 🚀 -- v{}", VERSION);
        process::exit(0);
    }

    // If write mode is requested OR denied,
    // adhere to the user's wish,
    // otherwise we decide based on whether we are running in CI.
    let write_mode = match args.flag_write {
        None => ci_env_set(),
        Some(ref flag) => string_to_bool(flag)
    };

    // We can only release, if we are allowed to write
    let release_mode = write_mode && string_to_bool(&args.flag_release);
    let repository_path = get_repository_path(&args);

    config_builder.write(write_mode);
    config_builder.release(release_mode);
    config_builder.branch(args.flag_branch.clone());
    config_builder.repository_path(repository_path.clone());
    config_builder.signature(get_signature(repository_path.clone()));
    let (user, repo) = get_user_and_repo(&repository_path);
    if user.is_some() {
        config_builder.user(user.unwrap());
    }
    if repo.is_some() {
        config_builder.user(repo.unwrap());
    }
    let gh_token = get_github_token(&repository_path);
    if gh_token.is_some() {
        config_builder.gh_token(gh_token.unwrap());
    }

    config_builder.cargo_token(get_cargo_token());
}
<|endoftext|>


<commit_before>#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]
extern crate serde_json;
extern crate itertools;
extern crate clap;
extern crate regex;
extern crate byteorder;
extern crate rustc_serialize;

use clap::{App, Arg};
use itertools::{Itertools, EitherOrBoth};

mod config;
use config::Config;
mod parse;
use parse::{Parser, OutputFormat};

use std::io::BufReader;
use std::fs::File;

fn main() {
    let matches = App::new("uarc-asm")
        .version("0.1.0")
        .author("Geordon Worley <vadixidav@gmail.com>")
        .about("A configurable, generic assembler")
        .arg(Arg::with_name("format")
            .long("format")
            .short("f")
            .takes_value(true)
            .possible_values(&["little-endian", "big-endian", "hex-list"])
            .default_value("little-endian")
            .help("The format of the output files"))
        .arg(Arg::with_name("config")
            .long("config")
            .short("c")
            .help("Input JSON configuration file")
            .takes_value(true)
            .required(true))
        .arg(Arg::with_name("outputs")
            .long("outputs")
            .short("o")
            .multiple(true)
            .number_of_values(1)
            .help("Output file names in order (one at a time)"))
        .arg(Arg::with_name("inputs")
            .index(1)
            .multiple(true)
            .help("List of the input assembly files in the order they are parsed"))
        .get_matches();

    let format = match matches.value_of("format").unwrap() {
        "little-endian" => OutputFormat::LittleEndian,
        "big-endian" => OutputFormat::BigEndian,
        "hex-list" => OutputFormat::HexList,
        v => panic!("Error: \"{}\" is not a valid format.", v),
    };

    let config_filename = matches.value_of("config").unwrap();
    let config = Config::new_from_filename(config_filename);

    let mut parser = Parser::new(&config);

    for name in matches.values_of("inputs")
        .map(|iter| iter.collect())
        .unwrap_or_else(|| Vec::new()) {
        parser.parse(BufReader::new(File::open(&name)
            .unwrap_or_else(|e| panic!("Error: Failed to open input file \"{}\": {}", name, e))));
    }

    // Link the program.
    parser.link();

    let outputs =
        matches.values_of("outputs").map(|iter| iter.collect()).unwrap_or_else(|| Vec::new());
    for (i, name) in (0..config.segment_widths.len())
        .zip_longest(outputs)
        .map(|v| {
            match v {
                EitherOrBoth::Both(_, specified) => specified.into(),
                EitherOrBoth::Left(n) => format!("oseg{}", n),
                EitherOrBoth::Right(specified) => {
                    panic!("Error: Output \"{}\" goes past the amount of segments for this \
                            architecture.",
                           specified);
                }
            }
        })
        .enumerate() {
        parser.output(format,
                      i,
                      &mut File::create(&name).unwrap_or_else(|e| {
                          panic!("Error: Failed to open output file \"{}\": {}", name, e)
                      }));
    }
}
<commit_msg>fix: couldnt specify output file names<commit_after>#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]
extern crate serde_json;
extern crate itertools;
extern crate clap;
extern crate regex;
extern crate byteorder;
extern crate rustc_serialize;

use clap::{App, Arg};
use itertools::{Itertools, EitherOrBoth};

mod config;
use config::Config;
mod parse;
use parse::{Parser, OutputFormat};

use std::io::BufReader;
use std::fs::File;

fn main() {
    let matches = App::new("uarc-asm")
        .version("0.1.0")
        .author("Geordon Worley <vadixidav@gmail.com>")
        .about("A configurable, generic assembler")
        .arg(Arg::with_name("format")
            .long("format")
            .short("f")
            .takes_value(true)
            .possible_values(&["little-endian", "big-endian", "hex-list"])
            .default_value("little-endian")
            .help("The format of the output files"))
        .arg(Arg::with_name("config")
            .long("config")
            .short("c")
            .help("Input JSON configuration file")
            .takes_value(true)
            .required(true))
        .arg(Arg::with_name("outputs")
            .long("outputs")
            .short("o")
            .multiple(true)
            .number_of_values(1)
            .takes_value(true)
            .help("Output file names in order (one at a time)"))
        .arg(Arg::with_name("inputs")
            .index(1)
            .multiple(true)
            .help("List of the input assembly files in the order they are parsed"))
        .get_matches();

    let format = match matches.value_of("format").unwrap() {
        "little-endian" => OutputFormat::LittleEndian,
        "big-endian" => OutputFormat::BigEndian,
        "hex-list" => OutputFormat::HexList,
        v => panic!("Error: \"{}\" is not a valid format.", v),
    };

    let config_filename = matches.value_of("config").unwrap();
    let config = Config::new_from_filename(config_filename);

    let mut parser = Parser::new(&config);

    for name in matches.values_of("inputs")
        .map(|iter| iter.collect())
        .unwrap_or_else(|| Vec::new()) {
        parser.parse(BufReader::new(File::open(&name)
            .unwrap_or_else(|e| panic!("Error: Failed to open input file \"{}\": {}", name, e))));
    }

    // Link the program.
    parser.link();

    let outputs =
        matches.values_of("outputs").map(|iter| iter.collect()).unwrap_or_else(|| Vec::new());
    for (i, name) in (0..config.segment_widths.len())
        .zip_longest(outputs)
        .map(|v| {
            match v {
                EitherOrBoth::Both(_, specified) => specified.into(),
                EitherOrBoth::Left(n) => format!("oseg{}", n),
                EitherOrBoth::Right(specified) => {
                    panic!("Error: Output \"{}\" goes past the amount of segments for this \
                            architecture.",
                           specified);
                }
            }
        })
        .enumerate() {
        parser.output(format,
                      i,
                      &mut File::create(&name).unwrap_or_else(|e| {
                          panic!("Error: Failed to open output file \"{}\": {}", name, e)
                      }));
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;

use std::collections::HashMap;
use std::io::Write;
use byteorder::{LittleEndian, WriteBytesExt};

#[derive(Hash, Eq, PartialEq, Clone, Debug)]
struct Frac(u64, u64);

type Memory = HashMap<Frac, f64>;

fn simplify(Frac(a, b): Frac) -> Frac {
    fn gcd(x: u64, y: u64) -> u64 {
        if y == 0 {
            x
        } else {
            gcd(y, x % y)
        }
    }

    let d = gcd(a, b);
    Frac(a/d, b/d)
}

/// judge a set of notes based on harmony.
/// range: floats in [0, 1] and lower is better.
fn judge_harmony(noteset: &[Frac], memory: &Memory) -> f64 {
    let mut harmony_sum = 0_f64;

    for &Frac(a1, b1) in noteset {
        for (&Frac(a2, b2), &familiarity) in memory.iter() {
            let Frac(a3, b3) = simplify(Frac(a1*b2, a2*b1));
            harmony_sum += familiarity * (a3 as f64) * (b3 as f64);
        }
    }
    let iterations = noteset.len()*memory.len();
    let avg_harmony = (harmony_sum as f64)/(iterations as f64);

    (1_f64 - 1_f64/(avg_harmony/5_f64).exp()).max(0_f64).min(1_f64)
}

/// judge a set of notes based on familiarity & novelty balance.
/// range: floats in [0, 1] and lower is better.
fn judge_novelty(noteset: &[Frac], memory: &Memory) -> f64 {
    if noteset.len() < 1 {
        panic!("judge_novelty: need at least 1 note");
    }

    let mut familiarity_sum = 0_f64;
    for note in noteset {
        let &familiarity = memory.get(note).unwrap_or(&0_f64);
        familiarity_sum += familiarity;
    }

    let avg_familiarity = familiarity_sum / (noteset.len() as f64);
    let target_familiarity = 0.1_f64;
    let disparity = (target_familiarity - avg_familiarity).abs();

    (1_f64 - 1_f64/disparity.exp()).max(0_f64).min(1_f64)
}

/// judge a set of notes.
/// range: floats in [0, 1] and lower is better.
fn judge(noteset: &[Frac], memory: &Memory) -> f64 {
    (judge_harmony(noteset, memory) + judge_novelty(noteset, memory))/2_f64
}

fn forget(memory: &mut Memory) {
    for (n, val) in memory.iter_mut() {
        *val *= 0.75;
    }
}

fn remember(note_set: &[Frac], memory: &mut Memory) {
    let increase = 0.1_f64;
    for note in note_set {
        let val = match memory.get(note) {
            Some(v) => v + increase,
            None => increase,
        };
        memory.insert(note.clone(), val);
    }
}

/// step to a set of notes that minimizes the judge function.
fn step_notes(note_set: &[Frac], memory: &Memory) -> Vec<Frac> {
    let mut best: Vec<Frac> = note_set.to_owned();
    let mut best_score = 1_f64;
    for i in 0..note_set.len() {
        for a in 1..12 {
            for b in 1..12 {
                let possibility = simplify(Frac(a, b));
                if (note_set.contains(&possibility)) {
                    continue;
                }
                let note_set2: Vec<Frac> = note_set[0..i].iter()
                                                         .chain(note_set[i+1..note_set.len()].iter())
                                                         .chain([possibility].iter())
                                                         .map(|n| n.clone())
                                                         .collect();
                let score = judge(&note_set2, memory);
                if score < best_score {
                    best = note_set2;
                    best_score = score;
                }
            }
        }
    }

    best
}

fn main_loop() {
    let mut notes = vec![Frac(1, 3), Frac(1, 2), Frac(1, 1)];
    let mut memory = Memory::new();

    loop {
        //println!("notes: {:?}", notes);
        let harmony = judge_harmony(&notes);
        //println!("{}", harmony);
        let novelty = judge_novelty(&notes, &memory);
        //println!("{}", novelty);
        let mut input = String::new();
        //std::io::stdin().read_line(&mut input).unwrap();

        forget(&mut memory);
        notes = step_notes(&notes, &memory);
        remember(&notes, &mut memory);
    }
}

type PCM_Sample = i16;
static PCM_HZ: u64 = 44100_u64;
static STEPS_PER_SEC: u64 = 4;
static BASE_NOTE: f64 = 250_f64;
type Endianness = LittleEndian;

fn sine_wave(freq: f64, step: u64) -> f64 {
    (2.0*std::f64::consts::PI*(step as f64)*freq/(PCM_HZ as f64)).sin()
}

fn sine_waves(base_note: f64, fractions: &[Frac], step: u64) -> f64 {
    let mut sum = 0_f64;
    for &Frac(a, b) in fractions {
        let freq = (base_note / (b as f64)) * (a as f64);
        sum += sine_wave(freq, step);
    }

    sum / (fractions.len() as f64)
}

fn linear_envelope(sample: f64, duration: u64, progress: u64) -> f64 {
    sample * (progress as f64) / (duration as f64)
}

fn output_pcm() {
    let mut notes = vec![Frac(1, 2), Frac(1, 1), Frac(1, 3), Frac(1, 5), Frac(1, 7)];
    let mut memory = Memory::new();

    let mut j=0;
    for i in (0_u64..u64::max_value()).cycle() {
        if notes.len() > 10 {
            panic!(format!("{:?}", notes));
        }
        let sample = sine_waves(BASE_NOTE, &notes, i) *
                     (PCM_Sample::max_value() as f64);

        let enveloped = linear_envelope(sample, j, PCM_HZ/STEPS_PER_SEC);

        let bounded = enveloped.min(PCM_Sample::max_value() as f64 - 1_f64)
                               .max(PCM_Sample::min_value() as f64 + 1_f64);

        let as_sample: PCM_Sample = bounded as PCM_Sample;
        std::io::stdout().write_i16::<Endianness>(as_sample).unwrap();

        j += 1;
        if j == PCM_HZ/STEPS_PER_SEC {
            j = 0;
            forget(&mut memory);
            notes = step_notes(&notes, &memory);
            remember(&notes, &mut memory);
        }
    }
}

fn main() {
    output_pcm();
}
<commit_msg>Oops, that code wasn't supposed to be there.<commit_after>extern crate byteorder;

use std::collections::HashMap;
use std::io::Write;
use byteorder::{LittleEndian, WriteBytesExt};

#[derive(Hash, Eq, PartialEq, Clone, Debug)]
struct Frac(u64, u64);

type Memory = HashMap<Frac, f64>;

fn simplify(Frac(a, b): Frac) -> Frac {
    fn gcd(x: u64, y: u64) -> u64 {
        if y == 0 {
            x
        } else {
            gcd(y, x % y)
        }
    }

    let d = gcd(a, b);
    Frac(a/d, b/d)
}

/// judge a set of notes based on harmony.
/// range: floats in [0, 1] and lower is better.
fn judge_harmony(noteset: &[Frac], memory: &Memory) -> f64 {
    let mut harmony_sum = 0_f64;

    for &Frac(a1, b1) in noteset {
        for (&Frac(a2, b2), &familiarity) in memory.iter() {
            let Frac(a3, b3) = simplify(Frac(a1*b2, a2*b1));
            harmony_sum += familiarity * (a3 as f64) * (b3 as f64);
        }
    }
    let iterations = noteset.len()*memory.len();
    let avg_harmony = (harmony_sum as f64)/(iterations as f64);

    (1_f64 - 1_f64/(avg_harmony/5_f64).exp()).max(0_f64).min(1_f64)
}

/// judge a set of notes based on familiarity & novelty balance.
/// range: floats in [0, 1] and lower is better.
fn judge_novelty(noteset: &[Frac], memory: &Memory) -> f64 {
    if noteset.len() < 1 {
        panic!("judge_novelty: need at least 1 note");
    }

    let mut familiarity_sum = 0_f64;
    for note in noteset {
        let &familiarity = memory.get(note).unwrap_or(&0_f64);
        familiarity_sum += familiarity;
    }

    let avg_familiarity = familiarity_sum / (noteset.len() as f64);
    let target_familiarity = 0.1_f64;
    let disparity = (target_familiarity - avg_familiarity).abs();

    (1_f64 - 1_f64/disparity.exp()).max(0_f64).min(1_f64)
}

/// judge a set of notes.
/// range: floats in [0, 1] and lower is better.
fn judge(noteset: &[Frac], memory: &Memory) -> f64 {
    (judge_harmony(noteset, memory) + judge_novelty(noteset, memory))/2_f64
}

fn forget(memory: &mut Memory) {
    for (n, val) in memory.iter_mut() {
        *val *= 0.75;
    }
}

fn remember(note_set: &[Frac], memory: &mut Memory) {
    let increase = 0.1_f64;
    for note in note_set {
        let val = match memory.get(note) {
            Some(v) => v + increase,
            None => increase,
        };
        memory.insert(note.clone(), val);
    }
}

/// step to a set of notes that minimizes the judge function.
fn step_notes(note_set: &[Frac], memory: &Memory) -> Vec<Frac> {
    let mut best: Vec<Frac> = note_set.to_owned();
    let mut best_score = 1_f64;
    for i in 0..note_set.len() {
        for a in 1..12 {
            for b in 1..12 {
                let possibility = simplify(Frac(a, b));
                if (note_set.contains(&possibility)) {
                    continue;
                }
                let note_set2: Vec<Frac> = note_set[0..i].iter()
                                                         .chain(note_set[i+1..note_set.len()].iter())
                                                         .chain([possibility].iter())
                                                         .map(|n| n.clone())
                                                         .collect();
                let score = judge(&note_set2, memory);
                if score < best_score {
                    best = note_set2;
                    best_score = score;
                }
            }
        }
    }

    best
}

type PCM_Sample = i16;
static PCM_HZ: u64 = 44100_u64;
static STEPS_PER_SEC: u64 = 4;
static BASE_NOTE: f64 = 250_f64;
type Endianness = LittleEndian;

fn sine_wave(freq: f64, step: u64) -> f64 {
    (2.0*std::f64::consts::PI*(step as f64)*freq/(PCM_HZ as f64)).sin()
}

fn sine_waves(base_note: f64, fractions: &[Frac], step: u64) -> f64 {
    let mut sum = 0_f64;
    for &Frac(a, b) in fractions {
        let freq = (base_note / (b as f64)) * (a as f64);
        sum += sine_wave(freq, step);
    }

    sum / (fractions.len() as f64)
}

fn linear_envelope(sample: f64, duration: u64, progress: u64) -> f64 {
    sample * (progress as f64) / (duration as f64)
}

fn output_pcm() {
    let mut notes = vec![Frac(1, 2), Frac(1, 1), Frac(1, 3), Frac(1, 5), Frac(1, 7)];
    let mut memory = Memory::new();

    let mut j=0;
    for i in (0_u64..u64::max_value()).cycle() {
        if notes.len() > 10 {
            panic!(format!("{:?}", notes));
        }
        let sample = sine_waves(BASE_NOTE, &notes, i) *
                     (PCM_Sample::max_value() as f64);

        let enveloped = linear_envelope(sample, j, PCM_HZ/STEPS_PER_SEC);

        let bounded = enveloped.min(PCM_Sample::max_value() as f64 - 1_f64)
                               .max(PCM_Sample::min_value() as f64 + 1_f64);

        let as_sample: PCM_Sample = bounded as PCM_Sample;
        std::io::stdout().write_i16::<Endianness>(as_sample).unwrap();

        j += 1;
        if j == PCM_HZ/STEPS_PER_SEC {
            j = 0;
            forget(&mut memory);
            notes = step_notes(&notes, &memory);
            remember(&notes, &mut memory);
        }
    }
}

fn main() {
    output_pcm();
}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]

extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
mod energy;
use energy::*;
use std::io::prelude::*;
extern crate argparse;
use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
mod trajectory;
use trajectory::*;

const LJ_EPS : f64 = 1.0;
const LJ_SIG : f64 = 1.0;

const TRIES_INTENDED : f64 = 3.0;
const DISP_SCALE_FACTOR : f64 = 0.1;

// easy printing to stderr
macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize,
                  NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64,
                  CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool,
                  OUTPUT_PREFIX: &mut String, OUTPUT_INTERVAL: &mut i64, OUTPUT_MINIM: &mut bool,
                  VACUUM_SLAB: &mut f64) {
    let mut ap = ArgumentParser::new();
    ap.set_description("LJ MC simulation.");
    ap.refer(NUM_STEPS)
        .add_option(&["-n", "--nsteps"], Store,
                    "Simulation steps: Number of steps for averaging" );
    ap.refer(NUM_MINIM_STEPS)
        .add_option(&["-m", "--nminimsteps"], Store,
                    "Minimization steps: Number of steps before averaging starts");
    ap.refer(NUM_PARTICLES)
        .add_option(&["-p", "--nparticles"], Store,
                    "Total number of particles");
    ap.refer(DENSITY)
        .add_option(&["-d", "--density"], Store,
                    "Particle density");
    ap.refer(TEMPERATURE)
        .add_option(&["-t", "--temperature"], Store,
                    "Temperature");
    ap.refer(CUTOFF)
        .add_option(&["--cutoff"], Store,
                    "Lennard jones cutoff radius in length of epsilon");
    ap.refer(MAX_DISP_START)
        .add_option(&["--displacement"], Store,
                    "Displacement per trial move");
    ap.refer(SCALE)
        .add_option(&["--nodisplacementscale"], StoreFalse,
                    "Disable displacement scaling");
    ap.refer(OUTPUT_PREFIX)
        .add_option(&["-o", "--output"], Store,
                    "Output file prefix");
    ap.refer(OUTPUT_INTERVAL)
        .add_option(&["--osteps"], Store,
                    "Number of steps between writing to the trajectory file. -1 only writes last frame");
    ap.refer(OUTPUT_MINIM)
        .add_option(&["--writeminimization"], StoreTrue,
                    "Enables writing of minimization step to trajectory");
    ap.refer(VACUUM_SLAB)
        .add_option(&["--vacuum"], Store,
                    "Dimension of vacuum space above the intial system relative to the rest of the system (0=no slab, 1=half filled system, 2=thrid filled system ...).");
    ap.refer(TAILCORR)
        .add_option(&["--notailcorr"], StoreFalse,
                    "Disable tailcorrection");
    ap.refer(SHIFT)
        .add_option(&["--noshift"], StoreFalse,
                    "Disable lj shifting");
    ap.parse_args_or_exit();
}

fn main() {

    // define all the stuff
    let mut minim_steps  = 1000000;
    let mut sample_steps = 100000;

    let mut num_particles: usize = 512;
    let mut density = 0.7;
    let mut temperature = 0.9;

    let mut cutoff = 3.0;
    let mut displacement = 0.1;

    let mut  TAILCORR : bool = true;
    let mut SHIFT: bool = true;
    let mut SCALE: bool = true;

    let mut vacuum_slab = 0.0;

    let mut output_prefix = "montecarlo".to_string();
    let mut output_interval : i64 = 100;
    let mut output_minim : bool = false;
    parse_cmd_args(&mut sample_steps, &mut minim_steps, &mut num_particles,
                   &mut density, &mut temperature,
                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT,
                   &mut output_prefix, &mut output_interval, &mut output_minim,
                   &mut vacuum_slab);

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##################  LJ Monte Carlo Simulation  #################");
    println_stderr!("################################################################");
    println_stderr!("");

    // initialize stuff
    let beta = 1.0/temperature;
    let mut volume = (num_particles as f64)/ density;
    let length  = volume.cbrt();
    let (l_x, l_y, mut l_z) = (length, length, length);
    if vacuum_slab > 0.0 { // increase space in z
        let scale = vacuum_slab + 1.0;
        l_z *= scale;
        volume *= scale;
        density /= scale;
    }

    let cutoff_squared = cutoff * cutoff;
    let max_displacement = length / 2.0;

    let mut rng = rand::thread_rng();
    let particle_range = Range::new(0, num_particles-1);

    let mut rx : Vec<f64> = vec![];
    let mut ry : Vec<f64> = vec![];
    let mut rz : Vec<f64> = vec![];
    loop {
        rx.push(l_x * rng.gen::<f64>());
        ry.push(l_y * rng.gen::<f64>());
        rz.push(l_z * rng.gen::<f64>());
        if rx.len() == num_particles { break; }
    }

    let e_shift = if SHIFT { 4.0 * LJ_EPS * ( (LJ_SIG/cutoff).powi(12) - (LJ_SIG/cutoff).powi(6) ) } else { 0.0 };
    let e_corr = if TAILCORR { 8.0/3.0*std::f64::consts::PI*density*LJ_EPS*LJ_SIG.powi(3)*((1.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
    let p_corr = if TAILCORR { 16.0/3.0*std::f64::consts::PI*density.powi(2)*LJ_EPS*LJ_SIG.powi(3)*((2.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };

    println_stderr!("Particles: {}, Density: {}, Temperature: {}", num_particles, density, temperature);
    println_stderr!("System volume: {:8.3}, Dimensions {:.3}/{:.3}/{:.3}", volume, l_x, l_y, l_z);
    println_stderr!("Minimization steps: {}, Sampling steps: {}", minim_steps, sample_steps);
    println_stderr!("LJ params eps: {}, sigma: {}, cutoff: {}", LJ_EPS, LJ_SIG, cutoff);
    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, e_shift, p_corr);

    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
    let mut energy_sum = 0.0;
    let mut virial_sum = 0.0;
    let mut step_counter = 0;
    let mut accept_counter = 0;

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("#####################  Energy Minimization  ####################");
    println_stderr!("################################################################");
    println_stderr!("");

    // prepare and write first trajectory frame
    let mut trajectory : XYZTrajectory = XYZTrajectory::new(&format!("{}.xyz", output_prefix));
    if output_minim { trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true); }

    for step in 0..minim_steps+sample_steps {

        // select rnd particle
        let rnd_index = particle_range.ind_sample(&mut rng);

        // store old position
        let oldX = rx[rnd_index];
        let oldY = ry[rnd_index];
        let oldZ = rz[rnd_index];

        // old particle energy
        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);

        // rnd displacement and PBC
        rx[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        ry[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        rz[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        if rx[rnd_index] < 0.0 { rx[rnd_index] += l_x }
        if rx[rnd_index] > l_x { rx[rnd_index] -= l_x }
        if ry[rnd_index] < 0.0 { ry[rnd_index] += l_y }
        if ry[rnd_index] > l_y { ry[rnd_index] -= l_y }
        if rz[rnd_index] < 0.0 { rz[rnd_index] += l_z }
        if rz[rnd_index] > l_z { rz[rnd_index] -= l_z }

        // calculate energy difference
        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
        let dE = new_particle_energy - old_particle_energy;

        //accept move
        if rng.gen::<f64>() < (-beta * dE).exp() {
            accept_counter += 1;
            if step % 1000 == 0 { // calculate total energy every 1000 steps to account for rounding errors
                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
                energy = e;
                virial = v;
            } else {
                energy += dE;
                virial += new_particle_virial - old_particle_virial;
            }
        } else { // or restore old position
            rx[rnd_index] = oldX;
            ry[rnd_index] = oldY;
            rz[rnd_index] = oldZ;
        }

        // update sums for averaging
        step_counter += 1;
        energy_sum += energy;
        virial_sum += virial;

        // print some output during minimization
        if step < minim_steps && step_counter % 5000 == 0 && step != 0 {
            let tries_per_step : f64 = step_counter as f64 /accept_counter as f64;
            let acceptance_rate = 1.0/tries_per_step * 100.0;
            println_stderr!("Minim {:<10} Energy: {:<30.3} Accept.: {:<4.1}%   dr: {:.3}", step+1, energy, acceptance_rate, displacement);

            if SCALE {
                let scale_factor = (TRIES_INTENDED/tries_per_step * DISP_SCALE_FACTOR).abs();
                if tries_per_step < TRIES_INTENDED - 0.2 && displacement < max_displacement {
                    displacement += displacement * scale_factor;
                } else if tries_per_step > TRIES_INTENDED + 0.2 && displacement > 0.0 {
                    displacement -= displacement * scale_factor;
                }
                step_counter = 0;
                accept_counter = 0;
            }
        }

        // reset sums for sampling
        if step == minim_steps-1 {
            println_stderr!("");
            println_stderr!("################################################################");
            println_stderr!("##########################  Sampling  ##########################");
            println_stderr!("################################################################");
            println_stderr!("");
            step_counter = 0;
            accept_counter = 0;
            energy_sum = 0.0;
            virial_sum = 0.0;
        }

        if step > minim_steps && step_counter % 5000 == 0 {
            println_stderr!("Step  {:<10} Energy: {:<30.3}", step_counter, energy);
        }

        // write trajectory maybe
        if step as i64 % output_interval == 0 {
            if step > minim_steps || output_minim {
                trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
            }
        }

    }

    let final_energy = energy_sum/step_counter as f64;
    let particle_energy = final_energy / num_particles as f64;
    let final_virial = virial_sum / 3.0 / step_counter as f64 / volume;
    let pressure = virial_sum / 3.0 / step_counter as f64 / volume + density * temperature + p_corr;
    let final_acceptance_rate = 1.0/((accept_counter as f64)/(step_counter as f64)) * 100.0;

    println_stderr!("Done sampling!");
    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##########################  Results  ###########################");
    println_stderr!("################################################################");
    println_stderr!("");
    println!(
"Minimization: {}
Steps: {}

# Lennard Jones Params
epsilon: {}
sigma: {}
cutoff: {}

# System
Particles: {}
Density: {}
Temperature: {}
Volume: {}
Box dimension: {:.3}/{:.3}/{:.3}
Max Displacement: {}

# Correction
Energy correction: {}
Shift: {}
P-Correction: {}

# Averages
Tries: {}
Accepted: {}
Acceptance: {:.2}%
Energy: {}
Energy per particle: {}
Virial: {}
Pressure: {}",
         minim_steps, sample_steps,
        LJ_EPS, LJ_SIG, cutoff,
        num_particles, density, temperature, volume, l_x, l_y, l_z, displacement,
        e_corr, e_shift, p_corr,
        step_counter, accept_counter, final_acceptance_rate, final_energy, particle_energy, final_virial, pressure);

    trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
}
<commit_msg>fixed vacuum space initial conditions<commit_after>#![allow(non_snake_case)]

extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
mod energy;
use energy::*;
use std::io::prelude::*;
extern crate argparse;
use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
mod trajectory;
use trajectory::*;

const LJ_EPS : f64 = 1.0;
const LJ_SIG : f64 = 1.0;

const TRIES_INTENDED : f64 = 3.0;
const DISP_SCALE_FACTOR : f64 = 0.1;

// easy printing to stderr
macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize,
                  NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64,
                  CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool,
                  OUTPUT_PREFIX: &mut String, OUTPUT_INTERVAL: &mut i64, OUTPUT_MINIM: &mut bool,
                  VACUUM_SLAB: &mut f64) {
    let mut ap = ArgumentParser::new();
    ap.set_description("LJ MC simulation.");
    ap.refer(NUM_STEPS)
        .add_option(&["-n", "--nsteps"], Store,
                    "Simulation steps: Number of steps for averaging" );
    ap.refer(NUM_MINIM_STEPS)
        .add_option(&["-m", "--nminimsteps"], Store,
                    "Minimization steps: Number of steps before averaging starts");
    ap.refer(NUM_PARTICLES)
        .add_option(&["-p", "--nparticles"], Store,
                    "Total number of particles");
    ap.refer(DENSITY)
        .add_option(&["-d", "--density"], Store,
                    "Particle density");
    ap.refer(TEMPERATURE)
        .add_option(&["-t", "--temperature"], Store,
                    "Temperature");
    ap.refer(CUTOFF)
        .add_option(&["--cutoff"], Store,
                    "Lennard jones cutoff radius in length of epsilon");
    ap.refer(MAX_DISP_START)
        .add_option(&["--displacement"], Store,
                    "Displacement per trial move");
    ap.refer(SCALE)
        .add_option(&["--nodisplacementscale"], StoreFalse,
                    "Disable displacement scaling");
    ap.refer(OUTPUT_PREFIX)
        .add_option(&["-o", "--output"], Store,
                    "Output file prefix");
    ap.refer(OUTPUT_INTERVAL)
        .add_option(&["--osteps"], Store,
                    "Number of steps between writing to the trajectory file. -1 only writes last frame");
    ap.refer(OUTPUT_MINIM)
        .add_option(&["--writeminimization"], StoreTrue,
                    "Enables writing of minimization step to trajectory");
    ap.refer(VACUUM_SLAB)
        .add_option(&["--vacuum"], Store,
                    "Dimension of vacuum space above the intial system relative to the rest of the system (0=no slab, 1=half filled system, 2=thrid filled system ...).");
    ap.refer(TAILCORR)
        .add_option(&["--notailcorr"], StoreFalse,
                    "Disable tailcorrection");
    ap.refer(SHIFT)
        .add_option(&["--noshift"], StoreFalse,
                    "Disable lj shifting");
    ap.parse_args_or_exit();
}

fn main() {

    // define all the stuff
    let mut minim_steps  = 1000000;
    let mut sample_steps = 100000;

    let mut num_particles: usize = 512;
    let mut density = 0.7;
    let mut temperature = 0.9;

    let mut cutoff = 3.0;
    let mut displacement = 0.1;

    let mut  TAILCORR : bool = true;
    let mut SHIFT: bool = true;
    let mut SCALE: bool = true;

    let mut vacuum_slab = 0.0;

    let mut output_prefix = "montecarlo".to_string();
    let mut output_interval : i64 = 100;
    let mut output_minim : bool = false;
    parse_cmd_args(&mut sample_steps, &mut minim_steps, &mut num_particles,
                   &mut density, &mut temperature,
                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT,
                   &mut output_prefix, &mut output_interval, &mut output_minim,
                   &mut vacuum_slab);

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##################  LJ Monte Carlo Simulation  #################");
    println_stderr!("################################################################");
    println_stderr!("");

    // initialize stuff
    let beta = 1.0/temperature;
    let mut volume = (num_particles as f64)/ density;
    let length  = volume.cbrt();
    let (l_x, l_y, mut l_z) = (length, length, length);

    let cutoff_squared = cutoff * cutoff;
    let max_displacement = length / 2.0;

    let mut rng = rand::thread_rng();
    let particle_range = Range::new(0, num_particles-1);

    let mut rx : Vec<f64> = vec![];
    let mut ry : Vec<f64> = vec![];
    let mut rz : Vec<f64> = vec![];
    loop {
        rx.push(l_x * rng.gen::<f64>());
        ry.push(l_y * rng.gen::<f64>());
        rz.push(l_z * rng.gen::<f64>());
        if rx.len() == num_particles { break; }
    }

    if vacuum_slab > 0.0 { // increase space in z
        let scale = vacuum_slab + 1.0;
        l_z *= scale;
        volume *= scale;
        density /= scale;
    }

    let e_shift = if SHIFT { 4.0 * LJ_EPS * ( (LJ_SIG/cutoff).powi(12) - (LJ_SIG/cutoff).powi(6) ) } else { 0.0 };
    let e_corr = if TAILCORR { 8.0/3.0*std::f64::consts::PI*density*LJ_EPS*LJ_SIG.powi(3)*((1.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
    let p_corr = if TAILCORR { 16.0/3.0*std::f64::consts::PI*density.powi(2)*LJ_EPS*LJ_SIG.powi(3)*((2.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };

    println_stderr!("Particles: {}, Density: {}, Temperature: {}", num_particles, density, temperature);
    println_stderr!("System volume: {:8.3}, Dimensions {:.3}/{:.3}/{:.3}", volume, l_x, l_y, l_z);
    println_stderr!("Minimization steps: {}, Sampling steps: {}", minim_steps, sample_steps);
    println_stderr!("LJ params eps: {}, sigma: {}, cutoff: {}", LJ_EPS, LJ_SIG, cutoff);
    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, e_shift, p_corr);

    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
    let mut energy_sum = 0.0;
    let mut virial_sum = 0.0;
    let mut step_counter = 0;
    let mut accept_counter = 0;

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("#####################  Energy Minimization  ####################");
    println_stderr!("################################################################");
    println_stderr!("");

    // prepare and write first trajectory frame
    let mut trajectory : XYZTrajectory = XYZTrajectory::new(&format!("{}.xyz", output_prefix));
    if output_minim { trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true); }

    for step in 0..minim_steps+sample_steps {

        // select rnd particle
        let rnd_index = particle_range.ind_sample(&mut rng);

        // store old position
        let oldX = rx[rnd_index];
        let oldY = ry[rnd_index];
        let oldZ = rz[rnd_index];

        // old particle energy
        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);

        // rnd displacement and PBC
        rx[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        ry[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        rz[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        if rx[rnd_index] < 0.0 { rx[rnd_index] += l_x }
        if rx[rnd_index] > l_x { rx[rnd_index] -= l_x }
        if ry[rnd_index] < 0.0 { ry[rnd_index] += l_y }
        if ry[rnd_index] > l_y { ry[rnd_index] -= l_y }
        if rz[rnd_index] < 0.0 { rz[rnd_index] += l_z }
        if rz[rnd_index] > l_z { rz[rnd_index] -= l_z }

        // calculate energy difference
        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
        let dE = new_particle_energy - old_particle_energy;

        //accept move
        if rng.gen::<f64>() < (-beta * dE).exp() {
            accept_counter += 1;
            if step % 1000 == 0 { // calculate total energy every 1000 steps to account for rounding errors
                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
                energy = e;
                virial = v;
            } else {
                energy += dE;
                virial += new_particle_virial - old_particle_virial;
            }
        } else { // or restore old position
            rx[rnd_index] = oldX;
            ry[rnd_index] = oldY;
            rz[rnd_index] = oldZ;
        }

        // update sums for averaging
        step_counter += 1;
        energy_sum += energy;
        virial_sum += virial;

        // print some output during minimization
        if step < minim_steps && step_counter % 5000 == 0 && step != 0 {
            let tries_per_step : f64 = step_counter as f64 /accept_counter as f64;
            let acceptance_rate = 1.0/tries_per_step * 100.0;
            println_stderr!("Minim {:<10} Energy: {:<30.3} Accept.: {:<4.1}%   dr: {:.3}", step+1, energy, acceptance_rate, displacement);

            if SCALE {
                let scale_factor = (TRIES_INTENDED/tries_per_step * DISP_SCALE_FACTOR).abs();
                if tries_per_step < TRIES_INTENDED - 0.2 && displacement < max_displacement {
                    displacement += displacement * scale_factor;
                } else if tries_per_step > TRIES_INTENDED + 0.2 && displacement > 0.0 {
                    displacement -= displacement * scale_factor;
                }
                step_counter = 0;
                accept_counter = 0;
            }
        }

        // reset sums for sampling
        if step == minim_steps-1 {
            println_stderr!("");
            println_stderr!("################################################################");
            println_stderr!("##########################  Sampling  ##########################");
            println_stderr!("################################################################");
            println_stderr!("");
            step_counter = 0;
            accept_counter = 0;
            energy_sum = 0.0;
            virial_sum = 0.0;
        }

        if step > minim_steps && step_counter % 5000 == 0 {
            println_stderr!("Step  {:<10} Energy: {:<30.3}", step_counter, energy);
        }

        // write trajectory maybe
        if step as i64 % output_interval == 0 {
            if step > minim_steps || output_minim {
                trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
            }
        }

    }

    let final_energy = energy_sum/step_counter as f64;
    let particle_energy = final_energy / num_particles as f64;
    let final_virial = virial_sum / 3.0 / step_counter as f64 / volume;
    let pressure = virial_sum / 3.0 / step_counter as f64 / volume + density * temperature + p_corr;
    let final_acceptance_rate = 1.0/((accept_counter as f64)/(step_counter as f64)) * 100.0;

    println_stderr!("Done sampling!");
    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##########################  Results  ###########################");
    println_stderr!("################################################################");
    println_stderr!("");
    println!(
"Minimization: {}
Steps: {}

# Lennard Jones Params
epsilon: {}
sigma: {}
cutoff: {}

# System
Particles: {}
Density: {}
Temperature: {}
Volume: {}
Box dimension: {:.3}/{:.3}/{:.3}
Max Displacement: {}

# Correction
Energy correction: {}
Shift: {}
P-Correction: {}

# Averages
Tries: {}
Accepted: {}
Acceptance: {:.2}%
Energy: {}
Energy per particle: {}
Virial: {}
Pressure: {}",
         minim_steps, sample_steps,
        LJ_EPS, LJ_SIG, cutoff,
        num_particles, density, temperature, volume, l_x, l_y, l_z, displacement,
        e_corr, e_shift, p_corr,
        step_counter, accept_counter, final_acceptance_rate, final_energy, particle_energy, final_virial, pressure);

    trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
}
<|endoftext|>