<commit_before>//! This module implements the root command of the CLI tool.

use std::env;
use std::fmt;
use std::process;

use clap::{App, AppSettings, Arg, ArgMatches};
use failure::{bail, Error};
use log::{debug, info};

use crate::api::Api;
use crate::config::{prepare_environment, Auth, Config};
use crate::constants::{ARCH, PLATFORM, VERSION};
use crate::utils::system::{print_error, QuietExit};
use crate::utils::update::run_sentrycli_update_nagger;

const ABOUT: &str = "
Command line utility for Sentry.

This tool helps you manage remote resources on a Sentry server like
sourcemaps, debug symbols or releases.  Use `--help` on the subcommands
to learn more about them.";

macro_rules! each_subcommand {
    ($mac:ident) => {
        $mac!(upload_dif);
        $mac!(upload_proguard);
        $mac!(releases);
        $mac!(issues);
        $mac!(repos);
        $mac!(projects);
        #[cfg(not(feature = "managed"))]
        $mac!(update);
        #[cfg(not(feature = "managed"))]
        $mac!(uninstall);
        $mac!(info);
        $mac!(login);
        $mac!(send_event);
        $mac!(react_native);
        $mac!(difutil);

        // these here exist for legacy reasons only.  They were moved
        // to subcommands of the react-native command.  Note that
        // codepush was never available on that level.
        #[cfg(target_os = "macos")]
        $mac!(react_native_xcode);
        $mac!(react_native_gradle);
    };
}

// commands we want to run the update nagger on
const UPDATE_NAGGER_CMDS: &[&str] = &[
    "releases", "issues", "repos", "projects", "info", "login", "difutil",
];

// it would be great if this could be a macro expansion as well
// but rust bug #37663 breaks location information then.
// FIXME: #37663 has been merged.
pub mod info;
pub mod issues;
pub mod login;
pub mod projects;
pub mod releases;
pub mod repos;
pub mod send_event;
#[cfg(not(feature = "managed"))]
pub mod uninstall;
#[cfg(not(feature = "managed"))]
pub mod update;
pub mod upload_dif;
pub mod upload_proguard;

pub mod react_native;
pub mod react_native_appcenter;
pub mod react_native_codepush;
pub mod react_native_gradle;
#[cfg(target_os = "macos")]
pub mod react_native_xcode;

pub mod difutil;
pub mod difutil_bundle_sources;
pub mod difutil_check;
pub mod difutil_find;
pub mod difutil_id;

fn preexecute_hooks() -> Result<bool, Error> {
    return sentry_react_native_xcode_wrap();

    #[cfg(target_os = "macos")]
    fn sentry_react_native_xcode_wrap() -> Result<bool, Error> {
        if let Ok(val) = env::var("__SENTRY_RN_WRAP_XCODE_CALL") {
            env::remove_var("__SENTRY_RN_WRAP_XCODE_CALL");
            if &val == "1" {
                react_native_xcode::wrap_call()?;
                return Ok(true);
            }
        }
        Ok(false)
    }

    #[allow(clippy::unnecessary_wraps)]
    #[cfg(not(target_os = "macos"))]
    fn sentry_react_native_xcode_wrap() -> Result<bool, Error> {
        Ok(false)
    }
}

fn configure_args(config: &mut Config, matches: &ArgMatches<'_>) -> Result<(), Error> {
    if let Some(url) = matches.value_of("url") {
        config.set_base_url(url);
    }

    if let Some(api_key) = matches.value_of("api_key") {
        config.set_auth(Auth::Key(api_key.to_owned()));
    }

    if let Some(auth_token) = matches.value_of("auth_token") {
        config.set_auth(Auth::Token(auth_token.to_owned()));
    }

    if let Some(level_str) = matches.value_of("log_level") {
        match level_str.parse() {
            Ok(level) => {
                config.set_log_level(level);
            }
            Err(_) => {
                bail!("Unknown log level: {}", level_str);
            }
        }
    }

    Ok(())
}

fn add_commands<'a, 'b>(mut app: App<'a, 'b>) -> App<'a, 'b> {
    macro_rules! add_subcommand {
        ($name:ident) => {{
            let mut cmd = $name::make_app(App::new(stringify!($name).replace("_", "-").as_str()));

            // for legacy reasons
            if stringify!($name).starts_with("react_native_") {
                cmd = cmd.setting(AppSettings::Hidden);
            }

            app = app.subcommand(cmd);
        }};
    }

    each_subcommand!(add_subcommand);
    app
}

#[allow(clippy::cognitive_complexity)]
fn run_command(matches: &ArgMatches<'_>) -> Result<(), Error> {
    macro_rules! execute_subcommand {
        ($name:ident) => {{
            let cmd = stringify!($name).replace("_", "-");
            if let Some(sub_matches) = matches.subcommand_matches(&cmd) {
                let rv = $name::execute(&sub_matches)?;
                if UPDATE_NAGGER_CMDS.iter().any(|x| x == &cmd) {
                    run_sentrycli_update_nagger();
                }
                return Ok(rv);
            }
        }};
    }

    each_subcommand!(execute_subcommand);
    unreachable!();
}

struct DebugArgs<'a>(Vec<&'a str>);

impl<'a> fmt::Display for DebugArgs<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (idx, arg) in self.0.iter().enumerate() {
            if idx > 0 {
                write!(f, " ")?;
            }
            write!(f, "{:?}", arg)?;
        }
        Ok(())
    }
}

/// Given an argument vector and a `Config` this executes the
/// command line and returns the result.
pub fn execute(args: &[String]) -> Result<(), Error> {
    let mut config = Config::from_cli_config()?;

    // special case for the xcode integration for react native.  For more
    // information see commands/react_native_xcode.rs
    if preexecute_hooks()? {
        return Ok(());
    }

    let mut app = App::new("sentry-cli")
        .help_message("Print this help message.")
        .version(VERSION)
        .version_message("Print version information.")
        .about(ABOUT)
        .max_term_width(100)
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .global_setting(AppSettings::UnifiedHelpMessage)
        .arg(Arg::with_name("url").value_name("URL").long("url").help(
            "Fully qualified URL to the Sentry server.{n}\
             [defaults to https://sentry.io/]",
        ))
        .arg(
            Arg::with_name("auth_token")
                .value_name("AUTH_TOKEN")
                .long("auth-token")
                .help("Use the given Sentry auth token."),
        )
        .arg(
            Arg::with_name("api_key")
                .value_name("API_KEY")
                .long("api-key")
                .help("The given Sentry API key."),
        )
        .arg(
            Arg::with_name("log_level")
                .value_name("LOG_LEVEL")
                .long("log-level")
                .possible_values(&["trace", "debug", "info", "warn", "error"])
                .case_insensitive(true)
                .global(true)
                .help("Set the log output verbosity."),
        );

    app = add_commands(app);
    let matches = app.get_matches_from_safe(args)?;
    configure_args(&mut config, &matches)?;

    // bind the config to the process and fetch an immutable reference to it
    config.bind_to_process();
    if Config::current().get_filename().exists() {
        info!(
            "Loaded config from {}",
            Config::current().get_filename().display()
        );
    }

    debug!(
        "sentry-cli version: {}, platform: \"{}\", architecture: \"{}\"",
        VERSION, PLATFORM, ARCH
    );

    info!(
        "sentry-cli was invoked with the following command line: {}",
        DebugArgs(args.iter().map(String::as_str).collect())
    );

    run_command(&matches)
}

fn run() -> Result<(), Error> {
    prepare_environment();
    execute(&env::args().collect::<Vec<String>>())
}

fn setup() {
    use crate::utils::logging::Logger;

    crate::utils::system::init_backtrace();
    crate::utils::system::load_dotenv();

    // we use debug internally but our log handler then rejects to a lower limit.
    // This is okay for our uses but not as efficient.
    log::set_max_level(log::LevelFilter::Debug);

    // if we work with crash reporting we initialize the sentry system.  This
    // also configures the logger.
    #[cfg(feature = "with_crash_reporting")]
    {
        crate::utils::crashreporting::setup(Box::new(Logger));
    }
    #[cfg(not(feature = "with_crash_reporting"))]
    {
        log::set_logger(&Logger).unwrap();
    }
}

/// Executes the command line application and exits the process.
pub fn main() {
    setup();
    let result = run();

    let status_code = match result {
        Ok(()) => 0,
        Err(err) => {
            let code = if let Some(&QuietExit(code)) = err.downcast_ref() {
                code
            } else {
                print_error(&err);
                #[cfg(feature = "with_crash_reporting")]
                {
                    crate::utils::crashreporting::try_report_to_sentry(err);
                }
                1
            };

            // if the user hit an error, it might be time to run the update
            // nagger because maybe they tried to do something only newer
            // versions support.
            run_sentrycli_update_nagger();

            code
        }
    };

    // before we shut down we unbind the api to give the connection pool
    // a chance to collect.  Not doing so has shown to cause hung threads
    // on windows.
    Api::dispose_pool();
    process::exit(status_code);
}
<commit_msg>ref: Remove react-native-gradle and react-native-xcode commands<commit_after>//! This module implements the root command of the CLI tool.

use std::env;
use std::fmt;
use std::process;

use clap::{App, AppSettings, Arg, ArgMatches};
use failure::{bail, Error};
use log::{debug, info};

use crate::api::Api;
use crate::config::{prepare_environment, Auth, Config};
use crate::constants::{ARCH, PLATFORM, VERSION};
use crate::utils::system::{print_error, QuietExit};
use crate::utils::update::run_sentrycli_update_nagger;

const ABOUT: &str = "
Command line utility for Sentry.

This tool helps you manage remote resources on a Sentry server like
sourcemaps, debug symbols or releases.  Use `--help` on the subcommands
to learn more about them.";

macro_rules! each_subcommand {
    ($mac:ident) => {
        $mac!(upload_dif);
        $mac!(upload_proguard);
        $mac!(releases);
        $mac!(issues);
        $mac!(repos);
        $mac!(projects);
        #[cfg(not(feature = "managed"))]
        $mac!(update);
        #[cfg(not(feature = "managed"))]
        $mac!(uninstall);
        $mac!(info);
        $mac!(login);
        $mac!(send_event);
        $mac!(react_native);
        $mac!(difutil);
    };
}

// commands we want to run the update nagger on
const UPDATE_NAGGER_CMDS: &[&str] = &[
    "releases", "issues", "repos", "projects", "info", "login", "difutil",
];

// it would be great if this could be a macro expansion as well
// but rust bug #37663 breaks location information then.
// FIXME: #37663 has been merged.
pub mod info;
pub mod issues;
pub mod login;
pub mod projects;
pub mod releases;
pub mod repos;
pub mod send_event;
#[cfg(not(feature = "managed"))]
pub mod uninstall;
#[cfg(not(feature = "managed"))]
pub mod update;
pub mod upload_dif;
pub mod upload_proguard;

pub mod react_native;
pub mod react_native_appcenter;
pub mod react_native_codepush;
pub mod react_native_gradle;
#[cfg(target_os = "macos")]
pub mod react_native_xcode;

pub mod difutil;
pub mod difutil_bundle_sources;
pub mod difutil_check;
pub mod difutil_find;
pub mod difutil_id;

fn preexecute_hooks() -> Result<bool, Error> {
    return sentry_react_native_xcode_wrap();

    #[cfg(target_os = "macos")]
    fn sentry_react_native_xcode_wrap() -> Result<bool, Error> {
        if let Ok(val) = env::var("__SENTRY_RN_WRAP_XCODE_CALL") {
            env::remove_var("__SENTRY_RN_WRAP_XCODE_CALL");
            if &val == "1" {
                react_native_xcode::wrap_call()?;
                return Ok(true);
            }
        }
        Ok(false)
    }

    #[allow(clippy::unnecessary_wraps)]
    #[cfg(not(target_os = "macos"))]
    fn sentry_react_native_xcode_wrap() -> Result<bool, Error> {
        Ok(false)
    }
}

fn configure_args(config: &mut Config, matches: &ArgMatches<'_>) -> Result<(), Error> {
    if let Some(url) = matches.value_of("url") {
        config.set_base_url(url);
    }

    if let Some(api_key) = matches.value_of("api_key") {
        config.set_auth(Auth::Key(api_key.to_owned()));
    }

    if let Some(auth_token) = matches.value_of("auth_token") {
        config.set_auth(Auth::Token(auth_token.to_owned()));
    }

    if let Some(level_str) = matches.value_of("log_level") {
        match level_str.parse() {
            Ok(level) => {
                config.set_log_level(level);
            }
            Err(_) => {
                bail!("Unknown log level: {}", level_str);
            }
        }
    }

    Ok(())
}

fn add_commands<'a, 'b>(mut app: App<'a, 'b>) -> App<'a, 'b> {
    macro_rules! add_subcommand {
        ($name:ident) => {{
            let cmd = $name::make_app(App::new(stringify!($name).replace("_", "-").as_str()));
            app = app.subcommand(cmd);
        }};
    }

    each_subcommand!(add_subcommand);
    app
}

#[allow(clippy::cognitive_complexity)]
fn run_command(matches: &ArgMatches<'_>) -> Result<(), Error> {
    macro_rules! execute_subcommand {
        ($name:ident) => {{
            let cmd = stringify!($name).replace("_", "-");
            if let Some(sub_matches) = matches.subcommand_matches(&cmd) {
                let rv = $name::execute(&sub_matches)?;
                if UPDATE_NAGGER_CMDS.iter().any(|x| x == &cmd) {
                    run_sentrycli_update_nagger();
                }
                return Ok(rv);
            }
        }};
    }

    each_subcommand!(execute_subcommand);
    unreachable!();
}

struct DebugArgs<'a>(Vec<&'a str>);

impl<'a> fmt::Display for DebugArgs<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (idx, arg) in self.0.iter().enumerate() {
            if idx > 0 {
                write!(f, " ")?;
            }
            write!(f, "{:?}", arg)?;
        }
        Ok(())
    }
}

/// Given an argument vector and a `Config` this executes the
/// command line and returns the result.
pub fn execute(args: &[String]) -> Result<(), Error> {
    let mut config = Config::from_cli_config()?;

    // special case for the xcode integration for react native.  For more
    // information see commands/react_native_xcode.rs
    if preexecute_hooks()? {
        return Ok(());
    }

    let mut app = App::new("sentry-cli")
        .help_message("Print this help message.")
        .version(VERSION)
        .version_message("Print version information.")
        .about(ABOUT)
        .max_term_width(100)
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .global_setting(AppSettings::UnifiedHelpMessage)
        .arg(Arg::with_name("url").value_name("URL").long("url").help(
            "Fully qualified URL to the Sentry server.{n}\
             [defaults to https://sentry.io/]",
        ))
        .arg(
            Arg::with_name("auth_token")
                .value_name("AUTH_TOKEN")
                .long("auth-token")
                .help("Use the given Sentry auth token."),
        )
        .arg(
            Arg::with_name("api_key")
                .value_name("API_KEY")
                .long("api-key")
                .help("The given Sentry API key."),
        )
        .arg(
            Arg::with_name("log_level")
                .value_name("LOG_LEVEL")
                .long("log-level")
                .possible_values(&["trace", "debug", "info", "warn", "error"])
                .case_insensitive(true)
                .global(true)
                .help("Set the log output verbosity."),
        );

    app = add_commands(app);
    let matches = app.get_matches_from_safe(args)?;
    configure_args(&mut config, &matches)?;

    // bind the config to the process and fetch an immutable reference to it
    config.bind_to_process();
    if Config::current().get_filename().exists() {
        info!(
            "Loaded config from {}",
            Config::current().get_filename().display()
        );
    }

    debug!(
        "sentry-cli version: {}, platform: \"{}\", architecture: \"{}\"",
        VERSION, PLATFORM, ARCH
    );

    info!(
        "sentry-cli was invoked with the following command line: {}",
        DebugArgs(args.iter().map(String::as_str).collect())
    );

    run_command(&matches)
}

fn run() -> Result<(), Error> {
    prepare_environment();
    execute(&env::args().collect::<Vec<String>>())
}

fn setup() {
    use crate::utils::logging::Logger;

    crate::utils::system::init_backtrace();
    crate::utils::system::load_dotenv();

    // we use debug internally but our log handler then rejects to a lower limit.
    // This is okay for our uses but not as efficient.
    log::set_max_level(log::LevelFilter::Debug);

    // if we work with crash reporting we initialize the sentry system.  This
    // also configures the logger.
    #[cfg(feature = "with_crash_reporting")]
    {
        crate::utils::crashreporting::setup(Box::new(Logger));
    }
    #[cfg(not(feature = "with_crash_reporting"))]
    {
        log::set_logger(&Logger).unwrap();
    }
}

/// Executes the command line application and exits the process.
pub fn main() {
    setup();
    let result = run();

    let status_code = match result {
        Ok(()) => 0,
        Err(err) => {
            let code = if let Some(&QuietExit(code)) = err.downcast_ref() {
                code
            } else {
                print_error(&err);
                #[cfg(feature = "with_crash_reporting")]
                {
                    crate::utils::crashreporting::try_report_to_sentry(err);
                }
                1
            };

            // if the user hit an error, it might be time to run the update
            // nagger because maybe they tried to do something only newer
            // versions support.
            run_sentrycli_update_nagger();

            code
        }
    };

    // before we shut down we unbind the api to give the connection pool
    // a chance to collect.  Not doing so has shown to cause hung threads
    // on windows.
    Api::dispose_pool();
    process::exit(status_code);
}
<|endoftext|>


<commit_before>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    #[structopt(long="merged", short="m", help = "Merged Output file", name="MERGED_OUT_FILE", default_value="-")]
    merged_outfile: Option<String>,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push(exon.start-1..exon.end);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    if exon.start < last.end && last.start < exon.end {
                        last.end = exon.end;
                        extend = true;
                    }
                }
                if !extend {
                    merged.push(exon.clone());
                }
            }
        }
    }

    write_exon_cov_to_file(options,
        total_reads,
        &unmerged_exons,
        &options.outfile,
        bamfiles,
        bamstrand,
        &tidmaps,
    )?;
    if let Some(ref merged_outfile) = options.merged_outfile {
        write_exon_cov_to_file(options,
            total_reads,
            &merged_exons,
            &merged_outfile,
            bamfiles,
            bamstrand,
            &tidmaps,
        )?;
    }
    Ok(())
}

fn write_exon_cov_to_file(
    options: &Options,
    total_reads: u64,
    exons: &HashMap<(String,String),Vec<Range<u64>>>,
    outfile: &str,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>,
    tidmaps: &Arc<HashMap<String,HashMap<String,u32>>>,
    )
    -> Result<()>
{
    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&outfile)?) });

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in exons.keys() {
        let exons = exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.start as u32, exon.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.end && exon.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.end) -
                                        std::cmp::max(e.start, exon.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.end-exon.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.start,
                    end: exon.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<commit_msg>add gene_id and transcript_id to output where appropriate<commit_after>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    #[structopt(long="merged", short="m", help = "Merged Output file", name="MERGED_OUT_FILE", default_value="-")]
    merged_outfile: Option<String>,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
    transcript_id: String,
    gene_id: String
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<(Range<u64>,Option<usize>)>>::new();
    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push((exon.start-1..exon.end, Some(*exon_row)));
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<(Range<u64>,Option<usize>)>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.0.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    if exon.0.start < last.0.end && last.0.start < exon.0.end {
                        last.0.end = exon.0.end;
                        extend = true;
                    }
                }
                if !extend {
                    merged.push((exon.0.clone(), None));
                }
            }
        }
    }

    write_exon_cov_to_file(options,
        total_reads,
        &unmerged_exons,
        &options.outfile,
        bamfiles,
        bamstrand,
        &tidmaps,
        &annot,
    )?;
    if let Some(ref merged_outfile) = options.merged_outfile {
        write_exon_cov_to_file(options,
            total_reads,
            &merged_exons,
            &merged_outfile,
            bamfiles,
            bamstrand,
            &tidmaps,
            &annot,
        )?;
    }
    Ok(())
}

fn write_exon_cov_to_file(
    options: &Options,
    total_reads: u64,
    exons: &HashMap<(String,String),Vec<(Range<u64>,Option<usize>)>>,
    outfile: &str,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>,
    tidmaps: &Arc<HashMap<String,HashMap<String,u32>>>,
    annot: &Arc<IndexedAnnotation>)
    -> Result<()>
{
    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&outfile)?) });

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in exons.keys() {
        let exons = exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();
            let annot = annot.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.0.start as u32, exon.0.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.0.end && exon.0.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.0.end) -
                                        std::cmp::max(e.start, exon.0.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.0.end-exon.0.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                let gene_id = match exon.1 {
                    Some(exon_row) => match get_gene_name(exon_row, &annot) {
                        Some(name) => name,
                        None => "".to_string(),
                    }
                    None => "".to_string(),
                };
                let transcript_id = match exon.1 {
                    Some(exon_row) => match get_name(exon_row, &annot) {
                        Some(name) => name,
                        None => "".to_string(),
                    }
                    None => "".to_string(),
                };
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.0.start,
                    end: exon.0.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                    transcript_id: transcript_id,
                    gene_id: gene_id,
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>use std::default::Default;
use std::ffi::OsStr;
use std::fs;
use std::io::Read;
use std::path;

use liquid;
use serde_json;
use serde_yaml;
use toml;

use error::*;
use files;

const DATA_DIR: &'static str = "_data";

#[derive(Debug, PartialEq, Clone)]
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SiteBuilder {
    pub name: Option<String>,
    pub description: Option<String>,
    pub base_url: Option<String>,
    #[serde(skip)]
    pub data_dir: &'static str,
}

impl Default for SiteBuilder {
    fn default() -> SiteBuilder {
        SiteBuilder {
            name: None,
            description: None,
            base_url: None,
            data_dir: DATA_DIR,
        }
    }
}

impl SiteBuilder {
    pub fn build(self, root: &path::Path) -> Result<Site> {
        let SiteBuilder {
            name,
            description,
            base_url,
            data_dir,
        } = self;
        let base_url = base_url.map(|mut l| {
                                        if l.ends_with('/') {
                                            l.pop();
                                        }
                                        l
                                    });

        let mut attributes = liquid::Object::new();
        let mut data = liquid::Object::new();
        insert_data_dir(&mut data, &root.join(data_dir))?;
        if !data.is_empty() {
            attributes.insert("data".to_owned(), liquid::Value::Object(data));
        }

        Ok(Site {
               name,
               description,
               base_url,
               attributes,
           })
    }
}

#[derive(Debug, PartialEq, Default, Clone)]
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Site {
    pub name: Option<String>,
    pub description: Option<String>,
    pub base_url: Option<String>,
    pub attributes: liquid::Object,
}

fn deep_insert(data_map: &mut liquid::Object,
               file_path: &path::Path,
               target_key: String,
               data: liquid::Value)
               -> Result<()> {
    // now find the nested map it is supposed to be in
    let target_map = if let Some(path) = file_path.parent() {
        let mut map = data_map;
        for part in path.iter() {
            let key = part.to_str().ok_or_else(|| {
                format!("The data from {:?} can't be loaded as it contains non utf-8 characters",
                        path)
            })?;
            let cur_map = map;
            map = cur_map
                .entry(String::from(key))
                .or_insert_with(|| liquid::Value::Object(liquid::Object::new()))
                .as_object_mut()
                .ok_or_else(|| {
                                format!("Aborting: Duplicate in data tree. Would overwrite {:?} ",
                                        path)
                            })?;
        }
        map
    } else {
        data_map
    };

    match target_map.insert(target_key, data) {
        None => Ok(()),
        _ => {
            Err(format!("The data from {:?} can't be loaded: the key already exists",
                        file_path)
                    .into())
        }
    }
}

fn load_data(data_path: &path::Path) -> Result<liquid::Value> {
    let ext = data_path.extension().unwrap_or_else(|| OsStr::new(""));

    let data: liquid::Value;

    if ext == OsStr::new("yml") || ext == OsStr::new("yaml") {
        let reader = fs::File::open(data_path)?;
        data = serde_yaml::from_reader(reader)?;
    } else if ext == OsStr::new("json") {
        let reader = fs::File::open(data_path)?;
        data = serde_json::from_reader(reader)?;
    } else if ext == OsStr::new("toml") {
        let mut reader = fs::File::open(data_path)?;
        let mut text = String::new();
        reader.read_to_string(&mut text)?;
        data = toml::from_str(&text)?;
    } else {
        bail!("Failed to load of data {:?}: unknown file type '{:?}'.\n\
              Supported data files extensions are: yml, yaml, json and toml.",
              data_path,
              ext);
    }

    Ok(data)
}

fn insert_data_dir(data: &mut liquid::Object, data_root: &path::Path) -> Result<()> {
    let data_files_builder = files::FilesBuilder::new(data_root)?;
    let data_files = data_files_builder.build()?;
    for df in data_files.files() {
        let file_stem = df.file_stem()
            .expect("Files will always return with a stem");
        let file_stem = String::from(file_stem.to_str().unwrap());

        let full_path = data_root.join(df.clone());

        let data_fragment = load_data(&full_path)?;

        deep_insert(data, &df, file_stem, data_fragment)?;
    }

    Ok(())
}
<commit_msg>feat: Expose config's site values<commit_after>use std::default::Default;
use std::ffi::OsStr;
use std::fs;
use std::io::Read;
use std::path;

use liquid;
use serde_json;
use serde_yaml;
use toml;

use error::*;
use files;

const DATA_DIR: &'static str = "_data";

#[derive(Debug, PartialEq, Clone)]
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SiteBuilder {
    pub name: Option<String>,
    pub description: Option<String>,
    pub base_url: Option<String>,
    #[serde(skip)]
    pub data_dir: &'static str,
}

impl Default for SiteBuilder {
    fn default() -> SiteBuilder {
        SiteBuilder {
            name: None,
            description: None,
            base_url: None,
            data_dir: DATA_DIR,
        }
    }
}

impl SiteBuilder {
    pub fn build(self, root: &path::Path) -> Result<Site> {
        let SiteBuilder {
            name,
            description,
            base_url,
            data_dir,
        } = self;
        let base_url = base_url.map(|mut l| {
                                        if l.ends_with('/') {
                                            l.pop();
                                        }
                                        l
                                    });

        let mut attributes = liquid::Object::new();
        if let Some(ref name) = name {
            attributes.insert("name".to_owned(), liquid::Value::str(name));
        }
        if let Some(ref description) = description {
            attributes.insert("description".to_owned(), liquid::Value::str(description));
        }
        if let Some(ref base_url) = base_url {
            attributes.insert("base_url".to_owned(), liquid::Value::str(base_url));
        }
        let mut data = liquid::Object::new();
        insert_data_dir(&mut data, &root.join(data_dir))?;
        if !data.is_empty() {
            attributes.insert("data".to_owned(), liquid::Value::Object(data));
        }

        Ok(Site {
               name,
               description,
               base_url,
               attributes,
           })
    }
}

#[derive(Debug, PartialEq, Default, Clone)]
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Site {
    pub name: Option<String>,
    pub description: Option<String>,
    pub base_url: Option<String>,
    pub attributes: liquid::Object,
}

fn deep_insert(data_map: &mut liquid::Object,
               file_path: &path::Path,
               target_key: String,
               data: liquid::Value)
               -> Result<()> {
    // now find the nested map it is supposed to be in
    let target_map = if let Some(path) = file_path.parent() {
        let mut map = data_map;
        for part in path.iter() {
            let key = part.to_str().ok_or_else(|| {
                format!("The data from {:?} can't be loaded as it contains non utf-8 characters",
                        path)
            })?;
            let cur_map = map;
            map = cur_map
                .entry(String::from(key))
                .or_insert_with(|| liquid::Value::Object(liquid::Object::new()))
                .as_object_mut()
                .ok_or_else(|| {
                                format!("Aborting: Duplicate in data tree. Would overwrite {:?} ",
                                        path)
                            })?;
        }
        map
    } else {
        data_map
    };

    match target_map.insert(target_key, data) {
        None => Ok(()),
        _ => {
            Err(format!("The data from {:?} can't be loaded: the key already exists",
                        file_path)
                    .into())
        }
    }
}

fn load_data(data_path: &path::Path) -> Result<liquid::Value> {
    let ext = data_path.extension().unwrap_or_else(|| OsStr::new(""));

    let data: liquid::Value;

    if ext == OsStr::new("yml") || ext == OsStr::new("yaml") {
        let reader = fs::File::open(data_path)?;
        data = serde_yaml::from_reader(reader)?;
    } else if ext == OsStr::new("json") {
        let reader = fs::File::open(data_path)?;
        data = serde_json::from_reader(reader)?;
    } else if ext == OsStr::new("toml") {
        let mut reader = fs::File::open(data_path)?;
        let mut text = String::new();
        reader.read_to_string(&mut text)?;
        data = toml::from_str(&text)?;
    } else {
        bail!("Failed to load of data {:?}: unknown file type '{:?}'.\n\
              Supported data files extensions are: yml, yaml, json and toml.",
              data_path,
              ext);
    }

    Ok(data)
}

fn insert_data_dir(data: &mut liquid::Object, data_root: &path::Path) -> Result<()> {
    let data_files_builder = files::FilesBuilder::new(data_root)?;
    let data_files = data_files_builder.build()?;
    for df in data_files.files() {
        let file_stem = df.file_stem()
            .expect("Files will always return with a stem");
        let file_stem = String::from(file_stem.to_str().unwrap());

        let full_path = data_root.join(df.clone());

        let data_fragment = load_data(&full_path)?;

        deep_insert(data, &df, file_stem, data_fragment)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use core::prelude::*;
use core::cmp::{min, max};
use core::mem::{size_of, zeroed, replace, swap, transmute};
use core::ptr;

/// The smallest number of elements that may be quicksorted.
/// Must be at least 9.
const MIN_QUICKSORT_ELEMS: uint = 10;

/// The maximum number of elements to be insertion sorted.
const MAX_INSERTION_SORT_ELEMS: uint = 42;

/// Controls the number of elements to be insertion sorted.
/// Higher values give more insertion sorted elements.
const INSERTION_SORT_FACTOR: uint = 450;

pub fn sort_by<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    if maybe_insertion_sort(v, compare) { return; }
    let heapsort_depth = (3 * log2(v.len())) / 2;
    do_introsort(v, compare, 0, heapsort_depth);
}

#[inline]
pub fn sort<T: Ord>(v: &mut [T]) {
    sort_by(v, &|&: a: &T, b| a.cmp(b));
}

fn introsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C, rec: u32, heapsort_depth: u32) {
    if maybe_insertion_sort(v, compare) { return; }
    do_introsort(v, compare, rec, heapsort_depth);
}

fn do_introsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C, rec: u32, heapsort_depth: u32) {
    // This works around some bugs with unboxed closures
    macro_rules! maybe_swap(
        ($v: expr, $a: expr, $b: expr, $compare: expr) => {
            if compare_idxs($v, *$a, *$b, $compare) == Greater {
                swap($a, $b);
            }
        }
    )

    if rec > heapsort_depth {
        heapsort(v, compare);
        return;
    }

    let n = v.len();

    // Pivot selection algorithm based on Java's DualPivotQuicksort.

    // Fast approximation of n / 7
    let seventh = (n / 8) + (n / 64) + 1;

    // Pick five element evenly spaced around the middle (inclusive) of the slice.
    let mut e3 = n / 2;
    let mut e2 = e3 - seventh;
    let mut e1 = e3 - 2*seventh;
    let mut e4 = e3 + seventh;
    let mut e5 = e3 + 2*seventh;

    // Sort them with a sorting network.
    unsafe {
        maybe_swap!(v, &mut e1, &mut e2, compare);
        maybe_swap!(v, &mut e4, &mut e5, compare);
        maybe_swap!(v, &mut e3, &mut e5, compare);
        maybe_swap!(v, &mut e3, &mut e4, compare);
        maybe_swap!(v, &mut e2, &mut e5, compare);
        maybe_swap!(v, &mut e1, &mut e4, compare);
        maybe_swap!(v, &mut e1, &mut e3, compare);
        maybe_swap!(v, &mut e2, &mut e4, compare);
        maybe_swap!(v, &mut e2, &mut e3, compare);
    }

    if unsafe { compare_idxs(v, e1, e2, compare) != Equal &&
                compare_idxs(v, e2, e3, compare) != Equal &&
                compare_idxs(v, e3, e4, compare) != Equal &&
                compare_idxs(v, e4, e5, compare) != Equal } {
        // No consecutive pivot candidates are the same, meaning there is some variaton.
        dual_pivot_sort(v, (e1, e2, e3, e4, e5), compare, rec, heapsort_depth);
    } else {
        // Two consecutive pivots candidates where the same.
        // There are probably many similar elements.
        single_pivot_sort(v, e3, compare, rec, heapsort_depth);
    }
}

fn maybe_insertion_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) -> bool {
    let n = v.len();
    if n <= 1 {
        return true;
    }

    let threshold = min(MAX_INSERTION_SORT_ELEMS,
                        max(MIN_QUICKSORT_ELEMS, INSERTION_SORT_FACTOR / size_of::<T>()));
    if n <= threshold {
        insertion_sort(v, compare);
        return true;
    }
    return false;
}

pub fn insertion_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    let mut i = 1;
    let n = v.len();
    while i < n {
        let mut j = i;
        while j > 0 && unsafe { compare_idxs(v, j-1, j, compare) } == Greater {
            unsafe { unsafe_swap(v, j, j-1); }
            j -= 1;
        }
        i += 1;
    }
}

fn dual_pivot_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivots: (uint, uint, uint, uint, uint),
                                                               compare: &C, rec: u32, heapsort_depth: u32) {
    let (pmin, p1, _pmid, p2, pmax) = pivots;
    let n = v.len();

    let lp = 0;
    let rp = n - 1;

    v.swap(p1, lp);
    v.swap(p2, rp);

    let mut lesser = 1;
    let mut greater = n - 2;

    unsafe {
        // Skip elements that are already in the correct position
        while compare_idxs(v, lesser, lp, compare) == Less { lesser += 1; }
        while compare_idxs(v, greater, rp, compare) == Greater { greater -= 1; }

        let mut k = lesser;
        // XXX We make some unecessary swaps since we can't leave uninitialized values
        // in `v` in case `compare` unwinds.
        while k <= greater {
            if compare_idxs(v, k, lp, compare) == Less {
                unsafe_swap(v, k, lesser);
                lesser += 1;
            } else {
                let cmp = compare_idxs(v, k, rp, compare);
                if cmp == Greater || cmp == Equal {
                    while k < greater && compare_idxs(v, greater, rp, compare) == Greater {
                        greater -= 1;
                    }
                    unsafe_swap(v, k, greater);
                    greater -= 1;
                    if compare_idxs(v, k, lp, compare) == Less {
                        unsafe_swap(v, k, lesser);
                        lesser += 1;
                    }
                }
            }
            k += 1;
        }
    }

    lesser -= 1;
    greater += 1;

    // Swap back pivots
    unsafe {
        unsafe_swap(v, lp, lesser);
        unsafe_swap(v, rp, greater);
    }

    introsort(v[mut ..lesser], compare, rec + 1, heapsort_depth);
    introsort(v[mut greater+1..], compare, rec + 1, heapsort_depth);
    introsort(v[mut left..right], compare, rec + 1, heapsort_depth);
}

fn single_pivot_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivot: uint, compare: &C, rec: u32, heapsort_depth: u32) {
    let (l, r) = fat_partition(v, pivot, compare);
    let n = v.len();
    if l > 1 {
        introsort(v[mut ..l], compare, rec + 1, heapsort_depth);
    }
    if r > 1 {
        introsort(v[mut n - r..], compare, rec + 1, heapsort_depth);
    }
}

/// Partitions elements, using the element at `pivot` as pivot.
/// After partitioning, the array looks as following:
/// <<<<<==>>>
/// Return (number of < elements, number of > elements)
fn fat_partition<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivot: uint, compare: &C) -> (uint, uint)  {
    let mut a = 0;
    let mut b = a;
    let mut c = v.len() - 1;
    let mut d = c;
    v.swap(0, pivot);
    loop {
        while b <= c {
            let r = compare_idxs_safe(v, b, 0, compare);
            if r == Greater { break; }
            if r == Equal {
                unsafe { unsafe_swap(v, a, b); }
                a += 1;
            }
            b += 1;
        }
        while c >= b {
            let r = compare_idxs_safe(v, c, 0, compare);
            if r == Less { break; }
            if r == Equal {
                unsafe { unsafe_swap(v, c, d); }
                d -= 1;
            }
            c -= 1;
        }
        if b > c { break; }
        unsafe { unsafe_swap(v, b, c); }
        b += 1;
        c -= 1;
    }

    let n = v.len();
    let l = min(a, b - a);
    unsafe { swap_many(v, 0, b - l, l); }
    let r = min(d - c, n - 1 - d);
    unsafe { swap_many(v, b, n - r, r); }

    return (b - a, d - c);
}

unsafe fn swap_many<T>(v: &mut [T], a: uint, b: uint, n: uint) {
    let mut i = 0;
    while i < n {
        unsafe_swap(v, a + i, b + i);
        i += 1;
    }
}

#[cold]
#[inline(never)]
pub fn heapsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    heapify(v, compare);
    let mut end = v.len();
    while end > 0 {
        end -= 1;
        v.swap(0, end);
        siftdown_range(v, 0, end, compare);
    }
}

#[inline]
fn heapify<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    let mut n = v.len() / 2;
    while n > 0 {
        n -= 1;
        siftdown(v, n, compare)
    }
}

#[inline]
fn siftup<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], start: uint, mut pos: uint, compare: &C) {
    unsafe {
        let new = replace(&mut v[pos], zeroed());

        while pos > start {
            let parent = (pos - 1) >> 1;
            // TODO: Get rid of transmute when high-rank lifetimes work
            if (*compare)(transmute(&new), transmute(v.unsafe_get(parent))) == Greater {
                let x = replace(&mut v[parent], zeroed());
                ptr::write(&mut v[pos], x);
                pos = parent;
                continue
            }
            break
        }
        ptr::write(&mut v[pos], new);
    }
}

#[inline]
fn siftdown_range<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], mut pos: uint, end: uint, compare: &C) {
    unsafe {
        let start = pos;
        let new = replace(&mut v[pos], zeroed());

        let mut child = 2 * pos + 1;
        while child < end {
            let right = child + 1;
            if right < end && compare_idxs(v, child, right, compare) != Greater {
                child = right;
            }
            let x = replace(&mut v[child], zeroed());
            ptr::write(&mut v[pos], x);
            pos = child;
            child = 2 * pos + 1;
        }

        ptr::write(&mut v[pos], new);
        siftup(v, start, pos, compare);
    }
}

#[inline]
fn siftdown<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pos: uint, compare: &C) {
    let len = v.len();
    siftdown_range(v, pos, len, compare);
}

fn log2(x: uint) -> u32 {
    if x <= 1 { return 0; }
    let n = if size_of::<uint>() == 8 {
        (unsafe { ::core::intrinsics::ctlz64(x as u64) }) as u32
    } else {
        unsafe { ::core::intrinsics::ctlz32(x as u32) }
    };
    size_of::<uint>() as u32 * 8 - n
}

// TODO Replace this function when unboxed closures work properly
// Blocked on https://github.com/rust-lang/rust/issues/17661
#[inline(always)]
unsafe fn compare_idxs<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &[T], a: uint, b: uint, compare: &C) -> Ordering {
    let x = v.unsafe_get(a);
    let y = v.unsafe_get(b);
    (*compare)(transmute(x), transmute(y))
}

#[inline(always)]
fn compare_idxs_safe<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &[T], a: uint, b: uint, compare: &C) -> Ordering {
    unsafe { (*compare)(transmute(&v[a]), transmute(&v[b])) }
}

#[inline(always)]
unsafe fn unsafe_swap<T>(v: &mut[T], a: uint, b: uint) {
    ptr::swap(v.unsafe_mut(a) as *mut T, v.unsafe_mut(b) as *mut T);
}
<commit_msg>Fix silly error<commit_after>use core::prelude::*;
use core::cmp::{min, max};
use core::mem::{size_of, zeroed, replace, swap, transmute};
use core::ptr;

/// The smallest number of elements that may be quicksorted.
/// Must be at least 9.
const MIN_QUICKSORT_ELEMS: uint = 10;

/// The maximum number of elements to be insertion sorted.
const MAX_INSERTION_SORT_ELEMS: uint = 42;

/// Controls the number of elements to be insertion sorted.
/// Higher values give more insertion sorted elements.
const INSERTION_SORT_FACTOR: uint = 450;

pub fn sort_by<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    if maybe_insertion_sort(v, compare) { return; }
    let heapsort_depth = (3 * log2(v.len())) / 2;
    do_introsort(v, compare, 0, heapsort_depth);
}

#[inline]
pub fn sort<T: Ord>(v: &mut [T]) {
    sort_by(v, &|&: a: &T, b| a.cmp(b));
}

fn introsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C, rec: u32, heapsort_depth: u32) {
    if maybe_insertion_sort(v, compare) { return; }
    do_introsort(v, compare, rec, heapsort_depth);
}

fn do_introsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C, rec: u32, heapsort_depth: u32) {
    // This works around some bugs with unboxed closures
    macro_rules! maybe_swap(
        ($v: expr, $a: expr, $b: expr, $compare: expr) => {
            if compare_idxs($v, *$a, *$b, $compare) == Greater {
                swap($a, $b);
            }
        }
    )

    if rec > heapsort_depth {
        heapsort(v, compare);
        return;
    }

    let n = v.len();

    // Pivot selection algorithm based on Java's DualPivotQuicksort.

    // Fast approximation of n / 7
    let seventh = (n / 8) + (n / 64) + 1;

    // Pick five element evenly spaced around the middle (inclusive) of the slice.
    let mut e3 = n / 2;
    let mut e2 = e3 - seventh;
    let mut e1 = e3 - 2*seventh;
    let mut e4 = e3 + seventh;
    let mut e5 = e3 + 2*seventh;

    // Sort them with a sorting network.
    unsafe {
        maybe_swap!(v, &mut e1, &mut e2, compare);
        maybe_swap!(v, &mut e4, &mut e5, compare);
        maybe_swap!(v, &mut e3, &mut e5, compare);
        maybe_swap!(v, &mut e3, &mut e4, compare);
        maybe_swap!(v, &mut e2, &mut e5, compare);
        maybe_swap!(v, &mut e1, &mut e4, compare);
        maybe_swap!(v, &mut e1, &mut e3, compare);
        maybe_swap!(v, &mut e2, &mut e4, compare);
        maybe_swap!(v, &mut e2, &mut e3, compare);
    }

    if unsafe { compare_idxs(v, e1, e2, compare) != Equal &&
                compare_idxs(v, e2, e3, compare) != Equal &&
                compare_idxs(v, e3, e4, compare) != Equal &&
                compare_idxs(v, e4, e5, compare) != Equal } {
        // No consecutive pivot candidates are the same, meaning there is some variaton.
        dual_pivot_sort(v, (e1, e2, e3, e4, e5), compare, rec, heapsort_depth);
    } else {
        // Two consecutive pivots candidates where the same.
        // There are probably many similar elements.
        single_pivot_sort(v, e3, compare, rec, heapsort_depth);
    }
}

fn maybe_insertion_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) -> bool {
    let n = v.len();
    if n <= 1 {
        return true;
    }

    let threshold = min(MAX_INSERTION_SORT_ELEMS,
                        max(MIN_QUICKSORT_ELEMS, INSERTION_SORT_FACTOR / size_of::<T>()));
    if n <= threshold {
        insertion_sort(v, compare);
        return true;
    }
    return false;
}

pub fn insertion_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    let mut i = 1;
    let n = v.len();
    while i < n {
        let mut j = i;
        while j > 0 && unsafe { compare_idxs(v, j-1, j, compare) } == Greater {
            unsafe { unsafe_swap(v, j, j-1); }
            j -= 1;
        }
        i += 1;
    }
}

fn dual_pivot_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivots: (uint, uint, uint, uint, uint),
                                                               compare: &C, rec: u32, heapsort_depth: u32) {
    let (_, p1, _, p2, _) = pivots;
    let n = v.len();

    let lp = 0;
    let rp = n - 1;

    v.swap(p1, lp);
    v.swap(p2, rp);

    let mut lesser = 1;
    let mut greater = n - 2;

    unsafe {
        // Skip elements that are already in the correct position
        while compare_idxs(v, lesser, lp, compare) == Less { lesser += 1; }
        while compare_idxs(v, greater, rp, compare) == Greater { greater -= 1; }

        let mut k = lesser;
        // XXX We make some unecessary swaps since we can't leave uninitialized values
        // in `v` in case `compare` unwinds.
        while k <= greater {
            if compare_idxs(v, k, lp, compare) == Less {
                unsafe_swap(v, k, lesser);
                lesser += 1;
            } else {
                let cmp = compare_idxs(v, k, rp, compare);
                if cmp == Greater || cmp == Equal {
                    while k < greater && compare_idxs(v, greater, rp, compare) == Greater {
                        greater -= 1;
                    }
                    unsafe_swap(v, k, greater);
                    greater -= 1;
                    if compare_idxs(v, k, lp, compare) == Less {
                        unsafe_swap(v, k, lesser);
                        lesser += 1;
                    }
                }
            }
            k += 1;
        }
    }

    lesser -= 1;
    greater += 1;

    // Swap back pivots
    unsafe {
        unsafe_swap(v, lp, lesser);
        unsafe_swap(v, rp, greater);
    }

    introsort(v[mut ..lesser], compare, rec + 1, heapsort_depth);
    introsort(v[mut lesser+1..greater], compare, rec + 1, heapsort_depth);
    introsort(v[mut greater+1..], compare, rec + 1, heapsort_depth);
}

fn single_pivot_sort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivot: uint, compare: &C, rec: u32, heapsort_depth: u32) {
    let (l, r) = fat_partition(v, pivot, compare);
    let n = v.len();
    if l > 1 {
        introsort(v[mut ..l], compare, rec + 1, heapsort_depth);
    }
    if r > 1 {
        introsort(v[mut n - r..], compare, rec + 1, heapsort_depth);
    }
}

/// Partitions elements, using the element at `pivot` as pivot.
/// After partitioning, the array looks as following:
/// <<<<<==>>>
/// Return (number of < elements, number of > elements)
fn fat_partition<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pivot: uint, compare: &C) -> (uint, uint)  {
    let mut a = 0;
    let mut b = a;
    let mut c = v.len() - 1;
    let mut d = c;
    v.swap(0, pivot);
    loop {
        while b <= c {
            let r = compare_idxs_safe(v, b, 0, compare);
            if r == Greater { break; }
            if r == Equal {
                unsafe { unsafe_swap(v, a, b); }
                a += 1;
            }
            b += 1;
        }
        while c >= b {
            let r = compare_idxs_safe(v, c, 0, compare);
            if r == Less { break; }
            if r == Equal {
                unsafe { unsafe_swap(v, c, d); }
                d -= 1;
            }
            c -= 1;
        }
        if b > c { break; }
        unsafe { unsafe_swap(v, b, c); }
        b += 1;
        c -= 1;
    }

    let n = v.len();
    let l = min(a, b - a);
    unsafe { swap_many(v, 0, b - l, l); }
    let r = min(d - c, n - 1 - d);
    unsafe { swap_many(v, b, n - r, r); }

    return (b - a, d - c);
}

unsafe fn swap_many<T>(v: &mut [T], a: uint, b: uint, n: uint) {
    let mut i = 0;
    while i < n {
        unsafe_swap(v, a + i, b + i);
        i += 1;
    }
}

#[cold]
#[inline(never)]
pub fn heapsort<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    heapify(v, compare);
    let mut end = v.len();
    while end > 0 {
        end -= 1;
        v.swap(0, end);
        siftdown_range(v, 0, end, compare);
    }
}

#[inline]
fn heapify<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], compare: &C) {
    let mut n = v.len() / 2;
    while n > 0 {
        n -= 1;
        siftdown(v, n, compare)
    }
}

#[inline]
fn siftup<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], start: uint, mut pos: uint, compare: &C) {
    unsafe {
        let new = replace(&mut v[pos], zeroed());

        while pos > start {
            let parent = (pos - 1) >> 1;
            // TODO: Get rid of transmute when high-rank lifetimes work
            if (*compare)(transmute(&new), transmute(v.unsafe_get(parent))) == Greater {
                let x = replace(&mut v[parent], zeroed());
                ptr::write(&mut v[pos], x);
                pos = parent;
                continue
            }
            break
        }
        ptr::write(&mut v[pos], new);
    }
}

#[inline]
fn siftdown_range<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], mut pos: uint, end: uint, compare: &C) {
    unsafe {
        let start = pos;
        let new = replace(&mut v[pos], zeroed());

        let mut child = 2 * pos + 1;
        while child < end {
            let right = child + 1;
            if right < end && compare_idxs(v, child, right, compare) != Greater {
                child = right;
            }
            let x = replace(&mut v[child], zeroed());
            ptr::write(&mut v[pos], x);
            pos = child;
            child = 2 * pos + 1;
        }

        ptr::write(&mut v[pos], new);
        siftup(v, start, pos, compare);
    }
}

#[inline]
fn siftdown<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &mut [T], pos: uint, compare: &C) {
    let len = v.len();
    siftdown_range(v, pos, len, compare);
}

fn log2(x: uint) -> u32 {
    if x <= 1 { return 0; }
    let n = if size_of::<uint>() == 8 {
        (unsafe { ::core::intrinsics::ctlz64(x as u64) }) as u32
    } else {
        unsafe { ::core::intrinsics::ctlz32(x as u32) }
    };
    size_of::<uint>() as u32 * 8 - n
}

// TODO Replace this function when unboxed closures work properly
// Blocked on https://github.com/rust-lang/rust/issues/17661
#[inline(always)]
unsafe fn compare_idxs<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &[T], a: uint, b: uint, compare: &C) -> Ordering {
    let x = v.unsafe_get(a);
    let y = v.unsafe_get(b);
    (*compare)(transmute(x), transmute(y))
}

#[inline(always)]
fn compare_idxs_safe<'a, T: 'a, C: Fn<(&'a T, &'a T), Ordering>>(v: &[T], a: uint, b: uint, compare: &C) -> Ordering {
    unsafe { (*compare)(transmute(&v[a]), transmute(&v[b])) }
}

#[inline(always)]
unsafe fn unsafe_swap<T>(v: &mut[T], a: uint, b: uint) {
    ptr::swap(v.unsafe_mut(a) as *mut T, v.unsafe_mut(b) as *mut T);
}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A tape is a symbol reader that can be rewound.
//!
//! It can be useful for matching a pattern repeatedly against an input, as the state machine may read many extra symbols
//! before it finds a match at an earlier position. To find another match after the first one, it's necessary to rewind
//! the stream.
//!

use super::symbol_reader::*;

///
/// Rewindable symbol reader
///
pub struct Tape<Symbol: Sized, SourceReader: SymbolReader<Symbol>> {
    /// Symbol reader where items not in the buffer are read from
    read_from: SourceReader,

    /// Circular buffer where stored symbols are kept
    buffer: Vec<Symbol>,

    /// Next symbol to read from the buffer
    read_index: usize,

    /// Position of the last symbol in the buffer
    last_symbol_index: usize,

    /// Position of the first symbol in the buffer
    first_symbol_index: usize,
}

impl<Symbol: Sized, SourceReader: SymbolReader<Symbol>> Tape<Symbol, SourceReader> {
    fn new(source: SourceReader) -> Tape<Symbol, SourceReader> {
        Tape { read_from: source, buffer: vec![], read_index: 0, last_symbol_index: 0, first_symbol_index: 0 }
    }
}<commit_msg>Finish up the tape implementation<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A tape is a symbol reader that can be rewound.
//!
//! It can be useful for matching a pattern repeatedly against an input, as the state machine may read many extra symbols
//! before it finds a match at an earlier position. To find another match after the first one, it's necessary to rewind
//! the stream.
//!
//! A tape implements `SymbolReader` so can be treated just like a normal symbol stream. However, it also adds two new methods.
//! `rewind` will go back a number of symbols (provided the tape contains that many symbols). `cut` removes any symbols prior
//! to the current one from the tape, making it impossible to rewind beyond the current point.
//!

use super::symbol_reader::*;

///
/// Rewindable symbol reader
///
pub struct Tape<Symbol: Sized, SourceReader: SymbolReader<Symbol>+Sized> {
    /// Symbol reader where items not in the buffer are read from
    read_from: SourceReader,

    /// Circular buffer where stored symbols are kept
    buffer: Vec<Option<Symbol>>,

    /// Next symbol to read from the buffer
    read_index: usize,

    /// Position of the last symbol in the buffer
    last_symbol_index: usize,

    /// Position of the first symbol in the buffer
    first_symbol_index: usize,

    /// True if the target reader has returned None
    end_of_reader: bool
}

impl<Symbol: Clone+Sized, SourceReader: SymbolReader<Symbol>> Tape<Symbol, SourceReader> {
    ///
    /// Creates a new tape from a symbol reader
    ///
    pub fn new(source: SourceReader) -> Tape<Symbol, SourceReader> {
        Tape { read_from: source, buffer: vec![None, None, None, None], read_index: 0, last_symbol_index: 0, first_symbol_index: 0, end_of_reader: false }
    }

    ///
    /// Resizes the buffer so that it can store at least one more symbol
    ///
    pub fn resize(&mut self) {
        let new_size        = self.buffer.len() * 2;
        let mut new_buffer  = Vec::with_capacity(new_size);

        // Copy all the occupied symbols from the old buffer
        let mut copy_index     = self.first_symbol_index;
        let mut new_read_index = 0;
        while copy_index != self.last_symbol_index {
            new_buffer.push(self.buffer[copy_index].clone());

            copy_index += 1;
            if copy_index >= self.buffer.len() { copy_index = 0; }

            if copy_index == self.read_index { new_read_index = new_buffer.len() }
        }

        // The current length is where the first 'new' symbol will go
        let new_last_symbol_index = new_buffer.len();

        // The rest of the buffer is filled with 'Nones'
        for _ in new_last_symbol_index..new_size {
            new_buffer.push(None);
        }

        // Replace the buffer
        self.buffer             = new_buffer;
        self.first_symbol_index = 0;
        self.last_symbol_index  = new_last_symbol_index;
        self.read_index         = new_read_index;
    }

    ///
    /// Trims the buffer so that it's not possible to rewind past the current read position
    ///
    pub fn cut(&mut self) {
        self.first_symbol_index = self.read_index;
    }

    ///
    /// Returns the number of symbols in the buffer
    ///
    #[inline]
    pub fn buffer_size(&mut self) -> usize {
        if self.first_symbol_index > self.last_symbol_index {
            (self.last_symbol_index + self.buffer.len()) - self.first_symbol_index
        } else {
            self.last_symbol_index - self.first_symbol_index
        }
    }

    ///
    /// Moves backwards by num_symbols
    ///
    pub fn rewind(&mut self, num_symbols: usize) {
        // Do nothing if we're over the maximum length of the buffer
        if num_symbols > self.buffer_size() {
            panic!("Can't rewind beyond the start of the tape");
        }

        // Move backwards num_symbols
        let mut new_read_index = (self.read_index+self.buffer.len()) - num_symbols;
        if new_read_index >= self.buffer.len() { new_read_index -= self.buffer.len(); }

        self.read_index = new_read_index;
    }
}

impl<Symbol: Clone+Sized, Reader: SymbolReader<Symbol>+Sized> SymbolReader<Symbol> for Tape<Symbol, Reader> {
    fn next_symbol(&mut self) -> Option<Symbol> {
        if self.read_index == self.last_symbol_index {
            // If the source reader has ended, there's nothing more to read
            if self.end_of_reader {
                return None;
            }

            // At end of buffer: need to fill it some more
            let maybe_symbol = self.read_from.next_symbol();
            match maybe_symbol {
                None => {
                    // Mark this stream as having finished (we won't read from the source stream any more)
                    self.end_of_reader = true;
                    return None;
                },

                _ => {
                    // Resize the buffer if needed
                    if (self.last_symbol_index+1)%self.buffer.len() == self.first_symbol_index {
                        self.resize();
                    }

                    // Store the next symbol (we'll immediately read it from the buffer later on)
                    self.buffer[self.last_symbol_index] = maybe_symbol;

                    self.last_symbol_index += 1;
                    if self.last_symbol_index >= self.buffer.len() { self.last_symbol_index = 0; }
                }
            }
        }

        // Result is whatever is in the buffer at the read position
        let result = self.buffer[self.read_index].clone();

        self.read_index += 1;
        if self.read_index >= self.buffer.len() { self.read_index = 0; }

        result
    }
}

#[cfg(test)]
mod test {
    use super::super::*;

    #[test]
    fn can_read_from_tape() {
        let source_vec    = vec![1,2,3,4,5,6];
        let source_stream = source_vec.read_symbols();
        let mut tape      = Tape::new(source_stream);

        assert!(tape.next_symbol() == Some(1));
        assert!(tape.next_symbol() == Some(2));
        assert!(tape.next_symbol() == Some(3));
        assert!(tape.next_symbol() == Some(4));
        assert!(tape.next_symbol() == Some(5));
        assert!(tape.next_symbol() == Some(6));
        assert!(tape.next_symbol() == None);
    }

    #[test]
    fn can_rewind_tape() {
        let source_vec    = vec![1,2,3,4,5,6];
        let source_stream = source_vec.read_symbols();
        let mut tape      = Tape::new(source_stream);

        assert!(tape.next_symbol() == Some(1));
        assert!(tape.next_symbol() == Some(2));
        assert!(tape.next_symbol() == Some(3));
        tape.rewind(3);

        assert!(tape.next_symbol() == Some(1));
        assert!(tape.next_symbol() == Some(2));
        assert!(tape.next_symbol() == Some(3));
        assert!(tape.next_symbol() == Some(4));
        assert!(tape.next_symbol() == Some(5));
        assert!(tape.next_symbol() == Some(6));
        assert!(tape.next_symbol() == None);

        tape.rewind(4);
        assert!(tape.next_symbol() == Some(3));
        assert!(tape.next_symbol() == Some(4));
        assert!(tape.next_symbol() == Some(5));
        assert!(tape.next_symbol() == Some(6));
        assert!(tape.next_symbol() == None);
    }

    #[test]
    fn can_cut_tape() {
        let source_vec    = vec![1,2,3,4,5,6,7,8,9];
        let source_stream = source_vec.read_symbols();
        let mut tape      = Tape::new(source_stream);

        assert!(tape.next_symbol() == Some(1));
        assert!(tape.next_symbol() == Some(2));
        assert!(tape.next_symbol() == Some(3));
        tape.cut();

        assert!(tape.next_symbol() == Some(4));
        assert!(tape.next_symbol() == Some(5));
        assert!(tape.next_symbol() == Some(6));
        tape.rewind(3);

        assert!(tape.next_symbol() == Some(4));
        assert!(tape.next_symbol() == Some(5));
        assert!(tape.next_symbol() == Some(6));
        assert!(tape.next_symbol() == Some(7));
        assert!(tape.next_symbol() == Some(8));
        assert!(tape.next_symbol() == Some(9));
        assert!(tape.next_symbol() == None);

        tape.rewind(3);
        tape.cut();
        assert!(tape.next_symbol() == Some(7));
        assert!(tape.next_symbol() == Some(8));
        assert!(tape.next_symbol() == Some(9));
        assert!(tape.next_symbol() == None);

        tape.rewind(2);
        assert!(tape.next_symbol() == Some(8));
        assert!(tape.next_symbol() == Some(9));
        assert!(tape.next_symbol() == None);
    }
}
<|endoftext|>


<commit_before>use crossbeam;
use floating_duration::TimeFormat;
use regex::Regex;
use std::fs::File;
use std::io::{Write, BufReader, BufRead, BufWriter};
use std::path;
use std::process;
use term;
use time;

pub trait Runnable {
    fn run(&self);
}

pub trait Task: Runnable {
    fn name(&self) -> String;
    fn should_run(&self, _: &path::Path) -> bool;
    fn start_delay(&self) -> time::Duration;
}

pub trait ShellTask: Task {
    fn command(&self) -> String;
    fn work_directory(&self) -> Option<path::PathBuf>;
    fn redirect_stdout(&self) -> Option<path::PathBuf>;
    fn redirect_stderr(&self) -> Option<path::PathBuf>;
    fn supress_stdout(&self) -> bool;
    fn supress_stderr(&self) -> bool;
}

fn handle_output<R: BufRead, W: Write>(reader: R, echo: bool, mut redirect: Option<BufWriter<W>>) {
    for line in reader.lines() {
        if !line.is_ok() {
            continue;
        }
        let line = line.unwrap();

        // TODO(sirver): My understanding is that \x1b[ should start every ANSI sequence, but I
        // also see \x1b( in my outputs - so I filtered that too. Figure out a more principled way
        // of removing color.
        lazy_static! {
            static ref RE: Regex = Regex::new("\u{1b}[\\[(][0-9;]*[a-zA-Z]").unwrap();
        }
        redirect.as_mut().map(|w| writeln!(w, "{}", RE.replace_all(&line, "")).unwrap());
        if echo {
            println!("{}", line);
        }
    }

}

impl<T: ShellTask> Runnable for T {
    /// Dispatches to 'program' with 'str'.
    fn run(&self) {
        let command = self.command();
        let args = command.split_whitespace().collect::<Vec<&str>>();
        let mut terminal = term::stdout().unwrap();
        terminal.fg(term::color::CYAN).unwrap();
        write!(terminal, "\x1b[2J").unwrap(); // Clear the screen.
        writeln!(terminal, "==> {}", self.name()).unwrap();
        terminal.reset().unwrap();
        terminal.flush().unwrap();

        let start_time = time::PreciseTime::now();

        let mut child = {
            let mut child = process::Command::new(args[0]);

            child.args(&args[1..])
                .stdin(process::Stdio::inherit())
                .stdout(process::Stdio::piped())
                .stderr(process::Stdio::piped());

            if let Some(path) = self.work_directory() {
                child.current_dir(path);
            }
            child.spawn()
                .unwrap_or_else(|e| panic!("failed to execute: {}", e))
        };

        {
            let stdout = BufReader::new(child.stdout.as_mut().unwrap());
            let echo_stdout = !self.supress_stdout();
            let redirect_stdout =
                self.redirect_stdout().map(|path| BufWriter::new(File::create(path).unwrap()));
            let stderr = BufReader::new(child.stderr.as_mut().unwrap());
            let echo_stderr = !self.supress_stderr();
            let redirect_stderr =
                self.redirect_stderr().map(|path| BufWriter::new(File::create(path).unwrap()));

            crossbeam::scope(|scope| {
                scope.spawn(move || { handle_output(stdout, echo_stdout, redirect_stdout); });
                scope.spawn(move || { handle_output(stderr, echo_stderr, redirect_stderr); });
            });
        }

        let success = match child.wait().unwrap().code() {
            Some(code) => code == 0,
            None => false,
        };

        let duration = start_time.to(time::PreciseTime::now()).to_std().unwrap();

        terminal.fg(term::color::CYAN).unwrap();
        write!(terminal, "==> {}: ", self.name()).unwrap();
        terminal.reset().unwrap();
        if success {
            terminal.fg(term::color::GREEN).unwrap();
            write!(terminal, "Success. ").unwrap();
        } else {
            terminal.fg(term::color::RED).unwrap();
            write!(terminal, "Failed. ").unwrap();
        }
        terminal.reset().unwrap();
        write!(terminal, "({})", TimeFormat(duration)).unwrap();

        writeln!(terminal, "").unwrap();
    }
}
<commit_msg>Better color regex.<commit_after>use crossbeam;
use floating_duration::TimeFormat;
use regex::Regex;
use std::fs::File;
use std::io::{Write, BufReader, BufRead, BufWriter};
use std::path;
use std::process;
use term;
use time;

pub trait Runnable {
    fn run(&self);
}

pub trait Task: Runnable {
    fn name(&self) -> String;
    fn should_run(&self, _: &path::Path) -> bool;
    fn start_delay(&self) -> time::Duration;
}

pub trait ShellTask: Task {
    fn command(&self) -> String;
    fn work_directory(&self) -> Option<path::PathBuf>;
    fn redirect_stdout(&self) -> Option<path::PathBuf>;
    fn redirect_stderr(&self) -> Option<path::PathBuf>;
    fn supress_stdout(&self) -> bool;
    fn supress_stderr(&self) -> bool;
}

fn handle_output<R: BufRead, W: Write>(reader: R, echo: bool, mut redirect: Option<BufWriter<W>>) {
    for line in reader.lines() {
        if !line.is_ok() {
            continue;
        }
        let line = line.unwrap();

        // TODO(sirver): My understanding is that \x1b[ should start every ANSI sequence, but I
        // also see \x1b( in my outputs - so I filtered that too. Figure out a more principled way
        // of removing color.
        // for now this is lifted from chalk/ansi-regex.
        lazy_static! {
            static ref RE: Regex = Regex::new("[\u{1b}\u{9b}][\\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]").unwrap();
        }
        redirect.as_mut().map(|w| writeln!(w, "{}", RE.replace_all(&line, "")).unwrap());
        if echo {
            println!("{}", line);
        }
    }

}

impl<T: ShellTask> Runnable for T {
    /// Dispatches to 'program' with 'str'.
    fn run(&self) {
        let command = self.command();
        let args = command.split_whitespace().collect::<Vec<&str>>();
        let mut terminal = term::stdout().unwrap();
        terminal.fg(term::color::CYAN).unwrap();
        write!(terminal, "\x1b[2J").unwrap(); // Clear the screen.
        writeln!(terminal, "==> {}", self.name()).unwrap();
        terminal.reset().unwrap();
        terminal.flush().unwrap();

        let start_time = time::PreciseTime::now();

        let mut child = {
            let mut child = process::Command::new(args[0]);

            child.args(&args[1..])
                .stdin(process::Stdio::inherit())
                .stdout(process::Stdio::piped())
                .stderr(process::Stdio::piped());

            if let Some(path) = self.work_directory() {
                child.current_dir(path);
            }
            child.spawn()
                .unwrap_or_else(|e| panic!("failed to execute: {}", e))
        };

        {
            let stdout = BufReader::new(child.stdout.as_mut().unwrap());
            let echo_stdout = !self.supress_stdout();
            let redirect_stdout =
                self.redirect_stdout().map(|path| BufWriter::new(File::create(path).unwrap()));
            let stderr = BufReader::new(child.stderr.as_mut().unwrap());
            let echo_stderr = !self.supress_stderr();
            let redirect_stderr =
                self.redirect_stderr().map(|path| BufWriter::new(File::create(path).unwrap()));

            crossbeam::scope(|scope| {
                scope.spawn(move || { handle_output(stdout, echo_stdout, redirect_stdout); });
                scope.spawn(move || { handle_output(stderr, echo_stderr, redirect_stderr); });
            });
        }

        let success = match child.wait().unwrap().code() {
            Some(code) => code == 0,
            None => false,
        };

        let duration = start_time.to(time::PreciseTime::now()).to_std().unwrap();

        terminal.fg(term::color::CYAN).unwrap();
        write!(terminal, "==> {}: ", self.name()).unwrap();
        terminal.reset().unwrap();
        if success {
            terminal.fg(term::color::GREEN).unwrap();
            write!(terminal, "Success. ").unwrap();
        } else {
            terminal.fg(term::color::RED).unwrap();
            write!(terminal, "Failed. ").unwrap();
        }
        terminal.reset().unwrap();
        write!(terminal, "({})", TimeFormat(duration)).unwrap();

        writeln!(terminal, "").unwrap();
    }
}
<|endoftext|>