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

extern crate ytdl;
extern crate clap;
extern crate url;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate reqwest;
extern crate pbr;
extern crate serde_json;
extern crate openssl_probe;
extern crate env_proxy;

use clap::{App, AppSettings, Arg};
use pbr::{ProgressBar, Units};
use reqwest::header::{ContentLength, Headers, Range};
use std::fs::OpenOptions;
use std::io::{Read, Write};
use std::str::FromStr;
use ytdl::format;
use ytdl::format_list::{Filter, FormatList};
use ytdl::video_info;

#[derive(Debug)]
struct Options {
    no_progress: bool,
    info_only: bool,
    silent: bool,
    debug: bool,
    append: bool,
    json: bool,
    download_url: bool,
    filter: Vec<String>,
    byte_range: String,
    output_file: String,
    start_offset: i32,
}

fn main() {
    openssl_probe::init_ssl_cert_env_vars();
    env_logger::init().unwrap();

    let flags = vec![
        Arg::with_name("output")
            .short("o")
            .long("output")
            .value_name("FILE")
            .help("write output to a file")
            .takes_value(true),
        Arg::with_name("no-progress")
            .long("no-progress")
            .help("write output to a file"),
        Arg::with_name("range")
            .short("r")
            .long("range")
            .value_name("RANGE")
            .help(
                "download a specific range of bytes of the video, [start]-[end]",
            )
            .takes_value(true),
        Arg::with_name("url")
            .help("youtube video url, short url or video id")
            .required(true)
            .index(1),
        Arg::with_name("download-url")
            .short("-u")
            .long("download-url")
            .help("prints download url to stdout"),
        Arg::with_name("json")
            .short("j")
            .long("json")
            .help("print info json to stdout"),
        Arg::with_name("debug")
            .short("d")
            .long("debug")
            .help("output debug log"),
        Arg::with_name("filter")
            .short("f")
            .long("filter")
            .value_name("FILTER")
            .multiple(true)
            .help("filter available formats, syntax: val1 val2 val3")
            .takes_value(true),
        Arg::with_name("append")
            .short("-a")
            .long("--append")
            .help("append to output file instead of overwriting"),
        Arg::with_name("start-offset")
            .long("start-offset")
            .value_name("STARTOFFSET")
            .help("offset the start of the video")
            .takes_value(true),
        Arg::with_name("silent")
            .short("s")
            .long("silent")
            .help("only output error, also diables progressbar"),
        Arg::with_name("info")
            .short("i")
            .long("info")
            .help("only output info"),
    ];

    let matches = App::new("ytdl")
        .setting(AppSettings::ArgRequiredElseHelp)
        .version("0.1.4")
        .about("download youtube videos")
        .args(&flags)
        .get_matches();

    let mut options = Options {
        no_progress: matches.is_present("no-progress"),
        info_only: matches.is_present("info"),
        silent: matches.is_present("silent"),
        debug: matches.is_present("debug"),
        append: matches.is_present("append"),
        json: matches.is_present("json"),
        download_url: matches.is_present("download-url"),
        filter: matches
            .values_of("filter")
            .unwrap_or_default()
            .map(|x| x.to_string())
            .collect(),
        output_file: matches.value_of("output").unwrap_or_default().to_string(),
        byte_range: matches.value_of("range").unwrap_or_default().to_string(),
        start_offset: matches
            .value_of("start-offset")
            .unwrap_or("0")
            .parse()
            .unwrap(),
    };

    let identifier = matches.value_of("url").unwrap_or_default();
    if options.filter.is_empty() {
        options.filter = vec![
            format!("{}:mp4", format::FORMAT_EXTENSION_KEY),
            format!("!{}:", format::FORMAT_VIDEO_ENCODING_KEY),
            format!("!{}:", format::FORMAT_AUDIO_ENCODING_KEY),
            format!("best"),
        ];
    }

    handler(identifier, &options);
}

fn handler(identifier: &str, options: &Options) {
    info!("fetching video info...");
    let info = match video_info::get_video_info(identifier) {
        Ok(i) => i,
        Err(e) => {
            println!("unable to fetch video info: {}", e.to_string());
            return;
        }
    };

    if options.info_only {
        println!("Author: {}", info.author);
        println!("Duration: {}s", info.duration);
        return;
    } else if options.json {
        println!(
            "{}",
            serde_json::to_string_pretty(&info).unwrap_or_default()
        );
        return;
    }

    let formats = filter_formats(&options.filter, &info.formats);
    if formats.is_empty() {
        println!("no formats available that match criteria");
        return;
    }

    let mut download_url = match video_info::get_download_url(&formats[0]) {
        Ok(u) => u,
        Err(e) => {
            println!("unable to get download url: {}", e.to_string());
            return;
        }
    };

    if options.start_offset != 0 {
        download_url
            .query_pairs_mut()
            .append_pair("begin", &format!("{}", &options.start_offset * 1000,));
    }

    if options.download_url {
        println!("{}", download_url.as_str());
    }

    let mut filename = if !options.output_file.is_empty() {
        format!("{}.{}", options.output_file, &formats[0].extension)
    } else {
        video_info::get_filename(&info, &formats[0])
    };
    filename = filename.replace("/", "-");

    let mut file = if options.append {
        OpenOptions::new()
            .write(true)
            .create(true)
            .append(true)
            .open(&filename)
            .expect("create output file fail")
    } else {
        OpenOptions::new()
            .write(true)
            .create(true)
            .open(&filename)
            .expect("create output file fail")
    };

    info!("download to {}", filename);
    let mut headers = Headers::new();
    if !options.byte_range.is_empty() {
        let r =
            Range::from_str(&format!("bytes={}", &options.byte_range)).expect("invalid range str");
        headers.set(r);
    }

    let client = video_info::get_client(download_url.as_str()).expect("get request client fail");
    let mut resp = client
        .get(download_url.as_str())
        .expect("download fail")
        .headers(headers)
        .send()
        .expect("download fail");

    let file_size = resp.headers()
        .get::<ContentLength>()
        .map(|l| **l)
        .unwrap_or(0);

    let mut pb = ProgressBar::new(file_size);
    pb.format("╢▌▌░╟");
    pb.set_units(Units::Bytes);
    pb.show_percent = true;
    pb.show_speed = true;
    pb.show_time_left = true;
    pb.set_width(Some(100));
    let mut buf = [0; 128 * 1024];

    loop {
        match resp.read(&mut buf) {
            Ok(len) => {
                if file.write_all(&buf[..len]).is_err() {
                    println!("write to file fail");
                    return;
                }
                if !options.silent && !options.no_progress {
                    pb.add(len as u64);
                }
                if len == 0 {
                    break;
                }
            }
            Err(_) => {
                println!("stream pipe file error");
                return;
            },
        };
    }

    println!("");
}

fn filter_formats(filters: &Vec<String>, formats: &FormatList) -> FormatList {
    let mut formats = formats.clone();
    for fi in filters {
        let filter_str = fi.as_str();
        formats = match filter_str {
            "best" | "worst" => {
                formats
                    .extremes(format::FORMAT_RESOLUTION_KEY, filter_str == "best")
                    .extremes(format::FORMAT_AUDIO_BITRATE_KEY, filter_str == "best")
            }
            "best-video" | "worst-video" => {
                formats.extremes(
                    format::FORMAT_RESOLUTION_KEY,
                    filter_str.starts_with("best"),
                )
            }
            "best-audio" | "worst-audio" => {
                formats.extremes(
                    format::FORMAT_AUDIO_BITRATE_KEY,
                    filter_str.starts_with("best"),
                )
            }
            _ => {
                let split = filter_str.splitn(2, ":").collect::<Vec<&str>>();
                if split.len() != 2 {
                    panic!("invalid filter key");
                }

                let mut key = split[0].to_string();
                let exclude = key.starts_with("!");
                let start = if exclude { 1 } else { 0 };
                let key: String = key.drain(start..).collect::<String>();
                let value = split[1].trim();
                if value == "best" || value == "worst" {
                    let f = formats.extremes(&key, value == "best");
                    formats = formats.subtract(&f);
                    return formats;
                }
                let mut vals = value.split(",").collect::<Vec<&str>>();
                for i in 0..vals.len() {
                    vals[i] = vals[i].trim();
                }
                let f = formats.filter(&key, &vals);
                if exclude {
                    formats.subtract(&f)
                } else {
                    f
                }
            }
        };
    }
    formats
}
<commit_msg>drop env logger<commit_after>#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

extern crate ytdl;
extern crate clap;
extern crate url;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate reqwest;
extern crate pbr;
extern crate serde_json;
extern crate openssl_probe;
extern crate env_proxy;

use clap::{App, AppSettings, Arg};
use pbr::{ProgressBar, Units};
use reqwest::header::{ContentLength, Headers, Range};
use std::fs::OpenOptions;
use std::io::{Read, Write};
use std::str::FromStr;
use ytdl::format;
use ytdl::format_list::{Filter, FormatList};
use ytdl::video_info;

#[derive(Debug)]
struct Options {
    no_progress: bool,
    info_only: bool,
    silent: bool,
    debug: bool,
    append: bool,
    json: bool,
    download_url: bool,
    filter: Vec<String>,
    byte_range: String,
    output_file: String,
    start_offset: i32,
}

fn main() {
    openssl_probe::init_ssl_cert_env_vars();
    drop(env_logger::init());

    let flags = vec![
        Arg::with_name("output")
            .short("o")
            .long("output")
            .value_name("FILE")
            .help("write output to a file")
            .takes_value(true),
        Arg::with_name("no-progress")
            .long("no-progress")
            .help("write output to a file"),
        Arg::with_name("range")
            .short("r")
            .long("range")
            .value_name("RANGE")
            .help(
                "download a specific range of bytes of the video, [start]-[end]",
            )
            .takes_value(true),
        Arg::with_name("url")
            .help("youtube video url, short url or video id")
            .required(true)
            .index(1),
        Arg::with_name("download-url")
            .short("-u")
            .long("download-url")
            .help("prints download url to stdout"),
        Arg::with_name("json")
            .short("j")
            .long("json")
            .help("print info json to stdout"),
        Arg::with_name("debug")
            .short("d")
            .long("debug")
            .help("output debug log"),
        Arg::with_name("filter")
            .short("f")
            .long("filter")
            .value_name("FILTER")
            .multiple(true)
            .help("filter available formats, syntax: val1 val2 val3")
            .takes_value(true),
        Arg::with_name("append")
            .short("-a")
            .long("--append")
            .help("append to output file instead of overwriting"),
        Arg::with_name("start-offset")
            .long("start-offset")
            .value_name("STARTOFFSET")
            .help("offset the start of the video")
            .takes_value(true),
        Arg::with_name("silent")
            .short("s")
            .long("silent")
            .help("only output error, also diables progressbar"),
        Arg::with_name("info")
            .short("i")
            .long("info")
            .help("only output info"),
    ];

    let matches = App::new("ytdl")
        .setting(AppSettings::ArgRequiredElseHelp)
        .version("0.1.4")
        .about("download youtube videos")
        .args(&flags)
        .get_matches();

    let mut options = Options {
        no_progress: matches.is_present("no-progress"),
        info_only: matches.is_present("info"),
        silent: matches.is_present("silent"),
        debug: matches.is_present("debug"),
        append: matches.is_present("append"),
        json: matches.is_present("json"),
        download_url: matches.is_present("download-url"),
        filter: matches
            .values_of("filter")
            .unwrap_or_default()
            .map(|x| x.to_string())
            .collect(),
        output_file: matches.value_of("output").unwrap_or_default().to_string(),
        byte_range: matches.value_of("range").unwrap_or_default().to_string(),
        start_offset: matches
            .value_of("start-offset")
            .unwrap_or("0")
            .parse()
            .unwrap(),
    };

    let identifier = matches.value_of("url").unwrap_or_default();
    if options.filter.is_empty() {
        options.filter = vec![
            format!("{}:mp4", format::FORMAT_EXTENSION_KEY),
            format!("!{}:", format::FORMAT_VIDEO_ENCODING_KEY),
            format!("!{}:", format::FORMAT_AUDIO_ENCODING_KEY),
            format!("best"),
        ];
    }

    handler(identifier, &options);
}

fn handler(identifier: &str, options: &Options) {
    info!("fetching video info...");
    let info = match video_info::get_video_info(identifier) {
        Ok(i) => i,
        Err(e) => {
            println!("unable to fetch video info: {}", e.to_string());
            return;
        }
    };

    if options.info_only {
        println!("Author: {}", info.author);
        println!("Duration: {}s", info.duration);
        return;
    } else if options.json {
        println!(
            "{}",
            serde_json::to_string_pretty(&info).unwrap_or_default()
        );
        return;
    }

    let formats = filter_formats(&options.filter, &info.formats);
    if formats.is_empty() {
        println!("no formats available that match criteria");
        return;
    }

    let mut download_url = match video_info::get_download_url(&formats[0]) {
        Ok(u) => u,
        Err(e) => {
            println!("unable to get download url: {}", e.to_string());
            return;
        }
    };

    if options.start_offset != 0 {
        download_url
            .query_pairs_mut()
            .append_pair("begin", &format!("{}", &options.start_offset * 1000,));
    }

    if options.download_url {
        println!("{}", download_url.as_str());
    }

    let mut filename = if !options.output_file.is_empty() {
        format!("{}.{}", options.output_file, &formats[0].extension)
    } else {
        video_info::get_filename(&info, &formats[0])
    };
    filename = filename.replace("/", "-");

    let mut file = if options.append {
        OpenOptions::new()
            .write(true)
            .create(true)
            .append(true)
            .open(&filename)
            .expect("create output file fail")
    } else {
        OpenOptions::new()
            .write(true)
            .create(true)
            .open(&filename)
            .expect("create output file fail")
    };

    info!("download to {}", filename);
    let mut headers = Headers::new();
    if !options.byte_range.is_empty() {
        let r =
            Range::from_str(&format!("bytes={}", &options.byte_range)).expect("invalid range str");
        headers.set(r);
    }

    let client = video_info::get_client(download_url.as_str()).expect("get request client fail");
    let mut resp = client
        .get(download_url.as_str())
        .expect("download fail")
        .headers(headers)
        .send()
        .expect("download fail");

    let file_size = resp.headers()
        .get::<ContentLength>()
        .map(|l| **l)
        .unwrap_or(0);

    let mut pb = ProgressBar::new(file_size);
    pb.format("╢▌▌░╟");
    pb.set_units(Units::Bytes);
    pb.show_percent = true;
    pb.show_speed = true;
    pb.show_time_left = true;
    pb.set_width(Some(100));
    let mut buf = [0; 128 * 1024];

    loop {
        match resp.read(&mut buf) {
            Ok(len) => {
                if file.write_all(&buf[..len]).is_err() {
                    println!("write to file fail");
                    return;
                }
                if !options.silent && !options.no_progress {
                    pb.add(len as u64);
                }
                if len == 0 {
                    break;
                }
            }
            Err(_) => {
                println!("stream pipe file error");
                return;
            }
        };
    }

    println!("");
}

fn filter_formats(filters: &Vec<String>, formats: &FormatList) -> FormatList {
    let mut formats = formats.clone();
    for fi in filters {
        let filter_str = fi.as_str();
        formats = match filter_str {
            "best" | "worst" => {
                formats
                    .extremes(format::FORMAT_RESOLUTION_KEY, filter_str == "best")
                    .extremes(format::FORMAT_AUDIO_BITRATE_KEY, filter_str == "best")
            }
            "best-video" | "worst-video" => {
                formats.extremes(
                    format::FORMAT_RESOLUTION_KEY,
                    filter_str.starts_with("best"),
                )
            }
            "best-audio" | "worst-audio" => {
                formats.extremes(
                    format::FORMAT_AUDIO_BITRATE_KEY,
                    filter_str.starts_with("best"),
                )
            }
            _ => {
                let split = filter_str.splitn(2, ":").collect::<Vec<&str>>();
                if split.len() != 2 {
                    panic!("invalid filter key");
                }

                let mut key = split[0].to_string();
                let exclude = key.starts_with("!");
                let start = if exclude { 1 } else { 0 };
                let key: String = key.drain(start..).collect::<String>();
                let value = split[1].trim();
                if value == "best" || value == "worst" {
                    let f = formats.extremes(&key, value == "best");
                    formats = formats.subtract(&f);
                    return formats;
                }
                let mut vals = value.split(",").collect::<Vec<&str>>();
                for i in 0..vals.len() {
                    vals[i] = vals[i].trim();
                }
                let f = formats.filter(&key, &vals);
                if exclude {
                    formats.subtract(&f)
                } else {
                    f
                }
            }
        };
    }
    formats
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_codegen_utils;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
extern crate log;

use rustc::session::Session;
use rustc::middle::cstore::CrateStore;
use rustc_driver::{Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_driver::driver::{CompileState, CompileController};
use rustc::session::config::{self, Input, ErrorOutputType};
use rustc::hir::{self, itemlikevisit};
use rustc::ty::TyCtxt;
use rustc_codegen_utils::codegen_backend::CodegenBackend;
use syntax::ast;
use std::path::PathBuf;

struct MiriCompilerCalls {
    default: Box<RustcDefaultCalls>,
    /// Whether to begin interpretation at the start_fn lang item or not
    ///
    /// If false, the interpretation begins at the `main` function
    start_fn: bool,
}

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default.early_callback(
            matches,
            sopts,
            cfg,
            descriptions,
            output,
        )
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default.no_input(
            matches,
            sopts,
            cfg,
            odir,
            ofile,
            descriptions,
        )
    }
    fn late_callback(
        &mut self,
        codegen_backend: &CodegenBackend,
        matches: &getopts::Matches,
        sess: &Session,
        cstore: &CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default.late_callback(codegen_backend, matches, sess, cstore, input, odir, ofile)
    }
    fn build_controller(
        self: Box<Self>,
        sess: &Session,
        matches: &getopts::Matches,
    ) -> CompileController<'a> {
        let this = *self;
        let mut control = this.default.build_controller(sess, matches);
        control.after_hir_lowering.callback = Box::new(after_hir_lowering);
        let start_fn = this.start_fn;
        control.after_analysis.callback = Box::new(move |state| after_analysis(state, start_fn));
        control.after_analysis.stop = Compilation::Stop;
        control
    }
}

fn after_hir_lowering(state: &mut CompileState) {
    let attr = (
        String::from("miri"),
        syntax::feature_gate::AttributeType::Whitelisted,
    );
    state.session.plugin_attributes.borrow_mut().push(attr);
}

fn after_analysis<'a, 'tcx>(state: &mut CompileState<'a, 'tcx>, use_start_fn: bool) {
    state.session.abort_if_errors();

    let tcx = state.tcx.unwrap();

    if std::env::args().any(|arg| arg == "--test") {
        struct Visitor<'a, 'tcx: 'a>(
            TyCtxt<'a, 'tcx, 'tcx>,
            &'a CompileState<'a, 'tcx>
        );
        impl<'a, 'tcx: 'a, 'hir> itemlikevisit::ItemLikeVisitor<'hir> for Visitor<'a, 'tcx> {
            fn visit_item(&mut self, i: &'hir hir::Item) {
                if let hir::Item_::ItemFn(.., body_id) = i.node {
                    if i.attrs.iter().any(|attr| {
                        attr.name() == "test"
                    })
                    {
                        let did = self.0.hir.body_owner_def_id(body_id);
                        println!(
                            "running test: {}",
                            self.0.def_path_debug_str(did),
                        );
                        miri::eval_main(self.0, did, None);
                        self.1.session.abort_if_errors();
                    }
                }
            }
            fn visit_trait_item(&mut self, _trait_item: &'hir hir::TraitItem) {}
            fn visit_impl_item(&mut self, _impl_item: &'hir hir::ImplItem) {}
        }
        state.hir_crate.unwrap().visit_all_item_likes(
            &mut Visitor(tcx, state),
        );
    } else if let Some((entry_node_id, _, _)) = *state.session.entry_fn.borrow() {
        let entry_def_id = tcx.hir.local_def_id(entry_node_id);
        // Use start_fn lang item if we have -Zmiri-start-fn set
        let start_wrapper = if use_start_fn {
            Some(tcx.lang_items().start_fn().unwrap())
        } else {
            None
        };
        miri::eval_main(tcx, entry_def_id, start_wrapper);

        state.session.abort_if_errors();
    } else {
        println!("no main function found, assuming auxiliary build");
    }
}

fn init_logger() {
    let format = |formatter: &mut env_logger::fmt::Formatter, record: &log::Record| {
        use std::io::Write;
        if record.level() == log::Level::Trace {
            // prepend frame number
            let indentation = log_settings::settings().indentation;
            writeln!(
                formatter,
                "{indentation}:{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.module_path().unwrap_or("<unknown module>"),
                indentation = indentation,
                text = record.args(),
            )
        } else {
            writeln!(
                formatter,
                "{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.module_path().unwrap_or("<unknown_module>"),
                text = record.args(),
            )
        }
    };

    let mut builder = env_logger::Builder::new();
    builder.format(format).filter(
        None,
        log::LevelFilter::Info,
    );

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

    builder.init();
}

fn find_sysroot() -> String {
    if let Ok(sysroot) = std::env::var("MIRI_SYSROOT") {
        return sysroot;
    }

    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect(
                    "need to specify RUST_SYSROOT env var or use rustup or multirust",
                )
                .to_owned()
        }
    }
}

fn main() {
    rustc_driver::init_rustc_env_logger();
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    let mut start_fn = false;
    args.retain(|arg| {
        if arg == "-Zmiri-start-fn" {
            start_fn = true;
            false
        } else {
            true
        }
    });

    rustc_driver::run_compiler(&args, Box::new(MiriCompilerCalls {
        default: Box::new(RustcDefaultCalls),
        start_fn,
    }), None, None);
}
<commit_msg>fix for latest rust nightly<commit_after>#![feature(rustc_private)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_codegen_utils;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
extern crate log;

use rustc::session::Session;
use rustc::middle::cstore::CrateStore;
use rustc_driver::{Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_driver::driver::{CompileState, CompileController};
use rustc::session::config::{self, Input, ErrorOutputType};
use rustc::hir::{self, itemlikevisit};
use rustc::ty::TyCtxt;
use rustc_codegen_utils::codegen_backend::CodegenBackend;
use syntax::ast;
use std::path::PathBuf;

struct MiriCompilerCalls {
    default: Box<RustcDefaultCalls>,
    /// Whether to begin interpretation at the start_fn lang item or not
    ///
    /// If false, the interpretation begins at the `main` function
    start_fn: bool,
}

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default.early_callback(
            matches,
            sopts,
            cfg,
            descriptions,
            output,
        )
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default.no_input(
            matches,
            sopts,
            cfg,
            odir,
            ofile,
            descriptions,
        )
    }
    fn late_callback(
        &mut self,
        codegen_backend: &CodegenBackend,
        matches: &getopts::Matches,
        sess: &Session,
        cstore: &CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default.late_callback(codegen_backend, matches, sess, cstore, input, odir, ofile)
    }
    fn build_controller(
        self: Box<Self>,
        sess: &Session,
        matches: &getopts::Matches,
    ) -> CompileController<'a> {
        let this = *self;
        let mut control = this.default.build_controller(sess, matches);
        control.after_hir_lowering.callback = Box::new(after_hir_lowering);
        let start_fn = this.start_fn;
        control.after_analysis.callback = Box::new(move |state| after_analysis(state, start_fn));
        control.after_analysis.stop = Compilation::Stop;
        control
    }
}

fn after_hir_lowering(state: &mut CompileState) {
    let attr = (
        String::from("miri"),
        syntax::feature_gate::AttributeType::Whitelisted,
    );
    state.session.plugin_attributes.borrow_mut().push(attr);
}

fn after_analysis<'a, 'tcx>(state: &mut CompileState<'a, 'tcx>, use_start_fn: bool) {
    state.session.abort_if_errors();

    let tcx = state.tcx.unwrap();

    if std::env::args().any(|arg| arg == "--test") {
        struct Visitor<'a, 'tcx: 'a>(
            TyCtxt<'a, 'tcx, 'tcx>,
            &'a CompileState<'a, 'tcx>
        );
        impl<'a, 'tcx: 'a, 'hir> itemlikevisit::ItemLikeVisitor<'hir> for Visitor<'a, 'tcx> {
            fn visit_item(&mut self, i: &'hir hir::Item) {
                if let hir::ItemKind::Fn(.., body_id) = i.node {
                    if i.attrs.iter().any(|attr| {
                        attr.name() == "test"
                    })
                    {
                        let did = self.0.hir.body_owner_def_id(body_id);
                        println!(
                            "running test: {}",
                            self.0.def_path_debug_str(did),
                        );
                        miri::eval_main(self.0, did, None);
                        self.1.session.abort_if_errors();
                    }
                }
            }
            fn visit_trait_item(&mut self, _trait_item: &'hir hir::TraitItem) {}
            fn visit_impl_item(&mut self, _impl_item: &'hir hir::ImplItem) {}
        }
        state.hir_crate.unwrap().visit_all_item_likes(
            &mut Visitor(tcx, state),
        );
    } else if let Some((entry_node_id, _, _)) = *state.session.entry_fn.borrow() {
        let entry_def_id = tcx.hir.local_def_id(entry_node_id);
        // Use start_fn lang item if we have -Zmiri-start-fn set
        let start_wrapper = if use_start_fn {
            Some(tcx.lang_items().start_fn().unwrap())
        } else {
            None
        };
        miri::eval_main(tcx, entry_def_id, start_wrapper);

        state.session.abort_if_errors();
    } else {
        println!("no main function found, assuming auxiliary build");
    }
}

fn init_logger() {
    let format = |formatter: &mut env_logger::fmt::Formatter, record: &log::Record| {
        use std::io::Write;
        if record.level() == log::Level::Trace {
            // prepend frame number
            let indentation = log_settings::settings().indentation;
            writeln!(
                formatter,
                "{indentation}:{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.module_path().unwrap_or("<unknown module>"),
                indentation = indentation,
                text = record.args(),
            )
        } else {
            writeln!(
                formatter,
                "{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.module_path().unwrap_or("<unknown_module>"),
                text = record.args(),
            )
        }
    };

    let mut builder = env_logger::Builder::new();
    builder.format(format).filter(
        None,
        log::LevelFilter::Info,
    );

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

    builder.init();
}

fn find_sysroot() -> String {
    if let Ok(sysroot) = std::env::var("MIRI_SYSROOT") {
        return sysroot;
    }

    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect(
                    "need to specify RUST_SYSROOT env var or use rustup or multirust",
                )
                .to_owned()
        }
    }
}

fn main() {
    rustc_driver::init_rustc_env_logger();
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    let mut start_fn = false;
    args.retain(|arg| {
        if arg == "-Zmiri-start-fn" {
            start_fn = true;
            false
        } else {
            true
        }
    });

    rustc_driver::run_compiler(&args, Box::new(MiriCompilerCalls {
        default: Box::new(RustcDefaultCalls),
        start_fn,
    }), None, None);
}
<|endoftext|>


<commit_before>extern crate nss;

use nss::crypto::symm;

fn doit() -> Result<(), String>
{
    let mut crypt = try!(symm::Crypter::new(symm::AES_128_CBC, true));
    let key : &[u8] = [0xe8, 0xa7, 0x7c, 0xe2, 0x05, 0x63, 0x6a, 0x31,0,0,0,0,0,0,0,0];
    let iv : &[u8] = [0xe4, 0xbb, 0x3b, 0xd3, 0xc3, 0x71, 0x2e, 0x58,0,0,0,0,0,0,0,0];
    let message = b"Encrypt me!\0";

    println!("input:     {}", String::from_utf8_lossy(message));

    try!(crypt.init(symm::Encrypt, key, iv));
    let mut out = try!(crypt.update(message));
    out.extend(try!(crypt.final()).into_iter());

    print!("encrypted:");
    for b in out.iter()
    {
        print!(" {:02x}", *b);
    }
    println!("");

    try!(crypt.init(symm::Decrypt, key, iv));
    let mut dec = try!(crypt.update(out.as_slice()));
    dec.extend(try!(crypt.final()).into_iter());

    println!("decrypted: {}", String::from_utf8_lossy(dec.as_slice()));

    Ok(())
}

fn main()
{
    nss::init().unwrap();
    let result = doit();
    let message = result.err().unwrap_or("Ok!".to_string());
    println!("{}", message);
    let _ = nss::close();
}
<commit_msg>Styling<commit_after>extern crate nss;

use nss::crypto::symm;

fn doit() -> Result<(), String>
{
    let mut crypt = try!(symm::Crypter::new(symm::AES_128_CBC, true));
    let key : &[u8] = [0xe8, 0xa7, 0x7c, 0xe2, 0x05, 0x63, 0x6a, 0x31, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
    let iv : &[u8] = [0xe4, 0xbb, 0x3b, 0xd3, 0xc3, 0x71, 0x2e, 0x58, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
    let message = b"Encrypt me!\0";

    println!("input:     {}", String::from_utf8_lossy(message));

    try!(crypt.init(symm::Encrypt, key, iv));
    let mut out = try!(crypt.update(message));
    out.extend(try!(crypt.final()).into_iter());

    print!("encrypted:");
    for b in out.iter()
    {
        print!(" {:02x}", *b);
    }
    println!("");

    try!(crypt.init(symm::Decrypt, key, iv));
    let mut dec = try!(crypt.update(out.as_slice()));
    dec.extend(try!(crypt.final()).into_iter());

    println!("decrypted: {}", String::from_utf8_lossy(dec.as_slice()));

    Ok(())
}

fn main()
{
    nss::init().unwrap();
    let result = doit();
    let message = result.err().unwrap_or("Ok!".to_string());
    println!("{}", message);
    let _ = nss::close();
}
<|endoftext|>


<commit_before>//! Bit level parsers
//!

#[macro_use]
mod macros;

pub mod complete;
pub mod streaming;

use crate::error::{ErrorKind, ParseError};
use crate::internal::{Err, IResult, Needed};
use crate::lib::std::ops::RangeFrom;
use crate::traits::{ErrorConvert, Slice};

/// Converts a byte-level input to a bit-level input, for consumption by a parser that uses bits.
///
/// Afterwards, the input is converted back to a byte-level parser, with any remaining bits thrown
/// away.
///
/// # Example
/// ```ignore
/// # #[macro_use] extern crate nom;
/// # use nom::IResult;
/// use nom::bits::bits;
/// use nom::bits::complete::take;
/// use nom::error::Error;
///
/// fn take_4_bits(input: &[u8]) -> IResult<&[u8], u64> {
///   bits(take::<_, _, _, Error<_>>(4usize))(input)
/// }
///
/// let input = vec![0xAB, 0xCD, 0xEF, 0x12];
/// let sl    = &input[..];
///
/// assert_eq!(take_4_bits( sl ), Ok( (&sl[1..], 0xA) ));
/// ```
pub fn bits<I, O, E1: ParseError<(I, usize)> + ErrorConvert<E2>, E2: ParseError<I>, P>(
  mut parser: P,
) -> impl FnMut(I) -> IResult<I, O, E2>
where
  I: Slice<RangeFrom<usize>>,
  P: FnMut((I, usize)) -> IResult<(I, usize), O, E1>,
{
  move |input: I| match parser((input, 0)) {
    Ok(((rest, offset), res)) => {
      let byte_index = offset / 8 + if offset % 8 == 0 { 0 } else { 1 };
      Ok((rest.slice(byte_index..), res))
    }
    Err(Err::Incomplete(n)) => Err(Err::Incomplete(n.map(|u| u.get() / 8 + 1))),
    Err(Err::Error(e)) => Err(Err::Error(e.convert())),
    Err(Err::Failure(e)) => Err(Err::Failure(e.convert())),
  }
}

#[doc(hidden)]
pub fn bitsc<I, O, E1: ParseError<(I, usize)> + ErrorConvert<E2>, E2: ParseError<I>, P>(
  input: I,
  parser: P,
) -> IResult<I, O, E2>
where
  I: Slice<RangeFrom<usize>>,
  P: FnMut((I, usize)) -> IResult<(I, usize), O, E1>,
{
  bits(parser)(input)
}

/// Counterpart to `bits`, `bytes` transforms its bit stream input into a byte slice for the underlying
/// parser, allowing byte-slice parsers to work on bit streams.
///
/// A partial byte remaining in the input will be ignored and the given parser will start parsing
/// at the next full byte.
///
/// ```ignore
/// # #[macro_use] extern crate nom;
/// # use nom::IResult;
/// # use nom::combinator::rest;
/// # use nom::sequence::tuple;
/// use nom::bits::{bits, bytes, streaming::take_bits};
///
/// fn parse(input: &[u8]) -> IResult<&[u8], (u8, u8, &[u8])> {
///   bits(tuple((
///     take_bits(4usize),
///     take_bits(8usize),
///     bytes(rest)
///   )))(input)
/// }
///
/// let input = &[0xde, 0xad, 0xbe, 0xaf];
///
/// assert_eq!(parse( input ), Ok(( &[][..], (0xd, 0xea, &[0xbe, 0xaf][..]) )));
/// ```
pub fn bytes<I, O, E1: ParseError<I> + ErrorConvert<E2>, E2: ParseError<(I, usize)>, P>(
  mut parser: P,
) -> impl FnMut((I, usize)) -> IResult<(I, usize), O, E2>
where
  I: Slice<RangeFrom<usize>> + Clone,
  P: FnMut(I) -> IResult<I, O, E1>,
{
  move |(input, offset): (I, usize)| {
    let inner = if offset % 8 != 0 {
      input.slice((1 + offset / 8)..)
    } else {
      input.slice((offset / 8)..)
    };
    let i = (input, offset);
    match parser(inner) {
      Ok((rest, res)) => Ok(((rest, 0), res)),
      Err(Err::Incomplete(Needed::Unknown)) => Err(Err::Incomplete(Needed::Unknown)),
      Err(Err::Incomplete(Needed::Size(sz))) => Err(match sz.get().checked_mul(8) {
        Some(v) => Err::Incomplete(Needed::new(v)),
        None => Err::Failure(E2::from_error_kind(i, ErrorKind::TooLarge)),
      }),
      Err(Err::Error(e)) => Err(Err::Error(e.convert())),
      Err(Err::Failure(e)) => Err(Err::Failure(e.convert())),
    }
  }
}

#[doc(hidden)]
pub fn bytesc<I, O, E1: ParseError<I> + ErrorConvert<E2>, E2: ParseError<(I, usize)>, P>(
  input: (I, usize),
  parser: P,
) -> IResult<(I, usize), O, E2>
where
  I: Slice<RangeFrom<usize>> + Clone,
  P: FnMut(I) -> IResult<I, O, E1>,
{
  bytes(parser)(input)
}
<commit_msg>Fix bits::bits example and add tests<commit_after>//! Bit level parsers
//!

#[macro_use]
mod macros;

pub mod complete;
pub mod streaming;

use crate::error::{ErrorKind, ParseError};
use crate::internal::{Err, IResult, Needed};
use crate::lib::std::ops::RangeFrom;
use crate::traits::{ErrorConvert, Slice};

/// Converts a byte-level input to a bit-level input, for consumption by a parser that uses bits.
///
/// Afterwards, the input is converted back to a byte-level parser, with any remaining bits thrown
/// away.
///
/// # Example
/// ```
/// use nom::bits::{bits, streaming::take};
/// use nom::error::Error;
/// use nom::sequence::tuple;
/// use nom::IResult;
///
/// fn parse(input: &[u8]) -> IResult<&[u8], (u8, u8)> {
///     bits::<_, _, Error<(&[u8], usize)>, _, _>(tuple((take(4usize), take(8usize))))(input)
/// }
///
/// let input = &[0x12, 0x34, 0xff, 0xff];
///
/// let output = parse(input).expect("We take 1.5 bytes and the input is longer than 2 bytes");
///
/// // The first byte is consumed, the second byte is partially consumed and dropped.
/// let remaining = output.0;
/// assert_eq!(remaining, [0xff, 0xff]);
///
/// let parsed = output.1;
/// assert_eq!(parsed.0, 0x01);
/// assert_eq!(parsed.1, 0x23);
/// ```
pub fn bits<I, O, E1, E2, P>(mut parser: P) -> impl FnMut(I) -> IResult<I, O, E2>
where
  E1: ParseError<(I, usize)> + ErrorConvert<E2>,
  E2: ParseError<I>,
  I: Slice<RangeFrom<usize>>,
  P: FnMut((I, usize)) -> IResult<(I, usize), O, E1>,
{
  move |input: I| match parser((input, 0)) {
    Ok(((rest, offset), result)) => {
      // If the next byte has been partially read, it will be sliced away as well.
      // The parser functions might already slice away all fully read bytes.
      // That's why `offset / 8` isn't necessarily needed at all times.
      let remaining_bytes_index = offset / 8 + if offset % 8 == 0 { 0 } else { 1 };
      Ok((rest.slice(remaining_bytes_index..), result))
    }
    Err(Err::Incomplete(n)) => Err(Err::Incomplete(n.map(|u| u.get() / 8 + 1))),
    Err(Err::Error(e)) => Err(Err::Error(e.convert())),
    Err(Err::Failure(e)) => Err(Err::Failure(e.convert())),
  }
}

#[doc(hidden)]
pub fn bitsc<I, O, E1: ParseError<(I, usize)> + ErrorConvert<E2>, E2: ParseError<I>, P>(
  input: I,
  parser: P,
) -> IResult<I, O, E2>
where
  I: Slice<RangeFrom<usize>>,
  P: FnMut((I, usize)) -> IResult<(I, usize), O, E1>,
{
  bits(parser)(input)
}

/// Counterpart to `bits`, `bytes` transforms its bit stream input into a byte slice for the underlying
/// parser, allowing byte-slice parsers to work on bit streams.
///
/// A partial byte remaining in the input will be ignored and the given parser will start parsing
/// at the next full byte.
///
/// ```ignore
/// # #[macro_use] extern crate nom;
/// # use nom::IResult;
/// # use nom::combinator::rest;
/// # use nom::sequence::tuple;
/// use nom::bits::{bits, bytes, streaming::take_bits};
///
/// fn parse(input: &[u8]) -> IResult<&[u8], (u8, u8, &[u8])> {
///   bits(tuple((
///     take_bits(4usize),
///     take_bits(8usize),
///     bytes(rest)
///   )))(input)
/// }
///
/// let input = &[0xde, 0xad, 0xbe, 0xaf];
///
/// assert_eq!(parse( input ), Ok(( &[][..], (0xd, 0xea, &[0xbe, 0xaf][..]) )));
/// ```
pub fn bytes<I, O, E1: ParseError<I> + ErrorConvert<E2>, E2: ParseError<(I, usize)>, P>(
  mut parser: P,
) -> impl FnMut((I, usize)) -> IResult<(I, usize), O, E2>
where
  I: Slice<RangeFrom<usize>> + Clone,
  P: FnMut(I) -> IResult<I, O, E1>,
{
  move |(input, offset): (I, usize)| {
    let inner = if offset % 8 != 0 {
      input.slice((1 + offset / 8)..)
    } else {
      input.slice((offset / 8)..)
    };
    let i = (input, offset);
    match parser(inner) {
      Ok((rest, res)) => Ok(((rest, 0), res)),
      Err(Err::Incomplete(Needed::Unknown)) => Err(Err::Incomplete(Needed::Unknown)),
      Err(Err::Incomplete(Needed::Size(sz))) => Err(match sz.get().checked_mul(8) {
        Some(v) => Err::Incomplete(Needed::new(v)),
        None => Err::Failure(E2::from_error_kind(i, ErrorKind::TooLarge)),
      }),
      Err(Err::Error(e)) => Err(Err::Error(e.convert())),
      Err(Err::Failure(e)) => Err(Err::Failure(e.convert())),
    }
  }
}

#[doc(hidden)]
pub fn bytesc<I, O, E1: ParseError<I> + ErrorConvert<E2>, E2: ParseError<(I, usize)>, P>(
  input: (I, usize),
  parser: P,
) -> IResult<(I, usize), O, E2>
where
  I: Slice<RangeFrom<usize>> + Clone,
  P: FnMut(I) -> IResult<I, O, E1>,
{
  bytes(parser)(input)
}

#[cfg(test)]
mod test {
  use super::*;
  use crate::bits::streaming::take;
  use crate::error::Error;
  use crate::sequence::tuple;

  #[test]
  /// Take the `bits` function and assert that remaining bytes are correctly returned, if the
  /// previous bytes are fully consumed
  fn test_complete_byte_consumption_bits() {
    let input = &[0x12, 0x34, 0x56, 0x78];

    // Take 3 bit slices with sizes [4, 8, 4].
    let result: IResult<&[u8], (u8, u8, u8)> =
      bits::<_, _, Error<(&[u8], usize)>, _, _>(tuple((take(4usize), take(8usize), take(4usize))))(
        input,
      );

    let output = result.expect("We take 2 bytes and the input is longer than 2 bytes");

    let remaining = output.0;
    assert_eq!(remaining, [0x56, 0x78]);

    let parsed = output.1;
    assert_eq!(parsed.0, 0x01);
    assert_eq!(parsed.1, 0x23);
    assert_eq!(parsed.2, 0x04);
  }

  #[test]
  /// Take the `bits` function and assert that remaining bytes are correctly returned, if the
  /// previous bytes are NOT fully consumed. Partially consumed bytes are supposed to be dropped.
  /// I.e. if we consume 1.5 bytes of 4 bytes, 2 bytes will be returned, bits 13-16 will be
  /// dropped.
  fn test_partial_byte_consumption_bits() {
    let input = &[0x12, 0x34, 0x56, 0x78];

    // Take bit slices with sizes [4, 8].
    let result: IResult<&[u8], (u8, u8)> =
      bits::<_, _, Error<(&[u8], usize)>, _, _>(tuple((take(4usize), take(8usize))))(input);

    let output = result.expect("We take 1.5 bytes and the input is longer than 2 bytes");

    let remaining = output.0;
    assert_eq!(remaining, [0x56, 0x78]);

    let parsed = output.1;
    assert_eq!(parsed.0, 0x01);
    assert_eq!(parsed.1, 0x23);
  }

  #[test]
  #[cfg(feature = "std")]
  /// Ensure that in Incomplete error is thrown, if too few bytes are passed for a given parser.
  fn test_incomplete_bits() {
    let input = &[0x12];

    // Take bit slices with sizes [4, 8].
    let result: IResult<&[u8], (u8, u8)> =
      bits::<_, _, Error<(&[u8], usize)>, _, _>(tuple((take(4usize), take(8usize))))(input);

    assert!(result.is_err());
    let error = result.err().unwrap();
    assert_eq!("Parsing requires 2 bytes/chars", error.to_string());
  }
}
<|endoftext|>


<commit_before>use {alloc, Bytes, SeqByteStr, MAX_CAPACITY};
use traits::{Buf, MutBuf, MutBufExt, ByteStr};
use std::{cmp, ptr};

/*
 *
 * ===== ByteBuf =====
 *
 */

/// A `Buf` backed by a contiguous region of memory.
///
/// This `Buf` is better suited for cases where there is a clear delineation
/// between reading and writing.
pub struct ByteBuf {
    mem: alloc::MemRef,
    cap: u32,
    pos: u32,
    lim: u32,
    mark: Option<u32>,
}

impl ByteBuf {
    /// Create a new `ByteBuf` by copying the contents of the given slice.
    pub fn from_slice(bytes: &[u8]) -> ByteBuf {
        let mut buf = ByteBuf::mut_with_capacity(bytes.len());
        buf.write(bytes).ok().expect("unexpected failure");
        buf.flip()
    }

    pub fn mut_with_capacity(capacity: usize) -> MutByteBuf {
        assert!(capacity <= MAX_CAPACITY);
        MutByteBuf { buf: ByteBuf::new(capacity as u32) }
    }

    pub fn none() -> ByteBuf {
        ByteBuf {
            mem: alloc::MemRef::none(),
            cap: 0,
            pos: 0,
            lim: 0,
            mark: None,
        }
    }

    pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32, lim: u32) -> ByteBuf {
        debug_assert!(pos <= lim && lim <= cap, "invalid arguments; cap={}; pos={}; lim={}", cap, pos, lim);

        ByteBuf {
            mem: mem,
            cap: cap,
            pos: pos,
            lim: lim,
            mark: None,
        }
    }

    fn new(mut capacity: u32) -> ByteBuf {
        // Handle 0 capacity case
        if capacity == 0 {
            return ByteBuf::none();
        }

        // Round the capacity to the closest power of 2
        capacity = capacity.next_power_of_two();

        // Allocate the memory
        let mem = alloc::heap(capacity as usize);

        // If the allocation failed, return a blank buf
        if mem.is_none() {
            return ByteBuf::none();
        }

        ByteBuf {
            mem: mem,
            cap: capacity,
            pos: 0,
            lim: capacity,
            mark: None,
        }
    }

    pub fn capacity(&self) -> usize {
        self.cap as usize
    }

    pub fn flip(self) -> MutByteBuf {
        let mut buf = MutByteBuf { buf: self };
        buf.clear();
        buf
    }

    /// Flips the buffer back to mutable, resetting the write position
    /// to the byte after the previous write.
    pub fn resume(mut self) -> MutByteBuf {
        self.pos = self.lim;
        self.lim = self.cap;
        MutByteBuf { buf: self }
    }

    pub fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        let len = cmp::min(dst.len(), self.remaining());
        let cnt = len as u32;

        unsafe {
            ptr::copy_nonoverlapping(
                self.mem.ptr().offset(self.pos as isize),
                dst.as_mut_ptr(),
                len);
        }

        self.pos += cnt;
        len
    }

    pub fn to_seq_byte_str(self) -> SeqByteStr {
        unsafe {
            let ByteBuf { mem, pos, lim, .. } = self;
            SeqByteStr::from_mem_ref(
                mem, pos, lim - pos)
        }
    }

    #[inline]
    pub fn to_bytes(self) -> Bytes {
        Bytes::of(self.to_seq_byte_str())
    }

    /// Marks the current read location.
    ///
    /// Together with `reset`, this can be used to read from a section of the
    /// buffer multiple times. The marked location will be cleared when the
    /// buffer is flipped.
    pub fn mark(&mut self) {
        self.mark = Some(self.pos);
    }

    /// Resets the read position to the previously marked position.
    ///
    /// Together with `mark`, this can be used to read from a section of the
    /// buffer multiple times.
    ///
    /// # Panics
    ///
    /// This method will panic if no mark has been set.
    pub fn reset(&mut self) {
        self.pos = self.mark.take().expect("no mark set");
    }

    #[inline]
    fn pos(&self) -> usize {
        self.pos as usize
    }

    #[inline]
    fn lim(&self) -> usize {
        self.lim as usize
    }

    #[inline]
    fn remaining_u32(&self) -> u32 {
        self.lim - self.pos
    }
}

impl Buf for ByteBuf {

    #[inline]
    fn remaining(&self) -> usize {
        self.remaining_u32() as usize
    }

    #[inline]
    fn bytes<'a>(&'a self) -> &'a [u8] {
        &self.mem.bytes()[self.pos()..self.lim()]
    }

    #[inline]
    fn advance(&mut self, mut cnt: usize) {
        cnt = cmp::min(cnt, self.remaining());
        self.pos += cnt as u32;
    }

    #[inline]
    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        ByteBuf::read_slice(self, dst)
    }
}

/*
 *
 * ===== ROByteBuf =====
 *
 */

/// Same as `ByteBuf` but cannot be flipped to a `MutByteBuf`.
pub struct ROByteBuf {
    buf: ByteBuf,
}

impl ROByteBuf {
    pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32, lim: u32) -> ROByteBuf {
        ROByteBuf {
            buf: ByteBuf::from_mem_ref(mem, cap, pos, lim)
        }
    }

    pub fn to_seq_byte_str(self) -> SeqByteStr {
        self.buf.to_seq_byte_str()
    }

    pub fn to_bytes(self) -> Bytes {
        self.buf.to_bytes()
    }

    /// Marks the current read location.
    ///
    /// Together with `reset`, this can be used to read from a section of the
    /// buffer multiple times.
    pub fn mark(&mut self) {
        self.buf.mark = Some(self.buf.pos);
    }

    /// Resets the read position to the previously marked position.
    ///
    /// Together with `mark`, this can be used to read from a section of the
    /// buffer multiple times.
    ///
    /// # Panics
    ///
    /// This method will panic if no mark has been set.
    pub fn reset(&mut self) {
        self.buf.pos = self.buf.mark.take().expect("no mark set");
    }
}

impl Buf for ROByteBuf {

    fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    fn bytes<'a>(&'a self) -> &'a [u8] {
        self.buf.bytes()
    }

    fn advance(&mut self, cnt: usize) {
        self.buf.advance(cnt)
    }

    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        self.buf.read_slice(dst)
    }
}

/*
 *
 * ===== MutByteBuf =====
 *
 */

pub struct MutByteBuf {
    buf: ByteBuf,
}

impl MutByteBuf {
    pub fn capacity(&self) -> usize {
        self.buf.capacity() as usize
    }

    pub fn flip(self) -> ByteBuf {
        let mut buf = self.buf;

        buf.lim = buf.pos;
        buf.pos = 0;
        buf
    }

    pub fn clear(&mut self) {
        self.buf.pos = 0;
        self.buf.lim = self.buf.cap;
    }

    #[inline]
    pub fn write_slice(&mut self, src: &[u8]) -> usize {
        let cnt = src.len() as u32;
        let rem = self.buf.remaining_u32();

        if rem < cnt {
            self.write_ptr(src.as_ptr(), rem)
        } else {
            self.write_ptr(src.as_ptr(), cnt)
        }
    }

    #[inline]
    fn write_ptr(&mut self, src: *const u8, len: u32) -> usize {
        unsafe {
            ptr::copy_nonoverlapping(
                src,
                self.buf.mem.ptr().offset(self.buf.pos as isize),
                len as usize);

            self.buf.pos += len;
            len as usize
        }
    }

    pub fn bytes<'a>(&'a self) -> &'a [u8] {
        &self.buf.mem.bytes()[..self.buf.pos()]
    }
}

impl MutBuf for MutByteBuf {
    fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    fn advance(&mut self, cnt: usize) {
        self.buf.advance(cnt)
    }

    unsafe fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8] {
        let pos = self.buf.pos();
        let lim = self.buf.lim();
        &mut self.buf.mem.bytes_mut()[pos..lim]
    }
}
<commit_msg>Implement Debug for the various ByteBuf types<commit_after>use {alloc, Bytes, SeqByteStr, MAX_CAPACITY};
use traits::{Buf, MutBuf, MutBufExt, ByteStr};
use std::{cmp, fmt, ptr};

/*
 *
 * ===== ByteBuf =====
 *
 */

/// A `Buf` backed by a contiguous region of memory.
///
/// This `Buf` is better suited for cases where there is a clear delineation
/// between reading and writing.
pub struct ByteBuf {
    mem: alloc::MemRef,
    cap: u32,
    pos: u32,
    lim: u32,
    mark: Option<u32>,
}

impl ByteBuf {
    /// Create a new `ByteBuf` by copying the contents of the given slice.
    pub fn from_slice(bytes: &[u8]) -> ByteBuf {
        let mut buf = ByteBuf::mut_with_capacity(bytes.len());
        buf.write(bytes).ok().expect("unexpected failure");
        buf.flip()
    }

    pub fn mut_with_capacity(capacity: usize) -> MutByteBuf {
        assert!(capacity <= MAX_CAPACITY);
        MutByteBuf { buf: ByteBuf::new(capacity as u32) }
    }

    pub fn none() -> ByteBuf {
        ByteBuf {
            mem: alloc::MemRef::none(),
            cap: 0,
            pos: 0,
            lim: 0,
            mark: None,
        }
    }

    pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32, lim: u32) -> ByteBuf {
        debug_assert!(pos <= lim && lim <= cap, "invalid arguments; cap={}; pos={}; lim={}", cap, pos, lim);

        ByteBuf {
            mem: mem,
            cap: cap,
            pos: pos,
            lim: lim,
            mark: None,
        }
    }

    fn new(mut capacity: u32) -> ByteBuf {
        // Handle 0 capacity case
        if capacity == 0 {
            return ByteBuf::none();
        }

        // Round the capacity to the closest power of 2
        capacity = capacity.next_power_of_two();

        // Allocate the memory
        let mem = alloc::heap(capacity as usize);

        // If the allocation failed, return a blank buf
        if mem.is_none() {
            return ByteBuf::none();
        }

        ByteBuf {
            mem: mem,
            cap: capacity,
            pos: 0,
            lim: capacity,
            mark: None,
        }
    }

    pub fn capacity(&self) -> usize {
        self.cap as usize
    }

    pub fn flip(self) -> MutByteBuf {
        let mut buf = MutByteBuf { buf: self };
        buf.clear();
        buf
    }

    /// Flips the buffer back to mutable, resetting the write position
    /// to the byte after the previous write.
    pub fn resume(mut self) -> MutByteBuf {
        self.pos = self.lim;
        self.lim = self.cap;
        MutByteBuf { buf: self }
    }

    pub fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        let len = cmp::min(dst.len(), self.remaining());
        let cnt = len as u32;

        unsafe {
            ptr::copy_nonoverlapping(
                self.mem.ptr().offset(self.pos as isize),
                dst.as_mut_ptr(),
                len);
        }

        self.pos += cnt;
        len
    }

    pub fn to_seq_byte_str(self) -> SeqByteStr {
        unsafe {
            let ByteBuf { mem, pos, lim, .. } = self;
            SeqByteStr::from_mem_ref(
                mem, pos, lim - pos)
        }
    }

    #[inline]
    pub fn to_bytes(self) -> Bytes {
        Bytes::of(self.to_seq_byte_str())
    }

    /// Marks the current read location.
    ///
    /// Together with `reset`, this can be used to read from a section of the
    /// buffer multiple times. The marked location will be cleared when the
    /// buffer is flipped.
    pub fn mark(&mut self) {
        self.mark = Some(self.pos);
    }

    /// Resets the read position to the previously marked position.
    ///
    /// Together with `mark`, this can be used to read from a section of the
    /// buffer multiple times.
    ///
    /// # Panics
    ///
    /// This method will panic if no mark has been set.
    pub fn reset(&mut self) {
        self.pos = self.mark.take().expect("no mark set");
    }

    #[inline]
    fn pos(&self) -> usize {
        self.pos as usize
    }

    #[inline]
    fn lim(&self) -> usize {
        self.lim as usize
    }

    #[inline]
    fn remaining_u32(&self) -> u32 {
        self.lim - self.pos
    }
}

impl Buf for ByteBuf {

    #[inline]
    fn remaining(&self) -> usize {
        self.remaining_u32() as usize
    }

    #[inline]
    fn bytes<'a>(&'a self) -> &'a [u8] {
        &self.mem.bytes()[self.pos()..self.lim()]
    }

    #[inline]
    fn advance(&mut self, mut cnt: usize) {
        cnt = cmp::min(cnt, self.remaining());
        self.pos += cnt as u32;
    }

    #[inline]
    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        ByteBuf::read_slice(self, dst)
    }
}

impl fmt::Debug for ByteBuf {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.bytes().fmt(fmt)
    }
}

/*
 *
 * ===== ROByteBuf =====
 *
 */

/// Same as `ByteBuf` but cannot be flipped to a `MutByteBuf`.
pub struct ROByteBuf {
    buf: ByteBuf,
}

impl ROByteBuf {
    pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32, lim: u32) -> ROByteBuf {
        ROByteBuf {
            buf: ByteBuf::from_mem_ref(mem, cap, pos, lim)
        }
    }

    pub fn to_seq_byte_str(self) -> SeqByteStr {
        self.buf.to_seq_byte_str()
    }

    pub fn to_bytes(self) -> Bytes {
        self.buf.to_bytes()
    }

    /// Marks the current read location.
    ///
    /// Together with `reset`, this can be used to read from a section of the
    /// buffer multiple times.
    pub fn mark(&mut self) {
        self.buf.mark = Some(self.buf.pos);
    }

    /// Resets the read position to the previously marked position.
    ///
    /// Together with `mark`, this can be used to read from a section of the
    /// buffer multiple times.
    ///
    /// # Panics
    ///
    /// This method will panic if no mark has been set.
    pub fn reset(&mut self) {
        self.buf.pos = self.buf.mark.take().expect("no mark set");
    }
}

impl Buf for ROByteBuf {

    fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    fn bytes<'a>(&'a self) -> &'a [u8] {
        self.buf.bytes()
    }

    fn advance(&mut self, cnt: usize) {
        self.buf.advance(cnt)
    }

    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        self.buf.read_slice(dst)
    }
}

impl fmt::Debug for ROByteBuf {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.bytes().fmt(fmt)
    }
}

/*
 *
 * ===== MutByteBuf =====
 *
 */

pub struct MutByteBuf {
    buf: ByteBuf,
}

impl MutByteBuf {
    pub fn capacity(&self) -> usize {
        self.buf.capacity() as usize
    }

    pub fn flip(self) -> ByteBuf {
        let mut buf = self.buf;

        buf.lim = buf.pos;
        buf.pos = 0;
        buf
    }

    pub fn clear(&mut self) {
        self.buf.pos = 0;
        self.buf.lim = self.buf.cap;
    }

    #[inline]
    pub fn write_slice(&mut self, src: &[u8]) -> usize {
        let cnt = src.len() as u32;
        let rem = self.buf.remaining_u32();

        if rem < cnt {
            self.write_ptr(src.as_ptr(), rem)
        } else {
            self.write_ptr(src.as_ptr(), cnt)
        }
    }

    #[inline]
    fn write_ptr(&mut self, src: *const u8, len: u32) -> usize {
        unsafe {
            ptr::copy_nonoverlapping(
                src,
                self.buf.mem.ptr().offset(self.buf.pos as isize),
                len as usize);

            self.buf.pos += len;
            len as usize
        }
    }

    pub fn bytes<'a>(&'a self) -> &'a [u8] {
        &self.buf.mem.bytes()[..self.buf.pos()]
    }
}

impl MutBuf for MutByteBuf {
    fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    fn advance(&mut self, cnt: usize) {
        self.buf.advance(cnt)
    }

    unsafe fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8] {
        let pos = self.buf.pos();
        let lim = self.buf.lim();
        &mut self.buf.mem.bytes_mut()[pos..lim]
    }
}

impl fmt::Debug for MutByteBuf {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.bytes().fmt(fmt)
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::env;

use crate::error::*;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    /// Note: enabling `try_parsing` can reduce performance it will try and parse
    /// each environment variable 3 times (bool, i64, f64)
    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: None,
            ignore_empty: false,
            try_parsing: false,
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        let separator = match self.separator {
            Some(ref separator) => separator,
            _ => "",
        };

        let group_separator = match self.separator {
            Some(ref separator) => separator,
            _ => "_",
        };

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix.clone(), group_separator));

        for (key, value) in env::vars() {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                continue;
            }

            let mut key = key.to_string();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key
                    .to_lowercase()
                    .starts_with(&prefix_pattern.to_lowercase())
                {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                // convert to lowercase because bool parsing expects all lowercase
                if let Ok(parsed) = value.to_lowercase().parse::<bool>() {
                    ValueKind::Boolean(parsed)
                } else if let Ok(parsed) = value.parse::<i64>() {
                    ValueKind::Integer(parsed)
                } else if let Ok(parsed) = value.parse::<f64>() {
                    ValueKind::Float(parsed)
                } else {
                    ValueKind::String(value)
                }
            } else {
                ValueKind::String(value)
            };

            m.insert(key.to_lowercase(), Value::new(Some(&uri), value));
        }

        Ok(m)
    }
}
<commit_msg>refactor(env): optimize and reduce allocations<commit_after>use std::collections::HashMap;
use std::env;

use crate::error::*;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    /// Note: enabling `try_parsing` can reduce performance it will try and parse
    /// each environment variable 3 times (bool, i64, f64)
    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: None,
            ignore_empty: false,
            try_parsing: false,
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        let separator = self.separator.as_deref().unwrap_or("");
        let group_separator = self.separator.as_deref().unwrap_or("_");

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix, group_separator).to_lowercase());

        for (key, value) in env::vars() {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                continue;
            }

            let mut key = key.to_lowercase();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key.starts_with(prefix_pattern) {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                // convert to lowercase because bool parsing expects all lowercase
                if let Ok(parsed) = value.to_lowercase().parse::<bool>() {
                    ValueKind::Boolean(parsed)
                } else if let Ok(parsed) = value.parse::<i64>() {
                    ValueKind::Integer(parsed)
                } else if let Ok(parsed) = value.parse::<f64>() {
                    ValueKind::Float(parsed)
                } else {
                    ValueKind::String(value)
                }
            } else {
                ValueKind::String(value)
            };

            m.insert(key, Value::new(Some(&uri), value));
        }

        Ok(m)
    }
}
<|endoftext|>


<commit_before>use ::{ NifEnv, NifTerm };
use wrapper::copy_term;
use wrapper::nif_interface::{ self, NIF_ENV, NIF_TERM };
pub use wrapper::nif_interface::ErlNifPid;
use std::mem;
use std::sync::{Arc, Weak};

impl<'a> NifEnv<'a> {
    /// Return the calling process's pid.
    ///
    /// # Panics
    ///
    /// Panics if this environment is process-independent.  (The only way to get such an
    /// environment is to use `OwnedEnv`.  The `NifEnv` that Rustler passes to NIFs when they're
    /// called is always associated with the calling Erlang process.)
    pub fn pid(self) -> ErlNifPid {
        let mut pid: ErlNifPid = unsafe { mem::uninitialized() };
        if unsafe { nif_interface::enif_self(self.as_c_arg(), &mut pid) }.is_null() {
            panic!("environment is process-independent");
        }
        pid
    }
}

// Helper function used by save() and load().
fn nif_term_into_env<'a>(term: NIF_TERM, src_env: NIF_ENV, dest_env: NIF_ENV) -> NIF_TERM {
    if src_env == dest_env {
        term
    } else {
        copy_term(dest_env, term)
    }
}

/// A process-independent environment.
///
/// An owned environment is a place where Erlang terms can be created outside of a NIF call. Rust
/// code can use an owned environment to build a message and send it to an Erlang process.
///
///     use rustler::env::{OwnedEnv, ErlNifPid};
///     use rustler::NifEncoder;
///
///     fn send_string_to_pid(data: &str, pid: ErlNifPid) {
///         let mut msg_env = OwnedEnv::new();
///         msg_env.send(pid, |env| data.encode(env));
///     }
///
/// There's no way to run Erlang code in an owned environment. It's not a process. It's just a
/// workspace for building terms.
pub struct OwnedEnv {
    env: Arc<NIF_ENV>
}

unsafe impl Send for OwnedEnv {}

impl OwnedEnv {
    /// Allocates a new process-independent environment.
    pub fn new() -> OwnedEnv {
        OwnedEnv {
            env: Arc::new(unsafe { nif_interface::enif_alloc_env() })
        }
    }

    /// Run some code in this environment.
    pub fn run<F, R>(&self, closure: F) -> R
        where F: for<'a> FnOnce(NifEnv<'a>) -> R
    {
        let env_lifetime = ();
        let env = unsafe { NifEnv::new(&env_lifetime, *self.env) };
        closure(env)
    }

    /// Run a closure that produces a term, then send the term to another process.
    ///
    /// After the closure runs and the message is sent, the environment is cleared as though by
    /// calling the `.clear()` method.
    pub fn send<F>(&mut self, recipient: ErlNifPid, closure: F)
        where F: for<'a> FnOnce(NifEnv<'a>) -> NifTerm<'a>
    {
        let env_lifetime = ();
        let c_env = *self.env;
        let env = unsafe { NifEnv::new(&env_lifetime, c_env) };
        let message = closure(env);

        self.env = Arc::new(c_env);
        unsafe {
            nif_interface::enif_send(*self.env, &recipient, *self.env, message.as_c_arg());
        }
    }

    /// Free all terms in this environment and clear it for reuse.
    ///
    /// This invalidates `SavedTerm`s that were saved in this environment;
    /// if you later try to `.load()` one, you'll get a panic.
    ///
    /// Unless you call this method after a call to `run()`, all terms created within the
    /// environment hang around in memory until the `OwnedEnv` is dropped: garbage collection does
    /// not continually happen as needed in a NIF environment.
    pub fn clear(&mut self) {
        let c_env = *self.env;
        self.env = Arc::new(c_env);
        unsafe { nif_interface::enif_clear_env(c_env); }
    }

    /// Save a term for use in a later call to `.run()` or `.send()`.
    ///
    /// For your safety, Rust doesn't let you save `NifTerm` values from one `.run()` call to a
    /// later `.run()` call. If you try, it'll complain about lifetimes.
    ///
    /// `.save()` offers a way to do this. For example, maybe you'd like to copy a term from the
    /// caller into an `OwnedEnv`, then use that term on another thread.
    ///
    ///     # use rustler::{ NifEnv, NifTerm };
    ///     use rustler::env::OwnedEnv;
    ///     use std::thread;
    ///
    ///     fn thread_example<'a>(env: NifEnv<'a>, term: NifTerm<'a>) {
    ///         // Copy `term` into a new OwnedEnv, for use on another thread.
    ///         let mut thread_env = OwnedEnv::new();
    ///         let saved_term = thread_env.save(term);
    ///
    ///         thread::spawn(move || {
    ///             // Now run some code on the thread, using the saved term.
    ///             thread_env.run(|env| {
    ///                 let term = saved_term.load(env);
    ///                 //... do stuff with term ...
    ///             });
    ///         });
    ///     }
    ///
    /// **Note: There is no way to save terms across `OwnedEnv::send()` or `clear()`.**
    /// If you try, the `.load()` call will panic.
    pub fn save<'a>(&self, term: NifTerm<'a>) -> SavedTerm {
        SavedTerm {
            term: nif_term_into_env(term.as_c_arg(), term.get_env().as_c_arg(), *self.env),
            env_generation: Arc::downgrade(&self.env),
        }
    }
}

impl Drop for OwnedEnv {
    fn drop(&mut self) {
        unsafe { nif_interface::enif_free_env(*self.env); }
    }
}

/// A term that was created in an `OwnedEnv` and saved for later use.
///
/// These are created by calling `OwnedEnv::save()`. See that method's documentation for an
/// example.
#[derive(Clone)]
pub struct SavedTerm {
    env_generation: Weak<NIF_ENV>,
    term: NIF_TERM,
}

unsafe impl Send for SavedTerm {}

impl SavedTerm {
    /// Load this saved term into any `NifEnv`.
    ///
    /// If `env` isn't in the same `OwnedEnv` where this term was saved, the term will be copied.
    ///
    /// # Panics
    /// Panics if the `OwnedEnv` where this term was saved has been cleared or dropped.
    pub fn load<'a>(&self, env: NifEnv<'a>) -> NifTerm<'a> {
        // Check that the saved term is still valid.
        match self.env_generation.upgrade() {
            None =>
                panic!("term is from a cleared or dropped OwnedEnv"),
            Some(env_arc) =>
                NifTerm::new(env, nif_term_into_env(self.term, *env_arc, env.as_c_arg()))
        }
    }
}
<commit_msg>Fix busted merge.<commit_after>use ::{ NifEnv, NifTerm };
use wrapper::copy_term;
use wrapper::nif_interface::{ self, NIF_ENV, NIF_TERM };
pub use wrapper::nif_interface::ErlNifPid;
use std::mem;
use std::sync::{Arc, Weak};

impl<'a> NifEnv<'a> {
    /// Return the calling process's pid.
    ///
    /// # Panics
    ///
    /// Panics if this environment is process-independent.  (The only way to get such an
    /// environment is to use `OwnedEnv`.  The `NifEnv` that Rustler passes to NIFs when they're
    /// called is always associated with the calling Erlang process.)
    pub fn pid(self) -> ErlNifPid {
        let mut pid: ErlNifPid = unsafe { mem::uninitialized() };
        if unsafe { nif_interface::enif_self(self.as_c_arg(), &mut pid) }.is_null() {
            panic!("environment is process-independent");
        }
        pid
    }
}

// Helper function used by save() and load().
unsafe fn nif_term_into_env<'a>(term: NIF_TERM, src_env: NIF_ENV, dest_env: NIF_ENV) -> NIF_TERM {
    if src_env == dest_env {
        term
    } else {
        copy_term(dest_env, term)
    }
}

/// A process-independent environment.
///
/// An owned environment is a place where Erlang terms can be created outside of a NIF call. Rust
/// code can use an owned environment to build a message and send it to an Erlang process.
///
///     use rustler::env::{OwnedEnv, ErlNifPid};
///     use rustler::NifEncoder;
///
///     fn send_string_to_pid(data: &str, pid: ErlNifPid) {
///         let mut msg_env = OwnedEnv::new();
///         msg_env.send(pid, |env| data.encode(env));
///     }
///
/// There's no way to run Erlang code in an owned environment. It's not a process. It's just a
/// workspace for building terms.
pub struct OwnedEnv {
    env: Arc<NIF_ENV>
}

unsafe impl Send for OwnedEnv {}

impl OwnedEnv {
    /// Allocates a new process-independent environment.
    pub fn new() -> OwnedEnv {
        OwnedEnv {
            env: Arc::new(unsafe { nif_interface::enif_alloc_env() })
        }
    }

    /// Run some code in this environment.
    pub fn run<F, R>(&self, closure: F) -> R
        where F: for<'a> FnOnce(NifEnv<'a>) -> R
    {
        let env_lifetime = ();
        let env = unsafe { NifEnv::new(&env_lifetime, *self.env) };
        closure(env)
    }

    /// Run a closure that produces a term, then send the term to another process.
    ///
    /// After the closure runs and the message is sent, the environment is cleared as though by
    /// calling the `.clear()` method.
    pub fn send<F>(&mut self, recipient: ErlNifPid, closure: F)
        where F: for<'a> FnOnce(NifEnv<'a>) -> NifTerm<'a>
    {
        let env_lifetime = ();
        let c_env = *self.env;
        let env = unsafe { NifEnv::new(&env_lifetime, c_env) };
        let message = closure(env);

        self.env = Arc::new(c_env);
        unsafe {
            nif_interface::enif_send(*self.env, &recipient, *self.env, message.as_c_arg());
        }
    }

    /// Free all terms in this environment and clear it for reuse.
    ///
    /// This invalidates `SavedTerm`s that were saved in this environment;
    /// if you later try to `.load()` one, you'll get a panic.
    ///
    /// Unless you call this method after a call to `run()`, all terms created within the
    /// environment hang around in memory until the `OwnedEnv` is dropped: garbage collection does
    /// not continually happen as needed in a NIF environment.
    pub fn clear(&mut self) {
        let c_env = *self.env;
        self.env = Arc::new(c_env);
        unsafe { nif_interface::enif_clear_env(c_env); }
    }

    /// Save a term for use in a later call to `.run()` or `.send()`.
    ///
    /// For your safety, Rust doesn't let you save `NifTerm` values from one `.run()` call to a
    /// later `.run()` call. If you try, it'll complain about lifetimes.
    ///
    /// `.save()` offers a way to do this. For example, maybe you'd like to copy a term from the
    /// caller into an `OwnedEnv`, then use that term on another thread.
    ///
    ///     # use rustler::{ NifEnv, NifTerm };
    ///     use rustler::env::OwnedEnv;
    ///     use std::thread;
    ///
    ///     fn thread_example<'a>(env: NifEnv<'a>, term: NifTerm<'a>) {
    ///         // Copy `term` into a new OwnedEnv, for use on another thread.
    ///         let mut thread_env = OwnedEnv::new();
    ///         let saved_term = thread_env.save(term);
    ///
    ///         thread::spawn(move || {
    ///             // Now run some code on the thread, using the saved term.
    ///             thread_env.run(|env| {
    ///                 let term = saved_term.load(env);
    ///                 //... do stuff with term ...
    ///             });
    ///         });
    ///     }
    ///
    /// **Note: There is no way to save terms across `OwnedEnv::send()` or `clear()`.**
    /// If you try, the `.load()` call will panic.
    pub fn save<'a>(&self, term: NifTerm<'a>) -> SavedTerm {
        SavedTerm {
            term: unsafe { nif_term_into_env(term.as_c_arg(), term.get_env().as_c_arg(), *self.env) },
            env_generation: Arc::downgrade(&self.env),
        }
    }
}

impl Drop for OwnedEnv {
    fn drop(&mut self) {
        unsafe { nif_interface::enif_free_env(*self.env); }
    }
}

/// A term that was created in an `OwnedEnv` and saved for later use.
///
/// These are created by calling `OwnedEnv::save()`. See that method's documentation for an
/// example.
#[derive(Clone)]
pub struct SavedTerm {
    env_generation: Weak<NIF_ENV>,
    term: NIF_TERM,
}

unsafe impl Send for SavedTerm {}

impl SavedTerm {
    /// Load this saved term into any `NifEnv`.
    ///
    /// If `env` isn't in the same `OwnedEnv` where this term was saved, the term will be copied.
    ///
    /// # Panics
    /// Panics if the `OwnedEnv` where this term was saved has been cleared or dropped.
    pub fn load<'a>(&self, env: NifEnv<'a>) -> NifTerm<'a> {
        // Check that the saved term is still valid.
        match self.env_generation.upgrade() {
            None =>
                panic!("term is from a cleared or dropped OwnedEnv"),
            Some(env_arc) =>
                NifTerm::new(env, unsafe { nif_term_into_env(self.term, *env_arc, env.as_c_arg()) })
        }
    }
}
<|endoftext|>