<commit_before>//! Implementation of compiling the compiler and standard library, in "check"-based modes.

use crate::builder::{Builder, Kind, RunConfig, ShouldRun, Step};
use crate::cache::Interned;
use crate::compile::{add_to_sysroot, run_cargo, rustc_cargo, rustc_cargo_env, std_cargo};
use crate::config::TargetSelection;
use crate::tool::{prepare_tool_cargo, SourceType};
use crate::INTERNER;
use crate::{Compiler, Mode, Subcommand};
use std::path::PathBuf;

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Std {
    pub target: TargetSelection,
}

/// Returns args for the subcommand itself (not for cargo)
fn args(builder: &Builder<'_>) -> Vec<String> {
    fn strings<'a>(arr: &'a [&str]) -> impl Iterator<Item = String> + 'a {
        arr.iter().copied().map(String::from)
    }

    if let Subcommand::Clippy { fix, .. } = builder.config.cmd {
        // disable the most spammy clippy lints
        let ignored_lints = vec![
            "many_single_char_names", // there are a lot in stdarch
            "collapsible_if",
            "type_complexity",
            "missing_safety_doc", // almost 3K warnings
            "too_many_arguments",
            "needless_lifetimes", // people want to keep the lifetimes
            "wrong_self_convention",
        ];
        let mut args = vec![];
        if fix {
            #[rustfmt::skip]
            args.extend(strings(&[
                "--fix", "-Zunstable-options",
                // FIXME: currently, `--fix` gives an error while checking tests for libtest,
                // possibly because libtest is not yet built in the sysroot.
                // As a workaround, avoid checking tests and benches when passed --fix.
                "--lib", "--bins", "--examples",
            ]));
        }
        args.extend(strings(&["--", "--cap-lints", "warn"]));
        args.extend(ignored_lints.iter().map(|lint| format!("-Aclippy::{}", lint)));
        args
    } else {
        vec![]
    }
}

fn cargo_subcommand(kind: Kind) -> &'static str {
    match kind {
        Kind::Check => "check",
        Kind::Clippy => "clippy",
        Kind::Fix => "fix",
        _ => unreachable!(),
    }
}

impl Step for Std {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.all_krates("test")
    }

    fn make_run(run: RunConfig<'_>) {
        run.builder.ensure(Std { target: run.target });
    }

    fn run(self, builder: &Builder<'_>) {
        let target = self.target;
        let compiler = builder.compiler(builder.top_stage, builder.config.build);

        let mut cargo = builder.cargo(
            compiler,
            Mode::Std,
            SourceType::InTree,
            target,
            cargo_subcommand(builder.kind),
        );
        std_cargo(builder, target, compiler.stage, &mut cargo);

        builder.info(&format!("Checking std artifacts ({} -> {})", &compiler.host, target));
        run_cargo(
            builder,
            cargo,
            args(builder),
            &libstd_stamp(builder, compiler, target),
            vec![],
            true,
        );

        // We skip populating the sysroot in non-zero stage because that'll lead
        // to rlib/rmeta conflicts if std gets built during this session.
        if compiler.stage == 0 {
            let libdir = builder.sysroot_libdir(compiler, target);
            let hostdir = builder.sysroot_libdir(compiler, compiler.host);
            add_to_sysroot(&builder, &libdir, &hostdir, &libstd_stamp(builder, compiler, target));
        }

        // Then run cargo again, once we've put the rmeta files for the library
        // crates into the sysroot. This is needed because e.g., core's tests
        // depend on `libtest` -- Cargo presumes it will exist, but it doesn't
        // since we initialize with an empty sysroot.
        //
        // Currently only the "libtest" tree of crates does this.

        if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
            let mut cargo = builder.cargo(
                compiler,
                Mode::Std,
                SourceType::InTree,
                target,
                cargo_subcommand(builder.kind),
            );
            std_cargo(builder, target, compiler.stage, &mut cargo);
            cargo.arg("--all-targets");

            // Explicitly pass -p for all dependencies krates -- this will force cargo
            // to also check the tests/benches/examples for these crates, rather
            // than just the leaf crate.
            for krate in builder.in_tree_crates("test", Some(target)) {
                cargo.arg("-p").arg(krate.name);
            }

            builder.info(&format!(
                "Checking std test/bench/example targets ({} -> {})",
                &compiler.host, target
            ));
            run_cargo(
                builder,
                cargo,
                args(builder),
                &libstd_test_stamp(builder, compiler, target),
                vec![],
                true,
            );
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustc {
    pub target: TargetSelection,
}

impl Step for Rustc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.all_krates("rustc-main")
    }

    fn make_run(run: RunConfig<'_>) {
        run.builder.ensure(Rustc { target: run.target });
    }

    /// Builds the compiler.
    ///
    /// This will build the compiler for a particular stage of the build using
    /// the `compiler` targeting the `target` architecture. The artifacts
    /// created will also be linked into the sysroot directory.
    fn run(self, builder: &Builder<'_>) {
        let compiler = builder.compiler(builder.top_stage, builder.config.build);
        let target = self.target;

        if compiler.stage != 0 {
            // If we're not in stage 0, then we won't have a std from the beta
            // compiler around. That means we need to make sure there's one in
            // the sysroot for the compiler to find. Otherwise, we're going to
            // fail when building crates that need to generate code (e.g., build
            // scripts and their dependencies).
            builder.ensure(crate::compile::Std { target: compiler.host, compiler });
            builder.ensure(crate::compile::Std { target, compiler });
        } else {
            builder.ensure(Std { target });
        }

        let mut cargo = builder.cargo(
            compiler,
            Mode::Rustc,
            SourceType::InTree,
            target,
            cargo_subcommand(builder.kind),
        );
        rustc_cargo(builder, &mut cargo, target);
        if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
            cargo.arg("--all-targets");
        }

        // Explicitly pass -p for all compiler krates -- this will force cargo
        // to also check the tests/benches/examples for these crates, rather
        // than just the leaf crate.
        for krate in builder.in_tree_crates("rustc-main", Some(target)) {
            cargo.arg("-p").arg(krate.name);
        }

        builder.info(&format!("Checking compiler artifacts ({} -> {})", &compiler.host, target));
        run_cargo(
            builder,
            cargo,
            args(builder),
            &librustc_stamp(builder, compiler, target),
            vec![],
            true,
        );

        let libdir = builder.sysroot_libdir(compiler, target);
        let hostdir = builder.sysroot_libdir(compiler, compiler.host);
        add_to_sysroot(&builder, &libdir, &hostdir, &librustc_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct CodegenBackend {
    pub target: TargetSelection,
    pub backend: Interned<String>,
}

impl Step for CodegenBackend {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.paths(&["compiler/rustc_codegen_cranelift", "rustc_codegen_cranelift"])
    }

    fn make_run(run: RunConfig<'_>) {
        for &backend in &[INTERNER.intern_str("cranelift")] {
            run.builder.ensure(CodegenBackend { target: run.target, backend });
        }
    }

    fn run(self, builder: &Builder<'_>) {
        let compiler = builder.compiler(builder.top_stage, builder.config.build);
        let target = self.target;
        let backend = self.backend;

        builder.ensure(Rustc { target });

        let mut cargo = builder.cargo(
            compiler,
            Mode::Codegen,
            SourceType::Submodule,
            target,
            cargo_subcommand(builder.kind),
        );
        cargo
            .arg("--manifest-path")
            .arg(builder.src.join(format!("compiler/rustc_codegen_{}/Cargo.toml", backend)));
        rustc_cargo_env(builder, &mut cargo, target);

        builder.info(&format!(
            "Checking {} artifacts ({} -> {})",
            backend, &compiler.host.triple, target.triple
        ));

        run_cargo(
            builder,
            cargo,
            args(builder),
            &codegen_backend_stamp(builder, compiler, target, backend),
            vec![],
            true,
        );
    }
}

macro_rules! tool_check_step {
    ($name:ident, $path:expr, $source_type:expr) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
        pub struct $name {
            pub target: TargetSelection,
        }

        impl Step for $name {
            type Output = ();
            const ONLY_HOSTS: bool = true;
            const DEFAULT: bool = true;

            fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
                run.path($path)
            }

            fn make_run(run: RunConfig<'_>) {
                run.builder.ensure($name { target: run.target });
            }

            fn run(self, builder: &Builder<'_>) {
                let compiler = builder.compiler(builder.top_stage, builder.config.build);
                let target = self.target;

                builder.ensure(Rustc { target });

                let mut cargo = prepare_tool_cargo(
                    builder,
                    compiler,
                    Mode::ToolRustc,
                    target,
                    cargo_subcommand(builder.kind),
                    $path,
                    $source_type,
                    &[],
                );

                if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
                    cargo.arg("--all-targets");
                }

                builder.info(&format!(
                    "Checking {} artifacts ({} -> {})",
                    stringify!($name).to_lowercase(),
                    &compiler.host.triple,
                    target.triple
                ));
                run_cargo(
                    builder,
                    cargo,
                    args(builder),
                    &stamp(builder, compiler, target),
                    vec![],
                    true,
                );

                let libdir = builder.sysroot_libdir(compiler, target);
                let hostdir = builder.sysroot_libdir(compiler, compiler.host);
                add_to_sysroot(&builder, &libdir, &hostdir, &stamp(builder, compiler, target));

                /// Cargo's output path in a given stage, compiled by a particular
                /// compiler for the specified target.
                fn stamp(
                    builder: &Builder<'_>,
                    compiler: Compiler,
                    target: TargetSelection,
                ) -> PathBuf {
                    builder
                        .cargo_out(compiler, Mode::ToolRustc, target)
                        .join(format!(".{}-check.stamp", stringify!($name).to_lowercase()))
                }
            }
        }
    };
}

tool_check_step!(Rustdoc, "src/tools/rustdoc", SourceType::InTree);
// Clippy is a hybrid. It is an external tool, but uses a git subtree instead
// of a submodule. Since the SourceType only drives the deny-warnings
// behavior, treat it as in-tree so that any new warnings in clippy will be
// rejected.
tool_check_step!(Clippy, "src/tools/clippy", SourceType::InTree);

tool_check_step!(Bootstrap, "src/bootstrap", SourceType::InTree);

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
fn libstd_stamp(builder: &Builder<'_>, compiler: Compiler, target: TargetSelection) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check.stamp")
}

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
fn libstd_test_stamp(
    builder: &Builder<'_>,
    compiler: Compiler,
    target: TargetSelection,
) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check-test.stamp")
}

/// Cargo's output path for librustc in a given stage, compiled by a particular
/// compiler for the specified target.
fn librustc_stamp(builder: &Builder<'_>, compiler: Compiler, target: TargetSelection) -> PathBuf {
    builder.cargo_out(compiler, Mode::Rustc, target).join(".librustc-check.stamp")
}

/// Cargo's output path for librustc_codegen_llvm in a given stage, compiled by a particular
/// compiler for the specified target and backend.
fn codegen_backend_stamp(
    builder: &Builder<'_>,
    compiler: Compiler,
    target: TargetSelection,
    backend: Interned<String>,
) -> PathBuf {
    builder
        .cargo_out(compiler, Mode::Codegen, target)
        .join(format!(".librustc_codegen_{}-check.stamp", backend))
}
<commit_msg>Print which stage is being checked (now that it may not be stage0)<commit_after>//! Implementation of compiling the compiler and standard library, in "check"-based modes.

use crate::builder::{Builder, Kind, RunConfig, ShouldRun, Step};
use crate::cache::Interned;
use crate::compile::{add_to_sysroot, run_cargo, rustc_cargo, rustc_cargo_env, std_cargo};
use crate::config::TargetSelection;
use crate::tool::{prepare_tool_cargo, SourceType};
use crate::INTERNER;
use crate::{Compiler, Mode, Subcommand};
use std::path::PathBuf;

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Std {
    pub target: TargetSelection,
}

/// Returns args for the subcommand itself (not for cargo)
fn args(builder: &Builder<'_>) -> Vec<String> {
    fn strings<'a>(arr: &'a [&str]) -> impl Iterator<Item = String> + 'a {
        arr.iter().copied().map(String::from)
    }

    if let Subcommand::Clippy { fix, .. } = builder.config.cmd {
        // disable the most spammy clippy lints
        let ignored_lints = vec![
            "many_single_char_names", // there are a lot in stdarch
            "collapsible_if",
            "type_complexity",
            "missing_safety_doc", // almost 3K warnings
            "too_many_arguments",
            "needless_lifetimes", // people want to keep the lifetimes
            "wrong_self_convention",
        ];
        let mut args = vec![];
        if fix {
            #[rustfmt::skip]
            args.extend(strings(&[
                "--fix", "-Zunstable-options",
                // FIXME: currently, `--fix` gives an error while checking tests for libtest,
                // possibly because libtest is not yet built in the sysroot.
                // As a workaround, avoid checking tests and benches when passed --fix.
                "--lib", "--bins", "--examples",
            ]));
        }
        args.extend(strings(&["--", "--cap-lints", "warn"]));
        args.extend(ignored_lints.iter().map(|lint| format!("-Aclippy::{}", lint)));
        args
    } else {
        vec![]
    }
}

fn cargo_subcommand(kind: Kind) -> &'static str {
    match kind {
        Kind::Check => "check",
        Kind::Clippy => "clippy",
        Kind::Fix => "fix",
        _ => unreachable!(),
    }
}

impl Step for Std {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.all_krates("test")
    }

    fn make_run(run: RunConfig<'_>) {
        run.builder.ensure(Std { target: run.target });
    }

    fn run(self, builder: &Builder<'_>) {
        let target = self.target;
        let compiler = builder.compiler(builder.top_stage, builder.config.build);

        let mut cargo = builder.cargo(
            compiler,
            Mode::Std,
            SourceType::InTree,
            target,
            cargo_subcommand(builder.kind),
        );
        std_cargo(builder, target, compiler.stage, &mut cargo);

        builder.info(&format!(
            "Checking stage{} std artifacts ({} -> {})",
            builder.top_stage, &compiler.host, target
        ));
        run_cargo(
            builder,
            cargo,
            args(builder),
            &libstd_stamp(builder, compiler, target),
            vec![],
            true,
        );

        // We skip populating the sysroot in non-zero stage because that'll lead
        // to rlib/rmeta conflicts if std gets built during this session.
        if compiler.stage == 0 {
            let libdir = builder.sysroot_libdir(compiler, target);
            let hostdir = builder.sysroot_libdir(compiler, compiler.host);
            add_to_sysroot(&builder, &libdir, &hostdir, &libstd_stamp(builder, compiler, target));
        }

        // Then run cargo again, once we've put the rmeta files for the library
        // crates into the sysroot. This is needed because e.g., core's tests
        // depend on `libtest` -- Cargo presumes it will exist, but it doesn't
        // since we initialize with an empty sysroot.
        //
        // Currently only the "libtest" tree of crates does this.

        if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
            let mut cargo = builder.cargo(
                compiler,
                Mode::Std,
                SourceType::InTree,
                target,
                cargo_subcommand(builder.kind),
            );
            std_cargo(builder, target, compiler.stage, &mut cargo);
            cargo.arg("--all-targets");

            // Explicitly pass -p for all dependencies krates -- this will force cargo
            // to also check the tests/benches/examples for these crates, rather
            // than just the leaf crate.
            for krate in builder.in_tree_crates("test", Some(target)) {
                cargo.arg("-p").arg(krate.name);
            }

            builder.info(&format!(
                "Checking stage{} std test/bench/example targets ({} -> {})",
                builder.top_stage, &compiler.host, target
            ));
            run_cargo(
                builder,
                cargo,
                args(builder),
                &libstd_test_stamp(builder, compiler, target),
                vec![],
                true,
            );
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustc {
    pub target: TargetSelection,
}

impl Step for Rustc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.all_krates("rustc-main")
    }

    fn make_run(run: RunConfig<'_>) {
        run.builder.ensure(Rustc { target: run.target });
    }

    /// Builds the compiler.
    ///
    /// This will build the compiler for a particular stage of the build using
    /// the `compiler` targeting the `target` architecture. The artifacts
    /// created will also be linked into the sysroot directory.
    fn run(self, builder: &Builder<'_>) {
        let compiler = builder.compiler(builder.top_stage, builder.config.build);
        let target = self.target;

        if compiler.stage != 0 {
            // If we're not in stage 0, then we won't have a std from the beta
            // compiler around. That means we need to make sure there's one in
            // the sysroot for the compiler to find. Otherwise, we're going to
            // fail when building crates that need to generate code (e.g., build
            // scripts and their dependencies).
            builder.ensure(crate::compile::Std { target: compiler.host, compiler });
            builder.ensure(crate::compile::Std { target, compiler });
        } else {
            builder.ensure(Std { target });
        }

        let mut cargo = builder.cargo(
            compiler,
            Mode::Rustc,
            SourceType::InTree,
            target,
            cargo_subcommand(builder.kind),
        );
        rustc_cargo(builder, &mut cargo, target);
        if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
            cargo.arg("--all-targets");
        }

        // Explicitly pass -p for all compiler krates -- this will force cargo
        // to also check the tests/benches/examples for these crates, rather
        // than just the leaf crate.
        for krate in builder.in_tree_crates("rustc-main", Some(target)) {
            cargo.arg("-p").arg(krate.name);
        }

        builder.info(&format!(
            "Checking stage{} compiler artifacts ({} -> {})",
            builder.top_stage, &compiler.host, target
        ));
        run_cargo(
            builder,
            cargo,
            args(builder),
            &librustc_stamp(builder, compiler, target),
            vec![],
            true,
        );

        let libdir = builder.sysroot_libdir(compiler, target);
        let hostdir = builder.sysroot_libdir(compiler, compiler.host);
        add_to_sysroot(&builder, &libdir, &hostdir, &librustc_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct CodegenBackend {
    pub target: TargetSelection,
    pub backend: Interned<String>,
}

impl Step for CodegenBackend {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
        run.paths(&["compiler/rustc_codegen_cranelift", "rustc_codegen_cranelift"])
    }

    fn make_run(run: RunConfig<'_>) {
        for &backend in &[INTERNER.intern_str("cranelift")] {
            run.builder.ensure(CodegenBackend { target: run.target, backend });
        }
    }

    fn run(self, builder: &Builder<'_>) {
        let compiler = builder.compiler(builder.top_stage, builder.config.build);
        let target = self.target;
        let backend = self.backend;

        builder.ensure(Rustc { target });

        let mut cargo = builder.cargo(
            compiler,
            Mode::Codegen,
            SourceType::Submodule,
            target,
            cargo_subcommand(builder.kind),
        );
        cargo
            .arg("--manifest-path")
            .arg(builder.src.join(format!("compiler/rustc_codegen_{}/Cargo.toml", backend)));
        rustc_cargo_env(builder, &mut cargo, target);

        builder.info(&format!(
            "Checking stage{} {} artifacts ({} -> {})",
            builder.top_stage, backend, &compiler.host.triple, target.triple
        ));

        run_cargo(
            builder,
            cargo,
            args(builder),
            &codegen_backend_stamp(builder, compiler, target, backend),
            vec![],
            true,
        );
    }
}

macro_rules! tool_check_step {
    ($name:ident, $path:expr, $source_type:expr) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
        pub struct $name {
            pub target: TargetSelection,
        }

        impl Step for $name {
            type Output = ();
            const ONLY_HOSTS: bool = true;
            const DEFAULT: bool = true;

            fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
                run.path($path)
            }

            fn make_run(run: RunConfig<'_>) {
                run.builder.ensure($name { target: run.target });
            }

            fn run(self, builder: &Builder<'_>) {
                let compiler = builder.compiler(builder.top_stage, builder.config.build);
                let target = self.target;

                builder.ensure(Rustc { target });

                let mut cargo = prepare_tool_cargo(
                    builder,
                    compiler,
                    Mode::ToolRustc,
                    target,
                    cargo_subcommand(builder.kind),
                    $path,
                    $source_type,
                    &[],
                );

                if let Subcommand::Check { all_targets: true, .. } = builder.config.cmd {
                    cargo.arg("--all-targets");
                }

                builder.info(&format!(
                    "Checking stage{} {} artifacts ({} -> {})",
                    builder.top_stage,
                    stringify!($name).to_lowercase(),
                    &compiler.host.triple,
                    target.triple
                ));
                run_cargo(
                    builder,
                    cargo,
                    args(builder),
                    &stamp(builder, compiler, target),
                    vec![],
                    true,
                );

                let libdir = builder.sysroot_libdir(compiler, target);
                let hostdir = builder.sysroot_libdir(compiler, compiler.host);
                add_to_sysroot(&builder, &libdir, &hostdir, &stamp(builder, compiler, target));

                /// Cargo's output path in a given stage, compiled by a particular
                /// compiler for the specified target.
                fn stamp(
                    builder: &Builder<'_>,
                    compiler: Compiler,
                    target: TargetSelection,
                ) -> PathBuf {
                    builder
                        .cargo_out(compiler, Mode::ToolRustc, target)
                        .join(format!(".{}-check.stamp", stringify!($name).to_lowercase()))
                }
            }
        }
    };
}

tool_check_step!(Rustdoc, "src/tools/rustdoc", SourceType::InTree);
// Clippy is a hybrid. It is an external tool, but uses a git subtree instead
// of a submodule. Since the SourceType only drives the deny-warnings
// behavior, treat it as in-tree so that any new warnings in clippy will be
// rejected.
tool_check_step!(Clippy, "src/tools/clippy", SourceType::InTree);

tool_check_step!(Bootstrap, "src/bootstrap", SourceType::InTree);

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
fn libstd_stamp(builder: &Builder<'_>, compiler: Compiler, target: TargetSelection) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check.stamp")
}

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
fn libstd_test_stamp(
    builder: &Builder<'_>,
    compiler: Compiler,
    target: TargetSelection,
) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check-test.stamp")
}

/// Cargo's output path for librustc in a given stage, compiled by a particular
/// compiler for the specified target.
fn librustc_stamp(builder: &Builder<'_>, compiler: Compiler, target: TargetSelection) -> PathBuf {
    builder.cargo_out(compiler, Mode::Rustc, target).join(".librustc-check.stamp")
}

/// Cargo's output path for librustc_codegen_llvm in a given stage, compiled by a particular
/// compiler for the specified target and backend.
fn codegen_backend_stamp(
    builder: &Builder<'_>,
    compiler: Compiler,
    target: TargetSelection,
    backend: Interned<String>,
) -> PathBuf {
    builder
        .cargo_out(compiler, Mode::Codegen, target)
        .join(format!(".librustc_codegen_{}-check.stamp", backend))
}
<|endoftext|>


<commit_before>use {Args, KvPair};
use quote::ToTokens;
use proc_macro2::TokenStream;
use syn::token::Comma;
use syn::punctuated::Punctuated;
use syn::{self, Expr, ExprClosure, FnArg, ArgCaptured};

impl Args {
    pub fn process(self, input: String) -> TokenStream {
        let mut body = TokenStream::new();

        let Args { elems, opts, closure } = self;

        if let Some(elems) = elems {
            process_elems(elems.0, &mut body);
        }

        if let Some(opts) = opts {
            process_opts(opts.0, &mut body);
        }

        if let Some(Expr::Closure(closure)) = closure {
            process_closure(closure, &mut body);
        }

        quote!({
            let mut args = Arg::new();
            args.set_string(&format!("args!({})", #input));
            #body
            args
        })
    }
}

fn process_elems(elems: impl IntoIterator<Item=Expr>, body: &mut TokenStream) {
    for expr in elems {
        if let Expr::Closure(closure) = expr {
            process_closure(closure, body);
            continue;
        }
        body.extend(quote! {
            args.add_arg(#expr.into_arg());
        });
    }
}

fn _process_list(exprs: impl IntoIterator<Item=Expr>, body: &mut TokenStream) {
    let mut arg = quote! {
        let mut list_arg = Arg::new();
    };
    for expr in exprs {
        arg.extend(quote! {
            list_arg.add_arg(#expr.into_arg());
        });
    }
    arg.extend(quote! {
        args.add_arg(list_arg.into_arg());
    });
    body.extend(arg);
}

fn process_opts(opts: impl IntoIterator<Item=KvPair>, body: &mut TokenStream) {
    for KvPair(key, val) in opts {
        let key = key.to_string();
        body.extend(quote! {
            match Arg::create_term_pair(#key, #val) {
                Ok(temp_pair) => args.add_opt(temp_pair),
                Err(error) => args.set_term(Err(error)),
            }
        });
    }
}

fn process_closure(mut closure: ExprClosure, body: &mut TokenStream) {
    let mut ids = quote! {
        #[allow(unused_mut)]
        let mut ids = Vec::new();
    };
    let mut args = Punctuated::new();
    for (i, mut arg) in closure.inputs.iter_mut().enumerate() {
        let var = var(i);
        ids.extend(quote! {
            for t in #var.term().unwrap().get_args() {
                ids.push(t.get_datum().clone());
            }
        });
        args.push(syn::parse(var.into()).unwrap());
        match arg {
            FnArg::Inferred(pat) => {
                let captured = ArgCaptured {
                    pat: pat.clone(),
                    colon_token: Default::default(),
                    ty: syn::parse_str("Client").unwrap(),
                };
                *arg = FnArg::Captured(captured);
            }
            arg => {
                let arg = arg.clone().into_token_stream().to_string();
                panic!(format!("`{}`: type annotations are not supported in ReQL closure arguments", arg.replace(" :", ":")));
            }
        }
    }
    ids.extend(quote!(ids));
    let func = func(closure, args, ids);
    body.extend(quote! {
        use reql::{Client, RepeatedField, Term, Datum, TT, DT};

        args.add_arg(#func.into_arg());
    });
}

fn func(closure: ExprClosure, args: Punctuated<Expr, Comma>, ids: TokenStream) -> TokenStream {
    quote! {{
        let res: Client = (#closure)(#args);
        let mut closure = Client::new();
        match res.term() {
            Ok(res) => {
                // ARRAY
                let mut array = Datum::new();
                array.set_field_type(DT::R_ARRAY);
                let args = RepeatedField::from_vec({#ids});
                array.set_r_array(args);
                // DATUM
                let mut datum = Term::new();
                datum.set_field_type(TT::DATUM);
                datum.set_datum(array);
                // FUNC
                let mut func = Term::new();
                func.set_field_type(TT::FUNC);
                let args = RepeatedField::from_vec(vec![datum, res.clone()]);
                func.set_args(args);
                closure.set_term(Ok(func));
            }
            Err(error) => {
                closure.set_term(Err(error));
            }
        }
        closure
    }}
}

fn var(idx: usize) -> TokenStream {
    quote! {{
        // ID
        let mut id = Datum::new();
        id.set_field_type(DT::R_NUM);
        id.set_r_num(#idx as f64);
        // DATUM
        let mut datum = Term::new();
        datum.set_field_type(TT::DATUM);
        datum.set_datum(id);
        // VAR
        let mut var = Term::new();
        var.set_field_type(TT::VAR);
        let args = RepeatedField::from_vec(vec![datum]);
        var.set_args(args);
        let mut client = Client::new();
        client.set_term(Ok(var));
        client
    }}
}
<commit_msg>use new error handling API<commit_after>use {Args, KvPair};
use proc_macro2::TokenStream;
use syn::token::Comma;
use syn::punctuated::Punctuated;
use syn::spanned::Spanned;
use syn::{self, Expr, ExprClosure, FnArg, ArgCaptured};

impl Args {
    pub fn process(self, input: String) -> TokenStream {
        let mut body = TokenStream::new();

        let Args { elems, opts, closure } = self;

        if let Some(elems) = elems {
            process_elems(elems.0, &mut body);
        }

        if let Some(opts) = opts {
            process_opts(opts.0, &mut body);
        }

        if let Some(Expr::Closure(closure)) = closure {
            process_closure(closure, &mut body);
        }

        quote!({
            let mut args = Arg::new();
            args.set_string(&format!("args!({})", #input));
            #body
            args
        })
    }
}

fn process_elems(elems: impl IntoIterator<Item=Expr>, body: &mut TokenStream) {
    for expr in elems {
        if let Expr::Closure(closure) = expr {
            process_closure(closure, body);
            continue;
        }
        body.extend(quote! {
            args.add_arg(#expr.into_arg());
        });
    }
}

fn _process_list(exprs: impl IntoIterator<Item=Expr>, body: &mut TokenStream) {
    let mut arg = quote! {
        let mut list_arg = Arg::new();
    };
    for expr in exprs {
        arg.extend(quote! {
            list_arg.add_arg(#expr.into_arg());
        });
    }
    arg.extend(quote! {
        args.add_arg(list_arg.into_arg());
    });
    body.extend(arg);
}

fn process_opts(opts: impl IntoIterator<Item=KvPair>, body: &mut TokenStream) {
    for KvPair(key, val) in opts {
        let key = key.to_string();
        body.extend(quote! {
            match Arg::create_term_pair(#key, #val) {
                Ok(temp_pair) => args.add_opt(temp_pair),
                Err(error) => args.set_term(Err(error)),
            }
        });
    }
}

fn process_closure(mut closure: ExprClosure, body: &mut TokenStream) {
    let mut ids = quote! {
        #[allow(unused_mut)]
        let mut ids = Vec::new();
    };
    let mut args = Punctuated::new();
    for (i, mut arg) in closure.inputs.iter_mut().enumerate() {
        let var = var(i);
        ids.extend(quote! {
            for t in #var.term().unwrap().get_args() {
                ids.push(t.get_datum().clone());
            }
        });
        args.push(syn::parse(var.into()).unwrap());
        match arg {
            FnArg::Inferred(pat) => {
                let captured = ArgCaptured {
                    pat: pat.clone(),
                    colon_token: Default::default(),
                    ty: syn::parse_str("Client").unwrap(),
                };
                *arg = FnArg::Captured(captured);
            }
            arg => {
                arg.span().unstable()
                    .error("unsupported ReQL closure argument")
                    .note("only arguments with no type annotations are supported")
                    .emit();
            }
        }
    }
    ids.extend(quote!(ids));
    let func = func(closure, args, ids);
    body.extend(quote! {
        use reql::{Client, RepeatedField, Term, Datum, TT, DT};

        args.add_arg(#func.into_arg());
    });
}

fn func(closure: ExprClosure, args: Punctuated<Expr, Comma>, ids: TokenStream) -> TokenStream {
    quote! {{
        let res: Client = (#closure)(#args);
        let mut closure = Client::new();
        match res.term() {
            Ok(res) => {
                // ARRAY
                let mut array = Datum::new();
                array.set_field_type(DT::R_ARRAY);
                let args = RepeatedField::from_vec({#ids});
                array.set_r_array(args);
                // DATUM
                let mut datum = Term::new();
                datum.set_field_type(TT::DATUM);
                datum.set_datum(array);
                // FUNC
                let mut func = Term::new();
                func.set_field_type(TT::FUNC);
                let args = RepeatedField::from_vec(vec![datum, res.clone()]);
                func.set_args(args);
                closure.set_term(Ok(func));
            }
            Err(error) => {
                closure.set_term(Err(error));
            }
        }
        closure
    }}
}

fn var(idx: usize) -> TokenStream {
    quote! {{
        // ID
        let mut id = Datum::new();
        id.set_field_type(DT::R_NUM);
        id.set_r_num(#idx as f64);
        // DATUM
        let mut datum = Term::new();
        datum.set_field_type(TT::DATUM);
        datum.set_datum(id);
        // VAR
        let mut var = Term::new();
        var.set_field_type(TT::VAR);
        let args = RepeatedField::from_vec(vec![datum]);
        var.set_args(args);
        let mut client = Client::new();
        client.set_term(Ok(var));
        client
    }}
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::io;
use std::rc::{Rc, Weak};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use mio::{Event, PollOpt, Ready, Token};
use mio::net::UdpSocket;

use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::datagram_buffer::DatagramBuffer;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;

const TAG: &'static str = "UDPConnection";

pub struct UDPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    socket: UdpSocket,
    token: Token,
    client_to_network: DatagramBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: Option<Box<[u8]>>, // used when the client buffer is full, should be rare
    closed: bool,
}

impl UDPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, reference_packet: &IPv4Packet) -> io::Result<Rc<RefCell<Self>>> {
        let socket = UDPConnection::create_socket(&id)?;
        let raw: &[u8] = reference_packet.raw();
        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            client: client,
            socket: socket,
            token: Token(0), // default value, will be set afterwards
            client_to_network: DatagramBuffer::new(4 * MAX_PACKET_LENGTH),
            network_to_client: Packetizer::new(raw, ipv4_header, transport_header),
            pending_packet_for_client: None,
            closed: false,
        }));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            let token = selector.register(&self_ref.socket, handler, Ready::readable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    fn create_socket(id: &ConnectionId) -> io::Result<UdpSocket> {
        let autobind_addr = SocketAddr::new(Ipv4Addr::new(0, 0, 0, 0).into(), 0);
        let udp_socket = UdpSocket::bind(&autobind_addr)?;
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        udp_socket.connect(rewritten_destination)?;
        Ok(udp_socket)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.disconnect();

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if let Err(err) = self.write() {
            error!(target: TAG, "Cannot write: {}", err);
            self.close(selector);
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        if let Err(err) = self.read() {
            error!(target: TAG, "Cannot read: {}", err);
            self.close(selector);
        }
        // TODO push to client
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let mut ready = Ready::empty();
        if self.may_read() {
            ready = Ready::readable();
        }
        if self.may_write() {
            ready = ready | Ready::writable();
        }
        selector.reregister(&self.socket, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        assert!(self.pending_packet_for_client.is_none());
        let packet = self.network_to_client.packetize(&mut self.socket)?;
        // TODO
        Ok(())
    }

    fn write(&mut self) -> io::Result<()> {
        self.client_to_network.write_to(&mut self.socket)?;
        Ok(())
    }

    fn may_read(&self) -> bool {
        // TODO
        true
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }
}

impl Connection for UDPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self) {
        // TODO
    }

    fn is_expired(&self) -> bool {
        // TODO
        false
    }
}
<commit_msg>Implement UDPConnection.process_receive()<commit_after>use std::cell::RefCell;
use std::io;
use std::rc::{Rc, Weak};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::UdpSocket;

use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::datagram_buffer::DatagramBuffer;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;

const TAG: &'static str = "UDPConnection";

pub struct UDPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    socket: UdpSocket,
    token: Token,
    client_to_network: DatagramBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: Option<Box<[u8]>>, // used when the client buffer is full, should be rare
    closed: bool,
}

impl UDPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, reference_packet: &IPv4Packet) -> io::Result<Rc<RefCell<Self>>> {
        let socket = UDPConnection::create_socket(&id)?;
        let raw: &[u8] = reference_packet.raw();
        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            client: client,
            socket: socket,
            token: Token(0), // default value, will be set afterwards
            client_to_network: DatagramBuffer::new(4 * MAX_PACKET_LENGTH),
            network_to_client: Packetizer::new(raw, ipv4_header, transport_header),
            pending_packet_for_client: None,
            closed: false,
        }));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            let token = selector.register(&self_ref.socket, handler, Ready::readable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    fn create_socket(id: &ConnectionId) -> io::Result<UdpSocket> {
        let autobind_addr = SocketAddr::new(Ipv4Addr::new(0, 0, 0, 0).into(), 0);
        let udp_socket = UdpSocket::bind(&autobind_addr)?;
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        udp_socket.connect(rewritten_destination)?;
        Ok(udp_socket)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.disconnect();

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if let Err(err) = self.write() {
            error!(target: TAG, "Cannot write: {}", err);
            self.close(selector);
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        if let Err(err) = self.read(selector) {
            error!(target: TAG, "Cannot read: {}", err);
            self.close(selector);
        }
    }

    fn read(&mut self, selector: &mut Selector) -> io::Result<()> {
        let ipv4_packet = self.network_to_client.packetize(&mut self.socket)?;
        let client_rc = self.client.upgrade().expect("expected client not found");
        match client_rc.borrow_mut().send_to_client(selector, &ipv4_packet) {
            Ok(_) => {
                debug!(target: TAG, "Packet ({} bytes) sent to client", ipv4_packet.length());
                if log_enabled!(target: TAG, LogLevel::Trace) {
                    // TODO log binary
                }
            },
            Err(err) => {
                // TODO copy packet locally for further sending

            }
        }
        Ok(())
    }

    fn write(&mut self) -> io::Result<()> {
        self.client_to_network.write_to(&mut self.socket)?;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let mut ready = Ready::empty();
        if self.may_read() {
            ready = Ready::readable();
        }
        if self.may_write() {
            ready = ready | Ready::writable();
        }
        selector.reregister(&self.socket, self.token, ready, PollOpt::level())
    }

    fn may_read(&self) -> bool {
        // TODO
        true
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }
}

impl Connection for UDPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self) {
        // TODO
    }

    fn is_expired(&self) -> bool {
        // TODO
        false
    }
}
<|endoftext|>


<commit_before>use futures::channel::oneshot;
use futures::executor::ThreadPool;
use futures::future::{self, ok, Future, FutureExt, TryFutureExt};
use futures::task::SpawnExt;
use std::sync::mpsc;
use std::thread;

fn run<F: Future + Send + 'static>(future: F) {
    let tp = ThreadPool::new().unwrap();
    tp.spawn(future.map(drop)).unwrap();
}

#[test]
fn join1() {
    let (tx, rx) = mpsc::channel();
    run(future::try_join(ok::<i32, i32>(1), ok(2)).map_ok(move |v| tx.send(v).unwrap()));
    assert_eq!(rx.recv(), Ok((1, 2)));
    assert!(rx.recv().is_err());
}

#[test]
fn join2() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    assert!(rx.try_recv().is_err());
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok((1, 2)));
    assert!(rx.recv().is_err());
}

#[test]
fn join3() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_err(move |_v| tx.send(1).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    assert_eq!(rx.recv(), Ok(1));
    assert!(rx.recv().is_err());
    drop(c2);
}

#[test]
fn join4() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_err(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    assert!(rx.recv().is_ok());
    drop(c2);
    assert!(rx.recv().is_err());
}

#[test]
fn join5() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (c3, p3) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(future::try_join(p1, p2), p3).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    assert!(rx.try_recv().is_err());
    c2.send(2).unwrap();
    assert!(rx.try_recv().is_err());
    c3.send(3).unwrap();
    assert_eq!(rx.recv(), Ok(((1, 2), 3)));
    assert!(rx.recv().is_err());
}

#[test]
fn select1() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    let (v, p2) = rx.recv().unwrap().into_inner();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_ok(move |v| tx.send(v).unwrap()));
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select2() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_err(move |v| tx.send((1, v.into_inner().1)).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    let (v, p2) = rx.recv().unwrap();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_ok(move |v| tx.send(v).unwrap()));
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select3() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_err(move |v| tx.send((1, v.into_inner().1)).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    let (v, p2) = rx.recv().unwrap();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_err(move |_v| tx.send(2).unwrap()));
    drop(c2);
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select4() {
    let (tx, rx) = mpsc::channel::<oneshot::Sender<i32>>();

    let t = thread::spawn(move || {
        for c in rx {
            c.send(1).unwrap();
        }
    });

    let (tx2, rx2) = mpsc::channel();
    for _ in 0..10000 {
        let (c1, p1) = oneshot::channel::<i32>();
        let (c2, p2) = oneshot::channel::<i32>();

        let tx3 = tx2.clone();
        run(future::try_select(p1, p2).map_ok(move |_| tx3.send(()).unwrap()));
        tx.send(c1).unwrap();
        rx2.recv().unwrap();
        drop(c2);
    }
    drop(tx);

    t.join().unwrap();
}
<commit_msg>futures: tests/eventual.rs : Disable w/o executor + thread-pool<commit_after>#![cfg(all(feature = "executor", feature = "thread-pool"))]
use futures::channel::oneshot;
use futures::executor::ThreadPool;
use futures::future::{self, ok, Future, FutureExt, TryFutureExt};
use futures::task::SpawnExt;
use std::sync::mpsc;
use std::thread;

fn run<F: Future + Send + 'static>(future: F) {
    let tp = ThreadPool::new().unwrap();
    tp.spawn(future.map(drop)).unwrap();
}

#[test]
fn join1() {
    let (tx, rx) = mpsc::channel();
    run(future::try_join(ok::<i32, i32>(1), ok(2)).map_ok(move |v| tx.send(v).unwrap()));
    assert_eq!(rx.recv(), Ok((1, 2)));
    assert!(rx.recv().is_err());
}

#[test]
fn join2() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    assert!(rx.try_recv().is_err());
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok((1, 2)));
    assert!(rx.recv().is_err());
}

#[test]
fn join3() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_err(move |_v| tx.send(1).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    assert_eq!(rx.recv(), Ok(1));
    assert!(rx.recv().is_err());
    drop(c2);
}

#[test]
fn join4() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(p1, p2).map_err(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    assert!(rx.recv().is_ok());
    drop(c2);
    assert!(rx.recv().is_err());
}

#[test]
fn join5() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (c3, p3) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_join(future::try_join(p1, p2), p3).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    assert!(rx.try_recv().is_err());
    c2.send(2).unwrap();
    assert!(rx.try_recv().is_err());
    c3.send(3).unwrap();
    assert_eq!(rx.recv(), Ok(((1, 2), 3)));
    assert!(rx.recv().is_err());
}

#[test]
fn select1() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_ok(move |v| tx.send(v).unwrap()));
    assert!(rx.try_recv().is_err());
    c1.send(1).unwrap();
    let (v, p2) = rx.recv().unwrap().into_inner();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_ok(move |v| tx.send(v).unwrap()));
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select2() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_err(move |v| tx.send((1, v.into_inner().1)).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    let (v, p2) = rx.recv().unwrap();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_ok(move |v| tx.send(v).unwrap()));
    c2.send(2).unwrap();
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select3() {
    let (c1, p1) = oneshot::channel::<i32>();
    let (c2, p2) = oneshot::channel::<i32>();
    let (tx, rx) = mpsc::channel();
    run(future::try_select(p1, p2).map_err(move |v| tx.send((1, v.into_inner().1)).unwrap()));
    assert!(rx.try_recv().is_err());
    drop(c1);
    let (v, p2) = rx.recv().unwrap();
    assert_eq!(v, 1);
    assert!(rx.recv().is_err());

    let (tx, rx) = mpsc::channel();
    run(p2.map_err(move |_v| tx.send(2).unwrap()));
    drop(c2);
    assert_eq!(rx.recv(), Ok(2));
    assert!(rx.recv().is_err());
}

#[test]
fn select4() {
    let (tx, rx) = mpsc::channel::<oneshot::Sender<i32>>();

    let t = thread::spawn(move || {
        for c in rx {
            c.send(1).unwrap();
        }
    });

    let (tx2, rx2) = mpsc::channel();
    for _ in 0..10000 {
        let (c1, p1) = oneshot::channel::<i32>();
        let (c2, p2) = oneshot::channel::<i32>();

        let tx3 = tx2.clone();
        run(future::try_select(p1, p2).map_ok(move |_| tx3.send(()).unwrap()));
        tx.send(c1).unwrap();
        rx2.recv().unwrap();
        drop(c2);
    }
    drop(tx);

    t.join().unwrap();
}
<|endoftext|>


<commit_before>//
// base64n.rs - new base64 module
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//      0 A            17 R            34 i            51 z
//      1 B            18 S            35 j            52 0
//      2 C            19 T            36 k            53 1
//      3 D            20 U            37 l            54 2
//      4 E            21 V            38 m            55 3
//      5 F            22 W            39 n            56 4
//      6 G            23 X            40 o            57 5
//      7 H            24 Y            41 p            58 6
//      8 I            25 Z            42 q            59 7
//      9 J            26 a            43 r            60 8
//     10 K            27 b            44 s            61 9
//     11 L            28 c            45 t            62 +
//     12 M            29 d            46 u            63 /
//     13 N            30 e            47 v
//     14 O            31 f            48 w         (pad) =
//     15 P            32 g            49 x
//     16 Q            33 h            50 y
//
// The "URL and Filename safe" Base 64 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 A            17 R            34 i            51 z
//     1 B            18 S            35 j            52 0
//     2 C            19 T            36 k            53 1
//     3 D            20 U            37 l            54 2
//     4 E            21 V            38 m            55 3
//     5 F            22 W            39 n            56 4
//     6 G            23 X            40 o            57 5
//     7 H            24 Y            41 p            58 6
//     8 I            25 Z            42 q            59 7
//     9 J            26 a            43 r            60 8
//    10 K            27 b            44 s            61 9
//    11 L            28 c            45 t            62 - (minus)
//    12 M            29 d            46 u            63 _ (underline)
//    13 N            30 e            47 v
//    14 O            31 f            48 w
//    15 P            32 g            49 x
//    16 Q            33 h            50 y         (pad) =
//

import vec::len;

const PAD: u8 = 61u8;

iface enc {
    fn encode(dst: [mutable u8], src: [u8]);
    fn encode_byte(src: [u8]) -> [u8];
    // FIXME `decode` should return desired length of `dst`
    fn decode(dst: [mutable u8], src: [u8]);
    fn decode_byte(src: [u8]) -> [u8];
}

fn mk() -> enc {
    type _enc = {table: [u8]};

    impl of enc for _enc {
        fn encode(dst: [mutable u8], src: [u8]) {
            b64encode(self.table, dst, src);
        }
        fn decode(dst: [mutable u8], src: [u8]) {
            // FIXME need self.decode_map instead of self.table.
            b64decode(self.table, dst, src);
        }
        fn encode_byte(src: [u8]) -> [u8] { [] }
        fn decode_byte(src: [u8]) -> [u8] { [] }
    }

    let table = vec::to_mut(vec::from_elem(62u, 0u8)), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }

    {table: vec::from_mut(table)} as enc
}

fn encoded_len(src_length: uint) -> uint {
    (src_length + 2u) / 3u * 4u
}

fn decoded_len(src_length: uint) -> uint {
    src_length / 4u * 3u
}

fn b64encode(table: [u8], dst: [mutable u8], src: [u8]) {
    if len(src) == 0u {
        ret;
    }

    let src_length = len(src);
    let dst_length = len(dst);
    let dst_curr = 0u;
    let src_curr = 0u;

    if dst_length % 4u != 0u {
        fail "dst's length should be divisible by 4";
    }

    while src_length > 0u {
        dst[dst_curr + 0u] = 0u8;
        dst[dst_curr + 1u] = 0u8;
        dst[dst_curr + 2u] = 0u8;
        dst[dst_curr + 3u] = 0u8;

        if src_length == 1u {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
        } else if src_length == 2u {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
            dst[dst_curr + 1u] |= (src[src_curr + 1u] >> 4u8);
            dst[dst_curr + 2u] |= (src[src_curr + 1u] << 2u8) & 0x3f_u8;
        } else {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
            dst[dst_curr + 1u] |= (src[src_curr + 1u] >> 4u8);
            dst[dst_curr + 2u] |= (src[src_curr + 1u] << 2u8) & 0x3f_u8;
            dst[dst_curr + 2u] |= (src[src_curr + 2u] >> 6u8);
            dst[dst_curr + 3u] |= (src[src_curr + 3u]) & 0x3f_u8;
        }

        dst[dst_curr + 0u] = table[dst[dst_curr + 0u]];
        dst[dst_curr + 1u] = table[dst[dst_curr + 1u]];
        dst[dst_curr + 2u] = table[dst[dst_curr + 2u]];
        dst[dst_curr + 3u] = table[dst[dst_curr + 3u]];

        if src_length < 3u {
            dst[dst_curr + 3u] = PAD;
            if src_length < 2u {
                dst[dst_curr + 2u] = PAD;
            }
            break;
        }

        src_length -= 3u;
        src_curr += 3u;
        dst_curr += 4u;
    }
}

fn b64decode(table: [u8], dst: [mutable u8], src: [u8]) {
    // FIXME write
}
<commit_msg>wip: base64n.rs<commit_after>//
// base64n.rs - new base64 module
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//      0 A            17 R            34 i            51 z
//      1 B            18 S            35 j            52 0
//      2 C            19 T            36 k            53 1
//      3 D            20 U            37 l            54 2
//      4 E            21 V            38 m            55 3
//      5 F            22 W            39 n            56 4
//      6 G            23 X            40 o            57 5
//      7 H            24 Y            41 p            58 6
//      8 I            25 Z            42 q            59 7
//      9 J            26 a            43 r            60 8
//     10 K            27 b            44 s            61 9
//     11 L            28 c            45 t            62 +
//     12 M            29 d            46 u            63 /
//     13 N            30 e            47 v
//     14 O            31 f            48 w         (pad) =
//     15 P            32 g            49 x
//     16 Q            33 h            50 y
//
// The "URL and Filename safe" Base 64 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 A            17 R            34 i            51 z
//     1 B            18 S            35 j            52 0
//     2 C            19 T            36 k            53 1
//     3 D            20 U            37 l            54 2
//     4 E            21 V            38 m            55 3
//     5 F            22 W            39 n            56 4
//     6 G            23 X            40 o            57 5
//     7 H            24 Y            41 p            58 6
//     8 I            25 Z            42 q            59 7
//     9 J            26 a            43 r            60 8
//    10 K            27 b            44 s            61 9
//    11 L            28 c            45 t            62 - (minus)
//    12 M            29 d            46 u            63 _ (underline)
//    13 N            30 e            47 v
//    14 O            31 f            48 w
//    15 P            32 g            49 x
//    16 Q            33 h            50 y         (pad) =
//

import vec::len;

const PAD: u8 = 61u8;

iface enc {
    fn encode(dst: [mutable u8], src: [u8]);
    fn encode_byte(src: [u8]) -> [u8];
    // FIXME `decode` should return desired length of `dst`
    fn decode(dst: [mutable u8], src: [u8]);
    fn decode_byte(src: [u8]) -> [u8];
}

fn mk() -> enc {
    type _enc = {table: [u8]};

    impl of enc for _enc {
        fn encode(dst: [mutable u8], src: [u8]) {
            b64encode(self.table, dst, src);
        }
        fn decode(dst: [mutable u8], src: [u8]) {
            // FIXME need self.decode_map instead of self.table.
            b64decode(self.table, dst, src);
        }
        fn encode_byte(src: [u8]) -> [u8] {
            let dst_length = encoded_len(len(src));
            let dst = vec::to_mut(vec::from_elem(dst_length, 0u8));
            self.encode(dst, src);
            vec::from_mut(dst)
        }
        fn decode_byte(src: [u8]) -> [u8] { [] }
    }

    let table = vec::to_mut(vec::from_elem(62u, 0u8)), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }

    {table: vec::from_mut(table)} as enc
}

fn encoded_len(src_length: uint) -> uint {
    (src_length + 2u) / 3u * 4u
}

fn decoded_len(src_length: uint) -> uint {
    src_length / 4u * 3u
}

fn b64encode(table: [u8], dst: [mutable u8], src: [u8]) {
    if len(src) == 0u {
        ret;
    }

    let src_length = len(src);
    let dst_length = len(dst);
    let dst_curr = 0u;
    let src_curr = 0u;

    if dst_length % 4u != 0u {
        fail "dst's length should be divisible by 4";
    }

    while src_length > 0u {
        dst[dst_curr + 0u] = 0u8;
        dst[dst_curr + 1u] = 0u8;
        dst[dst_curr + 2u] = 0u8;
        dst[dst_curr + 3u] = 0u8;

        if src_length == 1u {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
        } else if src_length == 2u {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
            dst[dst_curr + 1u] |= (src[src_curr + 1u] >> 4u8);
            dst[dst_curr + 2u] |= (src[src_curr + 1u] << 2u8) & 0x3f_u8;
        } else {
            dst[dst_curr + 0u] |= (src[src_curr + 0u]) >> 2u8;
            dst[dst_curr + 1u] |= (src[src_curr + 0u] << 4u8) & 0x3f_u8;
            dst[dst_curr + 1u] |= (src[src_curr + 1u] >> 4u8);
            dst[dst_curr + 2u] |= (src[src_curr + 1u] << 2u8) & 0x3f_u8;
            dst[dst_curr + 2u] |= (src[src_curr + 2u] >> 6u8);
            dst[dst_curr + 3u] |= (src[src_curr + 3u]) & 0x3f_u8;
        }

        dst[dst_curr + 0u] = table[dst[dst_curr + 0u]];
        dst[dst_curr + 1u] = table[dst[dst_curr + 1u]];
        dst[dst_curr + 2u] = table[dst[dst_curr + 2u]];
        dst[dst_curr + 3u] = table[dst[dst_curr + 3u]];

        if src_length < 3u {
            dst[dst_curr + 3u] = PAD;
            if src_length < 2u {
                dst[dst_curr + 2u] = PAD;
            }
            break;
        }

        src_length -= 3u;
        src_curr += 3u;
        dst_curr += 4u;
    }
}

fn b64decode(table: [u8], dst: [mutable u8], src: [u8]) {
    // FIXME write
}
<|endoftext|>


<commit_before>pub struct BitVector {
    value: u64,
}

impl BitVector {
    pub fn new() -> Self {
        BitVector { value: 0 }
    }

    pub fn set(&mut self, i: usize) {
        self.value |= 1 << i;
    }

    pub fn has(&self, i: usize) -> bool {
        self.value & 1 << i != 0
    }

    pub fn intersect(&self, other: &Self) -> Self {
        BitVector { value: self.value & other.value }
    }

    pub fn is_empty(&self) -> bool {
        self.value == 0
    }

    pub fn to_vec(&self) -> Vec<usize> {
        let mut vec = vec![];
        for i in 0..64 {
            if self.has(i) {
                vec.push(i);
            }
        }
        vec
    }
}

impl Clone for BitVector {
    fn clone(&self) -> Self {
        BitVector { value: self.value }
    }
}
<commit_msg>[independent_set/rust] Simplify BitVector to_vec<commit_after>pub struct BitVector {
    value: u64,
}

impl BitVector {
    pub fn new() -> Self {
        BitVector { value: 0 }
    }

    pub fn set(&mut self, i: usize) {
        self.value |= 1 << i;
    }

    pub fn has(&self, i: usize) -> bool {
        self.value & 1 << i != 0
    }

    pub fn intersect(&self, other: &Self) -> Self {
        BitVector { value: self.value & other.value }
    }

    pub fn is_empty(&self) -> bool {
        self.value == 0
    }

    pub fn to_vec(&self) -> Vec<usize> {
        (0..64).filter(|&i| self.has(i)).collect()
    }
}

impl Clone for BitVector {
    fn clone(&self) -> Self {
        BitVector { value: self.value }
    }
}
<|endoftext|>


<commit_before>use std::{cmp::min, ffi::OsStr, fs::File, io, mem, num::Wrapping, path::PathBuf, time::Duration};

use clap::Parser;
use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 100;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    endpoint: String,
    filename: PathBuf,
    client: reqwest::blocking::Client,
    rng: SmallRng,
    spinner_idx: Wrapping<usize>,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(endpoint: &str, filename: PathBuf) -> Importer {
        Importer {
            endpoint: endpoint.to_owned(),
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            spinner_idx: Wrapping(0),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put(format!("{}/import/lichess", self.endpoint))
            .json(&self.batch)
            .send()
            .expect("send batch");

        self.spinner_idx += Wrapping(1);
        let spinner = &['⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'];

        println!(
            "{} {:?}: {}: {} - {}",
            spinner[self.spinner_idx.0 % spinner.len()],
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(true, |name| name == "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.00,

                _ if rating >= 2500 => 1.00,

                Speed::Rapid if rating >= 2200 => 1.00,
                Speed::Rapid if rating >= 2000 => 0.50,
                Speed::Rapid if rating >= 1800 => 0.28,
                Speed::Rapid if rating >= 1600 => 0.24,

                Speed::Blitz if rating >= 2200 => 0.24,
                Speed::Blitz if rating >= 2000 => 0.11,
                Speed::Blitz if rating >= 1600 => 0.08,

                Speed::Bullet if rating >= 2200 => 0.30,
                Speed::Bullet if rating >= 2000 => 0.17,
                Speed::Bullet if rating >= 1600 => 0.11,

                Speed::UltraBullet => 1.00,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.00
            } else {
                0.50
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

#[derive(Parser)]
struct Args {
    #[clap(long, default_value = "http://localhost:9002")]
    endpoint: String,
    pgns: Vec<PathBuf>,
}

fn main() -> Result<(), io::Error> {
    let args = Args::parse();

    for arg in args.pgns {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.extension() == Some(OsStr::new("bz2")) {
            println!("Reading compressed {:?} ...", arg);
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            println!("Reading {:?} ...", arg);
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(&args.endpoint, arg);
        reader.read_all(&mut importer)?;
        importer.send();
    }

    Ok(())
}
<commit_msg>increase sampling rates<commit_after>use std::{cmp::min, ffi::OsStr, fs::File, io, mem, num::Wrapping, path::PathBuf, time::Duration};

use clap::Parser;
use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 100;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    endpoint: String,
    filename: PathBuf,
    client: reqwest::blocking::Client,
    rng: SmallRng,
    spinner_idx: Wrapping<usize>,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(endpoint: &str, filename: PathBuf) -> Importer {
        Importer {
            endpoint: endpoint.to_owned(),
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            spinner_idx: Wrapping(0),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put(format!("{}/import/lichess", self.endpoint))
            .json(&self.batch)
            .send()
            .expect("send batch");

        self.spinner_idx += Wrapping(1);
        let spinner = &['⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'];

        println!(
            "{} {:?}: {}: {} - {}",
            spinner[self.spinner_idx.0 % spinner.len()],
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(true, |name| name == "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.00,

                _ if rating >= 2500 => 1.00,

                Speed::Rapid if rating >= 2200 => 1.00,
                Speed::Rapid if rating >= 2000 => 0.83,
                Speed::Rapid if rating >= 1800 => 0.46,
                Speed::Rapid if rating >= 1600 => 0.39,

                Speed::Blitz if rating >= 2200 => 0.38,
                Speed::Blitz if rating >= 2000 => 0.18,
                Speed::Blitz if rating >= 1600 => 0.13,

                Speed::Bullet if rating >= 2200 => 0.48,
                Speed::Bullet if rating >= 2000 => 0.27,
                Speed::Bullet if rating >= 1800 => 0.19,
                Speed::Bullet if rating >= 1600 => 0.18,

                Speed::UltraBullet => 1.00,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.00
            } else {
                0.50
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

#[derive(Parser)]
struct Args {
    #[clap(long, default_value = "http://localhost:9002")]
    endpoint: String,
    pgns: Vec<PathBuf>,
}

fn main() -> Result<(), io::Error> {
    let args = Args::parse();

    for arg in args.pgns {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.extension() == Some(OsStr::new("bz2")) {
            println!("Reading compressed {:?} ...", arg);
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            println!("Reading {:?} ...", arg);
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(&args.endpoint, arg);
        reader.read_all(&mut importer)?;
        importer.send();
    }

    Ok(())
}
<|endoftext|>