<commit_before>// Copyright (C) 2016 Mark Lee

// 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 env_logger;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate redis;
extern crate regex;
extern crate rquery;
extern crate semver;
extern crate serde;
extern crate serde_json;

use hyper::{Client as HTTPClient, Get, Server};
use hyper::header::{Accept, ContentType};
use hyper::server::{Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use redis::Commands;
use regex::Regex;
use rquery::Document;
use semver::Version;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io;
use std::io::Read;

const DEFAULT_SERVER_PORT: u32 = 9000;
const MIN_BUNDLER_VERSION: &'static str = "1.11.0";
const RUBY_LANGPACK_RELEASES_URL: &'static str = "https://github.\
                                                  com/heroku/heroku-buildpack-ruby/releases.atom";

fn download_url(http: &HTTPClient, url: &str) -> io::Result<String> {
    let mut resp = http.get(url).send().expect("Could not send HTTP request");
    let mut body = String::new();
    match resp.read_to_string(&mut body) {
        Err(error) => Err(error),
        _ => Ok(body),
    }
}

fn download_latest_buildpack_release(http: &HTTPClient) -> String {
    let xml = download_url(&http, RUBY_LANGPACK_RELEASES_URL)
                  .expect("Could not download Atom releases!");
    let atom_doc = Document::new_from_xml_string(&xml[..]).expect("Could not parse Atom releases!");
    let latest_tag = atom_doc.select("entry title")
                             .expect("Could not find latest buildpack tag!");
    latest_tag.text().clone()
}

fn redis_cache_value(redis: &redis::Connection, key: &str, value: &str) -> String {
    let set_result: redis::RedisResult<()> = redis.set_ex(key, value, 3600);
    if let Err(_) = set_result {
        warn!("Cannot set {} in Redis, ignoring", key);
    }
    value.to_owned()
}

fn redis_cache_hash_value(redis: &redis::Connection, hash_name: &str, key: &str, value: &str) {
    let set_result: redis::RedisResult<()> = redis.hset_nx(hash_name, key, value);
    if let Err(_) = set_result {
        warn!("Cannot set {} in Redis, ignoring", key);
    }
}

fn latest_buildpack_release(http: &HTTPClient, maybe_redis: &Option<redis::Connection>) -> String {
    if let Some(ref redis) = *maybe_redis {
        if let Ok(buildpack_release) = redis.get("latest_buildpack_release") {
            buildpack_release
        } else {
            let buildpack_release = download_latest_buildpack_release(http);
            redis_cache_value(redis, "latest_buildpack_release", &buildpack_release[..])
        }
    } else {
        download_latest_buildpack_release(http)
    }
}

fn cached_version_from_buildpack_release(buildpack_release: &str,
                                         maybe_redis: &Option<redis::Connection>)
                                         -> Option<String> {
    if let Some(ref redis) = *maybe_redis {
        if let Ok(version) = redis.hget("bundler_version", buildpack_release) {
            Some(version)
        } else {
            None
        }
    } else {
        None
    }
}

fn bundler_version_from_ruby_buildpack(http: &HTTPClient,
                                       maybe_redis: Option<redis::Connection>)
                                       -> Option<String> {
    let buildpack_release = latest_buildpack_release(&http, &maybe_redis);
    if let Some(cached_version) = cached_version_from_buildpack_release(&buildpack_release[..],
                                                                        &maybe_redis) {
        return Some(cached_version);
    }
    let ruby_langpack_url = format!("https://raw.githubusercontent.com/\
                                     heroku/heroku-buildpack-ruby/{}/lib/language_pack/\
                                     ruby.rb",
                                    buildpack_release);
    let ruby_file = &download_url(&http, &ruby_langpack_url[..])
                         .expect("Could not download ruby.rb!")[..];
    let regex = Regex::new(r#"BUNDLER_VERSION += "(.+?)""#).expect("Invalid regular expression!");
    if regex.is_match(ruby_file) {
        let captures = regex.captures(ruby_file).expect("Could not match?!");
        let version = captures.at(1).expect("Capture not found?!");
        if let Some(redis) = maybe_redis {
            redis_cache_hash_value(&redis, "bundler_version", &buildpack_release[..], version);
            redis_cache_value(&redis, "latest_supported_bundler_version", version);
        }
        Some(version.to_owned())
    } else {
        None
    }
}

fn is_bundler_upgraded(http: &HTTPClient,
                       redis_result: redis::RedisResult<redis::Connection>)
                       -> bool {
    let bundler_version_result = if let Ok(redis) = redis_result {
        match redis.get("latest_supported_bundler_version") {
            Ok(version) => Some(version),
            _ => bundler_version_from_ruby_buildpack(&http, Some(redis)),
        }
    } else {
        bundler_version_from_ruby_buildpack(&http, None)
    };
    if let Some(buildpack_bundler_version_str) = bundler_version_result {
        let min_version = Version::parse(MIN_BUNDLER_VERSION)
                              .expect("Could not parse min version!");
        let new_version = Version::parse(&buildpack_bundler_version_str[..])
                              .expect("Could not parse new version!");
        min_version < new_version
    } else {
        false
    }
}

fn determine_content_type(req: &Request) -> ContentType {
    if let Some(accept) = req.headers.get::<Accept>() {
        let mut use_json = false;
        for qitem in &accept.0 {
            let qitem_str = &format!("{}", qitem)[..];
            if qitem_str == "application/json" {
                use_json = true;
                break;
            }
        }

        if use_json {
            ContentType::json()
        } else {
            ContentType::html()
        }
    } else {
        ContentType::html()
    }
}

fn result_to_json(result: bool) -> String {
    let mut map = HashMap::new();
    map.insert("result".to_owned(), result);
    serde_json::to_string(&map).expect("Could not serialize result!")
}

fn result_to_html(result: bool) -> String {
    let result_str = if result {
        r#"
Yes
"#
    } else {
        r#"
No
"#
    };
    match File::open("index.html") {
        Ok(mut html_file) => {
            let mut html = String::new();
            html_file.read_to_string(&mut html).expect("Could not read HTML file!");
            html.replace("{{ is_bundler_upgraded }}", result_str)
                .replace("{{ MIN_BUNDLER_VERSION }}", MIN_BUNDLER_VERSION)
        }
        Err(error) => format!("HTML NOT FOUND: {}", error),
    }
}

fn server_port() -> String {
    match env::var("PORT") {
        Ok(port) => port,
        _ => format!("{}", DEFAULT_SERVER_PORT),
    }
}

fn connect_to_redis() -> redis::RedisResult<redis::Connection> {
    if let Ok(redis_url) = env::var("REDIS_URL") {
        let client = redis::Client::open(&redis_url[..]).expect("Cannot connect to Redis");
        client.get_connection()
    } else {
        Err(redis::RedisError::from(io::Error::new(io::ErrorKind::Other, "REDIS_URL not found")))
    }
}

fn request_handler(req: Request, mut res: Response) {
    if let AbsolutePath(ref path) = req.uri {
        if *path == "/" {
            if req.method == Get {
                let http = HTTPClient::new();
                let redis = connect_to_redis();
                let result = is_bundler_upgraded(&http, redis);
                let content_type = determine_content_type(&req);
                let data = match &format!("{}", content_type)[..] {
                    "application/json; charset=utf-8" => result_to_json(result),
                    _ => result_to_html(result),
                };
                res.headers_mut().set(content_type);
                res.send(data.as_bytes()).expect("Could not set response body!");
            } else {
                *res.status_mut() = hyper::status::StatusCode::MethodNotAllowed;
            }
        } else {
            *res.status_mut() = hyper::NotFound;
        }
    } else {
        *res.status_mut() = hyper::BadRequest;
    }
}

fn main() {
    env_logger::init().expect("Could not initialize env_logger!");
    let server = Server::http(&format!("0.0.0.0:{}", server_port())[..])
                     .expect("Could not create server!");
    server.handle(request_handler).expect("Could not set up HTTP request handler!");
}
<commit_msg>Remove redundant Redis key; use RedisResult instead of Option<commit_after>// Copyright (C) 2016 Mark Lee

// 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 env_logger;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate redis;
extern crate regex;
extern crate rquery;
extern crate semver;
extern crate serde;
extern crate serde_json;

use hyper::{Client as HTTPClient, Get, Server};
use hyper::header::{Accept, ContentType};
use hyper::server::{Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use redis::{Commands, RedisError, RedisResult};
use regex::Regex;
use rquery::Document;
use semver::Version;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io;
use std::io::Read;

const DEFAULT_SERVER_PORT: u32 = 9000;
const MIN_BUNDLER_VERSION: &'static str = "1.11.0";
const RUBY_LANGPACK_RELEASES_URL: &'static str = "https://github.\
                                                  com/heroku/heroku-buildpack-ruby/releases.atom";

fn download_url(http: &HTTPClient, url: &str) -> io::Result<String> {
    let mut resp = http.get(url).send().expect("Could not send HTTP request");
    let mut body = String::new();
    match resp.read_to_string(&mut body) {
        Err(error) => Err(error),
        _ => Ok(body),
    }
}

fn download_latest_buildpack_release(http: &HTTPClient) -> String {
    let xml = download_url(&http, RUBY_LANGPACK_RELEASES_URL)
                  .expect("Could not download Atom releases!");
    let atom_doc = Document::new_from_xml_string(&xml[..]).expect("Could not parse Atom releases!");
    let latest_tag = atom_doc.select("entry title")
                             .expect("Could not find latest buildpack tag!");
    latest_tag.text().clone()
}

fn redis_cache_value(redis: &redis::Connection, key: &str, value: &str) -> String {
    let set_result: RedisResult<()> = redis.set_ex(key, value, 3600);
    if let Err(_) = set_result {
        warn!("Cannot set {} in Redis, ignoring", key);
    }
    value.to_owned()
}

fn redis_cache_hash_value(redis: &redis::Connection, hash_name: &str, key: &str, value: &str) {
    let set_result: RedisResult<()> = redis.hset_nx(hash_name, key, value);
    if let Err(_) = set_result {
        warn!("Cannot set {} in Redis, ignoring", key);
    }
}

fn latest_buildpack_release(http: &HTTPClient, redis_result: &RedisResult<redis::Connection>) -> String {
    if let Ok(ref redis) = *redis_result {
        if let Ok(buildpack_release) = redis.get("latest_buildpack_release") {
            buildpack_release
        } else {
            let buildpack_release = download_latest_buildpack_release(http);
            redis_cache_value(redis, "latest_buildpack_release", &buildpack_release[..])
        }
    } else {
        download_latest_buildpack_release(http)
    }
}

fn cached_version_from_buildpack_release(buildpack_release: &str,
                                         redis_result: &RedisResult<redis::Connection>)
                                         -> Option<String> {
    if let Ok(ref redis) = *redis_result {
        if let Ok(version) = redis.hget("bundler_version", buildpack_release) {
            Some(version)
        } else {
            None
        }
    } else {
        None
    }
}

fn bundler_version_from_ruby_buildpack(http: &HTTPClient,
                                       redis_result: &RedisResult<redis::Connection>)
                                       -> Option<String> {
    let buildpack_release = latest_buildpack_release(http, redis_result);
    if let Some(cached_version) = cached_version_from_buildpack_release(&buildpack_release[..],
                                                                        redis_result) {
        return Some(cached_version);
    }
    let ruby_langpack_url = format!("https://raw.githubusercontent.com/\
                                     heroku/heroku-buildpack-ruby/{}/lib/language_pack/\
                                     ruby.rb",
                                    buildpack_release);
    let ruby_file = &download_url(&http, &ruby_langpack_url[..])
                         .expect("Could not download ruby.rb!")[..];
    let regex = Regex::new(r#"BUNDLER_VERSION += "(.+?)""#).expect("Invalid regular expression!");
    if regex.is_match(ruby_file) {
        let captures = regex.captures(ruby_file).expect("Could not match?!");
        let version = captures.at(1).expect("Capture not found?!");
        if let Ok(ref redis) = *redis_result {
            redis_cache_hash_value(&redis, "bundler_version", &buildpack_release[..], version);
        }
        Some(version.to_owned())
    } else {
        None
    }
}

fn is_bundler_upgraded(http: &HTTPClient,
                       redis_result: RedisResult<redis::Connection>)
                       -> bool {
    let bundler_version_result = bundler_version_from_ruby_buildpack(&http, &redis_result);
    if let Some(buildpack_bundler_version_str) = bundler_version_result {
        let min_version = Version::parse(MIN_BUNDLER_VERSION)
                              .expect("Could not parse min version!");
        let new_version = Version::parse(&buildpack_bundler_version_str[..])
                              .expect("Could not parse new version!");
        min_version < new_version
    } else {
        false
    }
}

fn determine_content_type(req: &Request) -> ContentType {
    if let Some(accept) = req.headers.get::<Accept>() {
        let mut use_json = false;
        for qitem in &accept.0 {
            let qitem_str = &format!("{}", qitem)[..];
            if qitem_str == "application/json" {
                use_json = true;
                break;
            }
        }

        if use_json {
            ContentType::json()
        } else {
            ContentType::html()
        }
    } else {
        ContentType::html()
    }
}

fn result_to_json(result: bool) -> String {
    let mut map = HashMap::new();
    map.insert("result".to_owned(), result);
    serde_json::to_string(&map).expect("Could not serialize result!")
}

fn result_to_html(result: bool) -> String {
    let result_str = if result {
        r#"
Yes
"#
    } else {
        r#"
No
"#
    };
    match File::open("index.html") {
        Ok(mut html_file) => {
            let mut html = String::new();
            html_file.read_to_string(&mut html).expect("Could not read HTML file!");
            html.replace("{{ is_bundler_upgraded }}", result_str)
                .replace("{{ MIN_BUNDLER_VERSION }}", MIN_BUNDLER_VERSION)
        }
        Err(error) => format!("HTML NOT FOUND: {}", error),
    }
}

fn server_port() -> String {
    match env::var("PORT") {
        Ok(port) => port,
        _ => format!("{}", DEFAULT_SERVER_PORT),
    }
}

fn connect_to_redis() -> RedisResult<redis::Connection> {
    if let Ok(redis_url) = env::var("REDIS_URL") {
        let client = redis::Client::open(&redis_url[..]).expect("Cannot connect to Redis");
        client.get_connection()
    } else {
        Err(RedisError::from(io::Error::new(io::ErrorKind::Other, "REDIS_URL not found")))
    }
}

fn request_handler(req: Request, mut res: Response) {
    if let AbsolutePath(ref path) = req.uri {
        if *path == "/" {
            if req.method == Get {
                let http = HTTPClient::new();
                let redis = connect_to_redis();
                let result = is_bundler_upgraded(&http, redis);
                let content_type = determine_content_type(&req);
                let data = match &format!("{}", content_type)[..] {
                    "application/json; charset=utf-8" => result_to_json(result),
                    _ => result_to_html(result),
                };
                res.headers_mut().set(content_type);
                res.send(data.as_bytes()).expect("Could not set response body!");
            } else {
                *res.status_mut() = hyper::status::StatusCode::MethodNotAllowed;
            }
        } else {
            *res.status_mut() = hyper::NotFound;
        }
    } else {
        *res.status_mut() = hyper::BadRequest;
    }
}

fn main() {
    env_logger::init().expect("Could not initialize env_logger!");
    let server = Server::http(&format!("0.0.0.0:{}", server_port())[..])
                     .expect("Could not create server!");
    server.handle(request_handler).expect("Could not set up HTTP request handler!");
}
<|endoftext|>


<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;

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

// The version of the program.
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,
    current_value: Value,

    // `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 line is different). This is the
    // so called "MultiPV" mode.
    pv_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,
            current_value: VALUE_UNKNOWN,
            search: MultipvSearch::new(tt),
            play_when: PlayWhen::Never,
            pv_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 PV
    // to `self.queue`.
    fn queue_pv(&mut self) {
        let SearchStatus { depth,
                           value,
                           value_bound,
                           ref multipv,
                           searched_nodes,
                           duration_millis,
                           nps,
                           .. } = *self.search.status();
        let score_suffix = match value_bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &multipv[0] {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, score_suffix)),
            ("time".to_string(), format!("{}", duration_millis)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
            ("pv".to_string(), pv_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 multipv, .. } = *self.search.status();
        self.queue.push_back(EngineReply::BestMove {
            best_move: multipv[0].get(0).map_or("0000".to_string(), |m| m.notation()),
            ponder_move: multipv[0].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.pv_count = value.parse::<usize>().unwrap_or(1);
            }

            // 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.current_value = VALUE_UNKNOWN;
        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.pv_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 {
            let SearchStatus { done, depth, value, searched_nodes, duration_millis, .. } =
                *self.search.update_status();

            // Send the PV when changed.
            if depth > 0 && (depth > self.current_depth || value != self.current_value) {
                self.current_depth = depth;
                self.current_value = value;
                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>removed current_value<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;

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

// The version of the program.
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 line is different). This is the
    // so called "MultiPV" mode.
    pv_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,
            pv_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 PV
    // to `self.queue`.
    fn queue_pv(&mut self) {
        let SearchStatus { depth,
                           value,
                           value_bound,
                           ref multipv,
                           searched_nodes,
                           duration_millis,
                           nps,
                           .. } = *self.search.status();
        let score_suffix = match value_bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &multipv[0] {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, score_suffix)),
            ("time".to_string(), format!("{}", duration_millis)),
            ("nodes".to_string(), format!("{}", searched_nodes)),
            ("nps".to_string(), format!("{}", nps)),
            ("pv".to_string(), pv_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 multipv, .. } = *self.search.status();
        self.queue.push_back(EngineReply::BestMove {
            best_move: multipv[0].get(0).map_or("0000".to_string(), |m| m.notation()),
            ponder_move: multipv[0].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.pv_count = value.parse::<usize>().unwrap_or(1);
            }

            // 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.pv_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 {
            let SearchStatus { done, depth, searched_nodes, duration_millis, .. } =
                *self.search.update_status();

            // Send the 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>extern crate toml;
extern crate graphviz;

use graphviz as dot;
use std::io::File;
use std::str;

fn main() {
    let (name, indirect_deps) = read_cargo_lock("Cargo.lock");
    let direct_deps           = read_cargo_toml("Cargo.toml");
    let mut nodes    = vec!(name).append(direct_deps.as_slice());
    let mut edges    = range(1, nodes.len()).map(|n| (0, n)).collect();

    add_deps(&mut nodes, &mut edges, indirect_deps);

    let graph = Graph { nodes: nodes, edges: edges };
    let mut f    = File::create(&Path::new("Cargo.dot"));

    graph.render_to(&mut f);
}

fn read_cargo_toml(toml_file_name: &str) -> (Vec<String>) {
    let toml_str = File::open(&Path::new(toml_file_name)).read_to_string().unwrap();
    let toml     = toml::Parser::new(toml_str.as_slice()).parse().unwrap();
    let deps: Vec<String> =
        toml.find(&"dependencies".to_string()).unwrap().as_table().unwrap().keys().map(|s| s.to_string()).collect();
    deps
}

fn read_cargo_lock(_lock_file_name: &str) -> (String, Vec<(&str, Vec<&str>)>) {
    ("cargo".to_string(), vec!())
}

fn add_deps(nodes: &mut Vec<String>, edges: &mut Vec<(uint, uint)>, deps: Vec<(&str, Vec<&str>)>) {
    for (crat, crate_deps) in deps.move_iter() {
        let idl = add_or_find(nodes, crat);
        for dep in crate_deps.move_iter() {
            let idr = add_or_find(nodes, dep);
            edges.push((idl, idr));
        }
    }
}

fn add_or_find<'a>(nodes: &mut Vec<String>, new: &'a str) -> uint {
    for i in range(0, nodes.len()) {
        let ref s = (*nodes)[i];
        if s.as_slice() == new {
            return i
        }
    }
    nodes.push(new.to_string());
    nodes.len() - 1
}

type Nd = uint;
type Ed<'a> = &'a (uint, uint);
struct Graph { nodes: Vec<String>,
               edges: Vec<(uint,uint)>,
             }

impl Graph {
    fn render_to<W:Writer>(self, output: &mut W) {
        dot::render(&self, output).unwrap()
    }
}

impl<'a> dot::Labeller<'a, Nd, Ed<'a>> for Graph {
    fn graph_id(&'a self) -> dot::Id<'a> { dot::Id::new("example3") }
    fn node_id(&'a self, n: &Nd) -> dot::Id {
        dot::Id::new(format!("N{:u}", *n))
    }
    fn node_label<'a>(&'a self, i: &Nd) -> dot::LabelText<'a> {
        dot::LabelStr(str::Slice(self.nodes[*i].as_slice()))
    }
}

impl<'a> dot::GraphWalk<'a, Nd, Ed<'a>> for Graph {
    fn nodes(&'a self) -> dot::Nodes<'a,Nd> {
        range(0, self.nodes.len()).collect()
    }
    fn edges(&'a self) -> dot::Edges<'a,Ed<'a>> {
        self.edges.iter().collect()
    }
    fn source(&self, e: &Ed) -> Nd { let &(s,_) = *e; s }
    fn target(&self, e: &Ed) -> Nd { let &(_,t) = *e; t }
}
<commit_msg>Parse everything from Cargo.lock, it works!!<commit_after>extern crate toml;
extern crate graphviz;

use graphviz as dot;
use std::io::File;
use std::str;

fn main() {
    let (name, direct_deps, indirect_deps) = read_cargo_lock("Cargo.lock");
    let mut nodes = vec!(name).append(direct_deps.as_slice());
    let mut edges = range(1, nodes.len()).map(|n| (0, n)).collect();
    add_deps(&mut nodes, &mut edges, indirect_deps);

    let graph = Graph { nodes: nodes, edges: edges };
    let mut f    = File::create(&Path::new("Cargo.dot"));

    graph.render_to(&mut f);
}

fn read_cargo_lock(file_name: &str) -> (String, Vec<String>, Vec<(String, Vec<String>)>) {
    let toml = read_toml(file_name);
    let root = toml.find(&"root".to_string()).unwrap()
                   .as_table().unwrap();
    let name = root.find(&"name".to_string()).unwrap()
                   .as_str().unwrap().to_string();
    let direct_deps = extract_deps(root);
    let all_deps = toml.find(&"package".to_string()).unwrap()
                       .as_slice().unwrap().iter().map(other_deps).collect();
    (name, direct_deps, all_deps)
}

fn extract_deps(v: &toml::Table) -> Vec<String> {
    match v.find(&"dependencies".to_string()) {
        None     => vec!(),
        Some(ds) => ds.as_slice().unwrap().iter()
                      .map(|v| parse_dep(v.as_str().unwrap())).collect()
    }
}

fn other_deps(v: &toml::Value) -> (String, Vec<String>) {
    let t = v.as_table().unwrap();
    let name = t.find(&"name".to_string()).unwrap()
                .as_str().unwrap().to_string();
    let deps = extract_deps(t);
    (name, deps)
}

fn read_toml(file_name: &str) -> toml::Table {
    let toml_str = File::open(&Path::new(file_name)).read_to_string().unwrap();
    toml::Parser::new(toml_str.as_slice()).parse().unwrap()

}

fn parse_dep(s: &str) -> String {
    s.chars().take_while(|&a| a != ' ').collect()
}

fn add_deps(nodes: &mut Vec<String>, edges: &mut Vec<(uint, uint)>, deps: Vec<(String, Vec<String>)>) {
    for (crat, crate_deps) in deps.move_iter() {
        let idl = add_or_find(nodes, crat);
        for dep in crate_deps.move_iter() {
            let idr = add_or_find(nodes, dep);
            edges.push((idl, idr));
        }
    }
}

fn add_or_find(nodes: &mut Vec<String>, new: String) -> uint {
    for i in range(0, nodes.len()) {
        let ref s = (*nodes)[i];
        if *s == new {
            return i
        }
    }
    nodes.push(new);
    nodes.len() - 1
}

type Nd = uint;
type Ed<'a> = &'a (uint, uint);
struct Graph { nodes: Vec<String>,
               edges: Vec<(uint,uint)>,
             }

impl Graph {
    fn render_to<W:Writer>(self, output: &mut W) {
        dot::render(&self, output).unwrap()
    }
}

impl<'a> dot::Labeller<'a, Nd, Ed<'a>> for Graph {
    fn graph_id(&'a self) -> dot::Id<'a> { dot::Id::new("example3") }
    fn node_id(&'a self, n: &Nd) -> dot::Id {
        dot::Id::new(format!("N{:u}", *n))
    }
    fn node_label<'a>(&'a self, i: &Nd) -> dot::LabelText<'a> {
        dot::LabelStr(str::Slice(self.nodes[*i].as_slice()))
    }
}

impl<'a> dot::GraphWalk<'a, Nd, Ed<'a>> for Graph {
    fn nodes(&'a self) -> dot::Nodes<'a,Nd> {
        range(0, self.nodes.len()).collect()
    }
    fn edges(&'a self) -> dot::Edges<'a,Ed<'a>> {
        self.edges.iter().collect()
    }
    fn source(&self, e: &Ed) -> Nd { let &(s,_) = *e; s }
    fn target(&self, e: &Ed) -> Nd { let &(_,t) = *e; t }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate rand;
extern crate ultra;

use rand::Rng;
use ultra::{Enigma, decrypt};

fn main() {
    let app = clap_app!(ultra =>
        (version: crate_version!())
        (@setting ArgRequiredElseHelp)
        (@setting ColoredHelp)
        (@arg decrypt: --decrypt -d "Decrypt a given piece of ciphertext")
        (@arg randomize: --randomize -R "Encrypt a message with random Enigma settings")
        (@arg ROTORS: --rotor -w +takes_value "Rotor order (default: \"123\")")
        (@arg KEY: --key -k +takes_value "Key settings (default: \"AAA\")")
        (@arg RING: --ring -r +takes_value "Ring settings (default: \"AAA\")")
        (@arg PLUGBOARD: --plugboard -p +takes_value "Plugboard settings (default: \"\")")
        (@arg MESSAGE: +required "The message to encrypt/decrypt")
    );

    let matches = app.get_matches();
    let msg = matches.value_of("MESSAGE").unwrap();

    if matches.is_present("decrypt") {
        let (plaintext, rotors, key, ring) = decrypt(msg);
        println!("{}", plaintext);
        println!("(Rotors: {}, Key Setting: {}, Ring Setting: {})", rotors, key, ring);
    }

    else if matches.is_present("randomize") {
        let mut rng = rand::thread_rng();

        let rotors: String = {
            let mut rotor_pool: Vec<char> = "12345".chars().collect();
            let mut rotors = Vec::with_capacity(3);

            for _ in 0..3 {
                let len = rotor_pool.len();
                rotors.push(rotor_pool.remove(rng.gen_range(0, len)));
            }

            rotors.into_iter().collect()
        };

        let mut key = String::with_capacity(3);
        let mut ring = String::with_capacity(3);
        let alpha: Vec<char> = "ABCDEFGHIJKLMNOPQRSTUVWXYZ".chars().collect();

        for _ in 0..3 {
            key.push(*rng.choose(&alpha).unwrap());
            ring.push(*rng.choose(&alpha).unwrap());
        }

        let mut enigma = Enigma::new(&rotors, &key, &ring, 'B', "");
        println!("{}", enigma.encrypt(msg));
        println!("(Rotors: {}, Key Setting: {}, Ring Setting: {})", &rotors, &key, &ring);
    }

    else {
        let rotors = matches.value_of("ROTORS").unwrap_or("123");
        let key = matches.value_of("KEY").unwrap_or("AAA");
        let ring = matches.value_of("RING").unwrap_or("AAA");
        let plugboard = matches.value_of("PLUGBOARD").unwrap_or("");

        let mut enigma = Enigma::new(rotors, key, ring, 'B', plugboard);
        println!("{}", enigma.encrypt(msg));
    }
}
<commit_msg>Match output case to input message in CLI<commit_after>#[macro_use]
extern crate clap;
extern crate rand;
extern crate ultra;

use std::ascii::AsciiExt;

use rand::Rng;
use ultra::{Enigma, decrypt};

trait CasedString {
    fn with_case_of(&self, target: &str) -> String;
}

impl CasedString for String {
    fn with_case_of(&self, target: &str) -> String {
        self.chars().zip(target.chars()).map(|(s, t)| {
            match t.is_lowercase() {
                true => s.to_ascii_lowercase(),
                false => s
            }
        }).collect()
    }
}

fn main() {
    let app = clap_app!(ultra =>
        (version: crate_version!())
        (@setting ArgRequiredElseHelp)
        (@setting ColoredHelp)
        (@arg decrypt: --decrypt -d "Decrypt a given piece of ciphertext")
        (@arg randomize: --randomize -R "Encrypt a message with random Enigma settings")
        (@arg ROTORS: --rotor -w +takes_value "Rotor order (default: \"123\")")
        (@arg KEY: --key -k +takes_value "Key settings (default: \"AAA\")")
        (@arg RING: --ring -r +takes_value "Ring settings (default: \"AAA\")")
        (@arg PLUGBOARD: --plugboard -p +takes_value "Plugboard settings (default: \"\")")
        (@arg MESSAGE: +required "The message to encrypt/decrypt")
    );

    let matches = app.get_matches();
    let msg = matches.value_of("MESSAGE").unwrap();

    if matches.is_present("decrypt") {
        let (plaintext, rotors, key, ring) = decrypt(msg);
        println!("{}", plaintext.with_case_of(msg));
        println!("(Rotors: {}, Key Setting: {}, Ring Setting: {})", rotors, key, ring);
    }

    else if matches.is_present("randomize") {
        let mut rng = rand::thread_rng();

        let rotors: String = {
            let mut rotor_pool: Vec<char> = "12345".chars().collect();
            let mut rotors = Vec::with_capacity(3);

            for _ in 0..3 {
                let len = rotor_pool.len();
                rotors.push(rotor_pool.remove(rng.gen_range(0, len)));
            }

            rotors.into_iter().collect()
        };

        let mut key = String::with_capacity(3);
        let mut ring = String::with_capacity(3);
        let alpha: Vec<char> = "ABCDEFGHIJKLMNOPQRSTUVWXYZ".chars().collect();

        for _ in 0..3 {
            key.push(*rng.choose(&alpha).unwrap());
            ring.push(*rng.choose(&alpha).unwrap());
        }

        let mut enigma = Enigma::new(&rotors, &key, &ring, 'B', "");
        println!("{}", enigma.encrypt(msg).with_case_of(msg));
        println!("(Rotors: {}, Key Setting: {}, Ring Setting: {})", &rotors, &key, &ring);
    }

    else {
        let rotors = matches.value_of("ROTORS").unwrap_or("123");
        let key = matches.value_of("KEY").unwrap_or("AAA");
        let ring = matches.value_of("RING").unwrap_or("AAA");
        let plugboard = matches.value_of("PLUGBOARD").unwrap_or("");

        let mut enigma = Enigma::new(rotors, key, ring, 'B', plugboard);
        println!("{}", enigma.encrypt(msg).with_case_of(msg));
    }
}
<|endoftext|>


<commit_before>#![feature(managed_boxes,phase)]
#[phase(syntax, link)] extern crate log;

extern crate syntax;
extern crate time;
extern crate debug;

#[cfg(not(test))]
use racer::{getline,Match,do_file_search, do_external_search};
#[cfg(not(test))]
use std::io::File;
#[cfg(not(test))]
use std::io::BufferedReader;
#[cfg(not(test))]
use std::str;
#[cfg(not(test))]
use racer::scopes;

pub mod racer;

#[cfg(not(test))]
fn match_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords2(&m.filepath, m.point).unwrap();
    println!("MATCH {},{},{},{}", m.matchstr,
                                    linenum.to_str(),
                                    charnum.to_str(),
                                    m.filepath.as_str().unwrap());
}

#[cfg(not(test))]
fn complete() {
    let args = std::os::args();
    match std::uint::parse_bytes(std::os::args().as_slice()[2].as_bytes(), 10) {
        Some(linenum) => {
            // input: linenum, colnum, fname
            let charnum = std::uint::parse_bytes(std::os::args().as_slice()[3].as_bytes(), 10).unwrap();
            let fname = args.as_slice()[4].as_slice();
            let fpath = Path::new(fname);
            let filetxt = BufferedReader::new(File::open(&fpath)).read_to_end().unwrap();
            // print the start-end of the identifier being matched
            let src = str::from_utf8(filetxt.as_slice()).unwrap();
            let line = getline(&fpath, linenum);
            let (start, pos) = racer::expand_ident(line.as_slice(), charnum);
            println!("PREFIX {},{},{}", start, pos, line.as_slice().slice(start, pos));

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

            do_file_search(p.as_slice()[0], &Path::new("."), &mut |m| {
                if p.len() == 1 {
                    match_fn(m);
                } else {
                    do_external_search(p.slice_from(1), &m.filepath, m.point, racer::StartsWith, &mut |m| match_fn(m));
                }
            });

        }
    }
}

#[cfg(not(test))]
fn prefix() {
    let args_ = std::os::args();
    let args = args_.as_slice();
    let linenum = std::uint::parse_bytes(args[2].as_bytes(), 10).unwrap();
    let charnum = std::uint::parse_bytes(args[3].as_bytes(), 10).unwrap();
    let fname = args[4].as_slice();

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

#[cfg(not(test))]
fn find_definition() {
    let args_ = std::os::args();
    let args = args_.as_slice();
    let linenum = std::uint::parse_bytes(args[2].as_bytes(), 10).unwrap();
    let charnum = std::uint::parse_bytes(args[3].as_bytes(), 10).unwrap();
    let fname = args[4].as_slice();
    let fpath = Path::new(fname);
    let filetxt = BufferedReader::new(File::open(&fpath)).read_to_end().unwrap();
    let src = str::from_utf8(filetxt.as_slice()).unwrap();
    let pos = scopes::coords_to_point(src, linenum, charnum);

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

#[cfg(not(test))]
fn print_usage() {
    let program = std::os::args().as_slice()[0].clone();
    println!("usage: {} complete linenum charnum fname", program);
    println!("or:    {} find-definition linenum charnum fname", program);
    println!("or:    {} complete fullyqualifiedname   (e.g. std::io::)",program);
    println!("or:    {} prefix linenum charnum fname",program);
}

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

    let args_ = std::os::args();
    let args = args_.as_slice();

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

    let command = args[1].as_slice();
    match command.as_slice() {
        "prefix" => prefix(),
        "complete" => complete(),
        "find-definition" => find_definition(),
        "help" => print_usage(),
        _ => {
            println!("Sorry, I didn't understand command {}", command );
            print_usage();
            std::os::set_exit_status(1);
            return;
        }
    }
}
<commit_msg>phase syntax is now called plugin<commit_after>#![feature(managed_boxes,phase)]
#[phase(plugin, link)] extern crate log;

extern crate syntax;
extern crate time;
extern crate debug;

#[cfg(not(test))]
use racer::{getline,Match,do_file_search, do_external_search};
#[cfg(not(test))]
use std::io::File;
#[cfg(not(test))]
use std::io::BufferedReader;
#[cfg(not(test))]
use std::str;
#[cfg(not(test))]
use racer::scopes;

pub mod racer;

#[cfg(not(test))]
fn match_fn(m:Match) {
    let (linenum, charnum) = scopes::point_to_coords2(&m.filepath, m.point).unwrap();
    println!("MATCH {},{},{},{}", m.matchstr,
                                    linenum.to_str(),
                                    charnum.to_str(),
                                    m.filepath.as_str().unwrap());
}

#[cfg(not(test))]
fn complete() {
    let args = std::os::args();
    match std::uint::parse_bytes(std::os::args().as_slice()[2].as_bytes(), 10) {
        Some(linenum) => {
            // input: linenum, colnum, fname
            let charnum = std::uint::parse_bytes(std::os::args().as_slice()[3].as_bytes(), 10).unwrap();
            let fname = args.as_slice()[4].as_slice();
            let fpath = Path::new(fname);
            let filetxt = BufferedReader::new(File::open(&fpath)).read_to_end().unwrap();
            // print the start-end of the identifier being matched
            let src = str::from_utf8(filetxt.as_slice()).unwrap();
            let line = getline(&fpath, linenum);
            let (start, pos) = racer::expand_ident(line.as_slice(), charnum);
            println!("PREFIX {},{},{}", start, pos, line.as_slice().slice(start, pos));

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

            do_file_search(p.as_slice()[0], &Path::new("."), &mut |m| {
                if p.len() == 1 {
                    match_fn(m);
                } else {
                    do_external_search(p.slice_from(1), &m.filepath, m.point, racer::StartsWith, &mut |m| match_fn(m));
                }
            });

        }
    }
}

#[cfg(not(test))]
fn prefix() {
    let args_ = std::os::args();
    let args = args_.as_slice();
    let linenum = std::uint::parse_bytes(args[2].as_bytes(), 10).unwrap();
    let charnum = std::uint::parse_bytes(args[3].as_bytes(), 10).unwrap();
    let fname = args[4].as_slice();

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

#[cfg(not(test))]
fn find_definition() {
    let args_ = std::os::args();
    let args = args_.as_slice();
    let linenum = std::uint::parse_bytes(args[2].as_bytes(), 10).unwrap();
    let charnum = std::uint::parse_bytes(args[3].as_bytes(), 10).unwrap();
    let fname = args[4].as_slice();
    let fpath = Path::new(fname);
    let filetxt = BufferedReader::new(File::open(&fpath)).read_to_end().unwrap();
    let src = str::from_utf8(filetxt.as_slice()).unwrap();
    let pos = scopes::coords_to_point(src, linenum, charnum);

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

#[cfg(not(test))]
fn print_usage() {
    let program = std::os::args().as_slice()[0].clone();
    println!("usage: {} complete linenum charnum fname", program);
    println!("or:    {} find-definition linenum charnum fname", program);
    println!("or:    {} complete fullyqualifiedname   (e.g. std::io::)",program);
    println!("or:    {} prefix linenum charnum fname",program);
}

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

    let args_ = std::os::args();
    let args = args_.as_slice();

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

    let command = args[1].as_slice();
    match command.as_slice() {
        "prefix" => prefix(),
        "complete" => complete(),
        "find-definition" => find_definition(),
        "help" => print_usage(),
        _ => {
            println!("Sorry, I didn't understand command {}", command );
            print_usage();
            std::os::set_exit_status(1);
            return;
        }
    }
}
<|endoftext|>


<commit_before>extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate log4rs;
#[macro_use]
extern crate lazy_static;

mod lib;

use lib::config;
use lib::db;
use lib::logger;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::DownloadError;
use lib::converter::Converter;

use std::fs::{remove_file,remove_dir_all};

static VERSION : &'static str = "0.3"; // String not valid
static CONFIG_PATH : &'static str = "config.cfg";
static LOG_CONFIG: &'static str = "log.conf";
static LOG_PATTERN: &'static str = "%d{%d-%m-%Y %H:%M:%S}\t[%l]\t%f:%L \t%m";
static SLEEP_MS: u32 = 5000;
static CODE_SUCCESS: i8 = 2;
static CODE_SUCCESS_WARNINGS: i8 = 3; // finished with warnings
static CODE_FAILED_INTERNAL: i8 = 10; // internal error
static CODE_FAILED_QUALITY: i8 = 11; // qualitz not available
static CODE_FAILED_UNAVAILABLE: i8 = 12; // source unavailable (private / removed)

lazy_static! {
    pub static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
// enum QueryCodes {
//     Waiting = 0,
//     InProgress = 1,
//     Finished = 2,
//     Failed = 3,
// }

enum Thing { String(String), Bool(bool), None }

fn main() {
    logger::initialize();

    let pool = db::db_connect(db::mysql_options(), SLEEP_MS);

    let converter = Converter::new(&CONFIG.general.ffmpeg_bin_dir,&CONFIG.general.mp3_quality , pool.clone());
    let mut print_pause = true;
    loop {
        if let Some(result) = db::request_entry(& pool) {
            print_pause = true;
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            db::set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");
            db::set_query_state(&pool.clone(),&qid, "started", false);
            let action_result: Result<Thing,DownloadError>;
            {
                let mut left_files: Vec<String> = Vec::with_capacity(2);
                if result.playlist {
                    action_result = handle_playlist(result, &converter,&mut left_files);
                }else{
                    action_result = handle_download(result, &None, &converter,&mut left_files);
                }

                if !left_files.is_empty() {
                    trace!("cleaning up files");
                    for i in &left_files {
                        match remove_file(&i) {
                            Ok(_) => (trace!("cleaning up {}",i)),
                            Err(e) => warn!("unable to remove file '{}' {}",i,e),
                        }
                    }
                }
            }

            let code: i8 = match action_result {
                Ok(t) => {
                    match t {
                        Thing::Bool(b) => {
                            if b { CODE_SUCCESS_WARNINGS } else { CODE_SUCCESS }
                        }
                        _ => CODE_SUCCESS,
                    }
                },
                Err(e) => {
                    warn!("Error: {:?}", e);
                    match e {
                        DownloadError::NotAvailable => CODE_FAILED_UNAVAILABLE,
                        DownloadError::ExtractorError => CODE_FAILED_UNAVAILABLE,
                        DownloadError::QualityNotAvailable => CODE_FAILED_QUALITY,
                        _ => {
                            let details = match e {
                                DownloadError::DBError(s) => s,
                                DownloadError::DownloadError(s) => s,
                                DownloadError::FFMPEGError(s) => s,
                                DownloadError::InternalError(s) => s,
                                _ => unreachable!(),
                            };
                            db::add_query_status(&pool.clone(),&qid, &details);
                            CODE_FAILED_INTERNAL
                        },
                    }
                }
            };
            db::set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");
			db::set_null_state(&pool.clone(), &qid);

        } else {
            if print_pause { debug!("Pausing.."); print_pause = false; }
            std::thread::sleep_ms(SLEEP_MS);
        }
    }
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video separately and convert them together
///In case of a playlist download it depends on the target download folder & if it should be bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, the video but with the highest quality.
///Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, it's moved after a successful download, converted etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode'd name of the file is stored in the DB
fn handle_download<'a>(downl_db: DownloadDB, folder: &Option<String>, converter: &Converter, file_db: &mut Vec<String>) -> Result<Thing,DownloadError>{
    //update progress
    let is_zipped = match *folder {
        Some(_) => true,
        None => false,
    };
    if !downl_db.compress {
        db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, 0,false);
    }

    let download = Downloader::new(&downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded

    let temp_path = lib::format_file_path(&downl_db.qid, folder.clone(), false);
    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            trace!("DMCA error!");
            match download.lib_request_video(1,0, &temp_path) {
                Err(err) => { warn!("lib-call error {:?}", err); return Err(err); },
                Ok(v) => { dmca = true; v },
            }
        },
        Err(e) => { // unknown error / restricted source etc.. abort
            error!("Unknown error: {:?}", e);
            return Err(e);
        },
    };

    file_db.push(temp_path.clone());
    let save_path = lib::format_save_path(folder.clone(),&name, &download);

    trace!("Filename: {}", name);

    let is_splitted_video = if dmca {
        false
    } else {
        lib::is_split_container(&downl_db.quality)
    };
    let convert_audio = CONFIG.extensions.mp3.contains(&downl_db.quality);

    let total_steps = if dmca {
        if download.is_audio() {
            3
        } else {
            2
        }
    } else if is_splitted_video {
        4
    } else if download.is_audio() {
        3
    } else {
        2
    };

    if !dmca {
        if !downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, total_steps,false);
        }

        //download first file, download audio raw source if specified or video
        try!(download.download_file(&temp_path, convert_audio));

        if is_splitted_video {
            // download audio file & convert together
            if !downl_db.compress {
                db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 3, total_steps,false);
            }

            let audio_path = lib::format_file_path(&downl_db.qid, folder.clone(), true);
            file_db.push(audio_path.clone());

            trace!("Downloading audio.. {}", audio_path);
            try!(download.download_file(&audio_path, true));

            if !downl_db.compress {
                db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 4, total_steps,false);
            }

            match converter.merge_files(&downl_db.qid,&temp_path, &audio_path,&save_path.to_string_lossy(), !downl_db.compress) {
                Err(e) => {println!("merge error: {:?}",e); return Err(e);},
                Ok(()) => {},
            }

            try!(remove_file(&audio_path));
            file_db.pop();

        }else{ // we're already done, only need to copy if it's not raw audio source
            if !download.is_audio() {
                try!(lib::move_file(&temp_path, &save_path));
            }
        }
    }

    if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
        if !downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&downl_db.qid, total_steps, total_steps, false);
        }
        try!(converter.extract_audio(&temp_path, &save_path.to_string_lossy(), convert_audio));
        try!(remove_file(&temp_path));
    }else{
        if is_splitted_video {
            try!(remove_file(&temp_path));
        }else{
            try!(lib::move_file(&temp_path, &save_path));
        }
    }
    file_db.pop();

    if !is_zipped { // add file to list, except it's for zip-compression later (=folder set)
        db::add_file_entry(&downl_db.pool.clone(), &downl_db.qid,&save_path.file_name().unwrap().to_string_lossy(), &name);
    }
    if !downl_db.compress {
        db::update_steps(&downl_db.pool.clone(),&downl_db.qid, total_steps, total_steps, true);
    }

    if folder.is_some() {
        Ok(Thing::String(save_path.to_string_lossy().into_owned()))
    } else {
        Ok(Thing::None)
    }
}

///Handles a playlist request
///If zipping isn't requested the downloads will be split up,
///so for each video in the playlist an own query entry will be created
///if warnings occured (unavailable video etc) the return will be true
fn handle_playlist(mut downl_db: DownloadDB, converter: &Converter, file_db: &mut Vec<String>) -> Result<Thing, DownloadError>{
    let mut max_steps: i32 = if downl_db.compress { 4 } else { 3 };
    db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, max_steps,false);

    let pl_id: i64 = downl_db.qid;
    downl_db.update_folder(format!("{}/{}",&CONFIG.general.temp_dir,pl_id));
    trace!("Folder:  {}",downl_db.folder);

    let db_copy = downl_db.clone();
    let download = Downloader::new(&db_copy, &CONFIG.general);

    let mut playlist_name = String::new();
    if downl_db.compress {
        playlist_name = try!(download.get_playlist_name());
        playlist_name = lib::url_encode(&playlist_name);
        println!("pl name {}",playlist_name);
        try!(std::fs::create_dir(&downl_db.folder));
        file_db.push(downl_db.folder.clone());
    }
    db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, max_steps,false);
    trace!("retriving playlist videos..");
    let file_ids: Vec<String> = try!(download.get_playlist_ids());

    let handler_folder = if downl_db.compress {
        Some(pl_id.to_string())
    }else {
        None
    };

    trace!("got em");
    max_steps += file_ids.len() as i32;
    let mut current_step = 2;
    let mut warnings = false;
    let mut current_url: String;
    let mut failed_log: String = String::from("Following urls couldn't be downloaded: \n");
    // we'll store all files and delete em later, so we don't need rm -rf
    let mut file_delete_list: Vec<String> = Vec::with_capacity(if downl_db.compress { file_ids.len() } else { 0 });
    for id in file_ids.iter() {
        current_step += 1;
        if downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        }
        current_url = format!("https://www.youtube.com/watch?v={}",id);
        downl_db.update_video(current_url.clone(), current_step as i64);
        debug!("id: {}",id);
        let db_copy = downl_db.clone();
        match handle_download(db_copy, &handler_folder, converter, file_db) {
            Err(e) => { warn!("error downloading {}: {:?}",id,e);
                        failed_log.push_str(&format!("{} {:?}\n", current_url, e));
                        warnings = true;
            },
            Ok(e) => {  if downl_db.compress {
                            match e {
                                Thing::String(v) => file_delete_list.push(v),
                                _ => {error!("handle_download not returning a filename"); panic!();},
                            }
                        }
            },
        }
    }

    if warnings {
        db::add_query_status(&downl_db.pool.clone(),&pl_id, &failed_log);
    }

    trace!("downloaded all videos");
    if downl_db.compress { // zip to file, add to db & remove all sources
        current_step += 1;
        db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        let zip_name = format!("{}.zip",playlist_name);
        let zip_file = format!("{}/{}",&CONFIG.general.download_dir,zip_name);
        trace!("zip file: {} \n zip source {}",zip_file, &downl_db.folder);
        try!(lib::zip_folder(&downl_db.folder, &zip_file));
        db::add_file_entry(&downl_db.pool.clone(), &pl_id,&zip_name, &playlist_name);

        current_step += 1;
        db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        try!(lib::delete_files(file_delete_list));
        try!(remove_dir_all(downl_db.folder));
        file_db.pop();
    }

    Ok(Thing::Bool(warnings))
}
<commit_msg>add lib_use functionality closes #39<commit_after>extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate log4rs;
#[macro_use]
extern crate lazy_static;

mod lib;

use lib::config;
use lib::db;
use lib::logger;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::DownloadError;
use lib::converter::Converter;

use std::fs::{remove_file,remove_dir_all};

static VERSION : &'static str = "0.3"; // String not valid
static CONFIG_PATH : &'static str = "config.cfg";
static LOG_CONFIG: &'static str = "log.conf";
static LOG_PATTERN: &'static str = "%d{%d-%m-%Y %H:%M:%S}\t[%l]\t%f:%L \t%m";
static SLEEP_MS: u32 = 5000;
static CODE_SUCCESS: i8 = 2;
static CODE_SUCCESS_WARNINGS: i8 = 3; // finished with warnings
static CODE_FAILED_INTERNAL: i8 = 10; // internal error
static CODE_FAILED_QUALITY: i8 = 11; // qualitz not available
static CODE_FAILED_UNAVAILABLE: i8 = 12; // source unavailable (private / removed)

lazy_static! {
    pub static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
// enum QueryCodes {
//     Waiting = 0,
//     InProgress = 1,
//     Finished = 2,
//     Failed = 3,
// }

enum Thing { String(String), Bool(bool), None }

fn main() {
    logger::initialize();

    let pool = db::db_connect(db::mysql_options(), SLEEP_MS);

    let converter = Converter::new(&CONFIG.general.ffmpeg_bin_dir,&CONFIG.general.mp3_quality , pool.clone());
    let mut print_pause = true;
    loop {
        if let Some(result) = db::request_entry(& pool) {
            print_pause = true;
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            db::set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");
            db::set_query_state(&pool.clone(),&qid, "started", false);
            let action_result: Result<Thing,DownloadError>;
            {
                let mut left_files: Vec<String> = Vec::with_capacity(2);
                if result.playlist {
                    action_result = handle_playlist(result, &converter,&mut left_files);
                }else{
                    action_result = handle_download(result, &None, &converter,&mut left_files);
                }

                if !left_files.is_empty() {
                    trace!("cleaning up files");
                    for i in &left_files {
                        match remove_file(&i) {
                            Ok(_) => (trace!("cleaning up {}",i)),
                            Err(e) => warn!("unable to remove file '{}' {}",i,e),
                        }
                    }
                }
            }

            let code: i8 = match action_result {
                Ok(t) => {
                    match t {
                        Thing::Bool(b) => {
                            if b { CODE_SUCCESS_WARNINGS } else { CODE_SUCCESS }
                        }
                        _ => CODE_SUCCESS,
                    }
                },
                Err(e) => {
                    warn!("Error: {:?}", e);
                    match e {
                        DownloadError::NotAvailable => CODE_FAILED_UNAVAILABLE,
                        DownloadError::ExtractorError => CODE_FAILED_UNAVAILABLE,
                        DownloadError::QualityNotAvailable => CODE_FAILED_QUALITY,
                        _ => {
                            let details = match e {
                                DownloadError::DBError(s) => s,
                                DownloadError::DownloadError(s) => s,
                                DownloadError::FFMPEGError(s) => s,
                                DownloadError::InternalError(s) => s,
                                _ => unreachable!(),
                            };
                            db::add_query_status(&pool.clone(),&qid, &details);
                            CODE_FAILED_INTERNAL
                        },
                    }
                }
            };
            db::set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");
			db::set_null_state(&pool.clone(), &qid);

        } else {
            if print_pause { debug!("Pausing.."); print_pause = false; }
            std::thread::sleep_ms(SLEEP_MS);
        }
    }
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video separately and convert them together
///In case of a playlist download it depends on the target download folder & if it should be bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, the video but with the highest quality.
///Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, it's moved after a successful download, converted etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode'd name of the file is stored in the DB
fn handle_download<'a>(downl_db: DownloadDB, folder: &Option<String>, converter: &Converter, file_db: &mut Vec<String>) -> Result<Thing,DownloadError>{
    //update progress
    let is_zipped = match *folder {
        Some(_) => true,
        None => false,
    };
    if !downl_db.compress {
        db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, 0,false);
    }

    let download = Downloader::new(&downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded

    let temp_path = lib::format_file_path(&downl_db.qid, folder.clone(), false);
    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            info!("DMCA error!");
            if CONFIG.general.lib_use {
                match download.lib_request_video(1,0, &temp_path) {
                    Err(err) => { warn!("lib-call error {:?}", err); return Err(err); },
                    Ok(v) => { dmca = true; v },
                }
            } else {
                return Err(DownloadError::NotAvailable);
            }
        },
        Err(e) => { // unknown error / restricted source etc.. abort
            error!("Unknown error: {:?}", e);
            return Err(e);
        },
    };

    file_db.push(temp_path.clone());
    let save_path = lib::format_save_path(folder.clone(),&name, &download);

    trace!("Filename: {}", name);

    let is_splitted_video = if dmca {
        false
    } else {
        lib::is_split_container(&downl_db.quality)
    };
    let convert_audio = CONFIG.extensions.mp3.contains(&downl_db.quality);

    let total_steps = if dmca {
        if download.is_audio() {
            3
        } else {
            2
        }
    } else if is_splitted_video {
        4
    } else if download.is_audio() {
        3
    } else {
        2
    };

    if !dmca {
        if !downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, total_steps,false);
        }

        //download first file, download audio raw source if specified or video
        try!(download.download_file(&temp_path, convert_audio));

        if is_splitted_video {
            // download audio file & convert together
            if !downl_db.compress {
                db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 3, total_steps,false);
            }

            let audio_path = lib::format_file_path(&downl_db.qid, folder.clone(), true);
            file_db.push(audio_path.clone());

            trace!("Downloading audio.. {}", audio_path);
            try!(download.download_file(&audio_path, true));

            if !downl_db.compress {
                db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 4, total_steps,false);
            }

            match converter.merge_files(&downl_db.qid,&temp_path, &audio_path,&save_path.to_string_lossy(), !downl_db.compress) {
                Err(e) => {println!("merge error: {:?}",e); return Err(e);},
                Ok(()) => {},
            }

            try!(remove_file(&audio_path));
            file_db.pop();

        }else{ // we're already done, only need to copy if it's not raw audio source
            if !download.is_audio() {
                try!(lib::move_file(&temp_path, &save_path));
            }
        }
    }

    if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
        if !downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&downl_db.qid, total_steps, total_steps, false);
        }
        try!(converter.extract_audio(&temp_path, &save_path.to_string_lossy(), convert_audio));
        try!(remove_file(&temp_path));
    }else{
        if is_splitted_video {
            try!(remove_file(&temp_path));
        }else{
            try!(lib::move_file(&temp_path, &save_path));
        }
    }
    file_db.pop();

    if !is_zipped { // add file to list, except it's for zip-compression later (=folder set)
        db::add_file_entry(&downl_db.pool.clone(), &downl_db.qid,&save_path.file_name().unwrap().to_string_lossy(), &name);
    }
    if !downl_db.compress {
        db::update_steps(&downl_db.pool.clone(),&downl_db.qid, total_steps, total_steps, true);
    }

    if folder.is_some() {
        Ok(Thing::String(save_path.to_string_lossy().into_owned()))
    } else {
        Ok(Thing::None)
    }
}

///Handles a playlist request
///If zipping isn't requested the downloads will be split up,
///so for each video in the playlist an own query entry will be created
///if warnings occured (unavailable video etc) the return will be true
fn handle_playlist(mut downl_db: DownloadDB, converter: &Converter, file_db: &mut Vec<String>) -> Result<Thing, DownloadError>{
    let mut max_steps: i32 = if downl_db.compress { 4 } else { 3 };
    db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, max_steps,false);

    let pl_id: i64 = downl_db.qid;
    downl_db.update_folder(format!("{}/{}",&CONFIG.general.temp_dir,pl_id));
    trace!("Folder:  {}",downl_db.folder);

    let db_copy = downl_db.clone();
    let download = Downloader::new(&db_copy, &CONFIG.general);

    let mut playlist_name = String::new();
    if downl_db.compress {
        playlist_name = try!(download.get_playlist_name());
        playlist_name = lib::url_encode(&playlist_name);
        println!("pl name {}",playlist_name);
        try!(std::fs::create_dir(&downl_db.folder));
        file_db.push(downl_db.folder.clone());
    }
    db::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, max_steps,false);
    trace!("retriving playlist videos..");
    let file_ids: Vec<String> = try!(download.get_playlist_ids());

    let handler_folder = if downl_db.compress {
        Some(pl_id.to_string())
    }else {
        None
    };

    trace!("got em");
    max_steps += file_ids.len() as i32;
    let mut current_step = 2;
    let mut warnings = false;
    let mut current_url: String;
    let mut failed_log: String = String::from("Following urls couldn't be downloaded: \n");
    // we'll store all files and delete em later, so we don't need rm -rf
    let mut file_delete_list: Vec<String> = Vec::with_capacity(if downl_db.compress { file_ids.len() } else { 0 });
    for id in file_ids.iter() {
        current_step += 1;
        if downl_db.compress {
            db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        }
        current_url = format!("https://www.youtube.com/watch?v={}",id);
        downl_db.update_video(current_url.clone(), current_step as i64);
        debug!("id: {}",id);
        let db_copy = downl_db.clone();
        match handle_download(db_copy, &handler_folder, converter, file_db) {
            Err(e) => { warn!("error downloading {}: {:?}",id,e);
                        failed_log.push_str(&format!("{} {:?}\n", current_url, e));
                        warnings = true;
            },
            Ok(e) => {  if downl_db.compress {
                            match e {
                                Thing::String(v) => file_delete_list.push(v),
                                _ => {error!("handle_download not returning a filename"); panic!();},
                            }
                        }
            },
        }
    }

    if warnings {
        db::add_query_status(&downl_db.pool.clone(),&pl_id, &failed_log);
    }

    trace!("downloaded all videos");
    if downl_db.compress { // zip to file, add to db & remove all sources
        current_step += 1;
        db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        let zip_name = format!("{}.zip",playlist_name);
        let zip_file = format!("{}/{}",&CONFIG.general.download_dir,zip_name);
        trace!("zip file: {} \n zip source {}",zip_file, &downl_db.folder);
        try!(lib::zip_folder(&downl_db.folder, &zip_file));
        db::add_file_entry(&downl_db.pool.clone(), &pl_id,&zip_name, &playlist_name);

        current_step += 1;
        db::update_steps(&downl_db.pool.clone(),&pl_id, current_step, max_steps,false);
        try!(lib::delete_files(file_delete_list));
        try!(remove_dir_all(downl_db.folder));
        file_db.pop();
    }

    Ok(Thing::Bool(warnings))
}
<|endoftext|>