<commit_before>//! FIXME: write short doc here

use ra_syntax::{ast, SyntaxKind};

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

pub(super) fn complete_use_tree_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    // complete keyword "crate" in use stmt
    let source_range = ctx.source_range();
    match (ctx.use_item_syntax.as_ref(), ctx.path_prefix.as_ref()) {
        (Some(_), 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);
        }
        (Some(_), Some(_)) => {
            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);
        }
        _ => {}
    }
}

pub(super) fn complete_expr_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    if ctx.token.kind() == SyntaxKind::COMMENT {
        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 && !has_trait_or_impl_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 || ctx.block_expr_parent {
        add_keyword(ctx, acc, "fn", "fn $0() {}");
    }
    if (ctx.has_item_list_or_source_file_parent && !has_trait_or_impl_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 && !has_trait_or_impl_parent {
        add_keyword(ctx, acc, "enum", "enum $0 {}");
        add_keyword(ctx, acc, "struct", "struct $0 {}");
        add_keyword(ctx, acc, "union", "union $0 {}");
    }

    if ctx.block_expr_parent || ctx.is_match_arm {
        add_keyword(ctx, acc, "match", "match $0 {}");
        add_keyword(ctx, acc, "loop", "loop {$0}");
    }
    if ctx.block_expr_parent {
        add_keyword(ctx, acc, "while", "while $0 {}");
    }
    if ctx.if_is_prev || ctx.block_expr_parent {
        add_keyword(ctx, acc, "let", "let ");
    }
    if ctx.if_is_prev || ctx.block_expr_parent || ctx.is_match_arm {
        add_keyword(ctx, acc, "if", "if ");
        add_keyword(ctx, acc, "if let", "if 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 && !has_trait_or_impl_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 || 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 && !has_trait_or_impl_parent)
        || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "static", "static ");
    };
    if (ctx.has_item_list_or_source_file_parent && !has_trait_or_impl_parent)
        || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "extern", "extern ");
    }
    if ctx.has_item_list_or_source_file_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_trait_parent {
        add_keyword(ctx, acc, "pub", "pub ")
    }

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

fn keyword(ctx: &CompletionContext, kw: &str, snippet: &str) -> CompletionItem {
    let res = CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), kw)
        .kind(CompletionItemKind::Keyword);

    match ctx.config.snippet_cap {
        Some(cap) => res.insert_snippet(cap, snippet),
        _ => res.insert_text(if snippet.contains('$') { kw } else { snippet }),
    }
    .build()
}

fn add_keyword(ctx: &CompletionContext, acc: &mut Completions, kw: &str, snippet: &str) {
    acc.add(keyword(ctx, kw, snippet));
}

fn complete_return(
    ctx: &CompletionContext,
    fn_def: &ast::FnDef,
    can_be_stmt: bool,
) -> Option<CompletionItem> {
    let snip = match (can_be_stmt, fn_def.ret_type().is_some()) {
        (true, true) => "return $0;",
        (true, false) => "return;",
        (false, true) => "return $0",
        (false, false) => "return",
    };
    Some(keyword(ctx, "return", snip))
}

#[cfg(test)]
mod tests {
    use crate::completion::{test_utils::completion_list, CompletionKind};
    use insta::assert_snapshot;

    fn get_keyword_completions(code: &str) -> String {
        completion_list(code, CompletionKind::Keyword)
    }

    #[test]
    fn test_keywords_in_use_stmt() {
        assert_snapshot!(
            get_keyword_completions(r"use <|>"),
            @r###"
            kw crate::
            kw self
            kw super::
        "###
        );

        assert_snapshot!(
            get_keyword_completions(r"use a::<|>"),
            @r###"
            kw self
            kw super::
        "###
        );

        assert_snapshot!(
            get_keyword_completions(r"use a::{b, <|>}"),
            @r###"
            kw self
            kw super::
        "###
        );
    }

    #[test]
    fn test_keywords_at_source_file_level() {
        assert_snapshot!(
            get_keyword_completions(r"m<|>"),
            @r###"
            kw const
            kw enum
            kw extern
            kw fn
            kw impl
            kw mod
            kw pub
            kw static
            kw struct
            kw trait
            kw type
            kw union
            kw unsafe
            kw use
        "###
        );
    }

    #[test]
    fn test_keywords_in_function() {
        assert_snapshot!(
            get_keyword_completions(r"fn quux() { <|> }"),
            @r###"
            kw const
            kw extern
            kw fn
            kw if
            kw if let
            kw impl
            kw let
            kw loop
            kw match
            kw mod
            kw return
            kw static
            kw trait
            kw type
            kw unsafe
            kw use
            kw while
        "###
        );
    }

    #[test]
    fn test_keywords_inside_block() {
        assert_snapshot!(
            get_keyword_completions(r"fn quux() { if true { <|> } }"),
            @r###"
            kw const
            kw extern
            kw fn
            kw if
            kw if let
            kw impl
            kw let
            kw loop
            kw match
            kw mod
            kw return
            kw static
            kw trait
            kw type
            kw unsafe
            kw use
            kw while
        "###
        );
    }

    #[test]
    fn test_keywords_after_if() {
        assert_snapshot!(
            get_keyword_completions(
                r"
                fn quux() {
                    if true {
                        ()
                    } <|>
                }
                ",
            ),
            @r###"
            kw const
            kw else
            kw else if
            kw extern
            kw fn
            kw if
            kw if let
            kw impl
            kw let
            kw loop
            kw match
            kw mod
            kw return
            kw static
            kw trait
            kw type
            kw unsafe
            kw use
            kw while
        "###
        );
    }

    #[test]
    fn test_keywords_in_match_arm() {
        assert_snapshot!(
            get_keyword_completions(
                r"
                fn quux() -> i32 {
                    match () {
                        () => <|>
                    }
                }
                ",
            ),
            @r###"
            kw if
            kw if let
            kw loop
            kw match
            kw return
            kw unsafe
        "###
        );
    }

    #[test]
    fn test_keywords_in_trait_def() {
        assert_snapshot!(
            get_keyword_completions(r"trait My { <|> }"),
            @r###"
            kw const
            kw fn
            kw type
            kw unsafe
        "###
        );
    }

    #[test]
    fn test_keywords_in_impl_def() {
        assert_snapshot!(
            get_keyword_completions(r"impl My { <|> }"),
            @r###"
            kw const
            kw fn
            kw pub
            kw type
            kw unsafe
        "###
        );
    }

    #[test]
    fn test_keywords_in_loop() {
        assert_snapshot!(
            get_keyword_completions(r"fn my() { loop { <|> } }"),
            @r###"
            kw break
            kw const
            kw continue
            kw extern
            kw fn
            kw if
            kw if let
            kw impl
            kw let
            kw loop
            kw match
            kw mod
            kw return
            kw static
            kw trait
            kw type
            kw unsafe
            kw use
            kw while
        "###
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_item_list() {
        assert_snapshot!(
            get_keyword_completions(r"unsafe <|>"),
            @r###"
            kw fn
            kw impl
            kw trait
        "###
        );
    }

    #[test]
    fn test_keywords_after_unsafe_in_block_expr() {
        assert_snapshot!(
            get_keyword_completions(r"fn my_fn() { unsafe <|> }"),
            @r###"
            kw fn
            kw impl
            kw trait
        "###
        );
    }

    #[test]
    fn test_mut_in_ref_and_in_fn_parameters_list() {
        assert_snapshot!(
            get_keyword_completions(r"fn my_fn(&<|>) {}"),
            @r###"
            kw mut
        "###
        );
        assert_snapshot!(
            get_keyword_completions(r"fn my_fn(<|>) {}"),
            @r###"
            kw mut
        "###
        );
        assert_snapshot!(
            get_keyword_completions(r"fn my_fn() { let &<|> }"),
            @r###"
            kw mut
        "###
        );
    }

    #[test]
    fn test_where_keyword() {
        assert_snapshot!(
            get_keyword_completions(r"trait A <|>"),
            @r###"
            kw where
        "###
        );
        assert_snapshot!(
            get_keyword_completions(r"impl A <|>"),
            @r###"
            kw where
        "###
        );
    }
}
<commit_msg>Refactor keyword completion tests<commit_after>//! FIXME: write short doc here

use ra_syntax::{ast, SyntaxKind};

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

pub(super) fn complete_use_tree_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    // complete keyword "crate" in use stmt
    let source_range = ctx.source_range();
    match (ctx.use_item_syntax.as_ref(), ctx.path_prefix.as_ref()) {
        (Some(_), 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);
        }
        (Some(_), Some(_)) => {
            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);
        }
        _ => {}
    }
}

pub(super) fn complete_expr_keyword(acc: &mut Completions, ctx: &CompletionContext) {
    if ctx.token.kind() == SyntaxKind::COMMENT {
        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 && !has_trait_or_impl_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 || ctx.block_expr_parent {
        add_keyword(ctx, acc, "fn", "fn $0() {}");
    }
    if (ctx.has_item_list_or_source_file_parent && !has_trait_or_impl_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 && !has_trait_or_impl_parent {
        add_keyword(ctx, acc, "enum", "enum $0 {}");
        add_keyword(ctx, acc, "struct", "struct $0 {}");
        add_keyword(ctx, acc, "union", "union $0 {}");
    }

    if ctx.block_expr_parent || ctx.is_match_arm {
        add_keyword(ctx, acc, "match", "match $0 {}");
        add_keyword(ctx, acc, "loop", "loop {$0}");
    }
    if ctx.block_expr_parent {
        add_keyword(ctx, acc, "while", "while $0 {}");
    }
    if ctx.if_is_prev || ctx.block_expr_parent {
        add_keyword(ctx, acc, "let", "let ");
    }
    if ctx.if_is_prev || ctx.block_expr_parent || ctx.is_match_arm {
        add_keyword(ctx, acc, "if", "if ");
        add_keyword(ctx, acc, "if let", "if 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 && !has_trait_or_impl_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 || 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 && !has_trait_or_impl_parent)
        || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "static", "static ");
    };
    if (ctx.has_item_list_or_source_file_parent && !has_trait_or_impl_parent)
        || ctx.block_expr_parent
    {
        add_keyword(ctx, acc, "extern", "extern ");
    }
    if ctx.has_item_list_or_source_file_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_trait_parent {
        add_keyword(ctx, acc, "pub", "pub ")
    }

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

fn keyword(ctx: &CompletionContext, kw: &str, snippet: &str) -> CompletionItem {
    let res = CompletionItem::new(CompletionKind::Keyword, ctx.source_range(), kw)
        .kind(CompletionItemKind::Keyword);

    match ctx.config.snippet_cap {
        Some(cap) => res.insert_snippet(cap, snippet),
        _ => res.insert_text(if snippet.contains('$') { kw } else { snippet }),
    }
    .build()
}

fn add_keyword(ctx: &CompletionContext, acc: &mut Completions, kw: &str, snippet: &str) {
    acc.add(keyword(ctx, kw, snippet));
}

fn complete_return(
    ctx: &CompletionContext,
    fn_def: &ast::FnDef,
    can_be_stmt: bool,
) -> Option<CompletionItem> {
    let snip = match (can_be_stmt, fn_def.ret_type().is_some()) {
        (true, true) => "return $0;",
        (true, false) => "return;",
        (false, true) => "return $0",
        (false, false) => "return",
    };
    Some(keyword(ctx, "return", snip))
}

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

    use crate::completion::{test_utils::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 <|>",
            expect![[r#"
                kw crate::
                kw self
                kw super::
            "#]],
        );

        check(
            r"use a::<|>",
            expect![[r#"
            kw self
            kw super::
        "#]],
        );

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

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

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

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

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

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

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

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

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

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

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

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

    #[test]
    fn test_where_keyword() {
        check(
            r"trait A <|>",
            expect![[r#"
                kw where
            "#]],
        );
        check(
            r"impl A <|>",
            expect![[r#"
                kw where
            "#]],
        );
    }
}
<|endoftext|>


<commit_before>mod vfs_filter;

use std::sync::Arc;
use std::path::Path;
use std::collections::HashSet;

use rustc_hash::FxHashMap;

use ra_db::{
    CrateGraph, FileId, SourceRoot, SourceRootId, SourceDatabase, salsa,
};
use ra_hir::{db, HirInterner};
use ra_project_model::ProjectWorkspace;
use ra_vfs::{Vfs, VfsChange};
use vfs_filter::IncludeRustFiles;

type Result<T> = std::result::Result<T, failure::Error>;

#[salsa::database(ra_db::SourceDatabaseStorage, db::HirDatabaseStorage, db::DefDatabaseStorage)]
#[derive(Debug)]
pub struct BatchDatabase {
    runtime: salsa::Runtime<BatchDatabase>,
    interner: Arc<HirInterner>,
}

impl salsa::Database for BatchDatabase {
    fn salsa_runtime(&self) -> &salsa::Runtime<BatchDatabase> {
        &self.runtime
    }
}

impl AsRef<HirInterner> for BatchDatabase {
    fn as_ref(&self) -> &HirInterner {
        &self.interner
    }
}

fn vfs_file_to_id(f: ra_vfs::VfsFile) -> FileId {
    FileId(f.0.into())
}
fn vfs_root_to_id(r: ra_vfs::VfsRoot) -> SourceRootId {
    SourceRootId(r.0.into())
}

impl BatchDatabase {
    pub fn load(crate_graph: CrateGraph, vfs: &mut Vfs) -> BatchDatabase {
        let mut db =
            BatchDatabase { runtime: salsa::Runtime::default(), interner: Default::default() };
        db.set_crate_graph(Arc::new(crate_graph));

        // wait until Vfs has loaded all roots
        let receiver = vfs.task_receiver().clone();
        let mut roots_loaded = HashSet::new();
        for task in receiver {
            vfs.handle_task(task);
            let mut done = false;
            for change in vfs.commit_changes() {
                match change {
                    VfsChange::AddRoot { root, files } => {
                        let source_root_id = vfs_root_to_id(root);
                        log::debug!(
                            "loaded source root {:?} with path {:?}",
                            source_root_id,
                            vfs.root2path(root)
                        );
                        let mut file_map = FxHashMap::default();
                        for (vfs_file, path, text) in files {
                            let file_id = vfs_file_to_id(vfs_file);
                            db.set_file_text(file_id, text);
                            db.set_file_relative_path(file_id, path.clone());
                            db.set_file_source_root(file_id, source_root_id);
                            file_map.insert(path, file_id);
                        }
                        let source_root = SourceRoot { files: file_map };
                        db.set_source_root(source_root_id, Arc::new(source_root));
                        roots_loaded.insert(source_root_id);
                        if roots_loaded.len() == vfs.n_roots() {
                            done = true;
                        }
                    }
                    VfsChange::AddFile { .. }
                    | VfsChange::RemoveFile { .. }
                    | VfsChange::ChangeFile { .. } => {
                        // We just need the first scan, so just ignore these
                    }
                }
            }
            if done {
                break;
            }
        }

        db
    }

    pub fn load_cargo(root: impl AsRef<Path>) -> Result<(BatchDatabase, Vec<SourceRootId>)> {
        let root = std::env::current_dir()?.join(root);
        let ws = ProjectWorkspace::discover(root.as_ref())?;
        let mut roots = Vec::new();
        roots.push(IncludeRustFiles::member(root.clone()));
        roots.extend(IncludeRustFiles::from_roots(ws.to_roots()));
        let (mut vfs, roots) = Vfs::new(roots);
        let mut load = |path: &Path| {
            let vfs_file = vfs.load(path);
            log::debug!("vfs file {:?} -> {:?}", path, vfs_file);
            vfs_file.map(vfs_file_to_id)
        };
        let crate_graph = ws.to_crate_graph(&mut load);
        log::debug!("crate graph: {:?}", crate_graph);

        let local_roots = roots
            .into_iter()
            .filter(|r| vfs.root2path(*r).starts_with(&root))
            .map(vfs_root_to_id)
            .collect();

        let db = BatchDatabase::load(crate_graph, &mut vfs);
        Ok((db, local_roots))
    }
}

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

    #[test]
    fn test_loading_rust_analyzer() {
        let mut path = std::env::current_exe().unwrap();
        while !path.join("Cargo.toml").is_file() {
            path = path.parent().unwrap().to_owned();
        }
        let (db, roots) = BatchDatabase::load_cargo(path).unwrap();
        let mut n_crates = 0;
        for root in roots {
            for _krate in Crate::source_root_crates(&db, root) {
                n_crates += 1;
            }
        }

        // RA has quite a few crates, but the exact count doesn't matter
        assert!(n_crates > 20);
    }
}
<commit_msg>Make robust about target directory<commit_after>mod vfs_filter;

use std::sync::Arc;
use std::path::Path;
use std::collections::HashSet;

use rustc_hash::FxHashMap;

use ra_db::{
    CrateGraph, FileId, SourceRoot, SourceRootId, SourceDatabase, salsa,
};
use ra_hir::{db, HirInterner};
use ra_project_model::ProjectWorkspace;
use ra_vfs::{Vfs, VfsChange};
use vfs_filter::IncludeRustFiles;

type Result<T> = std::result::Result<T, failure::Error>;

#[salsa::database(ra_db::SourceDatabaseStorage, db::HirDatabaseStorage, db::DefDatabaseStorage)]
#[derive(Debug)]
pub struct BatchDatabase {
    runtime: salsa::Runtime<BatchDatabase>,
    interner: Arc<HirInterner>,
}

impl salsa::Database for BatchDatabase {
    fn salsa_runtime(&self) -> &salsa::Runtime<BatchDatabase> {
        &self.runtime
    }
}

impl AsRef<HirInterner> for BatchDatabase {
    fn as_ref(&self) -> &HirInterner {
        &self.interner
    }
}

fn vfs_file_to_id(f: ra_vfs::VfsFile) -> FileId {
    FileId(f.0.into())
}
fn vfs_root_to_id(r: ra_vfs::VfsRoot) -> SourceRootId {
    SourceRootId(r.0.into())
}

impl BatchDatabase {
    pub fn load(crate_graph: CrateGraph, vfs: &mut Vfs) -> BatchDatabase {
        let mut db =
            BatchDatabase { runtime: salsa::Runtime::default(), interner: Default::default() };
        db.set_crate_graph(Arc::new(crate_graph));

        // wait until Vfs has loaded all roots
        let receiver = vfs.task_receiver().clone();
        let mut roots_loaded = HashSet::new();
        for task in receiver {
            vfs.handle_task(task);
            let mut done = false;
            for change in vfs.commit_changes() {
                match change {
                    VfsChange::AddRoot { root, files } => {
                        let source_root_id = vfs_root_to_id(root);
                        log::debug!(
                            "loaded source root {:?} with path {:?}",
                            source_root_id,
                            vfs.root2path(root)
                        );
                        let mut file_map = FxHashMap::default();
                        for (vfs_file, path, text) in files {
                            let file_id = vfs_file_to_id(vfs_file);
                            db.set_file_text(file_id, text);
                            db.set_file_relative_path(file_id, path.clone());
                            db.set_file_source_root(file_id, source_root_id);
                            file_map.insert(path, file_id);
                        }
                        let source_root = SourceRoot { files: file_map };
                        db.set_source_root(source_root_id, Arc::new(source_root));
                        roots_loaded.insert(source_root_id);
                        if roots_loaded.len() == vfs.n_roots() {
                            done = true;
                        }
                    }
                    VfsChange::AddFile { .. }
                    | VfsChange::RemoveFile { .. }
                    | VfsChange::ChangeFile { .. } => {
                        // We just need the first scan, so just ignore these
                    }
                }
            }
            if done {
                break;
            }
        }

        db
    }

    pub fn load_cargo(root: impl AsRef<Path>) -> Result<(BatchDatabase, Vec<SourceRootId>)> {
        let root = std::env::current_dir()?.join(root);
        let ws = ProjectWorkspace::discover(root.as_ref())?;
        let mut roots = Vec::new();
        roots.push(IncludeRustFiles::member(root.clone()));
        roots.extend(IncludeRustFiles::from_roots(ws.to_roots()));
        let (mut vfs, roots) = Vfs::new(roots);
        let mut load = |path: &Path| {
            let vfs_file = vfs.load(path);
            log::debug!("vfs file {:?} -> {:?}", path, vfs_file);
            vfs_file.map(vfs_file_to_id)
        };
        let crate_graph = ws.to_crate_graph(&mut load);
        log::debug!("crate graph: {:?}", crate_graph);

        let local_roots = roots
            .into_iter()
            .filter(|r| vfs.root2path(*r).starts_with(&root))
            .map(vfs_root_to_id)
            .collect();

        let db = BatchDatabase::load(crate_graph, &mut vfs);
        Ok((db, local_roots))
    }
}

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

    #[test]
    fn test_loading_rust_analyzer() {
        let path = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap().parent().unwrap();
        let (db, roots) = BatchDatabase::load_cargo(path).unwrap();
        let mut n_crates = 0;
        for root in roots {
            for _krate in Crate::source_root_crates(&db, root) {
                n_crates += 1;
            }
        }

        // RA has quite a few crates, but the exact count doesn't matter
        assert!(n_crates > 20);
    }
}
<|endoftext|>


<commit_before>//! Epoch-based memory reclamation.
//!
//! An interesting problem concurrent collections deal with comes from the remove operation.
//! Suppose that a thread removes an element from a lock-free map, while another thread is reading
//! that same element at the same time. The first thread must wait until the second thread stops
//! reading the element. Only then it is safe to destruct it.
//!
//! Programming languages that come with garbage collectors solve this problem trivially. The
//! garbage collector will destruct the removed element when no thread can hold a reference to it
//! anymore.
//!
//! This crate implements a basic memory reclamation mechanism, which is based on epochs. When an
//! element gets removed from a concurrent collection, it is inserted into a pile of garbage and
//! marked with the current epoch. Every time a thread accesses a collection, it checks the current
//! epoch, attempts to increment it, and destructs some garbage that became so old that no thread
//! can be referencing it anymore.
//!
//! That is the general mechanism behind epoch-based memory reclamation, but the details are a bit
//! more complicated. Anyhow, memory reclamation is designed to be fully automatic and something
//! users of concurrent collections don't have to worry much about.
//!
//! # Pointers
//!
//! Concurrent collections are built using atomic pointers. This module provides [`Atomic`], which
//! is just a shared atomic pointer to a heap-allocated object. Loading an [`Atomic`] yields a
//! [`Shared`], which is an epoch-protected pointer through which the loaded object can be safely
//! read.
//!
//! # Pinning
//!
//! Before an [`Atomic`] can be loaded, a participant must be [`pin`]ned. By pinning a participant
//! we declare that any object that gets removed from now on must not be destructed just
//! yet. Garbage collection of newly removed objects is suspended until the participant gets
//! unpinned.
//!
//! # Garbage
//!
//! Objects that get removed from concurrent collections must be stashed away until all currently
//! pinned participants get unpinned. Such objects can be stored into a thread-local or global
//! storage, where they are kept until the right time for their destruction comes.
//!
//! There is a global shared instance of garbage queue. You can [`defer`](Guard::defer) the execution of an
//! arbitrary function until the global epoch is advanced enough. Most notably, concurrent data
//! structures may defer the deallocation of an object.
//!
//! # APIs
//!
//! For majority of use cases, just use the default garbage collector by invoking [`pin`]. If you
//! want to create your own garbage collector, use the [`Collector`] API.

#![doc(test(
    no_crate_inject,
    attr(
        deny(warnings, rust_2018_idioms),
        allow(dead_code, unused_assignments, unused_variables)
    )
))]
#![warn(
    missing_docs,
    missing_debug_implementations,
    rust_2018_idioms,
    unreachable_pub
)]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(feature = "nightly", feature(cfg_target_has_atomic))]
#![cfg_attr(feature = "nightly", feature(const_fn))]

#[cfg(crossbeam_loom)]
extern crate loom_crate as loom;

use cfg_if::cfg_if;

#[cfg(crossbeam_loom)]
#[allow(unused_imports, dead_code)]
mod primitive {
    pub(crate) mod cell {
        pub(crate) use loom::cell::UnsafeCell;
    }
    pub(crate) mod sync {
        pub(crate) mod atomic {
            use core::sync::atomic::Ordering;
            pub(crate) use loom::sync::atomic::AtomicUsize;
            pub(crate) fn fence(ord: Ordering) {
                if let Ordering::Acquire = ord {
                } else {
                    // FIXME: loom only supports acquire fences at the moment.
                    // https://github.com/tokio-rs/loom/issues/117
                    // let's at least not panic...
                    // this may generate some false positives (`SeqCst` is stronger than `Acquire`
                    // for example), and some false negatives (`Relaxed` is weaker than `Acquire`),
                    // but it's the best we can do for the time being.
                }
                loom::sync::atomic::fence(Ordering::Acquire)
            }

            // FIXME: loom does not support compiler_fence at the moment.
            // https://github.com/tokio-rs/loom/issues/117
            // we use fence as a stand-in for compiler_fence for the time being.
            // this may miss some races since fence is stronger than compiler_fence,
            // but it's the best we can do for the time being.
            pub(crate) use self::fence as compiler_fence;
        }
        pub(crate) use loom::sync::Arc;
    }
    pub(crate) use loom::lazy_static;
    pub(crate) use loom::thread_local;
}
#[cfg(not(crossbeam_loom))]
#[allow(unused_imports, dead_code)]
mod primitive {
    #[cfg(any(feature = "alloc", feature = "std"))]
    pub(crate) mod cell {
        #[derive(Debug)]
        #[repr(transparent)]
        pub(crate) struct UnsafeCell<T>(::core::cell::UnsafeCell<T>);

        // loom's UnsafeCell has a slightly different API than the standard library UnsafeCell.
        // Since we want the rest of the code to be agnostic to whether it's running under loom or
        // not, we write this small wrapper that provides the loom-supported API for the standard
        // library UnsafeCell. This is also what the loom documentation recommends:
        // https://github.com/tokio-rs/loom#handling-loom-api-differences
        impl<T> UnsafeCell<T> {
            #[inline]
            pub(crate) fn new(data: T) -> UnsafeCell<T> {
                UnsafeCell(::core::cell::UnsafeCell::new(data))
            }

            #[inline]
            pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R {
                f(self.0.get())
            }

            #[inline]
            pub(crate) fn with_mut<R>(&self, f: impl FnOnce(*mut T) -> R) -> R {
                f(self.0.get())
            }
        }
    }
    #[cfg(any(feature = "alloc", feature = "std"))]
    pub(crate) mod sync {
        pub(crate) mod atomic {
            pub(crate) use core::sync::atomic::compiler_fence;
            pub(crate) use core::sync::atomic::fence;
            pub(crate) use core::sync::atomic::AtomicUsize;
        }
        #[cfg_attr(feature = "nightly", cfg(target_has_atomic = "ptr"))]
        pub(crate) use alloc::sync::Arc;
    }

    #[cfg(feature = "std")]
    pub(crate) use std::thread_local;

    #[cfg(feature = "std")]
    pub(crate) use lazy_static::lazy_static;
}

#[cfg_attr(feature = "nightly", cfg(target_has_atomic = "ptr"))]
cfg_if! {
    if #[cfg(feature = "alloc")] {
        extern crate alloc;

        mod atomic;
        mod collector;
        mod deferred;
        mod epoch;
        mod guard;
        mod internal;
        mod sync;

        pub use self::atomic::{
            Pointable, Atomic, CompareExchangeError,
            Owned, Pointer, Shared,
        };
        pub use self::collector::{Collector, LocalHandle};
        pub use self::guard::{unprotected, Guard};

        #[allow(deprecated)]
        pub use self::atomic::{CompareAndSetError, CompareAndSetOrdering};
    }
}

cfg_if! {
    if #[cfg(feature = "std")] {
        mod default;
        pub use self::default::{default_collector, is_pinned, pin};
    }
}
<commit_msg>Use feature(const_fn_trait_bound) instead of feature(const_fn)<commit_after>//! Epoch-based memory reclamation.
//!
//! An interesting problem concurrent collections deal with comes from the remove operation.
//! Suppose that a thread removes an element from a lock-free map, while another thread is reading
//! that same element at the same time. The first thread must wait until the second thread stops
//! reading the element. Only then it is safe to destruct it.
//!
//! Programming languages that come with garbage collectors solve this problem trivially. The
//! garbage collector will destruct the removed element when no thread can hold a reference to it
//! anymore.
//!
//! This crate implements a basic memory reclamation mechanism, which is based on epochs. When an
//! element gets removed from a concurrent collection, it is inserted into a pile of garbage and
//! marked with the current epoch. Every time a thread accesses a collection, it checks the current
//! epoch, attempts to increment it, and destructs some garbage that became so old that no thread
//! can be referencing it anymore.
//!
//! That is the general mechanism behind epoch-based memory reclamation, but the details are a bit
//! more complicated. Anyhow, memory reclamation is designed to be fully automatic and something
//! users of concurrent collections don't have to worry much about.
//!
//! # Pointers
//!
//! Concurrent collections are built using atomic pointers. This module provides [`Atomic`], which
//! is just a shared atomic pointer to a heap-allocated object. Loading an [`Atomic`] yields a
//! [`Shared`], which is an epoch-protected pointer through which the loaded object can be safely
//! read.
//!
//! # Pinning
//!
//! Before an [`Atomic`] can be loaded, a participant must be [`pin`]ned. By pinning a participant
//! we declare that any object that gets removed from now on must not be destructed just
//! yet. Garbage collection of newly removed objects is suspended until the participant gets
//! unpinned.
//!
//! # Garbage
//!
//! Objects that get removed from concurrent collections must be stashed away until all currently
//! pinned participants get unpinned. Such objects can be stored into a thread-local or global
//! storage, where they are kept until the right time for their destruction comes.
//!
//! There is a global shared instance of garbage queue. You can [`defer`](Guard::defer) the execution of an
//! arbitrary function until the global epoch is advanced enough. Most notably, concurrent data
//! structures may defer the deallocation of an object.
//!
//! # APIs
//!
//! For majority of use cases, just use the default garbage collector by invoking [`pin`]. If you
//! want to create your own garbage collector, use the [`Collector`] API.

#![doc(test(
    no_crate_inject,
    attr(
        deny(warnings, rust_2018_idioms),
        allow(dead_code, unused_assignments, unused_variables)
    )
))]
#![warn(
    missing_docs,
    missing_debug_implementations,
    rust_2018_idioms,
    unreachable_pub
)]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(feature = "nightly", feature(cfg_target_has_atomic))]
#![cfg_attr(feature = "nightly", feature(const_fn_trait_bound))]

#[cfg(crossbeam_loom)]
extern crate loom_crate as loom;

use cfg_if::cfg_if;

#[cfg(crossbeam_loom)]
#[allow(unused_imports, dead_code)]
mod primitive {
    pub(crate) mod cell {
        pub(crate) use loom::cell::UnsafeCell;
    }
    pub(crate) mod sync {
        pub(crate) mod atomic {
            use core::sync::atomic::Ordering;
            pub(crate) use loom::sync::atomic::AtomicUsize;
            pub(crate) fn fence(ord: Ordering) {
                if let Ordering::Acquire = ord {
                } else {
                    // FIXME: loom only supports acquire fences at the moment.
                    // https://github.com/tokio-rs/loom/issues/117
                    // let's at least not panic...
                    // this may generate some false positives (`SeqCst` is stronger than `Acquire`
                    // for example), and some false negatives (`Relaxed` is weaker than `Acquire`),
                    // but it's the best we can do for the time being.
                }
                loom::sync::atomic::fence(Ordering::Acquire)
            }

            // FIXME: loom does not support compiler_fence at the moment.
            // https://github.com/tokio-rs/loom/issues/117
            // we use fence as a stand-in for compiler_fence for the time being.
            // this may miss some races since fence is stronger than compiler_fence,
            // but it's the best we can do for the time being.
            pub(crate) use self::fence as compiler_fence;
        }
        pub(crate) use loom::sync::Arc;
    }
    pub(crate) use loom::lazy_static;
    pub(crate) use loom::thread_local;
}
#[cfg(not(crossbeam_loom))]
#[allow(unused_imports, dead_code)]
mod primitive {
    #[cfg(any(feature = "alloc", feature = "std"))]
    pub(crate) mod cell {
        #[derive(Debug)]
        #[repr(transparent)]
        pub(crate) struct UnsafeCell<T>(::core::cell::UnsafeCell<T>);

        // loom's UnsafeCell has a slightly different API than the standard library UnsafeCell.
        // Since we want the rest of the code to be agnostic to whether it's running under loom or
        // not, we write this small wrapper that provides the loom-supported API for the standard
        // library UnsafeCell. This is also what the loom documentation recommends:
        // https://github.com/tokio-rs/loom#handling-loom-api-differences
        impl<T> UnsafeCell<T> {
            #[inline]
            pub(crate) fn new(data: T) -> UnsafeCell<T> {
                UnsafeCell(::core::cell::UnsafeCell::new(data))
            }

            #[inline]
            pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R {
                f(self.0.get())
            }

            #[inline]
            pub(crate) fn with_mut<R>(&self, f: impl FnOnce(*mut T) -> R) -> R {
                f(self.0.get())
            }
        }
    }
    #[cfg(any(feature = "alloc", feature = "std"))]
    pub(crate) mod sync {
        pub(crate) mod atomic {
            pub(crate) use core::sync::atomic::compiler_fence;
            pub(crate) use core::sync::atomic::fence;
            pub(crate) use core::sync::atomic::AtomicUsize;
        }
        #[cfg_attr(feature = "nightly", cfg(target_has_atomic = "ptr"))]
        pub(crate) use alloc::sync::Arc;
    }

    #[cfg(feature = "std")]
    pub(crate) use std::thread_local;

    #[cfg(feature = "std")]
    pub(crate) use lazy_static::lazy_static;
}

#[cfg_attr(feature = "nightly", cfg(target_has_atomic = "ptr"))]
cfg_if! {
    if #[cfg(feature = "alloc")] {
        extern crate alloc;

        mod atomic;
        mod collector;
        mod deferred;
        mod epoch;
        mod guard;
        mod internal;
        mod sync;

        pub use self::atomic::{
            Pointable, Atomic, CompareExchangeError,
            Owned, Pointer, Shared,
        };
        pub use self::collector::{Collector, LocalHandle};
        pub use self::guard::{unprotected, Guard};

        #[allow(deprecated)]
        pub use self::atomic::{CompareAndSetError, CompareAndSetOrdering};
    }
}

cfg_if! {
    if #[cfg(feature = "std")] {
        mod default;
        pub use self::default::{default_collector, is_pinned, pin};
    }
}
<|endoftext|>


<commit_before>//! An MPSC with futures.

use std::{error, fmt};
use std::any::Any;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize};
use std::sync::atomic::Ordering::SeqCst;
use crossbeam::sync::{AtomicOption, MsQueue};
use futures::Async;
use futures::stream::Stream;
use futures::task::{Task, park};
use void::Void;

//------------ channel ------------------------------------------------------

pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    Inner::new()
}

//------------ Inner --------------------------------------------------------

struct Inner<T> {
    queue: MsQueue<T>,
    tx_count: AtomicUsize,
    rx_dropped: AtomicBool,
    rx_task: AtomicOption<Task>,
}

impl<T> Inner<T> {
    fn new() -> (Sender<T>, Receiver<T>) {
        let inner = Arc::new(Inner {
            queue: MsQueue::new(),
            tx_count: AtomicUsize::new(1),
            rx_dropped: AtomicBool::new(false),
            rx_task: AtomicOption::new()
        });
        (Sender{inner: inner.clone()}, Receiver{inner: inner})
    }
}

//------------ Sender -------------------------------------------------------

pub struct Sender<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Sender<T> {
    pub fn send(&self, t: T) -> Result<(), T> {
        if self.inner.rx_dropped.load(SeqCst) {
            Err(t)
        }
        else {
            self.inner.queue.push(t);
            if let Some(task) = self.inner.rx_task.take(SeqCst) {
                task.unpark()
            }
            Ok(())
        }
    }
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Self {
        let inner = self.inner.clone();
        let _ = inner.tx_count.fetch_add(1, SeqCst);
        Sender{inner: inner}
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        let _ = self.inner.tx_count.fetch_sub(1, SeqCst);
    }
}

//------------ Receiver ------------------------------------------------------

pub struct Receiver<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Stream for Receiver<T> {
    type Item = T;
    type Error = Void;

    fn poll(&mut self) -> Result<Async<Option<T>>, Void> {
        if self.inner.tx_count.load(SeqCst) == 0 {
            Ok(Async::Ready(None))
        }
        else if let Some(t) = self.inner.queue.try_pop() {
            Ok(Async::Ready(Some(t)))
        }
        else {
            self.inner.rx_task.swap(park(), SeqCst);
            Ok(Async::NotReady)
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.rx_dropped.store(true, SeqCst)
    }
}

//------------ SendError ----------------------------------------------------

pub struct SendError<T>(T);

impl<T> fmt::Debug for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_tuple("SendError")
            .field(&"...")
            .finish()
    }
}

impl<T> fmt::Display for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "send failed because receiver is gone")
    }
}

impl<T: Any> error::Error for SendError<T> {
    fn description(&self) -> &str {
        "send failed because receiver is gone"
    }
}

<commit_msg>Test and fix for bit::net::mpsc.<commit_after>//! An MPSC with futures.

use std::{error, fmt};
use std::any::Any;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize};
use std::sync::atomic::Ordering::SeqCst;
use crossbeam::sync::{AtomicOption, MsQueue};
use futures::Async;
use futures::stream::Stream;
use futures::task::{Task, park};
use void::Void;

//------------ channel ------------------------------------------------------

pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    Inner::new()
}

//------------ Inner --------------------------------------------------------

struct Inner<T> {
    queue: MsQueue<T>,
    tx_count: AtomicUsize,
    rx_dropped: AtomicBool,
    rx_task: AtomicOption<Task>,
}

impl<T> Inner<T> {
    fn new() -> (Sender<T>, Receiver<T>) {
        let inner = Arc::new(Inner {
            queue: MsQueue::new(),
            tx_count: AtomicUsize::new(1),
            rx_dropped: AtomicBool::new(false),
            rx_task: AtomicOption::new()
        });
        (Sender{inner: inner.clone()}, Receiver{inner: inner})
    }
}

//------------ Sender -------------------------------------------------------

pub struct Sender<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Sender<T> {
    pub fn send(&self, t: T) -> Result<(), T> {
        if self.inner.rx_dropped.load(SeqCst) {
            Err(t)
        }
        else {
            self.inner.queue.push(t);
            if let Some(task) = self.inner.rx_task.take(SeqCst) {
                task.unpark()
            }
            Ok(())
        }
    }
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Self {
        let inner = self.inner.clone();
        let _ = inner.tx_count.fetch_add(1, SeqCst);
        Sender{inner: inner}
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        let _ = self.inner.tx_count.fetch_sub(1, SeqCst);
    }
}

//------------ Receiver ------------------------------------------------------

pub struct Receiver<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Stream for Receiver<T> {
    type Item = T;
    type Error = Void;

    fn poll(&mut self) -> Result<Async<Option<T>>, Void> {
        if let Some(t) = self.inner.queue.try_pop() {
            Ok(Async::Ready(Some(t)))
        }
        else if self.inner.tx_count.load(SeqCst) == 0 {
            Ok(Async::Ready(None))
        }
        else {
            self.inner.rx_task.swap(park(), SeqCst);
            Ok(Async::NotReady)
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.rx_dropped.store(true, SeqCst)
    }
}

//------------ SendError ----------------------------------------------------

pub struct SendError<T>(T);

impl<T> fmt::Debug for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_tuple("SendError")
            .field(&"...")
            .finish()
    }
}

impl<T> fmt::Display for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "send failed because receiver is gone")
    }
}

impl<T: Any> error::Error for SendError<T> {
    fn description(&self) -> &str {
        "send failed because receiver is gone"
    }
}

//============ Testing =======================================================

#[cfg(test)]
mod test {
    use std::thread;
    use futures::Future;
    use futures::stream::Stream;
    use super::*;

    #[test]
    fn send_then_drop() {
        let (tx, rx) = channel();
        tx.send(2).unwrap();
        let txc = tx.clone();
        thread::spawn(move || {
            txc.send(4).unwrap();
        });
        thread::spawn(move || {
            tx.send(4).unwrap();
        });
        assert_eq!(rx.collect().wait().unwrap(), [2, 4, 4]);
    }

    #[test]
    fn send_dropped() {
        let (tx, _) = channel();
        assert!(tx.send(4).is_err())
    }

    #[test]
    fn recv_dropped() {
        let (_, rx) = channel::<u8>();
        assert_eq!(rx.collect().wait().unwrap(), b"");
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use libc::{c_void};
use std::ffi::CStr;
use std::ffi::CString;
use std::{fmt, process, ptr};
use std::collections::HashMap;
use std::mem;
use std::mem::{transmute};
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};

use std::sync::{Arc, Mutex};

use super::super::couchbase::*;

use super::super::couchbase::types::response::format_error;

// Gets
pub type OperationResultGet = Result<response::Get, (Option<response::Get>, &'static str)>;
pub type OperationResultGetCallback = Box<Box<Fn(&response::Get)>>;
pub type OperationResultGetInternal<'a> = Result<&'a response::GetInternal, (Option<&'a response::GetInternal>, &'static str)>;
pub type OperationResultGetInternalCallback = Box<Box<Fn(&response::GetInternal)>>;

// Stores
pub type OperationResultStore = Result<response::Store, (Option<response::Store>, &'static str)>;
pub type OperationResultStoreCallback = Box<Box<Fn(&response::Store)>>;
pub type OperationResultStoreInternal<'a> = Result<&'a response::StoreInternal, (Option<&'a response::StoreInternal>, &'static str)>;
pub type OperationResultStoreInternalCallback = Box<Box<Fn(&response::StoreInternal)>>;

pub struct CouchbaseOperation<T> {
    counter: u64,
    map: HashMap<u64, T>
}

impl<T> fmt::Debug for CouchbaseOperation<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "CouchbaseOperation {{ counter: {}, length: {} }}", self.counter, self.map.len())
    }
}

impl<T> CouchbaseOperation<T> {
    pub fn new() -> CouchbaseOperation<T> {
        CouchbaseOperation {
            counter: 0,
            map: HashMap::new()
        }
    }

    pub fn increment_counter(&mut self) -> u64 {
        self.counter += 1;
        return self.counter;
    }
}

#[derive(Debug)]
pub struct CouchbaseOperations {
    get: CouchbaseOperation<cmd::Get>,
    store: CouchbaseOperation<cmd::Get>,
}

impl CouchbaseOperations {
    pub fn new() -> CouchbaseOperations {
        CouchbaseOperations {
            get: CouchbaseOperation::new(),
            store: CouchbaseOperation::new()
        }
    }
}

#[derive(Debug)]
pub struct Client {
    pub opts: Arc<Mutex<CreateSt>>,
    pub instance: Instance,
    pub uri: String,
    pub operations: CouchbaseOperations
}

impl Client {
    /// Constructor
    pub fn new(uri: &str) -> Client {
        let connstr = CString::new(uri).unwrap();

        let mut opts = CreateSt::default();
        opts.v3.connstr = connstr.as_ptr();

        let mut instance: Instance = ptr::null_mut();

        unsafe {
            let res = lcb_create(&mut instance as *mut Instance, &opts as *const CreateSt);
            if res != ErrorType::Success {
                error!("lcb_connect() failed - {:?}", res);
            }

            info!("Connecting to {}", uri);

            let res = lcb_connect(instance);
            if res != ErrorType::Success {
                error!("lcb_connect() failed - {:?}", res);
            }

            let res = lcb_wait(instance);
            if res != ErrorType::Success {
                error!("lcb_wait() failed - {:?}", res);
            }

            let res = lcb_get_bootstrap_status(instance);
            if res != ErrorType::Success {
                error!("lcb_get_bootstrap_status() failed - {:?}, \"{}\"",
                         res,
                         CStr::from_ptr(lcb_strerror(instance, res)).to_str().unwrap()
                );
                process::exit(-1);
            }

            lcb_install_callback3(instance, CallbackType::Get, op_callback);
            lcb_install_callback3(instance, CallbackType::Store, op_callback);

            Client {
                opts: Arc::new(Mutex::new(opts)),
                instance: instance,
                uri: uri.to_string(),
                operations: CouchbaseOperations::new()

            }
        }
    }

    ///  Will cause the operation to fail if the key already exists in the cluster.
    pub fn add<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Add, callback);
    }

    /// Rather than setting the contents of the entire document, take the value specified in value and _append_ it to the existing bytes in the value.
    pub fn append<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Append, callback);
    }

    /// Get document from database
    pub fn get<'a, F>(&'a mut self, key: &str, callback: F) -> &Client
        where F: Fn(OperationResultGet) + 'static
    {
        let mut gcmd = cmd::Get::default();

        gcmd.key._type = KvBufferType::Copy;
        gcmd.key.contig.bytes = key.as_ptr() as *const libc::c_void;
        gcmd.key.contig.nbytes = key.len() as u64;

        unsafe {
            let _id = self.operations.get.increment_counter();

            let boxed: OperationResultGetInternalCallback = Box::new(Box::new(move |result: &response::GetInternal| {
                match result.rc {
                    ErrorType::Success => {
                        debug!("{:?}", result);
                        callback(Ok(response::Get::new(result)));
                    },
                    _ => {
                        callback(Err((Some(response::Get::new(result)), "error" /* result.error(self.instance) */)));
                    }
                }
            }));

            debug!("Boxed box occupies {} bytes in the stack", mem::size_of_val(&boxed));

            // let user_data = transmute::<*const _, *mut c_void>(Box::into_raw(boxed));
            let user_data = Box::into_raw(boxed) as *mut Box<Fn(&response::GetInternal)> as *mut c_void;
            debug!("Setting get callback address {:?}", user_data);

            let res = lcb_get3(self.instance, user_data, &gcmd as *const cmd::Get);
            if res != ErrorType::Success {
                // callback(Err((None, format_error(self.instance, &res))));
            } else {
                let res = lcb_wait(self.instance);
                if res != ErrorType::Success {
                    // callback(Err((None, format_error(self.instance, &res))))
                }
            }
        }

        return self;
    }

    pub fn get_sync(&mut self, key: &str) -> OperationResultGet
    {
        let (tx, rx): (Sender<OperationResultGet>, Receiver<OperationResultGet>) = mpsc::channel();
        self.get(key, move |result: OperationResultGet| {
            let _ = tx.send(result);
        });
        return rx.recv().unwrap();
    }

    /// Like append, but prepends the new value to the existing value.
    pub fn prepend<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Prepend, callback);
    }

    /// Will cause the operation to fail _unless_ the key already exists in the cluster.
    pub fn replace<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Replace, callback);
    }

    /// Unconditionally store the item in the cluster
    pub fn set<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Set, callback);
    }

    /// Store document in database
    pub fn store<'a, F>(&'a mut self, key: &str, value: &str, operation: Operation, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        let mut gcmd = cmd::Store::default();
        gcmd.key._type = KvBufferType::Copy;
        gcmd.key.contig.bytes = key.as_bytes().as_ptr() as *const libc::c_void;
        gcmd.key.contig.nbytes = key.len() as u64;
        gcmd.value._type = KvBufferType::Copy;
        gcmd.value.contig.bytes = value.as_bytes().as_ptr() as *const libc::c_void;
        gcmd.value.contig.nbytes = value.len() as u64;
        gcmd.operation = operation;

        unsafe {
            let _id = self.operations.store.increment_counter();

            let boxed: OperationResultStoreInternalCallback = Box::new(Box::new(move |result: &response::StoreInternal| {
                match result.rc {
                    ErrorType::Success => {
                        debug!("{:?}", result);
                        callback(Ok(response::Store::new(result)));
                    },
                    _ => {
                        callback(Err((Some(response::Store::new(result)), "error" /* result.error(self.instance) */)));
                    }
                }
            }));

            debug!("Boxed box occupies {} bytes in the stack", mem::size_of_val(&boxed));

            // let user_data = transmute::<*const  _, *mut c_void>(Box::into_raw(boxed));
            let user_data = Box::into_raw(boxed) as *mut Box<Fn(&response::StoreInternal)> as *mut c_void;
            debug!("Setting store callback address {:?}", user_data);

            let res = lcb_store3(self.instance, user_data, &gcmd as *const cmd::Store);
            if res != ErrorType::Success {
                // callback(Err((None, format_error(self.instance, &res))))
            } else {
                let res = lcb_wait(self.instance);
                if res != ErrorType::Success {
                    // callback(Err((None, format_error(self.instance, &res))))
                }
            }
        }

        return self;
    }

    /// Behaviorally it is identical to set in that it will make the server unconditionally store the item, whether it exists or not.
    pub fn upsert<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Upsert, callback);
    }
}

impl Drop for Client {
    fn drop(&mut self) {
        unsafe {
            info!("Disconnecting from {}", self.uri);
            lcb_destroy(self.instance);
        }
    }
}

/// libcouchbse callback
unsafe extern "C" fn op_callback(_instance: Instance, cbtype: CallbackType, resp: *const response::Base) {
    match cbtype {
        CallbackType::Get => {
            let gresp = resp as *const response::GetInternal;
            debug!("{:?}", *gresp);

            let cookie = (*gresp).cookie;
            // let callback = transmute::<*mut c_void, *const Box<Fn(&response::GetInternal)>>(cookie);

            //let callback = cookie as *const Box<Fn(&response::GetInternal)>;
            let callback: Box<Box<Fn(&response::GetInternal)>> = Box::from_raw(cookie as *mut Box<Fn(&response::GetInternal)>);
            debug!("Retreived boxed box occupies {} bytes in the stack", mem::size_of_val(&callback));

            // debug!("Got get callback address {:?}", callback);

            (*callback)(&(*gresp));
        },
        CallbackType::Store => {
            let gresp = resp as *const response::StoreInternal;
            debug!("{:?}", *gresp);

            let cookie = (*gresp).cookie;
            // let callback = transmute::<*mut c_void, *const Box<Fn(&response::StoreInternal)>>(cookie);

            //let callback = cookie as *const Box<Fn(&response::StoreInternal)>;
            let callback: Box<Box<Fn(&response::StoreInternal)>> = Box::from_raw(cookie as *mut Box<Fn(&response::StoreInternal)>);
            debug!("Retreived boxed box occupies {} bytes in the stack", mem::size_of_val(&callback));

            // debug!("Got store callback address {:?}", callback);

            (*callback)(&(*gresp));
        },
        _ => error!("! Unknown Callback...")
    };
}
<commit_msg>Optio<Callback> -> Callback<commit_after>extern crate libc;

use libc::{c_void};
use std::ffi::CStr;
use std::ffi::CString;
use std::{fmt, process, ptr};
use std::collections::HashMap;
use std::mem;
use std::mem::{forget, transmute};
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};

use std::sync::{Arc, Mutex};

use super::super::couchbase::*;

use super::super::couchbase::types::response::format_error;

// Gets
pub type OperationResultGet = Result<response::Get, (Option<response::Get>, &'static str)>;
pub type OperationResultGetCallback = Box<Box<Fn(&response::Get)>>;
pub type OperationResultGetInternal<'a> = Result<&'a response::GetInternal, (Option<&'a response::GetInternal>, &'static str)>;
pub type OperationResultGetInternalCallback = Box<Box<Fn(&response::GetInternal)>>;

// Stores
pub type OperationResultStore = Result<response::Store, (Option<response::Store>, &'static str)>;
pub type OperationResultStoreCallback = Box<Box<Fn(&response::Store)>>;
pub type OperationResultStoreInternal<'a> = Result<&'a response::StoreInternal, (Option<&'a response::StoreInternal>, &'static str)>;
pub type OperationResultStoreInternalCallback = Box<Box<Fn(&response::StoreInternal)>>;

pub struct CouchbaseOperation<T> {
    counter: u64,
    map: HashMap<u64, T>
}

impl<T> fmt::Debug for CouchbaseOperation<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "CouchbaseOperation {{ counter: {}, length: {} }}", self.counter, self.map.len())
    }
}

impl<T> CouchbaseOperation<T> {
    pub fn new() -> CouchbaseOperation<T> {
        CouchbaseOperation {
            counter: 0,
            map: HashMap::new()
        }
    }

    pub fn increment_counter(&mut self) -> u64 {
        self.counter += 1;
        return self.counter;
    }
}

#[derive(Debug)]
pub struct CouchbaseOperations {
    get: CouchbaseOperation<cmd::Get>,
    store: CouchbaseOperation<cmd::Get>,
}

impl CouchbaseOperations {
    pub fn new() -> CouchbaseOperations {
        CouchbaseOperations {
            get: CouchbaseOperation::new(),
            store: CouchbaseOperation::new()
        }
    }
}

#[derive(Debug)]
pub struct Client {
    pub opts: Arc<Mutex<CreateSt>>,
    pub instance: Instance,
    pub uri: String,
    pub operations: CouchbaseOperations
}

impl Client {
    /// Constructor
    pub fn new(uri: &str) -> Client {
        let connstr = CString::new(uri).unwrap();

        let mut opts = CreateSt::default();
        opts.v3.connstr = connstr.as_ptr();

        let mut instance: Instance = ptr::null_mut();

        unsafe {
            let res = lcb_create(&mut instance as *mut Instance, &opts as *const CreateSt);
            if res != ErrorType::Success {
                error!("lcb_connect() failed - {:?}", res);
            }

            info!("Connecting to {}", uri);

            let res = lcb_connect(instance);
            if res != ErrorType::Success {
                error!("lcb_connect() failed - {:?}", res);
            }

            let res = lcb_wait(instance);
            if res != ErrorType::Success {
                error!("lcb_wait() failed - {:?}", res);
            }

            let res = lcb_get_bootstrap_status(instance);
            if res != ErrorType::Success {
                error!("lcb_get_bootstrap_status() failed - {:?}, \"{}\"",
                         res,
                         CStr::from_ptr(lcb_strerror(instance, res)).to_str().unwrap()
                );
                process::exit(-1);
            }

            lcb_install_callback3(instance, CallbackType::Get, op_callback);
            lcb_install_callback3(instance, CallbackType::Store, op_callback);

            Client {
                opts: Arc::new(Mutex::new(opts)),
                instance: instance,
                uri: uri.to_string(),
                operations: CouchbaseOperations::new()

            }
        }
    }

    ///  Will cause the operation to fail if the key already exists in the cluster.
    pub fn add<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Add, callback);
    }

    /// Rather than setting the contents of the entire document, take the value specified in value and _append_ it to the existing bytes in the value.
    pub fn append<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Append, callback);
    }

    /// Get document from database
    pub fn get<'a, F>(&'a mut self, key: &str, callback: F) -> &Client
        where F: Fn(OperationResultGet) + 'static
    {
        let key = key.to_owned();

        let mut gcmd = cmd::Get::default();

        gcmd.key._type = KvBufferType::Copy;
        gcmd.key.contig.bytes = key.as_ptr() as *const libc::c_void;
        gcmd.key.contig.nbytes = key.len() as u64;

        unsafe {
            let _id = self.operations.get.increment_counter();

            let boxed: OperationResultGetInternalCallback = Box::new(Box::new(move |result: &response::GetInternal| {
                match result.rc {
                    ErrorType::Success => {
                        debug!("{:?}", result);
                        callback(Ok(response::Get::new(result)));
                    },
                    _ => {
                        callback(Err((Some(response::Get::new(result)), "error" /* result.error(self.instance) */)));
                    }
                }
            }));

            debug!("Boxed box occupies {} bytes in the stack", mem::size_of_val(&boxed));

            // let user_data = transmute::<*const _, *mut c_void>(Box::into_raw(boxed));
            let user_data = Box::into_raw(boxed) as *mut Box<Fn(&response::GetInternal)> as *mut c_void;
            debug!("Setting get callback address {:?}", user_data);

            let res = lcb_get3(self.instance, user_data, &gcmd as *const cmd::Get);
            if res != ErrorType::Success {
                // callback(Err((None, format_error(self.instance, &res))));
            } else {
                let res = lcb_wait(self.instance);
                if res != ErrorType::Success {
                    // callback(Err((None, format_error(self.instance, &res))))
                }
            }
        }

        forget(key);

        return self;
    }

    pub fn get_sync(&mut self, key: &str) -> OperationResultGet
    {
        let (tx, rx): (Sender<OperationResultGet>, Receiver<OperationResultGet>) = mpsc::channel();
        self.get(key, move |result: OperationResultGet| {
            let _ = tx.send(result);
        });
        return rx.recv().unwrap();
    }

    /// Like append, but prepends the new value to the existing value.
    pub fn prepend<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Prepend, callback);
    }

    /// Will cause the operation to fail _unless_ the key already exists in the cluster.
    pub fn replace<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Replace, callback);
    }

    /// Unconditionally store the item in the cluster
    pub fn set<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Set, callback);
    }

    /// Store document in database
    pub fn store<'a, F>(&'a mut self, key: &str, value: &str, operation: Operation, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        let key = key.to_owned();

        let mut gcmd = cmd::Store::default();
        gcmd.key._type = KvBufferType::Copy;
        gcmd.key.contig.bytes = key.as_bytes().as_ptr() as *const libc::c_void;
        gcmd.key.contig.nbytes = key.len() as u64;
        gcmd.value._type = KvBufferType::Copy;
        gcmd.value.contig.bytes = value.as_bytes().as_ptr() as *const libc::c_void;
        gcmd.value.contig.nbytes = value.len() as u64;
        gcmd.operation = operation;

        unsafe {
            let _id = self.operations.store.increment_counter();

            let boxed: OperationResultStoreInternalCallback = Box::new(Box::new(move |result: &response::StoreInternal| {
                match result.rc {
                    ErrorType::Success => {
                        debug!("{:?}", result);
                        callback(Ok(response::Store::new(result)));
                    },
                    _ => {
                        callback(Err((Some(response::Store::new(result)), "error" /* result.error(self.instance) */)));
                    }
                }
            }));

            debug!("Boxed box occupies {} bytes in the stack", mem::size_of_val(&boxed));

            // let user_data = transmute::<*const  _, *mut c_void>(Box::into_raw(boxed));
            let user_data = Box::into_raw(boxed) as *mut Box<Fn(&response::StoreInternal)> as *mut c_void;
            debug!("Setting store callback address {:?}", user_data);

            let res = lcb_store3(self.instance, user_data, &gcmd as *const cmd::Store);
            if res != ErrorType::Success {
                // callback(Err((None, format_error(self.instance, &res))))
            } else {
                let res = lcb_wait(self.instance);
                if res != ErrorType::Success {
                    // callback(Err((None, format_error(self.instance, &res))))
                }
            }
        }

        return self;
    }

    /// Behaviorally it is identical to set in that it will make the server unconditionally store the item, whether it exists or not.
    pub fn upsert<'a, F>(&'a mut self, key: &str, value: &str, callback: F) -> &Client
        where F: Fn(OperationResultStore) + 'static
    {
        return self.store(key, value, Operation::Upsert, callback);
    }
}

impl Drop for Client {
    fn drop(&mut self) {
        unsafe {
            info!("Disconnecting from {}", self.uri);
            lcb_destroy(self.instance);
        }
    }
}

/// libcouchbse callback
unsafe extern "C" fn op_callback(_instance: Instance, cbtype: CallbackType, resp: *const response::Base) {
    match cbtype {
        CallbackType::Get => {
            let gresp = resp as *const response::GetInternal;
            debug!("{:?}", *gresp);

            let cookie = (*gresp).cookie;
            // let callback = transmute::<*mut c_void, *const Box<Fn(&response::GetInternal)>>(cookie);

            //let callback = cookie as *const Box<Fn(&response::GetInternal)>;
            let callback: Box<Box<Fn(&response::GetInternal)>> = Box::from_raw(cookie as *mut Box<Fn(&response::GetInternal)>);
            debug!("Retreived boxed box occupies {} bytes in the stack", mem::size_of_val(&callback));

            debug!("Got get callback address {:?}", callback);

            (*callback)(&(*gresp));
        },
        CallbackType::Store => {
            let gresp = resp as *const response::StoreInternal;
            debug!("{:?}", *gresp);

            let cookie = (*gresp).cookie;
            // let callback = transmute::<*mut c_void, *const Box<Fn(&response::StoreInternal)>>(cookie);

            //let callback = cookie as *const Box<Fn(&response::StoreInternal)>;
            let callback: Box<Box<Fn(&response::StoreInternal)>> = Box::from_raw(cookie as *mut Box<Fn(&response::StoreInternal)>);
            debug!("Retreived boxed box occupies {} bytes in the stack", mem::size_of_val(&callback));

            debug!("Got store callback address {:?}", callback);

            (*callback)(&(*gresp));
        },
        _ => error!("! Unknown Callback...")
    };
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::collections::HashMap;

use url::Url;
use hyper::{self, Client};
use hyper::client::{Request, Response};
use hyper::method::Method;
use rustc_serialize::json::{self, Json, Object};
use rustc_serialize::Encodable;
use multipart::client::Multipart;

use errors::WeChatError;
use types::WeChatResult;
use session::SessionStore;

#[derive(Debug, Clone)]
pub struct WeChatClient<T: SessionStore> {
    pub appid: String,
    pub secret: String,
    pub session: T,
}

impl<T: SessionStore> WeChatClient<T> {

    #[inline]
    pub fn new<S: Into<String>>(appid: S, secret: S, session: T) -> WeChatClient<T> {
        WeChatClient {
            appid: appid.into(),
            secret: secret.into(),
            session: session,
        }
    }

    #[inline]
    pub fn with_access_token<S: Into<String>>(appid: S, secret: S, access_token: S, session: T) -> WeChatClient<T> {
        let appid = appid.into();
        let secret = secret.into();
        let client = Self::new(appid.clone(), secret, session);
        let token_key = format!("{}_access_token", appid);
        client.session.set(&token_key, access_token.into(), None);
        client
    }

    #[inline]
    pub fn access_token(&self) -> String {
        let token_key = format!("{}_access_token", self.appid);
        let token: String = self.session.get(&token_key, Some("".to_owned())).unwrap();
        token
    }

    pub fn request<D: Encodable>(&self, method: Method, url: &str, params: Vec<(&str, &str)>, data: &D) -> WeChatResult<Response> {
        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let body = match json::encode(data) {
            Ok(text) => text,
            Err(_) => "".to_owned(),
        };
        let client = Client::new();
        let req = if method == Method::Post {
            client.post(http_url).body(&body)
        } else {
            client.get(http_url)
        };
        let res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        Ok(res)

    }

    pub fn upload_file<R: Read>(&self, url: &str, params: Vec<(&str, &str)>, files: &mut HashMap<String, &mut R>) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }

        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let request = Request::new(Method::Post, http_url).unwrap();
        let mut req = Multipart::from_request(request).unwrap();
        for (name, stream) in files.iter_mut() {
            req.write_stream(name, stream, None, None);
        }
        let mut res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        self.json_decode(&mut res)
    }

    #[inline]
    fn json_decode(&self, res: &mut Response) -> WeChatResult<Json> {
        let obj = match Json::from_reader(res) {
            Ok(decoded) => { decoded },
            Err(_) => {
                error!("Json decode error");
                return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() });
            }
        };
        match obj.find("errcode") {
            Some(code) => {
                let errcode = code.as_i64().unwrap();
                if errcode != 0 {
                    let errmsg = match obj.find("errmsg") {
                        Some(msg) => {
                            msg.as_string().unwrap()
                        },
                        None => { "" }
                    };
                    return Err(WeChatError::ClientError { errcode: errcode as i32, errmsg: errmsg.to_owned() });
                }
            },
            None => {},
        }
        Ok(obj)
    }

    #[inline]
    pub fn post<D: Encodable>(&self, url: &str, params: Vec<(&str, &str)>, data: &D) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        let mut res = try!(self.request(Method::Post, url, params, data));
        self.json_decode(&mut res)
    }

    #[inline]
    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        let mut res = try!(self.request(Method::Get, url, params, &Object::new()));
        self.json_decode(&mut res)
    }

    pub fn fetch_access_token(&self) -> Option<String> {
        let res = self.request(
            Method::Get,
            "token",
            vec![
                ("grant_type", "client_credential"),
                ("appid", &self.appid),
                ("secret", &self.secret),
            ],
            &Object::new()
        );
        let mut raw_data = match res {
            Ok(raw) => raw,
            Err(_) => { return None; },
        };
        let data = match self.json_decode(&mut raw_data) {
            Ok(data) => data,
            Err(_) => { return None; },
        };
        let token = match data.find("access_token") {
            Some(token) => token,
            None => { return None; }
        };
        let expires_in = match data.find("expires_in") {
            Some(expires) => expires.as_u64().unwrap() as usize,
            None => 7200usize,
        };
        let token_str = match *token {
            Json::String(ref v) => {
                let token_key = format!("{}_access_token", self.appid);
                self.session.set(&token_key, v.to_owned(), Some(expires_in));
                Some(format!("{}", v))
            },
            _ => None,
        };
        token_str
    }
}
<commit_msg>Record expires_at in session storage<commit_after>use std::io::Read;
use std::collections::HashMap;

use time;
use url::Url;
use hyper::{self, Client};
use hyper::client::{Request, Response};
use hyper::method::Method;
use rustc_serialize::json::{self, Json, Object};
use rustc_serialize::Encodable;
use multipart::client::Multipart;

use errors::WeChatError;
use types::WeChatResult;
use session::SessionStore;

#[derive(Debug, Clone)]
pub struct WeChatClient<T: SessionStore> {
    pub appid: String,
    pub secret: String,
    pub session: T,
}

impl<T: SessionStore> WeChatClient<T> {

    #[inline]
    pub fn new<S: Into<String>>(appid: S, secret: S, session: T) -> WeChatClient<T> {
        WeChatClient {
            appid: appid.into(),
            secret: secret.into(),
            session: session,
        }
    }

    #[inline]
    pub fn with_access_token<S: Into<String>>(appid: S, secret: S, access_token: S, session: T) -> WeChatClient<T> {
        let appid = appid.into();
        let secret = secret.into();
        let client = Self::new(appid.clone(), secret, session);
        let token_key = format!("{}_access_token", appid);
        client.session.set(&token_key, access_token.into(), None);
        client
    }

    #[inline]
    pub fn access_token(&self) -> String {
        let token_key = format!("{}_access_token", self.appid);
        let expires_key = format!("{}_expires_at", self.appid);
        let token: String = self.session.get(&token_key, Some("".to_owned())).unwrap();
        let timestamp = time::get_time().sec;
        let expires_at: i64 = self.session.get(&expires_key, Some(timestamp)).unwrap();
        if expires_at <= timestamp {
            "".to_owned()
        } else {
            token
        }
    }

    pub fn request<D: Encodable>(&self, method: Method, url: &str, params: Vec<(&str, &str)>, data: &D) -> WeChatResult<Response> {
        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let body = match json::encode(data) {
            Ok(text) => text,
            Err(_) => "".to_owned(),
        };
        let client = Client::new();
        let req = if method == Method::Post {
            client.post(http_url).body(&body)
        } else {
            client.get(http_url)
        };
        let res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        Ok(res)

    }

    pub fn upload_file<R: Read>(&self, url: &str, params: Vec<(&str, &str)>, files: &mut HashMap<String, &mut R>) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }

        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let request = Request::new(Method::Post, http_url).unwrap();
        let mut req = Multipart::from_request(request).unwrap();
        for (name, stream) in files.iter_mut() {
            req.write_stream(name, stream, None, None);
        }
        let mut res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        self.json_decode(&mut res)
    }

    #[inline]
    fn json_decode(&self, res: &mut Response) -> WeChatResult<Json> {
        let obj = match Json::from_reader(res) {
            Ok(decoded) => { decoded },
            Err(_) => {
                error!("Json decode error");
                return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() });
            }
        };
        match obj.find("errcode") {
            Some(code) => {
                let errcode = code.as_i64().unwrap();
                if errcode != 0 {
                    let errmsg = match obj.find("errmsg") {
                        Some(msg) => {
                            msg.as_string().unwrap()
                        },
                        None => { "" }
                    };
                    return Err(WeChatError::ClientError { errcode: errcode as i32, errmsg: errmsg.to_owned() });
                }
            },
            None => {},
        }
        Ok(obj)
    }

    #[inline]
    pub fn post<D: Encodable>(&self, url: &str, params: Vec<(&str, &str)>, data: &D) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        let mut res = try!(self.request(Method::Post, url, params, data));
        self.json_decode(&mut res)
    }

    #[inline]
    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> WeChatResult<Json> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        let mut res = try!(self.request(Method::Get, url, params, &Object::new()));
        self.json_decode(&mut res)
    }

    pub fn fetch_access_token(&self) -> Option<String> {
        let res = self.request(
            Method::Get,
            "token",
            vec![
                ("grant_type", "client_credential"),
                ("appid", &self.appid),
                ("secret", &self.secret),
            ],
            &Object::new()
        );
        let mut raw_data = match res {
            Ok(raw) => raw,
            Err(_) => { return None; },
        };
        let data = match self.json_decode(&mut raw_data) {
            Ok(data) => data,
            Err(_) => { return None; },
        };
        let token = match data.find("access_token") {
            Some(token) => token,
            None => { return None; }
        };
        let expires_in = match data.find("expires_in") {
            Some(expires) => expires.as_u64().unwrap() as usize,
            None => 7200usize,
        };
        let expires_at = time::get_time().sec + expires_in as i64;
        let token_str = match *token {
            Json::String(ref v) => {
                let token_key = format!("{}_access_token", self.appid);
                let expires_key = format!("{}_expires_at", self.appid);
                self.session.set(&token_key, v.to_owned(), Some(expires_in));
                self.session.set(&expires_key, expires_at, Some(expires_in));
                Some(format!("{}", v))
            },
            _ => None,
        };
        token_str
    }
}
<|endoftext|>