<commit_before>//! Completes keywords.

use syntax::SyntaxKind;
use test_utils::mark;

use crate::{CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions};

pub(crate) fn complete_use_tree_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    // complete keyword "crate" in use stmt
    let source_range = ctx.source_range();

    if ctx.use_item_syntax.is_some() {
        if ctx.path_qual.is_none() {
            CompletionItem::new(CompletionKind::Keyword, source_range, "crate::")
                .kind(CompletionItemKind::Keyword)
                .insert_text("crate::")
                .add_to(acc);
        }
        CompletionItem::new(CompletionKind::Keyword, source_range, "self")
            .kind(CompletionItemKind::Keyword)
            .add_to(acc);
        CompletionItem::new(CompletionKind::Keyword, source_range, "super::")
            .kind(CompletionItemKind::Keyword)
            .insert_text("super::")
            .add_to(acc);
    }

    // Suggest .await syntax for types that implement Future trait
    if let Some(receiver) = &ctx.dot_receiver {
        if let Some(ty) = ctx.sema.type_of_expr(receiver) {
            if ty.impls_future(ctx.db) {
                CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), "await")
                    .kind(CompletionItemKind::Keyword)
                    .detail("expr.await")
                    .insert_text("await")
                    .add_to(acc);
            }
        };
    }
}

pub(crate) fn complete_expr_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    if ctx.token.kind() == SyntaxKind::COMMENT {
        mark::hit!(no_keyword_completion_in_comments);
        return;
    }
    if ctx.record_lit_syntax.is_some() {
        mark::hit!(no_keyword_completion_in_record_lit);
        return;
    }

    let has_trait_or_impl_parent = ctx.has_impl_parent || ctx.has_trait_parent;
    if ctx.trait_as_prev_sibling || ctx.impl_as_prev_sibling {
        add_keyword(ctx, acc, "where", "where ");
        return;
    }
    if ctx.unsafe_is_prev {
        if ctx.has_item_list_or_source_file_parent || ctx.block_expr_parent {
            add_keyword(ctx, acc, "fn", "fn $0() {}")
        }

        if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
            add_keyword(ctx, acc, "trait", "trait $0 {}");
            add_keyword(ctx, acc, "impl", "impl $0 {}");
        }

        return;
    }
    if ctx.has_item_list_or_source_file_parent || has_trait_or_impl_parent || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "fn", "fn $0() {}");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "use", "use ");
        add_keyword(ctx, acc, "impl", "impl $0 {}");
        add_keyword(ctx, acc, "trait", "trait $0 {}");
    }

    if ctx.has_item_list_or_source_file_parent {
        add_keyword(ctx, acc, "enum", "enum $0 {}");
        add_keyword(ctx, acc, "struct", "struct $0");
        add_keyword(ctx, acc, "union", "union $0 {}");
    }

    if ctx.is_expr {
        add_keyword(ctx, acc, "match", "match $0 {}");
        add_keyword(ctx, acc, "while", "while $0 {}");
        add_keyword(ctx, acc, "loop", "loop {$0}");
        add_keyword(ctx, acc, "if", "if $0 {}");
        add_keyword(ctx, acc, "if let", "if let $1 = $0 {}");
        add_keyword(ctx, acc, "for", "for $0 in $1 {}");
    }

    if ctx.if_is_prev || ctx.block_expr_parent {
        add_keyword(ctx, acc, "let", "let ");
    }

    if ctx.after_if {
        add_keyword(ctx, acc, "else", "else {$0}");
        add_keyword(ctx, acc, "else if", "else if $0 {}");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "mod", "mod $0");
    }
    if ctx.bind_pat_parent || ctx.ref_pat_parent {
        add_keyword(ctx, acc, "mut", "mut ");
    }
    if ctx.has_item_list_or_source_file_parent || has_trait_or_impl_parent || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "const", "const ");
        add_keyword(ctx, acc, "type", "type ");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "static", "static ");
    };
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "extern", "extern ");
    }
    if ctx.has_item_list_or_source_file_parent
        || has_trait_or_impl_parent
        || ctx.block_expr_parent
        || ctx.is_match_arm
    {
        add_keyword(ctx, acc, "unsafe", "unsafe ");
    }
    if ctx.in_loop_body {
        if ctx.can_be_stmt {
            add_keyword(ctx, acc, "continue", "continue;");
            add_keyword(ctx, acc, "break", "break;");
        } else {
            add_keyword(ctx, acc, "continue", "continue");
            add_keyword(ctx, acc, "break", "break");
        }
    }
    if ctx.has_item_list_or_source_file_parent || ctx.has_impl_parent | ctx.has_field_list_parent {
        add_keyword(ctx, acc, "pub(crate)", "pub(crate) ");
        add_keyword(ctx, acc, "pub", "pub ");
    }

    if !ctx.is_trivial_path {
        return;
    }
    let fn_def = match &ctx.function_syntax {
        Some(it) => it,
        None => return,
    };

    add_keyword(
        ctx,
        acc,
        "return",
        match (ctx.can_be_stmt, fn_def.ret_type().is_some()) {
            (true, true) => "return $0;",
            (true, false) => "return;",
            (false, true) => "return $0",
            (false, false) => "return",
        },
    )
}

fn add_keyword(ctx: &CompletionContext, acc: &mut Completions, kw: &str, snippet: &str) {
    let builder = CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), kw)
        .kind(CompletionItemKind::Keyword);
    let builder = match ctx.config.snippet_cap {
        Some(cap) => {
            let tmp;
            let snippet = if snippet.ends_with('}') && ctx.incomplete_let {
                mark::hit!(let_semi);
                tmp = format!("{};", snippet);
                &tmp
            } else {
                snippet
            };
            builder.insert_snippet(cap, snippet)
        }
        None => builder.insert_text(if snippet.contains('$') { kw } else { snippet }),
    };
    acc.add(builder.build());
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};
    use test_utils::mark;

    use crate::{
        test_utils::{check_edit, completion_list},
        CompletionKind,
    };

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = completion_list(ra_fixture, CompletionKind::Keyword);
        expect.assert_eq(&actual)
    }

    #[test]
    fn test_keywords_in_use_stmt() {
        check(
            r"use $0",
            expect![[r#"
                kw crate::
                kw self
                kw super::
            "#]],
        );

        check(
            r"use a::$0",
            expect![[r#"
                kw self
                kw super::
            "#]],
        );

        check(
            r"use a::{b, $0}",
            expect![[r#"
                kw self
                kw super::
            "#]],
        );
    }

    #[test]
    fn test_keywords_at_source_file_level() {
        check(
            r"m$0",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw enum
                kw struct
                kw union
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw pub(crate)
                kw pub
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_function() {
        check(
            r"fn quux() { $0 }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_inside_block() {
        check(
            r"fn quux() { if true { $0 } }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_if() {
        check(
            r#"fn quux() { if true { () } $0 }"#,
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw let
                kw else
                kw else if
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
        check_edit(
            "else",
            r#"fn quux() { if true { () } $0 }"#,
            r#"fn quux() { if true { () } else {$0} }"#,
        );
    }

    #[test]
    fn test_keywords_in_match_arm() {
        check(
            r#"
fn quux() -> i32 {
    match () { () => $0 }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_trait_def() {
        check(
            r"trait My { $0 }",
            expect![[r#"
                kw fn
                kw const
                kw type
                kw unsafe
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_impl_def() {
        check(
            r"impl My { $0 }",
            expect![[r#"
                kw fn
                kw const
                kw type
                kw unsafe
                kw pub(crate)
                kw pub
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_loop() {
        check(
            r"fn my() { loop { $0 } }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw continue
                kw break
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_item_list() {
        check(
            r"unsafe $0",
            expect![[r#"
                kw fn
                kw trait
                kw impl
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_block_expr() {
        check(
            r"fn my_fn() { unsafe $0 }",
            expect![[r#"
                kw fn
                kw trait
                kw impl
            "#]],
        );
    }

    #[test]
    fn test_mut_in_ref_and_in_fn_parameters_list() {
        check(
            r"fn my_fn(&$0) {}",
            expect![[r#"
                kw mut
            "#]],
        );
        check(
            r"fn my_fn($0) {}",
            expect![[r#"
                kw mut
            "#]],
        );
        check(
            r"fn my_fn() { let &$0 }",
            expect![[r#"
                kw mut
            "#]],
        );
    }

    #[test]
    fn test_where_keyword() {
        check(
            r"trait A $0",
            expect![[r#"
                kw where
            "#]],
        );
        check(
            r"impl A $0",
            expect![[r#"
                kw where
            "#]],
        );
    }

    #[test]
    fn no_keyword_completion_in_comments() {
        mark::check!(no_keyword_completion_in_comments);
        check(
            r#"
fn test() {
    let x = 2; // A comment$0
}
"#,
            expect![[""]],
        );
        check(
            r#"
/*
Some multi-line comment$0
*/
"#,
            expect![[""]],
        );
        check(
            r#"
/// Some doc comment
/// let test$0 = 1
"#,
            expect![[""]],
        );
    }

    #[test]
    fn test_completion_await_impls_future() {
        check(
            r#"
//- /main.rs crate:main deps:std
use std::future::*;
struct A {}
impl Future for A {}
fn foo(a: A) { a.$0 }

//- /std/lib.rs crate:std
pub mod future {
    #[lang = "future_trait"]
    pub trait Future {}
}
"#,
            expect![[r#"
                kw await expr.await
            "#]],
        );

        check(
            r#"
//- /main.rs crate:main deps:std
use std::future::*;
fn foo() {
    let a = async {};
    a.$0
}

//- /std/lib.rs crate:std
pub mod future {
    #[lang = "future_trait"]
    pub trait Future {
        type Output;
    }
}
"#,
            expect![[r#"
                kw await expr.await
            "#]],
        )
    }

    #[test]
    fn after_let() {
        check(
            r#"fn main() { let _ = $0 }"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw return
            "#]],
        )
    }

    #[test]
    fn before_field() {
        check(
            r#"
struct Foo {
    $0
    pub f: i32,
}
"#,
            expect![[r#"
                kw pub(crate)
                kw pub
            "#]],
        )
    }

    #[test]
    fn skip_struct_initializer() {
        mark::check!(no_keyword_completion_in_record_lit);
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        $0
    }
}
"#,
            expect![[r#""#]],
        );
    }

    #[test]
    fn struct_initializer_field_expr() {
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        f: $0
    }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw return
            "#]],
        );
    }

    #[test]
    fn let_semi() {
        mark::check!(let_semi);
        check_edit(
            "match",
            r#"
fn main() { let x = $0 }
"#,
            r#"
fn main() { let x = match $0 {}; }
"#,
        );

        check_edit(
            "if",
            r#"
fn main() {
    let x = $0
    let y = 92;
}
"#,
            r#"
fn main() {
    let x = if $0 {};
    let y = 92;
}
"#,
        );

        check_edit(
            "loop",
            r#"
fn main() {
    let x = $0
    bar();
}
"#,
            r#"
fn main() {
    let x = loop {$0};
    bar();
}
"#,
        );
    }
}
<commit_msg>ci: update relevant test case expected values<commit_after>//! Completes keywords.

use syntax::SyntaxKind;
use test_utils::mark;

use crate::{CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions};

pub(crate) fn complete_use_tree_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    // complete keyword "crate" in use stmt
    let source_range = ctx.source_range();

    if ctx.use_item_syntax.is_some() {
        if ctx.path_qual.is_none() {
            CompletionItem::new(CompletionKind::Keyword, source_range, "crate::")
                .kind(CompletionItemKind::Keyword)
                .insert_text("crate::")
                .add_to(acc);
        }
        CompletionItem::new(CompletionKind::Keyword, source_range, "self")
            .kind(CompletionItemKind::Keyword)
            .add_to(acc);
        CompletionItem::new(CompletionKind::Keyword, source_range, "super::")
            .kind(CompletionItemKind::Keyword)
            .insert_text("super::")
            .add_to(acc);
    }

    // Suggest .await syntax for types that implement Future trait
    if let Some(receiver) = &ctx.dot_receiver {
        if let Some(ty) = ctx.sema.type_of_expr(receiver) {
            if ty.impls_future(ctx.db) {
                CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), "await")
                    .kind(CompletionItemKind::Keyword)
                    .detail("expr.await")
                    .insert_text("await")
                    .add_to(acc);
            }
        };
    }
}

pub(crate) fn complete_expr_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    if ctx.token.kind() == SyntaxKind::COMMENT {
        mark::hit!(no_keyword_completion_in_comments);
        return;
    }
    if ctx.record_lit_syntax.is_some() {
        mark::hit!(no_keyword_completion_in_record_lit);
        return;
    }

    let has_trait_or_impl_parent = ctx.has_impl_parent || ctx.has_trait_parent;
    if ctx.trait_as_prev_sibling || ctx.impl_as_prev_sibling {
        add_keyword(ctx, acc, "where", "where ");
        return;
    }
    if ctx.unsafe_is_prev {
        if ctx.has_item_list_or_source_file_parent || ctx.block_expr_parent {
            add_keyword(ctx, acc, "fn", "fn $0() {}")
        }

        if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
            add_keyword(ctx, acc, "trait", "trait $0 {}");
            add_keyword(ctx, acc, "impl", "impl $0 {}");
        }

        return;
    }
    if ctx.has_item_list_or_source_file_parent || has_trait_or_impl_parent || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "fn", "fn $0() {}");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "use", "use ");
        add_keyword(ctx, acc, "impl", "impl $0 {}");
        add_keyword(ctx, acc, "trait", "trait $0 {}");
    }

    if ctx.has_item_list_or_source_file_parent {
        add_keyword(ctx, acc, "enum", "enum $0 {}");
        add_keyword(ctx, acc, "struct", "struct $0");
        add_keyword(ctx, acc, "union", "union $0 {}");
    }

    if ctx.is_expr {
        add_keyword(ctx, acc, "match", "match $0 {}");
        add_keyword(ctx, acc, "while", "while $0 {}");
        add_keyword(ctx, acc, "loop", "loop {$0}");
        add_keyword(ctx, acc, "if", "if $0 {}");
        add_keyword(ctx, acc, "if let", "if let $1 = $0 {}");
        add_keyword(ctx, acc, "for", "for $0 in $1 {}");
    }

    if ctx.if_is_prev || ctx.block_expr_parent {
        add_keyword(ctx, acc, "let", "let ");
    }

    if ctx.after_if {
        add_keyword(ctx, acc, "else", "else {$0}");
        add_keyword(ctx, acc, "else if", "else if $0 {}");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "mod", "mod $0");
    }
    if ctx.bind_pat_parent || ctx.ref_pat_parent {
        add_keyword(ctx, acc, "mut", "mut ");
    }
    if ctx.has_item_list_or_source_file_parent || has_trait_or_impl_parent || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "const", "const ");
        add_keyword(ctx, acc, "type", "type ");
    }
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "static", "static ");
    };
    if (ctx.has_item_list_or_source_file_parent) || ctx.block_expr_parent {
        add_keyword(ctx, acc, "extern", "extern ");
    }
    if ctx.has_item_list_or_source_file_parent
        || has_trait_or_impl_parent
        || ctx.block_expr_parent
        || ctx.is_match_arm
    {
        add_keyword(ctx, acc, "unsafe", "unsafe ");
    }
    if ctx.in_loop_body {
        if ctx.can_be_stmt {
            add_keyword(ctx, acc, "continue", "continue;");
            add_keyword(ctx, acc, "break", "break;");
        } else {
            add_keyword(ctx, acc, "continue", "continue");
            add_keyword(ctx, acc, "break", "break");
        }
    }
    if ctx.has_item_list_or_source_file_parent || ctx.has_impl_parent | ctx.has_field_list_parent {
        add_keyword(ctx, acc, "pub(crate)", "pub(crate) ");
        add_keyword(ctx, acc, "pub", "pub ");
    }

    if !ctx.is_trivial_path {
        return;
    }
    let fn_def = match &ctx.function_syntax {
        Some(it) => it,
        None => return,
    };

    add_keyword(
        ctx,
        acc,
        "return",
        match (ctx.can_be_stmt, fn_def.ret_type().is_some()) {
            (true, true) => "return $0;",
            (true, false) => "return;",
            (false, true) => "return $0",
            (false, false) => "return",
        },
    )
}

fn add_keyword(ctx: &CompletionContext, acc: &mut Completions, kw: &str, snippet: &str) {
    let builder = CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), kw)
        .kind(CompletionItemKind::Keyword);
    let builder = match ctx.config.snippet_cap {
        Some(cap) => {
            let tmp;
            let snippet = if snippet.ends_with('}') && ctx.incomplete_let {
                mark::hit!(let_semi);
                tmp = format!("{};", snippet);
                &tmp
            } else {
                snippet
            };
            builder.insert_snippet(cap, snippet)
        }
        None => builder.insert_text(if snippet.contains('$') { kw } else { snippet }),
    };
    acc.add(builder.build());
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};
    use test_utils::mark;

    use crate::{
        test_utils::{check_edit, completion_list},
        CompletionKind,
    };

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = completion_list(ra_fixture, CompletionKind::Keyword);
        expect.assert_eq(&actual)
    }

    #[test]
    fn test_keywords_in_use_stmt() {
        check(
            r"use $0",
            expect![[r#"
                kw crate::
                kw self
                kw super::
            "#]],
        );

        check(
            r"use a::$0",
            expect![[r#"
                kw self
                kw super::
            "#]],
        );

        check(
            r"use a::{b, $0}",
            expect![[r#"
                kw self
                kw super::
            "#]],
        );
    }

    #[test]
    fn test_keywords_at_source_file_level() {
        check(
            r"m$0",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw enum
                kw struct
                kw union
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw pub(crate)
                kw pub
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_function() {
        check(
            r"fn quux() { $0 }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_inside_block() {
        check(
            r"fn quux() { if true { $0 } }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_if() {
        check(
            r#"fn quux() { if true { () } $0 }"#,
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw let
                kw else
                kw else if
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw return
            "#]],
        );
        check_edit(
            "else",
            r#"fn quux() { if true { () } $0 }"#,
            r#"fn quux() { if true { () } else {$0} }"#,
        );
    }

    #[test]
    fn test_keywords_in_match_arm() {
        check(
            r#"
fn quux() -> i32 {
    match () { () => $0 }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw unsafe
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_trait_def() {
        check(
            r"trait My { $0 }",
            expect![[r#"
                kw fn
                kw const
                kw type
                kw unsafe
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_impl_def() {
        check(
            r"impl My { $0 }",
            expect![[r#"
                kw fn
                kw const
                kw type
                kw unsafe
                kw pub(crate)
                kw pub
            "#]],
        );
    }

    #[test]
    fn test_keywords_in_loop() {
        check(
            r"fn my() { loop { $0 } }",
            expect![[r#"
                kw fn
                kw use
                kw impl
                kw trait
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw let
                kw mod
                kw const
                kw type
                kw static
                kw extern
                kw unsafe
                kw continue
                kw break
                kw return
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_item_list() {
        check(
            r"unsafe $0",
            expect![[r#"
                kw fn
                kw trait
                kw impl
            "#]],
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_block_expr() {
        check(
            r"fn my_fn() { unsafe $0 }",
            expect![[r#"
                kw fn
                kw trait
                kw impl
            "#]],
        );
    }

    #[test]
    fn test_mut_in_ref_and_in_fn_parameters_list() {
        check(
            r"fn my_fn(&$0) {}",
            expect![[r#"
                kw mut
            "#]],
        );
        check(
            r"fn my_fn($0) {}",
            expect![[r#"
                kw mut
            "#]],
        );
        check(
            r"fn my_fn() { let &$0 }",
            expect![[r#"
                kw mut
            "#]],
        );
    }

    #[test]
    fn test_where_keyword() {
        check(
            r"trait A $0",
            expect![[r#"
                kw where
            "#]],
        );
        check(
            r"impl A $0",
            expect![[r#"
                kw where
            "#]],
        );
    }

    #[test]
    fn no_keyword_completion_in_comments() {
        mark::check!(no_keyword_completion_in_comments);
        check(
            r#"
fn test() {
    let x = 2; // A comment$0
}
"#,
            expect![[""]],
        );
        check(
            r#"
/*
Some multi-line comment$0
*/
"#,
            expect![[""]],
        );
        check(
            r#"
/// Some doc comment
/// let test$0 = 1
"#,
            expect![[""]],
        );
    }

    #[test]
    fn test_completion_await_impls_future() {
        check(
            r#"
//- /main.rs crate:main deps:std
use std::future::*;
struct A {}
impl Future for A {}
fn foo(a: A) { a.$0 }

//- /std/lib.rs crate:std
pub mod future {
    #[lang = "future_trait"]
    pub trait Future {}
}
"#,
            expect![[r#"
                kw await expr.await
            "#]],
        );

        check(
            r#"
//- /main.rs crate:main deps:std
use std::future::*;
fn foo() {
    let a = async {};
    a.$0
}

//- /std/lib.rs crate:std
pub mod future {
    #[lang = "future_trait"]
    pub trait Future {
        type Output;
    }
}
"#,
            expect![[r#"
                kw await expr.await
            "#]],
        )
    }

    #[test]
    fn after_let() {
        check(
            r#"fn main() { let _ = $0 }"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw return
            "#]],
        )
    }

    #[test]
    fn before_field() {
        check(
            r#"
struct Foo {
    $0
    pub f: i32,
}
"#,
            expect![[r#"
                kw pub(crate)
                kw pub
            "#]],
        )
    }

    #[test]
    fn skip_struct_initializer() {
        mark::check!(no_keyword_completion_in_record_lit);
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        $0
    }
}
"#,
            expect![[r#""#]],
        );
    }

    #[test]
    fn struct_initializer_field_expr() {
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        f: $0
    }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw for
                kw return
            "#]],
        );
    }

    #[test]
    fn let_semi() {
        mark::check!(let_semi);
        check_edit(
            "match",
            r#"
fn main() { let x = $0 }
"#,
            r#"
fn main() { let x = match $0 {}; }
"#,
        );

        check_edit(
            "if",
            r#"
fn main() {
    let x = $0
    let y = 92;
}
"#,
            r#"
fn main() {
    let x = if $0 {};
    let y = 92;
}
"#,
        );

        check_edit(
            "loop",
            r#"
fn main() {
    let x = $0
    bar();
}
"#,
            r#"
fn main() {
    let x = loop {$0};
    bar();
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>use super::rpc::codec::RpcCodec;

use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
use std::io::{Error, ErrorKind};
use std::net::SocketAddr;

use tokio::net::TcpListener;
use tokio_codec::Framed;
use tokio::prelude::*;

use blockchain::Blockchain;
use producer::Producer;
use net::peer::*;
use net::rpc::*;
use fut_util::*;

pub fn start(addr: SocketAddr, blockchain: Arc<Blockchain>, producer: Arc<Option<Producer>>) {
    let listener = TcpListener::bind(&addr).unwrap();
    info!("Server binded to {:?}", &addr);
    let server = listener.incoming().for_each(move |socket| {
        let addr = socket.peer_addr().unwrap();
        info!("[{}] Accepted connection", addr);

        let frame = Framed::new(socket, RpcCodec::new());
        let hs = frame.into_future()
                        .map_err(|(e, _)| e)
                        .and_then(|(data, frame)| {
                            if let Some(data) = data { return Ok((data, frame)) }
                            Err(Error::from(ErrorKind::UnexpectedEof))
                        });
        let hs = hs.and_then(|(data, frame)| {
            if data.id != 0 {
                return Err(Error::new(ErrorKind::InvalidData, "id must be 0"))
            }
            if let Some(msg) = data.msg {
                let peer_type = match msg {
                    RpcMsg::Handshake(peer_type) => peer_type,
                    _ => return Err(Error::new(ErrorKind::InvalidData, "expected handshake msg"))
                };
                Ok((peer_type, frame))
            } else {
                Err(Error::new(ErrorKind::InvalidData, "expected handshake msg"))
            }
        });

        let blockchain = Arc::clone(&blockchain);
        let producer = Arc::clone(&producer);
        let client = hs.and_then(move |(peer_type, frame)| {
            let peer = Peer::new(peer_type, addr, frame);
            debug!("Handshake from client completed: {:?}", peer);
            peer.get_sender().send(RpcPayload {
                id: 0,
                msg: None
            });
            handle_peer(peer, blockchain, producer);
            Ok(())
        });

        ::tokio::spawn(client.map_err(move |e| {
            debug!("[{}] Handshake error: {:?}", addr, e);
        }));
        Ok(())
    }).map_err(|err| {
        error!("Server accept error: {:?}", err);
    });
    ::tokio::spawn(server);
}

fn handle_peer(peer: Peer, blockchain: Arc<Blockchain>, producer: Arc<Option<Producer>>) {
    macro_rules! quick_send {
        ($sender:expr, $rpc:expr, $msg:expr) => {
            $sender.send(RpcPayload {
                id: $rpc.id,
                msg: Some($msg)
            });
        };
        ($sender:expr, $rpc:expr) => {
            $sender.send(RpcPayload {
                id: $rpc.id,
                msg: None
            });
        };
    }

    let (tx, rx) = {
        let (tx, rx) = channel::unbounded();
        (tx, rx.map_err(|_| {
            Error::new(ErrorKind::Other, "rx error")
        }))
    };

    let addr = peer.addr;
    let force_close = Arc::new(AtomicBool::new(false));
    let sender = peer.get_sender();
    ::tokio::spawn(ZipEither::new(peer, rx).take_while({
        let force_close = Arc::clone(&force_close);
        move |_| {
            Ok(!force_close.load(Ordering::Acquire))
        }
    }).for_each(move |(rpc, _)| {
        if let Some(rpc) = rpc {
            if let Some(msg) = rpc.msg {
                match msg {
                    RpcMsg::Event(evt) => {
                        if let Some(producer) = &*producer {
                            match evt {
                                RpcEvent::Block(block) => {
                                    let _ = producer.add_block(block);
                                },
                                RpcEvent::Tx(tx) => {
                                    let _ = producer.add_tx(tx);
                                }
                            }
                        }
                    },
                    RpcMsg::Broadcast(tx) => {
                        if let Some(producer) = &*producer {
                            match producer.add_tx(tx) {
                                Ok(_) => {
                                    quick_send!(sender, rpc);
                                },
                                Err(s) => {
                                    quick_send!(sender, rpc, RpcMsg::Error(s));
                                }
                            }
                        }
                    },
                    RpcMsg::Properties(var) => {
                        match var.request() {
                            Some(_) => {
                                let props = blockchain.get_properties();
                                quick_send!(sender, rpc, RpcMsg::Properties(RpcVariant::Res(props)));
                            },
                            None => {
                                warn!("[{}] Property responses not supported", addr);

                            }
                        }
                    },
                    RpcMsg::Block(var) => {
                        match var.request() {
                            Some(height) => {
                                let block = match blockchain.get_block(*height) {
                                    Some(block) => Some((&*block).clone()),
                                    None => None
                                };
                                quick_send!(sender, rpc, RpcMsg::Block(RpcVariant::Res(block)));
                            },
                            None => {
                                warn!("[{}] Block responses not supported", addr);
                            }
                        }
                    },
                    RpcMsg::Balance(var) => {
                        match var.request() {
                            Some(addr) => {
                                let bal = blockchain.get_balance(addr);
                                quick_send!(sender, rpc, RpcMsg::Balance(RpcVariant::Res(bal)));
                            },
                            None => {
                                warn!("[{}] Balance responses not supported", addr);
                            }
                        }
                    },
                    RpcMsg::TotalFee(var) => {
                        match var.request() {
                            Some(addr) => {
                                let fee = blockchain.get_total_fee(addr);
                                match fee {
                                    Some(fee) => {
                                        quick_send!(sender, rpc, RpcMsg::TotalFee(RpcVariant::Res(fee)));
                                    },
                                    None => {
                                        let err = "failed to retrieve total fee".to_string();
                                        quick_send!(sender, rpc, RpcMsg::Error(err));
                                    }
                                }
                            },
                            None => {
                                warn!("[{}] Total fee responses not supported", addr);
                            }
                        }
                    },
                    _ => {
                        warn!("[{}] Invalid rpc message sent from peer", addr);
                        force_close.store(true, Ordering::Release);
                        tx.send(())
                    }
                }
            }
        }
        Ok(())
    }).and_then(move |_| {
        warn!("[{}] Client disconnected", addr);
        Ok(())
    }).map_err(move |e| {
        debug!("[{}] Error handling frame from client: {:?}", addr, e);
    }));
}
<commit_msg>Cleanup<commit_after>use super::rpc::codec::RpcCodec;

use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
use std::io::{Error, ErrorKind};
use std::net::SocketAddr;

use tokio::net::TcpListener;
use tokio_codec::Framed;
use tokio::prelude::*;

use blockchain::Blockchain;
use producer::Producer;
use net::peer::*;
use net::rpc::*;
use fut_util::*;

pub fn start(addr: SocketAddr, blockchain: Arc<Blockchain>, producer: Arc<Option<Producer>>) {
    let listener = TcpListener::bind(&addr).unwrap();
    info!("Server binded to {:?}", &addr);
    let server = listener.incoming().for_each(move |socket| {
        let addr = socket.peer_addr().unwrap();
        info!("[{}] Accepted connection", addr);

        let frame = Framed::new(socket, RpcCodec::new());
        let hs = frame.into_future()
                        .map_err(|(e, _)| e)
                        .and_then(|(data, frame)| {
                            if let Some(data) = data { return Ok((data, frame)) }
                            Err(Error::from(ErrorKind::UnexpectedEof))
                        });
        let hs = hs.and_then(|(data, frame)| {
            if data.id != 0 {
                return Err(Error::new(ErrorKind::InvalidData, "id must be 0"))
            }
            if let Some(msg) = data.msg {
                let peer_type = match msg {
                    RpcMsg::Handshake(peer_type) => peer_type,
                    _ => return Err(Error::new(ErrorKind::InvalidData, "expected handshake msg"))
                };
                Ok((peer_type, frame))
            } else {
                Err(Error::new(ErrorKind::InvalidData, "expected handshake msg"))
            }
        });

        let blockchain = Arc::clone(&blockchain);
        let producer = Arc::clone(&producer);
        let client = hs.and_then(move |(peer_type, frame)| {
            let peer = Peer::new(peer_type, addr, frame);
            debug!("Handshake from client completed: {:?}", peer);
            peer.get_sender().send(RpcPayload {
                id: 0,
                msg: None
            });
            handle_peer(peer, blockchain, producer);
            Ok(())
        });

        ::tokio::spawn(client.map_err(move |e| {
            debug!("[{}] Handshake error: {:?}", addr, e);
        }));
        Ok(())
    }).map_err(|err| {
        error!("Server accept error: {:?}", err);
    });
    ::tokio::spawn(server);
}

fn handle_peer(peer: Peer, blockchain: Arc<Blockchain>, producer: Arc<Option<Producer>>) {
    macro_rules! quick_send {
        ($sender:expr, $rpc:expr, $msg:expr) => {
            $sender.send(RpcPayload { id: $rpc.id, msg: Some($msg) });
        };
        ($sender:expr, $rpc:expr) => {
            $sender.send(RpcPayload { id: $rpc.id, msg: None });
        };
    }

    let (tx, rx) = {
        let (tx, rx) = channel::unbounded();
        (tx, rx.map_err(|_| {
            Error::new(ErrorKind::Other, "rx error")
        }))
    };

    let addr = peer.addr;
    let force_close = Arc::new(AtomicBool::new(false));
    let sender = peer.get_sender();
    ::tokio::spawn(ZipEither::new(peer, rx).take_while({
        let force_close = Arc::clone(&force_close);
        move |_| {
            Ok(!force_close.load(Ordering::Acquire))
        }
    }).for_each(move |(rpc, _)| {
        if let Some(rpc) = rpc {
            if let Some(msg) = rpc.msg {
                match msg {
                    RpcMsg::Event(evt) => {
                        if let Some(producer) = &*producer {
                            match evt {
                                RpcEvent::Block(block) => {
                                    let _ = producer.add_block(block);
                                },
                                RpcEvent::Tx(tx) => {
                                    let _ = producer.add_tx(tx);
                                }
                            }
                        }
                    },
                    RpcMsg::Broadcast(tx) => {
                        if let Some(producer) = &*producer {
                            match producer.add_tx(tx) {
                                Ok(_) => {
                                    quick_send!(sender, rpc);
                                },
                                Err(s) => {
                                    quick_send!(sender, rpc, RpcMsg::Error(s));
                                }
                            }
                        }
                    },
                    RpcMsg::Properties(var) => {
                        match var.request() {
                            Some(_) => {
                                let props = blockchain.get_properties();
                                quick_send!(sender, rpc, RpcMsg::Properties(RpcVariant::Res(props)));
                            },
                            None => {
                                warn!("[{}] Property responses not supported", addr);
                            }
                        }
                    },
                    RpcMsg::Block(var) => {
                        match var.request() {
                            Some(height) => {
                                let block = match blockchain.get_block(*height) {
                                    Some(block) => Some((&*block).clone()),
                                    None => None
                                };
                                quick_send!(sender, rpc, RpcMsg::Block(RpcVariant::Res(block)));
                            },
                            None => {
                                warn!("[{}] Block responses not supported", addr);
                            }
                        }
                    },
                    RpcMsg::Balance(var) => {
                        match var.request() {
                            Some(addr) => {
                                let bal = blockchain.get_balance(addr);
                                quick_send!(sender, rpc, RpcMsg::Balance(RpcVariant::Res(bal)));
                            },
                            None => {
                                warn!("[{}] Balance responses not supported", addr);
                            }
                        }
                    },
                    RpcMsg::TotalFee(var) => {
                        match var.request() {
                            Some(addr) => {
                                let fee = blockchain.get_total_fee(addr);
                                match fee {
                                    Some(fee) => {
                                        quick_send!(sender, rpc, RpcMsg::TotalFee(RpcVariant::Res(fee)));
                                    },
                                    None => {
                                        let err = "failed to retrieve total fee".to_string();
                                        quick_send!(sender, rpc, RpcMsg::Error(err));
                                    }
                                }
                            },
                            None => {
                                warn!("[{}] Total fee responses not supported", addr);
                            }
                        }
                    },
                    _ => {
                        warn!("[{}] Invalid rpc message sent from peer", addr);
                        force_close.store(true, Ordering::Release);
                        tx.send(())
                    }
                }
            }
        }
        Ok(())
    }).and_then(move |_| {
        warn!("[{}] Client disconnected", addr);
        Ok(())
    }).map_err(move |e| {
        debug!("[{}] Error handling frame from client: {:?}", addr, e);
    }));
}
<|endoftext|>


<commit_before>use hir::db::HirDatabase;
use ra_syntax::{
    ast::{self, AstNode},
    SmolStr, SyntaxElement,
    SyntaxKind::{NAME_REF, USE_ITEM},
    SyntaxNode,
};

use crate::{
    assist_ctx::{ActionBuilder, Assist, AssistCtx},
    auto_import_text_edit, AssistId, ImportsLocator,
};

// Assist: auto_import
//
// If the name is unresolved, provides all possible imports for it.
//
// ```
// fn main() {
//     let map = HashMap<|>::new();
// }
// ```
// ->
// ```
// use std::collections::HashMap;
//
// fn main() {
//     let map = HashMap<|>::new();
// }
// ```
pub(crate) fn auto_import<F: ImportsLocator>(
    ctx: AssistCtx<impl HirDatabase>,
    imports_locator: &mut F,
) -> Option<Assist> {
    let path: ast::Path = ctx.find_node_at_offset()?;
    let module = path.syntax().ancestors().find_map(ast::Module::cast);
    let position = match module.and_then(|it| it.item_list()) {
        Some(item_list) => item_list.syntax().clone(),
        None => {
            let current_file = path.syntax().ancestors().find_map(ast::SourceFile::cast)?;
            current_file.syntax().clone()
        }
    };
    let source_analyzer = ctx.source_analyzer(&position, None);
    let module_with_name_to_import = source_analyzer.module()?;
    let path_to_import = ctx.covering_element().ancestors().find_map(ast::Path::cast)?;
    if source_analyzer.resolve_path(ctx.db, &path_to_import).is_some() {
        return None;
    }

    let name_to_import = &find_applicable_name_ref(ctx.covering_element())?.syntax().to_string();
    let proposed_imports = imports_locator
        .find_imports(&name_to_import.to_string())
        .into_iter()
        .filter_map(|module_def| module_with_name_to_import.find_use_path(ctx.db, module_def))
        .filter(|use_path| !use_path.segments.is_empty())
        .take(20)
        .map(|import| import.to_string())
        .collect::<std::collections::BTreeSet<_>>();
    if proposed_imports.is_empty() {
        return None;
    }

    ctx.add_assist_group(AssistId("auto_import"), "auto import", || {
        proposed_imports
            .into_iter()
            .map(|import| import_to_action(import, &position, &path_to_import.syntax()))
            .collect()
    })
}

fn find_applicable_name_ref(element: SyntaxElement) -> Option<ast::NameRef> {
    if element.ancestors().find(|ancestor| ancestor.kind() == USE_ITEM).is_some() {
        None
    } else if element.kind() == NAME_REF {
        Some(element.as_node().cloned().and_then(ast::NameRef::cast)?)
    } else {
        let parent = element.parent()?;
        if parent.kind() == NAME_REF {
            Some(ast::NameRef::cast(parent)?)
        } else {
            None
        }
    }
}

fn import_to_action(import: String, position: &SyntaxNode, anchor: &SyntaxNode) -> ActionBuilder {
    let mut action_builder = ActionBuilder::default();
    action_builder.label(format!("Import `{}`", &import));
    auto_import_text_edit(
        position,
        anchor,
        &[SmolStr::new(import)],
        action_builder.text_edit_builder(),
    );
    action_builder
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{
        check_assist_with_imports_locator, check_assist_with_imports_locator_not_applicable,
        TestImportsLocator,
    };

    #[test]
    fn applicable_when_found_an_import() {
        check_assist_with_imports_locator(
            auto_import,
            TestImportsLocator::new,
            r"
            PubStruct<|>

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
            r"
            use PubMod::PubStruct;

            PubStruct<|>

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn applicable_when_found_multiple_imports() {
        check_assist_with_imports_locator(
            auto_import,
            TestImportsLocator::new,
            r"
            PubStruct<|>

            pub mod PubMod1 {
                pub struct PubStruct;
            }
            pub mod PubMod2 {
                pub struct PubStruct;
            }
            pub mod PubMod3 {
                pub struct PubStruct;
            }
            ",
            r"
            use PubMod1::PubStruct;

            PubStruct<|>

            pub mod PubMod1 {
                pub struct PubStruct;
            }
            pub mod PubMod2 {
                pub struct PubStruct;
            }
            pub mod PubMod3 {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_for_already_imported_types() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            use PubMod::PubStruct;

            PubStruct<|>

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_for_types_with_private_paths() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            PrivateStruct<|>

            pub mod PubMod {
                struct PrivateStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_when_no_imports_found() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            "
            PubStruct<|>",
        );
    }

    #[test]
    fn not_applicable_in_import_statements() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            use PubStruct<|>;

            pub mod PubMod {
                pub struct PubStruct;
            }",
        );
    }
}
<commit_msg>Properly select a target for auto importing<commit_after>use hir::db::HirDatabase;
use ra_syntax::{
    ast::{self, AstNode},
    SmolStr,
    SyntaxKind::USE_ITEM,
    SyntaxNode,
};

use crate::{
    assist_ctx::{ActionBuilder, Assist, AssistCtx},
    auto_import_text_edit, AssistId, ImportsLocator,
};

// Assist: auto_import
//
// If the name is unresolved, provides all possible imports for it.
//
// ```
// fn main() {
//     let map = HashMap<|>::new();
// }
// ```
// ->
// ```
// use std::collections::HashMap;
//
// fn main() {
//     let map = HashMap<|>::new();
// }
// ```
pub(crate) fn auto_import<F: ImportsLocator>(
    ctx: AssistCtx<impl HirDatabase>,
    imports_locator: &mut F,
) -> Option<Assist> {
    let path_to_import: ast::Path = ctx.find_node_at_offset()?;
    let path_to_import_syntax = path_to_import.syntax();
    if path_to_import_syntax.ancestors().find(|ancestor| ancestor.kind() == USE_ITEM).is_some() {
        return None;
    }

    let module = path_to_import_syntax.ancestors().find_map(ast::Module::cast);
    let position = match module.and_then(|it| it.item_list()) {
        Some(item_list) => item_list.syntax().clone(),
        None => {
            let current_file = path_to_import_syntax.ancestors().find_map(ast::SourceFile::cast)?;
            current_file.syntax().clone()
        }
    };
    let source_analyzer = ctx.source_analyzer(&position, None);
    let module_with_name_to_import = source_analyzer.module()?;
    if source_analyzer.resolve_path(ctx.db, &path_to_import).is_some() {
        return None;
    }

    let proposed_imports = imports_locator
        .find_imports(&path_to_import_syntax.to_string())
        .into_iter()
        .filter_map(|module_def| module_with_name_to_import.find_use_path(ctx.db, module_def))
        .filter(|use_path| !use_path.segments.is_empty())
        .take(20)
        .map(|import| import.to_string())
        .collect::<std::collections::BTreeSet<_>>();
    if proposed_imports.is_empty() {
        return None;
    }

    ctx.add_assist_group(AssistId("auto_import"), "auto import", || {
        proposed_imports
            .into_iter()
            .map(|import| import_to_action(import, &position, &path_to_import_syntax))
            .collect()
    })
}

fn import_to_action(import: String, position: &SyntaxNode, anchor: &SyntaxNode) -> ActionBuilder {
    let mut action_builder = ActionBuilder::default();
    action_builder.label(format!("Import `{}`", &import));
    auto_import_text_edit(
        position,
        anchor,
        &[SmolStr::new(import)],
        action_builder.text_edit_builder(),
    );
    action_builder
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{
        check_assist_with_imports_locator, check_assist_with_imports_locator_not_applicable,
        TestImportsLocator,
    };

    #[test]
    fn applicable_when_found_an_import() {
        check_assist_with_imports_locator(
            auto_import,
            TestImportsLocator::new,
            r"
            <|>PubStruct

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
            r"
            <|>use PubMod::PubStruct;

            PubStruct

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn applicable_when_found_multiple_imports() {
        check_assist_with_imports_locator(
            auto_import,
            TestImportsLocator::new,
            r"
            PubSt<|>ruct

            pub mod PubMod1 {
                pub struct PubStruct;
            }
            pub mod PubMod2 {
                pub struct PubStruct;
            }
            pub mod PubMod3 {
                pub struct PubStruct;
            }
            ",
            r"
            use PubMod1::PubStruct;

            PubSt<|>ruct

            pub mod PubMod1 {
                pub struct PubStruct;
            }
            pub mod PubMod2 {
                pub struct PubStruct;
            }
            pub mod PubMod3 {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_for_already_imported_types() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            use PubMod::PubStruct;

            PubStruct<|>

            pub mod PubMod {
                pub struct PubStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_for_types_with_private_paths() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            PrivateStruct<|>

            pub mod PubMod {
                struct PrivateStruct;
            }
            ",
        );
    }

    #[test]
    fn not_applicable_when_no_imports_found() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            "
            PubStruct<|>",
        );
    }

    #[test]
    fn not_applicable_in_import_statements() {
        check_assist_with_imports_locator_not_applicable(
            auto_import,
            TestImportsLocator::new,
            r"
            use PubStruct<|>;

            pub mod PubMod {
                pub struct PubStruct;
            }",
        );
    }
}
<|endoftext|>


<commit_before>use action;
use settings;
use glfw;
use Settings = piston::game::Settings;
use Game = piston::game::Game;
use piston::gl::Gl;
use graphics;
use graphics::*;
use object;
use Object = object::Object;
use text;
use game_state;

pub struct SnakeApp {
    settings: Settings,
    // Tells where the surface is.
    surface_y: Option<f64>,
    game_state: Option<game_state::GameState>,
    player_index: Option<uint>,
    blood_bar_index: Option<uint>,
    // Contains the game objects.
    objects: Vec<Object>,
}

impl Game for SnakeApp {
    fn get_settings<'a>(&'a self) -> &'a Settings { &self.settings }

    fn render(&self, c: &graphics::Context, gl: &mut Gl) {
        for obj in self.objects.iter() {
            obj.render(c, gl);
        }
    }

    fn update(&mut self, dt: f64) {
        if self.player_index == None { return; }

        // Update states of objects.
        let player_index = self.player_index.unwrap();
        let player_pos = self.objects.get(player_index).pos;
        let mut attack_damage: f64 = 0.0;
        for obj in self.objects.mut_iter() {
            match obj.update(dt, player_pos) {
                action::Passive => {},
                action::Attack(attack) => { attack_damage += attack; },
            }
        }

        // Decrease the players life with attacks.
        *self.objects.get_mut(player_index).blood_mut().unwrap() -= attack_damage;

        if self.blood_bar_index == None { return; }

        // Show blood.
        let player_blood = self.objects.get(player_index).blood().unwrap();
        let blood_bar_index = self.blood_bar_index.unwrap();
        let blood_bar = self.objects.get_mut(blood_bar_index);
        match blood_bar.data {
            object::BarData(ref mut bar) => {
                bar.value = player_blood;
            },
            _ => {},
        }
    }

    fn load(&mut self) {
        self.surface_y = Some(settings::SURFACE_Y);
        self.game_state = Some(settings::INITIAL_GAME_STATE);

        // Add player.
        self.objects.push(Object::player(
            settings::ORIGIN,
            settings::BLUE,
            settings::PLAYER_INITIAL_BLOOD,
            [settings::PLAYER_SPEED_LEFT, settings::PLAYER_SPEED_RIGHT],
            [settings::PLAYER_SPEED_UP, settings::PLAYER_SPEED_DOWN]
        ));
        self.player_index = Some(0);

        // Add blood and air bar.
        self.add_bars();

        // Add sharks.
        self.add_sharks();
    }

    fn key_press(&mut self, key: glfw::Key) {
        // TEST
        // println!("Key pressed {}", key);

        match (key, self.player_index) {
            (glfw::KeyRight, Some(player_index)) => {
                self.objects.get_mut(player_index).move_right();
            },
            (glfw::KeyUp, Some(player_index)) => {
                self.objects.get_mut(player_index).move_up();
            },
            (glfw::KeyLeft, Some(player_index)) => {
                self.objects.get_mut(player_index).move_left();
            },
            (glfw::KeyDown, Some(player_index)) => {
                self.objects.get_mut(player_index).move_down();
            },
            _ => {},
        }
    }

    fn key_release(&mut self, key: glfw::Key) {
        // TEST
        // println!("Key released {}", key);
    }
}

impl SnakeApp {
    pub fn new() -> SnakeApp {
        let exit_on_esc = true;
        let background_color = [1.0, 1.0, 1.0, 1.0];
        SnakeApp {
            settings: Settings::new(exit_on_esc, background_color),
            game_state: None,
            surface_y: None,
            objects: Vec::new(),
            player_index: None,
            blood_bar_index: None,
        }
    }

    pub fn add_bars(&mut self) {
        self.objects.push(Object::bar(
            settings::AIR_BAR_POS,
            "air",
            settings::AIR_BAR_TEXT_COLOR,
            settings::AIR_BAR_BACKGROUND_COLOR,
            settings::AIR_BAR_BAR_COLOR,
            settings::AIR_BAR_INITIAL_VALUE
        ));
        self.objects.push(Object::bar(
            settings::BLOOD_BAR_POS,
            "blood",
            settings::BLOOD_BAR_TEXT_COLOR,
            settings::BLOOD_BAR_BACKGROUND_COLOR,
            settings::BLOOD_BAR_BAR_COLOR,
            settings::BLOOD_BAR_INITIAL_VALUE
        ));
        self.blood_bar_index = Some(self.objects.len() - 1);
    }

    pub fn add_sharks(&mut self) {
        self.objects.push(Object::shark(settings::SHARK_1_POS, settings::SHARK_1_SETTINGS));
    }
}

<commit_msg>When player reaches surface, win<commit_after>use action;
use settings;
use glfw;
use Settings = piston::game::Settings;
use Game = piston::game::Game;
use piston::gl::Gl;
use graphics;
use graphics::*;
use object;
use Object = object::Object;
use text;
use game_state;

pub struct SnakeApp {
    settings: Settings,
    // Tells where the surface is.
    surface_y: Option<f64>,
    game_state: Option<game_state::GameState>,
    player_index: Option<uint>,
    blood_bar_index: Option<uint>,
    // Contains the game objects.
    objects: Vec<Object>,
}

impl Game for SnakeApp {
    fn get_settings<'a>(&'a self) -> &'a Settings { &self.settings }

    fn render(&self, c: &graphics::Context, gl: &mut Gl) {
        for obj in self.objects.iter() {
            obj.render(c, gl);
        }
    }

    fn update(&mut self, dt: f64) {
        if self.player_index == None { return; }

        // Update states of objects.
        let player_index = self.player_index.unwrap();
        let player_pos = self.objects.get(player_index).pos;
        let mut attack_damage: f64 = 0.0;
        for obj in self.objects.mut_iter() {
            match obj.update(dt, player_pos) {
                action::Passive => {},
                action::Attack(attack) => { attack_damage += attack; },
            }
        }

        // When player reaches surface, win.
        if player_pos[1] >= self.surface_y.unwrap() {
            self.game_state = Some(game_state::Win);
            return;
        }

        // Decrease the players life with attacks.
        *self.objects.get_mut(player_index).blood_mut().unwrap() -= attack_damage;

        if self.blood_bar_index == None { return; }

        // Show blood.
        let player_blood = self.objects.get(player_index).blood().unwrap();
        let blood_bar_index = self.blood_bar_index.unwrap();
        let blood_bar = self.objects.get_mut(blood_bar_index);
        match blood_bar.data {
            object::BarData(ref mut bar) => {
                bar.value = player_blood;
            },
            _ => {},
        }
    }

    fn load(&mut self) {
        self.surface_y = Some(settings::SURFACE_Y);
        self.game_state = Some(settings::INITIAL_GAME_STATE);

        // Add player.
        self.objects.push(Object::player(
            settings::ORIGIN,
            settings::BLUE,
            settings::PLAYER_INITIAL_BLOOD,
            [settings::PLAYER_SPEED_LEFT, settings::PLAYER_SPEED_RIGHT],
            [settings::PLAYER_SPEED_UP, settings::PLAYER_SPEED_DOWN]
        ));
        self.player_index = Some(0);

        // Add blood and air bar.
        self.add_bars();

        // Add sharks.
        self.add_sharks();
    }

    fn key_press(&mut self, key: glfw::Key) {
        // TEST
        // println!("Key pressed {}", key);

        match (key, self.player_index) {
            (glfw::KeyRight, Some(player_index)) => {
                self.objects.get_mut(player_index).move_right();
            },
            (glfw::KeyUp, Some(player_index)) => {
                self.objects.get_mut(player_index).move_up();
            },
            (glfw::KeyLeft, Some(player_index)) => {
                self.objects.get_mut(player_index).move_left();
            },
            (glfw::KeyDown, Some(player_index)) => {
                self.objects.get_mut(player_index).move_down();
            },
            _ => {},
        }
    }

    fn key_release(&mut self, key: glfw::Key) {
        // TEST
        // println!("Key released {}", key);
    }
}

impl SnakeApp {
    pub fn new() -> SnakeApp {
        let exit_on_esc = true;
        let background_color = [1.0, 1.0, 1.0, 1.0];
        SnakeApp {
            settings: Settings::new(exit_on_esc, background_color),
            game_state: None,
            surface_y: None,
            objects: Vec::new(),
            player_index: None,
            blood_bar_index: None,
        }
    }

    pub fn add_bars(&mut self) {
        self.objects.push(Object::bar(
            settings::AIR_BAR_POS,
            "air",
            settings::AIR_BAR_TEXT_COLOR,
            settings::AIR_BAR_BACKGROUND_COLOR,
            settings::AIR_BAR_BAR_COLOR,
            settings::AIR_BAR_INITIAL_VALUE
        ));
        self.objects.push(Object::bar(
            settings::BLOOD_BAR_POS,
            "blood",
            settings::BLOOD_BAR_TEXT_COLOR,
            settings::BLOOD_BAR_BACKGROUND_COLOR,
            settings::BLOOD_BAR_BAR_COLOR,
            settings::BLOOD_BAR_INITIAL_VALUE
        ));
        self.blood_bar_index = Some(self.objects.len() - 1);
    }

    pub fn add_sharks(&mut self) {
        self.objects.push(Object::shark(settings::SHARK_1_POS, settings::SHARK_1_SETTINGS));
    }
}

<|endoftext|>


<commit_before>use core::ops::Drop;

use common::memory::*;

use graphics::color::*;
use graphics::size::*;

pub struct BMP {
    pub data: usize,
    pub size: Size
}

impl BMP {
    pub fn new() -> BMP {
        BMP {
            data: 0,
            size: Size { width: 0, height: 0}
        }
    }

    pub fn from_data(file_data: usize) -> BMP {
        let data;
        let size;
        unsafe {
            if file_data > 0
                && *(file_data as *const u8) == 'B' as u8
                && *((file_data + 1) as *const u8) == 'M' as u8
            {
                let file_size = *((file_data + 0x2) as *const u32) as usize;
                let offset = *((file_data + 0xA) as *const u32) as usize;
                let width = *((file_data + 0x12) as *const u32) as usize;
                let height = *((file_data + 0x16) as *const u32) as usize;
                let depth = *((file_data + 0x1C) as *const u16) as usize;

                let bytes = (depth + 7)/8;
                let row_bytes = (depth * width + 31)/32 * 4;

                data = alloc(width * height * 4);
                size = Size {
                    width: width,
                    height: height
                };
                for y in 0..height {
                    for x in 0..width {
                        let pixel_offset = offset + (height - y - 1) * row_bytes + x * bytes;

                        let pixel_data;
                        if pixel_offset < file_size {
                            pixel_data = *((file_data + pixel_offset) as *const u32);
                        }else{
                            pixel_data = 0;
                        }

                        if bytes == 3 {
                            *((data + (y*width + x)*4) as *mut u32) = Color::new((pixel_data >> 16) as u8, (pixel_data >> 8) as u8, pixel_data as u8).data;
                        }else if bytes == 4 {
                            *((data + (y*width + x)*4) as *mut u32) = Color::alpha((pixel_data >> 24) as u8, (pixel_data >> 16) as u8, (pixel_data >> 8) as u8, pixel_data as u8).data;
                        }
                    }
                }
            }else{
                data = 0;
                size = Size {
                    width: 0,
                    height: 0
                };
            }
        }

        return BMP {
            data: data,
            size: size
        };
    }
}

impl Drop for BMP {
    fn drop(&mut self){
        unsafe {
            if self.data > 0 {
                unalloc(self.data);
                self.data = 0;
                self.size = Size {
                    width: 0,
                    height: 0
                };
            }
        }
    }
}
<commit_msg>Fix loading BMP files<commit_after>use core::ops::Drop;
use core::ptr;

use common::debug::*;
use common::memory::*;

use graphics::color::*;
use graphics::size::*;

pub struct BMP {
    pub data: usize,
    pub size: Size
}

impl BMP {
    pub fn new() -> BMP {
        BMP {
            data: 0,
            size: Size { width: 0, height: 0}
        }
    }

    pub fn from_data(file_data: usize) -> BMP {
        let data;
        let size;
        unsafe {
            if file_data > 0
                && *(file_data as *const u8) == 'B' as u8
                && ptr::read((file_data + 1) as *const u8) == 'M' as u8
            {
                let file_size = ptr::read((file_data + 0x2) as *const u32) as usize;
                let offset = ptr::read((file_data + 0xA) as *const u32) as usize;
                let header_size = ptr::read((file_data + 0xE) as *const u32) as usize;
                d("Size ");
                dd(header_size);
                dl();
                let width = ptr::read((file_data + 0x12) as *const u32) as usize;
                d("Width ");
                dd(width);
                dl();
                let height = ptr::read((file_data + 0x16) as *const u32) as usize;
                d("Height ");
                dd(height);
                dl();
                let depth = ptr::read((file_data + 0x1C) as *const u16) as usize;
                d("Depth ");
                dd(depth);
                dl();

                let bytes = (depth + 7)/8;
                let row_bytes = (depth * width + 31)/32 * 4;

                let mut red_mask = 0xFF0000;
                let mut green_mask = 0xFF00;
                let mut blue_mask = 0xFF;
                let mut alpha_mask = 0xFF000000;
                if ptr::read((file_data + 0x1E) as *const u32) == 3 {
                    red_mask = ptr::read((file_data + 0x36) as *const u32) as usize;
                    green_mask = ptr::read((file_data + 0x3A) as *const u32) as usize;
                    blue_mask = ptr::read((file_data + 0x3E) as *const u32) as usize;
                    alpha_mask = ptr::read((file_data + 0x42) as *const u32) as usize;
                }

                let mut red_shift = 0;
                while red_mask > 0 && red_shift < 32 && (red_mask >> red_shift) & 1 == 0 {
                    red_shift += 1;
                }

                let mut green_shift = 0;
                while green_mask > 0 && green_shift < 32 && (green_mask >> green_shift) & 1 == 0 {
                    green_shift += 1;
                }

                let mut blue_shift = 0;
                while blue_mask > 0 && blue_shift < 32 && (blue_mask >> blue_shift) & 1 == 0 {
                    blue_shift += 1;
                }

                let mut alpha_shift = 0;
                while alpha_mask > 0 && alpha_shift < 32 && (alpha_mask >> alpha_shift) & 1 == 0 {
                    alpha_shift += 1;
                }

                data = alloc(width * height * 4);
                size = Size {
                    width: width,
                    height: height
                };
                for y in 0..height {
                    for x in 0..width {
                        let pixel_offset = offset + (height - y - 1) * row_bytes + x * bytes;

                        let pixel_data;
                        if pixel_offset < file_size {
                            pixel_data = ptr::read((file_data + pixel_offset) as *const u32) as usize;
                        }else{
                            pixel_data = 0;
                        }

                        let red = ((pixel_data & red_mask) >> red_shift) as u8;
                        let green = ((pixel_data & green_mask) >> green_shift) as u8;
                        let blue = ((pixel_data & blue_mask) >> blue_shift) as u8;
                        if bytes == 3 {
                            ptr::write((data + (y*width + x)*4) as *mut u32, Color::new(red, green, blue).data);
                        }else if bytes == 4 {
                            let alpha = ((pixel_data & alpha_mask) >> alpha_shift) as u8;
                            ptr::write((data + (y*width + x)*4) as *mut u32, Color::alpha(red, green, blue, alpha).data);
                        }
                    }
                }
            }else{
                data = 0;
                size = Size {
                    width: 0,
                    height: 0
                };
            }
        }

        return BMP {
            data: data,
            size: size
        };
    }
}

impl Drop for BMP {
    fn drop(&mut self){
        unsafe {
            if self.data > 0 {
                unalloc(self.data);
                self.data = 0;
                self.size = Size {
                    width: 0,
                    height: 0
                };
            }
        }
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "use_std")]

use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Iterator;

/// Creates a new `GroupingMap` from `iter`
pub fn new<I, K, V>(iter: I) -> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    GroupingMap { iter }
}

/// `GroupingMap` is an intermediate struct for efficient group-and-fold operations.
/// It groups elements by their key and at the same time fold each group
/// using some aggregating operation.
///
/// No method on this struct performs temporary allocations.
pub struct GroupingMap<I> {
    iter: I,
}

impl<I, K, V> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    /// This is the generic way to perform any operations on a `GroupingMap`.
    /// It's suggested to use this method only to implement custom operations
    /// when the already provided ones are not enough.
    ///
    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in an `HashMap`.
    ///
    /// The `operation` function is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group if there is currently one;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being aggregated;
    /// If `operation` returns `Some(element)` then the accumulator is updated with `element`,
    /// otherwise the previous accumulation is discarded.
    ///
    /// Return a `HashMap` associating the key of each group with the result of aggregation of the group elements.
    /// If there's no result then there won't be an entry associated to that key.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let data = vec![2, 8, 5, 7, 9, 0, 4, 10];
    /// let lookup = data.into_iter()
    ///     .map(|n| (n % 4, n))
    ///     .into_grouping_map()
    ///     .aggregate(|acc, _, val| {
    ///         if val == 0 || val == 10 {
    ///             None
    ///         } else {
    ///             Some(acc.unwrap_or(0) + val)
    ///         }
    ///     });
    ///
    /// assert_eq!(lookup[&0], 4);        // 0 resets the accumulator so only 4 is summed
    /// assert_eq!(lookup[&1], 14);       // 5 + 9
    /// assert_eq!(lookup.get(&2), None); // 10 resets the accumulator and nothing is summed afterward
    /// assert_eq!(lookup[&3], 7);        // 7
    /// assert_eq!(lookup.len(), 3);      // The final keys are only 0, 1 and 2
    /// ```
    pub fn aggregate<FO, R>(self, mut operation: FO) -> HashMap<K, R>
    where
        FO: FnMut(Option<R>, &K, V) -> Option<R>,
    {
        let mut destination_map = HashMap::new();

        for (key, val) in self.iter {
            let acc = destination_map.remove(&key);
            if let Some(op_res) = operation(acc, &key, val) {
                destination_map.insert(key, op_res);
            }
        }

        destination_map
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// `init` is the value from which will be cloned the initial value of each accumulator.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold(0, |acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);   // 3 + 6
    /// assert_eq!(lookup[&1], 12);  // 1 + 4 + 7
    /// assert_eq!(lookup[&2], 7);   // 2 + 5
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold<FO, R>(self, init: R, mut operation: FO) -> HashMap<K, R>
    where
        R: Clone,
        FO: FnMut(R, &K, V) -> R,
    {
        self.aggregate(|acc, key, val| {
            let acc = acc.unwrap_or_else(|| init.clone());
            Some(operation(acc, key, val))
        })
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// This is similar to [`fold`] but the initial value of the accumulator is the first element of the group.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// [`fold`]: #tymethod.fold
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold_first(|acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);   // 3 + 6
    /// assert_eq!(lookup[&1], 12);  // 1 + 4 + 7
    /// assert_eq!(lookup[&2], 7);   // 2 + 5
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold_first<FO>(self, mut operation: FO) -> HashMap<K, V>
    where
        FO: FnMut(V, &K, V) -> V,
    {
        self.aggregate(|acc, key, val| {
            Some(match acc {
                Some(acc) => operation(acc, key, val),
                None => val,
            })
        })
    }

    /// Groups elements from the `GroupingMap` source by key and collects the elements of each group in
    /// an instance of `C`. The iteration order is preserved when inserting elements.
    ///
    /// Return a `HashMap` associating the key of each group with the collection containing the elements of that group.
    pub fn collect<C>(self) -> HashMap<K, C>
    where
        C: Default + Extend<V>,
    {
        self.aggregate(|acc, _, v| {
            let mut acc = acc.unwrap_or_else(C::default);
            acc.extend(Some(v));
            Some(acc)
        })
    }

    /// Groups elements from the `GroupingMap` source by key and counts them.
    ///
    /// Return a `HashMap` associating the key of each group with the number of elements in that group.
    pub fn count(self) -> HashMap<K, usize> {
        self.fold(0, |acc, _, _| acc + 1)
    }
}
<commit_msg>Added doctest for GroupingMap::collect<commit_after>#![cfg(feature = "use_std")]

use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Iterator;

/// Creates a new `GroupingMap` from `iter`
pub fn new<I, K, V>(iter: I) -> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    GroupingMap { iter }
}

/// `GroupingMap` is an intermediate struct for efficient group-and-fold operations.
/// It groups elements by their key and at the same time fold each group
/// using some aggregating operation.
///
/// No method on this struct performs temporary allocations.
pub struct GroupingMap<I> {
    iter: I,
}

impl<I, K, V> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    /// This is the generic way to perform any operations on a `GroupingMap`.
    /// It's suggested to use this method only to implement custom operations
    /// when the already provided ones are not enough.
    ///
    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in an `HashMap`.
    ///
    /// The `operation` function is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group if there is currently one;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being aggregated;
    /// If `operation` returns `Some(element)` then the accumulator is updated with `element`,
    /// otherwise the previous accumulation is discarded.
    ///
    /// Return a `HashMap` associating the key of each group with the result of aggregation of the group elements.
    /// If there's no result then there won't be an entry associated to that key.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let data = vec![2, 8, 5, 7, 9, 0, 4, 10];
    /// let lookup = data.into_iter()
    ///     .map(|n| (n % 4, n))
    ///     .into_grouping_map()
    ///     .aggregate(|acc, _, val| {
    ///         if val == 0 || val == 10 {
    ///             None
    ///         } else {
    ///             Some(acc.unwrap_or(0) + val)
    ///         }
    ///     });
    ///
    /// assert_eq!(lookup[&0], 4);        // 0 resets the accumulator so only 4 is summed
    /// assert_eq!(lookup[&1], 14);       // 5 + 9
    /// assert_eq!(lookup.get(&2), None); // 10 resets the accumulator and nothing is summed afterward
    /// assert_eq!(lookup[&3], 7);        // 7
    /// assert_eq!(lookup.len(), 3);      // The final keys are only 0, 1 and 2
    /// ```
    pub fn aggregate<FO, R>(self, mut operation: FO) -> HashMap<K, R>
    where
        FO: FnMut(Option<R>, &K, V) -> Option<R>,
    {
        let mut destination_map = HashMap::new();

        for (key, val) in self.iter {
            let acc = destination_map.remove(&key);
            if let Some(op_res) = operation(acc, &key, val) {
                destination_map.insert(key, op_res);
            }
        }

        destination_map
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// `init` is the value from which will be cloned the initial value of each accumulator.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold(0, |acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);   // 3 + 6
    /// assert_eq!(lookup[&1], 12);  // 1 + 4 + 7
    /// assert_eq!(lookup[&2], 7);   // 2 + 5
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold<FO, R>(self, init: R, mut operation: FO) -> HashMap<K, R>
    where
        R: Clone,
        FO: FnMut(R, &K, V) -> R,
    {
        self.aggregate(|acc, key, val| {
            let acc = acc.unwrap_or_else(|| init.clone());
            Some(operation(acc, key, val))
        })
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// This is similar to [`fold`] but the initial value of the accumulator is the first element of the group.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// [`fold`]: #tymethod.fold
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold_first(|acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);   // 3 + 6
    /// assert_eq!(lookup[&1], 12);  // 1 + 4 + 7
    /// assert_eq!(lookup[&2], 7);   // 2 + 5
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold_first<FO>(self, mut operation: FO) -> HashMap<K, V>
    where
        FO: FnMut(V, &K, V) -> V,
    {
        self.aggregate(|acc, key, val| {
            Some(match acc {
                Some(acc) => operation(acc, key, val),
                None => val,
            })
        })
    }

    /// Groups elements from the `GroupingMap` source by key and collects the elements of each group in
    /// an instance of `C`. The iteration order is preserved when inserting elements.
    ///
    /// Return a `HashMap` associating the key of each group with the collection containing the elements of that group.
    ///
    /// ```
    /// use itertools::Itertools;
    /// use std::collections::HashSet;
    ///
    /// let lookup = vec![0, 1, 2, 3, 4, 5, 6, 2, 3, 6].into_iter()
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .collect::<HashSet<_>>();
    ///
    /// assert_eq!(lookup[&0], vec![0, 3, 6].into_iter().collect());
    /// assert_eq!(lookup[&1], vec![1, 4].into_iter().collect());
    /// assert_eq!(lookup[&2], vec![2, 5].into_iter().collect());
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn collect<C>(self) -> HashMap<K, C>
    where
        C: Default + Extend<V>,
    {
        self.aggregate(|acc, _, v| {
            let mut acc = acc.unwrap_or_else(C::default);
            acc.extend(Some(v));
            Some(acc)
        })
    }

    /// Groups elements from the `GroupingMap` source by key and counts them.
    ///
    /// Return a `HashMap` associating the key of each group with the number of elements in that group.
    pub fn count(self) -> HashMap<K, usize> {
        self.fold(0, |acc, _, _| acc + 1)
    }
}
<|endoftext|>