<commit_before>#[macro_use]
extern crate log;
#[macro_use]
extern crate nom;

extern crate bip_bencode;
extern crate bip_metainfo;
extern crate bip_utracker;
extern crate chrono;
extern crate hyper;
extern crate pnet;
extern crate pnet_macros_support;
extern crate pretty_env_logger;
extern crate rustyline;
extern crate url;

mod packet;

use bip_bencode::Bencode;
use bip_metainfo::MetainfoFile;
use bip_utracker::contact::CompactPeersV4;
use chrono::{TimeZone, UTC};
use hyper::Client;
use hyper::header::Connection;
use packet::peer_pkt::MutablePeerHandshakePacket;
use pnet::packet::Packet;
use rustyline::Editor;
use rustyline::error::ReadlineError;
use url::Url;

use std::fs::File;
use std::io::prelude::*;
use std::net::TcpStream;
use std::str;
use std::string::String;
use std::time::{SystemTime, UNIX_EPOCH};
use std::thread;

const HISTORY_FILE: &'static str = ".rustyline.history";
const PEER_HANDSHAKE_STRUCT_SZ: usize = 68;

/// Print general information about the torrent.
fn print_metainfo_overview(bytes: &[u8]) {
    let metainfo = MetainfoFile::from_bytes(bytes).unwrap();
    let info = metainfo.info();
    let info_hash_hex = metainfo.info_hash()
        .as_ref()
        .iter()
        .map(|b| format!("{:02X}", b))
        .fold(String::new(), |mut acc, nex| {
            acc.push_str(&nex);
            acc
        });
    let utc_creation_date = metainfo.creation_date().map(|c| UTC.timestamp(c, 0));

    println!("------Metainfo File Overview-----");

    println!("InfoHash: {}", info_hash_hex);
    println!("Main Tracker: {:?}", metainfo.main_tracker());
    println!("Comment: {:?}", metainfo.comment());
    println!("Creator: {:?}", metainfo.created_by());
    println!("Creation Date: {:?}", utc_creation_date);

    println!("Directory: {:?}", info.directory());
    println!("Piece Length: {:?}", info.piece_length());
    println!("Number Of Pieces: {}", info.pieces().count());
    println!("Number Of Files: {}", info.files().count());
    println!("Total File Size: {}",
             info.files().fold(0, |acc, nex| acc + nex.length()));
}

fn connect_to_tracker(metainfo_file: MetainfoFile,
                      peer_id: &str,
                      port: u16)
                      -> Result<(Vec<String>, String), String> {
    debug!("connecting to tracker: {:?}", metainfo_file.main_tracker());
    let info_hash = metainfo_file.info_hash();
    // TODO figure can this conversion to url-encoded form be done safely?
    let info_hash_str = unsafe { str::from_utf8_unchecked(info_hash.as_ref()) };
    // TODO this needs to be calculated based on what we have
    let total_len = metainfo_file.info().files().fold(0, |acc, nex| acc + nex.length());

    let mut url = Url::parse(metainfo_file.main_tracker().unwrap()).unwrap();
    url.query_pairs_mut()
        .append_pair("info_hash", info_hash_str)
        .append_pair("peer_id", peer_id)
        .append_pair("port", &(port.to_string()))
        // TODO parametrize this
        .append_pair("uploaded", "0")
        // TODO parametrize this
        .append_pair("downloaded", "0")
        // TODO see note on total_len above
        .append_pair("left", &(total_len.to_string()))
        .append_pair("compact", "1")
        .append_pair("event", "started")
        .append_pair("supportcrypto", "0");
    trace!("URL {:?}", url);

    let client = Client::new();
    let mut res = client.get(url).header(Connection::close()).send().unwrap();
    let mut buffer = Vec::new();
    res.read_to_end(&mut buffer).unwrap();
    debug!("{:?}", res);
    let bencode = Bencode::decode(&buffer).unwrap();
    trace!("{:?}", bencode);
    let (_, peers) = CompactPeersV4::from_bytes(bencode.dict()
            .unwrap()
            .lookup("peers")
            .unwrap()
            .bytes()
            .unwrap())
        .unwrap();
    trace!("{:?}", peers);
    let mut ip_ports: Vec<String> = Vec::new();
    for peer in peers.iter() {
        debug!("{:?}", peer);
        ip_ports.push(peer.to_string());
    }
    Ok((ip_ports, info_hash_str.to_string()))
}

fn peer_connections(peer_ip_ports: Vec<String>,
                    info_hash: &str,
                    peer_id: &str)
                    -> Result<(), String> {
    for peer_ip_port in peer_ip_ports {
        let peer_ip_port_cl = peer_ip_port.clone();
        let info_hash_cl = info_hash.to_string().clone();
        let peer_id_cl = peer_id.to_string().clone();
        thread::spawn(move || {
            handshake_peer(&peer_ip_port_cl, &info_hash_cl, &peer_id_cl);
        });
    }
    Ok(())
}

fn handshake_peer(peer_ip_port: &str, info_hash: &str, peer_id: &str) -> Result<(), String> {
    let mut buf = vec![0u8; PEER_HANDSHAKE_STRUCT_SZ];
    let mut ph = MutablePeerHandshakePacket::new(&mut buf).unwrap();
    ph.set_pstrlen("BitTorrent protocol".len() as u8);
    ph.set_pstr(&String::from("BitTorrent protocol").into_bytes());
    ph.set_reserved(&[0; 8]);
    ph.set_info_hash(info_hash.as_bytes());
    ph.set_peer_id(peer_id.as_bytes());
    match TcpStream::connect(peer_ip_port) {
        Ok(mut stream) => {
            debug!("Connection to peer {:?} successful!", peer_ip_port);
            match stream.write(ph.packet()) {
                Ok(_) => {
                    debug!("Sending message successful!");
                    Ok(())
                }
                Err(_) => {
                    error!("Sending message failed!");
                    Err("Sending message failed!".to_owned())
                }
            }
        }
        Err(e) => {
            error!("Connection to peer {:?} failed! {:?}", peer_ip_port, e);
            Err("Connection to peer failed!".to_owned())
        }
    }
}

fn main() {
    pretty_env_logger::init();

    let mut rl = Editor::<()>::new();
    if rl.load_history(HISTORY_FILE).is_err() {
        info!("No previous history!");
    }

    let now = SystemTime::now();
    let duration = now.duration_since(UNIX_EPOCH).unwrap();
    let client_id = "-bittorrent-rs-".to_owned() + &format!("{}", duration.as_secs() % 100_000);

    loop {
        let readline = rl.readline("> ");
        match readline {
            Ok(cmd) => {
                rl.add_history_entry(&cmd);

                let cmd = cmd.trim().split(' ').collect::<Vec<&str>>();

                match cmd[0] {
                    "help" | "h" => {
                        println!("Commands:
parse/p <torrent file path>      - show Metainfo File Overview
connect/c <torrent file path>    - initiate connection to tracker, and handshake with peers
showfiles/sf <torrent file path> - show files in the torrent
help/h                           - show this help");
                    }
                    "parse" | "p" => {
                        if cmd.len() != 2 {
                            error!("usage: parse <torrent file>");
                        } else {
                            let path = cmd[1];
                            debug!("MetainfoFile: {:?}", path);

                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();
                                    print_metainfo_overview(&bytes);
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "connect" | "c" => {
                        if cmd.len() != 2 {
                            error!("usage: connect <torrent file>");
                        } else {
                            let path = cmd[1];
                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();

                                    // TODO: generate peer ID
                                    let result =
                                        connect_to_tracker(MetainfoFile::from_bytes(&bytes)
                                                               .unwrap(),
                                                           &client_id,
                                                           6882)
                                            .unwrap();
                                    peer_connections(result.0, &result.1, &client_id).unwrap();
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "" => {}
                    _ => {
                        println!("invalid command, see \"help\"");
                    }
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                error!("Error: {:?}", err);
                break;
            }
        }
    }
    rl.save_history(HISTORY_FILE).unwrap();
}
<commit_msg>Add showfiles command<commit_after>#[macro_use]
extern crate log;
#[macro_use]
extern crate nom;

extern crate bip_bencode;
extern crate bip_metainfo;
extern crate bip_utracker;
extern crate chrono;
extern crate hyper;
extern crate pnet;
extern crate pnet_macros_support;
extern crate pretty_env_logger;
extern crate rustyline;
extern crate url;

mod packet;

use bip_bencode::Bencode;
use bip_metainfo::MetainfoFile;
use bip_utracker::contact::CompactPeersV4;
use chrono::{TimeZone, UTC};
use hyper::Client;
use hyper::header::Connection;
use packet::peer_pkt::MutablePeerHandshakePacket;
use pnet::packet::Packet;
use rustyline::Editor;
use rustyline::error::ReadlineError;
use url::Url;

use std::fs::File;
use std::io::prelude::*;
use std::net::TcpStream;
use std::str;
use std::string::String;
use std::time::{SystemTime, UNIX_EPOCH};
use std::thread;

const HISTORY_FILE: &'static str = ".rustyline.history";
const PEER_HANDSHAKE_STRUCT_SZ: usize = 68;

/// Print file list from the torrent.
fn print_files(bytes: &[u8]) {
    let metainfo = MetainfoFile::from_bytes(bytes).unwrap();
    let info = metainfo.info();

    println!("File List:");
    println!("Size (bytes)\tPath");
    println!("------------\t----------------------------------------------");
    for file in info.files() {
        println!("{:12}\t{}",
                 file.length(),
                 file.paths().next().unwrap_or("<unknown>"));
    }
}

/// Print general information about the torrent.
fn print_metainfo_overview(bytes: &[u8]) {
    let metainfo = MetainfoFile::from_bytes(bytes).unwrap();
    let info = metainfo.info();
    let info_hash_hex = metainfo.info_hash()
        .as_ref()
        .iter()
        .map(|b| format!("{:02X}", b))
        .fold(String::new(), |mut acc, nex| {
            acc.push_str(&nex);
            acc
        });
    let utc_creation_date = metainfo.creation_date().map(|c| UTC.timestamp(c, 0));

    println!("------Metainfo File Overview-----");

    println!("InfoHash: {}", info_hash_hex);
    println!("Main Tracker: {:?}", metainfo.main_tracker());
    println!("Comment: {:?}", metainfo.comment());
    println!("Creator: {:?}", metainfo.created_by());
    println!("Creation Date: {:?}", utc_creation_date);

    println!("Directory: {:?}", info.directory());
    println!("Piece Length: {:?}", info.piece_length());
    println!("Number Of Pieces: {}", info.pieces().count());
    println!("Number Of Files: {}", info.files().count());
    println!("Total File Size: {}\n",
             info.files().fold(0, |acc, nex| acc + nex.length()));

    print_files(bytes);
}

fn connect_to_tracker(metainfo_file: MetainfoFile,
                      peer_id: &str,
                      port: u16)
                      -> Result<(Vec<String>, String), String> {
    debug!("connecting to tracker: {:?}", metainfo_file.main_tracker());
    let info_hash = metainfo_file.info_hash();
    // TODO figure can this conversion to url-encoded form be done safely?
    let info_hash_str = unsafe { str::from_utf8_unchecked(info_hash.as_ref()) };
    // TODO this needs to be calculated based on what we have
    let total_len = metainfo_file.info().files().fold(0, |acc, nex| acc + nex.length());

    let mut url = Url::parse(metainfo_file.main_tracker().unwrap()).unwrap();
    url.query_pairs_mut()
        .append_pair("info_hash", info_hash_str)
        .append_pair("peer_id", peer_id)
        .append_pair("port", &(port.to_string()))
        // TODO parametrize this
        .append_pair("uploaded", "0")
        // TODO parametrize this
        .append_pair("downloaded", "0")
        // TODO see note on total_len above
        .append_pair("left", &(total_len.to_string()))
        .append_pair("compact", "1")
        .append_pair("event", "started")
        .append_pair("supportcrypto", "0");
    trace!("URL {:?}", url);

    let client = Client::new();
    let mut res = client.get(url).header(Connection::close()).send().unwrap();
    let mut buffer = Vec::new();
    res.read_to_end(&mut buffer).unwrap();
    debug!("{:?}", res);
    let bencode = Bencode::decode(&buffer).unwrap();
    trace!("{:?}", bencode);
    let (_, peers) = CompactPeersV4::from_bytes(bencode.dict()
            .unwrap()
            .lookup("peers")
            .unwrap()
            .bytes()
            .unwrap())
        .unwrap();
    trace!("{:?}", peers);
    let mut ip_ports: Vec<String> = Vec::new();
    for peer in peers.iter() {
        debug!("{:?}", peer);
        ip_ports.push(peer.to_string());
    }
    Ok((ip_ports, info_hash_str.to_string()))
}

fn peer_connections(peer_ip_ports: Vec<String>,
                    info_hash: &str,
                    peer_id: &str)
                    -> Result<(), String> {
    for peer_ip_port in peer_ip_ports {
        let peer_ip_port_cl = peer_ip_port.clone();
        let info_hash_cl = info_hash.to_string().clone();
        let peer_id_cl = peer_id.to_string().clone();
        thread::spawn(move || {
            handshake_peer(&peer_ip_port_cl, &info_hash_cl, &peer_id_cl);
        });
    }
    Ok(())
}

fn handshake_peer(peer_ip_port: &str, info_hash: &str, peer_id: &str) -> Result<(), String> {
    let mut buf = vec![0u8; PEER_HANDSHAKE_STRUCT_SZ];
    let mut ph = MutablePeerHandshakePacket::new(&mut buf).unwrap();
    ph.set_pstrlen("BitTorrent protocol".len() as u8);
    ph.set_pstr(&String::from("BitTorrent protocol").into_bytes());
    ph.set_reserved(&[0; 8]);
    ph.set_info_hash(info_hash.as_bytes());
    ph.set_peer_id(peer_id.as_bytes());
    match TcpStream::connect(peer_ip_port) {
        Ok(mut stream) => {
            debug!("Connection to peer {:?} successful!", peer_ip_port);
            match stream.write(ph.packet()) {
                Ok(_) => {
                    debug!("Sending message successful!");
                    Ok(())
                }
                Err(_) => {
                    error!("Sending message failed!");
                    Err("Sending message failed!".to_owned())
                }
            }
        }
        Err(e) => {
            error!("Connection to peer {:?} failed! {:?}", peer_ip_port, e);
            Err("Connection to peer failed!".to_owned())
        }
    }
}

fn main() {
    pretty_env_logger::init();

    let mut rl = Editor::<()>::new();
    if rl.load_history(HISTORY_FILE).is_err() {
        info!("No previous history!");
    }

    let now = SystemTime::now();
    let duration = now.duration_since(UNIX_EPOCH).unwrap();
    let client_id = "-bittorrent-rs-".to_owned() + &format!("{}", duration.as_secs() % 100_000);

    loop {
        let readline = rl.readline("> ");
        match readline {
            Ok(cmd) => {
                rl.add_history_entry(&cmd);

                let cmd = cmd.trim().split(' ').collect::<Vec<&str>>();

                match cmd[0] {
                    "help" | "h" => {
                        println!("Commands:
parse/p <torrent file path>      - show Metainfo File Overview
connect/c <torrent file path>    - initiate connection to tracker, and handshake with peers
showfiles/sf <torrent file path> - show files in the torrent
help/h                           - show this help");
                    }
                    "parse" | "p" => {
                        if cmd.len() != 2 {
                            error!("usage: parse <torrent file>");
                        } else {
                            let path = cmd[1];
                            debug!("MetainfoFile: {:?}", path);

                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();
                                    print_metainfo_overview(&bytes);
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "connect" | "c" => {
                        if cmd.len() != 2 {
                            error!("usage: connect <torrent file>");
                        } else {
                            let path = cmd[1];
                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();

                                    // TODO: generate peer ID
                                    let result =
                                        connect_to_tracker(MetainfoFile::from_bytes(&bytes)
                                                               .unwrap(),
                                                           &client_id,
                                                           6882)
                                            .unwrap();
                                    peer_connections(result.0, &result.1, &client_id).unwrap();
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "showfiles" | "sf" => {
                        if cmd.len() != 2 {
                            error!("usage: connect <torrent file>");
                        } else {
                            let path = cmd[1];
                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();

                                    print_files(&bytes);
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "" => {}
                    _ => {
                        println!("invalid command, see \"help\"");
                    }
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                error!("Error: {:?}", err);
                break;
            }
        }
    }
    rl.save_history(HISTORY_FILE).unwrap();
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Alek Ratzloff
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::{env, thread};
use std::sync::mpsc::{channel, Receiver};
use getopts::Options;
use trtable::TrTable;

/// Permutes over a single term's possible alternate characters
///
fn permute(index: usize, term: &str, built: String, trtab: &TrTable) {
  // base case
  if index == term.len() {
    println!("{}", built);
    return;
  }

  let c: char = term.chars()
    .nth(index)
    .expect("");

  let list = trtab.get(c);
  if list == None {
    // just add on the next string to the built and keep going
    let mut new_built: String = (&built).to_string();
    new_built.push(c);
    permute(index + 1, &term, new_built, &trtab);
  } else {
    // go through all of the possibilities
    let u_list = list.unwrap();
    for append in u_list {
      let new_built: String = (&built).to_string() + append;
      permute(index + 1, &term, new_built, &trtab);
    }
  }
}

fn all_combos(chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    println!("{}", &built);
    return;
  }

  for letter in chars {
    let mut new_built: String = (&built).to_string();
    new_built.push(*letter);
    all_combos(&chars, count - 1, new_built);
  }
}

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

fn main() {
  let args: Vec<String> = env::args().collect();
  let program = args[0].clone();

  let mut opts = Options::new();
  opts.optopt("c", "combos", "iterate through all possible printable combinations of N letters. You may specify a comma separated list, e.g. 1,2,4,5",
    "N");
  opts.optflag("", "no-upper", "specific to -c. If this flag is set, all combinations will omit upper case letters.");
  opts.optflag("", "no-lower", "specific to -c. If this flag is set, all combinations will omit lower case letters.");
  opts.optflag("", "no-numbers", "specific to -c. If this flag is set, all combinations will omit numbers.");
  opts.optflag("", "no-symbols", "specific to -c. If this flag is set, all combinations will omit symbols.");
  opts.optopt("p", "permute", "permute over a term with similar letters. Use the -t option to specify a table file",
    "TERM");
  opts.optopt("t", "trtab", "translation table file, in YAML format. Required with the -p option.", "FILENAME");

  opts.optopt("j", "threads", "number of threads to use. Default is 1.", "NUM_THREADS");
  opts.optflag("h", "help", "print this menu");
  let matches = match opts.parse(&args[1..]) {
    Ok(m) => { m },
    Err(f) => { panic!(f.to_string()) },
  };

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

  let table_file = matches.opt_str("t");
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

  // permutation
  if let Some(permute_str) = matches.opt_str("p") {
    match table_file {
      Some(trtab_fname) => {
        // load the file
        let trtab = TrTable::load(&trtab_fname)
          .ok()
          .unwrap();
        permute(0, &permute_str, "".to_string(), &trtab);
      },
      None => println!("You must specify a table file with the permute option")
    }
  }
  // all combinations
  if let Some(combos) = matches.opt_str("c") {
    let counts = combos.split(",")
      .map(|n| n.parse::<usize>().unwrap());

    let mut all_chars = Vec::new();
    // upper case letters
    if !matches.opt_present("no-upper") {
      for l in char_iter::new('\x41', '\x5a') {
        all_chars.push(l);
      }
    }
    // lower case letters
    if !matches.opt_present("no-lower") {
      for l in char_iter::new('\x61', '\x7a') {
        all_chars.push(l);
      }
    }
    // numbers
    if !matches.opt_present("no-numbers") {
      for l in char_iter::new('\x30', '\x39') {
        all_chars.push(l);
      }
    }
    // symbols
    if !matches.opt_present("no-symbols") {
      // <space> to /
      for l in char_iter::new('\x20', '\x2F') {
        all_chars.push(l);
      }
      // : to @
      for l in char_iter::new('\x3a', '\x40') {
        all_chars.push(l);
      }
      // [ to `
      for l in char_iter::new('\x5b', '\x60') {
        all_chars.push(l);
      }
      // { to ~
      for l in char_iter::new('\x7b', '\x7e') {
        all_chars.push(l);
      }
    }

    // threads that are currently spawned
    let mut threads: Vec<(thread::JoinHandle<()>, Receiver<i32>)> = Vec::new();

    for n in counts {
      for letter in &all_chars {
        // get whether we have any thread slots open; if not, then
        // block until a thread finishes.
        while threads.len() == num_threads {
          let mut offset = 0;
          for index in 0 .. threads.len() {
            let remove = threads[index - offset].1
              .try_recv()
              .is_ok();

            if remove {
              threads.remove(index - offset);
              offset += 1;
            }
          }
          thread::sleep_ms(1);
        }

        // set up a new channel
        let (tx, rx) = channel();
        // create an intermediate var for the all_chars so we don't have to copy it fifty million times
        let all_chars_imdt = all_chars.clone();
        let l = *letter;
        let child = thread::spawn(move || {
          // shift ownership to the thread
          let all_chars_dup = all_chars_imdt;
          let mut new_built = String::new();

          new_built.push(l);
          all_combos(&all_chars_dup, n - 1, new_built);
          drop(tx.send(0));
        });
        threads.push((child, rx));
      }
    }

    // let the rest of the threads finish up
    for (th, _) in threads {
      drop(th.join());
    }
  }
}
<commit_msg>Cleaned up the code some, moved some stuff out of main so we don't have a megalith<commit_after>/*
 * Copyright (C) 2015 Alek Ratzloff
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::{env, thread};
use std::sync::mpsc::{channel, Receiver};
use getopts::{Options, Matches};
use trtable::TrTable;

/// Permutes over a single term's possible alternate characters
///
fn permute(index: usize, term: &str, built: String, trtab: &TrTable) {
  // base case
  if index == term.len() {
    println!("{}", built);
    return;
  }

  let c: char = term.chars()
    .nth(index)
    .expect("");

  let list = trtab.get(c);
  if list == None {
    // just add on the next string to the built and keep going
    let mut new_built: String = (&built).to_string();
    new_built.push(c);
    permute(index + 1, &term, new_built, &trtab);
  } else {
    // go through all of the possibilities
    let u_list = list.unwrap();
    for append in u_list {
      let new_built: String = (&built).to_string() + append;
      permute(index + 1, &term, new_built, &trtab);
    }
  }
}

fn all_combos(chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    println!("{}", &built);
    return;
  }

  for letter in chars {
    let mut new_built: String = (&built).to_string();
    new_built.push(*letter);
    all_combos(&chars, count - 1, new_built);
  }
}

fn parse_opts(args: &Vec<String>) -> (Options, Matches) {
  let mut opts = Options::new();
  opts.optopt("c", "combos", "iterate through all possible printable combinations of N letters. You may specify a comma separated list, e.g. 1,2,4,5",
    "N");
  opts.optflag("", "no-upper", "specific to -c. If this flag is set, all combinations will omit upper case letters.");
  opts.optflag("", "no-lower", "specific to -c. If this flag is set, all combinations will omit lower case letters.");
  opts.optflag("", "no-numbers", "specific to -c. If this flag is set, all combinations will omit numbers.");
  opts.optflag("", "no-symbols", "specific to -c. If this flag is set, all combinations will omit symbols.");
  opts.optopt("p", "permute", "permute over a term with similar letters. Use the -t option to specify a table file",
    "TERM");
  opts.optopt("t", "trtab", "translation table file, in YAML format. Required with the -p option.", "FILENAME");

  opts.optopt("j", "threads", "number of threads to use. Default is 1.", "NUM_THREADS");
  opts.optflag("h", "help", "print this menu");
  let matches = match opts.parse(&args[1..]) {
    Ok(m) => { m },
    Err(f) => { panic!(f.to_string()) },
  };
  (opts, matches)
}

fn main() {
  let args: Vec<String> = env::args().collect();
  let (opts, matches) = parse_opts(&args);

  // print help and exit
  if matches.opt_present("h") {
    let program = match args.get(0) {
      Some(s) => s,
      None => unreachable!(),
    };
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
    return;
  }

  // get the table file if it's there
  let table_file = matches.opt_str("t");
  // get the number of threads; if it doesn't exist, default to 1
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

  // permutation
  if let Some(permute_str) = matches.opt_str("p") {
    match table_file {
      Some(trtab_fname) => {
        // load the file
        let trtab = TrTable::load(&trtab_fname)
          .ok()
          .unwrap();
        permute(0, &permute_str, "".to_string(), &trtab);
      },
      None => println!("You must specify a table file with the permute option")
    }
  }
  // all combinations
  if let Some(combos) = matches.opt_str("c") {
    let counts = combos.split(",")
      .map(|n| n.parse::<usize>().unwrap());

    let mut all_chars = Vec::new();
    // upper case letters
    if !matches.opt_present("no-upper") {
      for l in char_iter::new('\x41', '\x5a') {
        all_chars.push(l);
      }
    }
    // lower case letters
    if !matches.opt_present("no-lower") {
      for l in char_iter::new('\x61', '\x7a') {
        all_chars.push(l);
      }
    }
    // numbers
    if !matches.opt_present("no-numbers") {
      for l in char_iter::new('\x30', '\x39') {
        all_chars.push(l);
      }
    }
    // symbols
    if !matches.opt_present("no-symbols") {
      // <space> to /
      for l in char_iter::new('\x20', '\x2F') {
        all_chars.push(l);
      }
      // : to @
      for l in char_iter::new('\x3a', '\x40') {
        all_chars.push(l);
      }
      // [ to `
      for l in char_iter::new('\x5b', '\x60') {
        all_chars.push(l);
      }
      // { to ~
      for l in char_iter::new('\x7b', '\x7e') {
        all_chars.push(l);
      }
    }

    // threads that are currently spawned
    let mut threads: Vec<(thread::JoinHandle<()>, Receiver<i32>)> = Vec::new();

    for n in counts {
      for letter in &all_chars {
        // get whether we have any thread slots open; if not, then
        // block until a thread finishes.
        while threads.len() == num_threads {
          let mut offset = 0;
          for index in 0 .. threads.len() {
            let remove = threads[index - offset].1
              .try_recv()
              .is_ok();

            if remove {
              threads.remove(index - offset);
              offset += 1;
            }
          }
          thread::sleep_ms(1);
        }

        // set up a new channel
        let (tx, rx) = channel();
        // create an intermediate var for the all_chars so we don't have to copy it fifty million times
        let all_chars_imdt = all_chars.clone();
        let l = *letter;
        let child = thread::spawn(move || {
          // shift ownership to the thread
          let all_chars_dup = all_chars_imdt;
          let mut new_built = String::new();

          new_built.push(l);
          all_combos(&all_chars_dup, n - 1, new_built);
          drop(tx.send(0));
        });
        threads.push((child, rx));
      }
    }

    // let the rest of the threads finish up
    for (th, _) in threads {
      drop(th.join());
    }
  }
}
<|endoftext|>


<commit_before>extern crate num_cpus;

use std::env;
use std::io::{self, Write, StderrLock};
use std::process::{Command, exit};
use std::thread::{self, JoinHandle};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/* TODO: Functionality can be increased to accept the following syntaxes from GNU Parallel:
 - Stdin support is currently missing.
 - {N}, {N.}, etc.
 - parallel command {1} {2} {3} ::: 1 2 3 ::: 4 5 6 ::: 7 8 9
 - paralllel command ::: a b c :::+ 1 2 3 ::: d e f :::+ 4 5 6
*/

fn main() {
    let stderr = io::stderr();
    let mut ncores = num_cpus::get();
    let mut command = String::new();
    let mut arg_tokens = Vec::new();
    let mut inputs = Vec::new();

    // Let's collect all parameters that we need from the program's arguments.
    // If an error is returned, this will handle that error as efficiently as possible.
    if let Err(why) = parse_arguments(&mut ncores, &mut command, &mut arg_tokens, &mut inputs) {
        let mut stderr = stderr.lock();
        let _ = stderr.write(b"parallel: parsing error: ");
        match why {
            ParseErr::JobsNaN(value) => {
                let _ = stderr.write(b"jobs parameter, '");
                let _ = stderr.write(value.as_bytes());
                let _ = stderr.write(b"', is not a number.\n");
            },
            _ => {
                let message: &[u8] = match why {
                    ParseErr::InputVarsNotDefined => b"input variables were not defined.\n",
                    ParseErr::JobsNoValue         => b"no jobs parameter was defined.\n",
                    _ => unreachable!()
                };
                let _ = stderr.write(message);
            }
        };
        exit(1);
    }

    // If no command was given, then the inputs are actually commands themselves.
    let input_is_command = command.is_empty();

    // It will be useful to know the number of inputs, to know when to quit.
    let num_inputs = inputs.len();

    // Stores the next input to be processed
    let shared_counter = Arc::new(AtomicUsize::new(0));

    // We will share the same list of inputs with each thread.
    let shared_input = Arc::new(inputs);

    // First we will create as many threads as `ncores` specifies.
    // The `threads` vector will contain the thread handles needed to
    // know when to quit the program.
    let mut threads: Vec<JoinHandle<()>> = Vec::with_capacity(ncores);
    for slot in 1..ncores+1 {
        // The command that each input variable will be sent to.
        let command = command.clone();
        // The arguments for the command.
        let argument_tokens = arg_tokens.clone();
        // Allow the thread to gain access to the list of inputs.
        let input = shared_input.clone();
        // Allow the thread to access the current command counter
        let counter = shared_counter.clone();
        // Allow the thread to know when it's time to stop.
        let num_inputs = num_inputs;

        // The actual thread where the work will happen on incoming data.
        let handle: JoinHandle<()> = thread::spawn(move || {
            let slot_number = slot;
            let stderr = io::stderr();
            loop {
                // Obtain the Nth item and it's job ID from the list of inputs.
                let (input_var, job_id) = {
                    // Atomically increment the counter
                    let old_counter = counter.fetch_add(1, Ordering::SeqCst);
                    if old_counter >= num_inputs {
                        break
                    } else {
                        let input_var = &input[old_counter];
                        let job_id = old_counter + 1;
                        (input_var, job_id)
                    }
                };

                if input_is_command {
                    // The inputs are actually the commands.
                    let mut iterator = input_var.split_whitespace();
                    let actual_command = iterator.next().unwrap();
                    let args = iterator.collect::<Vec<&str>>();
                    if let Err(_) = Command::new(actual_command).args(&args).status() {
                        let mut stderr = stderr.lock();
                        let _ = stderr.write(b"parallel: command error: ");
                        let _ = stderr.write(input_var.as_bytes());
                        let _ = stderr.write(b"\n");
                    }
                } else {
                    // Build a command by merging the command template with the input,
                    // and then execute that command.
                    let (slot, job) = (slot_number.to_string(), job_id.to_string());
                    if let Err(cmd_err) = cmd_builder(input_var, &command, &argument_tokens,
                        &slot, &job)
                    {
                        let mut stderr = stderr.lock();
                        cmd_err.handle(&mut stderr);
                    }
                }
            }
        });

        // After the thread has been created, add the important pieces needed by the
        // main thread to the `threads` vector.
        threads.push(handle);
    }

    for thread in threads.into_iter() { thread.join().unwrap(); }
}

enum CommandErr {
    Failed(String, Vec<String>)
}

impl CommandErr {
    fn handle(self, stderr: &mut StderrLock) {
        let _ = stderr.write(b"parallel: command error: ");
        match self {
            CommandErr::Failed(command, arguments) => {
                let _ = stderr.write(command.as_bytes());
                for arg in &arguments {
                    let _ = stderr.write(b" ");
                    let _ = stderr.write(arg.as_bytes());
                }
                let _ = stderr.write(b"\n");
            }
        }
    }
}

#[derive(Clone, PartialEq)]
enum Token {
    Character(char),
    Placeholder,
    RemoveExtension,
    Basename,
    Dirname,
    BaseAndExt,
    Slot,
    Job
}

fn tokenize(template: &str) -> Vec<Token> {
    let mut matching = false;
    let mut tokens = Vec::new();
    let mut pattern = String::new();
    for character in template.chars() {
        match (character, matching) {
            ('{', false) => matching = true,
            ('}', true)  => {
                matching = false;
                if pattern.is_empty() {
                    tokens.push(Token::Placeholder);
                } else {
                    match match_token(&pattern) {
                        Some(token) => tokens.push(token),
                        None => {
                            tokens.push(Token::Character('{'));
                            for character in pattern.chars() {
                                tokens.push(Token::Character(character));
                            }
                            tokens.push(Token::Character('}'));
                        }
                    }
                    pattern.clear();
                }
            }
            (_, false)  => tokens.push(Token::Character(character)),
            (_, true) => pattern.push(character)
        }
    }
    tokens
}

fn match_token(pattern: &str) -> Option<Token> {
    match pattern {
        "."  => Some(Token::RemoveExtension),
        "#"  => Some(Token::Job),
        "%"  => Some(Token::Slot),
        "/"  => Some(Token::Basename),
        "//" => Some(Token::Dirname),
        "/." => Some(Token::BaseAndExt),
        _    => None
    }

}

/// Builds the command and executes it
fn cmd_builder(input: &str, command: &str, arg_tokens: &[Token], slot_id: &str, job_id: &str)
    -> Result<(), CommandErr>
{
    let mut arguments = String::new();
    for arg in arg_tokens {
        match *arg {
            Token::Character(arg) => arguments.push(arg),
            Token::Basename => arguments.push_str(input),
            Token::BaseAndExt => arguments.push_str(basename(remove_extension(input))),
            Token::Dirname => arguments.push_str(dirname(input)),
            Token::Job => arguments.push_str(job_id),
            Token::Placeholder => arguments.push_str(input),
            Token::RemoveExtension => arguments.push_str(remove_extension(input)),
            Token::Slot => arguments.push_str(slot_id)
        }
    }

    let placeholder_exists = arg_tokens.iter().any(|ref x| {
        x == &&Token::BaseAndExt || x == &&Token::Basename || x == &&Token::Dirname ||
        x == &&Token::Job || x == &&Token::Placeholder || x == &&Token::RemoveExtension ||
        x == &&Token::Slot
    });

    if !placeholder_exists {
        arguments.push_str(input);
    }

    let arguments = arguments.split_whitespace().map(|x| x.to_owned()).collect::<Vec<String>>();

    if let Err(_) = Command::new(&command).args(&arguments).status() {
        return Err(CommandErr::Failed(String::from(command), arguments));
    }
    Ok(())
}

/// Removes the extension of a given input
fn remove_extension(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '.' { index = id; }
    }
    if index == 0 { input } else { &input[0..index] }
}

fn basename(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '/' { index = id; }
    }
    if index == 0 { input } else { &input[index+1..] }
}

fn dirname(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '/' { index = id; }
    }
    if index == 0 { input } else { &input[0..index] }
}

enum ParseErr {
    JobsNaN(String),
    JobsNoValue,
    InputVarsNotDefined,
}

// Parses input arguments and stores their values into their associated variabless.
fn parse_arguments(ncores: &mut usize, command: &mut String, arg_tokens: &mut Vec<Token>,
    input_variables: &mut Vec<String>) -> Result<(), ParseErr>
{
    let mut parsing_arguments = true;
    let mut command_is_set    = false;
    let mut raw_args = env::args().skip(1).peekable();
    let mut comm = String::new();
    while let Some(argument) = raw_args.next() {
        if parsing_arguments {
            match argument.as_str() {
                // Defines the number of jobs to run in parallel.
                "-j"  => {
                    match raw_args.peek() {
                        Some(val) => match val.parse::<usize>() {
                            Ok(val) => *ncores = val,
                            Err(_)  => return Err(ParseErr::JobsNaN(val.clone()))
                        },
                        None => return Err(ParseErr::JobsNoValue)
                    }
                    let _ = raw_args.next();
                },
                // Arguments after `:::` are input values.
                ":::" => parsing_arguments = false,
                _ => {
                    if command_is_set {
                        comm.push(' ');
                        comm.push_str(&argument);
                    } else {
                        comm.push_str(&argument);
                        command_is_set = true;
                    }

                }
            }
        } else {
            input_variables.push(argument);
        }
    }

    // This will fill in command and argument information needed by the threads.
    // If there is a space in the argument, then the command has arguments
    match comm.chars().position(|x| x == ' ') {
        Some(pos) => {
            *command    = String::from(&comm[0..pos]);
            *arg_tokens = tokenize(&comm[pos+1..]);
        },
        None => *command = comm
    }

    if input_variables.is_empty() { return Err(ParseErr::InputVarsNotDefined) }
    Ok(())
}
<commit_msg>Add a couple of tests<commit_after>extern crate num_cpus;

use std::env;
use std::io::{self, Write, StderrLock};
use std::process::{Command, exit};
use std::thread::{self, JoinHandle};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/* TODO: Functionality can be increased to accept the following syntaxes from GNU Parallel:
 - Stdin support is currently missing.
 - {N}, {N.}, etc.
 - parallel command {1} {2} {3} ::: 1 2 3 ::: 4 5 6 ::: 7 8 9
 - paralllel command ::: a b c :::+ 1 2 3 ::: d e f :::+ 4 5 6
*/

fn main() {
    let stderr = io::stderr();
    let mut ncores = num_cpus::get();
    let mut command = String::new();
    let mut arg_tokens = Vec::new();
    let mut inputs = Vec::new();

    // Let's collect all parameters that we need from the program's arguments.
    // If an error is returned, this will handle that error as efficiently as possible.
    if let Err(why) = parse_arguments(&mut ncores, &mut command, &mut arg_tokens, &mut inputs) {
        let mut stderr = stderr.lock();
        let _ = stderr.write(b"parallel: parsing error: ");
        match why {
            ParseErr::JobsNaN(value) => {
                let _ = stderr.write(b"jobs parameter, '");
                let _ = stderr.write(value.as_bytes());
                let _ = stderr.write(b"', is not a number.\n");
            },
            _ => {
                let message: &[u8] = match why {
                    ParseErr::InputVarsNotDefined => b"input variables were not defined.\n",
                    ParseErr::JobsNoValue         => b"no jobs parameter was defined.\n",
                    _ => unreachable!()
                };
                let _ = stderr.write(message);
            }
        };
        exit(1);
    }

    // If no command was given, then the inputs are actually commands themselves.
    let input_is_command = command.is_empty();

    // It will be useful to know the number of inputs, to know when to quit.
    let num_inputs = inputs.len();

    // Stores the next input to be processed
    let shared_counter = Arc::new(AtomicUsize::new(0));

    // We will share the same list of inputs with each thread.
    let shared_input = Arc::new(inputs);

    // First we will create as many threads as `ncores` specifies.
    // The `threads` vector will contain the thread handles needed to
    // know when to quit the program.
    let mut threads: Vec<JoinHandle<()>> = Vec::with_capacity(ncores);
    for slot in 1..ncores+1 {
        // The command that each input variable will be sent to.
        let command = command.clone();
        // The arguments for the command.
        let argument_tokens = arg_tokens.clone();
        // Allow the thread to gain access to the list of inputs.
        let input = shared_input.clone();
        // Allow the thread to access the current command counter
        let counter = shared_counter.clone();
        // Allow the thread to know when it's time to stop.
        let num_inputs = num_inputs;

        // The actual thread where the work will happen on incoming data.
        let handle: JoinHandle<()> = thread::spawn(move || {
            let slot_number = slot;
            let stderr = io::stderr();
            loop {
                // Obtain the Nth item and it's job ID from the list of inputs.
                let (input_var, job_id) = {
                    // Atomically increment the counter
                    let old_counter = counter.fetch_add(1, Ordering::SeqCst);
                    if old_counter >= num_inputs {
                        break
                    } else {
                        let input_var = &input[old_counter];
                        let job_id = old_counter + 1;
                        (input_var, job_id)
                    }
                };

                if input_is_command {
                    // The inputs are actually the commands.
                    let mut iterator = input_var.split_whitespace();
                    let actual_command = iterator.next().unwrap();
                    let args = iterator.collect::<Vec<&str>>();
                    if let Err(_) = Command::new(actual_command).args(&args).status() {
                        let mut stderr = stderr.lock();
                        let _ = stderr.write(b"parallel: command error: ");
                        let _ = stderr.write(input_var.as_bytes());
                        let _ = stderr.write(b"\n");
                    }
                } else {
                    // Build a command by merging the command template with the input,
                    // and then execute that command.
                    let (slot, job) = (slot_number.to_string(), job_id.to_string());
                    if let Err(cmd_err) = cmd_builder(input_var, &command, &argument_tokens,
                        &slot, &job)
                    {
                        let mut stderr = stderr.lock();
                        cmd_err.handle(&mut stderr);
                    }
                }
            }
        });

        // After the thread has been created, add the important pieces needed by the
        // main thread to the `threads` vector.
        threads.push(handle);
    }

    for thread in threads.into_iter() { thread.join().unwrap(); }
}

enum CommandErr {
    Failed(String, Vec<String>)
}

impl CommandErr {
    fn handle(self, stderr: &mut StderrLock) {
        let _ = stderr.write(b"parallel: command error: ");
        match self {
            CommandErr::Failed(command, arguments) => {
                let _ = stderr.write(command.as_bytes());
                for arg in &arguments {
                    let _ = stderr.write(b" ");
                    let _ = stderr.write(arg.as_bytes());
                }
                let _ = stderr.write(b"\n");
            }
        }
    }
}

#[derive(Clone, PartialEq, Debug)]
enum Token {
    Character(char),
    Placeholder,
    RemoveExtension,
    Basename,
    Dirname,
    BaseAndExt,
    Slot,
    Job
}

fn tokenize(template: &str) -> Vec<Token> {
    let mut matching = false;
    let mut tokens = Vec::new();
    let mut pattern = String::new();
    for character in template.chars() {
        match (character, matching) {
            ('{', false) => matching = true,
            ('}', true)  => {
                matching = false;
                if pattern.is_empty() {
                    tokens.push(Token::Placeholder);
                } else {
                    match match_token(&pattern) {
                        Some(token) => tokens.push(token),
                        None => {
                            tokens.push(Token::Character('{'));
                            for character in pattern.chars() {
                                tokens.push(Token::Character(character));
                            }
                            tokens.push(Token::Character('}'));
                        }
                    }
                    pattern.clear();
                }
            }
            (_, false)  => tokens.push(Token::Character(character)),
            (_, true) => pattern.push(character)
        }
    }
    tokens
}

fn match_token(pattern: &str) -> Option<Token> {
    match pattern {
        "."  => Some(Token::RemoveExtension),
        "#"  => Some(Token::Job),
        "%"  => Some(Token::Slot),
        "/"  => Some(Token::Basename),
        "//" => Some(Token::Dirname),
        "/." => Some(Token::BaseAndExt),
        _    => None
    }

}

/// Builds the command and executes it
fn cmd_builder(input: &str, command: &str, arg_tokens: &[Token], slot_id: &str, job_id: &str)
    -> Result<(), CommandErr>
{
    let mut arguments = String::new();
    for arg in arg_tokens {
        match *arg {
            Token::Character(arg) => arguments.push(arg),
            Token::Basename => arguments.push_str(basename(input)),
            Token::BaseAndExt => arguments.push_str(basename(remove_extension(input))),
            Token::Dirname => arguments.push_str(dirname(input)),
            Token::Job => arguments.push_str(job_id),
            Token::Placeholder => arguments.push_str(input),
            Token::RemoveExtension => arguments.push_str(remove_extension(input)),
            Token::Slot => arguments.push_str(slot_id)
        }
    }

    let placeholder_exists = arg_tokens.iter().any(|ref x| {
        x == &&Token::BaseAndExt || x == &&Token::Basename || x == &&Token::Dirname ||
        x == &&Token::Job || x == &&Token::Placeholder || x == &&Token::RemoveExtension ||
        x == &&Token::Slot
    });

    if !placeholder_exists {
        arguments.push_str(input);
    }

    let arguments = arguments.split_whitespace().map(|x| x.to_owned()).collect::<Vec<String>>();

    if let Err(_) = Command::new(&command).args(&arguments).status() {
        return Err(CommandErr::Failed(String::from(command), arguments));
    }
    Ok(())
}

/// Removes the extension of a given input
fn remove_extension(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '.' { index = id; }
    }
    if index == 0 { input } else { &input[0..index] }
}

fn basename(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '/' { index = id; }
    }
    if index == 0 { input } else { &input[index+1..] }
}

fn dirname(input: &str) -> &str {
    let mut index = 0;
    for (id, character) in input.chars().enumerate() {
        if character == '/' { index = id; }
    }
    if index == 0 { input } else { &input[0..index] }
}

enum ParseErr {
    JobsNaN(String),
    JobsNoValue,
    InputVarsNotDefined,
}

// Parses input arguments and stores their values into their associated variabless.
fn parse_arguments(ncores: &mut usize, command: &mut String, arg_tokens: &mut Vec<Token>,
    input_variables: &mut Vec<String>) -> Result<(), ParseErr>
{
    let mut parsing_arguments = true;
    let mut command_is_set    = false;
    let mut raw_args = env::args().skip(1).peekable();
    let mut comm = String::new();
    while let Some(argument) = raw_args.next() {
        if parsing_arguments {
            match argument.as_str() {
                // Defines the number of jobs to run in parallel.
                "-j"  => {
                    match raw_args.peek() {
                        Some(val) => match val.parse::<usize>() {
                            Ok(val) => *ncores = val,
                            Err(_)  => return Err(ParseErr::JobsNaN(val.clone()))
                        },
                        None => return Err(ParseErr::JobsNoValue)
                    }
                    let _ = raw_args.next();
                },
                // Arguments after `:::` are input values.
                ":::" => parsing_arguments = false,
                _ => {
                    if command_is_set {
                        comm.push(' ');
                        comm.push_str(&argument);
                    } else {
                        comm.push_str(&argument);
                        command_is_set = true;
                    }

                }
            }
        } else {
            input_variables.push(argument);
        }
    }

    // This will fill in command and argument information needed by the threads.
    // If there is a space in the argument, then the command has arguments
    match comm.chars().position(|x| x == ' ') {
        Some(pos) => {
            *command    = String::from(&comm[0..pos]);
            *arg_tokens = tokenize(&comm[pos+1..]);
        },
        None => *command = comm
    }

    if input_variables.is_empty() { return Err(ParseErr::InputVarsNotDefined) }
    Ok(())
}

#[test]
fn tokenizer_character() {
    assert_eq!(tokenize("foo"), vec![Token::Character('f'), Token::Character('o'), Token::Character('o')]);
}

#[test]
fn tokenizer_placeholder() {
    assert_eq!(tokenize("{}"), vec![Token::Placeholder]);
}

#[test]
fn tokenizer_remove_extension() {
    assert_eq!(tokenize("{.}"), vec![Token::RemoveExtension]);
}

#[test]
fn tokenizer_basename() {
    assert_eq!(tokenize("{/}"), vec![Token::Basename]);
}

#[test]
fn tokenizer_dirname() {
    assert_eq!(tokenize("{//}"), vec![Token::Dirname]);
}

#[test]
fn tokenizer_base_and_ext() {
    assert_eq!(tokenize("{/.}"), vec![Token::BaseAndExt]);
}

#[test]
fn tokenizer_slot() {
    assert_eq!(tokenize("{%}"), vec![Token::Slot]);
}

#[test]
fn tokenizer_job() {
    assert_eq!(tokenize("{#}"), vec![Token::Job]);
}

#[test]
fn tokenizer_multiple() {
    assert_eq!(tokenize("foo {} bar"), vec![Token::Character('f'), Token::Character('o'), Token::Character('o'), Token::Character(' '), Token::Placeholder, Token::Character(' '), Token::Character('b'), Token::Character('a'), Token::Character('r')]);
}

#[test]
fn tokenizer_no_space() {
    assert_eq!(tokenize("foo{}bar"), vec![Token::Character('f'), Token::Character('o'), Token::Character('o'), Token::Placeholder, Token::Character('b'), Token::Character('a'), Token::Character('r')]);
}

#[test]
fn path_remove_ext_simple() {
    assert_eq!(remove_extension("foo.txt"), "foo");
}

#[test]
fn path_remove_ext_dir() {
    assert_eq!(remove_extension("dir/foo.txt"), "dir/foo");
}

#[test]
fn path_remove_ext_empty() {
    assert_eq!(remove_extension(""), "");
}

#[test]
fn path_basename_simple() {
    assert_eq!(basename("foo.txt"), "foo.txt");
}

#[test]
fn path_basename_dir() {
    assert_eq!(basename("dir/foo.txt"), "foo.txt");
}

#[test]
fn path_basename_empty() {
    assert_eq!(basename(""), "");
}

#[test]
#[ignore]
fn path_dirname_simple() {
    assert_eq!(dirname("foo.txt"), ".");
}

#[test]
fn path_dirname_dir() {
    assert_eq!(dirname("dir/foo.txt"), "dir");
}

#[test]
fn path_dirname_empty() {
    assert_eq!(dirname(""), "");
}
<|endoftext|>


<commit_before>extern crate par;
extern crate hyper;
use std::io::Read;

fn main() {
    /*let bar = par::Bar::new(
        100, Some("example ")
    );
    for _ in 1..101 {
        bar.incr();
        bar.update();
        std::thread::sleep_ms(10);
    }*/

    let client = hyper::Client::new();
    let mut res = client.get("https://github.com/").send().unwrap();
    let mut buffer = String::new();
    let bar = par::Bar::new(16763, Some("dl "));
    let mut reader = par::ProgressReader::new(res, &bar);
    reader.read_to_string(&mut buffer);
}
<commit_msg>simple example for now<commit_after>extern crate par;

fn main() {
    let bar = par::Bar::new(
        1024*4
    );
    for _ in 1..(1024 * 4) {
        bar.incr();
        bar.update();
        std::thread::sleep_ms(10);
    }
    bar.finish_print("done");
}
<|endoftext|>


<commit_before>extern crate chan_signal;
#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate fuse;
extern crate libc;
#[macro_use]
extern crate log;
extern crate time;

use std::env;
use std::error::Error;
use std::ffi::{OsStr, OsString};
use std::os::unix::ffi::OsStrExt;
use std::io;
use std::path::Path;
use std::str::FromStr;
use std::thread;

use chan_signal::Signal;
use clap::{App, Arg};
use env_logger::LogBuilder;
use log::LogRecord;

mod catfs;
mod flags;
mod evicter;

use catfs::error;
use catfs::flags::{DiskSpace, FlagStorage};
use catfs::rlibc;

fn main() {
    if let Err(e) = main_internal() {
        error!("Cannot mount: {}", e);
        std::process::exit(1);
    }
}

fn main_internal() -> error::Result<()> {
    let format = |record: &LogRecord| {
        let t = time::now();
        format!(
            "{} {:5} - {}",
            time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
            record.level(),
            record.args()
        )
    };

    let mut builder = LogBuilder::new();
    builder.format(format);

    if env::var("RUST_LOG").is_ok() {
        builder.parse(&env::var("RUST_LOG").unwrap());
    }

    builder.init().unwrap();

    let mut flags: FlagStorage = Default::default();
    let mut test = false;

    let app = App::new("catfs")
        .about("Cache Anything FileSystem")
        .version(crate_version!());

    {
        fn diskspace_validator(s: String) -> Result<(), String> {
            DiskSpace::from_str(&s).map(|_| ()).map_err(
                |e| e.to_str().to_owned(),
            )
        }

        fn path_validator(s: String) -> Result<(), String> {
            Path::new(&s)
                .canonicalize()
                .map_err(|e| e.description().to_owned())
                .and_then(|p| if p.is_dir() {
                    Ok(())
                } else {
                    Err("is not a directory".to_owned())
                })
        }

        let mut args = [
            flags::Flag {
                arg: Arg::with_name("space")
                    .long("free")
                    .takes_value(true)
                    .help(
                        "Ensure filesystem has at least this much free space. (ex: 9.5%, 10G)",
                    )
                    .validator(diskspace_validator),
                value: &mut flags.free_space,
            },
            flags::Flag {
                arg: Arg::with_name("foreground").short("f").help(
                    "Run catfs in foreground.",
                ),
                value: &mut flags.foreground,
            },
            flags::Flag {
                arg: Arg::with_name("option")
                    .short("o")
                    .takes_value(true)
                    .multiple(true)
                    .help("Additional system-specific mount options. Be careful!"),
                value: &mut flags.mount_options,
            },
            flags::Flag {
                arg: Arg::with_name("test").long("test").help(
                    "Exit after parsing arguments",
                ),
                value: &mut test,
            },
            flags::Flag {
                arg: Arg::with_name("from")
                    .index(1)
                    .required(true)
                    .help("Cache files from this directory.")
                    .validator(path_validator),
                value: &mut flags.cat_from,
            },
            flags::Flag {
                arg: Arg::with_name("to")
                    .index(2)
                    .required(true)
                    .help("Cache files to this directory.")
                    .validator(path_validator),
                value: &mut flags.cat_to,
            },
            flags::Flag {
                arg: Arg::with_name("mountpoint")
                    .index(3)
                    .required(true)
                    .help("Expose the mount point at this directory.")
                    .validator(path_validator),
                value: &mut flags.mount_point,
            },
        ];

        flags::parse_options(app, &mut args);
    }

    if test {
        return Ok(());
    }

    flags.mount_options.push(
        OsString::from("default_permissions"),
    );

    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);
    let path_from = Path::new(&flags.cat_from).canonicalize()?;
    let path_to = Path::new(&flags.cat_to).canonicalize()?;
    let fs = catfs::CatFS::new(&path_from, &path_to)?;
    let cache_dir = fs.get_cache_dir()?;
    let mut options: Vec<&OsStr> = Vec::new();
    for i in 0..flags.mount_options.len() {
        options.push(&flags.mount_options[i]);
    }

    {
        let mut session = fuse::Session::new(fs, Path::new(&flags.mount_point), &options)?;
        thread::spawn(move || {
            if let Err(e) = session.run() {
                error!("session.run() = {}", e);
            }
            info!("{:?} unmounted", session.mountpoint());
            unsafe { libc::kill(libc::getpid(), libc::SIGTERM) };
        });

        let mut ev = evicter::Evicter::new(cache_dir, &flags.free_space);
        ev.run();
        // unmount after we get signaled becausep session will go out of scope
        let s = signal.recv().unwrap();
        info!(
            "Received {:?}, attempting to unmount {:?}",
            s,
            flags.mount_point
        );
        unmount(Path::new(&flags.mount_point))?;
    }
    rlibc::close(cache_dir)?;
    return Ok(());
}

use libc::{c_char, c_int};
use std::ffi::{CString, CStr};
/// Unmount an arbitrary mount point
pub fn unmount(mountpoint: &Path) -> io::Result<()> {
    // fuse_unmount_compat22 unfortunately doesn't return a status. Additionally,
    // it attempts to call realpath, which in turn calls into the filesystem. So
    // if the filesystem returns an error, the unmount does not take place, with
    // no indication of the error available to the caller. So we call unmount
    // directly, which is what osxfuse does anyway, since we already converted
    // to the real path when we first mounted.

    #[cfg(any(target_os = "macos", target_os = "freebsd", target_os = "dragonfly",
                target_os = "openbsd", target_os = "bitrig", target_os = "netbsd"))]
    #[inline]
    fn libc_umount(mnt: &CStr) -> c_int {
        unsafe { libc::unmount(mnt.as_ptr(), 0) }
    }

    #[cfg(not(any(target_os = "macos", target_os = "freebsd", target_os = "dragonfly",
                      target_os = "openbsd", target_os = "bitrig", target_os = "netbsd")))]
    #[inline]
    fn libc_umount(mnt: &CStr) -> c_int {
        use std::io::ErrorKind::PermissionDenied;

        let rc = unsafe { libc::umount(mnt.as_ptr()) };
        if rc < 0 && io::Error::last_os_error().kind() == PermissionDenied {
            // Linux always returns EPERM for non-root users.  We have to let the
            // library go through the setuid-root "fusermount -u" to unmount.
            unsafe {
                fuse_unmount_compat22(mnt.as_ptr());
            }
            0
        } else {
            rc
        }
    }

    let mnt = try!(CString::new(mountpoint.as_os_str().as_bytes()));
    let rc = libc_umount(&mnt);
    if rc < 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

extern "system" {
    pub fn fuse_unmount_compat22(mountpoint: *const c_char);
}
<commit_msg>pass default_permissions correctly<commit_after>extern crate chan_signal;
#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate fuse;
extern crate libc;
#[macro_use]
extern crate log;
extern crate time;

use std::env;
use std::error::Error;
use std::ffi::{OsStr, OsString};
use std::os::unix::ffi::OsStrExt;
use std::io;
use std::path::Path;
use std::str::FromStr;
use std::thread;

use chan_signal::Signal;
use clap::{App, Arg};
use env_logger::LogBuilder;
use log::LogRecord;

mod catfs;
mod flags;
mod evicter;

use catfs::error;
use catfs::flags::{DiskSpace, FlagStorage};
use catfs::rlibc;

fn main() {
    if let Err(e) = main_internal() {
        error!("Cannot mount: {}", e);
        std::process::exit(1);
    }
}

fn main_internal() -> error::Result<()> {
    let format = |record: &LogRecord| {
        let t = time::now();
        format!(
            "{} {:5} - {}",
            time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
            record.level(),
            record.args()
        )
    };

    let mut builder = LogBuilder::new();
    builder.format(format);

    if env::var("RUST_LOG").is_ok() {
        builder.parse(&env::var("RUST_LOG").unwrap());
    }

    builder.init().unwrap();

    let mut flags: FlagStorage = Default::default();
    let mut test = false;

    let app = App::new("catfs")
        .about("Cache Anything FileSystem")
        .version(crate_version!());

    {
        fn diskspace_validator(s: String) -> Result<(), String> {
            DiskSpace::from_str(&s).map(|_| ()).map_err(
                |e| e.to_str().to_owned(),
            )
        }

        fn path_validator(s: String) -> Result<(), String> {
            Path::new(&s)
                .canonicalize()
                .map_err(|e| e.description().to_owned())
                .and_then(|p| if p.is_dir() {
                    Ok(())
                } else {
                    Err("is not a directory".to_owned())
                })
        }

        let mut args = [
            flags::Flag {
                arg: Arg::with_name("space")
                    .long("free")
                    .takes_value(true)
                    .help(
                        "Ensure filesystem has at least this much free space. (ex: 9.5%, 10G)",
                    )
                    .validator(diskspace_validator),
                value: &mut flags.free_space,
            },
            flags::Flag {
                arg: Arg::with_name("foreground").short("f").help(
                    "Run catfs in foreground.",
                ),
                value: &mut flags.foreground,
            },
            flags::Flag {
                arg: Arg::with_name("option")
                    .short("o")
                    .takes_value(true)
                    .multiple(true)
                    .help("Additional system-specific mount options. Be careful!"),
                value: &mut flags.mount_options,
            },
            flags::Flag {
                arg: Arg::with_name("test").long("test").help(
                    "Exit after parsing arguments",
                ),
                value: &mut test,
            },
            flags::Flag {
                arg: Arg::with_name("from")
                    .index(1)
                    .required(true)
                    .help("Cache files from this directory.")
                    .validator(path_validator),
                value: &mut flags.cat_from,
            },
            flags::Flag {
                arg: Arg::with_name("to")
                    .index(2)
                    .required(true)
                    .help("Cache files to this directory.")
                    .validator(path_validator),
                value: &mut flags.cat_to,
            },
            flags::Flag {
                arg: Arg::with_name("mountpoint")
                    .index(3)
                    .required(true)
                    .help("Expose the mount point at this directory.")
                    .validator(path_validator),
                value: &mut flags.mount_point,
            },
        ];

        flags::parse_options(app, &mut args);
    }

    if test {
        return Ok(());
    }

    flags.mount_options.push(OsString::from("-o"));
    flags.mount_options.push(
        OsString::from("default_permissions"),
    );

    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);
    let path_from = Path::new(&flags.cat_from).canonicalize()?;
    let path_to = Path::new(&flags.cat_to).canonicalize()?;
    let fs = catfs::CatFS::new(&path_from, &path_to)?;
    let cache_dir = fs.get_cache_dir()?;
    let mut options: Vec<&OsStr> = Vec::new();
    for i in 0..flags.mount_options.len() {
        options.push(&flags.mount_options[i]);
    }

    {
        let mut session = fuse::Session::new(fs, Path::new(&flags.mount_point), &options)?;
        thread::spawn(move || {
            if let Err(e) = session.run() {
                error!("session.run() = {}", e);
            }
            info!("{:?} unmounted", session.mountpoint());
            unsafe { libc::kill(libc::getpid(), libc::SIGTERM) };
        });

        let mut ev = evicter::Evicter::new(cache_dir, &flags.free_space);
        ev.run();
        // unmount after we get signaled becausep session will go out of scope
        let s = signal.recv().unwrap();
        info!(
            "Received {:?}, attempting to unmount {:?}",
            s,
            flags.mount_point
        );
        unmount(Path::new(&flags.mount_point))?;
    }
    rlibc::close(cache_dir)?;
    return Ok(());
}

use libc::{c_char, c_int};
use std::ffi::{CString, CStr};
/// Unmount an arbitrary mount point
pub fn unmount(mountpoint: &Path) -> io::Result<()> {
    // fuse_unmount_compat22 unfortunately doesn't return a status. Additionally,
    // it attempts to call realpath, which in turn calls into the filesystem. So
    // if the filesystem returns an error, the unmount does not take place, with
    // no indication of the error available to the caller. So we call unmount
    // directly, which is what osxfuse does anyway, since we already converted
    // to the real path when we first mounted.

    #[cfg(any(target_os = "macos", target_os = "freebsd", target_os = "dragonfly",
                target_os = "openbsd", target_os = "bitrig", target_os = "netbsd"))]
    #[inline]
    fn libc_umount(mnt: &CStr) -> c_int {
        unsafe { libc::unmount(mnt.as_ptr(), 0) }
    }

    #[cfg(not(any(target_os = "macos", target_os = "freebsd", target_os = "dragonfly",
                      target_os = "openbsd", target_os = "bitrig", target_os = "netbsd")))]
    #[inline]
    fn libc_umount(mnt: &CStr) -> c_int {
        use std::io::ErrorKind::PermissionDenied;

        let rc = unsafe { libc::umount(mnt.as_ptr()) };
        if rc < 0 && io::Error::last_os_error().kind() == PermissionDenied {
            // Linux always returns EPERM for non-root users.  We have to let the
            // library go through the setuid-root "fusermount -u" to unmount.
            unsafe {
                fuse_unmount_compat22(mnt.as_ptr());
            }
            0
        } else {
            rc
        }
    }

    let mnt = try!(CString::new(mountpoint.as_os_str().as_bytes()));
    let rc = libc_umount(&mnt);
    if rc < 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

extern "system" {
    pub fn fuse_unmount_compat22(mountpoint: *const c_char);
}
<|endoftext|>


<commit_before>extern crate argparse;
extern crate glob;
extern crate regex;

use argparse::{ArgumentParser, Store, StoreOption, Print};
use std::io::BufReader;
use std::io::BufRead;
use std::fs::File;
use std::io::Write;
use glob::glob;
use regex::Regex;

fn identity<T>(x: T) -> T {x}

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

struct Options {
    path: String,
    file: Option<String>
}

fn get_possible_files_from_glob() -> Result<Vec<std::path::PathBuf>, glob::PatternError> {
    glob("**/*").map(|paths| paths.flat_map(identity).filter(|path| path.is_file()).collect())
}

fn get_filename_minus_extension(path_str: &String) -> String {
    std::path::Path::new(path_str).file_stem().unwrap().to_str().unwrap().to_string()
}

fn is_test_file(path: &String) -> bool {
    let re = Regex::new(r"^(features/|test/|spec/|tests/)").unwrap();
    re.is_match(path.as_str())
}

fn strip_test_words(filename: &String) -> String {
    let re = Regex::new(r"(test_)?(?P<p>\w+?)(_rake_spec|_spec|_test|_steps)?(\.rb|\.exs|\.ex|\.js|\.py)?$").unwrap();
    re.replace_all(filename.as_str(), "$p")
}

fn cleanse_path(path: &str) -> String {
    let s = path.to_string();
    if s.len() > 1 && s[0..2].to_string() == "./" {
        s[2..].to_string()
    } else {
        s
    }
}

fn find_longest_common_substring_length(s1: &String, s2: &String) -> i32 {
    // Currently this is implemented using a dynamic programming solution similar
    // to http://www.geeksforgeeks.org/longest-common-substring/. This is O(N*M)
    // where N is the length of one string and M is the length of the other
    // string.
    //
    // Another option would of course be to explore using something like a
    // suffix tree to solve this problem, something like, the following.
    // http://www.geeksforgeeks.org/suffix-tree-application-5-longest-common-substring-2/
    // This is O(M+N) to build a Generalized Suffix Tree and O(M+N) to find the
    // the longest common substring via depth first search.
    //
    // Beyond that we would have to explore not caring about longest substring
    // and moving to a similarity ranking algorithm that maybe cares about
    // subsequences rather that substrings, etc.
    if s1.is_empty() || s2.is_empty() {
        return 0
    }

    let mut m: Vec<Vec<i32>> = Vec::with_capacity(s1.len());
    for _ in 0..s1.len() {
        let v: Vec<i32> = vec![0; s2.len()];
        m.push(v);
    }

    let mut longest_length = 0;

    let s1_bytes = s1.as_bytes();
    let s2_bytes = s2.as_bytes();

    for i in 0..s1.len() {
        for j in 0..s2.len() {
            if s1_bytes[i] == s2_bytes[j] {
                m[i][j] = 1;
                if i > 0 && j > 0 {
                    m[i][j] += m[i-1][j-1];
                }
                if m[i][j] > longest_length {
                    longest_length = m[i][j];
                }
            }
        }
    }

    longest_length
}

fn score(s1: &String, s2: &String) -> f32 {
    let longest_match_length: f32 = find_longest_common_substring_length(s1, s2) as f32;
    (longest_match_length/s2.len() as f32) * (longest_match_length/s1.len() as f32)
}

fn find_alt(filename: &String, cleansed_path: &String, paths: Vec<String>, test_file: bool) -> String {
    let (_, alternate) = paths.iter()
        .map(|path| cleanse_path(&path))
        .filter(|path| path.contains(filename.as_str()))  // filter to paths that contain the filename
        .filter(|path| is_test_file(&path) != test_file)
        .fold((0 as f32, "".to_string()), |result, path| {
            let (highest_score, best_match) = result;
            let s = score(&path, &cleansed_path);
            if s > highest_score {
                (s, path)
            } else {
                (highest_score, best_match)
            }
        });
    alternate
}

fn main() {
    let mut options = Options {
        path: "".to_string(),
        file: None
    };

    { // block limits of borrows by refer() method calls
        let mut ap = ArgumentParser::new();
        ap.add_option(&["-v", "--version"], Print(env!("CARGO_PKG_VERSION").to_string()), "show version");
        ap.refer(&mut options.file).add_option(&["-f", "--file"], StoreOption, "possible alternates file, - for stdin");
        ap.refer(&mut options.path).add_argument("PATH", Store, "path to find alternate for").required();
        ap.parse_args_or_exit();
    }

    let cleansed_path = cleanse_path(&options.path);
    let mut filename = get_filename_minus_extension(&options.path);
    if is_test_file(&cleansed_path) {
        filename = strip_test_words(&filename);
    }

    let best_match = if let Some(unwrapped_file) = options.file {
        if unwrapped_file == "-" {
            let stdin = std::io::stdin();
            let paths: Vec<String> = stdin.lock().lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        } else {
            let f = match File::open(&unwrapped_file) {
                Ok(file) => file,
                Err(e) => {
                    printerr!("Failure occurred opening file {}, {}", &unwrapped_file, e);
                    std::process::exit(1)
                }
            };
            let file = BufReader::new(&f);
            let paths: Vec<String> = file.lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        }
    } else {
        match get_possible_files_from_glob() {
            Ok(paths) => {
                let unwrapped_paths:Vec<String> = paths.iter().map(|path| { path.to_str().unwrap().to_string() }).collect();
                find_alt(&filename, &cleansed_path, unwrapped_paths, is_test_file(&cleansed_path))
            },
            Err(e) => {
                printerr!("Error reading paths {}", e);
                std::process::exit(1)
            }
        }
    };
    print!("{}", best_match);
}
<commit_msg>Add support for Swift/iOS/Swift Package Manager<commit_after>extern crate argparse;
extern crate glob;
extern crate regex;

use argparse::{ArgumentParser, Store, StoreOption, Print};
use std::io::BufReader;
use std::io::BufRead;
use std::fs::File;
use std::io::Write;
use glob::glob;
use regex::Regex;

fn identity<T>(x: T) -> T {x}

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

struct Options {
    path: String,
    file: Option<String>
}

fn get_possible_files_from_glob() -> Result<Vec<std::path::PathBuf>, glob::PatternError> {
    glob("**/*").map(|paths| paths.flat_map(identity).filter(|path| path.is_file()).collect())
}

fn get_filename_minus_extension(path_str: &String) -> String {
    std::path::Path::new(path_str).file_stem().unwrap().to_str().unwrap().to_string()
}

fn is_test_file(path: &String) -> bool {
    let re = Regex::new(r"^(features/|test/|spec/|tests/|\w*Tests/)").unwrap();
    re.is_match(path.as_str())
}

fn strip_test_words(filename: &String) -> String {
    let re = Regex::new(r"(test_)?(?P<p>\w+?)(_rake_spec|_spec|_test|_steps|Tests|UITests|Specs|UISpecs)?(\.rb|\.exs|\.ex|\.js|\.py|\.swift)?$").unwrap();
    re.replace_all(filename.as_str(), "$p")
}

fn cleanse_path(path: &str) -> String {
    let s = path.to_string();
    if s.len() > 1 && s[0..2].to_string() == "./" {
        s[2..].to_string()
    } else {
        s
    }
}

fn find_longest_common_substring_length(s1: &String, s2: &String) -> i32 {
    // Currently this is implemented using a dynamic programming solution similar
    // to http://www.geeksforgeeks.org/longest-common-substring/. This is O(N*M)
    // where N is the length of one string and M is the length of the other
    // string.
    //
    // Another option would of course be to explore using something like a
    // suffix tree to solve this problem, something like, the following.
    // http://www.geeksforgeeks.org/suffix-tree-application-5-longest-common-substring-2/
    // This is O(M+N) to build a Generalized Suffix Tree and O(M+N) to find the
    // the longest common substring via depth first search.
    //
    // Beyond that we would have to explore not caring about longest substring
    // and moving to a similarity ranking algorithm that maybe cares about
    // subsequences rather that substrings, etc.
    if s1.is_empty() || s2.is_empty() {
        return 0
    }

    let mut m: Vec<Vec<i32>> = Vec::with_capacity(s1.len());
    for _ in 0..s1.len() {
        let v: Vec<i32> = vec![0; s2.len()];
        m.push(v);
    }

    let mut longest_length = 0;

    let s1_bytes = s1.as_bytes();
    let s2_bytes = s2.as_bytes();

    for i in 0..s1.len() {
        for j in 0..s2.len() {
            if s1_bytes[i] == s2_bytes[j] {
                m[i][j] = 1;
                if i > 0 && j > 0 {
                    m[i][j] += m[i-1][j-1];
                }
                if m[i][j] > longest_length {
                    longest_length = m[i][j];
                }
            }
        }
    }

    longest_length
}

fn score(s1: &String, s2: &String) -> f32 {
    let longest_match_length: f32 = find_longest_common_substring_length(s1, s2) as f32;
    (longest_match_length/s2.len() as f32) * (longest_match_length/s1.len() as f32)
}

fn find_alt(filename: &String, cleansed_path: &String, paths: Vec<String>, test_file: bool) -> String {
    let (_, alternate) = paths.iter()
        .map(|path| cleanse_path(&path))
        .filter(|path| path.contains(filename.as_str()))  // filter to paths that contain the filename
        .filter(|path| is_test_file(&path) != test_file)
        .fold((0 as f32, "".to_string()), |result, path| {
            let (highest_score, best_match) = result;
            let s = score(&path, &cleansed_path);
            if s > highest_score {
                (s, path)
            } else {
                (highest_score, best_match)
            }
        });
    alternate
}

fn main() {
    let mut options = Options {
        path: "".to_string(),
        file: None
    };

    { // block limits of borrows by refer() method calls
        let mut ap = ArgumentParser::new();
        ap.add_option(&["-v", "--version"], Print(env!("CARGO_PKG_VERSION").to_string()), "show version");
        ap.refer(&mut options.file).add_option(&["-f", "--file"], StoreOption, "possible alternates file, - for stdin");
        ap.refer(&mut options.path).add_argument("PATH", Store, "path to find alternate for").required();
        ap.parse_args_or_exit();
    }

    let cleansed_path = cleanse_path(&options.path);
    let mut filename = get_filename_minus_extension(&options.path);
    if is_test_file(&cleansed_path) {
        filename = strip_test_words(&filename);
    }

    let best_match = if let Some(unwrapped_file) = options.file {
        if unwrapped_file == "-" {
            let stdin = std::io::stdin();
            let paths: Vec<String> = stdin.lock().lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        } else {
            let f = match File::open(&unwrapped_file) {
                Ok(file) => file,
                Err(e) => {
                    printerr!("Failure occurred opening file {}, {}", &unwrapped_file, e);
                    std::process::exit(1)
                }
            };
            let file = BufReader::new(&f);
            let paths: Vec<String> = file.lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        }
    } else {
        match get_possible_files_from_glob() {
            Ok(paths) => {
                let unwrapped_paths:Vec<String> = paths.iter().map(|path| { path.to_str().unwrap().to_string() }).collect();
                find_alt(&filename, &cleansed_path, unwrapped_paths, is_test_file(&cleansed_path))
            },
            Err(e) => {
                printerr!("Error reading paths {}", e);
                std::process::exit(1)
            }
        }
    };
    print!("{}", best_match);
}
<|endoftext|>