<commit_before>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN_DEF`, `TYPE_ALIAS_DEF`, or `CONST_DEF` node
//! and an direct child of an `IMPL_DEF`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use hir::{self, Docs, HasSource};
use ra_assists::utils::get_missing_impl_items;
use ra_syntax::{
    ast::{self, edit},
    AstNode, SyntaxKind, SyntaxNode, TextRange,
};
use ra_text_edit::TextEdit;

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

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    let mut tokens = ctx.token.ancestors();
    let trigger = tokens.find(|p| match p.kind() {
        SyntaxKind::FN_DEF
        | SyntaxKind::TYPE_ALIAS_DEF
        | SyntaxKind::CONST_DEF
        | SyntaxKind::NAME_REF
        | SyntaxKind::BLOCK_EXPR => true,
        _ => false,
    });

    let impl_def = tokens
        .find(|p| match p.kind() {
            SyntaxKind::IMPL_DEF => true,
            _ => false,
        })
        .and_then(|n| ast::ImplDef::cast(n));

    if let (Some(trigger), Some(impl_def)) = (trigger, impl_def) {
        match trigger.kind() {
            SyntaxKind::NAME_REF => {
                get_missing_impl_items(&ctx.sema, &impl_def).iter().for_each(|item| match item {
                    hir::AssocItem::Function(fn_item) => {
                        add_function_impl(&trigger, acc, ctx, &fn_item)
                    }
                    hir::AssocItem::TypeAlias(type_item) => {
                        add_type_alias_impl(&trigger, acc, ctx, &type_item)
                    }
                    hir::AssocItem::Const(const_item) => {
                        add_const_impl(&trigger, acc, ctx, &const_item)
                    }
                })
            }

            SyntaxKind::FN_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Function(fn_item) => Some(fn_item),
                        _ => None,
                    },
                ) {
                    add_function_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::TYPE_ALIAS_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::TypeAlias(type_item) => Some(type_item),
                        _ => None,
                    },
                ) {
                    add_type_alias_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::CONST_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Const(const_item) => Some(const_item),
                        _ => None,
                    },
                ) {
                    add_const_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            _ => {}
        }
    }
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: &hir::Function,
) {
    let display = FunctionSignature::from_hir(ctx.db, *func);

    let fn_name = func.name(ctx.db).to_string();

    let label = if !func.params(ctx.db).is_empty() {
        format!("fn {}(..)", fn_name)
    } else {
        format!("fn {}()", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.has_self_param(ctx.db) {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };

    let snippet = format!("{} {{}}", display);

    let range = TextRange::from_to(fn_def_node.text_range().start(), ctx.source_range().end());

    builder.text_edit(TextEdit::replace(range, snippet)).kind(completion_kind).add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: &hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::from_to(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: &hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range =
            TextRange::from_to(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::ConstDef) -> String {
    let const_ = edit::strip_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == SyntaxKind::SEMI || s.kind() == SyntaxKind::EQ)
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::from_to(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use crate::completion::{do_completion, CompletionItem, CompletionKind};
    use insta::assert_debug_snapshot;

    fn complete(code: &str) -> Vec<CompletionItem> {
        do_completion(code, CompletionKind::Magic)
    }

    #[test]
    fn name_ref_function_type_const() {
        let completions = complete(
            r"
            trait Test {
                type TestType;
                const TEST_CONST: u16;
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const TEST_CONST: u16 = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "const TEST_CONST: u16 = ",
                kind: Const,
                lookup: "TEST_CONST",
            },
            CompletionItem {
                label: "fn test()",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
            CompletionItem {
                label: "type TestType = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "type TestType = ",
                kind: TypeAlias,
                lookup: "TestType",
            },
        ]
        "###);
    }

    #[test]
    fn name_ref_single_function() {
        let completions = complete(
            r"
            trait Test {
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn test()",
                source_range: [139; 140),
                delete: [139; 140),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
        ]
        "###);
    }

    #[test]
    fn single_function() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [141; 142),
                delete: [138; 142),
                insert: "fn foo() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn hide_implemented_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {}

                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo_bar()",
                source_range: [200; 201),
                delete: [197; 201),
                insert: "fn foo_bar() {}",
                kind: Function,
                lookup: "foo_bar",
            },
        ]
        "###);
    }

    #[test]
    fn completes_only_on_top_level() {
        let completions = complete(
            r"
            trait Test {
                fn foo();

                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {
                    <|>
                }
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"[]"###);
    }

    #[test]
    fn generic_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [144; 145),
                delete: [141; 145),
                insert: "fn foo<T>() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn generic_constrait_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>() where T: Into<String>;
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [166; 167),
                delete: [163; 167),
                insert: "fn foo<T>()\nwhere T: Into<String> {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn associated_type() {
        let completions = complete(
            r"
            trait Test {
                type SomeType;
            }

            impl Test for () {
                type S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "type SomeType = ",
                source_range: [124; 125),
                delete: [119; 125),
                insert: "type SomeType = ",
                kind: TypeAlias,
                lookup: "SomeType",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [133; 134),
                delete: [127; 134),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const_with_default() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16 = 42;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [138; 139),
                delete: [132; 139),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }
}
<commit_msg>Don't allow nested completions<commit_after>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN_DEF`, `TYPE_ALIAS_DEF`, or `CONST_DEF` node
//! and an direct child of an `IMPL_DEF`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use hir::{self, Docs, HasSource};
use ra_assists::utils::get_missing_impl_items;
use ra_syntax::{
    ast::{self, edit},
    AstNode, SyntaxKind, SyntaxNode, TextRange,
};
use ra_text_edit::TextEdit;

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

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    let mut tokens = ctx.token.ancestors();
    let completion_match = tokens
        .find(|p| match p.kind() {
            SyntaxKind::FN_DEF
            | SyntaxKind::TYPE_ALIAS_DEF
            | SyntaxKind::CONST_DEF
            | SyntaxKind::NAME_REF
            | SyntaxKind::BLOCK_EXPR => true,
            _ => false,
        })
        .and_then(|trigger| {
            for p in tokens {
                match p.kind() {
                    // No nested completions
                    SyntaxKind::FN_DEF | SyntaxKind::BLOCK => return None,
                    SyntaxKind::IMPL_DEF => return ast::ImplDef::cast(p).map(|p| (trigger, p)),
                    _ => {}
                }
            }
            None
        });

    if let Some((trigger, impl_def)) = completion_match {
        match trigger.kind() {
            SyntaxKind::NAME_REF => {
                get_missing_impl_items(&ctx.sema, &impl_def).iter().for_each(|item| match item {
                    hir::AssocItem::Function(fn_item) => {
                        add_function_impl(&trigger, acc, ctx, &fn_item)
                    }
                    hir::AssocItem::TypeAlias(type_item) => {
                        add_type_alias_impl(&trigger, acc, ctx, &type_item)
                    }
                    hir::AssocItem::Const(const_item) => {
                        add_const_impl(&trigger, acc, ctx, &const_item)
                    }
                })
            }

            SyntaxKind::FN_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Function(fn_item) => Some(fn_item),
                        _ => None,
                    },
                ) {
                    add_function_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::TYPE_ALIAS_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::TypeAlias(type_item) => Some(type_item),
                        _ => None,
                    },
                ) {
                    add_type_alias_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::CONST_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Const(const_item) => Some(const_item),
                        _ => None,
                    },
                ) {
                    add_const_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            _ => {}
        }
    }
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: &hir::Function,
) {
    let display = FunctionSignature::from_hir(ctx.db, *func);

    let fn_name = func.name(ctx.db).to_string();

    let label = if !func.params(ctx.db).is_empty() {
        format!("fn {}(..)", fn_name)
    } else {
        format!("fn {}()", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.has_self_param(ctx.db) {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };

    let snippet = format!("{} {{}}", display);

    let range = TextRange::from_to(fn_def_node.text_range().start(), ctx.source_range().end());

    builder.text_edit(TextEdit::replace(range, snippet)).kind(completion_kind).add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: &hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::from_to(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: &hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range =
            TextRange::from_to(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::ConstDef) -> String {
    let const_ = edit::strip_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == SyntaxKind::SEMI || s.kind() == SyntaxKind::EQ)
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::from_to(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use crate::completion::{do_completion, CompletionItem, CompletionKind};
    use insta::assert_debug_snapshot;

    fn complete(code: &str) -> Vec<CompletionItem> {
        do_completion(code, CompletionKind::Magic)
    }

    #[test]
    fn name_ref_function_type_const() {
        let completions = complete(
            r"
            trait Test {
                type TestType;
                const TEST_CONST: u16;
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const TEST_CONST: u16 = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "const TEST_CONST: u16 = ",
                kind: Const,
                lookup: "TEST_CONST",
            },
            CompletionItem {
                label: "fn test()",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
            CompletionItem {
                label: "type TestType = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "type TestType = ",
                kind: TypeAlias,
                lookup: "TestType",
            },
        ]
        "###);
    }

    #[test]
    fn name_ref_single_function() {
        let completions = complete(
            r"
            trait Test {
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn test()",
                source_range: [139; 140),
                delete: [139; 140),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
        ]
        "###);
    }

    #[test]
    fn single_function() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [141; 142),
                delete: [138; 142),
                insert: "fn foo() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn hide_implemented_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {}

                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo_bar()",
                source_range: [200; 201),
                delete: [197; 201),
                insert: "fn foo_bar() {}",
                kind: Function,
                lookup: "foo_bar",
            },
        ]
        "###);
    }

    #[test]
    fn completes_only_on_top_level() {
        let completions = complete(
            r"
            trait Test {
                fn foo();

                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {
                    <|>
                }
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"[]"###);
    }

    #[test]
    fn generic_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [144; 145),
                delete: [141; 145),
                insert: "fn foo<T>() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn generic_constrait_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>() where T: Into<String>;
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [166; 167),
                delete: [163; 167),
                insert: "fn foo<T>()\nwhere T: Into<String> {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn associated_type() {
        let completions = complete(
            r"
            trait Test {
                type SomeType;
            }

            impl Test for () {
                type S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "type SomeType = ",
                source_range: [124; 125),
                delete: [119; 125),
                insert: "type SomeType = ",
                kind: TypeAlias,
                lookup: "SomeType",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [133; 134),
                delete: [127; 134),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const_with_default() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16 = 42;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [138; 139),
                delete: [132; 139),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }
}
<|endoftext|>


<commit_before>use types;

use serde_json::value::ToJson;

use proto::{
    Term_TermType as TermType,
};

use types::Command as Cmd;

pub use types::{
    WithTableOpts, WithChangesOpts, WithGetAllOpts
};

#[allow(non_upper_case_globals)]
pub const r: Command = Command;

impl Client for Command {}

pub struct Command;

#[allow(dead_code)]
impl Command {
    pub fn db<T>(&self, arg: T) -> types::Db where
        T: Into<types::String>
        {
            Cmd::new(TermType::DB, None)
                .with_args(arg.into().into())
                .into()
        }

    pub fn uuid(&self) -> types::String {
        Cmd::new(TermType::UUID, None)
            .into()
    }

    pub fn map<T, F, O>(&self, arg: T, func: F) -> O where
        T: Into<types::String>,
        F: FnOnce(types::String),
        O: types::DataType
        {
            Cmd::new(TermType::MAP, None)
                .with_args(arg.into().into())
                .into()
        }
}

pub trait Client {
    fn table<T>(&self, arg: T) -> types::WithOpts<types::Table, types::TableOpts> where
        T: Into<types::String>
        {
            r.db("test").table(arg)
        }
}

pub trait Table where Self: types::DataType {
    fn table<T>(&self, arg: T) -> types::WithOpts<types::Table, types::TableOpts>
        where T: Into<types::String>
        {
            Cmd::new(TermType::TABLE, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}

pub trait Changes where Self: types::DataType {
    fn changes(&self) -> types::WithOpts<types::Stream, types::ChangesOpts<bool>>
        where types::ChangesOpts<bool>: Default + ToJson + Clone
        {
            Cmd::new(TermType::CHANGES, Some(self.clone().into()))
                .into()
        }
}

pub trait Get where Self: types::DataType {
    fn get<T>(&self, arg: T) -> types::ObjectSelection
        where T: Into<types::PrimaryKey>,
    {
        Cmd::new(TermType::GET, Some(self.clone().into()))
            .with_args(arg.into().into())
            .into()
    }
}

pub trait GetAll where Self: types::DataType {
    fn get_all<T>(&self, arg: T) -> types::WithOpts<types::StreamSelection, types::GetAllOpts>
        where T: Into<types::SecondaryKey>,
              types::GetAllOpts: Default + ToJson + Clone
    {
        Cmd::new(TermType::GET_ALL, Some(self.clone().into()))
            .with_args(arg.into().into())
            .into()
    }
}

pub trait GetField where Self: types::DataType {
    fn get_field<T, O>(&self, arg: T) -> O
        where T: Into<types::String>, O: types::DataType
        {
            Cmd::new(TermType::GET_FIELD, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}

pub trait GetFieldArray where Self: types::DataType {
    fn get_field<T>(&self, arg: T) -> types::Array
        where T: Into<types::String>
        {
            Cmd::new(TermType::GET_FIELD, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}
<commit_msg>Refactor<commit_after>use types;

use serde_json::value::ToJson;

use proto::{
    Term_TermType as TermType,
};

use types::Command as Cmd;

pub use types::{
    WithTableOpts, WithChangesOpts, WithGetAllOpts
};

#[allow(non_upper_case_globals)]
pub const r: Command = Command;

impl Client for Command {}

pub struct Command;

#[allow(dead_code)]
impl Command {
    pub fn db<T>(&self, arg: T) -> types::Db where
        T: Into<types::String>
        {
            Cmd::new(TermType::DB, None)
                .with_args(arg.into().into())
                .into()
        }

    pub fn uuid(&self) -> types::String {
        Cmd::new(TermType::UUID, None)
            .into()
    }

    pub fn map<T, F>(&self, arg: T, func: F) -> types::Stream where
        T: Into<types::String>,
        F: FnOnce(types::String),
        {
            Cmd::new(TermType::MAP, None)
                .with_args(arg.into().into())
                .into()
        }

    pub fn map_arr<T, F>(&self, arg: T, func: F) -> types::Array where
        T: Into<types::String>,
        F: FnOnce(types::String),
        {
            Cmd::new(TermType::MAP, None)
                .with_args(arg.into().into())
                .into()
        }
}

pub trait Client {
    fn table<T>(&self, arg: T) -> types::WithOpts<types::Table, types::TableOpts> where
        T: Into<types::String>
        {
            r.db("test").table(arg)
        }
}

pub trait Table where Self: types::DataType {
    fn table<T>(&self, arg: T) -> types::WithOpts<types::Table, types::TableOpts>
        where T: Into<types::String>
        {
            Cmd::new(TermType::TABLE, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}

pub trait Changes where Self: types::DataType {
    fn changes(&self) -> types::WithOpts<types::Stream, types::ChangesOpts<bool>>
        where types::ChangesOpts<bool>: Default + ToJson + Clone
        {
            Cmd::new(TermType::CHANGES, Some(self.clone().into()))
                .into()
        }
}

pub trait Get where Self: types::DataType {
    fn get<T>(&self, arg: T) -> types::ObjectSelection
        where T: Into<types::PrimaryKey>,
    {
        Cmd::new(TermType::GET, Some(self.clone().into()))
            .with_args(arg.into().into())
            .into()
    }
}

pub trait GetAll where Self: types::DataType {
    fn get_all<T>(&self, arg: T) -> types::WithOpts<types::StreamSelection, types::GetAllOpts>
        where T: Into<types::SecondaryKey>,
              types::GetAllOpts: Default + ToJson + Clone
    {
        Cmd::new(TermType::GET_ALL, Some(self.clone().into()))
            .with_args(arg.into().into())
            .into()
    }
}

pub trait GetField where Self: types::DataType {
    fn get_field<T, O>(&self, arg: T) -> O
        where T: Into<types::String>, O: types::DataType
        {
            Cmd::new(TermType::GET_FIELD, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}

pub trait GetFieldArray where Self: types::DataType {
    fn get_field<T>(&self, arg: T) -> types::Array
        where T: Into<types::String>
        {
            Cmd::new(TermType::GET_FIELD, Some(self.clone().into()))
                .with_args(arg.into().into())
                .into()
        }
}
<|endoftext|>


<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::fmt;
use std::error::Error as StdError;

pub struct RootCommand(Result<String>);

impl fmt::Display for RootCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            Ok(ref cmd) => write!(f, "{}", cmd),
            Err(_) => Err(fmt::Error),
        }
    }
}

pub struct Command;
pub struct Query;

pub trait IntoCommandArg : fmt::Display {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{}", self))
    }
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{:?}", self))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{:?}", self))
    }
}

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        match serde_json::to_string(self) {
            Ok(cmd) => Ok(cmd),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<String> {
        match self.0 {
            Ok(ref cmd) => Ok(cmd.to_string()),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

impl IntoCommandArg for bool { }
impl IntoCommandArg for char { }
impl IntoCommandArg for u8 { }
impl IntoCommandArg for u16 { }
impl IntoCommandArg for u32 { }
impl IntoCommandArg for u64 { }
impl IntoCommandArg for usize { }
impl IntoCommandArg for i8 { }
impl IntoCommandArg for i16 { }
impl IntoCommandArg for i32 { }
impl IntoCommandArg for i64 { }
impl IntoCommandArg for isize { }
impl IntoCommandArg for f32 { }
impl IntoCommandArg for f64 { }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(e.to_arg())
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None,
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None,
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None,
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE_CREATE,
                                     Some(name),
                                     None,
                                     Some(&commands)))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE,
                                     Some(name),
                                     None,
                                     Some(&commands)))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::INSERT,
                                     Some(data),
                                     None,
                                     Some(&commands)))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE, None as Option<&str>, None as Option<&str>, Some(&commands)))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::FILTER,
                                     Some(filter),
                                     None,
                                     Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                arguments: Option<T>,
                options: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
    {
        let mut cmds = format!("[{},", command.value());
        let mut args = String::new();
        if let Some(commands) = commands {
            args.push_str(commands);
            if arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&try!(arguments.to_arg()));
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", try!(options.to_arg())));
        }
        cmds.push(']');
        Ok(cmds)
    }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
                -> String {
        let mut qry = format!("[{}", query_type.value());
        if let Some(query) = query {
            qry.push_str(&format!(",{}", query));
        }
        if let Some(options) = options {
            qry.push_str(&format!(",{}", options));
        }
        qry.push_str("]");
        qry
    }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<commit_msg>Add support for options<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::fmt;
use std::error::Error as StdError;

#[derive(Debug)]
pub struct RootCommand(Result<String>);

impl fmt::Display for RootCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            Ok(ref cmd) => write!(f, "{}", cmd),
            Err(_) => Err(fmt::Error),
        }
    }
}

pub struct Command;
pub struct Query;

pub trait IntoCommandArg : fmt::Display + fmt::Debug {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{}", self)), None))
    }
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match serde_json::to_string(self) {
            Ok(cmd) => Ok((Some(cmd), None)),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match self.0 {
            Ok(ref cmd) => Ok((Some(cmd.to_string()), None)),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

impl IntoCommandArg for bool { }
impl IntoCommandArg for char { }
impl IntoCommandArg for u8 { }
impl IntoCommandArg for u16 { }
impl IntoCommandArg for u32 { }
impl IntoCommandArg for u64 { }
impl IntoCommandArg for usize { }
impl IntoCommandArg for i8 { }
impl IntoCommandArg for i16 { }
impl IntoCommandArg for i32 { }
impl IntoCommandArg for i64 { }
impl IntoCommandArg for isize { }
impl IntoCommandArg for f32 { }
impl IntoCommandArg for f64 { }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        match e.to_arg() {
            Ok(arg) => if let Some(arg) = arg.0 {
                RootCommand(Ok(arg))
            } else {
                RootCommand(Ok(String::new()))
            },
            Err(e) => RootCommand(Err(e)),
        }
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE_CREATE,
                                     Some(name),
                                     Some(&commands)))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE,
                                     Some(name),
                                     Some(&commands)))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::INSERT,
                                     Some(data),
                                     Some(&commands)))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE,
                                  None as Option<&str>,
                                  Some(&commands)))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::FILTER,
                                     Some(filter),
                                     Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                input: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
        {
            let mut cmds = format!("[{},", command.value());
            let mut args = String::new();
            if let Some(commands) = commands {
                args.push_str(commands);
                if input.is_some() {
                    args.push_str(",");
                }
            }
            let mut arguments: Option<String> = None;
            let mut options: Option<String> = None;
            if let Some(input) = input {
                let (args, opts) = try!(input.to_arg());
                arguments = args;
                options = opts;
            }
            if let Some(arguments) = arguments {
                args.push_str(&arguments);
            }
            cmds.push_str(&format!("[{}]", args));
            if let Some(options) = options {
                cmds.push_str(&format!(",{{{}}}", options));
            }
            cmds.push(']');
            Ok(cmds)
        }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
        -> String {
            let mut qry = format!("[{}", query_type.value());
            if let Some(query) = query {
                qry.push_str(&format!(",{}", query));
            }
            if let Some(options) = options {
                qry.push_str(&format!(",{}", options));
            }
            qry.push_str("]");
            qry
        }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

extern crate libc;

use crate::config::Config;
use crate::layout::{Layout, LayoutMsg};
use crate::xlib_window_system::XlibWindowSystem;
use crate::state::WmState;
use crate::workspace::MoveOp;
use self::libc::execvpe;
use std::{env, iter};
use std::ptr::null;
use std::ffi::CString;
use std::process::Command;
use std::path::Path;
use std::fs::OpenOptions;
use x11::xlib::Window;
use anyhow::{bail, Context, Result};

pub enum Cmd {
    Custom(Box<dyn Fn(&WmState) -> Result<Option<Cmd>, String>>),
    Exec(String, Vec<String>),
    SwitchWorkspace(usize),
    SwitchScreen(usize),
    MoveToWorkspace(usize),
    MoveToScreen(usize),
    SendLayoutMsg(LayoutMsg),
    NestLayout(Box<dyn Fn() -> Box<dyn Layout>>),
    Reload,
    Exit,
    KillClient,
    FocusParentUp,
    FocusParentDown,
    FocusParentMaster,
    FocusUp,
    FocusDown,
    FocusMaster,
    SwapUp,
    SwapDown,
    SwapMaster,
    SwapParentUp,
    SwapParentDown,
    SwapParentMaster,
}

impl Cmd {
    pub fn call(&self, xws: &XlibWindowSystem, state: &mut WmState, config: &Config) -> Result<()> {
        match self {
            Cmd::Custom(func) => {
                debug!("Cmd::Custom");
                match func(state) {
                    Ok(Some(cmd)) => cmd.call(xws, state, config)?,
                    Ok(None) => (),
                    Err(e) => error!("Cmd::Custom failed: {}", e),
                }
            }
            Cmd::Exec(ref cmd, ref args) => {
                debug!("Cmd::Exec: {} {:?}", cmd, args);
                exec(cmd.clone(), args.clone());
            }
            Cmd::SwitchWorkspace(index) => {
                debug!("Cmd::SwitchWorkspace: {}", index);
                state.switch_to_ws(xws, config, index - 1, true);
            }
            Cmd::SwitchScreen(screen) => {
                debug!("Cmd::SwitchScreen: {}", screen);
                state.switch_to_screen(xws, config, screen - 1);
            }
            Cmd::MoveToWorkspace(index) => {
                debug!("Cmd::MoveToWorkspace: {}", index);
                state.move_window_to_ws(xws, config, index - 1);
            }
            Cmd::MoveToScreen(screen) => {
                debug!("Cmd::MoveToScreen: {}", screen);
                state.move_window_to_screen(xws, config, screen - 1);
            }
            Cmd::SendLayoutMsg(ref msg) => {
                debug!("Cmd::SendLayoutMsg::{:?}", msg);
                state.current_ws_mut().send_layout_message(xws, msg.clone());
                state.current_ws().redraw(xws, config, state.get_screens());
            }
            Cmd::NestLayout(layout_fn) => {
                let layout = layout_fn();
                debug!("Cmd::NestLayout: {}", layout.name());
                state.current_ws_mut().nest_layout(layout);
            }
            Cmd::Reload => {
                debug!("Cmd::Reload");
                reload(state)
                    .context("failed to reload xr3wm")?;
            }
            Cmd::Exit => {
                debug!("Cmd::Exit");
                xws.close();
            }
            Cmd::KillClient => {
                if let Some(window) = state.current_ws().focused_window() {
                    debug!("Cmd::KillClient: {:#x}", window);
                    xws.kill_window(window);
                }
            }
            Cmd::FocusUp | Cmd::FocusDown | Cmd::FocusMaster | Cmd::FocusParentUp | Cmd::FocusParentDown | Cmd::FocusParentMaster => {
                if let Some(window) = state.current_ws().focused_window() {
                    let workspace = state.current_ws_mut();
                    let new_focus = match self {
                        Cmd::FocusUp => {
                            debug!("Cmd::FocusUp: {:#x}", window);
                            workspace.move_focus(MoveOp::Up)
                        }
                        Cmd::FocusDown => {
                            debug!("Cmd::FocusDown: {:#x}", window);
                            workspace.move_focus(MoveOp::Down)
                        }
                        Cmd::FocusMaster => {
                            debug!("Cmd::FocusMaster: {:#x}", window);
                            workspace.move_focus(MoveOp::Swap)
                        }
                        Cmd::FocusParentUp => {
                            debug!("Cmd::FocusParentUp: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Up)
                        }
                        Cmd::FocusParentDown => {
                            debug!("Cmd::FocusParentDown: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Down)
                        }
                        Cmd::FocusParentMaster => {
                            debug!("Cmd::FocusParentMaster: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Swap)
                        }
                        _ => None
                    };

                    if let Some(window) = new_focus {
                        xws.focus_window(window);
                    }
                }
            },
            Cmd::SwapUp | Cmd::SwapDown | Cmd::SwapMaster | Cmd::SwapParentUp | Cmd::SwapParentDown | Cmd::SwapParentMaster => {
                if let Some(window) = state.current_ws().focused_window() {
                    let workspace = state.current_ws_mut();
                    let new_focus = match self {
                        Cmd::SwapUp => {
                            debug!("Cmd::SwapUp: {:#x}", window);
                            workspace.move_window(MoveOp::Up)
                        }
                        Cmd::SwapDown => {
                            debug!("Cmd::SwapDown: {:#x}", window);
                            workspace.move_window(MoveOp::Down)
                        }
                        Cmd::SwapMaster => {
                            debug!("Cmd::SwapMaster: {:#x}", window);
                            workspace.move_window(MoveOp::Swap)
                        }
                        Cmd::SwapParentUp => {
                            debug!("Cmd::SwapParentUp: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Up)
                        }
                        Cmd::SwapParentDown => {
                            debug!("Cmd::SwapParentDown: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Down)
                        }
                        Cmd::SwapParentMaster => {
                            debug!("Cmd::SwapParentMaster: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Swap)
                        }
                        _ => false
                    };

                    if new_focus {
                        state.current_ws().redraw(xws, config, state.get_screens());
                    }
                }
            }
        }
        Ok(())
    }
}

fn reload(state: &WmState) -> Result<()> {
    info!("recompiling...");

    let cfg_dir = Config::get_dir()
        .context("failed to locate config directory")?;
    let mut cmd = Command::new("cargo");

    let output = cmd.current_dir(&format!("{}/.build", cfg_dir))
        .arg("build")
        .env("RUST_LOG", "none")
        .output()
        .context("failed to run cargo")?;

    if !output.status.success() {
        let stderr_msg = String::from_utf8(output.stderr)
            .context("failed to convert cargo stderr to UTF-8")?;
        bail!("failed to recompile: {}", stderr_msg)
    }

    debug!("Cmd::Reload: restarting xr3wm...");

    let path = Path::new(&cfg_dir).join(".state.tmp");

    // save current workspace states to restore on restart
    let file = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&path)
        .context("failed to open workspace state tmp file")?;

    serde_json::to_writer(file, state)
        .context("failed to serialize workspace states")?;

    let args: Vec<*const libc::c_char> = env::args()
        .filter_map(|x| CString::new(x).ok())
        .map(|x| x.into_raw() as *const libc::c_char)
        .chain(iter::once(null()))
        .collect();

    let envs: Vec<*const libc::c_char> = env::vars()
        .filter_map(|(k,v)| CString::new(format!("{}={}", k, v)).ok())
        .map(|x| x.into_raw() as *const libc::c_char)
        .chain(iter::once(null()))
        .collect();

    unsafe {
        execvpe(args[0] as *const libc::c_char, args.as_ptr(), envs.as_ptr());
        // execvp returns only if an error has occurred
        error!("failed to reload: {}", ::std::io::Error::last_os_error());
    }

    Ok(())
}

pub struct ManageHook {
    pub class_name: String,
    pub cmd: CmdManage,
}

pub enum CmdManage {
    Move(usize),
    Float,
    Fullscreen,
    Ignore,
}

impl CmdManage {
    pub fn call(&self,
                xws: &XlibWindowSystem,
                state: &mut WmState,
                config: &Config,
                window: Window) {
        match *self {
            CmdManage::Move(index) => {
                debug!("CmdManage::Move: {}, {}", window, index);
                state.add_window(Some(index - 1), xws, config, window);
            }
            CmdManage::Float => {
                debug!("CmdManage::Float");
                unimplemented!()
            }
            CmdManage::Fullscreen => {
                debug!("CmdManage::Fullscreen");
                unimplemented!()
            }
            CmdManage::Ignore => {
                debug!("CmdManage::Ignore");
                unimplemented!()
            }
        }
    }
}

fn exec(cmd: String, args: Vec<String>) {
    if !cmd.is_empty() {
        let mut cmd = Command::new(cmd);

        if !args.is_empty() {
            cmd.args(&args);
        }

        match cmd.envs(env::vars()).spawn() {
            Ok(_) => (),
            Err(e) => error!("failed to start \"{:?}\": {}", cmd, e),
        }
    }
}
<commit_msg>Prevent zombie child processes<commit_after>#![allow(dead_code)]

extern crate libc;

use crate::config::Config;
use crate::layout::{Layout, LayoutMsg};
use crate::xlib_window_system::XlibWindowSystem;
use crate::state::WmState;
use crate::workspace::MoveOp;
use self::libc::execvpe;
use std::{env, iter};
use std::ptr::null;
use std::ffi::CString;
use std::process::Command;
use std::path::Path;
use std::fs::OpenOptions;
use x11::xlib::Window;
use anyhow::{bail, Context, Result};

pub enum Cmd {
    Custom(Box<dyn Fn(&WmState) -> Result<Option<Cmd>, String>>),
    Exec(String, Vec<String>),
    SwitchWorkspace(usize),
    SwitchScreen(usize),
    MoveToWorkspace(usize),
    MoveToScreen(usize),
    SendLayoutMsg(LayoutMsg),
    NestLayout(Box<dyn Fn() -> Box<dyn Layout>>),
    Reload,
    Exit,
    KillClient,
    FocusParentUp,
    FocusParentDown,
    FocusParentMaster,
    FocusUp,
    FocusDown,
    FocusMaster,
    SwapUp,
    SwapDown,
    SwapMaster,
    SwapParentUp,
    SwapParentDown,
    SwapParentMaster,
}

impl Cmd {
    pub fn call(&self, xws: &XlibWindowSystem, state: &mut WmState, config: &Config) -> Result<()> {
        match self {
            Cmd::Custom(func) => {
                debug!("Cmd::Custom");
                match func(state) {
                    Ok(Some(cmd)) => cmd.call(xws, state, config)?,
                    Ok(None) => (),
                    Err(e) => error!("Cmd::Custom failed: {}", e),
                }
            }
            Cmd::Exec(ref cmd, ref args) => {
                debug!("Cmd::Exec: {} {:?}", cmd, args);
                exec(cmd.clone(), args.clone());
            }
            Cmd::SwitchWorkspace(index) => {
                debug!("Cmd::SwitchWorkspace: {}", index);
                state.switch_to_ws(xws, config, index - 1, true);
            }
            Cmd::SwitchScreen(screen) => {
                debug!("Cmd::SwitchScreen: {}", screen);
                state.switch_to_screen(xws, config, screen - 1);
            }
            Cmd::MoveToWorkspace(index) => {
                debug!("Cmd::MoveToWorkspace: {}", index);
                state.move_window_to_ws(xws, config, index - 1);
            }
            Cmd::MoveToScreen(screen) => {
                debug!("Cmd::MoveToScreen: {}", screen);
                state.move_window_to_screen(xws, config, screen - 1);
            }
            Cmd::SendLayoutMsg(ref msg) => {
                debug!("Cmd::SendLayoutMsg::{:?}", msg);
                state.current_ws_mut().send_layout_message(xws, msg.clone());
                state.current_ws().redraw(xws, config, state.get_screens());
            }
            Cmd::NestLayout(layout_fn) => {
                let layout = layout_fn();
                debug!("Cmd::NestLayout: {}", layout.name());
                state.current_ws_mut().nest_layout(layout);
            }
            Cmd::Reload => {
                debug!("Cmd::Reload");
                reload(state)
                    .context("failed to reload xr3wm")?;
            }
            Cmd::Exit => {
                debug!("Cmd::Exit");
                xws.close();
            }
            Cmd::KillClient => {
                if let Some(window) = state.current_ws().focused_window() {
                    debug!("Cmd::KillClient: {:#x}", window);
                    xws.kill_window(window);
                }
            }
            Cmd::FocusUp | Cmd::FocusDown | Cmd::FocusMaster | Cmd::FocusParentUp | Cmd::FocusParentDown | Cmd::FocusParentMaster => {
                if let Some(window) = state.current_ws().focused_window() {
                    let workspace = state.current_ws_mut();
                    let new_focus = match self {
                        Cmd::FocusUp => {
                            debug!("Cmd::FocusUp: {:#x}", window);
                            workspace.move_focus(MoveOp::Up)
                        }
                        Cmd::FocusDown => {
                            debug!("Cmd::FocusDown: {:#x}", window);
                            workspace.move_focus(MoveOp::Down)
                        }
                        Cmd::FocusMaster => {
                            debug!("Cmd::FocusMaster: {:#x}", window);
                            workspace.move_focus(MoveOp::Swap)
                        }
                        Cmd::FocusParentUp => {
                            debug!("Cmd::FocusParentUp: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Up)
                        }
                        Cmd::FocusParentDown => {
                            debug!("Cmd::FocusParentDown: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Down)
                        }
                        Cmd::FocusParentMaster => {
                            debug!("Cmd::FocusParentMaster: {:#x}", window);
                            workspace.move_parent_focus(MoveOp::Swap)
                        }
                        _ => None
                    };

                    if let Some(window) = new_focus {
                        xws.focus_window(window);
                    }
                }
            },
            Cmd::SwapUp | Cmd::SwapDown | Cmd::SwapMaster | Cmd::SwapParentUp | Cmd::SwapParentDown | Cmd::SwapParentMaster => {
                if let Some(window) = state.current_ws().focused_window() {
                    let workspace = state.current_ws_mut();
                    let new_focus = match self {
                        Cmd::SwapUp => {
                            debug!("Cmd::SwapUp: {:#x}", window);
                            workspace.move_window(MoveOp::Up)
                        }
                        Cmd::SwapDown => {
                            debug!("Cmd::SwapDown: {:#x}", window);
                            workspace.move_window(MoveOp::Down)
                        }
                        Cmd::SwapMaster => {
                            debug!("Cmd::SwapMaster: {:#x}", window);
                            workspace.move_window(MoveOp::Swap)
                        }
                        Cmd::SwapParentUp => {
                            debug!("Cmd::SwapParentUp: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Up)
                        }
                        Cmd::SwapParentDown => {
                            debug!("Cmd::SwapParentDown: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Down)
                        }
                        Cmd::SwapParentMaster => {
                            debug!("Cmd::SwapParentMaster: {:#x}", window);
                            workspace.move_parent_window(MoveOp::Swap)
                        }
                        _ => false
                    };

                    if new_focus {
                        state.current_ws().redraw(xws, config, state.get_screens());
                    }
                }
            }
        }
        Ok(())
    }
}

fn reload(state: &WmState) -> Result<()> {
    info!("recompiling...");

    let cfg_dir = Config::get_dir()
        .context("failed to locate config directory")?;
    let mut cmd = Command::new("cargo");

    let output = cmd.current_dir(&format!("{}/.build", cfg_dir))
        .arg("build")
        .env("RUST_LOG", "none")
        .output()
        .context("failed to run cargo")?;

    if !output.status.success() {
        let stderr_msg = String::from_utf8(output.stderr)
            .context("failed to convert cargo stderr to UTF-8")?;
        bail!("failed to recompile: {}", stderr_msg)
    }

    debug!("Cmd::Reload: restarting xr3wm...");

    let path = Path::new(&cfg_dir).join(".state.tmp");

    // save current workspace states to restore on restart
    let file = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&path)
        .context("failed to open workspace state tmp file")?;

    serde_json::to_writer(file, state)
        .context("failed to serialize workspace states")?;

    let args: Vec<*const libc::c_char> = env::args()
        .filter_map(|x| CString::new(x).ok())
        .map(|x| x.into_raw() as *const libc::c_char)
        .chain(iter::once(null()))
        .collect();

    let envs: Vec<*const libc::c_char> = env::vars()
        .filter_map(|(k,v)| CString::new(format!("{}={}", k, v)).ok())
        .map(|x| x.into_raw() as *const libc::c_char)
        .chain(iter::once(null()))
        .collect();

    unsafe {
        execvpe(args[0] as *const libc::c_char, args.as_ptr(), envs.as_ptr());
        // execvp returns only if an error has occurred
        error!("failed to reload: {}", ::std::io::Error::last_os_error());
    }

    Ok(())
}

pub struct ManageHook {
    pub class_name: String,
    pub cmd: CmdManage,
}

pub enum CmdManage {
    Move(usize),
    Float,
    Fullscreen,
    Ignore,
}

impl CmdManage {
    pub fn call(&self,
                xws: &XlibWindowSystem,
                state: &mut WmState,
                config: &Config,
                window: Window) {
        match *self {
            CmdManage::Move(index) => {
                debug!("CmdManage::Move: {}, {}", window, index);
                state.add_window(Some(index - 1), xws, config, window);
            }
            CmdManage::Float => {
                debug!("CmdManage::Float");
                unimplemented!()
            }
            CmdManage::Fullscreen => {
                debug!("CmdManage::Fullscreen");
                unimplemented!()
            }
            CmdManage::Ignore => {
                debug!("CmdManage::Ignore");
                unimplemented!()
            }
        }
    }
}

fn exec(cmd: String, args: Vec<String>) {
    if !cmd.is_empty() {
        let mut cmd = Command::new(cmd);

        if !args.is_empty() {
            cmd.args(&args);
        }

        std::thread::spawn(move || {
            match cmd.envs(env::vars()).spawn() {
                Ok(mut child) => {
                    child.wait().ok();
                },
                Err(e) => error!("failed to start \"{:?}\": {}", cmd, e),
            }
        });
    }
}
<|endoftext|>


<commit_before>use asm_ops::*;

use std::error::Error;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::process::Command;

pub fn compile(ops: &Vec<AsmOp>) -> Vec<u8>{
    let mut str = "use64\n".to_string();
    for op in ops{
        let temp_str = match *op {
            AsmOp::Add(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) =>  format!("add {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("add {:?}, {:?}\n", dest, i),
                _ => "".to_string()
            },
            AsmOp::Mul(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) => format!("imul {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("imul {:?}, {:?}\n", dest, i),
                _ => "".to_string()
            },
            AsmOp::Pop(ref dest) => match *dest {
                AsmOperand::RegisterOperand(ref dest) =>  format!("pop {:?}\n", dest),
                AsmOperand::Value(ref dest) => format!("pop {:?}\n", dest),
                AsmOperand::Memory(ref mem) => format!("popq [{}]\n", mem)
            },
            AsmOp::Push(ref dest) => match *dest {
                AsmOperand::RegisterOperand(ref dest) => format!("push {:?}\n", dest),
                AsmOperand::Value(ref dest) => format!("push {:?}\n", dest),
                AsmOperand::Memory(ref mem) => format!("pushq [{}]\n", mem)
            },
            AsmOp::Mov(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) => format!("mov {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("mov {:?}, {:?}\n", dest, i),
                AsmOperand::Memory(ref i) => format!("mov {:?}, [{:?}]\n", dest, i),
            },
            AsmOp::Out => "ret".to_string() ,
            // _ => {}
        };
        str = str + &temp_str;
    }
    println!("\nOutput assembly is:\n{}\n",str);

    write_asm(&str);
    assemble();
    read_bytes()
}

fn write_asm(str: &String) {
    let path = Path::new("target/temp.asm");
    let display = path.display();

    let mut file = match File::create(&path) {
        Err(why) => panic!("couldn't create {}: {}", display, Error::description(&why)),
        Ok(file) => file,
    };

    match file.write_all(str.as_bytes()) {
        Err(why) => {
            panic!("couldn't write to {}: {}", display, Error::description(&why))
        },
        Ok(_) => println!("Successfully wrote to {}", display),
    }
}

fn assemble() {
    let output = Command::new("fasm").arg("target/temp.asm").output().unwrap_or_else(|e| {
        panic!("failed to execute process: {}", e)
    });

    if !output.status.success() {
        println!("status: {}", output.status);
        println!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        println!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("Can't assemble!");
    } else {
        println!("Successfully assembled to temp.bin");
    }
}

fn read_bytes() -> Vec<u8>{
    let path = Path::new("target/temp.bin");
    let display = path.display();

    let mut file = match File::open(&path) {
       Err(why) => panic!("couldn't open {}: {}", display, Error::description(&why)),
       Ok(file) => file,
   };

    let mut contents: Vec<u8> = Vec::new();
    // Returns amount of bytes read and append the result to the buffer
    let result = file.read_to_end(&mut contents).unwrap();
    println!("Program consists of {} bytes", result);
    contents
}
<commit_msg>Unified pop and push instructions<commit_after>use asm_ops::*;

use std::error::Error;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::process::Command;

pub fn compile(ops: &Vec<AsmOp>) -> Vec<u8>{
    let mut str = "use64\n".to_string();
    for op in ops{
        let temp_str = match *op {
            AsmOp::Add(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) =>  format!("add {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("add {:?}, {:?}\n", dest, i),
                _ => "".to_string()
            },
            AsmOp::Mul(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) => format!("imul {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("imul {:?}, {:?}\n", dest, i),
                _ => "".to_string()
            },
            AsmOp::Pop(ref dest) => match *dest {
                AsmOperand::RegisterOperand(ref dest) =>  format!("popq {:?}\n", dest),
                AsmOperand::Value(ref dest) => format!("popq {:?}\n", dest),
                AsmOperand::Memory(ref mem) => format!("popq [{}]\n", mem)
            },
            AsmOp::Push(ref dest) => match *dest {
                AsmOperand::RegisterOperand(ref dest) => format!("pushq {:?}\n", dest),
                AsmOperand::Value(ref dest) => format!("pushq {:?}\n", dest),
                AsmOperand::Memory(ref mem) => format!("pushq [{}]\n", mem)
            },
            AsmOp::Mov(ref dest, ref operand) => match *operand{
                AsmOperand::RegisterOperand(ref source) => format!("mov {:?}, {:?}\n", dest, source),
                AsmOperand::Value(ref i) => format!("mov {:?}, {:?}\n", dest, i),
                AsmOperand::Memory(ref i) => format!("mov {:?}, [{:?}]\n", dest, i),
            },
            AsmOp::Out => "ret".to_string() ,
            // _ => {}
        };
        str = str + &temp_str;
    }
    println!("\nOutput assembly is:\n{}\n",str);

    write_asm(&str);
    assemble();
    read_bytes()
}

fn write_asm(str: &String) {
    let path = Path::new("target/temp.asm");
    let display = path.display();

    let mut file = match File::create(&path) {
        Err(why) => panic!("couldn't create {}: {}", display, Error::description(&why)),
        Ok(file) => file,
    };

    match file.write_all(str.as_bytes()) {
        Err(why) => {
            panic!("couldn't write to {}: {}", display, Error::description(&why))
        },
        Ok(_) => println!("Successfully wrote to {}", display),
    }
}

fn assemble() {
    let output = Command::new("fasm").arg("target/temp.asm").output().unwrap_or_else(|e| {
        panic!("failed to execute process: {}", e)
    });

    if !output.status.success() {
        println!("status: {}", output.status);
        println!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        println!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("Can't assemble!");
    } else {
        println!("Successfully assembled to temp.bin");
    }
}

fn read_bytes() -> Vec<u8>{
    let path = Path::new("target/temp.bin");
    let display = path.display();

    let mut file = match File::open(&path) {
       Err(why) => panic!("couldn't open {}: {}", display, Error::description(&why)),
       Ok(file) => file,
   };

    let mut contents: Vec<u8> = Vec::new();
    // Returns amount of bytes read and append the result to the buffer
    let result = file.read_to_end(&mut contents).unwrap();
    println!("Program consists of {} bytes", result);
    contents
}
<|endoftext|>