<commit_before>use std::sync::Arc;
use std::io;
use std::cell::{Cell as StdCell};
use bifrost::conshash::{ConsistentHashing, CHError};
use bifrost::raft::client::{RaftClient, ClientError};
use bifrost::raft;
use bifrost::rpc::{RPCError, DEFAULT_CLIENT_POOL, Server as RPCServer};
use bifrost::raft::state_machine::master::ExecError;
use bifrost::raft::state_machine::callback::server::{NotifyError};

use server::{transactions as txn_server, cell_rpc as plain_server};
use ram::types::Id;
use ram::cell::{Cell, Header, ReadError, WriteError};
use ram::schema::{sm as schema_sm};
use ram::schema::sm::client::{SMClient as SchemaClient};
use ram::schema::Schema;

use futures::Future;
use futures::prelude::*;

use self::transaction::*;

static TRANSACTION_MAX_RETRY: u32 = 500;

pub mod transaction;

#[derive(Debug)]
pub enum NebClientError {
    RaftClientError(ClientError),
    ConsistentHashtableError(CHError)
}

pub struct AsyncClientInner {
    pub conshash: Arc<ConsistentHashing>,
    pub raft_client: Arc<RaftClient>,
    pub schema_client: SchemaClient
}

impl AsyncClientInner {
    pub fn new<'a>(subscription_server: &Arc<RPCServer>, meta_servers: &Vec<String>, group: &'a str)
        -> Result<AsyncClientInner, NebClientError>
    {
        match RaftClient::new(meta_servers, raft::DEFAULT_SERVICE_ID) {
            Ok(raft_client) => {
                RaftClient::prepare_subscription(subscription_server);
                assert!(RaftClient::can_callback());
                match ConsistentHashing::new_client(group, &raft_client) {
                    Ok(chash) => Ok(AsyncClientInner {
                        conshash: chash,
                        raft_client: raft_client.clone(),
                        schema_client: SchemaClient::new(schema_sm::generate_sm_id(group), &raft_client)
                    }),
                    Err(err) => Err(NebClientError::ConsistentHashtableError(err))
                }
            },
            Err(err) => Err(NebClientError::RaftClientError(err))
        }
    }
    pub fn locate_server_address(&self, id: &Id) -> Result<String, RPCError> {
        match self.conshash.get_server(id.higher) {
            Some(n) => Ok(n),
            None => Err(RPCError::IOError(io::Error::new(io::ErrorKind::NotFound, "cannot locate")))
        }
    }
    pub fn locate_plain_server(&self, id: Id) -> Result<Arc<plain_server::AsyncServiceClient>, RPCError> {
        let address = self.locate_server_address(id)?;
        let client = match DEFAULT_CLIENT_POOL.get(&address) {
            Ok(c) => c,
            Err(e) => return Err(RPCError::IOError(e))
        };
        Ok(plain_server::AsyncServiceClient::new_async(plain_server::DEFAULT_SERVICE_ID, &client))
    }
    #[async]
    pub fn read_cell(this: Arc<Self>, id: Id) -> Result<Result<Cell, ReadError>, RPCError> {
        let client = this.locate_plain_server(id)?;
        await!(client.read_cell(&id))
    }
    #[async]
    pub fn write_cell(this: Arc<Self>, cell: Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = this.locate_plain_server(cell.id())?;
        await!(client.write_cell(&cell))
    }
    #[async]
    pub fn update_cell(this: Arc<Self>, cell: Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = this.locate_plain_server(cell.id())?;
        await!(client.update_cell(&cell))
    }
    #[async]
    pub fn remove_cell(this: Arc<Self>, id: Id) -> Result<Result<(), WriteError>, RPCError> {
        let client = this.locate_plain_server(id)?;
        await!(client.remove_cell(&id))
    }
    #[async]
    pub fn transaction<TFN, TR>(this: Arc<Self>, func: TFN) -> Result<TR, TxnError>
        where TFN: Fn(&Transaction) -> Result<TR, TxnError>
    {
        let server_name = match this.conshash.rand_server() {
            Some(name) => name,
            None => return Err(TxnError::CannotFindAServer)
        };
        let txn_client = match txn_server::new_client(&server_name) {
            Ok(client) => client,
            Err(e) => return Err(TxnError::IoError(e))
        };
        let mut txn_id: txn_server::TxnId;
        let mut retried = 0;
        while retried < TRANSACTION_MAX_RETRY {
            txn_id = match txn_client.begin() {
                Ok(Ok(id)) => id,
                _ => return Err(TxnError::CannotBegin)
            };
            let txn = Transaction{
                tid: txn_id,
                state: StdCell::new(txn_server::TxnState::Started),
                client: txn_client.clone(),
            };
            let exec_result = func(&txn);
            let mut exec_value = None;
            let mut txn_result = Ok(());
            match exec_result {
                Ok(val) => {
                    if txn.state.get() == txn_server::TxnState::Started {
                        txn_result = txn.prepare();
                        debug!("PREPARE STATE: {:?}", txn_result);
                    }
                    if txn_result.is_ok() && txn.state.get() == txn_server::TxnState::Prepared {
                        txn_result = txn.commit();
                        debug!("COMMIT STATE: {:?}", txn_result);
                    }
                    exec_value = Some(val);
                },
                Err(e) => {
                    txn_result = Err(e)
                }
            }
            debug!("TXN CONCLUSION: {:?}", txn_result);
            match txn_result {
                Ok(()) => {
                    return Ok(exec_value.unwrap());
                },
                Err(TxnError::NotRealizable) => {
                    let abort_result = txn.abort();  // continue the loop to retry
                    debug!("TXN NOT REALIZABLE, ABORT: {:?}", abort_result);
                },
                Err(e) => {
                    // abort will always be an error to achieve early break
                    let abort_result = txn.abort();
                    debug!("TXN ERROR, ABORT: {:?}", abort_result);
                    return Err(e);
                }
            }
            retried += 1;
            debug!("Client retry transaction, {:?} times", retried);
        }
        Err(TxnError::TooManyRetry)
    }
    #[async]
    pub fn new_schema_with_id(this: Arc<Self>, schema: Schema) -> Result<Result<(), NotifyError>, ExecError> {
        await!(this.schema_client.new_schema(&schema))
    }
    #[async]
    pub fn new_schema(this: Arc<Self>, schema: Schema) -> Result<Result<u32, NotifyError>, ExecError> {
        let schema_id = self.schema_client.next_id()?.unwrap();
        schema.id = schema_id;
        await!(self.new_schema_with_id(schema)).map(|r| r.map(|_| schema_id))
    }
    pub fn del_schema(&self, schema_id: &String) -> Result<Result<(), NotifyError>, ExecError> {
        self.schema_client.del_schema(schema_id)
    }
    pub fn get_all_schema(&self) -> Result<Vec<Schema>, ExecError> {
        Ok(self.schema_client.get_all()?.unwrap())
    }
}<commit_msg>export async client functions<commit_after>use std::sync::Arc;
use std::io;
use std::cell::{Cell as StdCell};
use bifrost::conshash::{ConsistentHashing, CHError};
use bifrost::raft::client::{RaftClient, ClientError};
use bifrost::raft;
use bifrost::rpc::{RPCError, DEFAULT_CLIENT_POOL, Server as RPCServer};
use bifrost::raft::state_machine::master::ExecError;
use bifrost::raft::state_machine::callback::server::{NotifyError};

use server::{transactions as txn_server, cell_rpc as plain_server};
use ram::types::Id;
use ram::cell::{Cell, Header, ReadError, WriteError};
use ram::schema::{sm as schema_sm};
use ram::schema::sm::client::{SMClient as SchemaClient};
use ram::schema::Schema;

use futures::Future;
use futures::prelude::*;

use self::transaction::*;

static TRANSACTION_MAX_RETRY: u32 = 500;

pub mod transaction;

#[derive(Debug)]
pub enum NebClientError {
    RaftClientError(ClientError),
    ConsistentHashtableError(CHError)
}

pub struct AsyncClientInner {
    pub conshash: Arc<ConsistentHashing>,
    pub raft_client: Arc<RaftClient>,
    pub schema_client: SchemaClient
}

impl AsyncClientInner {
    pub fn new<'a>(subscription_server: &Arc<RPCServer>, meta_servers: &Vec<String>, group: &'a str)
        -> Result<AsyncClientInner, NebClientError>
    {
        match RaftClient::new(meta_servers, raft::DEFAULT_SERVICE_ID) {
            Ok(raft_client) => {
                RaftClient::prepare_subscription(subscription_server);
                assert!(RaftClient::can_callback());
                match ConsistentHashing::new_client(group, &raft_client) {
                    Ok(chash) => Ok(AsyncClientInner {
                        conshash: chash,
                        raft_client: raft_client.clone(),
                        schema_client: SchemaClient::new(schema_sm::generate_sm_id(group), &raft_client)
                    }),
                    Err(err) => Err(NebClientError::ConsistentHashtableError(err))
                }
            },
            Err(err) => Err(NebClientError::RaftClientError(err))
        }
    }
    pub fn locate_server_address(&self, id: &Id) -> Result<String, RPCError> {
        match self.conshash.get_server(id.higher) {
            Some(n) => Ok(n),
            None => Err(RPCError::IOError(io::Error::new(io::ErrorKind::NotFound, "cannot locate")))
        }
    }
    pub fn locate_plain_server(&self, id: Id) -> Result<Arc<plain_server::AsyncServiceClient>, RPCError> {
        let address = self.locate_server_address(&id)?;
        let client = match DEFAULT_CLIENT_POOL.get(&address) {
            Ok(c) => c,
            Err(e) => return Err(RPCError::IOError(e))
        };
        Ok(plain_server::AsyncServiceClient::new(plain_server::DEFAULT_SERVICE_ID, &client))
    }
    #[async]
    pub fn read_cell(this: Arc<Self>, id: Id) -> Result<Result<Cell, ReadError>, RPCError> {
        let client = this.locate_plain_server(id)?;
        await!(client.read_cell(&id))
    }
    #[async]
    pub fn write_cell(this: Arc<Self>, cell: Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = this.locate_plain_server(cell.id())?;
        await!(client.write_cell(&cell))
    }
    #[async]
    pub fn update_cell(this: Arc<Self>, cell: Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = this.locate_plain_server(cell.id())?;
        await!(client.update_cell(&cell))
    }
    #[async]
    pub fn remove_cell(this: Arc<Self>, id: Id) -> Result<Result<(), WriteError>, RPCError> {
        let client = this.locate_plain_server(id)?;
        await!(client.remove_cell(&id))
    }
    #[async]
    pub fn transaction<TFN, TR>(this: Arc<Self>, func: TFN) -> Result<TR, TxnError>
        where TFN: Fn(&Transaction) -> Result<TR, TxnError>
    {
        let server_name = match this.conshash.rand_server() {
            Some(name) => name,
            None => return Err(TxnError::CannotFindAServer)
        };
        let txn_client = match txn_server::new_client(&server_name) {
            Ok(client) => client,
            Err(e) => return Err(TxnError::IoError(e))
        };
        let mut txn_id: txn_server::TxnId;
        let mut retried = 0;
        while retried < TRANSACTION_MAX_RETRY {
            txn_id = match txn_client.begin() {
                Ok(Ok(id)) => id,
                _ => return Err(TxnError::CannotBegin)
            };
            let txn = Transaction{
                tid: txn_id,
                state: StdCell::new(txn_server::TxnState::Started),
                client: txn_client.clone(),
            };
            let exec_result = func(&txn);
            let mut exec_value = None;
            let mut txn_result = Ok(());
            match exec_result {
                Ok(val) => {
                    if txn.state.get() == txn_server::TxnState::Started {
                        txn_result = txn.prepare();
                        debug!("PREPARE STATE: {:?}", txn_result);
                    }
                    if txn_result.is_ok() && txn.state.get() == txn_server::TxnState::Prepared {
                        txn_result = txn.commit();
                        debug!("COMMIT STATE: {:?}", txn_result);
                    }
                    exec_value = Some(val);
                },
                Err(e) => {
                    txn_result = Err(e)
                }
            }
            debug!("TXN CONCLUSION: {:?}", txn_result);
            match txn_result {
                Ok(()) => {
                    return Ok(exec_value.unwrap());
                },
                Err(TxnError::NotRealizable) => {
                    let abort_result = txn.abort();  // continue the loop to retry
                    debug!("TXN NOT REALIZABLE, ABORT: {:?}", abort_result);
                },
                Err(e) => {
                    // abort will always be an error to achieve early break
                    let abort_result = txn.abort();
                    debug!("TXN ERROR, ABORT: {:?}", abort_result);
                    return Err(e);
                }
            }
            retried += 1;
            debug!("Client retry transaction, {:?} times", retried);
        }
        Err(TxnError::TooManyRetry)
    }
    #[async]
    pub fn new_schema_with_id(this: Arc<Self>, schema: Schema) -> Result<Result<(), NotifyError>, ExecError> {
        this.schema_client.new_schema(&schema)
    }
    #[async]
    pub fn new_schema(this: Arc<Self>, schema: Schema) -> Result<Result<u32, NotifyError>, ExecError> {
        let schema_id = this.schema_client.next_id()?.unwrap();
        schema.id = schema_id;
        await!(Self::new_schema_with_id(this, schema)).map(|r| r.map(|_| schema_id))
    }
    #[async]
    pub fn del_schema(this: Arc<Self>, name: String) -> Result<Result<(), NotifyError>, ExecError> {
        this.schema_client.del_schema(&name)
    }
    #[async]
    pub fn get_all_schema(this: Arc<Self>) -> Result<Vec<Schema>, ExecError> {
        Ok(this.schema_client.get_all()?.unwrap())
    }
}

pub struct AsyncClient {
    inner: Arc<AsyncClientInner>
}

impl AsyncClient {
    pub fn new<'a>(subscription_server: &Arc<RPCServer>, meta_servers: &Vec<String>, group: &'a str)
                   -> Result<AsyncClient, NebClientError>
    {
        AsyncClientInner::new(subscription_server, meta_servers, group)
            .map(|inner| {
                AsyncClient {
                    inner: Arc::new(inner)
                }
            })
    }

    pub fn locate_server_address(&self, id: &Id) -> Result<String, RPCError> {
        self.inner.locate_server_address(id)
    }

    pub fn locate_plain_server(&self, id: Id) -> Result<Arc<plain_server::AsyncServiceClient>, RPCError> {
        self.inner.locate_plain_server(id)
    }

    pub fn read_cell(&self, id: Id)
        -> impl Future<Item = Result<Cell, ReadError>, Error = RPCError>
    {
        AsyncClientInner::read_cell(self.inner, id)
    }

    pub fn write_cell(&self, cell: Cell)
        -> impl Future<Item = Result<Header, WriteError>, Error = RPCError>
    {
        AsyncClientInner::write_cell(self.inner, cell)
    }

    pub fn update_cell(&self, cell: Cell)
        -> impl Future<Item = Result<Header, WriteError>, Error = RPCError>
    {
        AsyncClientInner::update_cell(self.inner, cell)
    }

    pub fn remove_cell(&self, id: Id)
        -> impl Future<Item = Result<(), WriteError>, Error = RPCError>
    {
        AsyncClientInner::remove_cell(self.inner, id)
    }

    pub fn transaction<TFN, TR>(&self, func: TFN)
        -> impl Future<Item = TR, Error = TxnError>
        where TFN: Fn(&Transaction) -> Result<TR, TxnError>
    {
        AsyncClientInner::transaction(self.inner, func)
    }
    pub fn new_schema_with_id(&self, schema: Schema)
        -> impl Future<Item = Result<(), NotifyError>, Error = ExecError>
    {
        AsyncClientInner::new_schema_with_id(self.inner, schema)
    }
    pub fn new_schema(&self, schema: Schema)
        -> impl Future<Item = Result<u32, NotifyError>, Error = ExecError>
    {
        AsyncClientInner::new_schema(self.inner, schema)
    }
    pub fn del_schema(&self, name: String)
        -> impl Future<Item = Result<(), NotifyError>, Error = ExecError>
    {
        AsyncClientInner::del_schema(self.inner, name)
    }
    pub fn get_all_schema(&self)
        -> impl Future<Item = Vec<Schema>, Error = ExecError>
    {
        AsyncClientInner::get_all_schema(self.inner)
    }
}<|endoftext|>


<commit_before>//! Postfix completions, like `Ok(10).ifl<|>` => `if let Ok() = Ok(10) { <|> }`.

mod format_like;

use ide_db::ty_filter::TryEnum;
use syntax::{
    ast::{self, AstNode, AstToken},
    SyntaxKind::BLOCK_EXPR,
    TextRange, TextSize,
};
use text_edit::TextEdit;

use self::format_like::add_format_like_completions;
use crate::{
    config::SnippetCap,
    context::CompletionContext,
    item::{Builder, CompletionKind},
    CompletionItem, CompletionItemKind, Completions,
};

pub(crate) fn complete_postfix(acc: &mut Completions, ctx: &CompletionContext) {
    if !ctx.config.enable_postfix_completions {
        return;
    }

    let dot_receiver = match &ctx.dot_receiver {
        Some(it) => it,
        None => return,
    };

    let receiver_text =
        get_receiver_text(dot_receiver, ctx.dot_receiver_is_ambiguous_float_literal);

    let receiver_ty = match ctx.sema.type_of_expr(&dot_receiver) {
        Some(it) => it,
        None => return,
    };

    let cap = match ctx.config.snippet_cap {
        Some(it) => it,
        None => return,
    };
    let try_enum = TryEnum::from_ty(&ctx.sema, &receiver_ty);
    if let Some(try_enum) = &try_enum {
        match try_enum {
            TryEnum::Result => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "ifl",
                    "if let Ok {}",
                    &format!("if let Ok($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);

                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "while",
                    "while let Ok {}",
                    &format!("while let Ok($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);
            }
            TryEnum::Option => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "ifl",
                    "if let Some {}",
                    &format!("if let Some($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);

                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "while",
                    "while let Some {}",
                    &format!("while let Some($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);
            }
        }
    } else if receiver_ty.is_bool() || receiver_ty.is_unknown() {
        postfix_snippet(
            ctx,
            cap,
            &dot_receiver,
            "if",
            "if expr {}",
            &format!("if {} {{\n    $0\n}}", receiver_text),
        )
        .add_to(acc);
        postfix_snippet(
            ctx,
            cap,
            &dot_receiver,
            "while",
            "while expr {}",
            &format!("while {} {{\n    $0\n}}", receiver_text),
        )
        .add_to(acc);
        postfix_snippet(ctx, cap, &dot_receiver, "not", "!expr", &format!("!{}", receiver_text))
            .add_to(acc);
    }

    postfix_snippet(ctx, cap, &dot_receiver, "ref", "&expr", &format!("&{}", receiver_text))
        .add_to(acc);
    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "refm",
        "&mut expr",
        &format!("&mut {}", receiver_text),
    )
    .add_to(acc);

    // The rest of the postfix completions create an expression that moves an argument,
    // so it's better to consider references now to avoid breaking the compilation
    let dot_receiver = include_references(dot_receiver);
    let receiver_text =
        get_receiver_text(&dot_receiver, ctx.dot_receiver_is_ambiguous_float_literal);

    match try_enum {
        Some(try_enum) => match try_enum {
            TryEnum::Result => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "match",
                    "match expr {}",
                    &format!("match {} {{\n    Ok(${{1:_}}) => {{$2}},\n    Err(${{3:_}}) => {{$0}},\n}}", receiver_text),
                )
                .add_to(acc);
            }
            TryEnum::Option => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "match",
                    "match expr {}",
                    &format!(
                        "match {} {{\n    Some(${{1:_}}) => {{$2}},\n    None => {{$0}},\n}}",
                        receiver_text
                    ),
                )
                .add_to(acc);
            }
        },
        None => {
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "match",
                "match expr {}",
                &format!("match {} {{\n    ${{1:_}} => {{$0}},\n}}", receiver_text),
            )
            .add_to(acc);
        }
    }

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "box",
        "Box::new(expr)",
        &format!("Box::new({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(ctx, cap, &dot_receiver, "ok", "Ok(expr)", &format!("Ok({})", receiver_text))
        .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "some",
        "Some(expr)",
        &format!("Some({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "dbg",
        "dbg!(expr)",
        &format!("dbg!({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "dbgr",
        "dbg!(&expr)",
        &format!("dbg!(&{})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "call",
        "function(expr)",
        &format!("${{1}}({})", receiver_text),
    )
    .add_to(acc);

    let parent_node = dot_receiver.syntax().parent().and_then(|p| p.parent());
    if let Some(parent) = parent_node {
        if parent.kind() == BLOCK_EXPR {
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "let",
                "let",
                &format!("let $0 = {};", receiver_text),
            )
            .add_to(acc);
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "letm",
                "let mut",
                &format!("let mut $0 = {};", receiver_text),
            )
            .add_to(acc);
        }
    }

    if let ast::Expr::Literal(literal) = dot_receiver.clone() {
        if let Some(literal_text) = ast::String::cast(literal.token()) {
            add_format_like_completions(acc, ctx, &dot_receiver, cap, &literal_text);
        }
    }
}

fn get_receiver_text(receiver: &ast::Expr, receiver_is_ambiguous_float_literal: bool) -> String {
    if receiver_is_ambiguous_float_literal {
        let text = receiver.syntax().text();
        let without_dot = ..text.len() - TextSize::of('.');
        text.slice(without_dot).to_string()
    } else {
        receiver.to_string()
    }
}

fn include_references(initial_element: &ast::Expr) -> ast::Expr {
    let mut resulting_element = initial_element.clone();
    while let Some(parent_ref_element) =
        resulting_element.syntax().parent().and_then(ast::RefExpr::cast)
    {
        resulting_element = ast::Expr::from(parent_ref_element);
    }
    resulting_element
}

fn postfix_snippet(
    ctx: &CompletionContext,
    cap: SnippetCap,
    receiver: &ast::Expr,
    label: &str,
    detail: &str,
    snippet: &str,
) -> Builder {
    let edit = {
        let receiver_syntax = receiver.syntax();
        let receiver_range = ctx.sema.original_range(receiver_syntax).range;
        let delete_range = TextRange::new(receiver_range.start(), ctx.source_range().end());
        TextEdit::replace(delete_range, snippet.to_string())
    };
    CompletionItem::new(CompletionKind::Postfix, ctx.source_range(), label)
        .detail(detail)
        .kind(CompletionItemKind::Snippet)
        .snippet_edit(cap, edit)
}

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

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

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

    #[test]
    fn postfix_completion_works_for_trivial_path_expression() {
        check(
            r#"
fn main() {
    let bar = true;
    bar.<|>
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn if    if expr {}
                sn let   let
                sn letm  let mut
                sn match match expr {}
                sn not   !expr
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
                sn while while expr {}
            "#]],
        );
    }

    #[test]
    fn postfix_completion_works_for_function_calln() {
        check(
            r#"
fn foo(elt: bool) -> bool {
    !elt
}

fn main() {
    let bar = true;
    foo(bar.<|>)
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn if    if expr {}
                sn match match expr {}
                sn not   !expr
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
                sn while while expr {}
            "#]],
        );
    }

    #[test]
    fn postfix_type_filtering() {
        check(
            r#"
fn main() {
    let bar: u8 = 12;
    bar.<|>
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn let   let
                sn letm  let mut
                sn match match expr {}
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
            "#]],
        )
    }

    #[test]
    fn option_iflet() {
        check_edit(
            "ifl",
            r#"
enum Option<T> { Some(T), None }

fn main() {
    let bar = Option::Some(true);
    bar.<|>
}
"#,
            r#"
enum Option<T> { Some(T), None }

fn main() {
    let bar = Option::Some(true);
    if let Some($1) = bar {
    $0
}
}
"#,
        );
    }

    #[test]
    fn result_match() {
        check_edit(
            "match",
            r#"
enum Result<T, E> { Ok(T), Err(E) }

fn main() {
    let bar = Result::Ok(true);
    bar.<|>
}
"#,
            r#"
enum Result<T, E> { Ok(T), Err(E) }

fn main() {
    let bar = Result::Ok(true);
    match bar {
    Ok(${1:_}) => {$2},
    Err(${3:_}) => {$0},
}
}
"#,
        );
    }

    #[test]
    fn postfix_completion_works_for_ambiguous_float_literal() {
        check_edit("refm", r#"fn main() { 42.<|> }"#, r#"fn main() { &mut 42 }"#)
    }

    #[test]
    fn works_in_simple_macro() {
        check_edit(
            "dbg",
            r#"
macro_rules! m { ($e:expr) => { $e } }
fn main() {
    let bar: u8 = 12;
    m!(bar.d<|>)
}
"#,
            r#"
macro_rules! m { ($e:expr) => { $e } }
fn main() {
    let bar: u8 = 12;
    m!(dbg!(bar))
}
"#,
        );
    }

    #[test]
    fn postfix_completion_for_references() {
        check_edit("dbg", r#"fn main() { &&42.<|> }"#, r#"fn main() { dbg!(&&42) }"#);
        check_edit("refm", r#"fn main() { &&42.<|> }"#, r#"fn main() { &&&mut 42 }"#);
    }

    #[test]
    fn postfix_completion_for_format_like_strings() {
        check_edit(
            "fmt",
            r#"fn main() { "{some_var:?}".<|> }"#,
            r#"fn main() { format!("{:?}", some_var) }"#,
        );
        check_edit(
            "panic",
            r#"fn main() { "Panic with {a}".<|> }"#,
            r#"fn main() { panic!("Panic with {}", a) }"#,
        );
        check_edit(
            "println",
            r#"fn main() { "{ 2+2 } { SomeStruct { val: 1, other: 32 } :?}".<|> }"#,
            r#"fn main() { println!("{} {:?}", 2+2, SomeStruct { val: 1, other: 32 }) }"#,
        );
        check_edit(
            "loge",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::error!("{}", 2+2) }"#,
        );
        check_edit(
            "logt",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::trace!("{}", 2+2) }"#,
        );
        check_edit(
            "logd",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::debug!("{}", 2+2) }"#,
        );
        check_edit(
            "logi",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::info!("{}", 2+2) }"#,
        );
        check_edit(
            "logw",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::warn!("{}", 2+2) }"#,
        );
        check_edit(
            "loge",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::error!("{}", 2+2) }"#,
        );
    }
}
<commit_msg>Trigger .let postfix snippet in the middle of the block<commit_after>//! Postfix completions, like `Ok(10).ifl<|>` => `if let Ok() = Ok(10) { <|> }`.

mod format_like;

use ide_db::ty_filter::TryEnum;
use syntax::{
    ast::{self, AstNode, AstToken},
    SyntaxKind::{BLOCK_EXPR, EXPR_STMT},
    TextRange, TextSize,
};
use text_edit::TextEdit;

use self::format_like::add_format_like_completions;
use crate::{
    config::SnippetCap,
    context::CompletionContext,
    item::{Builder, CompletionKind},
    CompletionItem, CompletionItemKind, Completions,
};

pub(crate) fn complete_postfix(acc: &mut Completions, ctx: &CompletionContext) {
    if !ctx.config.enable_postfix_completions {
        return;
    }

    let dot_receiver = match &ctx.dot_receiver {
        Some(it) => it,
        None => return,
    };

    let receiver_text =
        get_receiver_text(dot_receiver, ctx.dot_receiver_is_ambiguous_float_literal);

    let receiver_ty = match ctx.sema.type_of_expr(&dot_receiver) {
        Some(it) => it,
        None => return,
    };

    let cap = match ctx.config.snippet_cap {
        Some(it) => it,
        None => return,
    };
    let try_enum = TryEnum::from_ty(&ctx.sema, &receiver_ty);
    if let Some(try_enum) = &try_enum {
        match try_enum {
            TryEnum::Result => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "ifl",
                    "if let Ok {}",
                    &format!("if let Ok($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);

                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "while",
                    "while let Ok {}",
                    &format!("while let Ok($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);
            }
            TryEnum::Option => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "ifl",
                    "if let Some {}",
                    &format!("if let Some($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);

                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "while",
                    "while let Some {}",
                    &format!("while let Some($1) = {} {{\n    $0\n}}", receiver_text),
                )
                .add_to(acc);
            }
        }
    } else if receiver_ty.is_bool() || receiver_ty.is_unknown() {
        postfix_snippet(
            ctx,
            cap,
            &dot_receiver,
            "if",
            "if expr {}",
            &format!("if {} {{\n    $0\n}}", receiver_text),
        )
        .add_to(acc);
        postfix_snippet(
            ctx,
            cap,
            &dot_receiver,
            "while",
            "while expr {}",
            &format!("while {} {{\n    $0\n}}", receiver_text),
        )
        .add_to(acc);
        postfix_snippet(ctx, cap, &dot_receiver, "not", "!expr", &format!("!{}", receiver_text))
            .add_to(acc);
    }

    postfix_snippet(ctx, cap, &dot_receiver, "ref", "&expr", &format!("&{}", receiver_text))
        .add_to(acc);
    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "refm",
        "&mut expr",
        &format!("&mut {}", receiver_text),
    )
    .add_to(acc);

    // The rest of the postfix completions create an expression that moves an argument,
    // so it's better to consider references now to avoid breaking the compilation
    let dot_receiver = include_references(dot_receiver);
    let receiver_text =
        get_receiver_text(&dot_receiver, ctx.dot_receiver_is_ambiguous_float_literal);

    match try_enum {
        Some(try_enum) => match try_enum {
            TryEnum::Result => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "match",
                    "match expr {}",
                    &format!("match {} {{\n    Ok(${{1:_}}) => {{$2}},\n    Err(${{3:_}}) => {{$0}},\n}}", receiver_text),
                )
                .add_to(acc);
            }
            TryEnum::Option => {
                postfix_snippet(
                    ctx,
                    cap,
                    &dot_receiver,
                    "match",
                    "match expr {}",
                    &format!(
                        "match {} {{\n    Some(${{1:_}}) => {{$2}},\n    None => {{$0}},\n}}",
                        receiver_text
                    ),
                )
                .add_to(acc);
            }
        },
        None => {
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "match",
                "match expr {}",
                &format!("match {} {{\n    ${{1:_}} => {{$0}},\n}}", receiver_text),
            )
            .add_to(acc);
        }
    }

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "box",
        "Box::new(expr)",
        &format!("Box::new({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(ctx, cap, &dot_receiver, "ok", "Ok(expr)", &format!("Ok({})", receiver_text))
        .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "some",
        "Some(expr)",
        &format!("Some({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "dbg",
        "dbg!(expr)",
        &format!("dbg!({})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "dbgr",
        "dbg!(&expr)",
        &format!("dbg!(&{})", receiver_text),
    )
    .add_to(acc);

    postfix_snippet(
        ctx,
        cap,
        &dot_receiver,
        "call",
        "function(expr)",
        &format!("${{1}}({})", receiver_text),
    )
    .add_to(acc);

    if let Some(parent) = dot_receiver.syntax().parent().and_then(|p| p.parent()) {
        if matches!(parent.kind(), BLOCK_EXPR | EXPR_STMT) {
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "let",
                "let",
                &format!("let $0 = {};", receiver_text),
            )
            .add_to(acc);
            postfix_snippet(
                ctx,
                cap,
                &dot_receiver,
                "letm",
                "let mut",
                &format!("let mut $0 = {};", receiver_text),
            )
            .add_to(acc);
        }
    }

    if let ast::Expr::Literal(literal) = dot_receiver.clone() {
        if let Some(literal_text) = ast::String::cast(literal.token()) {
            add_format_like_completions(acc, ctx, &dot_receiver, cap, &literal_text);
        }
    }
}

fn get_receiver_text(receiver: &ast::Expr, receiver_is_ambiguous_float_literal: bool) -> String {
    if receiver_is_ambiguous_float_literal {
        let text = receiver.syntax().text();
        let without_dot = ..text.len() - TextSize::of('.');
        text.slice(without_dot).to_string()
    } else {
        receiver.to_string()
    }
}

fn include_references(initial_element: &ast::Expr) -> ast::Expr {
    let mut resulting_element = initial_element.clone();
    while let Some(parent_ref_element) =
        resulting_element.syntax().parent().and_then(ast::RefExpr::cast)
    {
        resulting_element = ast::Expr::from(parent_ref_element);
    }
    resulting_element
}

fn postfix_snippet(
    ctx: &CompletionContext,
    cap: SnippetCap,
    receiver: &ast::Expr,
    label: &str,
    detail: &str,
    snippet: &str,
) -> Builder {
    let edit = {
        let receiver_syntax = receiver.syntax();
        let receiver_range = ctx.sema.original_range(receiver_syntax).range;
        let delete_range = TextRange::new(receiver_range.start(), ctx.source_range().end());
        TextEdit::replace(delete_range, snippet.to_string())
    };
    CompletionItem::new(CompletionKind::Postfix, ctx.source_range(), label)
        .detail(detail)
        .kind(CompletionItemKind::Snippet)
        .snippet_edit(cap, edit)
}

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

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

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

    #[test]
    fn postfix_completion_works_for_trivial_path_expression() {
        check(
            r#"
fn main() {
    let bar = true;
    bar.<|>
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn if    if expr {}
                sn let   let
                sn letm  let mut
                sn match match expr {}
                sn not   !expr
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
                sn while while expr {}
            "#]],
        );
    }

    #[test]
    fn postfix_completion_works_for_function_calln() {
        check(
            r#"
fn foo(elt: bool) -> bool {
    !elt
}

fn main() {
    let bar = true;
    foo(bar.<|>)
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn if    if expr {}
                sn match match expr {}
                sn not   !expr
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
                sn while while expr {}
            "#]],
        );
    }

    #[test]
    fn postfix_type_filtering() {
        check(
            r#"
fn main() {
    let bar: u8 = 12;
    bar.<|>
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn let   let
                sn letm  let mut
                sn match match expr {}
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
            "#]],
        )
    }

    #[test]
    fn let_middle_block() {
        check(
            r#"
fn main() {
    baz.l<|>
    res
}
"#,
            expect![[r#"
                sn box   Box::new(expr)
                sn call  function(expr)
                sn dbg   dbg!(expr)
                sn dbgr  dbg!(&expr)
                sn if    if expr {}
                sn let   let
                sn letm  let mut
                sn match match expr {}
                sn not   !expr
                sn ok    Ok(expr)
                sn ref   &expr
                sn refm  &mut expr
                sn some  Some(expr)
                sn while while expr {}
            "#]],
        );
    }

    #[test]
    fn option_iflet() {
        check_edit(
            "ifl",
            r#"
enum Option<T> { Some(T), None }

fn main() {
    let bar = Option::Some(true);
    bar.<|>
}
"#,
            r#"
enum Option<T> { Some(T), None }

fn main() {
    let bar = Option::Some(true);
    if let Some($1) = bar {
    $0
}
}
"#,
        );
    }

    #[test]
    fn result_match() {
        check_edit(
            "match",
            r#"
enum Result<T, E> { Ok(T), Err(E) }

fn main() {
    let bar = Result::Ok(true);
    bar.<|>
}
"#,
            r#"
enum Result<T, E> { Ok(T), Err(E) }

fn main() {
    let bar = Result::Ok(true);
    match bar {
    Ok(${1:_}) => {$2},
    Err(${3:_}) => {$0},
}
}
"#,
        );
    }

    #[test]
    fn postfix_completion_works_for_ambiguous_float_literal() {
        check_edit("refm", r#"fn main() { 42.<|> }"#, r#"fn main() { &mut 42 }"#)
    }

    #[test]
    fn works_in_simple_macro() {
        check_edit(
            "dbg",
            r#"
macro_rules! m { ($e:expr) => { $e } }
fn main() {
    let bar: u8 = 12;
    m!(bar.d<|>)
}
"#,
            r#"
macro_rules! m { ($e:expr) => { $e } }
fn main() {
    let bar: u8 = 12;
    m!(dbg!(bar))
}
"#,
        );
    }

    #[test]
    fn postfix_completion_for_references() {
        check_edit("dbg", r#"fn main() { &&42.<|> }"#, r#"fn main() { dbg!(&&42) }"#);
        check_edit("refm", r#"fn main() { &&42.<|> }"#, r#"fn main() { &&&mut 42 }"#);
    }

    #[test]
    fn postfix_completion_for_format_like_strings() {
        check_edit(
            "fmt",
            r#"fn main() { "{some_var:?}".<|> }"#,
            r#"fn main() { format!("{:?}", some_var) }"#,
        );
        check_edit(
            "panic",
            r#"fn main() { "Panic with {a}".<|> }"#,
            r#"fn main() { panic!("Panic with {}", a) }"#,
        );
        check_edit(
            "println",
            r#"fn main() { "{ 2+2 } { SomeStruct { val: 1, other: 32 } :?}".<|> }"#,
            r#"fn main() { println!("{} {:?}", 2+2, SomeStruct { val: 1, other: 32 }) }"#,
        );
        check_edit(
            "loge",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::error!("{}", 2+2) }"#,
        );
        check_edit(
            "logt",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::trace!("{}", 2+2) }"#,
        );
        check_edit(
            "logd",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::debug!("{}", 2+2) }"#,
        );
        check_edit(
            "logi",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::info!("{}", 2+2) }"#,
        );
        check_edit(
            "logw",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::warn!("{}", 2+2) }"#,
        );
        check_edit(
            "loge",
            r#"fn main() { "{2+2}".<|> }"#,
            r#"fn main() { log::error!("{}", 2+2) }"#,
        );
    }
}
<|endoftext|>


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

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

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

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

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

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

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

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

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

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

            PubStruct

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

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

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

            PubSt<|>ruct

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

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

            PubStruct<|>

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

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

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

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

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

            pub mod PubMod {
                pub struct PubStruct;
            }",
        );
    }
}
<commit_msg>Improve auto import message<commit_after>use hir::db::HirDatabase;
use ra_syntax::{
    ast::{self, AstNode},
    SmolStr,
    SyntaxKind::USE_ITEM,
    SyntaxNode,
};

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

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

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

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

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

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

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

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

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

            PubStruct

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

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

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

            PubSt<|>ruct

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

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

            PubStruct<|>

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

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

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

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

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

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


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::conversions::native_from_handleobject;
use dom::bindings::conversions::{ToJSValConvertible};
use dom::bindings::js::{JS, Root};
use dom::bindings::proxyhandler::{get_property_descriptor, fill_property_descriptor};
use dom::bindings::utils::{Reflectable, WindowProxyHandler};
use dom::bindings::utils::get_array_index_from_id;
use dom::document::{Document, DocumentHelpers};
use dom::element::Element;
use dom::window::Window;
use dom::window::WindowHelpers;

use js::jsapi::{JSContext, JSObject, JSPropertyDescriptor, JSErrNum};
use js::jsapi::{HandleObject, HandleId, MutableHandle, MutableHandleValue};
use js::jsapi::{JS_AlreadyHasOwnPropertyById, JS_ForwardGetPropertyTo};
use js::jsapi::{JS_GetOwnPropertyDescriptorById, JS_DefinePropertyById6};
use js::jsapi::{JS_ForwardSetPropertyTo, ObjectOpResult, RootedObject, RootedValue, Handle, HandleValue, Heap};
use js::jsapi::{JSAutoRequest, JSAutoCompartment};
use js::jsval::{ObjectValue, UndefinedValue};
use js::glue::{GetProxyPrivate};
use js::glue::{WrapperNew, CreateWrapperProxyHandler, ProxyTraps};
use js::{JSTrue, JSFalse};

use std::ptr;
use std::default::Default;

#[derive(JSTraceable, HeapSizeOf)]
#[privatize]
#[allow(raw_pointer_derive)]
#[must_root]
pub struct BrowsingContext {
    history: Vec<SessionHistoryEntry>,
    active_index: usize,
    window_proxy: Heap<*mut JSObject>,
    frame_element: Option<JS<Element>>,
}

impl BrowsingContext {
    pub fn new(document: &Document, frame_element: Option<&Element>) -> BrowsingContext {
        BrowsingContext {
            history: vec!(SessionHistoryEntry::new(document)),
            active_index: 0,
            window_proxy: Heap::default(),
            frame_element: frame_element.map(JS::from_ref),
        }
    }

    pub fn active_document(&self) -> Root<Document> {
        self.history[self.active_index].document.root()
    }

    pub fn active_window(&self) -> Root<Window> {
        let doc = self.active_document();
        doc.r().window()
    }

    pub fn frame_element(&self) -> Option<Root<Element>> {
        self.frame_element.map(Root::from_rooted)
    }

    pub fn window_proxy(&self) -> *mut JSObject {
        assert!(!self.window_proxy.get().is_null());
        self.window_proxy.get()
    }

    #[allow(unsafe_code)]
    pub fn create_window_proxy(&mut self) {
        let win = self.active_window();
        let win = win.r();

        let WindowProxyHandler(handler) = win.windowproxy_handler();
        assert!(!handler.is_null());

        let cx = win.get_cx();
        let _ar = JSAutoRequest::new(cx);
        let parent = win.reflector().get_jsobject();
        let _ac = JSAutoCompartment::new(cx, parent.get());
        let wrapper = unsafe { WrapperNew(cx, parent, handler) };
        assert!(!wrapper.is_null());
        self.window_proxy.set(wrapper);
    }
}

// This isn't a DOM struct, just a convenience struct
// without a reflector, so we don't mark this as #[dom_struct]
#[must_root]
#[privatize]
#[derive(JSTraceable, HeapSizeOf)]
pub struct SessionHistoryEntry {
    document: JS<Document>,
    children: Vec<BrowsingContext>
}

impl SessionHistoryEntry {
    fn new(document: &Document) -> SessionHistoryEntry {
        SessionHistoryEntry {
            document: JS::from_ref(document),
            children: vec!()
        }
    }
}

#[allow(unsafe_code)]
unsafe fn GetSubframeWindow(cx: *mut JSContext, proxy: HandleObject, id: HandleId) -> Option<Root<Window>> {
    let index = get_array_index_from_id(cx, id);
    if let Some(index) = index {
        let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
        let win: Root<Window> = native_from_handleobject(target.handle()).unwrap();
        let mut found = false;
        return win.r().IndexedGetter(index, &mut found);
    }

    None
}

#[allow(unsafe_code)]
unsafe extern fn getOwnPropertyDescriptor(cx: *mut JSContext, proxy: HandleObject, id: HandleId,
                                          desc: MutableHandle<JSPropertyDescriptor>) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if let Some(window) = window {
        let mut val = RootedValue::new(cx, UndefinedValue());
        window.to_jsval(cx, val.handle_mut());
        (*desc.ptr).value = val.ptr;
        fill_property_descriptor(&mut *desc.ptr, *proxy.ptr, true);
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());

    if JS_GetOwnPropertyDescriptorById(cx, target.handle(), id, desc) == 0 {
        return JSFalse;
    }

    if (*desc.ptr).obj != target.ptr {
        // Not an own property
        (*desc.ptr).obj = ptr::null_mut();
    } else {
        (*desc.ptr).obj = *proxy.ptr;
    }

    JSTrue
}

#[allow(unsafe_code)]
unsafe extern fn defineProperty(cx: *mut JSContext,
                                proxy: HandleObject,
                                id: HandleId,
                                desc: Handle<JSPropertyDescriptor>,
                                res: *mut ObjectOpResult) -> u8 {
    if get_array_index_from_id(cx, id).is_some() {
        // Spec says to Reject whether this is a supported index or not,
        // since we have no indexed setter or indexed creator.  That means
        // throwing in strict mode (FIXME: Bug 828137), doing nothing in
        // non-strict mode.
       (*res).code_ = JSErrNum::JSMSG_CANT_DEFINE_WINDOW_ELEMENT as u32;
       return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    JS_DefinePropertyById6(cx, target.handle(), id, desc, res)
}

#[allow(unsafe_code)]
unsafe extern fn hasOwn(cx: *mut JSContext, proxy: HandleObject, id: HandleId, bp: *mut u8) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if window.is_some() {
        *bp = JSTrue;
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    let mut found = 0;
    if JS_AlreadyHasOwnPropertyById(cx, target.handle(), id, &mut found) == 0 {
        return JSFalse;
    }

    *bp = (found != 0) as u8;
    return JSTrue;
}

#[allow(unsafe_code)]
unsafe extern fn get(cx: *mut JSContext,
                     proxy: HandleObject,
                     receiver: HandleObject,
                     id: HandleId,
                     vp: MutableHandleValue) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if let Some(window) = window {
        window.to_jsval(cx, vp);
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    JS_ForwardGetPropertyTo(cx, target.handle(), id, receiver, vp)
}

#[allow(unsafe_code)]
unsafe extern fn set(cx: *mut JSContext,
                     proxy: HandleObject,
                     receiver: HandleObject,
                     id: HandleId,
                     vp: MutableHandleValue,
                     res: *mut ObjectOpResult) -> u8 {
    if get_array_index_from_id(cx, id).is_some() {
        // Reject (which means throw if and only if strict) the set.
        (*res).code_ = JSErrNum::JSMSG_READ_ONLY as u32;
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    let receiver = RootedValue::new(cx, ObjectValue(&**receiver.ptr));
    JS_ForwardSetPropertyTo(cx, target.handle(), id, HandleValue { ptr: vp.ptr }, receiver.handle(), res)
}

static PROXY_HANDLER: ProxyTraps = ProxyTraps {
    enter: None,
    getOwnPropertyDescriptor: Some(getOwnPropertyDescriptor),
    defineProperty: Some(defineProperty),
    ownPropertyKeys: None,
    delete_: None,
    enumerate: None,
    preventExtensions: None,
    isExtensible: None,
    has: None,
    get: Some(get),
    set: Some(set),
    call: None,
    construct: None,
    getPropertyDescriptor: Some(get_property_descriptor),
    hasOwn: Some(hasOwn),
    getOwnEnumerablePropertyKeys: None,
    nativeCall: None,
    hasInstance: None,
    objectClassIs: None,
    className: None,
    fun_toString: None,
    boxedValue_unbox: None,
    defaultValue: None,
    trace: None,
    finalize: None,
    objectMoved: None,
    isCallable: None,
    isConstructor: None,
};

#[allow(unsafe_code)]
pub fn new_window_proxy_handler() -> WindowProxyHandler {
    unsafe {
        WindowProxyHandler(CreateWrapperProxyHandler(&PROXY_HANDLER))
    }
}
<commit_msg>Replace code in getOwnPropertyDescriptor with an assertion<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::conversions::native_from_handleobject;
use dom::bindings::conversions::{ToJSValConvertible};
use dom::bindings::js::{JS, Root};
use dom::bindings::proxyhandler::{get_property_descriptor, fill_property_descriptor};
use dom::bindings::utils::{Reflectable, WindowProxyHandler};
use dom::bindings::utils::get_array_index_from_id;
use dom::document::{Document, DocumentHelpers};
use dom::element::Element;
use dom::window::Window;
use dom::window::WindowHelpers;

use js::jsapi::{JSContext, JSObject, JSPropertyDescriptor, JSErrNum};
use js::jsapi::{HandleObject, HandleId, MutableHandle, MutableHandleValue};
use js::jsapi::{JS_AlreadyHasOwnPropertyById, JS_ForwardGetPropertyTo};
use js::jsapi::{JS_GetOwnPropertyDescriptorById, JS_DefinePropertyById6};
use js::jsapi::{JS_ForwardSetPropertyTo, ObjectOpResult, RootedObject, RootedValue, Handle, HandleValue, Heap};
use js::jsapi::{JSAutoRequest, JSAutoCompartment};
use js::jsval::{ObjectValue, UndefinedValue};
use js::glue::{GetProxyPrivate};
use js::glue::{WrapperNew, CreateWrapperProxyHandler, ProxyTraps};
use js::{JSTrue, JSFalse};

use std::default::Default;

#[derive(JSTraceable, HeapSizeOf)]
#[privatize]
#[allow(raw_pointer_derive)]
#[must_root]
pub struct BrowsingContext {
    history: Vec<SessionHistoryEntry>,
    active_index: usize,
    window_proxy: Heap<*mut JSObject>,
    frame_element: Option<JS<Element>>,
}

impl BrowsingContext {
    pub fn new(document: &Document, frame_element: Option<&Element>) -> BrowsingContext {
        BrowsingContext {
            history: vec!(SessionHistoryEntry::new(document)),
            active_index: 0,
            window_proxy: Heap::default(),
            frame_element: frame_element.map(JS::from_ref),
        }
    }

    pub fn active_document(&self) -> Root<Document> {
        self.history[self.active_index].document.root()
    }

    pub fn active_window(&self) -> Root<Window> {
        let doc = self.active_document();
        doc.r().window()
    }

    pub fn frame_element(&self) -> Option<Root<Element>> {
        self.frame_element.map(Root::from_rooted)
    }

    pub fn window_proxy(&self) -> *mut JSObject {
        assert!(!self.window_proxy.get().is_null());
        self.window_proxy.get()
    }

    #[allow(unsafe_code)]
    pub fn create_window_proxy(&mut self) {
        let win = self.active_window();
        let win = win.r();

        let WindowProxyHandler(handler) = win.windowproxy_handler();
        assert!(!handler.is_null());

        let cx = win.get_cx();
        let _ar = JSAutoRequest::new(cx);
        let parent = win.reflector().get_jsobject();
        let _ac = JSAutoCompartment::new(cx, parent.get());
        let wrapper = unsafe { WrapperNew(cx, parent, handler) };
        assert!(!wrapper.is_null());
        self.window_proxy.set(wrapper);
    }
}

// This isn't a DOM struct, just a convenience struct
// without a reflector, so we don't mark this as #[dom_struct]
#[must_root]
#[privatize]
#[derive(JSTraceable, HeapSizeOf)]
pub struct SessionHistoryEntry {
    document: JS<Document>,
    children: Vec<BrowsingContext>
}

impl SessionHistoryEntry {
    fn new(document: &Document) -> SessionHistoryEntry {
        SessionHistoryEntry {
            document: JS::from_ref(document),
            children: vec!()
        }
    }
}

#[allow(unsafe_code)]
unsafe fn GetSubframeWindow(cx: *mut JSContext, proxy: HandleObject, id: HandleId) -> Option<Root<Window>> {
    let index = get_array_index_from_id(cx, id);
    if let Some(index) = index {
        let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
        let win: Root<Window> = native_from_handleobject(target.handle()).unwrap();
        let mut found = false;
        return win.r().IndexedGetter(index, &mut found);
    }

    None
}

#[allow(unsafe_code)]
unsafe extern fn getOwnPropertyDescriptor(cx: *mut JSContext, proxy: HandleObject, id: HandleId,
                                          desc: MutableHandle<JSPropertyDescriptor>) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if let Some(window) = window {
        let mut val = RootedValue::new(cx, UndefinedValue());
        window.to_jsval(cx, val.handle_mut());
        (*desc.ptr).value = val.ptr;
        fill_property_descriptor(&mut *desc.ptr, *proxy.ptr, true);
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());

    if JS_GetOwnPropertyDescriptorById(cx, target.handle(), id, desc) == 0 {
        return JSFalse;
    }

    assert!((*desc.ptr).obj.is_null() || (*desc.ptr).obj == *proxy.ptr);

    JSTrue
}

#[allow(unsafe_code)]
unsafe extern fn defineProperty(cx: *mut JSContext,
                                proxy: HandleObject,
                                id: HandleId,
                                desc: Handle<JSPropertyDescriptor>,
                                res: *mut ObjectOpResult) -> u8 {
    if get_array_index_from_id(cx, id).is_some() {
        // Spec says to Reject whether this is a supported index or not,
        // since we have no indexed setter or indexed creator.  That means
        // throwing in strict mode (FIXME: Bug 828137), doing nothing in
        // non-strict mode.
       (*res).code_ = JSErrNum::JSMSG_CANT_DEFINE_WINDOW_ELEMENT as u32;
       return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    JS_DefinePropertyById6(cx, target.handle(), id, desc, res)
}

#[allow(unsafe_code)]
unsafe extern fn hasOwn(cx: *mut JSContext, proxy: HandleObject, id: HandleId, bp: *mut u8) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if window.is_some() {
        *bp = JSTrue;
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    let mut found = 0;
    if JS_AlreadyHasOwnPropertyById(cx, target.handle(), id, &mut found) == 0 {
        return JSFalse;
    }

    *bp = (found != 0) as u8;
    return JSTrue;
}

#[allow(unsafe_code)]
unsafe extern fn get(cx: *mut JSContext,
                     proxy: HandleObject,
                     receiver: HandleObject,
                     id: HandleId,
                     vp: MutableHandleValue) -> u8 {
    let window = GetSubframeWindow(cx, proxy, id);
    if let Some(window) = window {
        window.to_jsval(cx, vp);
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    JS_ForwardGetPropertyTo(cx, target.handle(), id, receiver, vp)
}

#[allow(unsafe_code)]
unsafe extern fn set(cx: *mut JSContext,
                     proxy: HandleObject,
                     receiver: HandleObject,
                     id: HandleId,
                     vp: MutableHandleValue,
                     res: *mut ObjectOpResult) -> u8 {
    if get_array_index_from_id(cx, id).is_some() {
        // Reject (which means throw if and only if strict) the set.
        (*res).code_ = JSErrNum::JSMSG_READ_ONLY as u32;
        return JSTrue;
    }

    let target = RootedObject::new(cx, GetProxyPrivate(*proxy.ptr).to_object());
    let receiver = RootedValue::new(cx, ObjectValue(&**receiver.ptr));
    JS_ForwardSetPropertyTo(cx, target.handle(), id, HandleValue { ptr: vp.ptr }, receiver.handle(), res)
}

static PROXY_HANDLER: ProxyTraps = ProxyTraps {
    enter: None,
    getOwnPropertyDescriptor: Some(getOwnPropertyDescriptor),
    defineProperty: Some(defineProperty),
    ownPropertyKeys: None,
    delete_: None,
    enumerate: None,
    preventExtensions: None,
    isExtensible: None,
    has: None,
    get: Some(get),
    set: Some(set),
    call: None,
    construct: None,
    getPropertyDescriptor: Some(get_property_descriptor),
    hasOwn: Some(hasOwn),
    getOwnEnumerablePropertyKeys: None,
    nativeCall: None,
    hasInstance: None,
    objectClassIs: None,
    className: None,
    fun_toString: None,
    boxedValue_unbox: None,
    defaultValue: None,
    trace: None,
    finalize: None,
    objectMoved: None,
    isCallable: None,
    isConstructor: None,
};

#[allow(unsafe_code)]
pub fn new_window_proxy_handler() -> WindowProxyHandler {
    unsafe {
        WindowProxyHandler(CreateWrapperProxyHandler(&PROXY_HANDLER))
    }
}
<|endoftext|>


<commit_before>use std::f32::{NEG_INFINITY, INFINITY};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;

use time;

use identity::{Team, JobDescription, Agent};
use life::{Commit, WorldState};
use motion::CENTER_OF_THE_WORLD;
use space::Pinfield;

pub fn orientation(team: Team) -> f32 {
    match team {
        Team::Orange => 1.0,
        Team::Blue => -1.0,
    }
}

pub fn figurine_valuation(agent: Agent) -> f32 {
    let value = match agent.job_description {
        // en.wikipedia.org/wiki/
        // Chess_piece_relative_value#Hans_Berliner.27s_system
        JobDescription::Servant => 1.0,
        JobDescription::Pony => 3.2,
        JobDescription::Scholar => 3.3,
        JobDescription::Cop => 5.1,
        JobDescription::Princess => 8.8,
        JobDescription::Figurehead => 20000.0,
    };
    orientation(agent.team) * value
}

pub fn score(world: WorldState) -> f32 {
    let mut valuation = 0.0;

    for team in Team::league().into_iter() {
        for agent in Agent::dramatis_personæ(team).into_iter() {
            valuation += world.agent_to_pinfield_ref(
                agent).pincount() as f32 * figurine_valuation(agent);
        }
        // breadth of scholarship bonus
        if world.agent_to_pinfield_ref(Agent {
                    team: team,
                    job_description: JobDescription::Scholar,
                })
                .to_locales()
                .len() >= 2 {
            valuation += orientation(team) * 0.5
        }
    }

    // ponies and servants want to be in the center of the world's
    // action, maybe??
    let center = Pinfield(CENTER_OF_THE_WORLD);
    // cast to signed to avoid overflow
    let orange_centerism: i8 = world.orange_servants.union(
        world.orange_ponies).intersection(center).pincount() as i8;
    let blue_centerism: i8 = world.blue_servants.union(
        world.blue_ponies).intersection(center).pincount() as i8;
    valuation += (orange_centerism - blue_centerism) as f32 * 0.1;

    valuation
}

fn mmv_lva_heuristic(commit: &Commit) -> f32 {
    // https://chessprogramming.wikispaces.com/MVV-LVA
    match commit.hospitalization {
        Some(patient) => (
            figurine_valuation(patient) - figurine_valuation(commit.patch.star)),
        None => 0.0,
    }
}

fn order_moves(commits: &mut Vec<Commit>) {
    commits.sort_by(|a, b| {
        mmv_lva_heuristic(&b)
            .partial_cmp(&mmv_lva_heuristic(&a))
            .unwrap_or(Ordering::Equal)
    });
}

pub fn α_β_negamax_search(
    world: WorldState, depth: u8, mut α: f32, β: f32,
    memory_bank: Arc<Mutex<HashMap<WorldState, f32>>>) -> (Option<Commit>, f32) {
    let team = world.to_move;
    let potential_score = orientation(team) * score(world);
    let mut premonitions = world.reckless_lookahead();
    order_moves(&mut premonitions);
    if depth == 0 || premonitions.is_empty() {
        return (None, potential_score)
    };
    let mut optimum = NEG_INFINITY;
    let mut optimand = None;
    for premonition in premonitions.into_iter() {
        let mut value = NEG_INFINITY;  // can't hurt to be pessimistic
        let cached: bool;
        {
            let open_vault = memory_bank.lock().unwrap();
            let cached_value_maybe = open_vault.get(&premonition.tree);
            match cached_value_maybe {
                Some(&cached_value) => {
                    cached = true;
                    value = cached_value;
                }
                None => { cached = false; }
            };
        }

        if !cached {
            let (_, acquired_value) = α_β_negamax_search(
                premonition.tree, depth - 1,
                -β, -α, memory_bank.clone()
            );
            value = -acquired_value;
            memory_bank.lock().unwrap().insert(premonition.tree, value);
        }

        if value > optimum {
            optimum = value;
            optimand = Some(premonition);
        }
        if value > α {
            α = value;
        }
        if α >= β {
            break;
        }
    }

    (optimand, optimum)
}

pub fn potentially_timebound_kickoff(world: &WorldState, depth: u8,
                                     nihilistically: bool,
                                     deadline_maybe: Option<time::Timespec>)
                                     -> Option<Vec<(Commit, f32)>> {
    let déjà_vu_table: HashMap<WorldState, f32> = HashMap::new();
    let memory_bank = Arc::new(Mutex::new(déjà_vu_table));
    let mut premonitions;
    if nihilistically {
        premonitions = world.reckless_lookahead();
    } else {
        premonitions = world.lookahead();
    }
    order_moves(&mut premonitions);
    let mut forecasts = Vec::new();
    let (tx, rx) = mpsc::channel();
    for &premonition in &premonitions {
        let travel_bank = memory_bank.clone();
        let explorer_radio = tx.clone();
        thread::spawn(move || {
            let search_hit = α_β_negamax_search(
                premonition.tree, depth - 1,
                NEG_INFINITY, INFINITY,
                travel_bank
            );
            explorer_radio.send(search_hit).ok();
        });
        loop {  // polling for results
            if let Some(deadline) = deadline_maybe {
                if time::get_time() > deadline {
                    return None;
                }
            }
            let search_hit_maybe = rx.try_recv().ok();
            match search_hit_maybe {
                Some(search_hit) => {
                    let (_grandchild, mut value) = search_hit;
                    value = -value;
                    forecasts.push((premonition, value));
                    break;
                },
                None => { thread::sleep_ms(2); }
            }
        }
    }
    forecasts.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
    Some(forecasts)
}

pub fn kickoff(world: &WorldState, depth: u8,
               nihilistically: bool) -> Vec<(Commit, f32)> {
    potentially_timebound_kickoff(world, depth, nihilistically, None).unwrap()
}

pub fn iterative_deepening_kickoff(world: &WorldState, timeout: time::Duration,
                                   nihilistically: bool)
                                   -> (Vec<(Commit, f32)>, u8) {
    let deadline = time::get_time() + timeout;
    let mut depth = 1;
    let mut forecasts = potentially_timebound_kickoff(
        world, depth, nihilistically, None).unwrap();
    while let Some(prophecy) = potentially_timebound_kickoff(
            world, depth, nihilistically, Some(deadline)) {
        forecasts = prophecy;
        depth += 1;
    }
    (forecasts, depth-1)
}

#[cfg(test)]
mod tests {
    extern crate test;
    use self::test::Bencher;

    use time;
    use super::{kickoff, score};
    use space::Locale;
    use life::WorldState;
    use identity::Team;

    impl WorldState {
        fn no_castling_at_all(&mut self) {
            self.orange_east_service_eligibility = false;
            self.orange_west_service_eligibility = false;
            self.blue_east_service_eligibility = false;
            self.blue_west_service_eligibility = false;
        }
    }

    #[bench]
    fn benchmark_scoring(b: &mut Bencher) {
        b.iter(|| score(WorldState::new()));
    }

    #[bench]
    fn benchmark_kickoff_depth_1(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 1, true));
    }

    #[bench]
    fn benchmark_kickoff_depth_2_arbys(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 2, true));
    }

    #[bench]
    fn benchmark_kickoff_depth_2_carefully(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 2, false));
    }

    #[ignore]
    #[bench]
    fn benchmark_kickoff_depth_3(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 3, true));
    }

    #[test]
    #[ignore]  // more research is needed
    fn concerning_short_circuiting_upon_finding_critical_endangerment() {
        let ws = WorldState::reconstruct("7K/r7/1r6/8/8/8/8/7k b -".to_owned());
        let start = time::get_time();
        kickoff(&ws, 30, true);
        let duration = time::get_time() - start;
        assert!(duration.num_seconds() < 20);
    }

    #[test]
    #[allow(float_cmp)]
    fn concerning_fairness_of_the_initial_position() {
        // It's okay to assume this is really 0.0. Floats may be imprecise,
        // but they do have well-defined behavior.
        assert_eq!(0.0, score(WorldState::new()));
    }

    #[test]
    fn concerning_servant_ascension_choices() {
        let ws = WorldState::reconstruct("8/q1P1k/8/8/8/8/6PP/7K w -".to_owned());
        // looking ahead 3 moves allows the Leafline AI to catch the
        // split, whereby transforming into a pony (rather than
        // transitioning into a princess, as would usually be
        // expected) endangers both the blue princess and figurehead
        let (ref best_move, score) = kickoff(&ws, 3, true)[0];
        println!("{:?}", best_move);
        assert!(score > 0.0);
        assert_eq!(best_move.tree.preserve(), "2N5/q3k3/8/8/8/8/6PP/7K b -");
    }

    #[test]
    fn experimentally_about_kickoff() {
        let mut world = WorldState::new_except_empty();
        // SCENARIO: let's imagine Orange (to move) has separate attacks against
        // Blue's pony and servant, against which Blue has no defense but
        // to run away. We predict that Orange will take the pony, and
        // then Blue will move the servant out of the way.

        // scholar endangers pony
        world.blue_ponies = world.blue_ponies.alight(
            Locale { rank: 0, file: 0 });
        world.orange_scholars = world.orange_scholars.alight(
            Locale { rank: 2, file: 2 });

        // pony endangers servant
        world.blue_servants = world.blue_servants.alight(
            Locale { rank: 7, file: 1 });
        world.orange_ponies = world.orange_ponies.alight(
            Locale { rank: 5, file: 2 });

        // Blue has another servant sitting nowhere interesting
        world.blue_servants = world.blue_servants.alight(
            Locale { rank: 3, file: 6 });
        world.no_castling_at_all();

        let depth = 2;
        let advisory = kickoff(&world, depth, true);

        // taking the pony is the right thing to do
        assert_eq!(Locale { rank: 0, file: 0 }, advisory[0].0.patch.whither);

        // And, furthermore, the answer should be the same if we face the
        // same situation with the colors reversed
        //
        // XXX this would be tidier and less copy-pastey if I had more
        // general figurine-placing functions that were three rather than
        // two levels of abstraction above twiddling bits on an unsigned
        // int ... oh, well
        let mut negaworld = WorldState::new_except_empty();
        negaworld.to_move = Team::Blue;

        // scholar endangers pony
        negaworld.orange_ponies = negaworld.orange_ponies.alight(
            Locale { rank: 0, file: 0 });
        negaworld.blue_scholars = negaworld.blue_scholars.alight(
            Locale { rank: 2, file: 2 });

        // pony endangers servant
        negaworld.orange_servants = negaworld.orange_servants.alight(
            Locale { rank: 7, file: 1 });
        negaworld.blue_ponies = negaworld.blue_ponies.alight(
            Locale { rank: 5, file: 2 });

        // Orange has another servant sitting nowhere interesting
        negaworld.orange_servants = negaworld.orange_servants.alight(
            Locale { rank: 3, file: 6 });
        negaworld.to_move = Team::Blue;

        negaworld.no_castling_at_all();

        let negadvisory = kickoff(&negaworld, depth, true);

        // taking the pony is still the right thing to do, even in the
        // negaworld
        assert_eq!(Locale { rank: 0, file: 0 },
                   negadvisory[0].0.patch.whither);
    }

}
<commit_msg>use more cores<commit_after>use std::f32::{NEG_INFINITY, INFINITY};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;

use time;

use identity::{Team, JobDescription, Agent};
use life::{Commit, WorldState};
use motion::CENTER_OF_THE_WORLD;
use space::Pinfield;

pub fn orientation(team: Team) -> f32 {
    match team {
        Team::Orange => 1.0,
        Team::Blue => -1.0,
    }
}

pub fn figurine_valuation(agent: Agent) -> f32 {
    let value = match agent.job_description {
        // en.wikipedia.org/wiki/
        // Chess_piece_relative_value#Hans_Berliner.27s_system
        JobDescription::Servant => 1.0,
        JobDescription::Pony => 3.2,
        JobDescription::Scholar => 3.3,
        JobDescription::Cop => 5.1,
        JobDescription::Princess => 8.8,
        JobDescription::Figurehead => 20000.0,
    };
    orientation(agent.team) * value
}

pub fn score(world: WorldState) -> f32 {
    let mut valuation = 0.0;

    for team in Team::league().into_iter() {
        for agent in Agent::dramatis_personæ(team).into_iter() {
            valuation += world.agent_to_pinfield_ref(
                agent).pincount() as f32 * figurine_valuation(agent);
        }
        // breadth of scholarship bonus
        if world.agent_to_pinfield_ref(Agent {
                    team: team,
                    job_description: JobDescription::Scholar,
                })
                .to_locales()
                .len() >= 2 {
            valuation += orientation(team) * 0.5
        }
    }

    // ponies and servants want to be in the center of the world's
    // action, maybe??
    let center = Pinfield(CENTER_OF_THE_WORLD);
    // cast to signed to avoid overflow
    let orange_centerism: i8 = world.orange_servants.union(
        world.orange_ponies).intersection(center).pincount() as i8;
    let blue_centerism: i8 = world.blue_servants.union(
        world.blue_ponies).intersection(center).pincount() as i8;
    valuation += (orange_centerism - blue_centerism) as f32 * 0.1;

    valuation
}

fn mmv_lva_heuristic(commit: &Commit) -> f32 {
    // https://chessprogramming.wikispaces.com/MVV-LVA
    match commit.hospitalization {
        Some(patient) => (
            figurine_valuation(patient) - figurine_valuation(commit.patch.star)),
        None => 0.0,
    }
}

fn order_moves(commits: &mut Vec<Commit>) {
    commits.sort_by(|a, b| {
        mmv_lva_heuristic(&b)
            .partial_cmp(&mmv_lva_heuristic(&a))
            .unwrap_or(Ordering::Equal)
    });
}

pub fn α_β_negamax_search(
    world: WorldState, depth: u8, mut α: f32, β: f32,
    memory_bank: Arc<Mutex<HashMap<WorldState, f32>>>) -> (Option<Commit>, f32) {
    let team = world.to_move;
    let potential_score = orientation(team) * score(world);
    let mut premonitions = world.reckless_lookahead();
    order_moves(&mut premonitions);
    if depth == 0 || premonitions.is_empty() {
        return (None, potential_score)
    };
    let mut optimum = NEG_INFINITY;
    let mut optimand = None;
    for premonition in premonitions.into_iter() {
        let mut value = NEG_INFINITY;  // can't hurt to be pessimistic
        let cached: bool;
        {
            let open_vault = memory_bank.lock().unwrap();
            let cached_value_maybe = open_vault.get(&premonition.tree);
            match cached_value_maybe {
                Some(&cached_value) => {
                    cached = true;
                    value = cached_value;
                }
                None => { cached = false; }
            };
        }

        if !cached {
            let (_, acquired_value) = α_β_negamax_search(
                premonition.tree, depth - 1,
                -β, -α, memory_bank.clone()
            );
            value = -acquired_value;
            memory_bank.lock().unwrap().insert(premonition.tree, value);
        }

        if value > optimum {
            optimum = value;
            optimand = Some(premonition);
        }
        if value > α {
            α = value;
        }
        if α >= β {
            break;
        }
    }

    (optimand, optimum)
}

#[allow(type_complexity)]
pub fn potentially_timebound_kickoff(world: &WorldState, depth: u8,
                                     nihilistically: bool,
                                     deadline_maybe: Option<time::Timespec>)
                                     -> Option<Vec<(Commit, f32)>> {
    let déjà_vu_table: HashMap<WorldState, f32> = HashMap::new();
    let memory_bank = Arc::new(Mutex::new(déjà_vu_table));
    let mut premonitions;
    if nihilistically {
        premonitions = world.reckless_lookahead();
    } else {
        premonitions = world.lookahead();
    }
    order_moves(&mut premonitions);
    let mut forecasts = Vec::new();
    let mut time_radios: Vec<(Commit, mpsc::Receiver<(Option<Commit>, f32)>)> =
        Vec::new();
    for &premonition in &premonitions {
        let travel_bank = memory_bank.clone();
        let (tx, rx) = mpsc::channel();
        let explorer_radio = tx.clone();
        time_radios.push((premonition, rx));
        thread::spawn(move || {
            let search_hit: (Option<Commit>, f32) = α_β_negamax_search(
                premonition.tree, depth - 1,
                NEG_INFINITY, INFINITY,
                travel_bank
            );
            explorer_radio.send(search_hit).ok();
        });
    }
    while !time_radios.is_empty() {  // polling for results
        if let Some(deadline) = deadline_maybe {
            if time::get_time() > deadline {
                return None;
            }
        }
        // iterate over indices so that we can use swap_remove during the loop
        for i in (0..time_radios.len()).rev() {
            let premonition = time_radios[i].0;
            if let Some(search_hit) = time_radios[i].1.try_recv().ok() {
                let (_grandchild, mut value) = search_hit;
                value = -value;
                forecasts.push((premonition, value));
                time_radios.swap_remove(i);
            }
        }
        thread::sleep_ms(2);
    }
    forecasts.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
    Some(forecasts)
}

pub fn kickoff(world: &WorldState, depth: u8,
               nihilistically: bool) -> Vec<(Commit, f32)> {
    potentially_timebound_kickoff(world, depth, nihilistically, None).unwrap()
}

pub fn iterative_deepening_kickoff(world: &WorldState, timeout: time::Duration,
                                   nihilistically: bool)
                                   -> (Vec<(Commit, f32)>, u8) {
    let deadline = time::get_time() + timeout;
    let mut depth = 1;
    let mut forecasts = potentially_timebound_kickoff(
        world, depth, nihilistically, None).unwrap();
    while let Some(prophecy) = potentially_timebound_kickoff(
            world, depth, nihilistically, Some(deadline)) {
        forecasts = prophecy;
        depth += 1;
    }
    (forecasts, depth-1)
}

#[cfg(test)]
mod tests {
    extern crate test;
    use self::test::Bencher;

    use time;
    use super::{kickoff, score};
    use space::Locale;
    use life::WorldState;
    use identity::Team;

    impl WorldState {
        fn no_castling_at_all(&mut self) {
            self.orange_east_service_eligibility = false;
            self.orange_west_service_eligibility = false;
            self.blue_east_service_eligibility = false;
            self.blue_west_service_eligibility = false;
        }
    }

    #[bench]
    fn benchmark_scoring(b: &mut Bencher) {
        b.iter(|| score(WorldState::new()));
    }

    #[bench]
    fn benchmark_kickoff_depth_1(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 1, true));
    }

    #[bench]
    fn benchmark_kickoff_depth_2_arbys(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 2, true));
    }

    #[bench]
    fn benchmark_kickoff_depth_2_carefully(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 2, false));
    }

    #[bench]
    fn benchmark_kickoff_depth_3(b: &mut Bencher) {
        let ws = WorldState::new();
        b.iter(|| kickoff(&ws, 3, true));
    }

    #[test]
    #[ignore]  // more research is needed
    fn concerning_short_circuiting_upon_finding_critical_endangerment() {
        let ws = WorldState::reconstruct("7K/r7/1r6/8/8/8/8/7k b -".to_owned());
        let start = time::get_time();
        kickoff(&ws, 30, true);
        let duration = time::get_time() - start;
        assert!(duration.num_seconds() < 20);
    }

    #[test]
    #[allow(float_cmp)]
    fn concerning_fairness_of_the_initial_position() {
        // It's okay to assume this is really 0.0. Floats may be imprecise,
        // but they do have well-defined behavior.
        assert_eq!(0.0, score(WorldState::new()));
    }

    #[test]
    fn concerning_servant_ascension_choices() {
        let ws = WorldState::reconstruct("8/q1P1k/8/8/8/8/6PP/7K w -".to_owned());
        // looking ahead 3 moves allows the Leafline AI to catch the
        // split, whereby transforming into a pony (rather than
        // transitioning into a princess, as would usually be
        // expected) endangers both the blue princess and figurehead
        let (ref best_move, score) = kickoff(&ws, 3, true)[0];
        println!("{:?}", best_move);
        assert!(score > 0.0);
        assert_eq!(best_move.tree.preserve(), "2N5/q3k3/8/8/8/8/6PP/7K b -");
    }

    #[test]
    fn experimentally_about_kickoff() {
        let mut world = WorldState::new_except_empty();
        // SCENARIO: let's imagine Orange (to move) has separate attacks against
        // Blue's pony and servant, against which Blue has no defense but
        // to run away. We predict that Orange will take the pony, and
        // then Blue will move the servant out of the way.

        // scholar endangers pony
        world.blue_ponies = world.blue_ponies.alight(
            Locale { rank: 0, file: 0 });
        world.orange_scholars = world.orange_scholars.alight(
            Locale { rank: 2, file: 2 });

        // pony endangers servant
        world.blue_servants = world.blue_servants.alight(
            Locale { rank: 7, file: 1 });
        world.orange_ponies = world.orange_ponies.alight(
            Locale { rank: 5, file: 2 });

        // Blue has another servant sitting nowhere interesting
        world.blue_servants = world.blue_servants.alight(
            Locale { rank: 3, file: 6 });
        world.no_castling_at_all();

        let depth = 2;
        let advisory = kickoff(&world, depth, true);

        // taking the pony is the right thing to do
        assert_eq!(Locale { rank: 0, file: 0 }, advisory[0].0.patch.whither);

        // And, furthermore, the answer should be the same if we face the
        // same situation with the colors reversed
        //
        // XXX this would be tidier and less copy-pastey if I had more
        // general figurine-placing functions that were three rather than
        // two levels of abstraction above twiddling bits on an unsigned
        // int ... oh, well
        let mut negaworld = WorldState::new_except_empty();
        negaworld.to_move = Team::Blue;

        // scholar endangers pony
        negaworld.orange_ponies = negaworld.orange_ponies.alight(
            Locale { rank: 0, file: 0 });
        negaworld.blue_scholars = negaworld.blue_scholars.alight(
            Locale { rank: 2, file: 2 });

        // pony endangers servant
        negaworld.orange_servants = negaworld.orange_servants.alight(
            Locale { rank: 7, file: 1 });
        negaworld.blue_ponies = negaworld.blue_ponies.alight(
            Locale { rank: 5, file: 2 });

        // Orange has another servant sitting nowhere interesting
        negaworld.orange_servants = negaworld.orange_servants.alight(
            Locale { rank: 3, file: 6 });
        negaworld.to_move = Team::Blue;

        negaworld.no_castling_at_all();

        let negadvisory = kickoff(&negaworld, depth, true);

        // taking the pony is still the right thing to do, even in the
        // negaworld
        assert_eq!(Locale { rank: 0, file: 0 },
                   negadvisory[0].0.patch.whither);
    }

}
<|endoftext|>