<commit_before>//! Tests for future-incompat-report messages

use cargo_test_support::registry::Package;
use cargo_test_support::{basic_manifest, is_nightly, project, Project};

// An arbitrary lint (array_into_iter) that triggers a report.
const FUTURE_EXAMPLE: &'static str = "fn main() { [true].into_iter(); }";
// Some text that will be displayed when the lint fires.
const FUTURE_OUTPUT: &'static str = "[..]array_into_iter[..]";

fn simple_project() -> Project {
    project()
        .file("Cargo.toml", &basic_manifest("foo", "0.0.0"))
        .file("src/main.rs", FUTURE_EXAMPLE)
        .build()
}

#[cargo_test]
fn no_output_on_stable() {
    let p = simple_project();

    p.cargo("build")
        .with_stderr_contains(FUTURE_OUTPUT)
        .with_stderr_does_not_contain("[..]cargo report[..]")
        .run();
}

#[cargo_test]
fn gate_future_incompat_report() {
    let p = simple_project();

    p.cargo("build --future-incompat-report")
        .with_stderr_contains("error: the `--future-incompat-report` flag is unstable[..]")
        .with_status(101)
        .run();

    // Both `-Z future-incompat-report` and `-Z unstable-opts` are required
    p.cargo("build --future-incompat-report -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains("error: the `--future-incompat-report` flag is unstable[..]")
        .with_status(101)
        .run();

    p.cargo("build --future-incompat-report -Z unstable-options")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains(
            "error: Usage of `--future-incompat-report` requires `-Z future-incompat-report`",
        )
        .with_status(101)
        .run();

    p.cargo("report future-incompatibilities --id foo")
        .with_stderr_contains("error: `cargo report` can only be used on the nightly channel")
        .with_status(101)
        .run();
}

#[cargo_test]
fn test_zero_future_incompat() {
    if !is_nightly() {
        return;
    }

    let p = project()
        .file("Cargo.toml", &basic_manifest("foo", "0.0.0"))
        .file("src/main.rs", "fn main() {}")
        .build();

    // No note if --future-incompat-report is not specified.
    p.cargo("build -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stderr(
            "\
[COMPILING] foo v0.0.0 [..]
[FINISHED] [..]
",
        )
        .run();

    p.cargo("build --future-incompat-report -Z unstable-options -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stderr(
            "\
[FINISHED] [..]
note: 0 dependencies had future-incompatible warnings
",
        )
        .run();
}

#[cargo_test]
#[ignore] // Waiting on https://github.com/rust-lang/rust/pull/86478
fn test_single_crate() {
    if !is_nightly() {
        return;
    }

    let p = simple_project();

    for command in &["build", "check", "rustc", "test"] {
        p.cargo(command).arg("-Zfuture-incompat-report")
            .masquerade_as_nightly_cargo()
            .with_stderr_contains(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: foo v0.0.0 [..]")
            .with_stderr_does_not_contain("[..]incompatibility[..]")
            .run();

        p.cargo(command).arg("-Zfuture-incompat-report").arg("-Zunstable-options").arg("--future-incompat-report")
            .masquerade_as_nightly_cargo()
            .with_stderr_contains(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: foo v0.0.0 [..]")
            .with_stderr_contains("The package `foo v0.0.0 ([..])` currently triggers the following future incompatibility lints:")
            .run();
    }
}

#[cargo_test]
#[ignore] // Waiting on https://github.com/rust-lang/rust/pull/86478
fn test_multi_crate() {
    if !is_nightly() {
        return;
    }

    Package::new("first-dep", "0.0.1")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("second-dep", "0.0.2")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.0"

                [dependencies]
                first-dep = "*"
                second-dep = "*"
              "#,
        )
        .file("src/main.rs", "fn main() {}")
        .build();

    for command in &["build", "check", "rustc", "test"] {
        p.cargo(command).arg("-Zfuture-incompat-report")
            .masquerade_as_nightly_cargo()
            .with_stderr_does_not_contain(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: first-dep v0.0.1, second-dep v0.0.2")
            // Check that we don't have the 'triggers' message shown at the bottom of this loop
            .with_stderr_does_not_contain("[..]triggers[..]")
            .run();

        p.cargo(command).arg("-Zunstable-options").arg("-Zfuture-incompat-report").arg("--future-incompat-report")
            .masquerade_as_nightly_cargo()
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: first-dep v0.0.1, second-dep v0.0.2")
            .with_stderr_contains("The package `first-dep v0.0.1` currently triggers the following future incompatibility lints:")
            .with_stderr_contains("The package `second-dep v0.0.2` currently triggers the following future incompatibility lints:")
            .run();
    }

    // Test that passing the correct id via '--id' doesn't generate a warning message
    let output = p
        .cargo("build -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .exec_with_output()
        .unwrap();

    // Extract the 'id' from the stdout. We are looking
    // for the id in a line of the form "run `cargo report future-incompatibilities --id yZ7S`"
    // which is generated by Cargo to tell the user what command to run
    // This is just to test that passing the id suppresses the warning mesasge. Any users needing
    // access to the report from a shell script should use the `--future-incompat-report` flag
    let stderr = std::str::from_utf8(&output.stderr).unwrap();

    // Find '--id <ID>' in the output
    let mut iter = stderr.split(' ');
    iter.find(|w| *w == "--id").unwrap();
    let id = iter
        .next()
        .unwrap_or_else(|| panic!("Unexpected output:\n{}", stderr));
    // Strip off the trailing '`' included in the output
    let id: String = id.chars().take_while(|c| *c != '`').collect();

    p.cargo(&format!("report future-incompatibilities -Z future-incompat-report --id {}", id))
        .masquerade_as_nightly_cargo()
        .with_stdout_contains("The package `first-dep v0.0.1` currently triggers the following future incompatibility lints:")
        .with_stdout_contains("The package `second-dep v0.0.2` currently triggers the following future incompatibility lints:")
        .run();

    // Test without --id, and also the full output of the report.
    let output = p
        .cargo("report future-incompat -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .exec_with_output()
        .unwrap();
    let output = std::str::from_utf8(&output.stdout).unwrap();
    assert!(output.starts_with("The following warnings were discovered"));
    let mut lines = output
        .lines()
        // Skip the beginning of the per-package information.
        .skip_while(|line| !line.starts_with("The package"));
    for expected in &["first-dep v0.0.1", "second-dep v0.0.2"] {
        assert_eq!(
            &format!(
                "The package `{}` currently triggers the following future incompatibility lints:",
                expected
            ),
            lines.next().unwrap()
        );
        let mut count = 0;
        while let Some(line) = lines.next() {
            if line.is_empty() {
                break;
            }
            count += 1;
        }
        assert!(count > 0);
    }
    assert_eq!(lines.next(), None);
}

#[cargo_test]
#[ignore] // Waiting on https://github.com/rust-lang/rust/pull/86478
fn color() {
    if !is_nightly() {
        return;
    }

    let p = simple_project();

    p.cargo("check -Zfuture-incompat-report")
        .masquerade_as_nightly_cargo()
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stdout_does_not_contain("[..]\x1b[[..]")
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .env("CARGO_TERM_COLOR", "always")
        .with_stdout_contains("[..]\x1b[[..]")
        .run();
}

#[cargo_test]
#[ignore] // Waiting on https://github.com/rust-lang/rust/pull/86478
fn bad_ids() {
    if !is_nightly() {
        return;
    }

    let p = simple_project();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id 1")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr("error: no reports are currently available")
        .run();

    p.cargo("check -Zfuture-incompat-report")
        .masquerade_as_nightly_cargo()
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id foo")
        .masquerade_as_nightly_cargo()
        .with_status(1)
        .with_stderr("error: Invalid value: could not parse `foo` as a number")
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id 7")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: could not find report with ID 7
Available IDs are: 1
",
        )
        .run();
}

#[cargo_test]
#[ignore] // Waiting on https://github.com/rust-lang/rust/pull/86478
fn suggestions_for_updates() {
    if !is_nightly() {
        return;
    }

    Package::new("with_updates", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("big_update", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("without_updates", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.1.0"

                [dependencies]
                with_updates = "1"
                big_update = "1"
                without_updates = "1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("generate-lockfile").run();

    Package::new("with_updates", "1.0.1")
        .file("src/lib.rs", "")
        .publish();
    Package::new("with_updates", "1.0.2")
        .file("src/lib.rs", "")
        .publish();
    Package::new("big_update", "2.0.0")
        .file("src/lib.rs", "")
        .publish();

    // This is a hack to force cargo to update the index. Cargo can't do this
    // automatically because doing a network update on every build would be a
    // bad idea. Under normal circumstances, we'll hope the user has done
    // something else along the way to trigger an update (building some other
    // project or something). This could use some more consideration of how to
    // handle this better (maybe only trigger an update if it hasn't updated
    // in a long while?).
    p.cargo("update -p without_updates").run();

    p.cargo("check -Zfuture-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains("[..]cargo report future-incompatibilities --id 1[..]")
        .run();

    p.cargo("report future-incompatibilities")
        .masquerade_as_nightly_cargo()
        .with_stdout_contains(
            "\
The following packages appear to have newer versions available.
You may want to consider updating them to a newer version to see if the issue has been fixed.

big_update v1.0.0 has the following newer versions available: 2.0.0
with_updates v1.0.0 has the following newer versions available: 1.0.1, 1.0.2
",
        )
        .run();
}
<commit_msg>Re-enable future-incompatible tests.<commit_after>//! Tests for future-incompat-report messages
//!
//! Note that these tests use the -Zfuture-incompat-test for rustc.
//! This causes rustc to treat *every* lint as future-incompatible.
//! This is done because future-incompatible lints are inherently
//! ephemeral, but we don't want to continually update these tests.
//! So we pick some random lint that will likely always be the same
//! over time.

use cargo_test_support::registry::Package;
use cargo_test_support::{basic_manifest, is_nightly, project, Project};

// An arbitrary lint (unused_variables) that triggers a lint.
// We use a special flag to force it to generate a report.
const FUTURE_EXAMPLE: &'static str = "fn main() { let x = 1; }";
// Some text that will be displayed when the lint fires.
const FUTURE_OUTPUT: &'static str = "[..]unused_variables[..]";

fn simple_project() -> Project {
    project()
        .file("Cargo.toml", &basic_manifest("foo", "0.0.0"))
        .file("src/main.rs", FUTURE_EXAMPLE)
        .build()
}

#[cargo_test]
fn no_output_on_stable() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }
    let p = simple_project();

    p.cargo("check")
        // Even though rustc emits the report, cargo ignores it without -Z.
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .with_stderr_contains(FUTURE_OUTPUT)
        .with_stderr_does_not_contain("[..]cargo report[..]")
        .run();
}

#[cargo_test]
fn gate_future_incompat_report() {
    let p = simple_project();

    p.cargo("build --future-incompat-report")
        .with_stderr_contains("error: the `--future-incompat-report` flag is unstable[..]")
        .with_status(101)
        .run();

    // Both `-Z future-incompat-report` and `-Z unstable-opts` are required
    p.cargo("build --future-incompat-report -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains("error: the `--future-incompat-report` flag is unstable[..]")
        .with_status(101)
        .run();

    p.cargo("build --future-incompat-report -Z unstable-options")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains(
            "error: Usage of `--future-incompat-report` requires `-Z future-incompat-report`",
        )
        .with_status(101)
        .run();

    p.cargo("report future-incompatibilities --id foo")
        .with_stderr_contains("error: `cargo report` can only be used on the nightly channel")
        .with_status(101)
        .run();
}

#[cargo_test]
fn test_zero_future_incompat() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    let p = project()
        .file("Cargo.toml", &basic_manifest("foo", "0.0.0"))
        .file("src/main.rs", "fn main() {}")
        .build();

    // No note if --future-incompat-report is not specified.
    p.cargo("build -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .with_stderr(
            "\
[COMPILING] foo v0.0.0 [..]
[FINISHED] [..]
",
        )
        .run();

    p.cargo("build --future-incompat-report -Z unstable-options -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .with_stderr(
            "\
[FINISHED] [..]
note: 0 dependencies had future-incompatible warnings
",
        )
        .run();
}

#[cargo_test]
fn test_single_crate() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    let p = simple_project();

    for command in &["build", "check", "rustc", "test"] {
        p.cargo(command).arg("-Zfuture-incompat-report")
            .masquerade_as_nightly_cargo()
            .env("RUSTFLAGS", "-Zfuture-incompat-test")
            .with_stderr_contains(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: foo v0.0.0 [..]")
            .with_stderr_does_not_contain("[..]incompatibility[..]")
            .run();

        p.cargo(command).arg("-Zfuture-incompat-report").arg("-Zunstable-options").arg("--future-incompat-report")
            .masquerade_as_nightly_cargo()
            .env("RUSTFLAGS", "-Zfuture-incompat-test")
            .with_stderr_contains(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: foo v0.0.0 [..]")
            .with_stderr_contains("The package `foo v0.0.0 ([..])` currently triggers the following future incompatibility lints:")
            .run();
    }
}

#[cargo_test]
fn test_multi_crate() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    Package::new("first-dep", "0.0.1")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("second-dep", "0.0.2")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.0.0"

                [dependencies]
                first-dep = "*"
                second-dep = "*"
              "#,
        )
        .file("src/main.rs", "fn main() {}")
        .build();

    for command in &["build", "check", "rustc", "test"] {
        p.cargo(command).arg("-Zfuture-incompat-report")
            .masquerade_as_nightly_cargo()
            .env("RUSTFLAGS", "-Zfuture-incompat-test")
            .with_stderr_does_not_contain(FUTURE_OUTPUT)
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: first-dep v0.0.1, second-dep v0.0.2")
            // Check that we don't have the 'triggers' message shown at the bottom of this loop
            .with_stderr_does_not_contain("[..]triggers[..]")
            .run();

        p.cargo(command).arg("-Zunstable-options").arg("-Zfuture-incompat-report").arg("--future-incompat-report")
            .masquerade_as_nightly_cargo()
            .env("RUSTFLAGS", "-Zfuture-incompat-test")
            .with_stderr_contains("warning: the following packages contain code that will be rejected by a future version of Rust: first-dep v0.0.1, second-dep v0.0.2")
            .with_stderr_contains("The package `first-dep v0.0.1` currently triggers the following future incompatibility lints:")
            .with_stderr_contains("The package `second-dep v0.0.2` currently triggers the following future incompatibility lints:")
            .run();
    }

    // Test that passing the correct id via '--id' doesn't generate a warning message
    let output = p
        .cargo("build -Z future-incompat-report")
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .masquerade_as_nightly_cargo()
        .exec_with_output()
        .unwrap();

    // Extract the 'id' from the stdout. We are looking
    // for the id in a line of the form "run `cargo report future-incompatibilities --id yZ7S`"
    // which is generated by Cargo to tell the user what command to run
    // This is just to test that passing the id suppresses the warning mesasge. Any users needing
    // access to the report from a shell script should use the `--future-incompat-report` flag
    let stderr = std::str::from_utf8(&output.stderr).unwrap();

    // Find '--id <ID>' in the output
    let mut iter = stderr.split(' ');
    iter.find(|w| *w == "--id").unwrap();
    let id = iter
        .next()
        .unwrap_or_else(|| panic!("Unexpected output:\n{}", stderr));
    // Strip off the trailing '`' included in the output
    let id: String = id.chars().take_while(|c| *c != '`').collect();

    p.cargo(&format!("report future-incompatibilities -Z future-incompat-report --id {}", id))
        .masquerade_as_nightly_cargo()
        .with_stdout_contains("The package `first-dep v0.0.1` currently triggers the following future incompatibility lints:")
        .with_stdout_contains("The package `second-dep v0.0.2` currently triggers the following future incompatibility lints:")
        .run();

    // Test without --id, and also the full output of the report.
    let output = p
        .cargo("report future-incompat -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .exec_with_output()
        .unwrap();
    let output = std::str::from_utf8(&output.stdout).unwrap();
    assert!(output.starts_with("The following warnings were discovered"));
    let mut lines = output
        .lines()
        // Skip the beginning of the per-package information.
        .skip_while(|line| !line.starts_with("The package"));
    for expected in &["first-dep v0.0.1", "second-dep v0.0.2"] {
        assert_eq!(
            &format!(
                "The package `{}` currently triggers the following future incompatibility lints:",
                expected
            ),
            lines.next().unwrap()
        );
        let mut count = 0;
        while let Some(line) = lines.next() {
            if line.is_empty() {
                break;
            }
            count += 1;
        }
        assert!(count > 0);
    }
    assert_eq!(lines.next(), None);
}

#[cargo_test]
fn color() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    let p = simple_project();

    p.cargo("check -Zfuture-incompat-report")
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .masquerade_as_nightly_cargo()
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .with_stdout_does_not_contain("[..]\x1b[[..]")
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report")
        .masquerade_as_nightly_cargo()
        .env("CARGO_TERM_COLOR", "always")
        .with_stdout_contains("[..]\x1b[[..]")
        .run();
}

#[cargo_test]
fn bad_ids() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    let p = simple_project();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id 1")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr("error: no reports are currently available")
        .run();

    p.cargo("check -Zfuture-incompat-report")
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .masquerade_as_nightly_cargo()
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id foo")
        .masquerade_as_nightly_cargo()
        .with_status(1)
        .with_stderr("error: Invalid value: could not parse `foo` as a number")
        .run();

    p.cargo("report future-incompatibilities -Z future-incompat-report --id 7")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: could not find report with ID 7
Available IDs are: 1
",
        )
        .run();
}

#[cargo_test]
fn suggestions_for_updates() {
    if !is_nightly() {
        // -Zfuture-incompat-test requires nightly (permanently)
        return;
    }

    Package::new("with_updates", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("big_update", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();
    Package::new("without_updates", "1.0.0")
        .file("src/lib.rs", FUTURE_EXAMPLE)
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.1.0"

                [dependencies]
                with_updates = "1"
                big_update = "1"
                without_updates = "1"
            "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("generate-lockfile").run();

    Package::new("with_updates", "1.0.1")
        .file("src/lib.rs", "")
        .publish();
    Package::new("with_updates", "1.0.2")
        .file("src/lib.rs", "")
        .publish();
    Package::new("big_update", "2.0.0")
        .file("src/lib.rs", "")
        .publish();

    // This is a hack to force cargo to update the index. Cargo can't do this
    // automatically because doing a network update on every build would be a
    // bad idea. Under normal circumstances, we'll hope the user has done
    // something else along the way to trigger an update (building some other
    // project or something). This could use some more consideration of how to
    // handle this better (maybe only trigger an update if it hasn't updated
    // in a long while?).
    p.cargo("update -p without_updates").run();

    p.cargo("check -Zfuture-incompat-report")
        .masquerade_as_nightly_cargo()
        .env("RUSTFLAGS", "-Zfuture-incompat-test")
        .with_stderr_contains("[..]cargo report future-incompatibilities --id 1[..]")
        .run();

    p.cargo("report future-incompatibilities")
        .masquerade_as_nightly_cargo()
        .with_stdout_contains(
            "\
The following packages appear to have newer versions available.
You may want to consider updating them to a newer version to see if the issue has been fixed.

big_update v1.0.0 has the following newer versions available: 2.0.0
with_updates v1.0.0 has the following newer versions available: 1.0.1, 1.0.2
",
        )
        .run();
}
<|endoftext|>


<commit_before>extern crate evalbotlib as backend;
extern crate serde;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate log;
extern crate toml;
extern crate futures;
extern crate tokio;
extern crate telebot;
extern crate env_logger;

use backend::{EvalService, Language, util};

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use std::borrow::Cow;

use futures::{Future, Stream, IntoFuture};
use telebot::RcBot;
use telebot::functions::*;
use telebot::objects::*;
use futures::future::Either;

macro_rules! nullify_future {
    ($task:expr, $fut:expr) => ($fut.map(|_| ())
        .or_else(|e| Ok(error!("error {}: {}", $task, e))));
}

static WHITELIST_FILENAME: &'static str = "tgwhitelist.toml";

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgCfg {
    owners: HashSet<String>,
    msg_owner_id: Option<i64>,
    bot_id: String,
    lang_subst: HashMap<String, String>
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgWhitelist {
    priv_enabled: bool,
    group_enabled: bool,
    allowed: HashSet<i64>,
    blocked: HashSet<i64>
}

impl TgWhitelist {
    fn priv_ok(&self, id: i64) -> bool {
        (!self.priv_enabled || self.allowed.contains(&id)) && !self.blocked.contains(&id)
    }

    fn group_ok(&self, id: i64) -> bool {
        (!self.group_enabled || self.allowed.contains(&id)) && !self.blocked.contains(&id)
    }

    fn allow(&mut self, id: i64) {
        self.allowed.insert(id);
    }

    fn unallow(&mut self, id: i64) {
        self.allowed.remove(&id);
    }

    fn block(&mut self, id: i64) {
        self.blocked.insert(id);
    }

    fn unblock(&mut self, id: i64) {
        self.blocked.remove(&id);
    }

    fn save<'a, 'b>(&'a self, path: &'static str) -> impl Future<Item = (), Error = ()> + 'b {
        nullify_future!("saving whitelist", util::encode(self, path))
    }
}

struct TgSvc {
    config: TgCfg,
    whitelist: RwLock<TgWhitelist>,
    service: EvalService,
}

fn telegram_wrap_result(s: &str, group: bool) -> String {
    // FIXME configurable max-lines and max-bytes
    if s.is_empty() {
        "no output".to_owned()
    } else {
        let mut r = "
".to_owned();
        let input = s.as_bytes();
        let mut cut_input = String::from_utf8_lossy(&input[..512.min(input.len())]);
        if group {
            cut_input = Cow::Owned(cut_input.lines().take(10).collect::<Vec<_>>().join("\n"));
        }
        r.push_str(&cut_input
            .replace(|c: char| c =='\u{FFFD}' || (c.is_control() && c != '\n' && c != '\t'), "")
            .replace('&', "&")
            .replace('<', "<")
            .replace('>', ">")
            .replace('"', """));
        r.push_str("
");
        if cut_input.len() + 1 // we also cut off the trailing \n
            < input.len() {
            r.push_str("... (truncated)");
        }
        r
    }
}

fn is_from_owner(msg: &Message, tgsvc: &TgSvc) -> bool {
    msg.from.as_ref().and_then(|u| u.username.as_ref().map(|un| tgsvc.is_owner(un))).unwrap_or(false)
}

fn handle_update((tgbot, update): (RcBot, Update), tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    if let Some((chat, Some(new_user))) = update.message.map(|m| (m.chat, m.new_chat_member)) {
        let chat_id = chat.id;
        if let (Ok(id), Ok(wl)) = (tgbot.inner.id.read(), tgsvc.whitelist.read()) {
            if *id == Some(new_user.id) && !wl.group_ok(chat_id) {
                tokio::spawn(nullify_future!("leaving group",
                    tgbot.message(chat_id, format!("You or this group is not on the whitelist. Seek help. ID: {}", chat_id)).send()
                        .and_then(move |(tgbot, _)| tgbot.leave_chat(chat_id).send())));
            }
        }
    }
    Ok(())
}

fn handle_eval(tgsvc: &Arc<TgSvc>, tgbot: RcBot, msg: Message, lang: &Arc<Language>, is_hash: bool)
    -> impl Future<Item = (), Error = ()> {
    let chat_id = msg.chat.id;
    let group = msg.chat.kind != "private";
    if let Ok(wl) = tgsvc.whitelist.read() {
        if group && !wl.group_ok(chat_id)
            || !group && !wl.priv_ok(chat_id) {
            tokio::spawn(nullify_future!("sending message",
                tgbot.message(chat_id, format!("You or this group is not on the whitelist. Seek help. ID: {}", chat_id)).send()));
            return Ok(()).into_future();
        }
    } else {
        error!("Failed to acquire RwLock");
        tokio::spawn(nullify_future!("sending message",
                tgbot.message(chat_id, "Internal error occurred".to_owned()).send()));
        return Ok(()).into_future();
    }

    let no_limit = is_hash && is_from_owner(&msg, tgsvc);
    let msg_id = msg.message_id;
    info!("({}) evaluating from {:?}: {:?}", msg_id, msg.from, msg.text.as_ref().map(|x| x.as_str()).unwrap_or(""));
    let code = msg.text.map(|x| {
        let mut r = x.trim_left().to_owned();
        r.push_str("\n");
        r
    }).unwrap_or_else(|| "".to_owned());
    tokio::spawn(nullify_future!("sending message",
        lang.eval(code, if no_limit { Some(0) } else { None }, Some(format!("tg{}", chat_id)))
            .then(move |e| {
                info!("({}) result: {:?}", msg_id, e.as_ref().map(|x| x.as_str()).unwrap_or(""));
                tgbot.message(chat_id, telegram_wrap_result(&e.unwrap_or_else(|x| x), group))
                    .parse_mode(ParseMode::HTML)
                    .reply_to_message_id(msg_id)
                    .send()
            })));

    Ok(()).into_future()
}

#[derive(Clone, Copy)]
enum WhitelistToggleOp {
    TogglePrivate,
    ToggleGroup
}

#[derive(Clone, Copy)]
enum WhitelistModOp {
    Allow,
    Unallow,
    Block,
    Unblock
}

fn handle_whitelist_toggle(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message, op: WhitelistToggleOp)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let resp = match me.whitelist.write() {
        Ok(mut wl) => match op {
            WhitelistToggleOp::TogglePrivate => {
                wl.priv_enabled = !wl.priv_enabled;
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Private whitelist enabled: {}", wl.priv_enabled)
            }
            WhitelistToggleOp::ToggleGroup => {
                wl.group_enabled = !wl.group_enabled;
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Group whitelist enabled: {}", wl.group_enabled)
            }
        }
        Err(err) => {
            error!("error while acquiring RwLock: {}", err);
            "error acquiring RwLock".to_owned()
        }
    };
    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

fn handle_whitelist_mod(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message, op: WhitelistModOp)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let arg = msg.text.as_ref().and_then(|t| t.trim().split_whitespace()
        .nth(0).and_then(|arg| arg.parse().ok()));
    let resp = match (arg, me.whitelist.write()) {
        (Some(id), Ok(mut wl)) => match op {
            WhitelistModOp::Allow => {
                wl.allow(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Allowed {}", id)
            }
            WhitelistModOp::Unallow => {
                wl.unallow(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Unallowed {}", id)
            }
            WhitelistModOp::Block => {
                wl.block(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Blocked {}", id)
            }
            WhitelistModOp::Unblock => {
                wl.unblock(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Unblocked {}", id)
            }
        }
        (None, _) => "Invalid ID".to_owned(),
        (_, Err(err)) => {
            error!("error while acquiring RwLock: {}", err);
            "error acquiring RwLock".to_owned()
        }
    };

    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

fn handle_leave(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let arg = msg.text.as_ref().and_then(|t| t.trim().split_whitespace()
        .nth(0).and_then(|arg| arg.parse().ok()));
    let resp = match arg {
        Some(id) => {
            tokio::spawn(nullify_future!("leaving group", tgbot.leave_chat(id).send()));
            "OK"
        }
        None => "Invalid ID"
    }.to_owned();

    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

impl TgSvc {
    fn run() -> impl Future<Item = (), Error = ()> {
        let cfgf = util::decode::<TgCfg, _>("evalbot.tg.toml")
            .map(|cfg| {
                debug!("Loaded config: {:?}", cfg);
                cfg
            })
            .map_err(|e| {
                error!("failed to read evalbot.tg.toml: {}", e);
            });
        let wlf = util::decode::<TgWhitelist, _>(WHITELIST_FILENAME).or_else(|e| {
            warn!("failed to read whitelist: {}; using empty whitelist", e);
            Ok(TgWhitelist {
                priv_enabled: false,
                group_enabled: false,
                allowed: HashSet::new(),
                blocked: HashSet::new(),
            }).into_future()
        });
        cfgf.join(wlf).join(EvalService::from_toml_file("evalbot.toml")
            .map_err(|e| {
                error!("failed to read evalbot.toml: {}", e);
            }))
            .map(|((cfg, wl), es)| TgSvc {
                config: cfg,
                whitelist: RwLock::new(wl),
                service: es
            })
            .and_then(TgSvc::handle)
    }

    fn handle(self) -> impl Future<Item = (), Error = ()> {
        let bot = RcBot::new(&self.config.bot_id).expect("Failed to initialise Telegram bot");
        let me = Arc::new(self);
        bot.resolve_name();

        macro_rules! handle {
            ($cmd:expr, $handler:ident $(,$params:expr)*) => {{
                let me = me.clone();
                bot.register(bot.new_cmd($cmd)
                    .map_err(|e| error!("error in command processing: {}", e))
                    .and_then(move |(tgbot, msg)| $handler(&me, tgbot, msg $(,$params)*)));
            }};
        }

        for (name, lang) in me.service.langs() {
            {
                let lang = lang.clone();
                handle!(name, handle_eval, &lang, false);
            }
            {
                let lang = lang.clone();
                handle!(&format!("{}#", name), handle_eval, &lang, true);
            }
        }

        handle!("privwl", handle_whitelist_toggle, WhitelistToggleOp::TogglePrivate);
        handle!("groupwl", handle_whitelist_toggle, WhitelistToggleOp::ToggleGroup);
        handle!("allow", handle_whitelist_mod, WhitelistModOp::Allow);
        handle!("unallow", handle_whitelist_mod, WhitelistModOp::Unallow);
        handle!("block", handle_whitelist_mod, WhitelistModOp::Block);
        handle!("unblock", handle_whitelist_mod, WhitelistModOp::Unblock);
        handle!("leave", handle_leave);

        bot.get_stream()
            .map_err(|e| error!("{}", e))
            .for_each(move |tuple| handle_update(tuple, &me))
            .into_future()
    }

    fn is_owner(&self, name: &str) -> bool {
        self.config.owners.contains(name)
    }
}

fn main() {
    env_logger::init();
    tokio::run(TgSvc::run());
}
<commit_msg>tgbot: Don't need to wrap result for Errors<commit_after>extern crate evalbotlib as backend;
extern crate serde;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate log;
extern crate toml;
extern crate futures;
extern crate tokio;
extern crate telebot;
extern crate env_logger;

use backend::{EvalService, Language, util};

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use std::borrow::Cow;

use futures::{Future, Stream, IntoFuture};
use telebot::RcBot;
use telebot::functions::*;
use telebot::objects::*;
use futures::future::Either;

macro_rules! nullify_future {
    ($task:expr, $fut:expr) => ($fut.map(|_| ())
        .or_else(|e| Ok(error!("error {}: {}", $task, e))));
}

static WHITELIST_FILENAME: &'static str = "tgwhitelist.toml";

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgCfg {
    owners: HashSet<String>,
    msg_owner_id: Option<i64>,
    bot_id: String,
    lang_subst: HashMap<String, String>
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgWhitelist {
    priv_enabled: bool,
    group_enabled: bool,
    allowed: HashSet<i64>,
    blocked: HashSet<i64>
}

impl TgWhitelist {
    fn priv_ok(&self, id: i64) -> bool {
        (!self.priv_enabled || self.allowed.contains(&id)) && !self.blocked.contains(&id)
    }

    fn group_ok(&self, id: i64) -> bool {
        (!self.group_enabled || self.allowed.contains(&id)) && !self.blocked.contains(&id)
    }

    fn allow(&mut self, id: i64) {
        self.allowed.insert(id);
    }

    fn unallow(&mut self, id: i64) {
        self.allowed.remove(&id);
    }

    fn block(&mut self, id: i64) {
        self.blocked.insert(id);
    }

    fn unblock(&mut self, id: i64) {
        self.blocked.remove(&id);
    }

    fn save<'a, 'b>(&'a self, path: &'static str) -> impl Future<Item = (), Error = ()> + 'b {
        nullify_future!("saving whitelist", util::encode(self, path))
    }
}

struct TgSvc {
    config: TgCfg,
    whitelist: RwLock<TgWhitelist>,
    service: EvalService,
}

fn telegram_wrap_result(s: &str, group: bool) -> String {
    // FIXME configurable max-lines and max-bytes
    if s.is_empty() {
        "no output".to_owned()
    } else {
        let mut r = "
".to_owned();
        let input = s.as_bytes();
        let mut cut_input = String::from_utf8_lossy(&input[..512.min(input.len())]);
        if group {
            cut_input = Cow::Owned(cut_input.lines().take(10).collect::<Vec<_>>().join("\n"));
        }
        r.push_str(&cut_input
            .replace(|c: char| c =='\u{FFFD}' || (c.is_control() && c != '\n' && c != '\t'), "")
            .replace('&', "&")
            .replace('<', "<")
            .replace('>', ">")
            .replace('"', """));
        r.push_str("
");
        if cut_input.len() + 1 // we also cut off the trailing \n
            < input.len() {
            r.push_str("... (truncated)");
        }
        r
    }
}

fn is_from_owner(msg: &Message, tgsvc: &TgSvc) -> bool {
    msg.from.as_ref().and_then(|u| u.username.as_ref().map(|un| tgsvc.is_owner(un))).unwrap_or(false)
}

fn handle_update((tgbot, update): (RcBot, Update), tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    if let Some((chat, Some(new_user))) = update.message.map(|m| (m.chat, m.new_chat_member)) {
        let chat_id = chat.id;
        if let (Ok(id), Ok(wl)) = (tgbot.inner.id.read(), tgsvc.whitelist.read()) {
            if *id == Some(new_user.id) && !wl.group_ok(chat_id) {
                tokio::spawn(nullify_future!("leaving group",
                    tgbot.message(chat_id, format!("You or this group is not on the whitelist. Seek help. ID: {}", chat_id)).send()
                        .and_then(move |(tgbot, _)| tgbot.leave_chat(chat_id).send())));
            }
        }
    }
    Ok(())
}

fn handle_eval(tgsvc: &Arc<TgSvc>, tgbot: RcBot, msg: Message, lang: &Arc<Language>, is_hash: bool)
    -> impl Future<Item = (), Error = ()> {
    let chat_id = msg.chat.id;
    let group = msg.chat.kind != "private";
    if let Ok(wl) = tgsvc.whitelist.read() {
        if group && !wl.group_ok(chat_id)
            || !group && !wl.priv_ok(chat_id) {
            tokio::spawn(nullify_future!("sending message",
                tgbot.message(chat_id, format!("You or this group is not on the whitelist. Seek help. ID: {}", chat_id)).send()));
            return Ok(()).into_future();
        }
    } else {
        error!("Failed to acquire RwLock");
        tokio::spawn(nullify_future!("sending message",
                tgbot.message(chat_id, "Internal error occurred".to_owned()).send()));
        return Ok(()).into_future();
    }

    let no_limit = is_hash && is_from_owner(&msg, tgsvc);
    let msg_id = msg.message_id;
    info!("({}) evaluating from {:?}: {:?}", msg_id, msg.from, msg.text.as_ref().map(|x| x.as_str()).unwrap_or(""));
    let code = msg.text.map(|x| {
        let mut r = x.trim_left().to_owned();
        r.push_str("\n");
        r
    }).unwrap_or_else(|| "".to_owned());
    tokio::spawn(nullify_future!("sending message",
        lang.eval(code, if no_limit { Some(0) } else { None }, Some(format!("tg{}", chat_id)))
            .then(move |e| {
                let ok = e.is_ok();
                info!("({}) result: {:?}", msg_id, e.as_ref().map(|x| x.as_str()).unwrap_or(""));
                let msg = tgbot.message(chat_id, e.map(|r| telegram_wrap_result(&r, group)).unwrap_or_else(|e| e));
                if ok {
                    msg.parse_mode(ParseMode::HTML)
                } else {
                    msg
                }.reply_to_message_id(msg_id).send()
            })));

    Ok(()).into_future()
}

#[derive(Clone, Copy)]
enum WhitelistToggleOp {
    TogglePrivate,
    ToggleGroup
}

#[derive(Clone, Copy)]
enum WhitelistModOp {
    Allow,
    Unallow,
    Block,
    Unblock
}

fn handle_whitelist_toggle(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message, op: WhitelistToggleOp)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let resp = match me.whitelist.write() {
        Ok(mut wl) => match op {
            WhitelistToggleOp::TogglePrivate => {
                wl.priv_enabled = !wl.priv_enabled;
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Private whitelist enabled: {}", wl.priv_enabled)
            }
            WhitelistToggleOp::ToggleGroup => {
                wl.group_enabled = !wl.group_enabled;
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Group whitelist enabled: {}", wl.group_enabled)
            }
        }
        Err(err) => {
            error!("error while acquiring RwLock: {}", err);
            "error acquiring RwLock".to_owned()
        }
    };
    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

fn handle_whitelist_mod(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message, op: WhitelistModOp)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let arg = msg.text.as_ref().and_then(|t| t.trim().split_whitespace()
        .nth(0).and_then(|arg| arg.parse().ok()));
    let resp = match (arg, me.whitelist.write()) {
        (Some(id), Ok(mut wl)) => match op {
            WhitelistModOp::Allow => {
                wl.allow(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Allowed {}", id)
            }
            WhitelistModOp::Unallow => {
                wl.unallow(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Unallowed {}", id)
            }
            WhitelistModOp::Block => {
                wl.block(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Blocked {}", id)
            }
            WhitelistModOp::Unblock => {
                wl.unblock(id);
                tokio::spawn(wl.save(WHITELIST_FILENAME));
                format!("Unblocked {}", id)
            }
        }
        (None, _) => "Invalid ID".to_owned(),
        (_, Err(err)) => {
            error!("error while acquiring RwLock: {}", err);
            "error acquiring RwLock".to_owned()
        }
    };

    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

fn handle_leave(me: &Arc<TgSvc>, tgbot: RcBot, msg: Message)
    -> impl Future<Item = (), Error = ()> {
    if !is_from_owner(&msg, me) {
        return Either::A(Ok(()).into_future());
    }

    let arg = msg.text.as_ref().and_then(|t| t.trim().split_whitespace()
        .nth(0).and_then(|arg| arg.parse().ok()));
    let resp = match arg {
        Some(id) => {
            tokio::spawn(nullify_future!("leaving group", tgbot.leave_chat(id).send()));
            "OK"
        }
        None => "Invalid ID"
    }.to_owned();

    Either::B(nullify_future!("sending message", tgbot.message(msg.chat.id, resp)
        .reply_to_message_id(msg.message_id)
        .send()))
}

impl TgSvc {
    fn run() -> impl Future<Item = (), Error = ()> {
        let cfgf = util::decode::<TgCfg, _>("evalbot.tg.toml")
            .map(|cfg| {
                debug!("Loaded config: {:?}", cfg);
                cfg
            })
            .map_err(|e| {
                error!("failed to read evalbot.tg.toml: {}", e);
            });
        let wlf = util::decode::<TgWhitelist, _>(WHITELIST_FILENAME).or_else(|e| {
            warn!("failed to read whitelist: {}; using empty whitelist", e);
            Ok(TgWhitelist {
                priv_enabled: false,
                group_enabled: false,
                allowed: HashSet::new(),
                blocked: HashSet::new(),
            }).into_future()
        });
        cfgf.join(wlf).join(EvalService::from_toml_file("evalbot.toml")
            .map_err(|e| {
                error!("failed to read evalbot.toml: {}", e);
            }))
            .map(|((cfg, wl), es)| TgSvc {
                config: cfg,
                whitelist: RwLock::new(wl),
                service: es
            })
            .and_then(TgSvc::handle)
    }

    fn handle(self) -> impl Future<Item = (), Error = ()> {
        let bot = RcBot::new(&self.config.bot_id).expect("Failed to initialise Telegram bot");
        let me = Arc::new(self);
        bot.resolve_name();

        macro_rules! handle {
            ($cmd:expr, $handler:ident $(,$params:expr)*) => {{
                let me = me.clone();
                bot.register(bot.new_cmd($cmd)
                    .map_err(|e| error!("error in command processing: {}", e))
                    .and_then(move |(tgbot, msg)| $handler(&me, tgbot, msg $(,$params)*)));
            }};
        }

        for (name, lang) in me.service.langs() {
            {
                let lang = lang.clone();
                handle!(name, handle_eval, &lang, false);
            }
            {
                let lang = lang.clone();
                handle!(&format!("{}#", name), handle_eval, &lang, true);
            }
        }

        handle!("privwl", handle_whitelist_toggle, WhitelistToggleOp::TogglePrivate);
        handle!("groupwl", handle_whitelist_toggle, WhitelistToggleOp::ToggleGroup);
        handle!("allow", handle_whitelist_mod, WhitelistModOp::Allow);
        handle!("unallow", handle_whitelist_mod, WhitelistModOp::Unallow);
        handle!("block", handle_whitelist_mod, WhitelistModOp::Block);
        handle!("unblock", handle_whitelist_mod, WhitelistModOp::Unblock);
        handle!("leave", handle_leave);

        bot.get_stream()
            .map_err(|e| error!("{}", e))
            .for_each(move |tuple| handle_update(tuple, &me))
            .into_future()
    }

    fn is_owner(&self, name: &str) -> bool {
        self.config.owners.contains(name)
    }
}

fn main() {
    env_logger::init();
    tokio::run(TgSvc::run());
}
<|endoftext|>


<commit_before>use std::io;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::cmp::min;

pub use mio::{EventSet, PollOpt, Token};
use mio::Timeout as MioTimeout;
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;

use super::connection::{Connection, Direction};
use super::{Timeout, InternalTimeout, EventLoop};

use super::super::messages::{MessageBuffer, RawMessage};

pub type PeerId = Token;

const SERVER_ID: PeerId = Token(1);

#[derive(Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    pub listen_address: SocketAddr,
    pub max_connections: usize, // TODO: think more about config parameters
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u32>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

pub struct Network {
    config: NetworkConfiguration,
    listener: Option<TcpListener>,
    connections: Slab<Connection>,
    addresses: HashMap<SocketAddr, PeerId>,
    reconnects: HashMap<SocketAddr, MioTimeout>,
}

impl Network {
    pub fn with_config(config: NetworkConfiguration) -> Network {
        Network {
            config: config,
            listener: None,
            connections: Slab::new_starting_at(Token(2), config.max_connections),
            addresses: HashMap::new(),
            reconnects: HashMap::new(),
        }
    }

    pub fn address(&self) -> &SocketAddr {
        &self.config.listen_address
    }

    pub fn bind(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
        if let Some(_) = self.listener {
            return Err(io::Error::new(io::ErrorKind::Other, "Already binded"));
        }
        let listener = TcpListener::bind(&self.config.listen_address)?;
        let r = event_loop.register(&listener, SERVER_ID, EventSet::readable(), PollOpt::edge());
        match r {
            Ok(()) => {
                self.listener = Some(listener);
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    // TODO: Use ticks for fast reregistering sockets
    // TODO: Implement Connections collection with (re)registering

    pub fn io(&mut self,
              event_loop: &mut EventLoop,
              id: PeerId,
              set: EventSet)
              -> io::Result<Option<MessageBuffer>> {
        if id == SERVER_ID {
            // Accept new connections
            // FIXME: Fail-safe accepting of new connections?
            let pair = match self.listener {
                Some(ref listener) => listener.accept()?,
                None => None,
            };
            if let Some((socket, address)) = pair {
                let peer = Connection::new(socket, address, Direction::Incoming);
                self.add_connection(event_loop, peer)?;

                trace!("{}: Accepted incoming connection from {} id: {}",
                       self.address(),
                       address,
                       id.0);

                // TODO send event to node
            }
            return Ok(None);
        }

        if set.is_error() {
            trace!("{}: connection with {} closed with error",
                  self.address(),
                  self.connections[id].address());

            self.try_reconnect(event_loop, id);
            return Ok(None);
        }

        if set.is_hup() {
            trace!("{}: connection with addr {} closed",
                   self.address(),
                   self.connections[id].address());

            self.try_reconnect(event_loop, id); // for debug only
            //self.remove_connection(id);
            return Ok(None);
        }

        if set.is_writable() {
            let address = *self.connections[id].address();

            trace!("{}: Socket is writable {} id: {}",
                   self.address(),
                   address,
                   id.0);

            // Write data into socket
            if let Err(e) = self.connections[id].try_write() {
                warn!("An error when write to socket {} occured, error: {}",
                      address,
                      e);
                self.try_reconnect(event_loop, id);
                return Ok(None);
            }
            self.reregister_connection(event_loop, id, PollOpt::edge())?;

            self.mark_connected(event_loop, id);
            return Ok(None);
        }

        if set.is_readable() {
            let address = *self.connections[id].address();
            self.mark_connected(event_loop, id);

            trace!("{}: Socket is readable {} id: {}",
                   self.address(),
                   address,
                   id.0);

            return match self.connections[id].try_read() {
                Ok(buffer) => Ok(buffer),
                Err(e) => {
                    warn!("An error when read from socket {} occured, error: {}",
                          address,
                          e);
                    self.try_reconnect(event_loop, id);
                    Ok(None)
                }
            };
        }
        // FIXME: Can we be here?
        Ok(None)
    }

    pub fn tick(&mut self, _: &mut EventLoop) {}

    pub fn get_peer(&mut self,
                    event_loop: &mut EventLoop,
                    address: &SocketAddr)
                    -> io::Result<PeerId> {
        if let Some(id) = self.addresses.get(address) {
            return Ok(*id);
        };

        let peer = Connection::new(TcpStream::connect(address)?, *address, Direction::Outgoing);
        let id = self.add_connection(event_loop, peer)?;

        trace!("{}: Establish connection with {}, id: {}",
               self.address(),
               address,
               id.0);
        Ok(id)
    }

    pub fn send_to(&mut self,
                   event_loop: &mut EventLoop,
                   address: &SocketAddr,
                   message: RawMessage) {
        let r = match self.get_peer(event_loop, address) {
            Ok(id) => {
                self.connections[id]
                    .send(message)
                    .and_then(|_| {
                        event_loop.reregister(self.connections[id].socket(),
                                        id,
                                        self.connections[id].interest(),
                                        PollOpt::edge())?;
                        self.mark_connected(event_loop, id);
                        Ok(())
                    })
                    .or_else(|e| {
                        self.try_reconnect(event_loop, id);
                        Err(e)
                    })
            }
            Err(e) => Err(e),
        };

        if let Err(e) = r {
            warn!("An error occured when try to send a message (address: {}, error: {:?})",
                  address,
                  e);
            return;
        }
    }

    pub fn handle_timeout(&mut self, event_loop: &mut EventLoop, timeout: InternalTimeout) {
        match timeout {
            InternalTimeout::Reconnect(addr, delay) => {
                if self.reconnects.contains_key(&addr) {
                    self.get_peer(event_loop, &addr);

                    let delay = min(delay * 2, self.config.tcp_reconnect_timeout_max);
                    // TODO Fail-safe timeout error handling
                    self.add_reconnect_timeout(event_loop, addr, delay).unwrap();

                    trace!("Try to reconnect with delay {}", delay);
                }
            }
        }
    }

    fn remove_connection(&mut self, id: Token) {
        let address = *self.connections[id].address();
        self.addresses.remove(&address);
        self.connections.remove(id);
    }

    fn try_reconnect(&mut self, event_loop: &mut EventLoop, id: Token) {
        if self.connections[id].direction() == &Direction::Outgoing {
            let addr = *self.connections[id].address();
            match self.try_reconnect_addr(event_loop, addr) {
                Err(e) => {
                    error!("{}: An error during reconnect occured: {:?}",
                           self.address(),
                           e);
                }
                Ok(_) => {}
            }
        }
        self.remove_connection(id);
    }

    fn try_reconnect_addr(&mut self,
                          event_loop: &mut EventLoop,
                          address: SocketAddr)
                          -> io::Result<()> {
        if !self.reconnects.contains_key(&address) {
            let delay = self.config.tcp_reconnect_timeout;
            return self.add_reconnect_timeout(event_loop, address, delay);
        }
        Ok(())
    }

    fn add_reconnect_timeout(&mut self,
                             event_loop: &mut EventLoop,
                             address: SocketAddr,
                             delay: u64)
                             -> io::Result<()> {
        let reconnect = Timeout::Internal(InternalTimeout::Reconnect(address, delay));
        let timeout = event_loop.timeout_ms(reconnect, delay)
            .map_err(|e| {
                io::Error::new(io::ErrorKind::Other, format!("A mio error occured {:?}", e))
            })?;
        self.reconnects.insert(address, timeout);

        trace!("Add reconnect timeout {}", delay);
        Ok(())
    }

    fn add_connection(&mut self,
                      event_loop: &mut EventLoop,
                      mut connection: Connection)
                      -> io::Result<PeerId> {
        self.configure_stream(connection.socket_mut())?;
        let address = *connection.address();
        let id = self.connections
            .insert(connection)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "Maximum connections"))?;
        self.addresses.insert(address, id);

        self.register_connection(event_loop, id, PollOpt::edge())?;
        Ok(id)
    }

    fn mark_connected(&mut self, event_loop: &mut EventLoop, id: Token) {
        let address = *self.connections[id].address();
        self.clear_reconnect_request(event_loop, &address)
    }

    fn clear_reconnect_request(&mut self, event_loop: &mut EventLoop, addr: &SocketAddr) {
        if let Some(timeout) = self.reconnects.remove(addr) {
            trace!("{}: Clear reconnect timeout for {}", self.address(), addr);
            event_loop.clear_timeout(timeout);
        }
    }

    fn register_connection(&mut self,
                           event_loop: &mut EventLoop,
                           id: Token,
                           opts: PollOpt)
                           -> io::Result<()> {
        event_loop.register(self.connections[id].socket(),
                      id,
                      self.connections[id].interest(),
                      opts)
            .or_else(|e| {
                self.remove_connection(id);
                Err(e)
            })
    }

    fn reregister_connection(&mut self,
                             event_loop: &mut EventLoop,
                             id: Token,
                             opts: PollOpt)
                             -> io::Result<()> {
        event_loop.reregister(self.connections[id].socket(),
                        id,
                        self.connections[id].interest(),
                        opts)
            .or_else(|e| {
                self.remove_connection(id);
                Err(e)
            })
    }

    fn configure_stream(&self, stream: &mut TcpStream) -> io::Result<()> {
        stream.set_keepalive(self.config.tcp_keep_alive)?;
        stream.set_nodelay(self.config.tcp_nodelay)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {}
<commit_msg>Do not send data to disconnected socket<commit_after>use std::io;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::cmp::min;

pub use mio::{EventSet, PollOpt, Token};
use mio::Timeout as MioTimeout;
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;

use super::connection::{Connection, Direction};
use super::{Timeout, InternalTimeout, EventLoop};

use super::super::messages::{MessageBuffer, RawMessage};

pub type PeerId = Token;

const SERVER_ID: PeerId = Token(1);
const RECONNECT_GROW_FACTOR: f32 = 2.0;

#[derive(Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    pub listen_address: SocketAddr,
    pub max_connections: usize, // TODO: think more about config parameters
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u32>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

pub struct Network {
    config: NetworkConfiguration,
    listener: Option<TcpListener>,
    connections: Slab<Connection>,
    addresses: HashMap<SocketAddr, PeerId>,
    reconnects: HashMap<SocketAddr, MioTimeout>,
}

impl Network {
    pub fn with_config(config: NetworkConfiguration) -> Network {
        Network {
            config: config,
            listener: None,
            connections: Slab::new_starting_at(Token(2), config.max_connections),
            addresses: HashMap::new(),
            reconnects: HashMap::new(),
        }
    }

    pub fn address(&self) -> &SocketAddr {
        &self.config.listen_address
    }

    pub fn bind(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
        if let Some(_) = self.listener {
            return Err(io::Error::new(io::ErrorKind::Other, "Already binded"));
        }
        let listener = TcpListener::bind(&self.config.listen_address)?;
        let r = event_loop.register(&listener, SERVER_ID, EventSet::readable(), PollOpt::edge());
        match r {
            Ok(()) => {
                self.listener = Some(listener);
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    // TODO: Use ticks for fast reregistering sockets
    // TODO: Implement Connections collection with (re)registering

    pub fn io(&mut self,
              event_loop: &mut EventLoop,
              id: PeerId,
              set: EventSet)
              -> io::Result<Option<MessageBuffer>> {
        if id == SERVER_ID {
            // Accept new connections
            // FIXME: Fail-safe accepting of new connections?
            let pair = match self.listener {
                Some(ref listener) => listener.accept()?,
                None => None,
            };
            if let Some((socket, address)) = pair {
                let peer = Connection::new(socket, address, Direction::Incoming);
                self.add_connection(event_loop, peer)?;

                trace!("{}: Accepted incoming connection from {} id: {}",
                       self.address(),
                       address,
                       id.0);

                // TODO send event to node
            }
            return Ok(None);
        }

        if set.is_error() {
            trace!("{}: connection with {} closed with error",
                  self.address(),
                  self.connections[id].address());

            self.try_reconnect(event_loop, id);
            return Ok(None);
        }

        if set.is_hup() {
            trace!("{}: connection with addr {} closed",
                   self.address(),
                   self.connections[id].address());

            self.try_reconnect(event_loop, id); // for debug only
            //self.remove_connection(id);
            return Ok(None);
        }

        if set.is_writable() {
            let address = *self.connections[id].address();

            trace!("{}: Socket is writable {} id: {}",
                   self.address(),
                   address,
                   id.0);

            // Write data into socket
            if let Err(e) = self.connections[id].try_write() {
                warn!("{}: An error when write to socket {} occured, error: {}",
                      self.address(),
                      address,
                      e);
                self.try_reconnect(event_loop, id);
                return Ok(None);
            }
            self.reregister_connection(event_loop, id, PollOpt::edge())?;

            self.mark_connected(event_loop, id);
            return Ok(None);
        }

        if set.is_readable() {
            let address = *self.connections[id].address();
            self.mark_connected(event_loop, id);

            trace!("{}: Socket is readable {} id: {}",
                   self.address(),
                   address,
                   id.0);

            return match self.connections[id].try_read() {
                Ok(buffer) => Ok(buffer),
                Err(e) => {
                    warn!("{}: An error when read from socket {} occured, error: {}",
                          self.address(),
                          address,
                          e);
                    self.try_reconnect(event_loop, id);
                    Ok(None)
                }
            };
        }
        // FIXME: Can we be here?
        Ok(None)
    }

    pub fn tick(&mut self, _: &mut EventLoop) {}

    pub fn get_peer(&mut self,
                    event_loop: &mut EventLoop,
                    address: &SocketAddr)
                    -> io::Result<PeerId> {
        if let Some(id) = self.addresses.get(address) {
            return Ok(*id);
        };

        let peer = Connection::new(TcpStream::connect(address)?, *address, Direction::Outgoing);
        let id = self.add_connection(event_loop, peer)?;

        trace!("{}: Establish connection with {}, id: {}",
               self.address(),
               address,
               id.0);
        Ok(id)
    }

    pub fn send_to(&mut self,
                   event_loop: &mut EventLoop,
                   address: &SocketAddr,
                   message: RawMessage) {
        if self.reconnects.contains_key(address) {
            info!("{}: An attempt to send data socket {} which is not online.",
                  self.address(),
                  address);
            return;
        }

        let r = match self.get_peer(event_loop, address) {
            Ok(id) => {
                self.connections[id]
                    .send(message)
                    .and_then(|_| {
                        event_loop.reregister(self.connections[id].socket(),
                                        id,
                                        self.connections[id].interest(),
                                        PollOpt::edge())?;
                        self.mark_connected(event_loop, id);
                        Ok(())
                    })
                    .or_else(|e| {
                        self.try_reconnect(event_loop, id);
                        Err(e)
                    })
            }
            Err(e) => Err(e),
        };

        if let Err(e) = r {
            warn!("{}: An error occured when try to send a message (address: {}, error: {:?})",
                  self.address(),
                  address,
                  e);
            return;
        }
    }

    pub fn handle_timeout(&mut self, event_loop: &mut EventLoop, timeout: InternalTimeout) {
        match timeout {
            InternalTimeout::Reconnect(addr, delay) => {
                if self.reconnects.contains_key(&addr) {
                    if let Err(e) = self.get_peer(event_loop, &addr) {
                        error!("{}: Unable to create connection to addr {}, error {:?}",
                               self.address(),
                               addr,
                               e);
                    }

                    let delay = min((delay as f32 * RECONNECT_GROW_FACTOR) as u64, self.config.tcp_reconnect_timeout_max);
                    // TODO Fail-safe timeout error handling
                    self.add_reconnect_timeout(event_loop, addr, delay).unwrap();

                    trace!("Try to reconnect with delay {}", delay);
                }
            }
        }
    }

    fn remove_connection(&mut self, id: Token) {
        let address = *self.connections[id].address();
        self.addresses.remove(&address);
        self.connections.remove(id);
    }

    fn try_reconnect(&mut self, event_loop: &mut EventLoop, id: Token) {
        if self.connections[id].direction() == &Direction::Outgoing {
            let addr = *self.connections[id].address();
            match self.try_reconnect_addr(event_loop, addr) {
                Err(e) => {
                    error!("{}: An error during reconnect occured: {:?}",
                           self.address(),
                           e);
                }
                Ok(_) => {}
            }
        }
        self.remove_connection(id);
    }

    fn try_reconnect_addr(&mut self,
                          event_loop: &mut EventLoop,
                          address: SocketAddr)
                          -> io::Result<()> {
        if !self.reconnects.contains_key(&address) {
            let delay = self.config.tcp_reconnect_timeout;
            return self.add_reconnect_timeout(event_loop, address, delay);
        }
        Ok(())
    }

    fn add_reconnect_timeout(&mut self,
                             event_loop: &mut EventLoop,
                             address: SocketAddr,
                             delay: u64)
                             -> io::Result<()> {
        let reconnect = Timeout::Internal(InternalTimeout::Reconnect(address, delay));
        let timeout = event_loop.timeout_ms(reconnect, delay)
            .map_err(|e| {
                io::Error::new(io::ErrorKind::Other, format!("A mio error occured {:?}", e))
            })?;
        self.reconnects.insert(address, timeout);

        trace!("Add reconnect timeout {}", delay);
        Ok(())
    }

    fn add_connection(&mut self,
                      event_loop: &mut EventLoop,
                      mut connection: Connection)
                      -> io::Result<PeerId> {
        self.configure_stream(connection.socket_mut())?;
        let address = *connection.address();
        let id = self.connections
            .insert(connection)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "Maximum connections"))?;
        self.addresses.insert(address, id);

        self.register_connection(event_loop, id, PollOpt::edge())?;
        Ok(id)
    }

    fn mark_connected(&mut self, event_loop: &mut EventLoop, id: Token) {
        let address = *self.connections[id].address();
        self.clear_reconnect_request(event_loop, &address)
    }

    fn clear_reconnect_request(&mut self, event_loop: &mut EventLoop, addr: &SocketAddr) {
        if let Some(timeout) = self.reconnects.remove(addr) {
            trace!("{}: Clear reconnect timeout for {}", self.address(), addr);
            event_loop.clear_timeout(timeout);
        }
    }

    fn register_connection(&mut self,
                           event_loop: &mut EventLoop,
                           id: Token,
                           opts: PollOpt)
                           -> io::Result<()> {
        event_loop.register(self.connections[id].socket(),
                      id,
                      self.connections[id].interest() | EventSet::writable(),
                      opts)
            .or_else(|e| {
                self.remove_connection(id);
                Err(e)
            })
    }

    fn reregister_connection(&mut self,
                             event_loop: &mut EventLoop,
                             id: Token,
                             opts: PollOpt)
                             -> io::Result<()> {
        event_loop.reregister(self.connections[id].socket(),
                        id,
                        self.connections[id].interest(),
                        opts)
            .or_else(|e| {
                self.remove_connection(id);
                Err(e)
            })
    }

    fn configure_stream(&self, stream: &mut TcpStream) -> io::Result<()> {
        stream.set_keepalive(self.config.tcp_keep_alive)?;
        stream.set_nodelay(self.config.tcp_nodelay)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {}
<|endoftext|>


<commit_before>use std::iter;

use either::Either;
use hir::{AsName, EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use rustc_hash::{FxHashMap, FxHashSet};
use syntax::{
    algo::{find_node_at_offset, SyntaxRewriter},
    ast::{self, edit::IndentLevel, make, AstNode, NameOwner, VisibilityOwner},
    SourceFile, SyntaxElement, SyntaxNode, T,
};

use crate::{
    utils::{insert_use, mod_path_to_ast, ImportScope},
    AssistContext, AssistId, AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = extract_field_list_if_applicable(&variant)?;

    let variant_name = variant.name()?;
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_definition(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }

    let enum_ast = variant.parent_enum();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let variant_hir_name = variant_hir.name(ctx.db());
            let enum_module_def = ModuleDef::from(enum_hir);
            let usages =
                Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)).usages(&ctx.sema).all();

            let mut visited_modules_set = FxHashSet::default();
            let current_module = enum_hir.module(ctx.db());
            visited_modules_set.insert(current_module);
            let mut rewriters = FxHashMap::default();
            for reference in usages {
                let rewriter = rewriters
                    .entry(reference.file_range.file_id)
                    .or_insert_with(SyntaxRewriter::default);
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    rewriter,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            let mut rewriter =
                rewriters.remove(&ctx.frange.file_id).unwrap_or_else(SyntaxRewriter::default);
            for (file_id, rewriter) in rewriters {
                builder.edit_file(file_id);
                builder.rewrite(rewriter);
            }
            builder.edit_file(ctx.frange.file_id);
            update_variant(&mut rewriter, &variant);
            extract_struct_def(
                &mut rewriter,
                &enum_ast,
                variant_name.clone(),
                &field_list,
                &variant.parent_enum().syntax().clone().into(),
                enum_ast.visibility(),
            );
            builder.rewrite(rewriter);
        },
    )
}

fn extract_field_list_if_applicable(
    variant: &ast::Variant,
) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
    match variant.kind() {
        ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
            Some(Either::Left(field_list))
        }
        ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
            Some(Either::Right(field_list))
        }
        _ => None,
    }
}

fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .filter(|(_, def)| match def {
            // only check type-namespace
            hir::ScopeDef::ModuleDef(def) => matches!(def,
                ModuleDef::Module(_) | ModuleDef::Adt(_) |
                ModuleDef::EnumVariant(_) | ModuleDef::Trait(_) |
                ModuleDef::TypeAlias(_) | ModuleDef::BuiltinType(_)
            ),
            _ => false,
        })
        .any(|(name, _)| name == variant_name.as_name())
}

fn insert_import(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    scope_node: &SyntaxNode,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path(db, enum_module_def.clone());
    if let Some(mut mod_path) = mod_path.filter(|path| path.len() > 1) {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        let scope = ImportScope::find_insert_use_container(scope_node, ctx)?;

        *rewriter += insert_use(&scope, mod_path_to_ast(&mod_path), ctx.config.insert_use.merge);
    }
    Some(())
}

fn extract_struct_def(
    rewriter: &mut SyntaxRewriter,
    enum_: &ast::Enum,
    variant_name: ast::Name,
    field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
    start_offset: &SyntaxElement,
    visibility: Option<ast::Visibility>,
) -> Option<()> {
    let pub_vis = Some(make::visibility_pub());
    let field_list = match field_list {
        Either::Left(field_list) => {
            make::record_field_list(field_list.fields().flat_map(|field| {
                Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
            }))
            .into()
        }
        Either::Right(field_list) => make::tuple_field_list(
            field_list
                .fields()
                .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
        )
        .into(),
    };

    rewriter.insert_before(
        start_offset,
        make::struct_(visibility, variant_name, None, field_list).syntax(),
    );
    rewriter.insert_before(start_offset, &make::tokens::blank_line());

    if let indent_level @ 1..=usize::MAX = IndentLevel::from_node(enum_.syntax()).0 as usize {
        rewriter
            .insert_before(start_offset, &make::tokens::whitespace(&" ".repeat(4 * indent_level)));
    }
    Some(())
}

fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Option<()> {
    let name = variant.name()?;
    let tuple_field = make::tuple_field(None, make::ty(name.text()));
    let replacement = make::variant(
        name,
        Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
    );
    rewriter.replace(variant.syntax(), replacement.syntax());
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let offset = reference.file_range.range.start();
    let (segment, expr) = if let Some(path_expr) =
        find_node_at_offset::<ast::PathExpr>(source_file.syntax(), offset)
    {
        // tuple variant
        (path_expr.path()?.segment()?, path_expr.syntax().parent()?.clone())
    } else if let Some(record_expr) =
        find_node_at_offset::<ast::RecordExpr>(source_file.syntax(), offset)
    {
        // record variant
        (record_expr.path()?.segment()?, record_expr.syntax().clone())
    } else {
        return None;
    };

    let module = ctx.sema.scope(&expr).module()?;
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, rewriter, &expr, &module, enum_module_def, variant_hir_name).is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    rewriter.insert_after(segment.syntax(), &make::token(T!['(']));
    rewriter.insert_after(segment.syntax(), segment.syntax());
    rewriter.insert_after(&expr, &make::token(T![')']));
    Some(())
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields_tuple() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_several_fields_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32, bar: u32 } }",
            r#"struct One{ pub foo: u32, pub bar: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32 } }",
            r#"struct One{ pub foo: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_enum_variant_name_value_namespace() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"const One: () = ();
enum A { <|>One(u32, u32) }"#,
            r#"const One: () = ();
struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    #[test]
    fn extract_record_fix_references() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
enum E {
    <|>V { i: i32, j: i32 }
}

fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}

fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    #[test]
    fn test_several_files() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V(i32, i32)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(9, 2);
}
"#,
            r#"
//- /main.rs
struct V(pub i32, pub i32);

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(V(9, 2));
}
"#,
        )
    }

    #[test]
    fn test_several_files_record() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V { i: i32, j: i32 }
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
//- /main.rs
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8, u32) }"#,
        );
    }

    #[test]
    fn test_extract_not_applicable_one_field() {
        check_not_applicable(r"enum A { <|>One(u32) }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_tuple() {
        check_not_applicable(r"enum A { <|>None() }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_named() {
        check_not_applicable(r"enum A { <|>None {} }");
    }
}
<commit_msg>Use Module::find_use_path_prefixed in extract_struct_from_enum_variant assist<commit_after>use std::iter;

use either::Either;
use hir::{AsName, EnumVariant, Module, ModuleDef, Name};
use ide_db::{defs::Definition, search::Reference, RootDatabase};
use rustc_hash::{FxHashMap, FxHashSet};
use syntax::{
    algo::{find_node_at_offset, SyntaxRewriter},
    ast::{self, edit::IndentLevel, make, AstNode, NameOwner, VisibilityOwner},
    SourceFile, SyntaxElement, SyntaxNode, T,
};

use crate::{
    utils::{insert_use, mod_path_to_ast, ImportScope},
    AssistContext, AssistId, AssistKind, Assists,
};

// Assist: extract_struct_from_enum_variant
//
// Extracts a struct from enum variant.
//
// ```
// enum A { <|>One(u32, u32) }
// ```
// ->
// ```
// struct One(pub u32, pub u32);
//
// enum A { One(One) }
// ```
pub(crate) fn extract_struct_from_enum_variant(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let field_list = extract_field_list_if_applicable(&variant)?;

    let variant_name = variant.name()?;
    let variant_hir = ctx.sema.to_def(&variant)?;
    if existing_definition(ctx.db(), &variant_name, &variant_hir) {
        return None;
    }

    let enum_ast = variant.parent_enum();
    let enum_hir = ctx.sema.to_def(&enum_ast)?;
    let target = variant.syntax().text_range();
    acc.add(
        AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
        "Extract struct from enum variant",
        target,
        |builder| {
            let variant_hir_name = variant_hir.name(ctx.db());
            let enum_module_def = ModuleDef::from(enum_hir);
            let usages =
                Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)).usages(&ctx.sema).all();

            let mut visited_modules_set = FxHashSet::default();
            let current_module = enum_hir.module(ctx.db());
            visited_modules_set.insert(current_module);
            let mut rewriters = FxHashMap::default();
            for reference in usages {
                let rewriter = rewriters
                    .entry(reference.file_range.file_id)
                    .or_insert_with(SyntaxRewriter::default);
                let source_file = ctx.sema.parse(reference.file_range.file_id);
                update_reference(
                    ctx,
                    rewriter,
                    reference,
                    &source_file,
                    &enum_module_def,
                    &variant_hir_name,
                    &mut visited_modules_set,
                );
            }
            let mut rewriter =
                rewriters.remove(&ctx.frange.file_id).unwrap_or_else(SyntaxRewriter::default);
            for (file_id, rewriter) in rewriters {
                builder.edit_file(file_id);
                builder.rewrite(rewriter);
            }
            builder.edit_file(ctx.frange.file_id);
            update_variant(&mut rewriter, &variant);
            extract_struct_def(
                &mut rewriter,
                &enum_ast,
                variant_name.clone(),
                &field_list,
                &variant.parent_enum().syntax().clone().into(),
                enum_ast.visibility(),
            );
            builder.rewrite(rewriter);
        },
    )
}

fn extract_field_list_if_applicable(
    variant: &ast::Variant,
) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
    match variant.kind() {
        ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
            Some(Either::Left(field_list))
        }
        ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
            Some(Either::Right(field_list))
        }
        _ => None,
    }
}

fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
    variant
        .parent_enum(db)
        .module(db)
        .scope(db, None)
        .into_iter()
        .filter(|(_, def)| match def {
            // only check type-namespace
            hir::ScopeDef::ModuleDef(def) => matches!(def,
                ModuleDef::Module(_) | ModuleDef::Adt(_) |
                ModuleDef::EnumVariant(_) | ModuleDef::Trait(_) |
                ModuleDef::TypeAlias(_) | ModuleDef::BuiltinType(_)
            ),
            _ => false,
        })
        .any(|(name, _)| name == variant_name.as_name())
}

fn insert_import(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    scope_node: &SyntaxNode,
    module: &Module,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
) -> Option<()> {
    let db = ctx.db();
    let mod_path = module.find_use_path_prefixed(
        db,
        enum_module_def.clone(),
        ctx.config.insert_use.prefix_kind,
    );
    if let Some(mut mod_path) = mod_path.filter(|path| !path.is_ident()) {
        mod_path.segments.pop();
        mod_path.segments.push(variant_hir_name.clone());
        let scope = ImportScope::find_insert_use_container(scope_node, ctx)?;

        *rewriter += insert_use(&scope, mod_path_to_ast(&mod_path), ctx.config.insert_use.merge);
    }
    Some(())
}

fn extract_struct_def(
    rewriter: &mut SyntaxRewriter,
    enum_: &ast::Enum,
    variant_name: ast::Name,
    field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
    start_offset: &SyntaxElement,
    visibility: Option<ast::Visibility>,
) -> Option<()> {
    let pub_vis = Some(make::visibility_pub());
    let field_list = match field_list {
        Either::Left(field_list) => {
            make::record_field_list(field_list.fields().flat_map(|field| {
                Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
            }))
            .into()
        }
        Either::Right(field_list) => make::tuple_field_list(
            field_list
                .fields()
                .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
        )
        .into(),
    };

    rewriter.insert_before(
        start_offset,
        make::struct_(visibility, variant_name, None, field_list).syntax(),
    );
    rewriter.insert_before(start_offset, &make::tokens::blank_line());

    if let indent_level @ 1..=usize::MAX = IndentLevel::from_node(enum_.syntax()).0 as usize {
        rewriter
            .insert_before(start_offset, &make::tokens::whitespace(&" ".repeat(4 * indent_level)));
    }
    Some(())
}

fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Option<()> {
    let name = variant.name()?;
    let tuple_field = make::tuple_field(None, make::ty(name.text()));
    let replacement = make::variant(
        name,
        Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
    );
    rewriter.replace(variant.syntax(), replacement.syntax());
    Some(())
}

fn update_reference(
    ctx: &AssistContext,
    rewriter: &mut SyntaxRewriter,
    reference: Reference,
    source_file: &SourceFile,
    enum_module_def: &ModuleDef,
    variant_hir_name: &Name,
    visited_modules_set: &mut FxHashSet<Module>,
) -> Option<()> {
    let offset = reference.file_range.range.start();
    let (segment, expr) = if let Some(path_expr) =
        find_node_at_offset::<ast::PathExpr>(source_file.syntax(), offset)
    {
        // tuple variant
        (path_expr.path()?.segment()?, path_expr.syntax().parent()?.clone())
    } else if let Some(record_expr) =
        find_node_at_offset::<ast::RecordExpr>(source_file.syntax(), offset)
    {
        // record variant
        (record_expr.path()?.segment()?, record_expr.syntax().clone())
    } else {
        return None;
    };

    let module = ctx.sema.scope(&expr).module()?;
    if !visited_modules_set.contains(&module) {
        if insert_import(ctx, rewriter, &expr, &module, enum_module_def, variant_hir_name).is_some()
        {
            visited_modules_set.insert(module);
        }
    }
    rewriter.insert_after(segment.syntax(), &make::token(T!['(']));
    rewriter.insert_after(segment.syntax(), segment.syntax());
    rewriter.insert_after(&expr, &make::token(T![')']));
    Some(())
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_assist, check_assist_not_applicable},
        utils::FamousDefs,
    };

    use super::*;

    #[test]
    fn test_extract_struct_several_fields_tuple() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One(u32, u32) }",
            r#"struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_several_fields_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32, bar: u32 } }",
            r#"struct One{ pub foo: u32, pub bar: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_one_field_named() {
        check_assist(
            extract_struct_from_enum_variant,
            "enum A { <|>One { foo: u32 } }",
            r#"struct One{ pub foo: u32 }

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_enum_variant_name_value_namespace() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"const One: () = ();
enum A { <|>One(u32, u32) }"#,
            r#"const One: () = ();
struct One(pub u32, pub u32);

enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_pub_visibility() {
        check_assist(
            extract_struct_from_enum_variant,
            "pub enum A { <|>One(u32, u32) }",
            r#"pub struct One(pub u32, pub u32);

pub enum A { One(One) }"#,
        );
    }

    #[test]
    fn test_extract_struct_with_complex_imports() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"mod my_mod {
    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(1, 1);
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(1, 1);
        }

        pub enum MyEnum {
            <|>MyField(u8, u8),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
}"#,
            r#"use my_mod::my_other_mod::MyField;

mod my_mod {
    use self::my_other_mod::MyField;

    fn another_fn() {
        let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
    }

    pub mod my_other_mod {
        fn another_fn() {
            let m = MyEnum::MyField(MyField(1, 1));
        }

        pub struct MyField(pub u8, pub u8);

        pub enum MyEnum {
            MyField(MyField),
        }
    }
}

fn another_fn() {
    let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
}"#,
        );
    }

    #[test]
    fn extract_record_fix_references() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
enum E {
    <|>V { i: i32, j: i32 }
}

fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}

fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    #[test]
    fn test_several_files() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V(i32, i32)
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V(9, 2);
}
"#,
            r#"
//- /main.rs
struct V(pub i32, pub i32);

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V(9, 2));
}
"#,
        )
    }

    #[test]
    fn test_several_files_record() {
        check_assist(
            extract_struct_from_enum_variant,
            r#"
//- /main.rs
enum E {
    <|>V { i: i32, j: i32 }
}
mod foo;

//- /foo.rs
use crate::E;
fn f() {
    let e = E::V { i: 9, j: 2 };
}
"#,
            r#"
//- /main.rs
struct V{ pub i: i32, pub j: i32 }

enum E {
    V(V)
}
mod foo;

//- /foo.rs
use crate::{E, V};
fn f() {
    let e = E::V(V { i: 9, j: 2 });
}
"#,
        )
    }

    fn check_not_applicable(ra_fixture: &str) {
        let fixture =
            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
        check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
    }

    #[test]
    fn test_extract_enum_not_applicable_for_element_with_no_fields() {
        check_not_applicable("enum A { <|>One }");
    }

    #[test]
    fn test_extract_enum_not_applicable_if_struct_exists() {
        check_not_applicable(
            r#"struct One;
        enum A { <|>One(u8, u32) }"#,
        );
    }

    #[test]
    fn test_extract_not_applicable_one_field() {
        check_not_applicable(r"enum A { <|>One(u32) }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_tuple() {
        check_not_applicable(r"enum A { <|>None() }");
    }

    #[test]
    fn test_extract_not_applicable_no_field_named() {
        check_not_applicable(r"enum A { <|>None {} }");
    }
}
<|endoftext|>