<commit_before>use regex::Regex;
use std::fmt::Write as _;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::path::MAIN_SEPARATOR;
use std::str;

use protobuf_test_common::build::copy_tests_v2_v3;
use protobuf_test_common::build::glob_simple;
use std::process::Command;
use std::process::Stdio;

fn to_paths(iter: impl IntoIterator<Item = impl Into<String>>) -> Vec<PathBuf> {
    iter.into_iter()
        .map(|item| item.into().replace(MAIN_SEPARATOR, "/").into())
        .collect()
}

#[derive(Default, Debug)]
struct TestStats {
    passed: u32,
    passed_marked_skip: u32,
    skipped: u32,
    failed: u32,
}

fn normalize_generated_file(contents: &str) -> String {
    let parsed_by = Regex::new("^// \\.proto file is parsed by.*").unwrap();

    let mut r = String::new();
    let mut inside_descriptor_data = false;
    for line in contents.lines() {
        let line = if inside_descriptor_data {
            if line == "\";" {
                inside_descriptor_data = false;
                line.to_owned()
            } else {
                continue;
            }
        } else if line.starts_with("static file_descriptor_proto_data") {
            inside_descriptor_data = true;
            line.to_owned()
        } else {
            parsed_by.replace(line, "").into_owned()
        };

        writeln!(&mut r, "{}", line).unwrap();
    }

    // sanity check
    assert!(!inside_descriptor_data);

    r
}

fn normalize_generated_file_in_place(path: &Path) {
    let contents = fs::read_to_string(path).unwrap();
    let contents = normalize_generated_file(&contents);
    fs::write(path, &contents).unwrap();
}

fn print_diff(dir: &Path, a: &Path, b: &Path) {
    if cfg!(windows) {
        // likely we don't have `diff` command on Windows
        return;
    }

    // Not using difference crate because it's slow for large files
    let output = Command::new("diff")
        .current_dir(dir)
        .arg("-u")
        .arg(a)
        .arg(b)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .output()
        .unwrap();

    print!("{}", str::from_utf8(&output.stdout).unwrap());
    print!("{}", str::from_utf8(&output.stderr).unwrap());
}

fn test_diff_in<F>(root: &str, s: &str, include: &str, should_skip: F, stats: &mut TestStats)
where
    F: Fn(&str) -> bool,
{
    let include_full = format!("{}/{}", root, include);
    let s_full = format!("{}/{}", root, s);

    let inputs_glob = format!("{}/*.proto", s_full);
    let inputs = to_paths(glob_simple(&inputs_glob));
    assert!(!inputs.is_empty(), "glob is empty: {}", inputs_glob);
    let includes = to_paths(vec![include_full.as_str(), "../proto"]);

    let temp_dir = tempfile::Builder::new()
        .prefix("protobuf-codegen-identical-test")
        .tempdir()
        .unwrap();

    let protoc_dir = format!("{}/protoc", temp_dir.path().to_str().unwrap());
    let pure_dir = format!("{}/pure", temp_dir.path().to_str().unwrap());

    fs::create_dir(&protoc_dir).unwrap();
    fs::create_dir(&pure_dir).unwrap();

    protoc_rust::Codegen::new()
        .inputs(&inputs)
        .includes(&includes)
        .out_dir(&protoc_dir)
        .run()
        .unwrap();

    protobuf_codegen_pure::Codegen::new()
        .inputs(&inputs)
        .includes(&includes)
        .out_dir(&pure_dir)
        .run()
        .unwrap();

    for input in &inputs {
        let label = input.display().to_string();
        assert!(label.starts_with(&format!("{}/", root)));
        let label = &label[format!("{}/", root).len()..];
        let proto_name = input.file_name().unwrap().to_str().unwrap();
        let rs_name = protobuf_codegen::proto_name_to_rs(proto_name);
        let protoc_rs = format!("{}/{}", protoc_dir, rs_name);
        let pure_rs = format!("{}/{}", pure_dir, rs_name);

        normalize_generated_file_in_place(Path::new(&protoc_rs));
        normalize_generated_file_in_place(Path::new(&pure_rs));

        let protoc_rs_contents =
            fs::read_to_string(&protoc_rs).expect(&format!("while reading {}", protoc_rs));
        let pure_rs_contents =
            fs::read_to_string(&pure_rs).expect(&format!("while reading {}", pure_rs));
        let skip = should_skip(input.to_str().unwrap());
        if protoc_rs_contents == pure_rs_contents {
            if !skip {
                stats.passed += 1;
                println!("{}: PASSED", label);
            } else {
                stats.passed_marked_skip += 1;
                println!("{}: PASSED BUT MARKED SKIP", label);
            }
        } else {
            if skip {
                stats.skipped += 1;
                println!("{} SKIPPED", label);
            } else {
                stats.failed += 1;
                println!("{} FAILED", label);
            }

            print_diff(temp_dir.path(), Path::new(&protoc_rs), Path::new(&pure_rs));
        }
    }
}

#[test]
fn test_diff() {
    let should_skip = |path: &str| {
        fs::read_to_string(path)
            .unwrap()
            .contains("@skip-codegen-identical-test")
    };

    let common_v3_root = tempfile::Builder::new()
        .prefix("common-v3")
        .tempdir()
        .unwrap();
    fs::create_dir_all(format!(
        "{}/src/common/v3",
        common_v3_root.path().to_str().unwrap()
    ))
    .unwrap();

    copy_tests_v2_v3(
        "../protobuf-test/src/common/v2",
        &format!("{}/src/common/v3", common_v3_root.path().to_str().unwrap()),
    );

    let mut stats = TestStats {
        passed: 0,
        passed_marked_skip: 0,
        skipped: 0,
        failed: 0,
    };
    test_diff_in(
        "../protobuf-test",
        "src/v2",
        "src/v2",
        should_skip,
        &mut stats,
    );
    test_diff_in(
        "../protobuf-test",
        "src/v3",
        "src/v3",
        should_skip,
        &mut stats,
    );
    test_diff_in(
        "../protobuf-test",
        "src/common/v2",
        "src/common/v2",
        should_skip,
        &mut stats,
    );
    test_diff_in(
        common_v3_root.path().to_str().unwrap(),
        "src/common/v3",
        "src/common/v3",
        should_skip,
        &mut stats,
    );
    test_diff_in(
        "../protobuf-test",
        "../interop/cxx",
        "../interop/cxx",
        should_skip,
        &mut stats,
    );
    test_diff_in(
        "../protobuf-test",
        "src/google/protobuf",
        "src",
        |_| true,
        &mut stats,
    );

    println!("{:?}", stats);
    assert!(stats.passed != 0, "sanity check");
    assert!(stats.failed == 0, "at least one test failed");
}
<commit_msg>Split diff tests into several smaller tests<commit_after>use regex::Regex;
use std::fmt::Write as _;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::path::MAIN_SEPARATOR;
use std::str;

use protobuf_test_common::build::copy_tests_v2_v3;
use protobuf_test_common::build::glob_simple;
use std::process::Command;
use std::process::Stdio;

fn to_paths(iter: impl IntoIterator<Item = impl Into<String>>) -> Vec<PathBuf> {
    iter.into_iter()
        .map(|item| item.into().replace(MAIN_SEPARATOR, "/").into())
        .collect()
}

#[derive(Default, Debug)]
struct TestStats {
    passed: u32,
    passed_marked_skip: u32,
    skipped: u32,
    failed: u32,
}

fn normalize_generated_file(contents: &str) -> String {
    let parsed_by = Regex::new("^// \\.proto file is parsed by.*").unwrap();

    let mut r = String::new();
    let mut inside_descriptor_data = false;
    for line in contents.lines() {
        let line = if inside_descriptor_data {
            if line == "\";" {
                inside_descriptor_data = false;
                line.to_owned()
            } else {
                continue;
            }
        } else if line.starts_with("static file_descriptor_proto_data") {
            inside_descriptor_data = true;
            line.to_owned()
        } else {
            parsed_by.replace(line, "").into_owned()
        };

        writeln!(&mut r, "{}", line).unwrap();
    }

    // sanity check
    assert!(!inside_descriptor_data);

    r
}

fn normalize_generated_file_in_place(path: &Path) {
    let contents = fs::read_to_string(path).unwrap();
    let contents = normalize_generated_file(&contents);
    fs::write(path, &contents).unwrap();
}

fn print_diff(dir: &Path, a: &Path, b: &Path) {
    if cfg!(windows) {
        // likely we don't have `diff` command on Windows
        return;
    }

    // Not using difference crate because it's slow for large files
    let output = Command::new("diff")
        .current_dir(dir)
        .arg("-u")
        .arg(a)
        .arg(b)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .output()
        .unwrap();

    print!("{}", str::from_utf8(&output.stdout).unwrap());
    print!("{}", str::from_utf8(&output.stderr).unwrap());
}

fn test_diff_in<F>(root: &str, s: &str, include: &str, should_skip: F)
where
    F: Fn(&str) -> bool,
{
    let mut stats = TestStats {
        passed: 0,
        passed_marked_skip: 0,
        skipped: 0,
        failed: 0,
    };

    let include_full = format!("{}/{}", root, include);
    let s_full = format!("{}/{}", root, s);

    let inputs_glob = format!("{}/*.proto", s_full);
    let inputs = to_paths(glob_simple(&inputs_glob));
    assert!(!inputs.is_empty(), "glob is empty: {}", inputs_glob);
    let includes = to_paths(vec![include_full.as_str(), "../proto"]);

    let temp_dir = tempfile::Builder::new()
        .prefix("protobuf-codegen-identical-test")
        .tempdir()
        .unwrap();

    let protoc_dir = format!("{}/protoc", temp_dir.path().to_str().unwrap());
    let pure_dir = format!("{}/pure", temp_dir.path().to_str().unwrap());

    fs::create_dir(&protoc_dir).unwrap();
    fs::create_dir(&pure_dir).unwrap();

    protoc_rust::Codegen::new()
        .inputs(&inputs)
        .includes(&includes)
        .out_dir(&protoc_dir)
        .run()
        .unwrap();

    protobuf_codegen_pure::Codegen::new()
        .inputs(&inputs)
        .includes(&includes)
        .out_dir(&pure_dir)
        .run()
        .unwrap();

    for input in &inputs {
        let label = input.display().to_string();
        assert!(label.starts_with(&format!("{}/", root)));
        let label = &label[format!("{}/", root).len()..];
        let proto_name = input.file_name().unwrap().to_str().unwrap();
        let rs_name = protobuf_codegen::proto_name_to_rs(proto_name);
        let protoc_rs = format!("{}/{}", protoc_dir, rs_name);
        let pure_rs = format!("{}/{}", pure_dir, rs_name);

        normalize_generated_file_in_place(Path::new(&protoc_rs));
        normalize_generated_file_in_place(Path::new(&pure_rs));

        let protoc_rs_contents =
            fs::read_to_string(&protoc_rs).expect(&format!("while reading {}", protoc_rs));
        let pure_rs_contents =
            fs::read_to_string(&pure_rs).expect(&format!("while reading {}", pure_rs));
        let skip = should_skip(input.to_str().unwrap());
        if protoc_rs_contents == pure_rs_contents {
            if !skip {
                stats.passed += 1;
                println!("{}: PASSED", label);
            } else {
                stats.passed_marked_skip += 1;
                println!("{}: PASSED BUT MARKED SKIP", label);
            }
        } else {
            if skip {
                stats.skipped += 1;
                println!("{} SKIPPED", label);
            } else {
                stats.failed += 1;
                println!("{} FAILED", label);
            }

            print_diff(temp_dir.path(), Path::new(&protoc_rs), Path::new(&pure_rs));
        }
    }

    println!("{:?}", stats);
    assert!(
        stats.passed != 0 || s == "src/google/protobuf",
        "sanity check"
    );
    assert!(stats.failed == 0, "at least one test failed");
}

fn should_skip_with_marker(path: &str) -> bool {
    fs::read_to_string(path)
        .unwrap()
        .contains("@skip-codegen-identical-test")
}

#[test]
fn common_v2() {
    test_diff_in(
        "../protobuf-test",
        "src/common/v2",
        "src/common/v2",
        should_skip_with_marker,
    );
}

#[test]
fn common_v3() {
    let common_v3_root = tempfile::Builder::new()
        .prefix("common-v3")
        .tempdir()
        .unwrap();
    fs::create_dir_all(format!(
        "{}/src/common/v3",
        common_v3_root.path().to_str().unwrap()
    ))
    .unwrap();

    copy_tests_v2_v3(
        "../protobuf-test/src/common/v2",
        &format!("{}/src/common/v3", common_v3_root.path().to_str().unwrap()),
    );

    test_diff_in(
        common_v3_root.path().to_str().unwrap(),
        "src/common/v3",
        "src/common/v3",
        should_skip_with_marker,
    );
}

#[test]
fn v2() {
    test_diff_in(
        "../protobuf-test",
        "src/v2",
        "src/v2",
        should_skip_with_marker,
    );
}

#[test]
fn v3() {
    test_diff_in(
        "../protobuf-test",
        "src/v3",
        "src/v3",
        should_skip_with_marker,
    );
}

#[test]
fn interop() {
    test_diff_in(
        "../protobuf-test",
        "../interop/cxx",
        "../interop/cxx",
        should_skip_with_marker,
    );
}

#[test]
fn google() {
    test_diff_in("../protobuf-test", "src/google/protobuf", "src", |_| true);
}
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;

pub use stats::{Prob, LogProb};

/// A factor to convert log-probabilities to PHRED-scale (phred = p * LOG_TO_PHRED_FACTOR).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * PHRED_TO_LOG_FACTOR).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
pub fn ln_1m_exp(p: LogProb) -> LogProb {
    if p < -0.693 {
        (-p.exp()).ln_1p()
    }
    else {
        (-p.exp_m1()).ln()
    }
}

/// Convert log scale probability to PHRED scale.
pub fn log_to_phred(p: LogProb) -> f64 {
    p * LOG_TO_PHRED_FACTOR
}

/// Convert PHRED scale probability to log scale.
pub fn phred_to_log(p: f64) -> LogProb {
    p * PHRED_TO_LOG_FACTOR
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_sum(probs: &[LogProb]) -> LogProb {
    if probs.is_empty() {
        f64::NEG_INFINITY
    }
    else {
        let mut pmax = probs[0];
        let mut imax = 0;
        for (i, &p) in probs.iter().enumerate().skip(1) {
            if p > pmax {
                pmax = p;
                imax = i;
            }
        }
        if pmax == f64::NEG_INFINITY {
            f64::NEG_INFINITY
        }
        else {
            // TODO use sum() once it has been stabilized: .sum::<usize>()
            pmax + (probs.iter().enumerate().filter_map(|(i, p)| if i != imax { Some((p - pmax).exp()) } else { None }).fold(0.0, |s, e| s + e)).ln_1p()
        }
    }
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_add(mut p0: LogProb, mut p1: LogProb) -> LogProb {
    if p1 > p0 {
        mem::swap(&mut p0, &mut p1);
    }
    if p0 == f64::NEG_INFINITY {
        f64::NEG_INFINITY
    }
    else {
        p0 + (p1 - p0).exp().ln_1p()
    }
}

/// Calcualte the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_cumsum(probs: &[LogProb]) -> Vec<LogProb> {
    probs.iter().scan(f64::NEG_INFINITY, |s, p| {
        *s = log_prob_add(*s, *p);
        Some(*s)
    }).collect()
}

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

    #[test]
    fn test_log_prob_sum() {
        let probs = [f64::NEG_INFINITY, 0.0, f64::NEG_INFINITY];
        assert_eq!(log_prob_sum(&probs), 0.0);
    }

    #[test]
    fn test_empty_log_prob_sum() {
        assert_eq!(log_prob_sum(&[]), f64::NEG_INFINITY);
    }
}
<commit_msg>Handle a corner case in logprob sums.<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;

pub use stats::{Prob, LogProb};

/// A factor to convert log-probabilities to PHRED-scale (phred = p * LOG_TO_PHRED_FACTOR).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * PHRED_TO_LOG_FACTOR).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
pub fn ln_1m_exp(p: LogProb) -> LogProb {
    if p < -0.693 {
        (-p.exp()).ln_1p()
    }
    else {
        (-p.exp_m1()).ln()
    }
}

/// Convert log scale probability to PHRED scale.
pub fn log_to_phred(p: LogProb) -> f64 {
    p * LOG_TO_PHRED_FACTOR
}

/// Convert PHRED scale probability to log scale.
pub fn phred_to_log(p: f64) -> LogProb {
    p * PHRED_TO_LOG_FACTOR
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_sum(probs: &[LogProb]) -> LogProb {
    if probs.is_empty() {
        f64::NEG_INFINITY
    }
    else {
        let mut pmax = probs[0];
        let mut imax = 0;
        for (i, &p) in probs.iter().enumerate().skip(1) {
            if p > pmax {
                pmax = p;
                imax = i;
            }
        }
        if pmax == f64::NEG_INFINITY {
            f64::NEG_INFINITY
        }
        else if pmax == f64::INFINITY {
            f64::INFINITY
        }
        else {
            // TODO use sum() once it has been stabilized: .sum::<usize>()
            pmax + (probs.iter().enumerate().filter_map(|(i, p)| if i != imax { Some((p - pmax).exp()) } else { None }).fold(0.0, |s, e| s + e)).ln_1p()
        }
    }
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_add(mut p0: LogProb, mut p1: LogProb) -> LogProb {
    if p1 > p0 {
        mem::swap(&mut p0, &mut p1);
    }
    if p0 == f64::NEG_INFINITY {
        f64::NEG_INFINITY
    }
    else if p0 == f64::INFINITY {
        f64::INFINITY
    }
    else {
        p0 + (p1 - p0).exp().ln_1p()
    }
}

/// Calcualte the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_cumsum(probs: &[LogProb]) -> Vec<LogProb> {
    probs.iter().scan(f64::NEG_INFINITY, |s, p| {
        *s = log_prob_add(*s, *p);
        Some(*s)
    }).collect()
}

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

    #[test]
    fn test_log_prob_sum() {
        let probs = [f64::NEG_INFINITY, 0.0, f64::NEG_INFINITY];
        assert_eq!(log_prob_sum(&probs), 0.0);
    }

    #[test]
    fn test_empty_log_prob_sum() {
        assert_eq!(log_prob_sum(&[]), f64::NEG_INFINITY);
    }
}
<|endoftext|>


<commit_before>//! Initialization logic for a generic instance of the `Allocate` channel allocation trait.

use std::thread;
#[cfg(feature = "getopts")]
use std::io::BufRead;
#[cfg(feature = "getopts")]
use getopts;
use std::sync::Arc;

use std::any::Any;

use crate::allocator::thread::ThreadBuilder;
use crate::allocator::{AllocateBuilder, Process, Generic, GenericBuilder};
use crate::allocator::zero_copy::initialize::initialize_networking;

use crate::logging::{CommunicationSetup, CommunicationEvent};
use logging_core::Logger;

/// Possible configurations for the communication infrastructure.
pub enum Configuration {
    /// Use one thread.
    Thread,
    /// Use one process with an indicated number of threads.
    Process(usize),
    /// Expect multiple processes.
    Cluster {
        /// Number of per-process worker threads
        threads: usize,
        /// Identity of this process
        process: usize,
        /// Addresses of all processes
        addresses: Vec<String>,
        /// Verbosely report connection process
        report: bool,
        /// Closure to create a new logger for a communication thread
        log_fn: Box<Fn(CommunicationSetup) -> Option<Logger<CommunicationEvent, CommunicationSetup>> + Send + Sync>,
    }
}

#[cfg(feature = "getopts")]
impl Configuration {

    /// Returns a `getopts::Options` struct that can be used to print
    /// usage information in higher-level systems.
    pub fn options() -> getopts::Options {
        let mut opts = getopts::Options::new();
        opts.optopt("w", "threads", "number of per-process worker threads", "NUM");
        opts.optopt("p", "process", "identity of this process", "IDX");
        opts.optopt("n", "processes", "number of processes", "NUM");
        opts.optopt("h", "hostfile", "text file whose lines are process addresses", "FILE");
        opts.optflag("r", "report", "reports connection progress");

        opts
    }

    /// Constructs a new configuration by parsing supplied text arguments.
    ///
    /// Most commonly, this uses `std::env::Args()` as the supplied iterator.
    pub fn from_args<I: Iterator<Item=String>>(args: I) -> Result<Configuration,String> {
        let opts = Configuration::options();

        opts.parse(args)
            .map_err(|e| format!("{:?}", e))
            .map(|matches| {

            // let mut config = Configuration::new(1, 0, Vec::new());
            let threads = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let process = matches.opt_str("p").map(|x| x.parse().unwrap_or(0)).unwrap_or(0);
            let processes = matches.opt_str("n").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let report = matches.opt_present("report");

            assert!(process < processes);

            if processes > 1 {
                let mut addresses = Vec::new();
                if let Some(hosts) = matches.opt_str("h") {
                    let reader = ::std::io::BufReader::new(::std::fs::File::open(hosts.clone()).unwrap());
                    for x in reader.lines().take(processes) {
                        addresses.push(x.unwrap());
                    }
                    if addresses.len() < processes {
                        panic!("could only read {} addresses from {}, but -n: {}", addresses.len(), hosts, processes);
                    }
                }
                else {
                    for index in 0..processes {
                        addresses.push(format!("localhost:{}", 2101 + index));
                    }
                }

                assert!(processes == addresses.len());
                Configuration::Cluster {
                    threads,
                    process,
                    addresses,
                    report,
                    log_fn: Box::new( | _ | None),
                }
            }
            else if threads > 1 { Configuration::Process(threads) }
            else { Configuration::Thread }
        })
    }

    /// Attempts to assemble the described communication infrastructure.
    pub fn try_build(self) -> Result<(Vec<GenericBuilder>, Box<Any>), String> {
        match self {
            Configuration::Thread => {
                Ok((vec![GenericBuilder::Thread(ThreadBuilder)], Box::new(())))
            },
            Configuration::Process(threads) => {
                Ok((Process::new_vector(threads).into_iter().map(|x| GenericBuilder::Process(x)).collect(), Box::new(())))
            },
            Configuration::Cluster { threads, process, addresses, report, log_fn } => {
                match initialize_networking(addresses, process, threads, report, log_fn) {
                    Ok((stuff, guard)) => {
                        Ok((stuff.into_iter().map(|x| GenericBuilder::ZeroCopy(x)).collect(), Box::new(guard)))
                    },
                    Err(err) => Err(format!("failed to initialize networking: {}", err))
                }
            },
        }
    }
}

/// Initializes communication and executes a distributed computation.
///
/// This method allocates an `allocator::Generic` for each thread, spawns local worker threads,
/// and invokes the supplied function with the allocator.
/// The method returns a `WorkerGuards<T>` which can be `join`ed to retrieve the return values
/// (or errors) of the workers.
///
///
/// # Examples
/// ```
/// use timely_communication::Allocate;
///
/// // configure for two threads, just one process.
/// let config = timely_communication::Configuration::Process(2);
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize(config, |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver) = allocator.allocate(0);
///
///     // send typed data along each channel
///     use timely_communication::Message;
///     senders[0].send(Message::from_typed(format!("hello, {}", 0)));
///     senders[1].send(Message::from_typed(format!("hello, {}", 1)));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         allocator.receive();
///         if let Some(message) = receiver.recv() {
///             use std::ops::Deref;
///             println!("worker {}: received: <{}>", allocator.index(), message.deref());
///             expecting -= 1;
///         }
///         allocator.release();
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
///
/// The should produce output like:
///
/// ```ignore
/// worker 0 started
/// worker 1 started
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// result: Ok(0)
/// result: Ok(1)
/// ```
pub fn initialize<T:Send+'static, F: Fn(Generic)->T+Send+Sync+'static>(
    config: Configuration,
    func: F,
) -> Result<WorkerGuards<T>,String> {
    let (allocators, others) = config.try_build()?;
    initialize_from(allocators, others, func)
}

/// Initializes computation and runs a distributed computation.
///
/// This version of `initialize` allows you to explicitly specify the allocators that
/// you want to use, by providing an explicit list of allocator builders. Additionally,
/// you provide `others`, a `Box<Any>` which will be held by the resulting worker guard
/// and dropped when it is dropped, which allows you to join communication threads.
///
/// # Examples
/// ```
/// use timely_communication::Allocate;
///
/// // configure for two threads, just one process.
/// let builders = timely_communication::allocator::process::Process::new_vector(2);
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize_from(builders, Box::new(()), |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver) = allocator.allocate(0);
///
///     // send typed data along each channel
///     use timely_communication::Message;
///     senders[0].send(Message::from_typed(format!("hello, {}", 0)));
///     senders[1].send(Message::from_typed(format!("hello, {}", 1)));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         allocator.receive();
///         if let Some(message) = receiver.recv() {
///             use std::ops::Deref;
///             println!("worker {}: received: <{}>", allocator.index(), message.deref());
///             expecting -= 1;
///         }
///         allocator.release();
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
pub fn initialize_from<A, T, F>(
    builders: Vec<A>,
    _others: Box<Any>,
    func: F,
) -> Result<WorkerGuards<T>,String>
where
    A: AllocateBuilder+'static,
    T: Send+'static,
    F: Fn(<A as AllocateBuilder>::Allocator)->T+Send+Sync+'static
{
    let logic = Arc::new(func);
    let mut guards = Vec::new();
    for (index, builder) in builders.into_iter().enumerate() {
        let clone = logic.clone();
        guards.push(thread::Builder::new()
                            .name(format!("worker thread {}", index))
                            .spawn(move || {
                                let communicator = builder.build();
                                (*clone)(communicator)
                            })
                            .map_err(|e| format!("{:?}", e))?);
    }

    Ok(WorkerGuards { guards, _others })
}

/// Maintains `JoinHandle`s for worker threads.
pub struct WorkerGuards<T:Send+'static> {
    guards: Vec<::std::thread::JoinHandle<T>>,
    _others: Box<Any>,
}

impl<T:Send+'static> WorkerGuards<T> {

    /// Returns the number of guards.
    pub fn len(&self) -> usize { self.guards.len() }

    /// Returns a reference to the indexed guard.
    pub fn guard(&self, index: usize) -> &std::thread::JoinHandle<T> {
        &self.guards[index]
    }

    /// Waits on the worker threads and returns the results they produce.
    pub fn join(mut self) -> Vec<Result<T, String>> {
        self.guards
            .drain(..)
            .map(|guard| guard.join().map_err(|e| format!("{:?}", e)))
            .collect()
    }
}

impl<T:Send+'static> Drop for WorkerGuards<T> {
    fn drop(&mut self) {
        for guard in self.guards.drain(..) {
            guard.join().expect("Worker panic");
        }
        // println!("WORKER THREADS JOINED");
    }
}
<commit_msg>tidy api<commit_after>//! Initialization logic for a generic instance of the `Allocate` channel allocation trait.

use std::thread;
#[cfg(feature = "getopts")]
use std::io::BufRead;
#[cfg(feature = "getopts")]
use getopts;
use std::sync::Arc;

use std::any::Any;

use crate::allocator::thread::ThreadBuilder;
use crate::allocator::{AllocateBuilder, Process, Generic, GenericBuilder};
use crate::allocator::zero_copy::initialize::initialize_networking;

use crate::logging::{CommunicationSetup, CommunicationEvent};
use logging_core::Logger;

/// Possible configurations for the communication infrastructure.
pub enum Configuration {
    /// Use one thread.
    Thread,
    /// Use one process with an indicated number of threads.
    Process(usize),
    /// Expect multiple processes.
    Cluster {
        /// Number of per-process worker threads
        threads: usize,
        /// Identity of this process
        process: usize,
        /// Addresses of all processes
        addresses: Vec<String>,
        /// Verbosely report connection process
        report: bool,
        /// Closure to create a new logger for a communication thread
        log_fn: Box<Fn(CommunicationSetup) -> Option<Logger<CommunicationEvent, CommunicationSetup>> + Send + Sync>,
    }
}

#[cfg(feature = "getopts")]
impl Configuration {

    /// Returns a `getopts::Options` struct that can be used to print
    /// usage information in higher-level systems.
    pub fn options() -> getopts::Options {
        let mut opts = getopts::Options::new();
        opts.optopt("w", "threads", "number of per-process worker threads", "NUM");
        opts.optopt("p", "process", "identity of this process", "IDX");
        opts.optopt("n", "processes", "number of processes", "NUM");
        opts.optopt("h", "hostfile", "text file whose lines are process addresses", "FILE");
        opts.optflag("r", "report", "reports connection progress");

        opts
    }

    /// Constructs a new configuration by parsing supplied text arguments.
    ///
    /// Most commonly, this uses `std::env::Args()` as the supplied iterator.
    pub fn from_args<I: Iterator<Item=String>>(args: I) -> Result<Configuration,String> {
        let opts = Configuration::options();

        opts.parse(args)
            .map_err(|e| format!("{:?}", e))
            .map(|matches| {

            // let mut config = Configuration::new(1, 0, Vec::new());
            let threads = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let process = matches.opt_str("p").map(|x| x.parse().unwrap_or(0)).unwrap_or(0);
            let processes = matches.opt_str("n").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let report = matches.opt_present("report");

            assert!(process < processes);

            if processes > 1 {
                let mut addresses = Vec::new();
                if let Some(hosts) = matches.opt_str("h") {
                    let reader = ::std::io::BufReader::new(::std::fs::File::open(hosts.clone()).unwrap());
                    for x in reader.lines().take(processes) {
                        addresses.push(x.unwrap());
                    }
                    if addresses.len() < processes {
                        panic!("could only read {} addresses from {}, but -n: {}", addresses.len(), hosts, processes);
                    }
                }
                else {
                    for index in 0..processes {
                        addresses.push(format!("localhost:{}", 2101 + index));
                    }
                }

                assert!(processes == addresses.len());
                Configuration::Cluster {
                    threads,
                    process,
                    addresses,
                    report,
                    log_fn: Box::new( | _ | None),
                }
            }
            else if threads > 1 { Configuration::Process(threads) }
            else { Configuration::Thread }
        })
    }

    /// Attempts to assemble the described communication infrastructure.
    pub fn try_build(self) -> Result<(Vec<GenericBuilder>, Box<Any>), String> {
        match self {
            Configuration::Thread => {
                Ok((vec![GenericBuilder::Thread(ThreadBuilder)], Box::new(())))
            },
            Configuration::Process(threads) => {
                Ok((Process::new_vector(threads).into_iter().map(|x| GenericBuilder::Process(x)).collect(), Box::new(())))
            },
            Configuration::Cluster { threads, process, addresses, report, log_fn } => {
                match initialize_networking(addresses, process, threads, report, log_fn) {
                    Ok((stuff, guard)) => {
                        Ok((stuff.into_iter().map(|x| GenericBuilder::ZeroCopy(x)).collect(), Box::new(guard)))
                    },
                    Err(err) => Err(format!("failed to initialize networking: {}", err))
                }
            },
        }
    }
}

/// Initializes communication and executes a distributed computation.
///
/// This method allocates an `allocator::Generic` for each thread, spawns local worker threads,
/// and invokes the supplied function with the allocator.
/// The method returns a `WorkerGuards<T>` which can be `join`ed to retrieve the return values
/// (or errors) of the workers.
///
///
/// # Examples
/// ```
/// use timely_communication::Allocate;
///
/// // configure for two threads, just one process.
/// let config = timely_communication::Configuration::Process(2);
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize(config, |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver) = allocator.allocate(0);
///
///     // send typed data along each channel
///     use timely_communication::Message;
///     senders[0].send(Message::from_typed(format!("hello, {}", 0)));
///     senders[1].send(Message::from_typed(format!("hello, {}", 1)));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         allocator.receive();
///         if let Some(message) = receiver.recv() {
///             use std::ops::Deref;
///             println!("worker {}: received: <{}>", allocator.index(), message.deref());
///             expecting -= 1;
///         }
///         allocator.release();
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
///
/// The should produce output like:
///
/// ```ignore
/// worker 0 started
/// worker 1 started
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// result: Ok(0)
/// result: Ok(1)
/// ```
pub fn initialize<T:Send+'static, F: Fn(Generic)->T+Send+Sync+'static>(
    config: Configuration,
    func: F,
) -> Result<WorkerGuards<T>,String> {
    let (allocators, others) = config.try_build()?;
    initialize_from(allocators, others, func)
}

/// Initializes computation and runs a distributed computation.
///
/// This version of `initialize` allows you to explicitly specify the allocators that
/// you want to use, by providing an explicit list of allocator builders. Additionally,
/// you provide `others`, a `Box<Any>` which will be held by the resulting worker guard
/// and dropped when it is dropped, which allows you to join communication threads.
///
/// # Examples
/// ```
/// use timely_communication::Allocate;
///
/// // configure for two threads, just one process.
/// let builders = timely_communication::allocator::process::Process::new_vector(2);
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize_from(builders, Box::new(()), |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver) = allocator.allocate(0);
///
///     // send typed data along each channel
///     use timely_communication::Message;
///     senders[0].send(Message::from_typed(format!("hello, {}", 0)));
///     senders[1].send(Message::from_typed(format!("hello, {}", 1)));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         allocator.receive();
///         if let Some(message) = receiver.recv() {
///             use std::ops::Deref;
///             println!("worker {}: received: <{}>", allocator.index(), message.deref());
///             expecting -= 1;
///         }
///         allocator.release();
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
pub fn initialize_from<A, T, F>(
    builders: Vec<A>,
    _others: Box<Any>,
    func: F,
) -> Result<WorkerGuards<T>,String>
where
    A: AllocateBuilder+'static,
    T: Send+'static,
    F: Fn(<A as AllocateBuilder>::Allocator)->T+Send+Sync+'static
{
    let logic = Arc::new(func);
    let mut guards = Vec::new();
    for (index, builder) in builders.into_iter().enumerate() {
        let clone = logic.clone();
        guards.push(thread::Builder::new()
                            .name(format!("worker thread {}", index))
                            .spawn(move || {
                                let communicator = builder.build();
                                (*clone)(communicator)
                            })
                            .map_err(|e| format!("{:?}", e))?);
    }

    Ok(WorkerGuards { guards, _others })
}

/// Maintains `JoinHandle`s for worker threads.
pub struct WorkerGuards<T:Send+'static> {
    guards: Vec<::std::thread::JoinHandle<T>>,
    _others: Box<Any>,
}

impl<T:Send+'static> WorkerGuards<T> {

    /// Returns a reference to the indexed guard.
    pub fn guards(&self) -> &[std::thread::JoinHandle<T>] {
        &self.guards[..]
    }

    /// Waits on the worker threads and returns the results they produce.
    pub fn join(mut self) -> Vec<Result<T, String>> {
        self.guards
            .drain(..)
            .map(|guard| guard.join().map_err(|e| format!("{:?}", e)))
            .collect()
    }
}

impl<T:Send+'static> Drop for WorkerGuards<T> {
    fn drop(&mut self) {
        for guard in self.guards.drain(..) {
            guard.join().expect("Worker panic");
        }
        // println!("WORKER THREADS JOINED");
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::BufRead;
use std::cmp::Ordering;

fn candies(ratings: &Vec<i32>) -> i32 {
	//println!("ratings:");
	//println!("{:?}", ratings);
	let mut last_candies = 1i32;
	let mut last_rating = ratings[0];
	let mut candies = Vec::with_capacity(ratings.len());
	candies.push(1);
	for index in 1..ratings.len() {
		let current_rating = ratings[index];
		match (current_rating).cmp(&last_rating) {
            Ordering::Less => {
            	if last_candies == 1 {
            		//println!("candies:");
            		//println!("{:?}", candies);

					let mut next_rating = current_rating;
					let mut next_candy = 1;
					for backward_index in (0..index).rev() {
	            		let mut current_candy = candies[backward_index];
            			let current_rating = ratings[backward_index];

						// println!("rating:{:?}, candy:{:?}",
						// 	current_rating, current_candy);
						if current_rating > next_rating
							&& current_candy <= next_candy {
								current_candy = next_candy + 1;
							}
						candies[backward_index] = current_candy;
				        next_candy = current_candy;
            			next_rating = current_rating;
					}
            	}
            	last_candies = 1;
            },
            Ordering::Equal => {
            	last_candies = 1;
            },
            Ordering::Greater => {
				last_candies += 1;
            },
        }
    	candies.push(last_candies);
		last_rating = current_rating;
	}
	//println!("candies:");
	//println!("{:?}", candies);
	candies.iter().fold(0, |sum, i| sum + *i)
}

#[allow(dead_code)]
fn main() {
    let stdin = io::stdin();
    let mut lines = stdin.lock().lines();
    let line_count = lines.next().unwrap().unwrap();
    let count = line_count.parse::<usize>().unwrap();
    let ratings: Vec<i32> = lines.take(count)
        .map(|value| value.unwrap().parse::<i32>().unwrap())
        .collect();
    let change_count = candies(&ratings);
    println!("{}", change_count);
}

#[cfg(test)]
mod tests {
    #[test]
    fn candies_test0() {
    	assert_eq!(4, super::candies(&vec![1, 2, 2]));
    }
    #[test]
    fn candies_test1() {
    	assert_eq!(19, super::candies(&vec![2, 4, 2, 6, 1, 7, 8, 9, 2, 1]));
    }
    #[test]
    fn candies_incr_decr() {
    	assert_eq!(16, super::candies(&vec![1, 2, 3, 4, 3, 2, 1]));
	}
    #[test]
    fn candies_decr_incr_decr() {
    	assert_eq!(9, super::candies(&vec![6, 5, 4, 5, 4]));
	}
}
<commit_msg>simpler faster algorithm for candies<commit_after>use std::io;
use std::io::BufRead;

fn candies(ratings: &Vec<i32>) -> i32 {
	let mut candies = vec![1i32; ratings.len()];
	for i in 1..ratings.len() {
		if ratings[i] > ratings[i-1] {
			candies[i] = candies[i-1] + 1;
		}
	}
	for i in (1..ratings.len()).rev() {
		if ratings[i-1] > ratings[i] {
			candies[i-1] = std::cmp::max(candies[i-1], candies[i] + 1);
		}
	}
	candies.iter().fold(0, |sum, i| sum + *i)
}

#[allow(dead_code)]
fn main() {
    let stdin = io::stdin();
    let mut lines = stdin.lock().lines();
    let line_count = lines.next().unwrap().unwrap();
    let count = line_count.parse::<usize>().unwrap();
    let ratings: Vec<i32> = lines.take(count)
        .map(|value| value.unwrap().parse::<i32>().unwrap())
        .collect();
    let change_count = candies(&ratings);
    println!("{}", change_count);
}

#[cfg(test)]
mod tests {
    #[test]
    fn candies_test0() {
    	assert_eq!(4, super::candies(&vec![1, 2, 2]));
    }
    #[test]
    fn candies_test1() {
    	assert_eq!(19, super::candies(&vec![2, 4, 2, 6, 1, 7, 8, 9, 2, 1]));
    }
    #[test]
    fn candies_incr_decr() {
    	assert_eq!(16, super::candies(&vec![1, 2, 3, 4, 3, 2, 1]));
	}
    #[test]
    fn candies_decr_incr_decr() {
    	assert_eq!(9, super::candies(&vec![6, 5, 4, 5, 4]));
	}
}
<|endoftext|>


<commit_before>use std::io::Result as IoResult;
use std::io::Write;
use nom::IError;

#[derive(Debug)]
struct Field<'a> {
    name: &'a str,
    repeated: bool,
    optional: bool,
    required: bool,
    typ: &'a str,
    number: i32,
    default: Option<&'a str>,
}

impl<'a> Field<'a> {

    fn rust_type(&self) -> &str {
        match self.typ {
            "int32" | "sint32" | "sfixed32" => "i32",
            "int64" | "sint64" | "sfixed64" => "i64",
            "uint32" | "fixed32" => "u32",
            "uint64" | "fixed64" => "u64",
            "float" => "f32",
            "double" => "f64",
            "string" => "String",
            "bytes" => "Vec<u8>",
            t => t,
        }
    }

    fn wire_type(&self, enums: &[&str]) -> &str {
        match self.typ {
            "int32" | "sint32" | "int64" | "sint64" |
                "uint32" | "uint64" | "bool" | "enum" => "WireType::Varint",
            "fixed64" | "sfixed64" | "double" => "WireType::Fixed64",
            "fixed32" | "sfixed32" | "float" => "WireType::Fixed32",
            t if enums.contains(&t) => "WireType::Varint",
            _ => "WireType::LengthDelimited",
        }
    }

    fn read_fn(&self, enums: &[&str]) -> &str {
        match self.typ {
            "int32" | "sint32" | "int64" | "sint64" |
                "uint32" | "uint64" | "bool" | "fixed64" |
                "sfixed64" | "double" | "fixed32" | "sfixed32" |
                "float" | "string" | "bytes" => self.typ,
            t if enums.contains(&t) => "enum",
            _ => "message",
        }
    }

    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        if self.optional {
            writeln!(w, "    pub {}: Option<{}>,", self.name, self.rust_type())
        } else if self.repeated {
            writeln!(w, "    pub {}: Vec<{}>,", self.name, self.rust_type())
        } else {
            writeln!(w, "    pub {}: {},", self.name, self.rust_type())
        }
    }

    fn write_match_tag<W: Write>(&self, w: &mut W, enums: &[&str]) -> IoResult<()> {
        if self.optional {
            writeln!(w, "({}, {}) => msg.{} = Some(r.read_{}()?),",
                self.number, self.wire_type(enums), self.name, self.read_fn(enums))
        } else if self.repeated {
            writeln!(w, "({}, {}) => msg.{}.push(r.read_{}()?),",
                self.number, self.wire_type(enums), self.name, self.read_fn(enums))
        } else {
            writeln!(w, "({}, {}) => msg.{} = r.read_{}()?,",
                self.number, self.wire_type(enums), self.name, self.read_fn(enums))
        }
    }

    fn write_match_not_tag<W: Write>(&self, w: &mut W, name: &str) -> IoResult<()> {
        writeln!(w, "({}, _) => return Err(ErrorKind::InvalidMessage(tag, \"{} {}\").into()),",
            self.number, name, self.typ)
    }
}

#[derive(Debug)]
pub struct Message<'a> {
    name: &'a str,
    fields: Vec<Field<'a>>,
}

impl<'a> Message<'a> {
    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "#[derive(Debug, Default, PartialEq, Clone)]")?;
        writeln!(w, "pub struct {} {{", self.name)?;
        for f in &self.fields {
            f.write_definition(w)?;
        }
        writeln!(w, "}}")
    }

    fn write_impl_message<W: Write>(&self, w: &mut W, enums: &[&str]) -> IoResult<()> {
        writeln!(w, "impl Message for {} {{", self.name)?;
        writeln!(w, "    fn from_reader<R: BufRead>(mut r: &mut Reader<R>) -> Result<Self> {{")?;
        writeln!(w, "        let mut msg = Self::default();")?;
        writeln!(w, "        loop {{")?;
        writeln!(w, "            let tag = match r.next_tag() {{")?;
        writeln!(w, "                None => break,")?;
        writeln!(w, "                Some(Err(e)) => return Err(e),")?;
        writeln!(w, "                Some(Ok(tag)) => tag,")?;
        writeln!(w, "            }};")?;
        writeln!(w, "            match tag.unpack() {{")?;
        for f in &self.fields {
            write!(w, "                 ")?;
            f.write_match_tag(w, enums)?;
        }
        for f in &self.fields {
            write!(w, "                 ")?;
            f.write_match_not_tag(w, self.name)?;
        }
        writeln!(w, "                (_, wire_type) => r.read_unknown(wire_type)?,")?;
        writeln!(w, "            }}")?;
        writeln!(w, "        }}")?;
        writeln!(w, "        Ok(msg)")?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }
}

#[derive(Debug)]
pub struct Enumerator<'a> {
    name: &'a str,
    fields: Vec<(&'a str, u32)>,
}

impl<'a> Enumerator<'a> {
    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "#[derive(Debug, PartialEq, Eq, Clone)]")?;
        writeln!(w, "pub enum {} {{", self.name)?;
        for &(f, _) in &self.fields {
            writeln!(w, "    {},", f)?;
        }
        writeln!(w, "}}")
    }

    fn write_impl_default<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "impl Default for {} {{", self.name)?;
        writeln!(w, "    fn default() -> Self {{")?;
        // TODO: check with default field and return error if there is no field
        writeln!(w, "        {}::{}", self.name, self.fields[0].0)?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }

    fn write_from_u64<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "impl From<u64> for {} {{", self.name)?;
        writeln!(w, "    fn from(i: u64) -> Self {{")?;
        writeln!(w, "        match i {{")?;
        for &(f, number) in &self.fields {
            writeln!(w, "            {} => {}::{},", number, self.name, f)?;
        }
        writeln!(w, "            _ => Self::default(),")?;
        writeln!(w, "        }}")?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }
}

#[derive(Debug)]
pub enum MessageOrEnum<'a> {
    Msg(Message<'a>),
    Enum(Enumerator<'a>),
    Ignore,
}

#[derive(Debug)]
pub struct FileDescriptor<'a> {
    message_and_enums: Vec<MessageOrEnum<'a>>,
}

impl<'a> FileDescriptor<'a> {
    pub fn from_str(f: &'a str) -> Result<FileDescriptor<'a>, IError<u32>> {
        FileDescriptor::from_bytes(f.as_bytes())

    }
    pub fn from_bytes(b: &'a [u8]) -> Result<FileDescriptor<'a>, IError<u32>> {
        self::parser::file_descriptor(b).to_full_result()
    }

    pub fn write<W: Write>(&self, w: &mut W, filename: &str) -> IoResult<()> {

        writeln!(w, "//! Automatically generated rust module for '{}' file", filename)?;
        writeln!(w, "")?;
        writeln!(w, "#![allow(non_snake_case)]")?;
        writeln!(w, "#![allow(non_upper_case_globals)]")?;
        writeln!(w, "")?;
        writeln!(w, "use std::io::BufRead;")?;
        writeln!(w, "use quick_protobuf::errors::{{Result, ErrorKind}};")?;
        writeln!(w, "use quick_protobuf::message::Message;")?;
        writeln!(w, "use quick_protobuf::reader::Reader;")?;
        writeln!(w, "use quick_protobuf::types::WireType;")?;

        let enums = self.message_and_enums.iter().filter_map(|m| {
            if let &MessageOrEnum::Enum(ref e) = m {
                Some(e.name)
            } else {
                None
            }
        }).collect::<Vec<_>>();
        for m in &self.message_and_enums {
            match m {
                &MessageOrEnum::Msg(ref m) => {
                    writeln!(w, "")?;
                    m.write_definition(w)?;
                    writeln!(w, "")?;
                    m.write_impl_message(w, &enums)?;
                },
                &MessageOrEnum::Enum(ref m) => {
                    writeln!(w, "")?;
                    m.write_definition(w)?;
                    writeln!(w, "")?;
                    m.write_impl_default(w)?;
                    writeln!(w, "")?;
                    m.write_from_u64(w)?;
                },
                &MessageOrEnum::Ignore => continue,
            }
        }
        Ok(())
    }
}

mod parser {

    use super::*;
    use std::str;
    use nom::{alphanumeric, multispace, digit};

    named!(comment<()>, do_parse!(
        many0!(multispace) >> tag!("//") >> take_until_and_consume!("\n") >> ()));

    named!(default_value<&str>, do_parse!(
        many0!(multispace) >> tag!("[") >> many0!(multispace) >> tag!("default") >> many0!(multispace) >> tag!("=") >>
        many0!(multispace) >> default: map_res!(alphanumeric, str::from_utf8) >> many0!(multispace) >> tag!("]") >> many0!(multispace) >>
        (default)));

    named!(message_field<Field>, do_parse!(
        many0!(multispace) >>
        repeated: opt!(do_parse!(tag!("repeated") >> many1!(multispace) >> ())) >>
        optional: opt!(do_parse!(tag!("optional") >> many1!(multispace) >> ())) >>
        required: opt!(do_parse!(tag!("required") >> many1!(multispace) >> ())) >>
        typ: map_res!(alphanumeric, str::from_utf8) >> many1!(multispace) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(multispace) >>
        tag!("=") >> many0!(multispace) >>
        number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >>
        many0!(multispace) >> default: opt!(default_value) >>
        tag!(";") >> opt!(comment) >>
        (Field {
           name: name,
           repeated: repeated.is_some(),
           typ: typ,
           number: number,
           optional: optional.is_some(),
           required: required.is_some(),
           default: default,
        })));

    named!(message<Message>, do_parse!(
        tag!("message") >> opt!(multispace) >>
        name: map_res!(alphanumeric, str::from_utf8) >> opt!(multispace) >>
        tag!("{") >>
        fields: many0!(message_field) >>
        many0!(multispace) >> tag!("}") >>
        (Message {
           name: name,
           fields: fields,
        })));

    named!(enum_field<(&str, u32)>, do_parse!(
        many0!(multispace) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(multispace) >>
        tag!("=") >> many0!(multispace) >>
        number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(multispace) >>
        tag!(";") >>
        ((name, number))));

    named!(enumerator<Enumerator>, do_parse!(
        tag!("enum") >> many1!(multispace) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(multispace) >>
        tag!("{") >>
        fields: many0!(enum_field) >>
        many0!(multispace) >> tag!("}") >>
        (Enumerator {
            name: name,
            fields: fields,
        })));

    named!(ignore<()>, do_parse!(
        alt!(tag!("package") | tag!("option")) >> many1!(multispace) >>
        take_until_and_consume!(";") >> ()));

    named!(message_or_enum<MessageOrEnum>, do_parse!(many0!(multispace) >>
        msg_or_enum: alt!(
             message => { |m| MessageOrEnum::Msg(m) } |
             enumerator => { |e| MessageOrEnum::Enum(e) } |
             ignore => { |_| MessageOrEnum::Ignore } ) >>
        (msg_or_enum)));

    named!(pub file_descriptor<FileDescriptor>, do_parse!(
        message_and_enums: many0!(message_or_enum) >> many0!(multispace) >>
        (FileDescriptor {
            message_and_enums: message_and_enums
        })));

    #[test]
    fn test_message() {
        let msg = r#"message ReferenceData
    {
        repeated ScenarioInfo  scenarioSet = 1;
        repeated CalculatedObjectInfo calculatedObjectSet = 2;
        repeated RiskFactorList riskFactorListSet = 3;
        repeated RiskMaturityInfo riskMaturitySet = 4;
        repeated IndicatorInfo indicatorSet = 5;
        repeated RiskStrikeInfo riskStrikeSet = 6;
        repeated FreeProjectionList freeProjectionListSet = 7;
        repeated ValidationProperty ValidationSet = 8;
        repeated CalcProperties calcPropertiesSet = 9;
        repeated MaturityInfo maturitySet = 10;
    }"#;

        let mess = message(msg.as_bytes());
        if let ::nom::IResult::Done(_, mess) = mess {
            assert_eq!(10, mess.fields.len());
        }
    }

    #[test]
    fn test_enum() {
        let msg = r#"enum PairingStatus {
                DEALPAIRED        = 0;
                INVENTORYORPHAN   = 1;
                CALCULATEDORPHAN  = 2;
                CANCELED          = 3;
    }"#;

        let mess = enumerator(msg.as_bytes());
        if let ::nom::IResult::Done(_, mess) = mess {
            assert_eq!(4, mess.fields.len());
        }
    }

    #[test]
    fn test_ignore() {
        let msg = r#"package com.test.v0;

    option optimize_for = SPEED;
    "#;

        match ignore(msg.as_bytes()) {
            ::nom::IResult::Done(_, _) => (),
            e => panic!("Expecting done {:?}", e),
        }
    }
}
<commit_msg>simpler protorust_codegen<commit_after>use std::io::Result as IoResult;
use std::io::Write;
use nom::IError;

#[derive(Debug)]
enum Frequency {
    Optional,
    Repeated,
    Required,
}

#[derive(Debug)]
struct Field<'a> {
    name: &'a str,
    frequency: Frequency,
    typ: &'a str,
    number: i32,
    default: Option<&'a str>,
}

impl<'a> Field<'a> {

    fn rust_type(&self) -> &str {
        match self.typ {
            "int32" | "sint32" | "sfixed32" => "i32",
            "int64" | "sint64" | "sfixed64" => "i64",
            "uint32" | "fixed32" => "u32",
            "uint64" | "fixed64" => "u64",
            "float" => "f32",
            "double" => "f64",
            "string" => "String",
            "bytes" => "Vec<u8>",
            t => t,
        }
    }

    fn wire_type(&self, enums: &[&str]) -> &str {
        match self.typ {
            "int32" | "sint32" | "int64" | "sint64" |
                "uint32" | "uint64" | "bool" | "enum" => "WireType::Varint",
            "fixed64" | "sfixed64" | "double" => "WireType::Fixed64",
            "fixed32" | "sfixed32" | "float" => "WireType::Fixed32",
            t if enums.contains(&t) => "WireType::Varint",
            _ => "WireType::LengthDelimited",
        }
    }

    fn read_fn(&self, enums: &[&str]) -> &str {
        match self.typ {
            "int32" | "sint32" | "int64" | "sint64" |
                "uint32" | "uint64" | "bool" | "fixed64" |
                "sfixed64" | "double" | "fixed32" | "sfixed32" |
                "float" | "string" | "bytes" => self.typ,
            t if enums.contains(&t) => "enum",
            _ => "message",
        }
    }

    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        match self.frequency {
            Frequency::Optional => writeln!(w, "    pub {}: Option<{}>,", self.name, self.rust_type()),
            Frequency::Repeated => writeln!(w, "    pub {}: Vec<{}>,", self.name, self.rust_type()),
            Frequency::Required => writeln!(w, "    pub {}: {},", self.name, self.rust_type()),
        }
    }

    fn write_match_tag<W: Write>(&self, w: &mut W, enums: &[&str]) -> IoResult<()> {
        match self.frequency {
            Frequency::Optional => {
                writeln!(w, "({}, {}) => msg.{} = Some(r.read_{}()?),",
                    self.number, self.wire_type(enums), self.name, self.read_fn(enums))
            }
            Frequency::Repeated => {
                writeln!(w, "({}, {}) => msg.{}.push(r.read_{}()?),",
                    self.number, self.wire_type(enums), self.name, self.read_fn(enums))
            }
            Frequency::Required => {
                writeln!(w, "({}, {}) => msg.{} = r.read_{}()?,",
                    self.number, self.wire_type(enums), self.name, self.read_fn(enums))
            }
        }
    }

    fn write_match_not_tag<W: Write>(&self, w: &mut W, name: &str) -> IoResult<()> {
        writeln!(w, "({}, _) => return Err(ErrorKind::InvalidMessage(tag, \"{} {}\").into()),",
            self.number, name, self.typ)
    }
}

#[derive(Debug)]
pub struct Message<'a> {
    name: &'a str,
    fields: Vec<Field<'a>>,
}

impl<'a> Message<'a> {
    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "#[derive(Debug, Default, PartialEq, Clone)]")?;
        writeln!(w, "pub struct {} {{", self.name)?;
        for f in &self.fields {
            f.write_definition(w)?;
        }
        writeln!(w, "}}")
    }

    fn write_impl_message<W: Write>(&self, w: &mut W, enums: &[&str]) -> IoResult<()> {
        writeln!(w, "impl Message for {} {{", self.name)?;
        writeln!(w, "    fn from_reader<R: BufRead>(mut r: &mut Reader<R>) -> Result<Self> {{")?;
        writeln!(w, "        let mut msg = Self::default();")?;
        writeln!(w, "        loop {{")?;
        writeln!(w, "            let tag = match r.next_tag() {{")?;
        writeln!(w, "                None => break,")?;
        writeln!(w, "                Some(Err(e)) => return Err(e),")?;
        writeln!(w, "                Some(Ok(tag)) => tag,")?;
        writeln!(w, "            }};")?;
        writeln!(w, "            match tag.unpack() {{")?;
        for f in &self.fields {
            write!(w, "                 ")?;
            f.write_match_tag(w, enums)?;
        }
        for f in &self.fields {
            write!(w, "                 ")?;
            f.write_match_not_tag(w, self.name)?;
        }
        writeln!(w, "                (_, wire_type) => r.read_unknown(wire_type)?,")?;
        writeln!(w, "            }}")?;
        writeln!(w, "        }}")?;
        writeln!(w, "        Ok(msg)")?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }
}

#[derive(Debug)]
pub struct Enumerator<'a> {
    name: &'a str,
    fields: Vec<(&'a str, i32)>,
}

impl<'a> Enumerator<'a> {
    fn write_definition<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "#[derive(Debug, PartialEq, Eq, Clone)]")?;
        writeln!(w, "pub enum {} {{", self.name)?;
        for &(f, _) in &self.fields {
            writeln!(w, "    {},", f)?;
        }
        writeln!(w, "}}")
    }

    fn write_impl_default<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "impl Default for {} {{", self.name)?;
        writeln!(w, "    fn default() -> Self {{")?;
        // TODO: check with default field and return error if there is no field
        writeln!(w, "        {}::{}", self.name, self.fields[0].0)?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }

    fn write_from_u64<W: Write>(&self, w: &mut W) -> IoResult<()> {
        writeln!(w, "impl From<u64> for {} {{", self.name)?;
        writeln!(w, "    fn from(i: u64) -> Self {{")?;
        writeln!(w, "        match i {{")?;
        for &(f, number) in &self.fields {
            writeln!(w, "            {} => {}::{},", number, self.name, f)?;
        }
        writeln!(w, "            _ => Self::default(),")?;
        writeln!(w, "        }}")?;
        writeln!(w, "    }}")?;
        writeln!(w, "}}")
    }
}

#[derive(Debug)]
pub enum MessageOrEnum<'a> {
    Msg(Message<'a>),
    Enum(Enumerator<'a>),
    Ignore,
}

#[derive(Debug)]
pub struct FileDescriptor<'a> {
    message_and_enums: Vec<MessageOrEnum<'a>>,
}

impl<'a> FileDescriptor<'a> {
    pub fn from_str(f: &'a str) -> Result<FileDescriptor<'a>, IError<u32>> {
        FileDescriptor::from_bytes(f.as_bytes())

    }
    pub fn from_bytes(b: &'a [u8]) -> Result<FileDescriptor<'a>, IError<u32>> {
        self::parser::file_descriptor(b).to_full_result()
    }

    pub fn write<W: Write>(&self, w: &mut W, filename: &str) -> IoResult<()> {

        writeln!(w, "//! Automatically generated rust module for '{}' file", filename)?;
        writeln!(w, "")?;
        writeln!(w, "#![allow(non_snake_case)]")?;
        writeln!(w, "#![allow(non_upper_case_globals)]")?;
        writeln!(w, "")?;
        writeln!(w, "use std::io::BufRead;")?;
        writeln!(w, "use quick_protobuf::errors::{{Result, ErrorKind}};")?;
        writeln!(w, "use quick_protobuf::message::Message;")?;
        writeln!(w, "use quick_protobuf::reader::Reader;")?;
        writeln!(w, "use quick_protobuf::types::WireType;")?;

        let enums = self.message_and_enums.iter().filter_map(|m| {
            if let &MessageOrEnum::Enum(ref e) = m {
                Some(e.name)
            } else {
                None
            }
        }).collect::<Vec<_>>();
        for m in &self.message_and_enums {
            match m {
                &MessageOrEnum::Msg(ref m) => {
                    writeln!(w, "")?;
                    m.write_definition(w)?;
                    writeln!(w, "")?;
                    m.write_impl_message(w, &enums)?;
                },
                &MessageOrEnum::Enum(ref m) => {
                    writeln!(w, "")?;
                    m.write_definition(w)?;
                    writeln!(w, "")?;
                    m.write_impl_default(w)?;
                    writeln!(w, "")?;
                    m.write_from_u64(w)?;
                },
                &MessageOrEnum::Ignore => continue,
            }
        }
        Ok(())
    }
}

mod parser {

    use super::*;
    use std::str;
    use nom::{alphanumeric, multispace, digit};

    named!(comment<()>, do_parse!(tag!("//") >> take_until_and_consume!("\n") >> ()));

    /// word break: multispace or comment
    named!(br<()>, alt!(map!(multispace, |_| ()) | comment));

    named!(default_value<&str>, do_parse!(
        tag!("[") >> many0!(br) >> tag!("default") >> many0!(br) >> tag!("=") >> many0!(br) >>
        default: map_res!(alphanumeric, str::from_utf8) >> many0!(br) >> tag!("]") >>
        (default)));

    named!(message_field<Field>, do_parse!(
        frequency: alt!(map!(tag!("optional"), |_| Frequency::Optional) |
                        map!(tag!("repeated"), |_| Frequency::Repeated) |
                        map!(tag!("required"), |_| Frequency::Required)) >> many1!(br) >>
        typ: map_res!(alphanumeric, str::from_utf8) >> many1!(br) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(br) >>
        tag!("=") >> many0!(br) >>
        number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
        default: opt!(default_value) >> many0!(br) >> tag!(";") >> many0!(br) >>
        (Field {
           name: name,
           frequency: frequency,
           typ: typ,
           number: number,
           default: default,
        })));

    named!(message<Message>, do_parse!(
        tag!("message") >> many0!(br) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(br) >>
        tag!("{") >> many0!(br) >>
        fields: many0!(message_field) >>
        tag!("}") >> many0!(br) >>
        (Message { name: name, fields: fields })));

    named!(enum_field<(&str, i32)>, do_parse!(
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(br) >>
        tag!("=") >> many0!(br) >>
        number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
        tag!(";") >> many0!(br) >>
        ((name, number))));

    named!(enumerator<Enumerator>, do_parse!(
        tag!("enum") >> many1!(br) >>
        name: map_res!(alphanumeric, str::from_utf8) >> many0!(br) >>
        tag!("{") >> many0!(br) >>
        fields: many0!(enum_field) >>
        tag!("}") >> many0!(br) >>
        (Enumerator { name: name, fields: fields })));

    named!(ignore<()>, do_parse!(
        alt!(tag!("package") | tag!("option")) >> many1!(br) >>
        take_until_and_consume!(";") >> many0!(br) >> ()));

    named!(message_or_enum<MessageOrEnum>, alt!(
             message => { |m| MessageOrEnum::Msg(m) } |
             enumerator => { |e| MessageOrEnum::Enum(e) } |
             ignore => { |_| MessageOrEnum::Ignore } ));

    named!(pub file_descriptor<FileDescriptor>, do_parse!(
        many0!(br) >>
        message_and_enums: many0!(message_or_enum) >>
        (FileDescriptor {
            message_and_enums: message_and_enums
        })));

    #[test]
    fn test_message() {
        let msg = r#"message ReferenceData
    {
        repeated ScenarioInfo  scenarioSet = 1;
        repeated CalculatedObjectInfo calculatedObjectSet = 2;
        repeated RiskFactorList riskFactorListSet = 3;
        repeated RiskMaturityInfo riskMaturitySet = 4;
        repeated IndicatorInfo indicatorSet = 5;
        repeated RiskStrikeInfo riskStrikeSet = 6;
        repeated FreeProjectionList freeProjectionListSet = 7;
        repeated ValidationProperty ValidationSet = 8;
        repeated CalcProperties calcPropertiesSet = 9;
        repeated MaturityInfo maturitySet = 10;
    }"#;

        let mess = message(msg.as_bytes());
        if let ::nom::IResult::Done(_, mess) = mess {
            assert_eq!(10, mess.fields.len());
        }
    }

    #[test]
    fn test_enum() {
        let msg = r#"enum PairingStatus {
                DEALPAIRED        = 0;
                INVENTORYORPHAN   = 1;
                CALCULATEDORPHAN  = 2;
                CANCELED          = 3;
    }"#;

        let mess = enumerator(msg.as_bytes());
        if let ::nom::IResult::Done(_, mess) = mess {
            assert_eq!(4, mess.fields.len());
        }
    }

    #[test]
    fn test_ignore() {
        let msg = r#"package com.test.v0;

    option optimize_for = SPEED;
    "#;

        match ignore(msg.as_bytes()) {
            ::nom::IResult::Done(_, _) => (),
            e => panic!("Expecting done {:?}", e),
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
use stdsimd_test::assert_instr;

extern "C" {
    #[link_name = "llvm.x86.xbegin"]
    fn x86_xbegin() -> i32;
    #[link_name = "llvm.x86.xend"]
    fn x86_xend() -> ();
    #[link_name = "llvm.x86.xabort"]
    fn x86_xabort(imm8: i8) -> ();
    #[link_name = "llvm.x86.xtest"]
    fn x86_xtest() -> i32;
}

/// Transaction successfully started.
pub const _XBEGIN_STARTED: u32 = !0;

/// Transaction explicitly aborted with xabort. The parameter passed to xabort is available with
/// _xabort_code(status).
pub const _XABORT_EXPLICIT: u32 = 1 << 0;

/// Transaction retry is possible.
pub const _XABORT_RETRY: u32 = 1 << 1;

/// Transaction abort due to a memory conflict with another thread.
pub const _XABORT_CONFLICT: u32 = 1 << 2;

/// Transaction abort due to the transaction using too much memory.
pub const _XABORT_CAPACITY: u32 = 1 << 3;

/// Transaction abort due to a debug trap.
pub const _XABORT_DEBUG: u32 = 1 << 4;

/// Transaction abort in a inner nested transaction.
pub const _XABORT_NESTED: u32 = 1 << 5;

/// Specifies the start of a restricted transactional memory (RTM) code region and returns a value
/// indicating status.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xbegin).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xbegin))]
pub unsafe fn _xbegin() -> u32 {
    x86_xbegin() as _
}

/// Specifies the end of a restricted transactional memory (RTM) code region.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xend).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xend))]
pub unsafe fn _xend() {
    x86_xend()
}

/// Forces a restricted transactional memory (RTM) region to abort.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xabort).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xabort))]
pub unsafe fn _xabort(imm8: u32) {
    macro_rules! call {
        ($imm8:expr) => {
            x86_xabort($imm8)
        };
    }
    constify_imm8!(imm8, call)
}

/// Queries whether the processor is executing in a transactional region identified by restricted
/// transactional memory (RTM) or hardware lock elision (HLE).
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xtest).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xtest))]
pub unsafe fn _xtest() -> bool {
    x86_xtest() != 0
}

/// Retrieves the parameter passed to [`_xabort`] when [`_xbegin`]'s status has the `_XABORT_EXPLICIT` flag set.
#[inline]
pub fn _xabort_code(status: u32) -> u32 {
    (status >> 24) & 0xFF
}

#[cfg(test)]
mod tests {
    use crate::core_arch::x86::*;

    #[test]
    fn test_xbegin_xend() {
        unsafe {
            let mut x = 0;
            for _ in 0..10 {
                let code = rtm::_xbegin();
                if code == _XBEGIN_STARTED {
                    x += 1;
                    rtm::_xend();
                    assert_eq!(x, 1);
                    break
                }
                assert_eq!(x, 0);
            }
        }
    }

    #[test]
    fn test_xabort() {
        unsafe {
            // aborting with outside a transactional region does nothing
            _xabort(0);

            for abort_code in 0..10 {
                let mut x = 0;
                let code = rtm::_xbegin();
                if code == _XBEGIN_STARTED {
                    x += 1;
                    rtm::_xabort(abort_code);
                } else if code & _XABORT_EXPLICIT != 0 {
                    let test_abort_code = rtm::_xabort_code(code);
                    assert_eq!(test_abort_code, abort_code);
                }
                assert_eq!(x, 0);
            }
        }
    }

    #[test]
    fn test_xtest() {
        unsafe {
            assert_eq!(_xtest(), false);

            for _ in 0..10 {
                let code = rtm::_xbegin();
                if code == _XBEGIN_STARTED {
                    let in_tx = _xtest();
                    rtm::_xend();

                    // putting the assert inside the transaction would abort the transaction on fail
                    // without any output/panic/etc
                    assert_eq!(in_tx, true);
                    break
                }
            }
        }
    }
}<commit_msg>make _xabort_code const add top level comment use simd_test<commit_after>//! Intel's Restricted Transactional Memory (RTM).
//!
//! This CPU feature is available on Intel Broadwell or later CPUs (and some Haswell).
//!
//! [Wikipedia][wikipedia_rtm] provides a quick overview of the assembly instructions, and
//! Intel's [programming considerations][intel_consid] details what sorts of instructions within a
//! transaction are likely to cause an abort.
//!
//! [wikipedia_rtm]: https://en.wikipedia.org/wiki/Transactional_Synchronization_Extensions#Restricted_Transactional_Memory
//! [intel_consid]: https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-intel-transactional-synchronization-extensions-intel-tsx-programming-considerations

#[cfg(test)]
use stdsimd_test::assert_instr;

extern "C" {
    #[link_name = "llvm.x86.xbegin"]
    fn x86_xbegin() -> i32;
    #[link_name = "llvm.x86.xend"]
    fn x86_xend() -> ();
    #[link_name = "llvm.x86.xabort"]
    fn x86_xabort(imm8: i8) -> ();
    #[link_name = "llvm.x86.xtest"]
    fn x86_xtest() -> i32;
}

/// Transaction successfully started.
pub const _XBEGIN_STARTED: u32 = !0;

/// Transaction explicitly aborted with xabort. The parameter passed to xabort is available with
/// _xabort_code(status).
pub const _XABORT_EXPLICIT: u32 = 1 << 0;

/// Transaction retry is possible.
pub const _XABORT_RETRY: u32 = 1 << 1;

/// Transaction abort due to a memory conflict with another thread.
pub const _XABORT_CONFLICT: u32 = 1 << 2;

/// Transaction abort due to the transaction using too much memory.
pub const _XABORT_CAPACITY: u32 = 1 << 3;

/// Transaction abort due to a debug trap.
pub const _XABORT_DEBUG: u32 = 1 << 4;

/// Transaction abort in a inner nested transaction.
pub const _XABORT_NESTED: u32 = 1 << 5;

/// Specifies the start of a restricted transactional memory (RTM) code region and returns a value
/// indicating status.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xbegin).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xbegin))]
pub unsafe fn _xbegin() -> u32 {
    x86_xbegin() as _
}

/// Specifies the end of a restricted transactional memory (RTM) code region.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xend).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xend))]
pub unsafe fn _xend() {
    x86_xend()
}

/// Forces a restricted transactional memory (RTM) region to abort.
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xabort).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xabort))]
pub unsafe fn _xabort(imm8: u32) {
    macro_rules! call {
        ($imm8:expr) => {
            x86_xabort($imm8)
        };
    }
    constify_imm8!(imm8, call)
}

/// Queries whether the processor is executing in a transactional region identified by restricted
/// transactional memory (RTM) or hardware lock elision (HLE).
///
/// [Intel's documentation](https://software.intel.com/en-us/cpp-compiler-developer-guide-and-reference-xtest).
#[inline]
#[target_feature(enable = "rtm")]
#[cfg_attr(test, assert_instr(xtest))]
pub unsafe fn _xtest() -> bool {
    x86_xtest() != 0
}

/// Retrieves the parameter passed to [`_xabort`] when [`_xbegin`]'s status has the
/// `_XABORT_EXPLICIT` flag set.
#[inline]
pub const fn _xabort_code(status: u32) -> u32 {
    (status >> 24) & 0xFF
}

#[cfg(test)]
mod tests {
    use stdsimd_test::simd_test;

    use crate::core_arch::x86::*;

    #[simd_test(enable = "rtm")]
    unsafe fn test_xbegin_xend() {
        let mut x = 0;
        for _ in 0..10 {
            let code = rtm::_xbegin();
            if code == _XBEGIN_STARTED {
                x += 1;
                rtm::_xend();
                assert_eq!(x, 1);
                break
            }
            assert_eq!(x, 0);
        }
    }

    #[simd_test(enable = "rtm")]
    unsafe fn test_xabort() {
        // aborting outside a transactional region does nothing
        _xabort(0);

        for abort_code in 0..10 {
            let mut x = 0;
            let code = rtm::_xbegin();
            if code == _XBEGIN_STARTED {
                x += 1;
                rtm::_xabort(abort_code);
            } else if code & _XABORT_EXPLICIT != 0 {
                let test_abort_code = rtm::_xabort_code(code);
                assert_eq!(test_abort_code, abort_code);
            }
            assert_eq!(x, 0);
        }
    }

    #[simd_test(enable = "rtm")]
    unsafe fn test_xtest() {
        assert_eq!(_xtest(), false);

        for _ in 0..10 {
            let code = rtm::_xbegin();
            if code == _XBEGIN_STARTED {
                let in_tx = _xtest();
                rtm::_xend();

                // putting the assert inside the transaction would abort the transaction on fail
                // without any output/panic/etc
                assert_eq!(in_tx, true);
                break
            }
        }
    }
}<|endoftext|>