<commit_before>//! Module converting command-line arguments into test configuration.

use std::env;
use std::path::PathBuf;

use super::helpers::isatty;
use super::options::{ColorConfig, Options, OutputFormat, RunIgnored};
use super::time::TestTimeOptions;

#[derive(Debug)]
pub struct TestOpts {
    pub list: bool,
    pub filter: Option<String>,
    pub filter_exact: bool,
    pub force_run_in_process: bool,
    pub exclude_should_panic: bool,
    pub run_ignored: RunIgnored,
    pub run_tests: bool,
    pub bench_benchmarks: bool,
    pub logfile: Option<PathBuf>,
    pub nocapture: bool,
    pub color: ColorConfig,
    pub format: OutputFormat,
    pub test_threads: Option<usize>,
    pub skip: Vec<String>,
    pub time_options: Option<TestTimeOptions>,
    pub options: Options,
}

impl TestOpts {
    pub fn use_color(&self) -> bool {
        match self.color {
            ColorConfig::AutoColor => !self.nocapture && isatty::stdout_isatty(),
            ColorConfig::AlwaysColor => true,
            ColorConfig::NeverColor => false,
        }
    }
}

/// Result of parsing the options.
pub type OptRes = Result<TestOpts, String>;
/// Result of parsing the option part.
type OptPartRes<T> = Result<T, String>;

fn optgroups() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optflag("", "include-ignored", "Run ignored and not ignored tests")
        .optflag("", "ignored", "Run only ignored tests")
        .optflag("", "force-run-in-process", "Forces tests to run in-process when panic=abort")
        .optflag("", "exclude-should-panic", "Excludes tests marked as should_panic")
        .optflag("", "test", "Run tests and not benchmarks")
        .optflag("", "bench", "Run benchmarks instead of tests")
        .optflag("", "list", "List all tests and benchmarks")
        .optflag("h", "help", "Display this message (longer with --help)")
        .optopt(
            "",
            "logfile",
            "Write logs to the specified file instead \
             of stdout",
            "PATH",
        )
        .optflag(
            "",
            "nocapture",
            "don't capture stdout/stderr of each \
             task, allow printing directly",
        )
        .optopt(
            "",
            "test-threads",
            "Number of threads used for running tests \
             in parallel",
            "n_threads",
        )
        .optmulti(
            "",
            "skip",
            "Skip tests whose names contain FILTER (this flag can \
             be used multiple times)",
            "FILTER",
        )
        .optflag(
            "q",
            "quiet",
            "Display one character per test instead of one line. \
             Alias to --format=terse",
        )
        .optflag("", "exact", "Exactly match filters rather than by substring")
        .optopt(
            "",
            "color",
            "Configure coloring of output:
            auto   = colorize if stdout is a tty and tests are run on serially (default);
            always = always colorize output;
            never  = never colorize output;",
            "auto|always|never",
        )
        .optopt(
            "",
            "format",
            "Configure formatting of output:
            pretty = Print verbose output;
            terse  = Display one character per test;
            json   = Output a json document",
            "pretty|terse|json",
        )
        .optflag("", "show-output", "Show captured stdout of successful tests")
        .optopt(
            "Z",
            "",
            "Enable nightly-only flags:
            unstable-options = Allow use of experimental features",
            "unstable-options",
        )
        .optflagopt(
            "",
            "report-time",
            "Show execution time of each test. Awailable values:
            plain   = do not colorize the execution time (default);
            colored = colorize output according to the `color` parameter value;

            Threshold values for colorized output can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.
            Durations must be specified in milliseconds, e.g. `500,2000` means that the warn time
            is 0.5 seconds, and the critical time is 2 seconds.

            Not available for --format=terse",
            "plain|colored",
        )
        .optflag(
            "",
            "ensure-time",
            "Treat excess of the test execution time limit as error.

            Threshold values for this option can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.

            `CRITICAL_TIME` here means the limit that should not be exceeded by test.
            ",
        );
    opts
}

fn usage(binary: &str, options: &getopts::Options) {
    let message = format!("Usage: {} [OPTIONS] [FILTER]", binary);
    println!(
        r#"{usage}

The FILTER string is tested against the name of all tests, and only those
tests whose names contain the filter are run.

By default, all tests are run in parallel. This can be altered with the
--test-threads flag or the RUST_TEST_THREADS environment variable when running
tests (set it to 1).

All tests have their standard output and standard error captured by default.
This can be overridden with the --nocapture flag or setting RUST_TEST_NOCAPTURE
environment variable to a value other than "0". Logging is not captured by default.

Test Attributes:

    `#[test]`        - Indicates a function is a test to be run. This function
                       takes no arguments.
    `#[bench]`       - Indicates a function is a benchmark to be run. This
                       function takes one argument (test::Bencher).
    `#[should_panic]` - This function (also labeled with `#[test]`) will only pass if
                        the code causes a panic (an assertion failure or panic!)
                        A message may be provided, which the failure string must
                        contain: #[should_panic(expected = "foo")].
    `#[ignore]`       - When applied to a function which is already attributed as a
                        test, then the test runner will ignore these tests during
                        normal test runs. Running with --ignored or --include-ignored will run
                        these tests."#,
        usage = options.usage(&message)
    );
}

/// Parses command line arguments into test options.
/// Returns `None` if help was requested (since we only show help message and don't run tests),
/// returns `Some(Err(..))` if provided arguments are incorrect,
/// otherwise creates a `TestOpts` object and returns it.
pub fn parse_opts(args: &[String]) -> Option<OptRes> {
    // Parse matches.
    let opts = optgroups();
    let args = args.get(1..).unwrap_or(args);
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(f) => return Some(Err(f.to_string())),
    };

    // Check if help was requested.
    if matches.opt_present("h") {
        // Show help and do nothing more.
        usage(&args[0], &opts);
        return None;
    }

    // Actually parse the opts.
    let opts_result = parse_opts_impl(matches);

    Some(opts_result)
}

// Gets the option value and checks if unstable features are enabled.
macro_rules! unstable_optflag {
    ($matches:ident, $allow_unstable:ident, $option_name:literal) => {{
        let opt = $matches.opt_present($option_name);
        if !$allow_unstable && opt {
            return Err(format!(
                "The \"{}\" flag is only accepted on the nightly compiler",
                $option_name
            ));
        }

        opt
    }};
}

// Implementation of `parse_opts` that doesn't care about help message
// and returns a `Result`.
fn parse_opts_impl(matches: getopts::Matches) -> OptRes {
    let allow_unstable = get_allow_unstable(&matches)?;

    // Unstable flags
    let force_run_in_process = unstable_optflag!(matches, allow_unstable, "force-run-in-process");
    let exclude_should_panic = unstable_optflag!(matches, allow_unstable, "exclude-should-panic");
    let include_ignored = unstable_optflag!(matches, allow_unstable, "include-ignored");
    let time_options = get_time_options(&matches, allow_unstable)?;

    let quiet = matches.opt_present("quiet");
    let exact = matches.opt_present("exact");
    let list = matches.opt_present("list");
    let skip = matches.opt_strs("skip");

    let bench_benchmarks = matches.opt_present("bench");
    let run_tests = !bench_benchmarks || matches.opt_present("test");

    let logfile = get_log_file(&matches)?;
    let run_ignored = get_run_ignored(&matches, include_ignored)?;
    let filter = get_filter(&matches)?;
    let nocapture = get_nocapture(&matches)?;
    let test_threads = get_test_threads(&matches)?;
    let color = get_color_config(&matches)?;
    let format = get_format(&matches, quiet, allow_unstable)?;

    let options = Options::new().display_output(matches.opt_present("show-output"));

    let test_opts = TestOpts {
        list,
        filter,
        filter_exact: exact,
        force_run_in_process,
        exclude_should_panic,
        run_ignored,
        run_tests,
        bench_benchmarks,
        logfile,
        nocapture,
        color,
        format,
        test_threads,
        skip,
        time_options,
        options,
    };

    Ok(test_opts)
}

// FIXME: Copied from librustc_ast until linkage errors are resolved. Issue #47566
fn is_nightly() -> bool {
    // Whether this is a feature-staged build, i.e., on the beta or stable channel
    let disable_unstable_features = option_env!("CFG_DISABLE_UNSTABLE_FEATURES").is_some();
    // Whether we should enable unstable features for bootstrapping
    let bootstrap = env::var("RUSTC_BOOTSTRAP").is_ok();

    bootstrap || !disable_unstable_features
}

// Gets the CLI options associated with `report-time` feature.
fn get_time_options(
    matches: &getopts::Matches,
    allow_unstable: bool,
) -> OptPartRes<Option<TestTimeOptions>> {
    let report_time = unstable_optflag!(matches, allow_unstable, "report-time");
    let colored_opt_str = matches.opt_str("report-time");
    let mut report_time_colored = report_time && colored_opt_str == Some("colored".into());
    let ensure_test_time = unstable_optflag!(matches, allow_unstable, "ensure-time");

    // If `ensure-test-time` option is provided, time output is enforced,
    // so user won't be confused if any of tests will silently fail.
    let options = if report_time || ensure_test_time {
        if ensure_test_time && !report_time {
            report_time_colored = true;
        }
        Some(TestTimeOptions::new_from_env(ensure_test_time, report_time_colored))
    } else {
        None
    };

    Ok(options)
}

fn get_test_threads(matches: &getopts::Matches) -> OptPartRes<Option<usize>> {
    let test_threads = match matches.opt_str("test-threads") {
        Some(n_str) => match n_str.parse::<usize>() {
            Ok(0) => return Err("argument for --test-threads must not be 0".to_string()),
            Ok(n) => Some(n),
            Err(e) => {
                return Err(format!(
                    "argument for --test-threads must be a number > 0 \
                     (error: {})",
                    e
                ));
            }
        },
        None => None,
    };

    Ok(test_threads)
}

fn get_format(
    matches: &getopts::Matches,
    quiet: bool,
    allow_unstable: bool,
) -> OptPartRes<OutputFormat> {
    let format = match matches.opt_str("format").as_ref().map(|s| &**s) {
        None if quiet => OutputFormat::Terse,
        Some("pretty") | None => OutputFormat::Pretty,
        Some("terse") => OutputFormat::Terse,
        Some("json") => {
            if !allow_unstable {
                return Err("The \"json\" format is only accepted on the nightly compiler".into());
            }
            OutputFormat::Json
        }

        Some(v) => {
            return Err(format!(
                "argument for --format must be pretty, terse, or json (was \
                 {})",
                v
            ));
        }
    };

    Ok(format)
}

fn get_color_config(matches: &getopts::Matches) -> OptPartRes<ColorConfig> {
    let color = match matches.opt_str("color").as_ref().map(|s| &**s) {
        Some("auto") | None => ColorConfig::AutoColor,
        Some("always") => ColorConfig::AlwaysColor,
        Some("never") => ColorConfig::NeverColor,

        Some(v) => {
            return Err(format!(
                "argument for --color must be auto, always, or never (was \
                 {})",
                v
            ));
        }
    };

    Ok(color)
}

fn get_nocapture(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut nocapture = matches.opt_present("nocapture");
    if !nocapture {
        nocapture = match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        };
    }

    Ok(nocapture)
}

fn get_run_ignored(matches: &getopts::Matches, include_ignored: bool) -> OptPartRes<RunIgnored> {
    let run_ignored = match (include_ignored, matches.opt_present("ignored")) {
        (true, true) => {
            return Err("the options --include-ignored and --ignored are mutually exclusive".into());
        }
        (true, false) => RunIgnored::Yes,
        (false, true) => RunIgnored::Only,
        (false, false) => RunIgnored::No,
    };

    Ok(run_ignored)
}

fn get_filter(matches: &getopts::Matches) -> OptPartRes<Option<String>> {
    let filter = if !matches.free.is_empty() { Some(matches.free[0].clone()) } else { None };

    Ok(filter)
}

fn get_allow_unstable(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut allow_unstable = false;

    if let Some(opt) = matches.opt_str("Z") {
        if !is_nightly() {
            return Err("the option `Z` is only accepted on the nightly compiler".into());
        }

        match &*opt {
            "unstable-options" => {
                allow_unstable = true;
            }
            _ => {
                return Err("Unrecognized option to `Z`".into());
            }
        }
    };

    Ok(allow_unstable)
}

fn get_log_file(matches: &getopts::Matches) -> OptPartRes<Option<PathBuf>> {
    let logfile = matches.opt_str("logfile").map(|s| PathBuf::from(&s));

    Ok(logfile)
}
<commit_msg>Rollup merge of #70350 - workingjubilee:patch-1, r=Dylan-DPC<commit_after>//! Module converting command-line arguments into test configuration.

use std::env;
use std::path::PathBuf;

use super::helpers::isatty;
use super::options::{ColorConfig, Options, OutputFormat, RunIgnored};
use super::time::TestTimeOptions;

#[derive(Debug)]
pub struct TestOpts {
    pub list: bool,
    pub filter: Option<String>,
    pub filter_exact: bool,
    pub force_run_in_process: bool,
    pub exclude_should_panic: bool,
    pub run_ignored: RunIgnored,
    pub run_tests: bool,
    pub bench_benchmarks: bool,
    pub logfile: Option<PathBuf>,
    pub nocapture: bool,
    pub color: ColorConfig,
    pub format: OutputFormat,
    pub test_threads: Option<usize>,
    pub skip: Vec<String>,
    pub time_options: Option<TestTimeOptions>,
    pub options: Options,
}

impl TestOpts {
    pub fn use_color(&self) -> bool {
        match self.color {
            ColorConfig::AutoColor => !self.nocapture && isatty::stdout_isatty(),
            ColorConfig::AlwaysColor => true,
            ColorConfig::NeverColor => false,
        }
    }
}

/// Result of parsing the options.
pub type OptRes = Result<TestOpts, String>;
/// Result of parsing the option part.
type OptPartRes<T> = Result<T, String>;

fn optgroups() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optflag("", "include-ignored", "Run ignored and not ignored tests")
        .optflag("", "ignored", "Run only ignored tests")
        .optflag("", "force-run-in-process", "Forces tests to run in-process when panic=abort")
        .optflag("", "exclude-should-panic", "Excludes tests marked as should_panic")
        .optflag("", "test", "Run tests and not benchmarks")
        .optflag("", "bench", "Run benchmarks instead of tests")
        .optflag("", "list", "List all tests and benchmarks")
        .optflag("h", "help", "Display this message (longer with --help)")
        .optopt(
            "",
            "logfile",
            "Write logs to the specified file instead \
             of stdout",
            "PATH",
        )
        .optflag(
            "",
            "nocapture",
            "don't capture stdout/stderr of each \
             task, allow printing directly",
        )
        .optopt(
            "",
            "test-threads",
            "Number of threads used for running tests \
             in parallel",
            "n_threads",
        )
        .optmulti(
            "",
            "skip",
            "Skip tests whose names contain FILTER (this flag can \
             be used multiple times)",
            "FILTER",
        )
        .optflag(
            "q",
            "quiet",
            "Display one character per test instead of one line. \
             Alias to --format=terse",
        )
        .optflag("", "exact", "Exactly match filters rather than by substring")
        .optopt(
            "",
            "color",
            "Configure coloring of output:
            auto   = colorize if stdout is a tty and tests are run on serially (default);
            always = always colorize output;
            never  = never colorize output;",
            "auto|always|never",
        )
        .optopt(
            "",
            "format",
            "Configure formatting of output:
            pretty = Print verbose output;
            terse  = Display one character per test;
            json   = Output a json document",
            "pretty|terse|json",
        )
        .optflag("", "show-output", "Show captured stdout of successful tests")
        .optopt(
            "Z",
            "",
            "Enable nightly-only flags:
            unstable-options = Allow use of experimental features",
            "unstable-options",
        )
        .optflagopt(
            "",
            "report-time",
            "Show execution time of each test. Awailable values:
            plain   = do not colorize the execution time (default);
            colored = colorize output according to the `color` parameter value;

            Threshold values for colorized output can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.
            Durations must be specified in milliseconds, e.g. `500,2000` means that the warn time
            is 0.5 seconds, and the critical time is 2 seconds.

            Not available for --format=terse",
            "plain|colored",
        )
        .optflag(
            "",
            "ensure-time",
            "Treat excess of the test execution time limit as error.

            Threshold values for this option can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.

            `CRITICAL_TIME` here means the limit that should not be exceeded by test.
            ",
        );
    opts
}

fn usage(binary: &str, options: &getopts::Options) {
    let message = format!("Usage: {} [OPTIONS] [FILTER]", binary);
    println!(
        r#"{usage}

The FILTER string is tested against the name of all tests, and only those
tests whose names contain the filter are run.

By default, all tests are run in parallel. This can be altered with the
--test-threads flag or the RUST_TEST_THREADS environment variable when running
tests (set it to 1).

All tests have their standard output and standard error captured by default.
This can be overridden with the --nocapture flag or setting RUST_TEST_NOCAPTURE
environment variable to a value other than "0". Logging is not captured by default.

Test Attributes:

    `#[test]`        - Indicates a function is a test to be run. This function
                       takes no arguments.
    `#[bench]`       - Indicates a function is a benchmark to be run. This
                       function takes one argument (test::Bencher).
    `#[should_panic]` - This function (also labeled with `#[test]`) will only pass if
                        the code causes a panic (an assertion failure or panic!)
                        A message may be provided, which the failure string must
                        contain: #[should_panic(expected = "foo")].
    `#[ignore]`       - When applied to a function which is already attributed as a
                        test, then the test runner will ignore these tests during
                        normal test runs. Running with --ignored or --include-ignored will run
                        these tests."#,
        usage = options.usage(&message)
    );
}

/// Parses command line arguments into test options.
/// Returns `None` if help was requested (since we only show help message and don't run tests),
/// returns `Some(Err(..))` if provided arguments are incorrect,
/// otherwise creates a `TestOpts` object and returns it.
pub fn parse_opts(args: &[String]) -> Option<OptRes> {
    // Parse matches.
    let opts = optgroups();
    let args = args.get(1..).unwrap_or(args);
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(f) => return Some(Err(f.to_string())),
    };

    // Check if help was requested.
    if matches.opt_present("h") {
        // Show help and do nothing more.
        usage(&args[0], &opts);
        return None;
    }

    // Actually parse the opts.
    let opts_result = parse_opts_impl(matches);

    Some(opts_result)
}

// Gets the option value and checks if unstable features are enabled.
macro_rules! unstable_optflag {
    ($matches:ident, $allow_unstable:ident, $option_name:literal) => {{
        let opt = $matches.opt_present($option_name);
        if !$allow_unstable && opt {
            return Err(format!(
                "The \"{}\" flag is only accepted on the nightly compiler with -Z unstable-options",
                $option_name
            ));
        }

        opt
    }};
}

// Implementation of `parse_opts` that doesn't care about help message
// and returns a `Result`.
fn parse_opts_impl(matches: getopts::Matches) -> OptRes {
    let allow_unstable = get_allow_unstable(&matches)?;

    // Unstable flags
    let force_run_in_process = unstable_optflag!(matches, allow_unstable, "force-run-in-process");
    let exclude_should_panic = unstable_optflag!(matches, allow_unstable, "exclude-should-panic");
    let include_ignored = unstable_optflag!(matches, allow_unstable, "include-ignored");
    let time_options = get_time_options(&matches, allow_unstable)?;

    let quiet = matches.opt_present("quiet");
    let exact = matches.opt_present("exact");
    let list = matches.opt_present("list");
    let skip = matches.opt_strs("skip");

    let bench_benchmarks = matches.opt_present("bench");
    let run_tests = !bench_benchmarks || matches.opt_present("test");

    let logfile = get_log_file(&matches)?;
    let run_ignored = get_run_ignored(&matches, include_ignored)?;
    let filter = get_filter(&matches)?;
    let nocapture = get_nocapture(&matches)?;
    let test_threads = get_test_threads(&matches)?;
    let color = get_color_config(&matches)?;
    let format = get_format(&matches, quiet, allow_unstable)?;

    let options = Options::new().display_output(matches.opt_present("show-output"));

    let test_opts = TestOpts {
        list,
        filter,
        filter_exact: exact,
        force_run_in_process,
        exclude_should_panic,
        run_ignored,
        run_tests,
        bench_benchmarks,
        logfile,
        nocapture,
        color,
        format,
        test_threads,
        skip,
        time_options,
        options,
    };

    Ok(test_opts)
}

// FIXME: Copied from librustc_ast until linkage errors are resolved. Issue #47566
fn is_nightly() -> bool {
    // Whether this is a feature-staged build, i.e., on the beta or stable channel
    let disable_unstable_features = option_env!("CFG_DISABLE_UNSTABLE_FEATURES").is_some();
    // Whether we should enable unstable features for bootstrapping
    let bootstrap = env::var("RUSTC_BOOTSTRAP").is_ok();

    bootstrap || !disable_unstable_features
}

// Gets the CLI options associated with `report-time` feature.
fn get_time_options(
    matches: &getopts::Matches,
    allow_unstable: bool,
) -> OptPartRes<Option<TestTimeOptions>> {
    let report_time = unstable_optflag!(matches, allow_unstable, "report-time");
    let colored_opt_str = matches.opt_str("report-time");
    let mut report_time_colored = report_time && colored_opt_str == Some("colored".into());
    let ensure_test_time = unstable_optflag!(matches, allow_unstable, "ensure-time");

    // If `ensure-test-time` option is provided, time output is enforced,
    // so user won't be confused if any of tests will silently fail.
    let options = if report_time || ensure_test_time {
        if ensure_test_time && !report_time {
            report_time_colored = true;
        }
        Some(TestTimeOptions::new_from_env(ensure_test_time, report_time_colored))
    } else {
        None
    };

    Ok(options)
}

fn get_test_threads(matches: &getopts::Matches) -> OptPartRes<Option<usize>> {
    let test_threads = match matches.opt_str("test-threads") {
        Some(n_str) => match n_str.parse::<usize>() {
            Ok(0) => return Err("argument for --test-threads must not be 0".to_string()),
            Ok(n) => Some(n),
            Err(e) => {
                return Err(format!(
                    "argument for --test-threads must be a number > 0 \
                     (error: {})",
                    e
                ));
            }
        },
        None => None,
    };

    Ok(test_threads)
}

fn get_format(
    matches: &getopts::Matches,
    quiet: bool,
    allow_unstable: bool,
) -> OptPartRes<OutputFormat> {
    let format = match matches.opt_str("format").as_ref().map(|s| &**s) {
        None if quiet => OutputFormat::Terse,
        Some("pretty") | None => OutputFormat::Pretty,
        Some("terse") => OutputFormat::Terse,
        Some("json") => {
            if !allow_unstable {
                return Err("The \"json\" format is only accepted on the nightly compiler".into());
            }
            OutputFormat::Json
        }

        Some(v) => {
            return Err(format!(
                "argument for --format must be pretty, terse, or json (was \
                 {})",
                v
            ));
        }
    };

    Ok(format)
}

fn get_color_config(matches: &getopts::Matches) -> OptPartRes<ColorConfig> {
    let color = match matches.opt_str("color").as_ref().map(|s| &**s) {
        Some("auto") | None => ColorConfig::AutoColor,
        Some("always") => ColorConfig::AlwaysColor,
        Some("never") => ColorConfig::NeverColor,

        Some(v) => {
            return Err(format!(
                "argument for --color must be auto, always, or never (was \
                 {})",
                v
            ));
        }
    };

    Ok(color)
}

fn get_nocapture(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut nocapture = matches.opt_present("nocapture");
    if !nocapture {
        nocapture = match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        };
    }

    Ok(nocapture)
}

fn get_run_ignored(matches: &getopts::Matches, include_ignored: bool) -> OptPartRes<RunIgnored> {
    let run_ignored = match (include_ignored, matches.opt_present("ignored")) {
        (true, true) => {
            return Err("the options --include-ignored and --ignored are mutually exclusive".into());
        }
        (true, false) => RunIgnored::Yes,
        (false, true) => RunIgnored::Only,
        (false, false) => RunIgnored::No,
    };

    Ok(run_ignored)
}

fn get_filter(matches: &getopts::Matches) -> OptPartRes<Option<String>> {
    let filter = if !matches.free.is_empty() { Some(matches.free[0].clone()) } else { None };

    Ok(filter)
}

fn get_allow_unstable(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut allow_unstable = false;

    if let Some(opt) = matches.opt_str("Z") {
        if !is_nightly() {
            return Err("the option `Z` is only accepted on the nightly compiler".into());
        }

        match &*opt {
            "unstable-options" => {
                allow_unstable = true;
            }
            _ => {
                return Err("Unrecognized option to `Z`".into());
            }
        }
    };

    Ok(allow_unstable)
}

fn get_log_file(matches: &getopts::Matches) -> OptPartRes<Option<PathBuf>> {
    let logfile = matches.opt_str("logfile").map(|s| PathBuf::from(&s));

    Ok(logfile)
}
<|endoftext|>


<commit_before>//! Module converting command-line arguments into test configuration.

use std::env;
use std::path::PathBuf;

use super::helpers::isatty;
use super::options::{ColorConfig, Options, OutputFormat, RunIgnored};
use super::time::TestTimeOptions;

#[derive(Debug)]
pub struct TestOpts {
    pub list: bool,
    pub filter: Option<String>,
    pub filter_exact: bool,
    pub force_run_in_process: bool,
    pub exclude_should_panic: bool,
    pub run_ignored: RunIgnored,
    pub run_tests: bool,
    pub bench_benchmarks: bool,
    pub logfile: Option<PathBuf>,
    pub nocapture: bool,
    pub color: ColorConfig,
    pub format: OutputFormat,
    pub test_threads: Option<usize>,
    pub skip: Vec<String>,
    pub time_options: Option<TestTimeOptions>,
    pub options: Options,
}

impl TestOpts {
    pub fn use_color(&self) -> bool {
        match self.color {
            ColorConfig::AutoColor => !self.nocapture && isatty::stdout_isatty(),
            ColorConfig::AlwaysColor => true,
            ColorConfig::NeverColor => false,
        }
    }
}

/// Result of parsing the options.
pub type OptRes = Result<TestOpts, String>;
/// Result of parsing the option part.
type OptPartRes<T> = Result<T, String>;

fn optgroups() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optflag("", "include-ignored", "Run ignored and not ignored tests")
        .optflag("", "ignored", "Run only ignored tests")
        .optflag("", "force-run-in-process", "Forces tests to run in-process when panic=abort")
        .optflag("", "exclude-should-panic", "Excludes tests marked as should_panic")
        .optflag("", "test", "Run tests and not benchmarks")
        .optflag("", "bench", "Run benchmarks instead of tests")
        .optflag("", "list", "List all tests and benchmarks")
        .optflag("h", "help", "Display this message (longer with --help)")
        .optopt(
            "",
            "logfile",
            "Write logs to the specified file instead \
             of stdout",
            "PATH",
        )
        .optflag(
            "",
            "nocapture",
            "don't capture stdout/stderr of each \
             task, allow printing directly",
        )
        .optopt(
            "",
            "test-threads",
            "Number of threads used for running tests \
             in parallel",
            "n_threads",
        )
        .optmulti(
            "",
            "skip",
            "Skip tests whose names contain FILTER (this flag can \
             be used multiple times)",
            "FILTER",
        )
        .optflag(
            "q",
            "quiet",
            "Display one character per test instead of one line. \
             Alias to --format=terse",
        )
        .optflag("", "exact", "Exactly match filters rather than by substring")
        .optopt(
            "",
            "color",
            "Configure coloring of output:
            auto   = colorize if stdout is a tty and tests are run on serially (default);
            always = always colorize output;
            never  = never colorize output;",
            "auto|always|never",
        )
        .optopt(
            "",
            "format",
            "Configure formatting of output:
            pretty = Print verbose output;
            terse  = Display one character per test;
            json   = Output a json document",
            "pretty|terse|json",
        )
        .optflag("", "show-output", "Show captured stdout of successful tests")
        .optopt(
            "Z",
            "",
            "Enable nightly-only flags:
            unstable-options = Allow use of experimental features",
            "unstable-options",
        )
        .optflagopt(
            "",
            "report-time",
            "Show execution time of each test. Awailable values:
            plain   = do not colorize the execution time (default);
            colored = colorize output according to the `color` parameter value;

            Threshold values for colorized output can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.
            Durations must be specified in milliseconds, e.g. `500,2000` means that the warn time
            is 0.5 seconds, and the critical time is 2 seconds.

            Not available for --format=terse",
            "plain|colored",
        )
        .optflag(
            "",
            "ensure-time",
            "Treat excess of the test execution time limit as error.

            Threshold values for this option can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.

            `CRITICAL_TIME` here means the limit that should not be exceeded by test.
            ",
        );
    opts
}

fn usage(binary: &str, options: &getopts::Options) {
    let message = format!("Usage: {} [OPTIONS] [FILTER]", binary);
    println!(
        r#"{usage}

The FILTER string is tested against the name of all tests, and only those
tests whose names contain the filter are run.

By default, all tests are run in parallel. This can be altered with the
--test-threads flag or the RUST_TEST_THREADS environment variable when running
tests (set it to 1).

All tests have their standard output and standard error captured by default.
This can be overridden with the --nocapture flag or setting RUST_TEST_NOCAPTURE
environment variable to a value other than "0". Logging is not captured by default.

Test Attributes:

    `#[test]`        - Indicates a function is a test to be run. This function
                       takes no arguments.
    `#[bench]`       - Indicates a function is a benchmark to be run. This
                       function takes one argument (test::Bencher).
    `#[should_panic]` - This function (also labeled with `#[test]`) will only pass if
                        the code causes a panic (an assertion failure or panic!)
                        A message may be provided, which the failure string must
                        contain: #[should_panic(expected = "foo")].
    `#[ignore]`       - When applied to a function which is already attributed as a
                        test, then the test runner will ignore these tests during
                        normal test runs. Running with --ignored or --include-ignored will run
                        these tests."#,
        usage = options.usage(&message)
    );
}

/// Parses command line arguments into test options.
/// Returns `None` if help was requested (since we only show help message and don't run tests),
/// returns `Some(Err(..))` if provided arguments are incorrect,
/// otherwise creates a `TestOpts` object and returns it.
pub fn parse_opts(args: &[String]) -> Option<OptRes> {
    // Parse matches.
    let opts = optgroups();
    let args = args.get(1..).unwrap_or(args);
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(f) => return Some(Err(f.to_string())),
    };

    // Check if help was requested.
    if matches.opt_present("h") {
        // Show help and do nothing more.
        usage(&args[0], &opts);
        return None;
    }

    // Actually parse the opts.
    let opts_result = parse_opts_impl(matches);

    Some(opts_result)
}

// Gets the option value and checks if unstable features are enabled.
macro_rules! unstable_optflag {
    ($matches:ident, $allow_unstable:ident, $option_name:literal) => {{
        let opt = $matches.opt_present($option_name);
        if !$allow_unstable && opt {
            return Err(format!(
                "The \"{}\" flag is only accepted on the nightly compiler",
                $option_name
            ));
        }

        opt
    }};
}

// Implementation of `parse_opts` that doesn't care about help message
// and returns a `Result`.
fn parse_opts_impl(matches: getopts::Matches) -> OptRes {
    let allow_unstable = get_allow_unstable(&matches)?;

    // Unstable flags
    let force_run_in_process = unstable_optflag!(matches, allow_unstable, "force-run-in-process");
    let exclude_should_panic = unstable_optflag!(matches, allow_unstable, "exclude-should-panic");
    let include_ignored = unstable_optflag!(matches, allow_unstable, "include-ignored");
    let time_options = get_time_options(&matches, allow_unstable)?;

    let quiet = matches.opt_present("quiet");
    let exact = matches.opt_present("exact");
    let list = matches.opt_present("list");
    let skip = matches.opt_strs("skip");

    let bench_benchmarks = matches.opt_present("bench");
    let run_tests = !bench_benchmarks || matches.opt_present("test");

    let logfile = get_log_file(&matches)?;
    let run_ignored = get_run_ignored(&matches, include_ignored)?;
    let filter = get_filter(&matches)?;
    let nocapture = get_nocapture(&matches)?;
    let test_threads = get_test_threads(&matches)?;
    let color = get_color_config(&matches)?;
    let format = get_format(&matches, quiet, allow_unstable)?;

    let options = Options::new().display_output(matches.opt_present("show-output"));

    let test_opts = TestOpts {
        list,
        filter,
        filter_exact: exact,
        force_run_in_process,
        exclude_should_panic,
        run_ignored,
        run_tests,
        bench_benchmarks,
        logfile,
        nocapture,
        color,
        format,
        test_threads,
        skip,
        time_options,
        options,
    };

    Ok(test_opts)
}

// FIXME: Copied from librustc_ast until linkage errors are resolved. Issue #47566
fn is_nightly() -> bool {
    // Whether this is a feature-staged build, i.e., on the beta or stable channel
    let disable_unstable_features = option_env!("CFG_DISABLE_UNSTABLE_FEATURES").is_some();
    // Whether we should enable unstable features for bootstrapping
    let bootstrap = env::var("RUSTC_BOOTSTRAP").is_ok();

    bootstrap || !disable_unstable_features
}

// Gets the CLI options associated with `report-time` feature.
fn get_time_options(
    matches: &getopts::Matches,
    allow_unstable: bool,
) -> OptPartRes<Option<TestTimeOptions>> {
    let report_time = unstable_optflag!(matches, allow_unstable, "report-time");
    let colored_opt_str = matches.opt_str("report-time");
    let mut report_time_colored = report_time && colored_opt_str == Some("colored".into());
    let ensure_test_time = unstable_optflag!(matches, allow_unstable, "ensure-time");

    // If `ensure-test-time` option is provided, time output is enforced,
    // so user won't be confused if any of tests will silently fail.
    let options = if report_time || ensure_test_time {
        if ensure_test_time && !report_time {
            report_time_colored = true;
        }
        Some(TestTimeOptions::new_from_env(ensure_test_time, report_time_colored))
    } else {
        None
    };

    Ok(options)
}

fn get_test_threads(matches: &getopts::Matches) -> OptPartRes<Option<usize>> {
    let test_threads = match matches.opt_str("test-threads") {
        Some(n_str) => match n_str.parse::<usize>() {
            Ok(0) => return Err("argument for --test-threads must not be 0".to_string()),
            Ok(n) => Some(n),
            Err(e) => {
                return Err(format!(
                    "argument for --test-threads must be a number > 0 \
                     (error: {})",
                    e
                ));
            }
        },
        None => None,
    };

    Ok(test_threads)
}

fn get_format(
    matches: &getopts::Matches,
    quiet: bool,
    allow_unstable: bool,
) -> OptPartRes<OutputFormat> {
    let format = match matches.opt_str("format").as_ref().map(|s| &**s) {
        None if quiet => OutputFormat::Terse,
        Some("pretty") | None => OutputFormat::Pretty,
        Some("terse") => OutputFormat::Terse,
        Some("json") => {
            if !allow_unstable {
                return Err("The \"json\" format is only accepted on the nightly compiler".into());
            }
            OutputFormat::Json
        }

        Some(v) => {
            return Err(format!(
                "argument for --format must be pretty, terse, or json (was \
                 {})",
                v
            ));
        }
    };

    Ok(format)
}

fn get_color_config(matches: &getopts::Matches) -> OptPartRes<ColorConfig> {
    let color = match matches.opt_str("color").as_ref().map(|s| &**s) {
        Some("auto") | None => ColorConfig::AutoColor,
        Some("always") => ColorConfig::AlwaysColor,
        Some("never") => ColorConfig::NeverColor,

        Some(v) => {
            return Err(format!(
                "argument for --color must be auto, always, or never (was \
                 {})",
                v
            ));
        }
    };

    Ok(color)
}

fn get_nocapture(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut nocapture = matches.opt_present("nocapture");
    if !nocapture {
        nocapture = match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        };
    }

    Ok(nocapture)
}

fn get_run_ignored(matches: &getopts::Matches, include_ignored: bool) -> OptPartRes<RunIgnored> {
    let run_ignored = match (include_ignored, matches.opt_present("ignored")) {
        (true, true) => {
            return Err("the options --include-ignored and --ignored are mutually exclusive".into());
        }
        (true, false) => RunIgnored::Yes,
        (false, true) => RunIgnored::Only,
        (false, false) => RunIgnored::No,
    };

    Ok(run_ignored)
}

fn get_filter(matches: &getopts::Matches) -> OptPartRes<Option<String>> {
    let filter = if !matches.free.is_empty() { Some(matches.free[0].clone()) } else { None };

    Ok(filter)
}

fn get_allow_unstable(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut allow_unstable = false;

    if let Some(opt) = matches.opt_str("Z") {
        if !is_nightly() {
            return Err("the option `Z` is only accepted on the nightly compiler".into());
        }

        match &*opt {
            "unstable-options" => {
                allow_unstable = true;
            }
            _ => {
                return Err("Unrecognized option to `Z`".into());
            }
        }
    };

    Ok(allow_unstable)
}

fn get_log_file(matches: &getopts::Matches) -> OptPartRes<Option<PathBuf>> {
    let logfile = matches.opt_str("logfile").map(|s| PathBuf::from(&s));

    Ok(logfile)
}
<commit_msg>Request "-Z unstable-options" for unstable options<commit_after>//! Module converting command-line arguments into test configuration.

use std::env;
use std::path::PathBuf;

use super::helpers::isatty;
use super::options::{ColorConfig, Options, OutputFormat, RunIgnored};
use super::time::TestTimeOptions;

#[derive(Debug)]
pub struct TestOpts {
    pub list: bool,
    pub filter: Option<String>,
    pub filter_exact: bool,
    pub force_run_in_process: bool,
    pub exclude_should_panic: bool,
    pub run_ignored: RunIgnored,
    pub run_tests: bool,
    pub bench_benchmarks: bool,
    pub logfile: Option<PathBuf>,
    pub nocapture: bool,
    pub color: ColorConfig,
    pub format: OutputFormat,
    pub test_threads: Option<usize>,
    pub skip: Vec<String>,
    pub time_options: Option<TestTimeOptions>,
    pub options: Options,
}

impl TestOpts {
    pub fn use_color(&self) -> bool {
        match self.color {
            ColorConfig::AutoColor => !self.nocapture && isatty::stdout_isatty(),
            ColorConfig::AlwaysColor => true,
            ColorConfig::NeverColor => false,
        }
    }
}

/// Result of parsing the options.
pub type OptRes = Result<TestOpts, String>;
/// Result of parsing the option part.
type OptPartRes<T> = Result<T, String>;

fn optgroups() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optflag("", "include-ignored", "Run ignored and not ignored tests")
        .optflag("", "ignored", "Run only ignored tests")
        .optflag("", "force-run-in-process", "Forces tests to run in-process when panic=abort")
        .optflag("", "exclude-should-panic", "Excludes tests marked as should_panic")
        .optflag("", "test", "Run tests and not benchmarks")
        .optflag("", "bench", "Run benchmarks instead of tests")
        .optflag("", "list", "List all tests and benchmarks")
        .optflag("h", "help", "Display this message (longer with --help)")
        .optopt(
            "",
            "logfile",
            "Write logs to the specified file instead \
             of stdout",
            "PATH",
        )
        .optflag(
            "",
            "nocapture",
            "don't capture stdout/stderr of each \
             task, allow printing directly",
        )
        .optopt(
            "",
            "test-threads",
            "Number of threads used for running tests \
             in parallel",
            "n_threads",
        )
        .optmulti(
            "",
            "skip",
            "Skip tests whose names contain FILTER (this flag can \
             be used multiple times)",
            "FILTER",
        )
        .optflag(
            "q",
            "quiet",
            "Display one character per test instead of one line. \
             Alias to --format=terse",
        )
        .optflag("", "exact", "Exactly match filters rather than by substring")
        .optopt(
            "",
            "color",
            "Configure coloring of output:
            auto   = colorize if stdout is a tty and tests are run on serially (default);
            always = always colorize output;
            never  = never colorize output;",
            "auto|always|never",
        )
        .optopt(
            "",
            "format",
            "Configure formatting of output:
            pretty = Print verbose output;
            terse  = Display one character per test;
            json   = Output a json document",
            "pretty|terse|json",
        )
        .optflag("", "show-output", "Show captured stdout of successful tests")
        .optopt(
            "Z",
            "",
            "Enable nightly-only flags:
            unstable-options = Allow use of experimental features",
            "unstable-options",
        )
        .optflagopt(
            "",
            "report-time",
            "Show execution time of each test. Awailable values:
            plain   = do not colorize the execution time (default);
            colored = colorize output according to the `color` parameter value;

            Threshold values for colorized output can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.
            Durations must be specified in milliseconds, e.g. `500,2000` means that the warn time
            is 0.5 seconds, and the critical time is 2 seconds.

            Not available for --format=terse",
            "plain|colored",
        )
        .optflag(
            "",
            "ensure-time",
            "Treat excess of the test execution time limit as error.

            Threshold values for this option can be configured via
            `RUST_TEST_TIME_UNIT`, `RUST_TEST_TIME_INTEGRATION` and
            `RUST_TEST_TIME_DOCTEST` environment variables.

            Expected format of environment variable is `VARIABLE=WARN_TIME,CRITICAL_TIME`.

            `CRITICAL_TIME` here means the limit that should not be exceeded by test.
            ",
        );
    opts
}

fn usage(binary: &str, options: &getopts::Options) {
    let message = format!("Usage: {} [OPTIONS] [FILTER]", binary);
    println!(
        r#"{usage}

The FILTER string is tested against the name of all tests, and only those
tests whose names contain the filter are run.

By default, all tests are run in parallel. This can be altered with the
--test-threads flag or the RUST_TEST_THREADS environment variable when running
tests (set it to 1).

All tests have their standard output and standard error captured by default.
This can be overridden with the --nocapture flag or setting RUST_TEST_NOCAPTURE
environment variable to a value other than "0". Logging is not captured by default.

Test Attributes:

    `#[test]`        - Indicates a function is a test to be run. This function
                       takes no arguments.
    `#[bench]`       - Indicates a function is a benchmark to be run. This
                       function takes one argument (test::Bencher).
    `#[should_panic]` - This function (also labeled with `#[test]`) will only pass if
                        the code causes a panic (an assertion failure or panic!)
                        A message may be provided, which the failure string must
                        contain: #[should_panic(expected = "foo")].
    `#[ignore]`       - When applied to a function which is already attributed as a
                        test, then the test runner will ignore these tests during
                        normal test runs. Running with --ignored or --include-ignored will run
                        these tests."#,
        usage = options.usage(&message)
    );
}

/// Parses command line arguments into test options.
/// Returns `None` if help was requested (since we only show help message and don't run tests),
/// returns `Some(Err(..))` if provided arguments are incorrect,
/// otherwise creates a `TestOpts` object and returns it.
pub fn parse_opts(args: &[String]) -> Option<OptRes> {
    // Parse matches.
    let opts = optgroups();
    let args = args.get(1..).unwrap_or(args);
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(f) => return Some(Err(f.to_string())),
    };

    // Check if help was requested.
    if matches.opt_present("h") {
        // Show help and do nothing more.
        usage(&args[0], &opts);
        return None;
    }

    // Actually parse the opts.
    let opts_result = parse_opts_impl(matches);

    Some(opts_result)
}

// Gets the option value and checks if unstable features are enabled.
macro_rules! unstable_optflag {
    ($matches:ident, $allow_unstable:ident, $option_name:literal) => {{
        let opt = $matches.opt_present($option_name);
        if !$allow_unstable && opt {
            return Err(format!(
                "The \"{}\" flag is only accepted on the nightly compiler with -Z unstable-options",
                $option_name
            ));
        }

        opt
    }};
}

// Implementation of `parse_opts` that doesn't care about help message
// and returns a `Result`.
fn parse_opts_impl(matches: getopts::Matches) -> OptRes {
    let allow_unstable = get_allow_unstable(&matches)?;

    // Unstable flags
    let force_run_in_process = unstable_optflag!(matches, allow_unstable, "force-run-in-process");
    let exclude_should_panic = unstable_optflag!(matches, allow_unstable, "exclude-should-panic");
    let include_ignored = unstable_optflag!(matches, allow_unstable, "include-ignored");
    let time_options = get_time_options(&matches, allow_unstable)?;

    let quiet = matches.opt_present("quiet");
    let exact = matches.opt_present("exact");
    let list = matches.opt_present("list");
    let skip = matches.opt_strs("skip");

    let bench_benchmarks = matches.opt_present("bench");
    let run_tests = !bench_benchmarks || matches.opt_present("test");

    let logfile = get_log_file(&matches)?;
    let run_ignored = get_run_ignored(&matches, include_ignored)?;
    let filter = get_filter(&matches)?;
    let nocapture = get_nocapture(&matches)?;
    let test_threads = get_test_threads(&matches)?;
    let color = get_color_config(&matches)?;
    let format = get_format(&matches, quiet, allow_unstable)?;

    let options = Options::new().display_output(matches.opt_present("show-output"));

    let test_opts = TestOpts {
        list,
        filter,
        filter_exact: exact,
        force_run_in_process,
        exclude_should_panic,
        run_ignored,
        run_tests,
        bench_benchmarks,
        logfile,
        nocapture,
        color,
        format,
        test_threads,
        skip,
        time_options,
        options,
    };

    Ok(test_opts)
}

// FIXME: Copied from librustc_ast until linkage errors are resolved. Issue #47566
fn is_nightly() -> bool {
    // Whether this is a feature-staged build, i.e., on the beta or stable channel
    let disable_unstable_features = option_env!("CFG_DISABLE_UNSTABLE_FEATURES").is_some();
    // Whether we should enable unstable features for bootstrapping
    let bootstrap = env::var("RUSTC_BOOTSTRAP").is_ok();

    bootstrap || !disable_unstable_features
}

// Gets the CLI options associated with `report-time` feature.
fn get_time_options(
    matches: &getopts::Matches,
    allow_unstable: bool,
) -> OptPartRes<Option<TestTimeOptions>> {
    let report_time = unstable_optflag!(matches, allow_unstable, "report-time");
    let colored_opt_str = matches.opt_str("report-time");
    let mut report_time_colored = report_time && colored_opt_str == Some("colored".into());
    let ensure_test_time = unstable_optflag!(matches, allow_unstable, "ensure-time");

    // If `ensure-test-time` option is provided, time output is enforced,
    // so user won't be confused if any of tests will silently fail.
    let options = if report_time || ensure_test_time {
        if ensure_test_time && !report_time {
            report_time_colored = true;
        }
        Some(TestTimeOptions::new_from_env(ensure_test_time, report_time_colored))
    } else {
        None
    };

    Ok(options)
}

fn get_test_threads(matches: &getopts::Matches) -> OptPartRes<Option<usize>> {
    let test_threads = match matches.opt_str("test-threads") {
        Some(n_str) => match n_str.parse::<usize>() {
            Ok(0) => return Err("argument for --test-threads must not be 0".to_string()),
            Ok(n) => Some(n),
            Err(e) => {
                return Err(format!(
                    "argument for --test-threads must be a number > 0 \
                     (error: {})",
                    e
                ));
            }
        },
        None => None,
    };

    Ok(test_threads)
}

fn get_format(
    matches: &getopts::Matches,
    quiet: bool,
    allow_unstable: bool,
) -> OptPartRes<OutputFormat> {
    let format = match matches.opt_str("format").as_ref().map(|s| &**s) {
        None if quiet => OutputFormat::Terse,
        Some("pretty") | None => OutputFormat::Pretty,
        Some("terse") => OutputFormat::Terse,
        Some("json") => {
            if !allow_unstable {
                return Err("The \"json\" format is only accepted on the nightly compiler".into());
            }
            OutputFormat::Json
        }

        Some(v) => {
            return Err(format!(
                "argument for --format must be pretty, terse, or json (was \
                 {})",
                v
            ));
        }
    };

    Ok(format)
}

fn get_color_config(matches: &getopts::Matches) -> OptPartRes<ColorConfig> {
    let color = match matches.opt_str("color").as_ref().map(|s| &**s) {
        Some("auto") | None => ColorConfig::AutoColor,
        Some("always") => ColorConfig::AlwaysColor,
        Some("never") => ColorConfig::NeverColor,

        Some(v) => {
            return Err(format!(
                "argument for --color must be auto, always, or never (was \
                 {})",
                v
            ));
        }
    };

    Ok(color)
}

fn get_nocapture(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut nocapture = matches.opt_present("nocapture");
    if !nocapture {
        nocapture = match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        };
    }

    Ok(nocapture)
}

fn get_run_ignored(matches: &getopts::Matches, include_ignored: bool) -> OptPartRes<RunIgnored> {
    let run_ignored = match (include_ignored, matches.opt_present("ignored")) {
        (true, true) => {
            return Err("the options --include-ignored and --ignored are mutually exclusive".into());
        }
        (true, false) => RunIgnored::Yes,
        (false, true) => RunIgnored::Only,
        (false, false) => RunIgnored::No,
    };

    Ok(run_ignored)
}

fn get_filter(matches: &getopts::Matches) -> OptPartRes<Option<String>> {
    let filter = if !matches.free.is_empty() { Some(matches.free[0].clone()) } else { None };

    Ok(filter)
}

fn get_allow_unstable(matches: &getopts::Matches) -> OptPartRes<bool> {
    let mut allow_unstable = false;

    if let Some(opt) = matches.opt_str("Z") {
        if !is_nightly() {
            return Err("the option `Z` is only accepted on the nightly compiler".into());
        }

        match &*opt {
            "unstable-options" => {
                allow_unstable = true;
            }
            _ => {
                return Err("Unrecognized option to `Z`".into());
            }
        }
    };

    Ok(allow_unstable)
}

fn get_log_file(matches: &getopts::Matches) -> OptPartRes<Option<PathBuf>> {
    let logfile = matches.opt_str("logfile").map(|s| PathBuf::from(&s));

    Ok(logfile)
}
<|endoftext|>


<commit_before>#![allow(unused_variables)]

use std::thread::Thread;

use super::piece::King;
use super::moves::Move;
use super::pos::Position;
use super::gen;
use super::make_move::make_move_mut;

pub fn receive_legal(p: Position) -> Receiver<Move> {
    let (tx,rx) = sync_channel(0);
    Thread::spawn(move || gen_legal(&p, tx)).detach();
    rx
}

pub fn gen_legal(p: &Position, out: SyncSender<Move>) {
    let (tx, rx) = sync_channel::<Move>(1);
    let temp = (*p).clone();
    Thread::spawn(move || gen::gen_psudo_legal(&temp, tx)).detach();
    filter_legal(p, out, rx);
}

fn filter_legal(p: &Position, out: SyncSender<Move>, rx: Receiver<Move>) {
    for curr_move in rx.iter() {
        if is_legal(p.clone(), &curr_move) {
            if out.send_opt(curr_move).is_err() {
                return;
            }
        }
    }
}

pub fn is_legal(mut p: Position, curr_move: &Move) -> bool {
    //TODO Implement is_legal() for castling.
    make_move_mut(&mut p, curr_move);
    !can_take_king(&p)
}

pub fn can_take_king(p: &Position) -> bool {
    for m in gen::receive_psudo_legal(p.clone()).iter() {
        let dest = p.at(m.to());
        if dest.is_some() && dest.unwrap().piece_type() == King {
            return true;
        }
    }
    false
}
<commit_msg>Implement legality check for castling.<commit_after>#![allow(unused_variables)]

use std::thread::Thread;

use super::piece::King;
use super::square::Square;
use super::castle;
use super::moves::Move;
use super::pos::Position;
use super::gen;
use super::make_move::make_move_mut;

pub fn receive_legal(p: Position) -> Receiver<Move> {
    let (tx,rx) = sync_channel(0);
    Thread::spawn(move || gen_legal(&p, tx)).detach();
    rx
}

pub fn gen_legal(p: &Position, out: SyncSender<Move>) {
    let (tx, rx) = sync_channel::<Move>(1);
    let temp = (*p).clone();
    Thread::spawn(move || gen::gen_psudo_legal(&temp, tx)).detach();
    filter_legal(p, out, rx);
}

fn filter_legal(p: &Position, out: SyncSender<Move>, rx: Receiver<Move>) {
    for curr_move in rx.iter() {
        if is_legal(p.clone(), &curr_move) {
            if out.send_opt(curr_move).is_err() {
                return;
            }
        }
    }
}

pub fn is_legal(mut p: Position, curr_move: &Move) -> bool {
    let c = p.side_to_move();
    make_move_mut(&mut p, curr_move);
    match curr_move.castle() {
        None => {
            !can_take_king(&p)
        },
        Some(side) => {
            // Check for castling out of check, through check, and into check.
            let check_squares: Vec<Square> = castle::require_no_attack(side, c);
            gen::receive_psudo_legal(p).iter().all(
                |m| {
                    check_squares.iter().all( |x| *x != m.to() )
                }
            )
        }
    }
}

pub fn can_take_king(p: &Position) -> bool {
    for m in gen::receive_psudo_legal(p.clone()).iter() {
        let dest = p.at(m.to());
        if dest.is_some() && dest.unwrap().piece_type() == King {
            return true;
        }
    }
    false
}
<|endoftext|>


<commit_before>//! Hyphenating iterators.

use std::borrow::Cow;

use unicode_segmentation::UnicodeSegmentation;

use language::{Corpus};
use utilia::{Interspersable, Intersperse};

pub trait Hyphenation<Hyphenator> where Hyphenator : Iterator {
    /// Returns the indices of valid hyphenation points within the text.
    fn opportunities(self, corp: &Corpus) -> Vec<usize>;

    /// Returns an iterator over segments of the text separated by valid
    /// hyphenation points.
    ///
    /// When iterating over a word, such segments coincide with orthographic
    /// syllables. Note that, in some orthographies, the syllables of a hyphenated
    /// word are not necessarily substrings of the original word.
    fn hyphenate(self, corp: &Corpus) -> Hyphenator;
}

/// The `Standard` hyphenator iterates over a string, returning slices
/// delimited by string boundaries and valid hyphenation points.
///
/// For individual words, such slices coincide with orthographic syllables.
#[derive(Clone, Debug)]
pub struct Standard<'a> {
    text: &'a str,
    opportunities: Vec<usize>,
    prior: usize,
    i: usize
}

impl<'a> Standard<'a> {
    /// Inserts a soft hyphen at hyphenation points.
    pub fn punctuate(self) -> Intersperse<Self> {
        self.intersperse("\u{ad}")
    }
}

impl<'a> Iterator for Standard<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        let start = self.prior;
        let i = self.i;

        match self.opportunities.get(i) {
            Some(&end) => {
                self.prior = end;
                self.i = i + 1;
                Some(&self.text[start .. end])
            },
            None => {
                if i <= self.opportunities.len() {
                    self.i = i + 1;
                    Some(&self.text[start ..])
                } else {
                    None
                }
            }
        }
    }
}

impl<'a> Hyphenation<Standard<'a>> for &'a str {
    /// Returns the byte indices of valid hyphenation points within the string.
    fn opportunities(self, corp: &Corpus) -> Vec<usize> {
        let (l_min, r_min) = (corp.left_min, corp.right_min);
        let length_min = l_min + r_min;

        if self.chars().count() < length_min {
            return vec![];
        }

        let by_word = self.split_word_bound_indices();

        by_word.flat_map(|(i, word)| {
            let pts = match corp.exceptions.iter()
                                .filter_map(|exs| exs.score(word))
                                .next() {
                    Some(vec) => Cow::Borrowed(vec),
                    None => Cow::Owned(corp.patterns.score(word))
            }.into_owned();
            let length = pts.len();
            let l = l_min;
            let r = if length >= length_min { length - l_min - r_min + 1 } else { 0 };

            word.char_indices().skip(l)
                .zip(pts.into_iter().skip(l).take(r))
                .filter(|&(_, p)| p % 2 != 0)
                .map(move |((i1, _), _)| i1 + i)
        }).collect()
    }

    /// Returns an iterator over string slices separated by valid hyphenation
    /// points.
    fn hyphenate(self, corp: &Corpus) -> Standard<'a> {
        let os = self.opportunities(corp);

        Standard {
            text: self,
            opportunities: os,
            prior: 0,
            i: 0
        }
    }
}
<commit_msg>Add punctuate_with()<commit_after>//! Hyphenating iterators.

use std::borrow::Cow;

use unicode_segmentation::UnicodeSegmentation;

use language::{Corpus};
use utilia::{Interspersable, Intersperse};

pub trait Hyphenation<Hyphenator> where Hyphenator : Iterator {
    /// Returns the indices of valid hyphenation points within the text.
    fn opportunities(self, corp: &Corpus) -> Vec<usize>;

    /// Returns an iterator over segments of the text separated by valid
    /// hyphenation points.
    ///
    /// When iterating over a word, such segments coincide with orthographic
    /// syllables. Note that, in some orthographies, the syllables of a hyphenated
    /// word are not necessarily substrings of the original word.
    fn hyphenate(self, corp: &Corpus) -> Hyphenator;
}

/// The `Standard` hyphenator iterates over a string, returning slices
/// delimited by string boundaries and valid hyphenation points.
///
/// For individual words, such slices coincide with orthographic syllables.
#[derive(Clone, Debug)]
pub struct Standard<'a> {
    text: &'a str,
    opportunities: Vec<usize>,
    prior: usize,
    i: usize
}

impl<'a> Standard<'a> {
    /// Inserts a soft hyphen at hyphenation points.
    pub fn punctuate(self) -> Intersperse<Self> {
        self.intersperse("\u{ad}")
    }

    /// Inserts a given string at hyphenation points.
    pub fn punctuate_with(self, mark: &'a str) -> Intersperse<Self> {
        self.intersperse(mark)
    }
}

impl<'a> Iterator for Standard<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        let start = self.prior;
        let i = self.i;

        match self.opportunities.get(i) {
            Some(&end) => {
                self.prior = end;
                self.i = i + 1;
                Some(&self.text[start .. end])
            },
            None => {
                if i <= self.opportunities.len() {
                    self.i = i + 1;
                    Some(&self.text[start ..])
                } else {
                    None
                }
            }
        }
    }
}

impl<'a> Hyphenation<Standard<'a>> for &'a str {
    /// Returns the byte indices of valid hyphenation points within the string.
    fn opportunities(self, corp: &Corpus) -> Vec<usize> {
        let (l_min, r_min) = (corp.left_min, corp.right_min);
        let length_min = l_min + r_min;

        if self.chars().count() < length_min {
            return vec![];
        }

        let by_word = self.split_word_bound_indices();

        by_word.flat_map(|(i, word)| {
            let pts = match corp.exceptions.iter()
                                .filter_map(|exs| exs.score(word))
                                .next() {
                    Some(vec) => Cow::Borrowed(vec),
                    None => Cow::Owned(corp.patterns.score(word))
            }.into_owned();
            let length = pts.len();
            let l = l_min;
            let r = if length >= length_min { length - l_min - r_min + 1 } else { 0 };

            word.char_indices().skip(l)
                .zip(pts.into_iter().skip(l).take(r))
                .filter(|&(_, p)| p % 2 != 0)
                .map(move |((i1, _), _)| i1 + i)
        }).collect()
    }

    /// Returns an iterator over string slices separated by valid hyphenation
    /// points.
    fn hyphenate(self, corp: &Corpus) -> Standard<'a> {
        let os = self.opportunities(corp);

        Standard {
            text: self,
            opportunities: os,
            prior: 0,
            i: 0
        }
    }
}
<|endoftext|>


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

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    thing.boundary_pm_final(max_bits - 1);

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail,
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Some(self.lists[index - 1].lookahead1),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }

    fn boundary_pm_final(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        // Count of last chain of list.
        let last_count = self.lists[index].lookahead1.count;

        let weight_sum = {
            let previous_list = &self.lists[index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
            let new_chain = self.node_arena.alloc(Node {
               weight: 0,
               count: last_count + 1,
               tail: self.lists[index].lookahead1.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            // ???
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Can't modify nodes in the lists [not compiling]<commit_after>use std::cmp::Ordering;

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    thing.boundary_pm_final(max_bits - 1);

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail,
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Some(self.lists[index - 1].lookahead1),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }

    fn boundary_pm_final(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        // Count of last chain of list.
        let last_count = self.lists[index].lookahead1.count;

        let weight_sum = {
            let previous_list = &self.lists[index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
            let new_chain = self.node_arena.alloc(Node {
               weight: 0,
               count: last_count + 1,
               tail: self.lists[index].lookahead1.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            // ???
            self.lists[index].lookahead1.tail = Some(self.lists[index - 1].lookahead1);
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>//! Merkle tree implementation
//!

use rayon::prelude::*;
use hash::*;

/// This halves the merkle tree leaves, taking it one level up
///
/// Calls itself recursively until one is left
fn shrink_merkle_tree(hashes: Vec<Hash32Buf>) -> Vec<Hash32Buf> {

    if hashes.len() == 1 {
        return hashes;
    }

    // the result is half the size rounded up
    let count = (hashes.len() + 1 ) / 2;

    let result = (0..count).into_iter().map(|n| {
        let ref first = hashes[n*2];

        // we double the first one if we have an odd number of hashes
        // in this layer
        let ref second = hashes.get(n*2+1).unwrap_or(&first);

        Hash32Buf::double_sha256_from_pair(
            first.as_ref(),
            second.as_ref()
        )
    }).collect();

    shrink_merkle_tree(result)
}

/// Calculates the merkle root for the given set of hashes
pub fn get_merkle_root(hashes: Vec<Hash32Buf>) -> Hash32Buf {

    shrink_merkle_tree(hashes)[0]
}

#[cfg(test)]
mod tests {

    use super::*;
    use util::*;

    #[test]
    fn test_merkle1() {

        const HASH1: &'static str = "212300e77d897f2f059366ed03c8bf2757bc2b1dd30df15d34f6f1ee521e58e8";
        const HASH2: &'static str = "4feec9316077e49b59bc23173303e13be9e9f5f9fa0660a58112a04a65a84ef1";
        const EXP_MERKLE: &'static str = "03b750bf691caf40b7e33d8e15f64dd16becf944b39a82710d6d257159361b93";

        let hash1 = Hash32Buf::from_slice(&from_hex_rev(HASH1));
        let hash2 = Hash32Buf::from_slice(&from_hex_rev(HASH2));
        let exp_merkle = Hash32Buf::from_slice(&from_hex_rev(EXP_MERKLE));

        let mut merkle = Vec::new();
        merkle.push(hash1);
        merkle.push(hash2);

        let merkle_root = get_merkle_root(merkle);

        assert_eq!(exp_merkle, merkle_root);

    }

    #[test]
    fn test_merkle2() {

        const HASH_SET : [&'static str;12] = [
            "30803bc3fefa999bf187cda4fff3647a78db6b957fcf5a579270c0535ec1601e",
            "101d83a3e4739640fbb6279883478bb6a2814e6fcd58322f0b1d3bf03983268a",
            "d646c47be5581891fc8d098e0db6e288efa22962d175f6d8c77913c2f898c0aa",
            "61419189bed89e85689442ee144e0aafc8b1ae6ed813e15f241b0750b97886ec",
            "09e5c37fe017605a15ae1fb139b2c186c76013c8cb78e6524261b29aed0b0424",
            "8226f6778f0f900ed38d9ca7314cfc55b28cc9c809e99570e3499ffeaa57759f",
            "0568ba35848086f4b3352fb468338ff727411fbf0e49fbe46eaf926929d205fe",
            "4e9dc6455ee50181a12e43b46de1bc72b201e627294dea7abe2f928b3e86bdb3",
            "6780b0e7801554d92d409f30b19328948ee3bcfd53764be30db7f1146fc12890",
            "dbcf135133afb7c3a256f87d09ae15994c5358bffd3e3993b56efb43742dce1e",
            "29454d128096bf66ebefe3f80a38edba4befac78c0eebfd351b7099f01933e2f",
            "c8cb5078d0faae9dc7bfcd150207c15d331e12161ec0e7c66c744d1201e08f3b"];

        const HASH_SET_MERKLE: &'static str =
            "1ba9abf54ae4cb022f53e767669931bacdd42c783fe063c5738ca49d29f1fbe3";

        let exp_merkle = Hash32Buf::from_slice(&from_hex_rev(HASH_SET_MERKLE));

        let mut merkle = Vec::new();
        for tx in HASH_SET.iter() {
            let txh = Hash32Buf::from_slice(&from_hex_rev(tx));
            merkle.push(txh);
        }

        let merkle_root = get_merkle_root(merkle);

        assert_eq!(exp_merkle, merkle_root);

    }
}<commit_msg>Parellelize merkle tree<commit_after>//! Merkle tree implementation
//!

use rayon::prelude::*;
use hash::*;

/// This halves the merkle tree leaves, taking it one level up
///
/// Calls itself recursively until one is left
fn shrink_merkle_tree(hashes: Vec<Hash32Buf>) -> Vec<Hash32Buf> {

    if hashes.len() == 1 {
        return hashes;
    }

    // the result is half the size rounded up
    let count = (hashes.len() + 1 ) / 2;

    let reduce = |n| {
        let ref first: Hash32Buf = hashes[n * 2];

        // we double the first one if we have an odd number of hashes
        // in this layer
        let ref second = hashes.get(n * 2 + 1).unwrap_or(first);

        Hash32Buf::double_sha256_from_pair(
            first.as_ref(),
            second.as_ref()
        )
    };

    let result = if count > 20
        { (0..count).into_par_iter().map(reduce).collect() }
    else
        { (0..count).into_iter().map(reduce).collect() };

    shrink_merkle_tree(result)
}

/// Calculates the merkle root for the given set of hashes
pub fn get_merkle_root(hashes: Vec<Hash32Buf>) -> Hash32Buf {

    shrink_merkle_tree(hashes)[0]
}

#[cfg(test)]
mod tests {

    use super::*;
    use util::*;

    #[test]
    fn test_merkle1() {

        const HASH1: &'static str = "212300e77d897f2f059366ed03c8bf2757bc2b1dd30df15d34f6f1ee521e58e8";
        const HASH2: &'static str = "4feec9316077e49b59bc23173303e13be9e9f5f9fa0660a58112a04a65a84ef1";
        const EXP_MERKLE: &'static str = "03b750bf691caf40b7e33d8e15f64dd16becf944b39a82710d6d257159361b93";

        let hash1 = Hash32Buf::from_slice(&from_hex_rev(HASH1));
        let hash2 = Hash32Buf::from_slice(&from_hex_rev(HASH2));
        let exp_merkle = Hash32Buf::from_slice(&from_hex_rev(EXP_MERKLE));

        let mut merkle = Vec::new();
        merkle.push(hash1);
        merkle.push(hash2);

        let merkle_root = get_merkle_root(merkle);

        assert_eq!(exp_merkle, merkle_root);

    }

    #[test]
    fn test_merkle2() {

        const HASH_SET : [&'static str;12] = [
            "30803bc3fefa999bf187cda4fff3647a78db6b957fcf5a579270c0535ec1601e",
            "101d83a3e4739640fbb6279883478bb6a2814e6fcd58322f0b1d3bf03983268a",
            "d646c47be5581891fc8d098e0db6e288efa22962d175f6d8c77913c2f898c0aa",
            "61419189bed89e85689442ee144e0aafc8b1ae6ed813e15f241b0750b97886ec",
            "09e5c37fe017605a15ae1fb139b2c186c76013c8cb78e6524261b29aed0b0424",
            "8226f6778f0f900ed38d9ca7314cfc55b28cc9c809e99570e3499ffeaa57759f",
            "0568ba35848086f4b3352fb468338ff727411fbf0e49fbe46eaf926929d205fe",
            "4e9dc6455ee50181a12e43b46de1bc72b201e627294dea7abe2f928b3e86bdb3",
            "6780b0e7801554d92d409f30b19328948ee3bcfd53764be30db7f1146fc12890",
            "dbcf135133afb7c3a256f87d09ae15994c5358bffd3e3993b56efb43742dce1e",
            "29454d128096bf66ebefe3f80a38edba4befac78c0eebfd351b7099f01933e2f",
            "c8cb5078d0faae9dc7bfcd150207c15d331e12161ec0e7c66c744d1201e08f3b"];

        const HASH_SET_MERKLE: &'static str =
            "1ba9abf54ae4cb022f53e767669931bacdd42c783fe063c5738ca49d29f1fbe3";

        let exp_merkle = Hash32Buf::from_slice(&from_hex_rev(HASH_SET_MERKLE));

        let mut merkle = Vec::new();
        for tx in HASH_SET.iter() {
            let txh = Hash32Buf::from_slice(&from_hex_rev(tx));
            merkle.push(txh);
        }

        let merkle_root = get_merkle_root(merkle);

        assert_eq!(exp_merkle, merkle_root);

    }
}<|endoftext|>