<commit_before>use ast::make;
use hir::{HasSource, PathResolution};
use ide_db::{defs::Definition, search::FileReference};
use itertools::izip;
use syntax::{
    ast::{self, edit::AstNodeEdit, ArgListOwner},
    ted, AstNode,
};

use crate::{
    assist_context::{AssistContext, Assists},
    AssistId, AssistKind,
};

// Assist: inline_call
//
// Inlines a function or method body creating a `let` statement per parameter unless the parameter
// can be inlined. The parameter will be inlined either if it the supplied argument is a simple local
// or if the parameter is only accessed inside the function body once.
//
// ```
// # //- minicore: option
// fn foo(name: Option<&str>) {
//     let name = name.unwrap$0();
// }
// ```
// ->
// ```
// fn foo(name: Option<&str>) {
//     let name = match name {
//             Some(val) => val,
//             None => panic!("called `Option::unwrap()` on a `None` value"),
//         };
// }
// ```
pub(crate) fn inline_call(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let (label, function, arguments, expr) =
        if let Some(path_expr) = ctx.find_node_at_offset::<ast::PathExpr>() {
            let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
            let path = path_expr.path()?;

            let function = match ctx.sema.resolve_path(&path)? {
                PathResolution::Def(hir::ModuleDef::Function(f))
                | PathResolution::AssocItem(hir::AssocItem::Function(f)) => f,
                _ => return None,
            };
            (
                format!("Inline `{}`", path),
                function,
                call.arg_list()?.args().collect(),
                ast::Expr::CallExpr(call),
            )
        } else {
            let name_ref: ast::NameRef = ctx.find_node_at_offset()?;
            let call = name_ref.syntax().parent().and_then(ast::MethodCallExpr::cast)?;
            let receiver = call.receiver()?;
            let function = ctx.sema.resolve_method_call(&call)?;
            let mut arguments = vec![receiver];
            arguments.extend(call.arg_list()?.args());
            (format!("Inline `{}`", name_ref), function, arguments, ast::Expr::MethodCallExpr(call))
        };

    inline_(acc, ctx, label, function, arguments, expr)
}

pub(crate) fn inline_(
    acc: &mut Assists,
    ctx: &AssistContext,
    label: String,
    function: hir::Function,
    arg_list: Vec<ast::Expr>,
    expr: ast::Expr,
) -> Option<()> {
    let hir::InFile { value: function_source, file_id } = function.source(ctx.db())?;
    let param_list = function_source.param_list()?;
    let mut assoc_fn_params = function.assoc_fn_params(ctx.sema.db).into_iter();

    let mut params = Vec::new();
    if let Some(self_param) = param_list.self_param() {
        // FIXME this should depend on the receiver as well as the self_param
        params.push((
            make::ident_pat(
                self_param.amp_token().is_some(),
                self_param.mut_token().is_some(),
                make::name("this"),
            )
            .into(),
            assoc_fn_params.next()?,
        ));
    }
    for param in param_list.params() {
        params.push((param.pat()?, assoc_fn_params.next()?));
    }

    if arg_list.len() != params.len() {
        // Can't inline the function because they've passed the wrong number of
        // arguments to this function
        cov_mark::hit!(inline_call_incorrect_number_of_arguments);
        return None;
    }

    let body = function_source.body()?;

    acc.add(
        AssistId("inline_call", AssistKind::RefactorInline),
        label,
        expr.syntax().text_range(),
        |builder| {
            let body = body.clone_for_update();

            let file_id = file_id.original_file(ctx.sema.db);
            let usages_for_locals = |local| {
                Definition::Local(local)
                    .usages(&ctx.sema)
                    .all()
                    .references
                    .remove(&file_id)
                    .unwrap_or_default()
                    .into_iter()
            };
            // Contains the nodes of usages of parameters.
            // If the inner Vec for a parameter is empty it either means there are no usages or that the parameter
            // has a pattern that does not allow inlining
            let param_use_nodes: Vec<Vec<_>> = params
                .iter()
                .map(|(pat, param)| {
                    if !matches!(pat, ast::Pat::IdentPat(pat) if pat.is_simple_ident()) {
                        return Vec::new();
                    }
                    usages_for_locals(param.as_local(ctx.sema.db))
                        .map(|FileReference { name, range, .. }| match name {
                            ast::NameLike::NameRef(_) => body
                                .syntax()
                                .covering_element(range)
                                .ancestors()
                                .nth(3)
                                .and_then(ast::PathExpr::cast),
                            _ => None,
                        })
                        .collect::<Option<Vec<_>>>()
                        .unwrap_or_default()
                })
                .collect();

            // Rewrite `self` to `this`
            if param_list.self_param().is_some() {
                let this = || make::name_ref("this").syntax().clone_for_update();
                usages_for_locals(params[0].1.as_local(ctx.sema.db))
                    .flat_map(|FileReference { name, range, .. }| match name {
                        ast::NameLike::NameRef(_) => Some(body.syntax().covering_element(range)),
                        _ => None,
                    })
                    .for_each(|it| {
                        ted::replace(it, &this());
                    })
            }

            // Inline parameter expressions or generate `let` statements depending on whether inlining works or not.
            for ((pat, _), usages, expr) in izip!(params, param_use_nodes, arg_list).rev() {
                match &*usages {
                    // inline single use parameters
                    [usage] => {
                        let expr = if matches!(expr, ast::Expr::ClosureExpr(_))
                            && usage.syntax().parent().and_then(ast::Expr::cast).is_some()
                        {
                            make::expr_paren(expr)
                        } else {
                            expr
                        };
                        ted::replace(usage.syntax(), expr.syntax().clone_for_update());
                    }
                    // inline parameters whose expression is a simple local reference
                    [_, ..]
                        if matches!(&expr,
                            ast::Expr::PathExpr(expr)
                                if expr.path().and_then(|path| path.as_single_name_ref()).is_some()
                        ) =>
                    {
                        usages.into_iter().for_each(|usage| {
                            ted::replace(usage.syntax(), &expr.syntax().clone_for_update());
                        });
                    }
                    // cant inline, emit a let statement
                    // FIXME: emit type ascriptions when a coercion happens?
                    _ => body.push_front(make::let_stmt(pat, Some(expr)).clone_for_update().into()),
                }
            }

            let original_indentation = expr.indent_level();
            let replacement = body.reset_indent().indent(original_indentation);

            let replacement = match replacement.tail_expr() {
                Some(expr) if replacement.statements().next().is_none() => expr,
                _ => ast::Expr::BlockExpr(replacement),
            };
            builder.replace_ast(expr, replacement);
        },
    )
}

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

    use super::*;

    #[test]
    fn no_args_or_return_value_gets_inlined_without_block() {
        check_assist(
            inline_call,
            r#"
fn foo() { println!("Hello, World!"); }
fn main() {
    fo$0o();
}
"#,
            r#"
fn foo() { println!("Hello, World!"); }
fn main() {
    { println!("Hello, World!"); };
}
"#,
        );
    }

    #[test]
    fn not_applicable_when_incorrect_number_of_parameters_are_provided() {
        cov_mark::check!(inline_call_incorrect_number_of_arguments);
        check_assist_not_applicable(
            inline_call,
            r#"
fn add(a: u32, b: u32) -> u32 { a + b }
fn main() { let x = add$0(42); }
"#,
        );
    }

    #[test]
    fn args_with_side_effects() {
        check_assist(
            inline_call,
            r#"
fn foo(name: String) {
    println!("Hello, {}!", name);
}
fn main() {
    foo$0(String::from("Michael"));
}
"#,
            r#"
fn foo(name: String) {
    println!("Hello, {}!", name);
}
fn main() {
    {
        let name = String::from("Michael");
        println!("Hello, {}!", name);
    };
}
"#,
        );
    }

    #[test]
    fn function_with_multiple_statements() {
        check_assist(
            inline_call,
            r#"
fn foo(a: u32, b: u32) -> u32 {
    let x = a + b;
    let y = x - b;
    x * y
}

fn main() {
    let x = foo$0(1, 2);
}
"#,
            r#"
fn foo(a: u32, b: u32) -> u32 {
    let x = a + b;
    let y = x - b;
    x * y
}

fn main() {
    let x = {
        let b = 2;
        let x = 1 + b;
        let y = x - b;
        x * y
    };
}
"#,
        );
    }

    #[test]
    fn function_with_self_param() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo::add$0(Foo(3), 2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(Foo(3).0 + 2);
}
"#,
        );
    }

    #[test]
    fn method_by_val() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(3).add$0(2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(Foo(3).0 + 2);
}
"#,
        );
    }

    #[test]
    fn method_by_ref() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(&self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(3).add$0(2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(&self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = {
        let ref this = Foo(3);
        Foo(this.0 + 2)
    };
}
"#,
        );
    }

    #[test]
    fn method_by_ref_mut() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn clear(&mut self) {
        self.0 = 0;
    }
}

fn main() {
    let mut foo = Foo(3);
    foo.clear$0();
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn clear(&mut self) {
        self.0 = 0;
    }
}

fn main() {
    let mut foo = Foo(3);
    {
        let ref mut this = foo;
        this.0 = 0;
    };
}
"#,
        );
    }

    #[test]
    fn function_single_use_expr_in_param() {
        check_assist(
            inline_call,
            r#"
fn double(x: u32) -> u32 {
    2 * x
}
fn main() {
    let x = 51;
    let x = double$0(10 + x);
}
"#,
            r#"
fn double(x: u32) -> u32 {
    2 * x
}
fn main() {
    let x = 51;
    let x = 2 * 10 + x;
}
"#,
        );
    }

    #[test]
    fn function_multi_use_expr_in_param() {
        check_assist(
            inline_call,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let x = 51;
    let y = square$0(10 + x);
}
"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let x = 51;
    let y = {
        let x = 10 + x;
        x * x
    };
}
"#,
        );
    }

    #[test]
    fn function_multi_use_local_in_param() {
        check_assist(
            inline_call,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let local = 51;
    let y = square$0(local);
}
"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let local = 51;
    let y = local * local;
}
"#,
        );
    }

    #[test]
    fn method_in_impl() {
        check_assist(
            inline_call,
            r#"
struct Foo;
impl Foo {
    fn foo(&self) {
        self;
        self;
    }
    fn bar(&self) {
        self.foo$0();
    }
}
"#,
            r#"
struct Foo;
impl Foo {
    fn foo(&self) {
        self;
        self;
    }
    fn bar(&self) {
        {
            let ref this = self;
            this;
            this;
        };
    }
}
"#,
        );
    }

    #[test]
    fn wraps_closure_in_paren() {
        check_assist(
            inline_call,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    foo$0(|| {})
}
"#,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    {
        (|| {})();
    }
}
"#,
        );
        check_assist(
            inline_call,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    foo$0(main)
}
"#,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    {
        main();
    }
}
"#,
        );
    }
}
<commit_msg>Only inline closure, literal and local arguments when used once<commit_after>use ast::make;
use hir::{HasSource, PathResolution};
use ide_db::{defs::Definition, search::FileReference};
use itertools::izip;
use syntax::{
    ast::{self, edit::AstNodeEdit, ArgListOwner},
    ted, AstNode,
};

use crate::{
    assist_context::{AssistContext, Assists},
    AssistId, AssistKind,
};

// Assist: inline_call
//
// Inlines a function or method body creating a `let` statement per parameter unless the parameter
// can be inlined. The parameter will be inlined either if it the supplied argument is a simple local
// or if the parameter is only accessed inside the function body once.
//
// ```
// # //- minicore: option
// fn foo(name: Option<&str>) {
//     let name = name.unwrap$0();
// }
// ```
// ->
// ```
// fn foo(name: Option<&str>) {
//     let name = match name {
//             Some(val) => val,
//             None => panic!("called `Option::unwrap()` on a `None` value"),
//         };
// }
// ```
pub(crate) fn inline_call(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let (label, function, arguments, expr) =
        if let Some(path_expr) = ctx.find_node_at_offset::<ast::PathExpr>() {
            let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
            let path = path_expr.path()?;

            let function = match ctx.sema.resolve_path(&path)? {
                PathResolution::Def(hir::ModuleDef::Function(f))
                | PathResolution::AssocItem(hir::AssocItem::Function(f)) => f,
                _ => return None,
            };
            (
                format!("Inline `{}`", path),
                function,
                call.arg_list()?.args().collect(),
                ast::Expr::CallExpr(call),
            )
        } else {
            let name_ref: ast::NameRef = ctx.find_node_at_offset()?;
            let call = name_ref.syntax().parent().and_then(ast::MethodCallExpr::cast)?;
            let receiver = call.receiver()?;
            let function = ctx.sema.resolve_method_call(&call)?;
            let mut arguments = vec![receiver];
            arguments.extend(call.arg_list()?.args());
            (format!("Inline `{}`", name_ref), function, arguments, ast::Expr::MethodCallExpr(call))
        };

    inline_(acc, ctx, label, function, arguments, expr)
}

pub(crate) fn inline_(
    acc: &mut Assists,
    ctx: &AssistContext,
    label: String,
    function: hir::Function,
    arg_list: Vec<ast::Expr>,
    expr: ast::Expr,
) -> Option<()> {
    let hir::InFile { value: function_source, file_id } = function.source(ctx.db())?;
    let param_list = function_source.param_list()?;
    let mut assoc_fn_params = function.assoc_fn_params(ctx.sema.db).into_iter();

    let mut params = Vec::new();
    if let Some(self_param) = param_list.self_param() {
        // FIXME this should depend on the receiver as well as the self_param
        params.push((
            make::ident_pat(
                self_param.amp_token().is_some(),
                self_param.mut_token().is_some(),
                make::name("this"),
            )
            .into(),
            assoc_fn_params.next()?,
        ));
    }
    for param in param_list.params() {
        params.push((param.pat()?, assoc_fn_params.next()?));
    }

    if arg_list.len() != params.len() {
        // Can't inline the function because they've passed the wrong number of
        // arguments to this function
        cov_mark::hit!(inline_call_incorrect_number_of_arguments);
        return None;
    }

    let body = function_source.body()?;

    acc.add(
        AssistId("inline_call", AssistKind::RefactorInline),
        label,
        expr.syntax().text_range(),
        |builder| {
            let body = body.clone_for_update();

            let file_id = file_id.original_file(ctx.sema.db);
            let usages_for_locals = |local| {
                Definition::Local(local)
                    .usages(&ctx.sema)
                    .all()
                    .references
                    .remove(&file_id)
                    .unwrap_or_default()
                    .into_iter()
            };
            // Contains the nodes of usages of parameters.
            // If the inner Vec for a parameter is empty it either means there are no usages or that the parameter
            // has a pattern that does not allow inlining
            let param_use_nodes: Vec<Vec<_>> = params
                .iter()
                .map(|(pat, param)| {
                    if !matches!(pat, ast::Pat::IdentPat(pat) if pat.is_simple_ident()) {
                        return Vec::new();
                    }
                    usages_for_locals(param.as_local(ctx.sema.db))
                        .map(|FileReference { name, range, .. }| match name {
                            ast::NameLike::NameRef(_) => body
                                .syntax()
                                .covering_element(range)
                                .ancestors()
                                .nth(3)
                                .and_then(ast::PathExpr::cast),
                            _ => None,
                        })
                        .collect::<Option<Vec<_>>>()
                        .unwrap_or_default()
                })
                .collect();

            // Rewrite `self` to `this`
            if param_list.self_param().is_some() {
                let this = || make::name_ref("this").syntax().clone_for_update();
                usages_for_locals(params[0].1.as_local(ctx.sema.db))
                    .flat_map(|FileReference { name, range, .. }| match name {
                        ast::NameLike::NameRef(_) => Some(body.syntax().covering_element(range)),
                        _ => None,
                    })
                    .for_each(|it| {
                        ted::replace(it, &this());
                    })
            }

            // Inline parameter expressions or generate `let` statements depending on whether inlining works or not.
            for ((pat, _), usages, expr) in izip!(params, param_use_nodes, arg_list).rev() {
                let expr_is_name_ref = matches!(&expr,
                    ast::Expr::PathExpr(expr)
                        if expr.path().and_then(|path| path.as_single_name_ref()).is_some()
                );
                match &*usages {
                    // inline single use closure arguments
                    [usage]
                        if matches!(expr, ast::Expr::ClosureExpr(_))
                            && usage.syntax().parent().and_then(ast::Expr::cast).is_some() =>
                    {
                        cov_mark::hit!(inline_call_inline_closure);
                        let expr = make::expr_paren(expr);
                        ted::replace(usage.syntax(), expr.syntax().clone_for_update());
                    }
                    // inline single use literals
                    [usage] if matches!(expr, ast::Expr::Literal(_)) => {
                        cov_mark::hit!(inline_call_inline_literal);
                        ted::replace(usage.syntax(), expr.syntax().clone_for_update());
                    }
                    // inline direct local arguments
                    [_, ..] if expr_is_name_ref => {
                        cov_mark::hit!(inline_call_inline_locals);
                        usages.into_iter().for_each(|usage| {
                            ted::replace(usage.syntax(), &expr.syntax().clone_for_update());
                        });
                    }
                    // cant inline, emit a let statement
                    // FIXME: emit type ascriptions when a coercion happens?
                    _ => body.push_front(make::let_stmt(pat, Some(expr)).clone_for_update().into()),
                }
            }

            let original_indentation = expr.indent_level();
            let replacement = body.reset_indent().indent(original_indentation);

            let replacement = match replacement.tail_expr() {
                Some(expr) if replacement.statements().next().is_none() => expr,
                _ => ast::Expr::BlockExpr(replacement),
            };
            builder.replace_ast(expr, replacement);
        },
    )
}

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

    use super::*;

    #[test]
    fn no_args_or_return_value_gets_inlined_without_block() {
        check_assist(
            inline_call,
            r#"
fn foo() { println!("Hello, World!"); }
fn main() {
    fo$0o();
}
"#,
            r#"
fn foo() { println!("Hello, World!"); }
fn main() {
    { println!("Hello, World!"); };
}
"#,
        );
    }

    #[test]
    fn not_applicable_when_incorrect_number_of_parameters_are_provided() {
        cov_mark::check!(inline_call_incorrect_number_of_arguments);
        check_assist_not_applicable(
            inline_call,
            r#"
fn add(a: u32, b: u32) -> u32 { a + b }
fn main() { let x = add$0(42); }
"#,
        );
    }

    #[test]
    fn args_with_side_effects() {
        check_assist(
            inline_call,
            r#"
fn foo(name: String) {
    println!("Hello, {}!", name);
}
fn main() {
    foo$0(String::from("Michael"));
}
"#,
            r#"
fn foo(name: String) {
    println!("Hello, {}!", name);
}
fn main() {
    {
        let name = String::from("Michael");
        println!("Hello, {}!", name);
    };
}
"#,
        );
    }

    #[test]
    fn function_with_multiple_statements() {
        check_assist(
            inline_call,
            r#"
fn foo(a: u32, b: u32) -> u32 {
    let x = a + b;
    let y = x - b;
    x * y
}

fn main() {
    let x = foo$0(1, 2);
}
"#,
            r#"
fn foo(a: u32, b: u32) -> u32 {
    let x = a + b;
    let y = x - b;
    x * y
}

fn main() {
    let x = {
        let b = 2;
        let x = 1 + b;
        let y = x - b;
        x * y
    };
}
"#,
        );
    }

    #[test]
    fn function_with_self_param() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo::add$0(Foo(3), 2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = {
        let this = Foo(3);
        Foo(this.0 + 2)
    };
}
"#,
        );
    }

    #[test]
    fn method_by_val() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(3).add$0(2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = {
        let this = Foo(3);
        Foo(this.0 + 2)
    };
}
"#,
        );
    }

    #[test]
    fn method_by_ref() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn add(&self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = Foo(3).add$0(2);
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn add(&self, a: u32) -> Self {
        Foo(self.0 + a)
    }
}

fn main() {
    let x = {
        let ref this = Foo(3);
        Foo(this.0 + 2)
    };
}
"#,
        );
    }

    #[test]
    fn method_by_ref_mut() {
        check_assist(
            inline_call,
            r#"
struct Foo(u32);

impl Foo {
    fn clear(&mut self) {
        self.0 = 0;
    }
}

fn main() {
    let mut foo = Foo(3);
    foo.clear$0();
}
"#,
            r#"
struct Foo(u32);

impl Foo {
    fn clear(&mut self) {
        self.0 = 0;
    }
}

fn main() {
    let mut foo = Foo(3);
    {
        let ref mut this = foo;
        this.0 = 0;
    };
}
"#,
        );
    }

    #[test]
    fn function_multi_use_expr_in_param() {
        check_assist(
            inline_call,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let x = 51;
    let y = square$0(10 + x);
}
"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let x = 51;
    let y = {
        let x = 10 + x;
        x * x
    };
}
"#,
        );
    }

    #[test]
    fn function_use_local_in_param() {
        cov_mark::check!(inline_call_inline_locals);
        check_assist(
            inline_call,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let local = 51;
    let y = square$0(local);
}
"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}
fn main() {
    let local = 51;
    let y = local * local;
}
"#,
        );
    }

    #[test]
    fn method_in_impl() {
        check_assist(
            inline_call,
            r#"
struct Foo;
impl Foo {
    fn foo(&self) {
        self;
        self;
    }
    fn bar(&self) {
        self.foo$0();
    }
}
"#,
            r#"
struct Foo;
impl Foo {
    fn foo(&self) {
        self;
        self;
    }
    fn bar(&self) {
        {
            let ref this = self;
            this;
            this;
        };
    }
}
"#,
        );
    }

    #[test]
    fn wraps_closure_in_paren() {
        cov_mark::check!(inline_call_inline_closure);
        check_assist(
            inline_call,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    foo$0(|| {})
}
"#,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    {
        (|| {})();
    }
}
"#,
        );
        check_assist(
            inline_call,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    foo$0(main)
}
"#,
            r#"
fn foo(x: fn()) {
    x();
}

fn main() {
    {
        main();
    }
}
"#,
        );
    }

    #[test]
    fn inline_single_literal_expr() {
        cov_mark::check!(inline_call_inline_literal);
        check_assist(
            inline_call,
            r#"
fn foo(x: u32) -> u32{
    x
}

fn main() {
    foo$0(222);
}
"#,
            r#"
fn foo(x: u32) -> u32{
    x
}

fn main() {
    222;
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>use crate::{AssistContext, AssistId, Assists};

use ast::LoopBodyOwner;
use ra_fmt::unwrap_trivial_block;
use ra_syntax::{ast, match_ast, AstNode, TextRange, T};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
    let block = ast::BlockExpr::cast(l_curly_token.parent())?;
    let parent = block.syntax().parent()?;
    let (expr, expr_to_unwrap) = match_ast! {
        match parent {
            ast::IfExpr(if_expr) => {
                let expr_to_unwrap = if_expr.blocks().find_map(|expr| extract_expr(ctx.frange.range, expr));
                let expr_to_unwrap = expr_to_unwrap?;
                // Find if we are in a else if block
                let ancestor = if_expr.syntax().parent().and_then(ast::IfExpr::cast);

                match ancestor {
                    None => (ast::Expr::IfExpr(if_expr), expr_to_unwrap),
                    Some(ancestor) => (ast::Expr::IfExpr(ancestor), expr_to_unwrap),
                }
            },
            ast::ForExpr(for_expr) => {
                let block_expr = for_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::ForExpr(for_expr), expr_to_unwrap)
            },
            ast::WhileExpr(while_expr) => {
                let block_expr = while_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::WhileExpr(while_expr), expr_to_unwrap)
            },
            ast::LoopExpr(loop_expr) => {
                let block_expr = loop_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap)
            },
            _ => return None,
        }
    };

    let target = expr_to_unwrap.syntax().text_range();
    acc.add(AssistId("unwrap_block"), "Unwrap block", target, |edit| {
        edit.set_cursor(expr.syntax().text_range().start());

        let pat_start: &[_] = &[' ', '{', '\n'];
        let expr_to_unwrap = expr_to_unwrap.to_string();
        let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
        let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
        expr_string_lines.pop(); // Delete last line

        let expr_string = expr_string_lines
            .into_iter()
            .map(|line| line.replacen("    ", "", 1)) // Delete indentation
            .collect::<Vec<String>>()
            .join("\n");

        edit.replace(expr.syntax().text_range(), expr_string);
    })
}

fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::Expr> {
    let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);

    if cursor_in_range {
        Some(unwrap_trivial_block(block))
    } else {
        None
    }
}

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

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                <|>println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    <|>foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<commit_msg>do not remove then block when you unwrap else block #4361<commit_after>use crate::{AssistContext, AssistId, Assists};

use ast::{ElseBranch, Expr, LoopBodyOwner};
use ra_fmt::unwrap_trivial_block;
use ra_syntax::{ast, match_ast, AstNode, TextRange, T};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
    let block = ast::BlockExpr::cast(l_curly_token.parent())?;
    let parent = block.syntax().parent()?;
    let assist_id = AssistId("unwrap_block");
    let assist_label = "Unwrap block";

    let (expr, expr_to_unwrap) = match_ast! {
        match parent {
            ast::ForExpr(for_expr) => {
                let block_expr = for_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::ForExpr(for_expr), expr_to_unwrap)
            },
            ast::WhileExpr(while_expr) => {
                let block_expr = while_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::WhileExpr(while_expr), expr_to_unwrap)
            },
            ast::LoopExpr(loop_expr) => {
                let block_expr = loop_expr.loop_body()?;
                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
                (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap)
            },
            ast::IfExpr(if_expr) => {
                let mut resp = None;

                let then_branch = if_expr.then_branch()?;
                if then_branch.l_curly_token()?.text_range().contains_range(ctx.frange.range) {
                    if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
                        // For `else if` blocks
                        let ancestor_then_branch = ancestor.then_branch()?;
                        let l_curly_token = then_branch.l_curly_token()?;

                        let target = then_branch.syntax().text_range();
                        return acc.add(assist_id, assist_label, target, |edit| {
                            let range_to_del_else_if = TextRange::new(ancestor_then_branch.syntax().text_range().end(), l_curly_token.text_range().start());
                            let range_to_del_rest = TextRange::new(then_branch.syntax().text_range().end(), if_expr.syntax().text_range().end());

                            edit.set_cursor(ancestor_then_branch.syntax().text_range().end());
                            edit.delete(range_to_del_rest);
                            edit.delete(range_to_del_else_if);
                            edit.replace(target, update_expr_string(then_branch.to_string(), &[' ', '{']));
                        });
                    } else {
                        resp = Some((ast::Expr::IfExpr(if_expr.clone()), Expr::BlockExpr(then_branch)));
                    }
                } else if let Some(else_branch) = if_expr.else_branch() {
                    match else_branch {
                        ElseBranch::Block(else_block) => {
                            let l_curly_token = else_block.l_curly_token()?;
                            if l_curly_token.text_range().contains_range(ctx.frange.range) {
                                let target = else_block.syntax().text_range();
                                return acc.add(assist_id, assist_label, target, |edit| {
                                    let range_to_del = TextRange::new(then_branch.syntax().text_range().end(), l_curly_token.text_range().start());

                                    edit.set_cursor(then_branch.syntax().text_range().end());
                                    edit.delete(range_to_del);
                                    edit.replace(target, update_expr_string(else_block.to_string(), &[' ', '{']));
                                });
                            }
                        },
                        ElseBranch::IfExpr(_) => {},
                    }
                }

                resp?
            },
            _ => return None,
        }
    };

    let target = expr_to_unwrap.syntax().text_range();
    acc.add(assist_id, assist_label, target, |edit| {
        edit.set_cursor(expr.syntax().text_range().start());

        edit.replace(
            expr.syntax().text_range(),
            update_expr_string(expr_to_unwrap.to_string(), &[' ', '{', '\n']),
        );
    })
}

fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::Expr> {
    let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);

    if cursor_in_range {
        Some(unwrap_trivial_block(block))
    } else {
        None
    }
}

fn update_expr_string(expr_str: String, trim_start_pat: &[char]) -> String {
    let expr_string = expr_str.trim_start_matches(trim_start_pat);
    let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
    expr_string_lines.pop(); // Delete last line

    expr_string_lines
        .into_iter()
        .map(|line| line.replacen("    ", "", 1)) // Delete indentation
        .collect::<Vec<String>>()
        .join("\n")
}

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

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                <|>foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                }<|>
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                }<|>
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }<|>
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                } else {<|>
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                }<|>
                println!("else");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_middle() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                } else {
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }<|>
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    <|>foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                <|>if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A module for working with borrowed data.

#![stable(feature = "rust1", since = "1.0.0")]

use core::cmp::Ordering;
use core::hash::{Hash, Hasher};
use core::ops::{Add, AddAssign, Deref};

use fmt;

use self::Cow::*;

#[stable(feature = "rust1", since = "1.0.0")]
pub use core::borrow::{Borrow, BorrowMut};

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Borrow<B> for Cow<'a, B>
    where B: ToOwned,
          <B as ToOwned>::Owned: 'a
{
    fn borrow(&self) -> &B {
        &**self
    }
}

/// A generalization of `Clone` to borrowed data.
///
/// Some types make it possible to go from borrowed to owned, usually by
/// implementing the `Clone` trait. But `Clone` works only for going from `&T`
/// to `T`. The `ToOwned` trait generalizes `Clone` to construct owned data
/// from any borrow of a given type.
#[stable(feature = "rust1", since = "1.0.0")]
pub trait ToOwned {
    #[stable(feature = "rust1", since = "1.0.0")]
    type Owned: Borrow<Self>;

    /// Creates owned data from borrowed data, usually by cloning.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let s = "a"; // &str
    /// let ss = s.to_owned(); // String
    ///
    /// let v = &[1, 2]; // slice
    /// let vv = v.to_owned(); // Vec
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn to_owned(&self) -> Self::Owned;
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> ToOwned for T where T: Clone {
    type Owned = T;
    fn to_owned(&self) -> T {
        self.clone()
    }
}

/// A clone-on-write smart pointer.
///
/// The type `Cow` is a smart pointer providing clone-on-write functionality: it
/// can enclose and provide immutable access to borrowed data, and clone the
/// data lazily when mutation or ownership is required. The type is designed to
/// work with general borrowed data via the `Borrow` trait.
///
/// `Cow` implements `Deref`, which means that you can call
/// non-mutating methods directly on the data it encloses. If mutation
/// is desired, `to_mut` will obtain a mutable reference to an owned
/// value, cloning if necessary.
///
/// # Examples
///
/// ```
/// use std::borrow::Cow;
///
/// # #[allow(dead_code)]
/// fn abs_all(input: &mut Cow<[i32]>) {
///     for i in 0..input.len() {
///         let v = input[i];
///         if v < 0 {
///             // clones into a vector the first time (if not already owned)
///             input.to_mut()[i] = -v;
///         }
///     }
/// }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub enum Cow<'a, B: ?Sized + 'a>
    where B: ToOwned
{
    /// Borrowed data.
    #[stable(feature = "rust1", since = "1.0.0")]
    Borrowed(#[stable(feature = "rust1", since = "1.0.0")] &'a B),

    /// Owned data.
    #[stable(feature = "rust1", since = "1.0.0")]
    Owned(
        #[stable(feature = "rust1", since = "1.0.0")] <B as ToOwned>::Owned
    ),
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Clone for Cow<'a, B> where B: ToOwned {
    fn clone(&self) -> Cow<'a, B> {
        match *self {
            Borrowed(b) => Borrowed(b),
            Owned(ref o) => {
                let b: &B = o.borrow();
                Owned(b.to_owned())
            }
        }
    }
}

impl<'a, B: ?Sized> Cow<'a, B> where B: ToOwned {
    /// Acquires a mutable reference to the owned form of the data.
    ///
    /// Clones the data if it is not already owned.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::borrow::Cow;
    ///
    /// let mut cow: Cow<[_]> = Cow::Owned(vec![1, 2, 3]);
    ///
    /// let hello = cow.to_mut();
    ///
    /// assert_eq!(hello, &[1, 2, 3]);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn to_mut(&mut self) -> &mut <B as ToOwned>::Owned {
        match *self {
            Borrowed(borrowed) => {
                *self = Owned(borrowed.to_owned());
                self.to_mut()
            }
            Owned(ref mut owned) => owned,
        }
    }

    /// Extracts the owned data.
    ///
    /// Clones the data if it is not already owned.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::borrow::Cow;
    ///
    /// let cow: Cow<[_]> = Cow::Owned(vec![1, 2, 3]);
    ///
    /// let hello = cow.into_owned();
    ///
    /// assert_eq!(vec![1, 2, 3], hello);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn into_owned(self) -> <B as ToOwned>::Owned {
        match self {
            Borrowed(borrowed) => borrowed.to_owned(),
            Owned(owned) => owned,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Deref for Cow<'a, B> where B: ToOwned {
    type Target = B;

    fn deref(&self) -> &B {
        match *self {
            Borrowed(borrowed) => borrowed,
            Owned(ref owned) => owned.borrow(),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Eq for Cow<'a, B> where B: Eq + ToOwned {}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Ord for Cow<'a, B> where B: Ord + ToOwned {
    #[inline]
    fn cmp(&self, other: &Cow<'a, B>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, 'b, B: ?Sized, C: ?Sized> PartialEq<Cow<'b, C>> for Cow<'a, B>
    where B: PartialEq<C> + ToOwned,
          C: ToOwned
{
    #[inline]
    fn eq(&self, other: &Cow<'b, C>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> PartialOrd for Cow<'a, B> where B: PartialOrd + ToOwned {
    #[inline]
    fn partial_cmp(&self, other: &Cow<'a, B>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> fmt::Debug for Cow<'a, B>
    where B: fmt::Debug + ToOwned,
          <B as ToOwned>::Owned: fmt::Debug
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Borrowed(ref b) => fmt::Debug::fmt(b, f),
            Owned(ref o) => fmt::Debug::fmt(o, f),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> fmt::Display for Cow<'a, B>
    where B: fmt::Display + ToOwned,
          <B as ToOwned>::Owned: fmt::Display
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Borrowed(ref b) => fmt::Display::fmt(b, f),
            Owned(ref o) => fmt::Display::fmt(o, f),
        }
    }
}

#[stable(feature = "default", since = "1.11.0")]
impl<'a, B: ?Sized> Default for Cow<'a, B>
    where B: ToOwned,
          <B as ToOwned>::Owned: Default
{
    /// Creates an owned Cow<'a, B> with the default value for the contained owned value.
    fn default() -> Cow<'a, B> {
        Owned(<B as ToOwned>::Owned::default())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Hash for Cow<'a, B> where B: Hash + ToOwned {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        Hash::hash(&**self, state)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
#[allow(deprecated)]
impl<'a, T: ?Sized + ToOwned> AsRef<T> for Cow<'a, T> {
    fn as_ref(&self) -> &T {
        self
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> Add<&'a str> for Cow<'a, str> {
    type Output = Cow<'a, str>;

    fn add(self, rhs: &'a str) -> Self {
        if self == "" {
            Cow::Borrowed(rhs)
        } else if rhs == "" {
            self
        } else {
            Cow::Owned(self.into_owned() + rhs)
        }
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> Add<Cow<'a, str>> for Cow<'a, str> {
    type Output = Cow<'a, str>;

    fn add(self, rhs: Cow<'a, str>) -> Self {
        if self == "" {
            rhs
        } else if rhs == "" {
            self
        } else {
            Cow::Owned(self.into_owned() + rhs.borrow())
        }
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> AddAssign<&'a str> for Cow<'a, str> {
    fn add_assign(&mut self, rhs: &'a str) {
        if rhs == "" { return; }
        self.to_mut().push_str(rhs);
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> AddAssign<Cow<'a, str>> for Cow<'a, str> {
    fn add_assign(&mut self, rhs: Cow<'a, str>) {
        if rhs == "" { return; }
        self.to_mut().push_str(rhs.borrow());
    }
}
<commit_msg>Improve doc example for `std::borrow::Cow`.<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A module for working with borrowed data.

#![stable(feature = "rust1", since = "1.0.0")]

use core::cmp::Ordering;
use core::hash::{Hash, Hasher};
use core::ops::{Add, AddAssign, Deref};

use fmt;

use self::Cow::*;

#[stable(feature = "rust1", since = "1.0.0")]
pub use core::borrow::{Borrow, BorrowMut};

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Borrow<B> for Cow<'a, B>
    where B: ToOwned,
          <B as ToOwned>::Owned: 'a
{
    fn borrow(&self) -> &B {
        &**self
    }
}

/// A generalization of `Clone` to borrowed data.
///
/// Some types make it possible to go from borrowed to owned, usually by
/// implementing the `Clone` trait. But `Clone` works only for going from `&T`
/// to `T`. The `ToOwned` trait generalizes `Clone` to construct owned data
/// from any borrow of a given type.
#[stable(feature = "rust1", since = "1.0.0")]
pub trait ToOwned {
    #[stable(feature = "rust1", since = "1.0.0")]
    type Owned: Borrow<Self>;

    /// Creates owned data from borrowed data, usually by cloning.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let s = "a"; // &str
    /// let ss = s.to_owned(); // String
    ///
    /// let v = &[1, 2]; // slice
    /// let vv = v.to_owned(); // Vec
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn to_owned(&self) -> Self::Owned;
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> ToOwned for T where T: Clone {
    type Owned = T;
    fn to_owned(&self) -> T {
        self.clone()
    }
}

/// A clone-on-write smart pointer.
///
/// The type `Cow` is a smart pointer providing clone-on-write functionality: it
/// can enclose and provide immutable access to borrowed data, and clone the
/// data lazily when mutation or ownership is required. The type is designed to
/// work with general borrowed data via the `Borrow` trait.
///
/// `Cow` implements `Deref`, which means that you can call
/// non-mutating methods directly on the data it encloses. If mutation
/// is desired, `to_mut` will obtain a mutable reference to an owned
/// value, cloning if necessary.
///
/// # Examples
///
/// ```
/// use std::borrow::Cow;
///
/// fn abs_all(input: &mut Cow<[i32]>) {
///     for i in 0..input.len() {
///         let v = input[i];
///         if v < 0 {
///             // Clones into a vector if not already owned.
///             input.to_mut()[i] = -v;
///         }
///     }
/// }
///
/// // No clone occurs because `input` doesn't need to be mutated.
/// let slice = [0, 1, 2];
/// let mut input = Cow::from(&slice[..]);
/// abs_all(&mut input);
///
/// // Clone occurs because `input` needs to be mutated.
/// let slice = [-1, 0, 1];
/// let mut input = Cow::from(&slice[..]);
/// abs_all(&mut input);
///
/// // No clone occurs because `input` is already owned.
/// let mut input = Cow::from(vec![-1, 0, 1]);
/// abs_all(&mut input);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub enum Cow<'a, B: ?Sized + 'a>
    where B: ToOwned
{
    /// Borrowed data.
    #[stable(feature = "rust1", since = "1.0.0")]
    Borrowed(#[stable(feature = "rust1", since = "1.0.0")] &'a B),

    /// Owned data.
    #[stable(feature = "rust1", since = "1.0.0")]
    Owned(
        #[stable(feature = "rust1", since = "1.0.0")] <B as ToOwned>::Owned
    ),
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Clone for Cow<'a, B> where B: ToOwned {
    fn clone(&self) -> Cow<'a, B> {
        match *self {
            Borrowed(b) => Borrowed(b),
            Owned(ref o) => {
                let b: &B = o.borrow();
                Owned(b.to_owned())
            }
        }
    }
}

impl<'a, B: ?Sized> Cow<'a, B> where B: ToOwned {
    /// Acquires a mutable reference to the owned form of the data.
    ///
    /// Clones the data if it is not already owned.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::borrow::Cow;
    ///
    /// let mut cow: Cow<[_]> = Cow::Owned(vec![1, 2, 3]);
    ///
    /// let hello = cow.to_mut();
    ///
    /// assert_eq!(hello, &[1, 2, 3]);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn to_mut(&mut self) -> &mut <B as ToOwned>::Owned {
        match *self {
            Borrowed(borrowed) => {
                *self = Owned(borrowed.to_owned());
                self.to_mut()
            }
            Owned(ref mut owned) => owned,
        }
    }

    /// Extracts the owned data.
    ///
    /// Clones the data if it is not already owned.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::borrow::Cow;
    ///
    /// let cow: Cow<[_]> = Cow::Owned(vec![1, 2, 3]);
    ///
    /// let hello = cow.into_owned();
    ///
    /// assert_eq!(vec![1, 2, 3], hello);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn into_owned(self) -> <B as ToOwned>::Owned {
        match self {
            Borrowed(borrowed) => borrowed.to_owned(),
            Owned(owned) => owned,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Deref for Cow<'a, B> where B: ToOwned {
    type Target = B;

    fn deref(&self) -> &B {
        match *self {
            Borrowed(borrowed) => borrowed,
            Owned(ref owned) => owned.borrow(),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Eq for Cow<'a, B> where B: Eq + ToOwned {}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Ord for Cow<'a, B> where B: Ord + ToOwned {
    #[inline]
    fn cmp(&self, other: &Cow<'a, B>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, 'b, B: ?Sized, C: ?Sized> PartialEq<Cow<'b, C>> for Cow<'a, B>
    where B: PartialEq<C> + ToOwned,
          C: ToOwned
{
    #[inline]
    fn eq(&self, other: &Cow<'b, C>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> PartialOrd for Cow<'a, B> where B: PartialOrd + ToOwned {
    #[inline]
    fn partial_cmp(&self, other: &Cow<'a, B>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> fmt::Debug for Cow<'a, B>
    where B: fmt::Debug + ToOwned,
          <B as ToOwned>::Owned: fmt::Debug
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Borrowed(ref b) => fmt::Debug::fmt(b, f),
            Owned(ref o) => fmt::Debug::fmt(o, f),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> fmt::Display for Cow<'a, B>
    where B: fmt::Display + ToOwned,
          <B as ToOwned>::Owned: fmt::Display
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Borrowed(ref b) => fmt::Display::fmt(b, f),
            Owned(ref o) => fmt::Display::fmt(o, f),
        }
    }
}

#[stable(feature = "default", since = "1.11.0")]
impl<'a, B: ?Sized> Default for Cow<'a, B>
    where B: ToOwned,
          <B as ToOwned>::Owned: Default
{
    /// Creates an owned Cow<'a, B> with the default value for the contained owned value.
    fn default() -> Cow<'a, B> {
        Owned(<B as ToOwned>::Owned::default())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<'a, B: ?Sized> Hash for Cow<'a, B> where B: Hash + ToOwned {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        Hash::hash(&**self, state)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
#[allow(deprecated)]
impl<'a, T: ?Sized + ToOwned> AsRef<T> for Cow<'a, T> {
    fn as_ref(&self) -> &T {
        self
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> Add<&'a str> for Cow<'a, str> {
    type Output = Cow<'a, str>;

    fn add(self, rhs: &'a str) -> Self {
        if self == "" {
            Cow::Borrowed(rhs)
        } else if rhs == "" {
            self
        } else {
            Cow::Owned(self.into_owned() + rhs)
        }
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> Add<Cow<'a, str>> for Cow<'a, str> {
    type Output = Cow<'a, str>;

    fn add(self, rhs: Cow<'a, str>) -> Self {
        if self == "" {
            rhs
        } else if rhs == "" {
            self
        } else {
            Cow::Owned(self.into_owned() + rhs.borrow())
        }
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> AddAssign<&'a str> for Cow<'a, str> {
    fn add_assign(&mut self, rhs: &'a str) {
        if rhs == "" { return; }
        self.to_mut().push_str(rhs);
    }
}

#[stable(feature = "cow_add", since = "1.13.0")]
impl<'a> AddAssign<Cow<'a, str>> for Cow<'a, str> {
    fn add_assign(&mut self, rhs: Cow<'a, str>) {
        if rhs == "" { return; }
        self.to_mut().push_str(rhs.borrow());
    }
}
<|endoftext|>


<commit_before>extern crate websocket;
extern crate futures;
extern crate futures_cpupool;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
extern crate dotenv;
extern crate native_tls;

use self::dotenv::dotenv;

use self::websocket::message::OwnedMessage;
use self::websocket::server::{WsServer, InvalidConnection};
use self::websocket::async::Server;

use self::tokio_core::reactor::{Handle, Remote, Core};
use self::tokio_core::net::TcpListener;

use self::futures::{Future, Sink, Stream};
use self::futures::sync::mpsc;
use self::futures_cpupool::CpuPool;

use self::native_tls::{Pkcs12, TlsAcceptor};

use std::sync::{RwLock, Arc};
use std::rc::Rc;
use std::fmt::Debug;
use std::ops::Deref;
use std::collections::HashMap;
use std::cell::RefCell;
use std::fs::File;
use std::io::Read;

use model::ClientState;
use util::read_env_var;

pub type Id = u32;

type SinkContent = websocket::client::async::Framed<
    websocket::client::async::TlsStream<tokio_core::net::TcpStream>,
    websocket::async::MessageCodec<OwnedMessage>,
>;
type SplitSink = futures::stream::SplitSink<SinkContent>;

#[derive(Clone)]
pub struct Engine {
    pub send_channel: mpsc::UnboundedSender<(Id, Arc<RwLock<ClientState>>)>,
    pub remote: Remote,
}

#[derive(Serialize, Debug)]
pub struct Msg {
    pub id: Id,
    pub content: String,
}

pub trait EventHandler {
    fn new(engine: Engine) -> Self;
    fn main_loop(&self);
    fn message(&self, msg: &Msg);
    fn connect(&self, id: Id) -> bool;
    fn disconnect(&self, id: Id);
}

pub fn execute<T>()
where
    T: EventHandler + Send + Sync + 'static,
{
    dotenv().ok();

    let mut core = Core::new().expect("Failed to create Tokio event loop");
    let handle = core.handle();
    let remote = core.remote();

    let server = build_server(&handle);
    let pool = Arc::new(RwLock::new(CpuPool::new_num_cpus()));
    let connections = Arc::new(RwLock::new(HashMap::new()));
    let (receive_channel_out, receive_channel_in) = mpsc::unbounded();
    let (send_channel_out, send_channel_in) = mpsc::unbounded();

    let engine = Engine {
        send_channel: send_channel_out.clone(),
        remote: remote.clone(),
    };
    let event_handler = T::new(engine);
    let event_handler = Arc::new(event_handler);

    let conn_id = Rc::new(RefCell::new(Counter::new()));
    let connections_inner = connections.clone();
    let event_handler_inner = event_handler.clone();
    // Handle new connection
    let connection_handler = server
        .incoming()
        .map(|(upgrade, addr)| Some((upgrade, addr)))
        .or_else(|InvalidConnection { error, .. }| {
            println!("Connection dropped: {}", error);
            Ok(None)
        })
        .for_each(move |conn| {
            if conn.is_none() {
                return Ok(());
            }
            let (upgrade, addr) = conn.unwrap();
            let event_handler = event_handler_inner.clone();
            let connections_inner = connections_inner.clone();
            println!("Got a connection from: {}", addr);
            let channel = receive_channel_out.clone();
            let conn_id = conn_id.clone();
            let f = upgrade.accept().and_then(move |(framed, _)| {
                let id = conn_id.borrow_mut().next().expect(
                    "maximum amount of ids reached",
                );
                if !event_handler.connect(id) {
                    return Ok(());
                }
                let (sink, stream) = framed.split();
                channel.send((id, stream)).wait().unwrap();
                connections_inner.write().unwrap().insert(id, sink);
                Ok(())
            });
            spawn_future(f, "Handle new connection", &handle);
            Ok(())
        })
        .map_err(|e: ()| e);

    // Handle receiving messages from a client
    let remote_inner = remote.clone();
    let connections_inner = connections.clone();
    let event_handler_inner = event_handler.clone();
    let receive_handler = pool.read()
        .unwrap()
        .spawn_fn(|| {
            receive_channel_in.for_each(move |(id, stream)| {
                let connections = connections_inner.clone();
                let event_handler = event_handler_inner.clone();
                remote_inner.spawn(move |_| {
                    let connections = connections.clone();
                    let event_handler = event_handler.clone();
                    stream
                        .for_each(move |msg| {
                            process_message(id, &msg, &*event_handler, connections.clone());
                            Ok(())
                        })
                        .map_err(|e| println!("Error while receiving messages: {}", e))
                });
                Ok(())
            })
        })
        .map_err(|e| println!("Error while receiving messages: {:?}", e));

    // Handle sending messages to a client
    let connections_inner = connections.clone();
    let send_handler = pool.read()
        .unwrap()
        .spawn_fn(move || {
            let connections = connections_inner.clone();
            send_channel_in.for_each(move |(id, state): (Id, Arc<RwLock<ClientState>>)| {
                let mut sink_guard = connections.write().unwrap();
                // Todo: don't even send invalid ids
                if let Some(mut sink) = sink_guard.get_mut(&id) {
                    let msg = serde_json::to_string(state.read().unwrap().deref()).unwrap();
                    //println!("Sending message: {}", msg);
                    sink.start_send(OwnedMessage::Text(msg)).expect(
                        "Failed to start sending message",
                    );
                    sink.poll_complete().expect("Failed to send message");
                }
                Ok(())
            })
        })
        .map_err(|e| println!("Error while sending messages: {:?}", e));

    // Run maIn loop
    let main_fn = pool.read()
        .unwrap()
        .spawn_fn(move || {
            event_handler.main_loop();
            Ok::<(), ()>(())
        })
        .map_err(|e| println!("Error in main callback function: {:?}", e));
    let handlers = main_fn.select2(connection_handler.select2(
        receive_handler.select(send_handler),
    ));
    core.run(handlers)
        .map_err(|_| println!("Unspecified error while running core loop"))
        .unwrap();
}

pub fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
where
    F: Future<Item = I, Error = E> + 'static,
    E: Debug,
{
    handle.spawn(
        f.map_err(move |e| println!("Error in {}: '{:?}'", desc, e))
            .map(move |_| ()),
    );
}

fn build_server(handle: &Handle) -> WsServer<TlsAcceptor, TcpListener> {
    let mut file = File::open(&read_env_var("CERT_FILE")).unwrap();
    let mut pkcs12 = vec![];
    file.read_to_end(&mut pkcs12).unwrap();
    let pkcs12 = Pkcs12::from_der(&pkcs12, &read_env_var("CERT_PW")).unwrap();

    let address = format!("localhost:{}", read_env_var("CORE_PORT"));
    let acceptor = TlsAcceptor::builder(pkcs12).unwrap().build().unwrap();
    Server::bind_secure(address, acceptor, &handle).expect("Failed to create server")
}

fn process_message<T>(
    id: Id,
    msg: &OwnedMessage,
    event_handler: &T,
    connections: Arc<RwLock<HashMap<Id, SplitSink>>>,
) where
    T: EventHandler,
{
    match *msg {
        OwnedMessage::Text(ref content) => {
            let msg = Msg {
                id,
                content: content.clone(),
            };
            event_handler.message(&msg);
        }
        OwnedMessage::Close(_) => {
            connections
                .write()
                .unwrap()
                .remove(&id)
                .and_then(|_| Some(println!("Client with id {} disconnected", id)))
                .expect("Tried to remove id that was not in list");
            event_handler.disconnect(id);
        }
        _ => {}
    };
}

struct Counter {
    count: Id,
}
impl Counter {
    fn new() -> Self {
        Counter { count: 0 }
    }
}

impl Iterator for Counter {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.count != Id::max_value() {
            self.count += 1;
            Some(self.count)
        } else {
            None
        }
    }
}
<commit_msg>Core: Add some error messages for server creation<commit_after>extern crate websocket;
extern crate futures;
extern crate futures_cpupool;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
extern crate dotenv;
extern crate native_tls;

use self::dotenv::dotenv;

use self::websocket::message::OwnedMessage;
use self::websocket::server::{WsServer, InvalidConnection};
use self::websocket::async::Server;

use self::tokio_core::reactor::{Handle, Remote, Core};
use self::tokio_core::net::TcpListener;

use self::futures::{Future, Sink, Stream};
use self::futures::sync::mpsc;
use self::futures_cpupool::CpuPool;

use self::native_tls::{Pkcs12, TlsAcceptor};

use std::sync::{RwLock, Arc};
use std::rc::Rc;
use std::fmt::Debug;
use std::ops::Deref;
use std::collections::HashMap;
use std::cell::RefCell;
use std::fs::File;
use std::io::Read;

use model::ClientState;
use util::read_env_var;

pub type Id = u32;

type SinkContent = websocket::client::async::Framed<
    websocket::client::async::TlsStream<tokio_core::net::TcpStream>,
    websocket::async::MessageCodec<OwnedMessage>,
>;
type SplitSink = futures::stream::SplitSink<SinkContent>;

#[derive(Clone)]
pub struct Engine {
    pub send_channel: mpsc::UnboundedSender<(Id, Arc<RwLock<ClientState>>)>,
    pub remote: Remote,
}

#[derive(Serialize, Debug)]
pub struct Msg {
    pub id: Id,
    pub content: String,
}

pub trait EventHandler {
    fn new(engine: Engine) -> Self;
    fn main_loop(&self);
    fn message(&self, msg: &Msg);
    fn connect(&self, id: Id) -> bool;
    fn disconnect(&self, id: Id);
}

pub fn execute<T>()
where
    T: EventHandler + Send + Sync + 'static,
{
    dotenv().ok();

    let mut core = Core::new().expect("Failed to create Tokio event loop");
    let handle = core.handle();
    let remote = core.remote();

    let server = build_server(&handle);
    let pool = Arc::new(RwLock::new(CpuPool::new_num_cpus()));
    let connections = Arc::new(RwLock::new(HashMap::new()));
    let (receive_channel_out, receive_channel_in) = mpsc::unbounded();
    let (send_channel_out, send_channel_in) = mpsc::unbounded();

    let engine = Engine {
        send_channel: send_channel_out.clone(),
        remote: remote.clone(),
    };
    let event_handler = T::new(engine);
    let event_handler = Arc::new(event_handler);

    let conn_id = Rc::new(RefCell::new(Counter::new()));
    let connections_inner = connections.clone();
    let event_handler_inner = event_handler.clone();
    // Handle new connection
    let connection_handler = server
        .incoming()
        .map(|(upgrade, addr)| Some((upgrade, addr)))
        .or_else(|InvalidConnection { error, .. }| {
            println!("Connection dropped: {}", error);
            Ok(None)
        })
        .for_each(move |conn| {
            if conn.is_none() {
                return Ok(());
            }
            let (upgrade, addr) = conn.unwrap();
            let event_handler = event_handler_inner.clone();
            let connections_inner = connections_inner.clone();
            println!("Got a connection from: {}", addr);
            let channel = receive_channel_out.clone();
            let conn_id = conn_id.clone();
            let f = upgrade.accept().and_then(move |(framed, _)| {
                let id = conn_id.borrow_mut().next().expect(
                    "maximum amount of ids reached",
                );
                if !event_handler.connect(id) {
                    return Ok(());
                }
                let (sink, stream) = framed.split();
                channel.send((id, stream)).wait().unwrap();
                connections_inner.write().unwrap().insert(id, sink);
                Ok(())
            });
            spawn_future(f, "Handle new connection", &handle);
            Ok(())
        })
        .map_err(|e: ()| e);

    // Handle receiving messages from a client
    let remote_inner = remote.clone();
    let connections_inner = connections.clone();
    let event_handler_inner = event_handler.clone();
    let receive_handler = pool.read()
        .unwrap()
        .spawn_fn(|| {
            receive_channel_in.for_each(move |(id, stream)| {
                let connections = connections_inner.clone();
                let event_handler = event_handler_inner.clone();
                remote_inner.spawn(move |_| {
                    let connections = connections.clone();
                    let event_handler = event_handler.clone();
                    stream
                        .for_each(move |msg| {
                            process_message(id, &msg, &*event_handler, connections.clone());
                            Ok(())
                        })
                        .map_err(|e| println!("Error while receiving messages: {}", e))
                });
                Ok(())
            })
        })
        .map_err(|e| println!("Error while receiving messages: {:?}", e));

    // Handle sending messages to a client
    let connections_inner = connections.clone();
    let send_handler = pool.read()
        .unwrap()
        .spawn_fn(move || {
            let connections = connections_inner.clone();
            send_channel_in.for_each(move |(id, state): (Id, Arc<RwLock<ClientState>>)| {
                let mut sink_guard = connections.write().unwrap();
                // Todo: don't even send invalid ids
                if let Some(mut sink) = sink_guard.get_mut(&id) {
                    let msg = serde_json::to_string(state.read().unwrap().deref()).unwrap();
                    //println!("Sending message: {}", msg);
                    sink.start_send(OwnedMessage::Text(msg)).expect(
                        "Failed to start sending message",
                    );
                    sink.poll_complete().expect("Failed to send message");
                }
                Ok(())
            })
        })
        .map_err(|e| println!("Error while sending messages: {:?}", e));

    // Run maIn loop
    let main_fn = pool.read()
        .unwrap()
        .spawn_fn(move || {
            event_handler.main_loop();
            Ok::<(), ()>(())
        })
        .map_err(|e| println!("Error in main callback function: {:?}", e));
    let handlers = main_fn.select2(connection_handler.select2(
        receive_handler.select(send_handler),
    ));
    core.run(handlers)
        .map_err(|_| println!("Unspecified error while running core loop"))
        .unwrap();
}

pub fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
where
    F: Future<Item = I, Error = E> + 'static,
    E: Debug,
{
    handle.spawn(
        f.map_err(move |e| println!("Error in {}: '{:?}'", desc, e))
            .map(move |_| ()),
    );
}

fn build_server(handle: &Handle) -> WsServer<TlsAcceptor, TcpListener> {
    let mut file = File::open(&read_env_var("CERT_FILE")).expect("Failed to open certificate file");
    let mut pkcs12 = vec![];
    file.read_to_end(&mut pkcs12).unwrap();
    let pkcs12 = Pkcs12::from_der(&pkcs12, &read_env_var("CERT_PW")).expect(
        "Failed to open certificate file with given password",
    );

    let address = format!("localhost:{}", read_env_var("CORE_PORT"));
    let acceptor = TlsAcceptor::builder(pkcs12).unwrap().build().expect(
        "Failed to build TLS acceptor",
    );
    Server::bind_secure(address, acceptor, &handle).expect("Failed to create server")
}

fn process_message<T>(
    id: Id,
    msg: &OwnedMessage,
    event_handler: &T,
    connections: Arc<RwLock<HashMap<Id, SplitSink>>>,
) where
    T: EventHandler,
{
    match *msg {
        OwnedMessage::Text(ref content) => {
            let msg = Msg {
                id,
                content: content.clone(),
            };
            event_handler.message(&msg);
        }
        OwnedMessage::Close(_) => {
            connections
                .write()
                .unwrap()
                .remove(&id)
                .and_then(|_| Some(println!("Client with id {} disconnected", id)))
                .expect("Tried to remove id that was not in list");
            event_handler.disconnect(id);
        }
        _ => {}
    };
}

struct Counter {
    count: Id,
}
impl Counter {
    fn new() -> Self {
        Counter { count: 0 }
    }
}

impl Iterator for Counter {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.count != Id::max_value() {
            self.count += 1;
            Some(self.count)
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::iter::FromIterator;
use std::ops::Deref;

use syntax::ast;
use syntax::ptr::P;
use syntax::ast_util::empty_generics;
use syntax::codemap::{respan, mk_sp};
use syntax::ext::base::ExtCtxt;
use syntax::ext::build::AstBuilder;
use syntax::ext::quote::rt::ToTokens;

use super::Builder;
use super::utils;
use super::super::node;

/// A visitor to build the struct for each register
pub struct BuildRegStructs<'a> {
  builder: &'a mut Builder,
  cx: &'a ExtCtxt<'a>,
}

impl<'a> node::RegVisitor for BuildRegStructs<'a> {
  fn visit_prim_reg(&mut self, path: &Vec<String>, reg: &node::Reg,
                    fields: &Vec<node::Field>) {
    let width = match reg.ty {
      node::RegType::RegPrim(ref width, _) => width.node,
      _ => panic!("visit_prim_reg called with non-primitive register"),
    };
    for field in fields.iter() {
      for item in build_field_type(self.cx, path, reg, field).into_iter() {
        self.builder.push_item(item);
      }
    }

    for item in build_reg_struct(self.cx, path, reg, &width).into_iter() {
      self.builder.push_item(item);
    }
  }
}

impl<'a> BuildRegStructs<'a> {
  pub fn new(builder: &'a mut Builder, cx: &'a ExtCtxt<'a>)
             -> BuildRegStructs<'a> {
    BuildRegStructs {builder: builder, cx: cx}
  }
}

/// Build a field type if necessary (e.g. in the case of an `EnumField`)
fn build_field_type(cx: &ExtCtxt, path: &Vec<String>,
                    reg: &node::Reg, field: &node::Field)
                    -> Vec<P<ast::Item>> {
  match field.ty.node {
    node::FieldType::EnumField { ref variants, .. } => {
      // FIXME(bgamari): We construct a path, then only take the last
      // segment, this could be more efficient
      let name: ast::Ident =
        utils::field_type_path(cx, path, reg, field)
        .segments.last().unwrap().identifier;
      let enum_def: ast::EnumDef = ast::EnumDef {
        variants: FromIterator::from_iter(
          variants.iter().map(|v| P(build_enum_variant(cx, v)))),
      };
      let mut attrs: Vec<ast::Attribute> = vec!(
        utils::list_attribute(cx, "derive",
                              vec!("PartialEq"),
                              field.name.span),
        utils::list_attribute(cx, "allow",
                              vec!("dead_code",
                                   "non_camel_case_types",
                                   "missing_docs"),
                              field.name.span));

      // setting the enum's repr to the width of the register
      // (unless there's only 1 variant; in which case we omit
      // the repr attribute due to E0083
      if variants.len() > 1 {
        let enum_repr = utils::reg_primitive_type_name(reg)
          .expect("Unexpected non-primitive reg");

        attrs.push(utils::list_attribute(cx, "repr",
                                         vec!(enum_repr),
                                         field.name.span));
      }

      let ty_item: P<ast::Item> = P(ast::Item {
        ident: name,
        id: ast::DUMMY_NODE_ID,
        node: ast::ItemEnum(enum_def, empty_generics()),
        vis: ast::Public,
        attrs: attrs,
        span: field.ty.span,
      });
      vec!(ty_item)
    },
    _ => Vec::new()
  }
}

/// Produce a register struct if necessary (for primitive typed registers).
/// In this case `None` indicates no struct is necessary, not failure.
/// For instance,
///
///     pub struct REG {_value: u32}
fn build_reg_struct(cx: &ExtCtxt, path: &Vec<String>,
    reg: &node::Reg, _width: &node::RegWidth) -> Vec<P<ast::Item>> {
  let packed_ty =
    utils::reg_primitive_type(cx, reg)
    .expect("Unexpected non-primitive reg");

  let reg_doc = match reg.docstring {
    Some(d) => d.node.name.to_string(),
    None => "no documentation".to_string(),
  };
  let docstring = format!("Register `{}`: {}",
                          reg.name.node,
                          reg_doc);
  let doc_attr = utils::doc_attribute(cx, utils::intern_string(cx, docstring));

  let ty_name = utils::path_ident(cx, path);
  let item = quote_item!(cx,
    $doc_attr
    #[derive(Clone)]
    #[allow(non_camel_case_types)]
    #[repr(C)]
    pub struct $ty_name {
      value: VolatileCell<$packed_ty>,
    }
  );
  let mut item: ast::Item = item.unwrap().deref().clone();
  item.span = reg.name.span;
  let copy_impl = quote_item!(cx, impl ::core::marker::Copy for $ty_name {}).unwrap();
  vec!(P(item), copy_impl)
}

/// Build a variant of an `EnumField`
fn build_enum_variant(cx: &ExtCtxt, variant: &node::Variant)
                      -> ast::Variant {
  let doc = match variant.docstring {
    Some(d) => d.node.name.to_string(),
    None => "no documentation".to_string(),
  };
  let docstring = format!("`0x{:x}`. {}",
                          variant.value.node,
                          doc);
  let doc_attr = utils::doc_attribute(cx, utils::intern_string(cx, docstring));
  respan(
    mk_sp(variant.name.span.lo, variant.value.span.hi),
    ast::Variant_ {
      name: cx.ident_of(variant.name.node.as_str()),
      attrs: vec!(doc_attr),
      kind: ast::TupleVariantKind(Vec::new()),
      id: ast::DUMMY_NODE_ID,
      disr_expr: Some(utils::expr_int(cx, respan(variant.value.span,
                                                 variant.value.node as i64))),
      vis: ast::Inherited,
    }
  )
}
<commit_msg>Fixed ast::Variant_ layout<commit_after>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::iter::FromIterator;
use std::ops::Deref;

use syntax::ast;
use syntax::ptr::P;
use syntax::ast_util::empty_generics;
use syntax::codemap::{respan, mk_sp};
use syntax::ext::base::ExtCtxt;
use syntax::ext::build::AstBuilder;
use syntax::ext::quote::rt::ToTokens;

use super::Builder;
use super::utils;
use super::super::node;

/// A visitor to build the struct for each register
pub struct BuildRegStructs<'a> {
  builder: &'a mut Builder,
  cx: &'a ExtCtxt<'a>,
}

impl<'a> node::RegVisitor for BuildRegStructs<'a> {
  fn visit_prim_reg(&mut self, path: &Vec<String>, reg: &node::Reg,
                    fields: &Vec<node::Field>) {
    let width = match reg.ty {
      node::RegType::RegPrim(ref width, _) => width.node,
      _ => panic!("visit_prim_reg called with non-primitive register"),
    };
    for field in fields.iter() {
      for item in build_field_type(self.cx, path, reg, field).into_iter() {
        self.builder.push_item(item);
      }
    }

    for item in build_reg_struct(self.cx, path, reg, &width).into_iter() {
      self.builder.push_item(item);
    }
  }
}

impl<'a> BuildRegStructs<'a> {
  pub fn new(builder: &'a mut Builder, cx: &'a ExtCtxt<'a>)
             -> BuildRegStructs<'a> {
    BuildRegStructs {builder: builder, cx: cx}
  }
}

/// Build a field type if necessary (e.g. in the case of an `EnumField`)
fn build_field_type(cx: &ExtCtxt, path: &Vec<String>,
                    reg: &node::Reg, field: &node::Field)
                    -> Vec<P<ast::Item>> {
  match field.ty.node {
    node::FieldType::EnumField { ref variants, .. } => {
      // FIXME(bgamari): We construct a path, then only take the last
      // segment, this could be more efficient
      let name: ast::Ident =
        utils::field_type_path(cx, path, reg, field)
        .segments.last().unwrap().identifier;
      let enum_def: ast::EnumDef = ast::EnumDef {
        variants: FromIterator::from_iter(
          variants.iter().map(|v| P(build_enum_variant(cx, v)))),
      };
      let mut attrs: Vec<ast::Attribute> = vec!(
        utils::list_attribute(cx, "derive",
                              vec!("PartialEq"),
                              field.name.span),
        utils::list_attribute(cx, "allow",
                              vec!("dead_code",
                                   "non_camel_case_types",
                                   "missing_docs"),
                              field.name.span));

      // setting the enum's repr to the width of the register
      // (unless there's only 1 variant; in which case we omit
      // the repr attribute due to E0083
      if variants.len() > 1 {
        let enum_repr = utils::reg_primitive_type_name(reg)
          .expect("Unexpected non-primitive reg");

        attrs.push(utils::list_attribute(cx, "repr",
                                         vec!(enum_repr),
                                         field.name.span));
      }

      let ty_item: P<ast::Item> = P(ast::Item {
        ident: name,
        id: ast::DUMMY_NODE_ID,
        node: ast::ItemEnum(enum_def, empty_generics()),
        vis: ast::Public,
        attrs: attrs,
        span: field.ty.span,
      });
      vec!(ty_item)
    },
    _ => Vec::new()
  }
}

/// Produce a register struct if necessary (for primitive typed registers).
/// In this case `None` indicates no struct is necessary, not failure.
/// For instance,
///
///     pub struct REG {_value: u32}
fn build_reg_struct(cx: &ExtCtxt, path: &Vec<String>,
    reg: &node::Reg, _width: &node::RegWidth) -> Vec<P<ast::Item>> {
  let packed_ty =
    utils::reg_primitive_type(cx, reg)
    .expect("Unexpected non-primitive reg");

  let reg_doc = match reg.docstring {
    Some(d) => d.node.name.to_string(),
    None => "no documentation".to_string(),
  };
  let docstring = format!("Register `{}`: {}",
                          reg.name.node,
                          reg_doc);
  let doc_attr = utils::doc_attribute(cx, utils::intern_string(cx, docstring));

  let ty_name = utils::path_ident(cx, path);
  let item = quote_item!(cx,
    $doc_attr
    #[derive(Clone)]
    #[allow(non_camel_case_types)]
    #[repr(C)]
    pub struct $ty_name {
      value: VolatileCell<$packed_ty>,
    }
  );
  let mut item: ast::Item = item.unwrap().deref().clone();
  item.span = reg.name.span;
  let copy_impl = quote_item!(cx, impl ::core::marker::Copy for $ty_name {}).unwrap();
  vec!(P(item), copy_impl)
}

/// Build a variant of an `EnumField`
fn build_enum_variant(cx: &ExtCtxt, variant: &node::Variant)
                      -> ast::Variant {
  let doc = match variant.docstring {
    Some(d) => d.node.name.to_string(),
    None => "no documentation".to_string(),
  };
  let docstring = format!("`0x{:x}`. {}",
                          variant.value.node,
                          doc);
  let doc_attr = utils::doc_attribute(cx, utils::intern_string(cx, docstring));
  respan(
    mk_sp(variant.name.span.lo, variant.value.span.hi),
    ast::Variant_ {
      name: cx.ident_of(variant.name.node.as_str()),
      attrs: vec!(doc_attr),
      data: P(ast::VariantData::Tuple(Vec::new(), ast::DUMMY_NODE_ID)),
      disr_expr: Some(utils::expr_int(cx, respan(variant.value.span,
                                                 variant.value.node as i64))),
    }
  )
}
<|endoftext|>