<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use super::super::{ActivateError, ActivateResult, Queue, VirtioDevice};
use super::handle::*;
use super::*;

use memory_model::GuestMemory;
use sys_util::EventFd;
use vhost_backend::Vhost;
use vhost_backend::Vsock as VhostVsockFd;
use vhost_gen::vhost::*;
use virtio_gen::virtio_config::*;
use virtio_gen::virtio_ring::*;

use byteorder::{ByteOrder, LittleEndian};
use epoll;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

const QUEUE_SIZE: u16 = 256;
const NUM_QUEUES: usize = 3;
const QUEUE_SIZES: &'static [u16] = &[QUEUE_SIZE; NUM_QUEUES];

impl std::convert::From<super::Error> for ActivateError {
    fn from(error: super::Error) -> Self {
        ActivateError::BadVhostActivate(error)
    }
}

pub struct Vsock {
    vsock_fd: Option<VhostVsockFd>,
    cid: u64,
    avail_features: u64,
    acked_features: u64,
    config_space: Vec<u8>,
    epoll_config: VhostEpollConfig,
    interrupt: Option<EventFd>,
}

impl Vsock {
    /// Create a new virtio-vsock device with the given VM cid.
    pub fn new(cid: u64, mem: &GuestMemory, epoll_config: VhostEpollConfig) -> Result<Vsock> {
        let fd = VhostVsockFd::new(mem).map_err(Error::VhostOpen)?;

        let avail_features = 1 << VIRTIO_F_NOTIFY_ON_EMPTY
            | 1 << VIRTIO_RING_F_INDIRECT_DESC
            | 1 << VIRTIO_RING_F_EVENT_IDX
            | 1 << VHOST_F_LOG_ALL
            | 1 << VIRTIO_F_ANY_LAYOUT
            | 1 << VIRTIO_F_VERSION_1;

        Ok(Vsock {
            vsock_fd: Some(fd),
            cid,
            avail_features,
            acked_features: 0,
            config_space: Vec::new(),
            epoll_config,
            interrupt: Some(EventFd::new().map_err(Error::VhostIrqCreate)?),
        })
    }
}

impl VirtioDevice for Vsock {
    fn device_type(&self) -> u32 {
        TYPE_VSOCK
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn features(&self, page: u32) -> u32 {
        match page {
            // Get the lower 32-bits of the features bitfield.
            0 => self.avail_features as u32,
            // Get the upper 32-bits of the features bitfield.
            1 => (self.avail_features >> 32) as u32,
            _ => {
                warn!(
                    "vsock: virtio-vsock got request for features page: {}",
                    page
                );
                0u32
            }
        }
    }

    fn ack_features(&mut self, page: u32, value: u32) {
        let mut v = match page {
            0 => value as u64,
            1 => (value as u64) << 32,
            _ => {
                warn!(
                    "vsock: virtio-vsock device cannot ack unknown feature page: {}",
                    page
                );
                0u64
            }
        };

        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("vsock: virtio-vsock got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        match offset {
            0 if data.len() == 8 => LittleEndian::write_u64(data, self.cid),
            0 if data.len() == 4 => LittleEndian::write_u32(data, (self.cid & 0xffffffff) as u32),
            4 if data.len() == 4 => {
                LittleEndian::write_u32(data, ((self.cid >> 32) & 0xffffffff) as u32)
            }
            _ => warn!(
                "vsock: virtio-vsock received invalid read request of {} bytes at offset {}",
                data.len(),
                offset
            ),
        }
    }

    fn write_config(&mut self, offset: u64, data: &[u8]) {
        let data_len = data.len() as u64;
        let config_len = self.config_space.len() as u64;
        if offset + data_len > config_len {
            error!("Failed to write config space");
            return;
        }
        let (_, right) = self.config_space.split_at_mut(offset as usize);
        right.copy_from_slice(&data[..]);
    }

    fn activate(
        &mut self,
        _: GuestMemory,
        interrupt_evt: EventFd,
        interrupt_status: Arc<AtomicUsize>,
        queues: Vec<Queue>,
        queue_evts: Vec<EventFd>,
    ) -> ActivateResult {
        if queues.len() != NUM_QUEUES || queue_evts.len() != NUM_QUEUES {
            error!(
                "Cannot perform activate. Expected {} queue(s), got {}",
                NUM_QUEUES,
                queues.len()
            );
            return Err(ActivateError::BadActivate);
        }

        if let Some(vsock_fd) = self.vsock_fd.take() {
            if let Some(interrupt) = self.interrupt.take() {
                let cid = self.cid;

                // The third vq is an event-only vq that is not handled by the vhost
                // subsystem (but still needs to exist).  Split it off here.
                let vhost_queues = queues[..2].to_vec();

                // Preliminary setup for vhost net.
                vsock_fd.set_owner().map_err(Error::VhostSetOwner)?;

                let avail_features = vsock_fd.get_features().map_err(Error::VhostGetFeatures)?;
                let features: u64 = self.acked_features & avail_features;
                vsock_fd
                    .set_features(features)
                    .map_err(Error::VhostSetFeatures)?;

                vsock_fd.set_mem_table().map_err(Error::VhostSetMemTable)?;

                for (queue_index, ref queue) in vhost_queues.iter().enumerate() {
                    vsock_fd
                        .set_vring_num(queue_index, queue.get_max_size())
                        .map_err(Error::VhostSetVringNum)?;

                    vsock_fd
                        .set_vring_addr(
                            QUEUE_SIZES[queue_index],
                            queue.actual_size(),
                            queue_index,
                            0,
                            queue.desc_table,
                            queue.used_ring,
                            queue.avail_ring,
                            None,
                        )
                        .map_err(Error::VhostSetVringAddr)?;
                    vsock_fd
                        .set_vring_base(queue_index, 0)
                        .map_err(Error::VhostSetVringBase)?;
                    vsock_fd
                        .set_vring_call(queue_index, &interrupt)
                        .map_err(Error::VhostSetVringCall)?;
                    vsock_fd
                        .set_vring_kick(queue_index, &queue_evts[queue_index])
                        .map_err(Error::VhostSetVringKick)?;
                }

                let handler =
                    VhostEpollHandler::new(vsock_fd, interrupt_status, interrupt_evt, interrupt);

                // vsock specific ioctl setup for running device.
                handler
                    .get_device()
                    .set_guest_cid(cid)
                    .map_err(Error::VhostVsockSetCid)?;
                handler
                    .get_device()
                    .start()
                    .map_err(Error::VhostVsockStart)?;

                let queue_evt_raw_fd = handler.get_queue_evt();
                //channel should be open and working
                self.epoll_config
                    .get_sender()
                    .send(Box::new(handler))
                    .unwrap();

                epoll::ctl(
                    self.epoll_config.get_raw_epoll_fd(),
                    epoll::EPOLL_CTL_ADD,
                    queue_evt_raw_fd,
                    epoll::Event::new(epoll::EPOLLIN, self.epoll_config.get_queue_evt_token()),
                )
                .map_err(ActivateError::EpollCtl)?;

                return Ok(());
            }
        }
        Err(ActivateError::BadActivate)
    }
}
<commit_msg>vhost-vsock: fix build after epoll crate update<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use super::super::{ActivateError, ActivateResult, Queue, VirtioDevice};
use super::handle::*;
use super::*;

use memory_model::GuestMemory;
use sys_util::EventFd;
use vhost_backend::Vhost;
use vhost_backend::Vsock as VhostVsockFd;
use vhost_gen::vhost::*;
use virtio_gen::virtio_config::*;
use virtio_gen::virtio_ring::*;

use byteorder::{ByteOrder, LittleEndian};
use epoll;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

const QUEUE_SIZE: u16 = 256;
const NUM_QUEUES: usize = 3;
const QUEUE_SIZES: &'static [u16] = &[QUEUE_SIZE; NUM_QUEUES];

impl std::convert::From<super::Error> for ActivateError {
    fn from(error: super::Error) -> Self {
        ActivateError::BadVhostActivate(error)
    }
}

pub struct Vsock {
    vsock_fd: Option<VhostVsockFd>,
    cid: u64,
    avail_features: u64,
    acked_features: u64,
    config_space: Vec<u8>,
    epoll_config: VhostEpollConfig,
    interrupt: Option<EventFd>,
}

impl Vsock {
    /// Create a new virtio-vsock device with the given VM cid.
    pub fn new(cid: u64, mem: &GuestMemory, epoll_config: VhostEpollConfig) -> Result<Vsock> {
        let fd = VhostVsockFd::new(mem).map_err(Error::VhostOpen)?;

        let avail_features = 1 << VIRTIO_F_NOTIFY_ON_EMPTY
            | 1 << VIRTIO_RING_F_INDIRECT_DESC
            | 1 << VIRTIO_RING_F_EVENT_IDX
            | 1 << VHOST_F_LOG_ALL
            | 1 << VIRTIO_F_ANY_LAYOUT
            | 1 << VIRTIO_F_VERSION_1;

        Ok(Vsock {
            vsock_fd: Some(fd),
            cid,
            avail_features,
            acked_features: 0,
            config_space: Vec::new(),
            epoll_config,
            interrupt: Some(EventFd::new().map_err(Error::VhostIrqCreate)?),
        })
    }
}

impl VirtioDevice for Vsock {
    fn device_type(&self) -> u32 {
        TYPE_VSOCK
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn features(&self, page: u32) -> u32 {
        match page {
            // Get the lower 32-bits of the features bitfield.
            0 => self.avail_features as u32,
            // Get the upper 32-bits of the features bitfield.
            1 => (self.avail_features >> 32) as u32,
            _ => {
                warn!(
                    "vsock: virtio-vsock got request for features page: {}",
                    page
                );
                0u32
            }
        }
    }

    fn ack_features(&mut self, page: u32, value: u32) {
        let mut v = match page {
            0 => value as u64,
            1 => (value as u64) << 32,
            _ => {
                warn!(
                    "vsock: virtio-vsock device cannot ack unknown feature page: {}",
                    page
                );
                0u64
            }
        };

        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("vsock: virtio-vsock got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        match offset {
            0 if data.len() == 8 => LittleEndian::write_u64(data, self.cid),
            0 if data.len() == 4 => LittleEndian::write_u32(data, (self.cid & 0xffffffff) as u32),
            4 if data.len() == 4 => {
                LittleEndian::write_u32(data, ((self.cid >> 32) & 0xffffffff) as u32)
            }
            _ => warn!(
                "vsock: virtio-vsock received invalid read request of {} bytes at offset {}",
                data.len(),
                offset
            ),
        }
    }

    fn write_config(&mut self, offset: u64, data: &[u8]) {
        let data_len = data.len() as u64;
        let config_len = self.config_space.len() as u64;
        if offset + data_len > config_len {
            error!("Failed to write config space");
            return;
        }
        let (_, right) = self.config_space.split_at_mut(offset as usize);
        right.copy_from_slice(&data[..]);
    }

    fn activate(
        &mut self,
        _: GuestMemory,
        interrupt_evt: EventFd,
        interrupt_status: Arc<AtomicUsize>,
        queues: Vec<Queue>,
        queue_evts: Vec<EventFd>,
    ) -> ActivateResult {
        if queues.len() != NUM_QUEUES || queue_evts.len() != NUM_QUEUES {
            error!(
                "Cannot perform activate. Expected {} queue(s), got {}",
                NUM_QUEUES,
                queues.len()
            );
            return Err(ActivateError::BadActivate);
        }

        if let Some(vsock_fd) = self.vsock_fd.take() {
            if let Some(interrupt) = self.interrupt.take() {
                let cid = self.cid;

                // The third vq is an event-only vq that is not handled by the vhost
                // subsystem (but still needs to exist).  Split it off here.
                let vhost_queues = queues[..2].to_vec();

                // Preliminary setup for vhost net.
                vsock_fd.set_owner().map_err(Error::VhostSetOwner)?;

                let avail_features = vsock_fd.get_features().map_err(Error::VhostGetFeatures)?;
                let features: u64 = self.acked_features & avail_features;
                vsock_fd
                    .set_features(features)
                    .map_err(Error::VhostSetFeatures)?;

                vsock_fd.set_mem_table().map_err(Error::VhostSetMemTable)?;

                for (queue_index, ref queue) in vhost_queues.iter().enumerate() {
                    vsock_fd
                        .set_vring_num(queue_index, queue.get_max_size())
                        .map_err(Error::VhostSetVringNum)?;

                    vsock_fd
                        .set_vring_addr(
                            QUEUE_SIZES[queue_index],
                            queue.actual_size(),
                            queue_index,
                            0,
                            queue.desc_table,
                            queue.used_ring,
                            queue.avail_ring,
                            None,
                        )
                        .map_err(Error::VhostSetVringAddr)?;
                    vsock_fd
                        .set_vring_base(queue_index, 0)
                        .map_err(Error::VhostSetVringBase)?;
                    vsock_fd
                        .set_vring_call(queue_index, &interrupt)
                        .map_err(Error::VhostSetVringCall)?;
                    vsock_fd
                        .set_vring_kick(queue_index, &queue_evts[queue_index])
                        .map_err(Error::VhostSetVringKick)?;
                }

                let handler =
                    VhostEpollHandler::new(vsock_fd, interrupt_status, interrupt_evt, interrupt);

                // vsock specific ioctl setup for running device.
                handler
                    .get_device()
                    .set_guest_cid(cid)
                    .map_err(Error::VhostVsockSetCid)?;
                handler
                    .get_device()
                    .start()
                    .map_err(Error::VhostVsockStart)?;

                let queue_evt_raw_fd = handler.get_queue_evt();
                //channel should be open and working
                self.epoll_config
                    .get_sender()
                    .send(Box::new(handler))
                    .unwrap();

                epoll::ctl(
                    self.epoll_config.get_raw_epoll_fd(),
                    epoll::ControlOptions::EPOLL_CTL_ADD,
                    queue_evt_raw_fd,
                    epoll::Event::new(
                        epoll::Events::EPOLLIN,
                        self.epoll_config.get_queue_evt_token(),
                    ),
                )
                .map_err(ActivateError::EpollCtl)?;

                return Ok(());
            }
        }
        Err(ActivateError::BadActivate)
    }
}
<|endoftext|>


<commit_before>//! Completes keywords.

use syntax::{ast, SyntaxKind};
use test_utils::mark;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    #[test]
    fn struct_initializer_field_expr() {
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        f: $0
    }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw return
            "#]],
        );
    }
}
<commit_msg>Add braces when completing ifs<commit_after>//! Completes keywords.

use syntax::{ast, SyntaxKind};
use test_utils::mark;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    #[test]
    fn struct_initializer_field_expr() {
        check(
            r#"
struct Foo {
    pub f: i32,
}
fn foo() {
    Foo {
        f: $0
    }
}
"#,
            expect![[r#"
                kw match
                kw while
                kw loop
                kw if
                kw if let
                kw return
            "#]],
        );
    }
}
<|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/.
*
*  Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
* */
use std::collections::HashMap;
use futures::sync::mpsc::UnboundedSender;
use peer::set_peer_cid;

pub struct MlesDb {
    channels: Option<HashMap<u64, UnboundedSender<Vec<u8>>>>,
    messages: Vec<Vec<u8>>,
    peer_tx: Option<UnboundedSender<UnboundedSender<Vec<u8>>>>,
    history_limit: usize,
    tx_db: Vec<UnboundedSender<Vec<u8>>>,
}

impl MlesDb {
    pub fn new(hlim: usize) -> MlesDb {
        MlesDb {
            channels: None,
            messages: Vec::new(),
            peer_tx: None,
            history_limit: hlim,
            tx_db: Vec::new(),
        }
    }

    pub fn get_channels(&self) -> Option<&HashMap<u64, UnboundedSender<Vec<u8>>>> {
        self.channels.as_ref()
    }

    pub fn get_messages(&self) -> &Vec<Vec<u8>> {
        &self.messages
    }

    pub fn get_messages_len(&self) -> usize {
        self.messages.len()
    }

    pub fn get_tx_db(&self) -> &Vec<UnboundedSender<Vec<u8>>> {
        &self.tx_db
    }

    pub fn clear_tx_db(&mut self) {
        self.tx_db = Vec::new();
    }

    pub fn add_tx_db(&mut self, tx: UnboundedSender<Vec<u8>>) {
        self.tx_db.push(tx);
    }

    pub fn get_history_limit(&self) -> usize {
        self.history_limit
    }

    pub fn get_peer_tx(&self) -> Option<&UnboundedSender<UnboundedSender<Vec<u8>>>> {
        self.peer_tx.as_ref()
    }

    pub fn add_channel(&mut self, cid: u64, sender: UnboundedSender<Vec<u8>>) {
        if self.channels.is_none() {
            self.channels = Some(HashMap::new());
        }
        if let Some(ref mut channels) = self.channels {
            channels.insert(cid, sender);
        }
    }

    pub fn check_for_duplicate_cid(&self, cid: u32) -> bool {
        if let Some(ref channels) = self.channels {
            if channels.contains_key(&(cid as u64)) {
                return true;
            }
            if channels.contains_key(&set_peer_cid(cid)) {
                return true;
            }
        }
        false
    }

    pub fn rem_channel(&mut self, cid: u64) {
        if let Some(ref mut channels) = self.channels {
            channels.remove(&cid);
        }
    }

    pub fn get_channels_len(&mut self) -> usize {
        if let Some(ref channels) = self.channels {
            return channels.len();
        }
        0
    }

    pub fn add_message(&mut self, message: Vec<u8>) {
        if 0 == self.get_history_limit() {
            return;
        }
        if self.messages.len() == self.get_history_limit() {
            self.messages.remove(0);
        }
        self.messages.push(message);
    }

    pub fn set_peer_tx(&mut self, peer_tx: UnboundedSender<UnboundedSender<Vec<u8>>>) {
        self.peer_tx = Some(peer_tx);
    }
}

pub struct MlesPeerDb {
    channels: Vec<UnboundedSender<Vec<u8>>>,
    messages: Vec<Vec<u8>>,
    history_limit: usize
}

impl MlesPeerDb {
    pub fn new(hlim: usize) -> MlesPeerDb {
        MlesPeerDb {
            channels: Vec::new(),
            messages: Vec::new(),
            history_limit: hlim,
        }
    }

    pub fn get_channels(&self) -> &Vec<UnboundedSender<Vec<u8>>> {
        &self.channels
    }

    pub fn add_channel(&mut self, channel: UnboundedSender<Vec<u8>>) {
        self.channels.push(channel);
    }

    pub fn clear_channels(&mut self) {
        self.channels = Vec::new();
    }

    pub fn add_message(&mut self, message: Vec<u8>) {
        if 0 == self.get_history_limit() {
            return;
        }
        if self.messages.len() == self.get_history_limit() {
            self.messages.remove(0);
        }
        self.messages.push(message);
    }

    pub fn get_messages(&self) -> &Vec<Vec<u8>> {
        &self.messages
    }

    pub fn get_history_limit(&mut self) -> usize {
        self.history_limit
    }

    pub fn get_messages_len(&mut self) -> usize {
        self.messages.len()
    }

}

#[cfg(test)]

mod tests {
    use super::*;
    const HISTLIMIT: usize = 100;

    #[test]
    fn test_new_db() {
        let msg = "Message".to_string().into_bytes();
        let mut mles_db = MlesDb::new(HISTLIMIT);
        mles_db.add_message(msg);
        assert_eq!(1, mles_db.get_messages().len());
        assert_eq!(0, mles_db.get_channels_len());
        let channel = mles_db.get_channels();
        assert_eq!(true, channel.is_none());
    }

    #[test]
    fn test_new_peer_db() {
        let msg = "Message".to_string().into_bytes();
        let mut mles_peer = MlesPeerDb::new(HISTLIMIT);
        assert_eq!(0, mles_peer.get_messages_len());
        mles_peer.add_message(msg);
        assert_eq!(1, mles_peer.get_messages_len());
    }

    #[test]
    fn test_db_history_limit() {
        let mut limit = HISTLIMIT;
        let msg = "Message".to_string().into_bytes();
        let mut mles_peer = MlesPeerDb::new(limit);
        assert_eq!(0, mles_peer.get_messages_len());
        mles_peer.add_message(msg.clone());
        assert_eq!(1, mles_peer.get_messages_len());
        while limit > 0 {
            mles_peer.add_message(msg.clone());
            limit -= 1;
        }
        assert_eq!(HISTLIMIT, mles_peer.get_messages_len());
        assert_eq!(HISTLIMIT, mles_peer.get_history_limit());
    }
}
<commit_msg>Remove peer cid checks on duplicate identification<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/.
*
*  Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
* */
use std::collections::HashMap;
use futures::sync::mpsc::UnboundedSender;
use peer::set_peer_cid;

pub struct MlesDb {
    channels: Option<HashMap<u64, UnboundedSender<Vec<u8>>>>,
    messages: Vec<Vec<u8>>,
    peer_tx: Option<UnboundedSender<UnboundedSender<Vec<u8>>>>,
    history_limit: usize,
    tx_db: Vec<UnboundedSender<Vec<u8>>>,
}

impl MlesDb {
    pub fn new(hlim: usize) -> MlesDb {
        MlesDb {
            channels: None,
            messages: Vec::new(),
            peer_tx: None,
            history_limit: hlim,
            tx_db: Vec::new(),
        }
    }

    pub fn get_channels(&self) -> Option<&HashMap<u64, UnboundedSender<Vec<u8>>>> {
        self.channels.as_ref()
    }

    pub fn get_messages(&self) -> &Vec<Vec<u8>> {
        &self.messages
    }

    pub fn get_messages_len(&self) -> usize {
        self.messages.len()
    }

    pub fn get_tx_db(&self) -> &Vec<UnboundedSender<Vec<u8>>> {
        &self.tx_db
    }

    pub fn clear_tx_db(&mut self) {
        self.tx_db = Vec::new();
    }

    pub fn add_tx_db(&mut self, tx: UnboundedSender<Vec<u8>>) {
        self.tx_db.push(tx);
    }

    pub fn get_history_limit(&self) -> usize {
        self.history_limit
    }

    pub fn get_peer_tx(&self) -> Option<&UnboundedSender<UnboundedSender<Vec<u8>>>> {
        self.peer_tx.as_ref()
    }

    pub fn add_channel(&mut self, cid: u64, sender: UnboundedSender<Vec<u8>>) {
        if self.channels.is_none() {
            self.channels = Some(HashMap::new());
        }
        if let Some(ref mut channels) = self.channels {
            channels.insert(cid, sender);
        }
    }

    pub fn check_for_duplicate_cid(&self, cid: u32) -> bool {
        if let Some(ref channels) = self.channels {
            if channels.contains_key(&(cid as u64)) {
                return true;
            }
        }
        false
    }

    pub fn rem_channel(&mut self, cid: u64) {
        if let Some(ref mut channels) = self.channels {
            channels.remove(&cid);
        }
    }

    pub fn get_channels_len(&mut self) -> usize {
        if let Some(ref channels) = self.channels {
            return channels.len();
        }
        0
    }

    pub fn add_message(&mut self, message: Vec<u8>) {
        if 0 == self.get_history_limit() {
            return;
        }
        if self.messages.len() == self.get_history_limit() {
            self.messages.remove(0);
        }
        self.messages.push(message);
    }

    pub fn set_peer_tx(&mut self, peer_tx: UnboundedSender<UnboundedSender<Vec<u8>>>) {
        self.peer_tx = Some(peer_tx);
    }
}

pub struct MlesPeerDb {
    channels: Vec<UnboundedSender<Vec<u8>>>,
    messages: Vec<Vec<u8>>,
    history_limit: usize
}

impl MlesPeerDb {
    pub fn new(hlim: usize) -> MlesPeerDb {
        MlesPeerDb {
            channels: Vec::new(),
            messages: Vec::new(),
            history_limit: hlim,
        }
    }

    pub fn get_channels(&self) -> &Vec<UnboundedSender<Vec<u8>>> {
        &self.channels
    }

    pub fn add_channel(&mut self, channel: UnboundedSender<Vec<u8>>) {
        self.channels.push(channel);
    }

    pub fn clear_channels(&mut self) {
        self.channels = Vec::new();
    }

    pub fn add_message(&mut self, message: Vec<u8>) {
        if 0 == self.get_history_limit() {
            return;
        }
        if self.messages.len() == self.get_history_limit() {
            self.messages.remove(0);
        }
        self.messages.push(message);
    }

    pub fn get_messages(&self) -> &Vec<Vec<u8>> {
        &self.messages
    }

    pub fn get_history_limit(&mut self) -> usize {
        self.history_limit
    }

    pub fn get_messages_len(&mut self) -> usize {
        self.messages.len()
    }

}

#[cfg(test)]

mod tests {
    use super::*;
    const HISTLIMIT: usize = 100;

    #[test]
    fn test_new_db() {
        let msg = "Message".to_string().into_bytes();
        let mut mles_db = MlesDb::new(HISTLIMIT);
        mles_db.add_message(msg);
        assert_eq!(1, mles_db.get_messages().len());
        assert_eq!(0, mles_db.get_channels_len());
        let channel = mles_db.get_channels();
        assert_eq!(true, channel.is_none());
    }

    #[test]
    fn test_new_peer_db() {
        let msg = "Message".to_string().into_bytes();
        let mut mles_peer = MlesPeerDb::new(HISTLIMIT);
        assert_eq!(0, mles_peer.get_messages_len());
        mles_peer.add_message(msg);
        assert_eq!(1, mles_peer.get_messages_len());
    }

    #[test]
    fn test_db_history_limit() {
        let mut limit = HISTLIMIT;
        let msg = "Message".to_string().into_bytes();
        let mut mles_peer = MlesPeerDb::new(limit);
        assert_eq!(0, mles_peer.get_messages_len());
        mles_peer.add_message(msg.clone());
        assert_eq!(1, mles_peer.get_messages_len());
        while limit > 0 {
            mles_peer.add_message(msg.clone());
            limit -= 1;
        }
        assert_eq!(HISTLIMIT, mles_peer.get_messages_len());
        assert_eq!(HISTLIMIT, mles_peer.get_history_limit());
    }
}
<|endoftext|>


<commit_before>use card_strategy::CardStrategy;

use card::Card;
use game_status::GameStatus;
use player::PlayerName;

use std::collections::BTreeMap;

#[derive(Debug)]
pub struct DefensiveCardStrategy;

/*
    Play 2 of clubs if in hand.
    If going to win a deal, do so with highest ranking card.
*/
impl DefensiveCardStrategy {
    fn score_card<'a>(card: &'a Card, game_status: &'a GameStatus, player_name: &PlayerName) -> (i32, i32, i32, i32) {
        let card_penalty_to_me = Self::card_penalty_to_me(card, game_status, player_name);
        let card_penalty = 0 - Self::card_penalty(card, game_status);
        let trouble = Self::trouble_score(card, game_status);
        let card_rank = 0 - (u32::from(card.rank) as i32);
        (card_penalty_to_me, card_penalty, trouble, card_rank)
    }

    fn card_penalty(card: &Card, game_status: &GameStatus) -> i32 {
        game_status.round_parameters.card_points.get(card)
            .map(|penalty| *penalty)
            .unwrap_or_default()
    }

    fn card_penalty_to_me(card: &Card, game_status: &GameStatus, player_name: &PlayerName) -> i32 {
        if Self::will_win_deal(card, game_status) {
            Self::card_penalty(card, game_status)
        } else {
            0
        }
    }

    fn trouble_score(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::will_win_deal(card, game_status) {
            u32::from(card.rank) as i32
        } else {
            0
        }
    }

    fn will_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        game_status.my_in_progress_deal.as_ref().and_then(|deal|
            deal.suit.and_then(|suit|
                deal.deal_cards.iter()
                    .map(|deal_card| deal_card.card)
                    .filter(|card| card.suit == suit)
                    .max()
                    .map(|winning_card| card.suit == suit && card.rank > winning_card.rank))).unwrap_or(true)
    }
}

impl CardStrategy for DefensiveCardStrategy {

    fn pass_cards<'a>(&mut self, game_status: &'a GameStatus) -> Vec<&'a Card> {
        // Need to order cards by potential to win lowest penalty hands
        let mut initial_hand: Vec<&'a Card> = game_status.my_initial_hand.iter().collect();
        initial_hand.reverse();
        initial_hand.into_iter().take(3).collect()
    }

    fn play_card<'a>(&mut self, game_status: &'a GameStatus, player_name: &PlayerName) -> &'a Card {
        let current_suit = game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit);
        let mut valid_cards: Vec<&'a Card> = game_status.my_current_hand.iter().filter(|card| Some(card.suit) == current_suit).collect();

        if valid_cards.is_empty() {
            valid_cards.extend(&game_status.my_current_hand);
        }

        valid_cards.into_iter()
            .map(|card| ((Self::score_card(card, game_status, player_name), card), card))
            .collect::<BTreeMap<((i32, i32, i32, i32), &Card), &Card>>()
            .values()
            .next()
            .expect("No valid cards to play!")
    }

}<commit_msg>cleanup and notes<commit_after>use card_strategy::CardStrategy;

use card::Card;
use game_status::GameStatus;
use player::PlayerName;

use std::collections::BTreeMap;

#[derive(Debug)]
pub struct DefensiveCardStrategy;

/*
    Play 2 of clubs if in hand.
    If going to win a deal, do so with highest ranking card.
    Need strategy for winning negetive penalty hands
*/
impl DefensiveCardStrategy {
    fn score_card<'a>(card: &'a Card, game_status: &'a GameStatus, player_name: &PlayerName) -> (i32, i32, i32, i32) {
        let card_penalty_to_me = Self::card_penalty_to_me(card, game_status, player_name);
        let card_penalty = 0 - Self::card_penalty(card, game_status);
        let trouble = Self::trouble_score(card, game_status);
        let card_rank = 0 - (u32::from(card.rank) as i32);
        (card_penalty_to_me, card_penalty, trouble, card_rank)
    }

    fn card_penalty(card: &Card, game_status: &GameStatus) -> i32 {
        game_status.round_parameters.card_points.get(card)
            .map(|penalty| *penalty)
            .unwrap_or_default()
    }

    fn card_penalty_to_me(card: &Card, game_status: &GameStatus, player_name: &PlayerName) -> i32 {
        if Self::will_win_deal(card, game_status) {
            Self::card_penalty(card, game_status)
        } else {
            0
        }
    }

    fn trouble_score(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::will_win_deal(card, game_status) {
            u32::from(card.rank) as i32
        } else {
            0
        }
    }

    fn will_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        game_status.my_in_progress_deal.as_ref().and_then(|deal|
            deal.suit.and_then(|suit|
                deal.deal_cards.iter()
                    .map(|deal_card| deal_card.card)
                    .filter(|card| card.suit == suit)
                    .max()
                    .map(|winning_card| card.suit == suit && card.rank > winning_card.rank))).unwrap_or(true)
    }
}

impl CardStrategy for DefensiveCardStrategy {

    fn pass_cards<'a>(&mut self, game_status: &'a GameStatus) -> Vec<&'a Card> {
        // Need to order cards by potential to win lowest penalty hands
        let mut initial_hand: Vec<&Card> = game_status.my_initial_hand.iter().collect();
        initial_hand.reverse();
        initial_hand.into_iter().take(3).collect()
    }

    fn play_card<'a>(&mut self, game_status: &'a GameStatus, player_name: &PlayerName) -> &'a Card {
        let current_suit = game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit);
        let mut valid_cards: Vec<&Card> = game_status.my_current_hand.iter().filter(|card| Some(card.suit) == current_suit).collect();

        if valid_cards.is_empty() {
            valid_cards.extend(&game_status.my_current_hand);
        }

        valid_cards.into_iter()
            .map(|card| ((Self::score_card(card, game_status, player_name), card), card))
            .collect::<BTreeMap<((i32, i32, i32, i32), &Card), &Card>>()
            .values()
            .next()
            .expect("No valid cards to play!")
    }

}<|endoftext|>


<commit_before>#[cfg(feature = "backend_gbm")]
pub mod gbm;
#[cfg(feature = "backend_drm")]
pub mod dumb;
pub mod dmabuf;

mod swapchain;
pub use swapchain::{Slot, Swapchain};

pub use drm_fourcc::{DrmFormat as Format, DrmFourcc as Fourcc, DrmModifier as Modifier, DrmVendor as Vendor, UnrecognizedFourcc, UnrecognizedVendor};

pub trait Buffer {
    fn width(&self) -> u32;
    fn height(&self) -> u32;
    fn size(&self) -> (u32, u32) { (self.width(), self.height()) }
    fn format(&self) -> Format;
}

pub trait Allocator<B: Buffer> {
    type Error: std::error::Error;

    fn create_buffer(&mut self, width: u32, height: u32, format: Format) -> Result<B, Self::Error>;
}<commit_msg>swapchain: export Error type<commit_after>#[cfg(feature = "backend_gbm")]
pub mod gbm;
#[cfg(feature = "backend_drm")]
pub mod dumb;
pub mod dmabuf;

mod swapchain;
pub use swapchain::{Slot, Swapchain, SwapchainError};

pub use drm_fourcc::{DrmFormat as Format, DrmFourcc as Fourcc, DrmModifier as Modifier, DrmVendor as Vendor, UnrecognizedFourcc, UnrecognizedVendor};

pub trait Buffer {
    fn width(&self) -> u32;
    fn height(&self) -> u32;
    fn size(&self) -> (u32, u32) { (self.width(), self.height()) }
    fn format(&self) -> Format;
}

pub trait Allocator<B: Buffer> {
    type Error: std::error::Error;

    fn create_buffer(&mut self, width: u32, height: u32, format: Format) -> Result<B, Self::Error>;
}<|endoftext|>


<commit_before>use std::result::Result;
use std::marker::PhantomData;
use std::error::Error;
use serde::*;
use serde_json::*;

use super::super::tool::*;
use super::super::environment::*;

///
/// Represents a tool made from a function
///
pub struct FnTool<TIn: Deserialize, TOut: Serialize, TErr: Serialize> {
    function: Box<Fn(TIn, &Environment) -> Result<TOut, TErr> + Send + Sync>
}

///
/// Creates a Tool from a function that can produce an error and uses an environment
///
pub fn make_dynamic_tool<TIn: Deserialize, TOut: Serialize, TErr: Serialize, F: 'static+Send+Sync+Fn(TIn, &Environment) -> Result<TOut, TErr>>(function: F) -> FnTool<TIn, TOut, TErr> {
    FnTool { function: Box::new(function) }
}

///
/// Creates a Tool from a function that can produce an error but does not use an environment
///
pub fn make_tool<TIn: Deserialize, TOut: Serialize, TErr: Serialize, F: 'static+Send+Sync+Fn(TIn) -> Result<TOut, TErr>>(function: F) -> FnTool<TIn, TOut, TErr> {
    make_dynamic_tool(move |input, _| function(input))
}

///
/// Creates a Tool from a function that cannot produce an error and doesn't use an environment
///
pub fn make_pure_tool<TIn: Deserialize, TOut: Serialize, F: 'static+Send+Sync+Fn(TIn) -> TOut>(function: F) -> FnTool<TIn, TOut, ()> {
    make_dynamic_tool(move |input, _| Ok(function(input)))
}

impl<TIn, TOut, TErr> Tool for FnTool<TIn, TOut, TErr>
where TIn: Deserialize, TOut: Serialize, TErr: Serialize {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        // Decode
        let input_decoded = from_value::<TIn>(input);

        // Chain into the tool itself
        match input_decoded {
            Ok(input_decoded) => {
                // Successfully decoded as the tool's input format
                let result = (self.function)(input_decoded, environment);

                // Encode the error or the result
                // The encoding can go wrong, and we need to handle it slightly differently for the error or the success case
                match result {
                    Ok(res) => {
                        let encoded = to_value(res);
                        match encoded {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "JSON encode failed",
                                "description":  erm.description()
                            }])
                        }
                    },

                    Err(erm) => {
                        let encoded = to_value(erm);
                        match encoded {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "Error encode failed",
                                "description":  erm.description()
                            }])
                        }
                    }
                }
            },

            Err(input_error) => {
                // Input does not match the format expected by the tool
                Err(json![{
                    "error":        "JSON input decode failed",
                    "description":  input_error.description(),
                }])
            }
        }
    }
}

///
/// Represents a tool with Rust types
///
pub struct TypedTool<TIn: Serialize, TOut: Deserialize> {
    param1: PhantomData<TIn>,
    param2: PhantomData<TOut>,
    tool: Box<Tool>

    // tool could be &'a Tool and it would work as Box<Tool> implements Tool - except rust's deref behaviour
    // gets in the way: trying to pass a reference to a box actually passes a reference to the boxed thing, which
    // then doesn't work because the size is unknown
}

impl<TIn: Serialize, TOut: Deserialize> TypedTool<TIn, TOut> {
    ///
    /// Creates an object that can be used to invoke a tool with typed parameters instead of pure JSON
    ///
    pub fn from(tool: Box<Tool>) -> TypedTool<TIn, TOut> {
        TypedTool { param1: PhantomData, param2: PhantomData, tool: tool }
    }

    ///
    /// Invokes this tool
    ///
    pub fn invoke(&self, input: TIn, environment: &Environment) -> Result<TOut, Value> {
        let json_input = to_value(input);

        match json_input {
            Ok(json_input) => {
                let json_output = self.tool.invoke_json(json_input, environment);

                match json_output {
                    Ok(json_output) => {
                        let result = from_value::<TOut>(json_output);
                        match result {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "Result decode failed",
                                "description":  erm.description()
                            }])
                        }
                    },

                    Err(json_error) => {
                        Err(json_error)
                    }
                }
            },

            Err(erm) => {
                Err(json![{
                    "error":        "Input encode failed",
                    "description":  erm.description()
                }])
            }
        }
    }
}

///
/// Extension trait for environments that makes it possible to get a Rust typed tool immediately
///
pub trait TypedEnvironment {
    ///
    /// Retrieves a tool using a Rust interface by name
    ///
    fn get_typed_tool<TIn: Serialize, TOut: Deserialize>(&self, name: &str) -> Result<TypedTool<TIn, TOut>, RetrieveToolError>;
}

impl<TEnv> TypedEnvironment for TEnv where TEnv: Environment {
    ///
    /// Retrieves a tool using a Rust interface by name
    ///
    fn get_typed_tool<TIn: Serialize, TOut: Deserialize>(&self, name: &str) -> Result<TypedTool<TIn, TOut>, RetrieveToolError> {
        self.get_json_tool(name).map(|tool| TypedTool::from(tool))
    }
}

#[cfg(test)]
mod test {
    use super::super::super::*;
    use super::super::*;
    use std::result::Result;
    use serde_json::*;              // Rust says unused, but the json! macro will not work without this

    #[derive(Serialize, Deserialize)]
    struct TestIn {
        input: i32
    }

    #[derive(Serialize, Deserialize, PartialEq, Eq)]
    struct TestOut {
        output: i32
    }

    fn test_tool(x: TestIn) -> Result<TestOut, ()> {
        Ok(TestOut { output: x.input + 1 })
    }

    #[test]
    fn can_call_tool_via_json_interface() {
        let tool        = make_tool(test_tool);
        let environment = EmptyEnvironment::new();
        let result      = tool.invoke_json(json![{ "input": 4 }], &environment);

        assert!(result == Ok(json![{ "output": 5 }]));
    }

    #[test]
    fn can_call_tool_via_typed_interface() {
        let tool        = make_pure_tool(|x: i32| { x+1 });
        let environment = EmptyEnvironment::new();

        let typed_tool  = TypedTool::from(Box::new(tool));
        let result      = typed_tool.invoke(4, &environment);

        assert!(result == Ok(5));
    }
}
<commit_msg>Deserialize now needs a lifetime parameter<commit_after>use std::result::Result;
use std::marker::PhantomData;
use std::error::Error;
use serde::*;
use serde_json::*;

use super::super::tool::*;
use super::super::environment::*;

///
/// Represents a tool made from a function
///
pub struct FnTool<TIn: Deserialize<'static>, TOut: Serialize, TErr: Serialize> {
    function: Box<Fn(TIn, &Environment) -> Result<TOut, TErr> + Send + Sync>
}

///
/// Creates a Tool from a function that can produce an error and uses an environment
///
pub fn make_dynamic_tool<TIn: Deserialize<'static>, TOut: Serialize, TErr: Serialize, F: 'static+Send+Sync+Fn(TIn, &Environment) -> Result<TOut, TErr>>(function: F) -> FnTool<TIn, TOut, TErr> {
    FnTool { function: Box::new(function) }
}

///
/// Creates a Tool from a function that can produce an error but does not use an environment
///
pub fn make_tool<TIn: Deserialize<'static>, TOut: Serialize, TErr: Serialize, F: 'static+Send+Sync+Fn(TIn) -> Result<TOut, TErr>>(function: F) -> FnTool<TIn, TOut, TErr> {
    make_dynamic_tool(move |input, _| function(input))
}

///
/// Creates a Tool from a function that cannot produce an error and doesn't use an environment
///
pub fn make_pure_tool<TIn: Deserialize<'static>, TOut: Serialize, F: 'static+Send+Sync+Fn(TIn) -> TOut>(function: F) -> FnTool<TIn, TOut, ()> {
    make_dynamic_tool(move |input, _| Ok(function(input)))
}

impl<TIn, TOut, TErr> Tool for FnTool<TIn, TOut, TErr>
where for<'de> TIn: Deserialize<'de>, TOut: Serialize, TErr: Serialize {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        // Decode
        let input_decoded = from_value::<TIn>(input);

        // Chain into the tool itself
        match input_decoded {
            Ok(input_decoded) => {
                // Successfully decoded as the tool's input format
                let result = (self.function)(input_decoded, environment);

                // Encode the error or the result
                // The encoding can go wrong, and we need to handle it slightly differently for the error or the success case
                match result {
                    Ok(res) => {
                        let encoded = to_value(res);
                        match encoded {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "JSON encode failed",
                                "description":  erm.description()
                            }])
                        }
                    },

                    Err(erm) => {
                        let encoded = to_value(erm);
                        match encoded {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "Error encode failed",
                                "description":  erm.description()
                            }])
                        }
                    }
                }
            },

            Err(input_error) => {
                // Input does not match the format expected by the tool
                Err(json![{
                    "error":        "JSON input decode failed",
                    "description":  input_error.description(),
                }])
            }
        }
    }
}

///
/// Represents a tool with Rust types
///
pub struct TypedTool<TIn: Serialize, TOut: Deserialize<'static>> {
    param1: PhantomData<TIn>,
    param2: PhantomData<TOut>,
    tool: Box<Tool>

    // tool could be &'a Tool and it would work as Box<Tool> implements Tool - except rust's deref behaviour
    // gets in the way: trying to pass a reference to a box actually passes a reference to the boxed thing, which
    // then doesn't work because the size is unknown
}

impl<TIn: Serialize, TOut> TypedTool<TIn, TOut>
where for<'de> TOut: Deserialize<'de> {
    ///
    /// Creates an object that can be used to invoke a tool with typed parameters instead of pure JSON
    ///
    pub fn from(tool: Box<Tool>) -> TypedTool<TIn, TOut> {
        TypedTool { param1: PhantomData, param2: PhantomData, tool: tool }
    }

    ///
    /// Invokes this tool
    ///
    pub fn invoke(&self, input: TIn, environment: &Environment) -> Result<TOut, Value> {
        let json_input = to_value(input);

        match json_input {
            Ok(json_input) => {
                let json_output = self.tool.invoke_json(json_input, environment);

                match json_output {
                    Ok(json_output) => {
                        let result = from_value::<TOut>(json_output);
                        match result {
                            Ok(final_value) => Ok(final_value),
                            Err(erm)        => Err(json![{
                                "error":        "Result decode failed",
                                "description":  erm.description()
                            }])
                        }
                    },

                    Err(json_error) => {
                        Err(json_error)
                    }
                }
            },

            Err(erm) => {
                Err(json![{
                    "error":        "Input encode failed",
                    "description":  erm.description()
                }])
            }
        }
    }
}

///
/// Extension trait for environments that makes it possible to get a Rust typed tool immediately
///
pub trait TypedEnvironment {
    ///
    /// Retrieves a tool using a Rust interface by name
    ///
    fn get_typed_tool<TIn: Serialize, TOut>(&self, name: &str) -> Result<TypedTool<TIn, TOut>, RetrieveToolError>
    where for<'de> TOut: Deserialize<'de>;
}

impl<TEnv> TypedEnvironment for TEnv where TEnv: Environment {
    ///
    /// Retrieves a tool using a Rust interface by name
    ///
    fn get_typed_tool<TIn: Serialize, TOut>(&self, name: &str) -> Result<TypedTool<TIn, TOut>, RetrieveToolError>
    where for<'de> TOut: Deserialize<'de> {
        self.get_json_tool(name).map(|tool| TypedTool::from(tool))
    }
}

#[cfg(test)]
mod test {
    use super::super::super::*;
    use super::super::*;
    use std::result::Result;
    use serde_json::*;              // Rust says unused, but the json! macro will not work without this

    #[derive(Serialize, Deserialize)]
    struct TestIn {
        input: i32
    }

    #[derive(Serialize, Deserialize, PartialEq, Eq)]
    struct TestOut {
        output: i32
    }

    fn test_tool(x: TestIn) -> Result<TestOut, ()> {
        Ok(TestOut { output: x.input + 1 })
    }

    #[test]
    fn can_call_tool_via_json_interface() {
        let tool        = make_tool(test_tool);
        let environment = EmptyEnvironment::new();
        let result      = tool.invoke_json(json![{ "input": 4 }], &environment);

        assert!(result == Ok(json![{ "output": 5 }]));
    }

    #[test]
    fn can_call_tool_via_typed_interface() {
        let tool        = make_pure_tool(|x: i32| { x+1 });
        let environment = EmptyEnvironment::new();

        let typed_tool  = TypedTool::from(Box::new(tool));
        let result      = typed_tool.invoke(4, &environment);

        assert!(result == Ok(5));
    }
}
<|endoftext|>


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

/// Utility to track the rate of data transfers
use std::time::{Duration, SystemTime};

/// A rate counter tracks the number of transfers, the amount of data
/// exchanged and the rate of transfer (via a few timers) over the last
/// minute. The counter does not try to be accurate and update times
/// proactively, instead it only does so lazily. As a result, produced
/// rates are worst-case estimates.
pub struct RateCounter {
	last_min_bytes: Vec<u64>,
	last_min_times: Vec<u64>,
}

impl RateCounter {
	/// Instantiate a new rate counter
	pub fn new() -> RateCounter {
		RateCounter {
			last_min_bytes: vec![],
			last_min_times: vec![],
		}
	}

	/// Increments number of bytes transferred, updating counts and rates.
	pub fn inc(&mut self, bytes: u64) {
		let now_millis = millis_since_epoch();
		self.last_min_times.push(now_millis);
		self.last_min_bytes.push(bytes);
		while self.last_min_times.len() > 0 && self.last_min_times[0] > now_millis + 60000 {
			self.last_min_times.remove(0);
			self.last_min_bytes.remove(0);
		}
	}

	/// Number of bytes counted in the last minute
	pub fn bytes_per_min(&self) -> u64 {
		self.last_min_bytes.iter().sum()
	}

	/// Count of increases in the last minute
	pub fn count_per_min(&self) -> u64 {
		self.last_min_bytes.len() as u64
	}
}

// turns out getting the millisecs since epoch in Rust isn't as easy as it
// could be
fn millis_since_epoch() -> u64 {
	let since_epoch = SystemTime::now()
		.duration_since(SystemTime::UNIX_EPOCH)
		.unwrap_or(Duration::new(0, 0));
	since_epoch.as_secs() * 1000 + since_epoch.subsec_millis() as u64
}
<commit_msg>Time comparisons are difficult. Or not.<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// Utility to track the rate of data transfers
use std::time::{Duration, SystemTime};

/// A rate counter tracks the number of transfers, the amount of data
/// exchanged and the rate of transfer (via a few timers) over the last
/// minute. The counter does not try to be accurate and update times
/// proactively, instead it only does so lazily. As a result, produced
/// rates are worst-case estimates.
pub struct RateCounter {
	last_min_bytes: Vec<u64>,
	last_min_times: Vec<u64>,
}

impl RateCounter {
	/// Instantiate a new rate counter
	pub fn new() -> RateCounter {
		RateCounter {
			last_min_bytes: vec![],
			last_min_times: vec![],
		}
	}

	/// Increments number of bytes transferred, updating counts and rates.
	pub fn inc(&mut self, bytes: u64) {
		let now_millis = millis_since_epoch();
		self.last_min_times.push(now_millis);
		self.last_min_bytes.push(bytes);
		while self.last_min_times.len() > 0 && self.last_min_times[0] + 60000 < now_millis {
			self.last_min_times.remove(0);
			self.last_min_bytes.remove(0);
		}
	}

	/// Number of bytes counted in the last minute
	pub fn bytes_per_min(&self) -> u64 {
		self.last_min_bytes.iter().sum()
	}

	/// Count of increases in the last minute
	pub fn count_per_min(&self) -> u64 {
		self.last_min_bytes.len() as u64
	}
}

// turns out getting the millisecs since epoch in Rust isn't as easy as it
// could be
fn millis_since_epoch() -> u64 {
	let since_epoch = SystemTime::now()
		.duration_since(SystemTime::UNIX_EPOCH)
		.unwrap_or(Duration::new(0, 0));
	since_epoch.as_secs() * 1000 + since_epoch.subsec_millis() as u64
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::vec::Vec;
use autograd::variable::*;
use tensor::*;
use ::*;

thread_local! {
    pub static FUNC_TABLE: RefCell<VecDeque<FuncImpl>> = RefCell::new(VecDeque::new());
}
pub type FuncId = i32;
pub enum RootKind {
    RootVar(VarKind),
    RootFunc(Function),
}
impl RootKind {
    pub fn requires_grad(&self) -> bool {
        match *self {
            RootKind::RootVar(ref v) => v.requires_grad(),
            RootKind::RootFunc(ref f) => f.requires_grad(),
        }
    }
}

impl_func!(FuncStub);
impl FuncIntf for FuncStub {
    fn forward(&mut self, input: &mut TensorKindList) -> TensorKindList {
        unreachable!()
    }
    fn backward(&mut self, input: &mut TensorKindList) -> TensorKindList {
        unreachable!()
    }
}
pub struct FuncImpl {
    previous_functions: Vec<(RootKind, i32)>,
    saved_variables: Vec<VarId>,
    needs_input_grad: Vec<bool>,
    non_differentiable: Vec<TensorId>,
    output_ids: HashMap<VarId, usize>,
    to_save: Vec<TensorId>,
    requires_grad: bool,
    owner: RcMut<FuncIntf>,
}
impl Default for FuncImpl {
    fn default() -> Self {
        FuncImpl {
            previous_functions: Vec::new(),
            saved_variables: Vec::new(),
            needs_input_grad: Vec::new(),
            non_differentiable: Vec::new(),
            output_ids: HashMap::new(),
            to_save: Vec::new(),
            requires_grad: false,
            owner: FuncStub::new().value,
        }
    }
}

impl FuncImpl {
    fn _call_hooks(&self, input: &TensorKindList, output: &TensorKindList) {
        unimplemented!();
    }
    fn init(&mut self, intf: RcMut<FuncIntf>) {
        self.owner = intf;
    }
}

#[derive(Clone)]
pub struct Function {
    pub id: FuncId,
}

impl Default for Function {
    fn default() -> Self {
        Function { id: -1 }
    }
}

pub struct FIWrap<T> {
    value: RcMut<T>,
}
impl<T: 'static + FuncIntf> FIWrap<T> {
    pub fn new(arg: T) -> Self {
        let t = FIWrap { value: RcMutNew(arg) };
        t.value.borrow_mut().delegate().init(t.value.clone());
        t
    }
    pub fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        self.value.borrow_mut().f(input_)
    }
}

impl Function {
    pub fn new() -> Self {
        use std::usize;
        let mut id = usize::MAX;
        FUNC_TABLE.with(|f| {
                            let mut table = f.borrow_mut();
                            id = table.len();
                            table.push_back(FuncImpl::default());
                        });
        Function { id: id as i32 }
    }
    pub fn init(&self, intf: RcMut<FuncIntf>) {
        //FUNC_INTF_TABLE.with(|m| m.borrow_mut().insert(self.id, intf));
        self.access().init(intf)
    }
    pub fn previous_functions(&self) -> &Vec<(RootKind, i32)> {
        &self.access().previous_functions
    }
    pub fn output_ids(&self) -> &HashMap<VarId, usize> {
        &self.access().output_ids
    }
    pub fn requires_grad(&self) -> bool {
        self.access().requires_grad
    }
    fn access(&self) -> &mut FuncImpl {
        let vecp = FUNC_TABLE.with(|f| f.as_ptr());
        let vec = unsafe { &mut *vecp };
        &mut vec[self.id as usize]
    }
    pub fn from(id: FuncId) -> Self {
        Function { id: id }
    }
    pub fn saved_tensors(&mut self) -> TensorKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v).data())
            .collect()
    }
    pub fn saved_variables(&mut self) -> VarKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v))
            .collect()
    }
    pub fn save_for_backward(&mut self, input: &TensorKindList) {
        self.access().to_save = input.iter().map(|t| t.id()).collect();
    }
    pub fn _do_backward(&mut self,
                        grad_output: &mut TensorKindList,
                        retain_variables: bool)
                        -> TensorKindList {
        let inner = self.access();
        if inner.saved_variables.is_empty() {
            panic!("Trying to backward through the graph second \
                    time, but the buffers have already been freed. Please \
                    specify retain_variables=True when calling backward for \
                    the first time.");
        };
        let grad_input = inner.owner.borrow_mut().backward(grad_output);
        inner._call_hooks(&grad_input, grad_output);
        if !retain_variables {
            inner.saved_variables.clear();
        };
        grad_input.into()
    }
}

pub trait FuncDelegate {
    fn delegate(&mut self) -> &mut Function;
}

pub trait FuncIntf: FuncDelegate {
    fn forward(&mut self, input: &mut TensorKindList) -> TensorKindList;
    fn backward(&mut self, input: &mut TensorKindList) -> TensorKindList;
    fn save_for_backward(&mut self, input: &TensorKindList) {
        self.delegate().save_for_backward(input)
    }
    fn saved_variables(&mut self) -> VarKindList {
        self.delegate().saved_variables()
    }
    fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        let is_volatile = input_.iter().any(|v| v.is_volatile());
        {
            // do start graph stuff with f
            let f = self.delegate();
            let mut inner = f.access();
            if !is_volatile {
                inner.previous_functions = input_
                    .iter()
                    .map(|v| if let Some(grad_fn) = v.grad_fn() {
                             (RootKind::RootFunc(grad_fn), v.varid())
                         } else {
                             (RootKind::RootVar(v.clone()), v.varid())
                         })
                    .collect();
                inner.needs_input_grad = input_.iter().map(|v| v.requires_grad()).collect();
                inner.requires_grad = inner.needs_input_grad.iter().any(|v| *v);
            }
        }
        let v;
        {
            let mut input_tensors = input_.iter_mut().map(|v| v.data()).collect();
            v = self.forward(&mut input_tensors);
        }
        let f = self.delegate();
        let mut inner = f.access();
        let output = if is_volatile {
            let args = VariableArgs::build().volatile(true).done();
            v.into_iter().map(|t| VarKind::new_args(t, &args)).collect()
        } else {
            let args = VariableArgs::build()
                .creator(Some(f.clone()))
                .requires_grad(inner.requires_grad)
                .done();
            let mut output: VarKindList =
                v.into_iter().map(|t| VarKind::new_args(t, &args)).collect();
            for (i, v) in output.iter().enumerate() {
                inner.output_ids.insert(v.varid(), i);
            }
            if !inner.to_save.is_empty() {
                /* if a tensor was modified in place replace the old variable with the new one */
                let mut t2v = HashMap::new();
                for ref mut var in input_.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for ref mut var in &mut output.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for t in inner.to_save.iter() {
                    inner.saved_variables.push(t2v[t]);
                }
                inner.to_save.clear();
            };
            if !inner.non_differentiable.is_empty() {
                for ref mut var in &mut output {
                    if inner.non_differentiable.contains(&var.tid()) {
                        var.requires_nograd()
                    }
                }
                inner.non_differentiable.clear();
            };
            output
        };
        output
    }
}
<commit_msg>export needs_input_grad<commit_after>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::vec::Vec;
use autograd::variable::*;
use tensor::*;
use ::*;

thread_local! {
    pub static FUNC_TABLE: RefCell<VecDeque<FuncImpl>> = RefCell::new(VecDeque::new());
}
pub type FuncId = i32;
pub enum RootKind {
    RootVar(VarKind),
    RootFunc(Function),
}
impl RootKind {
    pub fn requires_grad(&self) -> bool {
        match *self {
            RootKind::RootVar(ref v) => v.requires_grad(),
            RootKind::RootFunc(ref f) => f.requires_grad(),
        }
    }
}

impl_func!(FuncStub);
impl FuncIntf for FuncStub {
    fn forward(&mut self, input: &mut TensorKindList) -> TensorKindList {
        unreachable!()
    }
    fn backward(&mut self, input: &mut TensorKindList) -> TensorKindList {
        unreachable!()
    }
}
pub struct FuncImpl {
    previous_functions: Vec<(RootKind, i32)>,
    saved_variables: Vec<VarId>,
    needs_input_grad: Vec<bool>,
    non_differentiable: Vec<TensorId>,
    output_ids: HashMap<VarId, usize>,
    to_save: Vec<TensorId>,
    requires_grad: bool,
    owner: RcMut<FuncIntf>,
}
impl Default for FuncImpl {
    fn default() -> Self {
        FuncImpl {
            previous_functions: Vec::new(),
            saved_variables: Vec::new(),
            needs_input_grad: Vec::new(),
            non_differentiable: Vec::new(),
            output_ids: HashMap::new(),
            to_save: Vec::new(),
            requires_grad: false,
            owner: FuncStub::new().value,
        }
    }
}

impl FuncImpl {
    fn _call_hooks(&self, input: &TensorKindList, output: &TensorKindList) {
        unimplemented!();
    }
    fn init(&mut self, intf: RcMut<FuncIntf>) {
        self.owner = intf;
    }
}

#[derive(Clone)]
pub struct Function {
    pub id: FuncId,
}

impl Default for Function {
    fn default() -> Self {
        Function { id: -1 }
    }
}

pub struct FIWrap<T> {
    value: RcMut<T>,
}
impl<T: 'static + FuncIntf> FIWrap<T> {
    pub fn new(arg: T) -> Self {
        let t = FIWrap { value: RcMutNew(arg) };
        t.value.borrow_mut().delegate().init(t.value.clone());
        t
    }
    pub fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        self.value.borrow_mut().f(input_)
    }
}

impl Function {
    pub fn new() -> Self {
        use std::usize;
        let mut id = usize::MAX;
        FUNC_TABLE.with(|f| {
                            let mut table = f.borrow_mut();
                            id = table.len();
                            table.push_back(FuncImpl::default());
                        });
        Function { id: id as i32 }
    }
    pub fn init(&self, intf: RcMut<FuncIntf>) {
        //FUNC_INTF_TABLE.with(|m| m.borrow_mut().insert(self.id, intf));
        self.access().init(intf)
    }
    pub fn previous_functions(&self) -> &Vec<(RootKind, i32)> {
        &self.access().previous_functions
    }
    pub fn output_ids(&self) -> &HashMap<VarId, usize> {
        &self.access().output_ids
    }
    pub fn requires_grad(&self) -> bool {
        self.access().requires_grad
    }
    pub fn needs_input_grad(&self) -> &Vec<bool> {
        &self.access().needs_input_grad
    }
    fn access(&self) -> &mut FuncImpl {
        let vecp = FUNC_TABLE.with(|f| f.as_ptr());
        let vec = unsafe { &mut *vecp };
        &mut vec[self.id as usize]
    }
    pub fn from(id: FuncId) -> Self {
        Function { id: id }
    }
    pub fn saved_tensors(&mut self) -> TensorKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v).data())
            .collect()
    }
    pub fn saved_variables(&mut self) -> VarKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v))
            .collect()
    }
    pub fn save_for_backward(&mut self, input: &TensorKindList) {
        self.access().to_save = input.iter().map(|t| t.id()).collect();
    }
    pub fn _do_backward(&mut self,
                        grad_output: &mut TensorKindList,
                        retain_variables: bool)
                        -> TensorKindList {
        let inner = self.access();
        if inner.saved_variables.is_empty() {
            panic!("Trying to backward through the graph second \
                    time, but the buffers have already been freed. Please \
                    specify retain_variables=True when calling backward for \
                    the first time.");
        };
        let grad_input = inner.owner.borrow_mut().backward(grad_output);
        inner._call_hooks(&grad_input, grad_output);
        if !retain_variables {
            inner.saved_variables.clear();
        };
        grad_input.into()
    }
}

pub trait FuncDelegate {
    fn delegate(&mut self) -> &mut Function;
}

pub trait FuncIntf: FuncDelegate {
    fn forward(&mut self, input: &mut TensorKindList) -> TensorKindList;
    fn backward(&mut self, input: &mut TensorKindList) -> TensorKindList;
    fn save_for_backward(&mut self, input: &TensorKindList) {
        self.delegate().save_for_backward(input)
    }
    fn saved_variables(&mut self) -> VarKindList {
        self.delegate().saved_variables()
    }
    fn needs_input_grad(&mut self) -> &Vec<bool> {
        self.delegate().needs_input_grad()
    }
    fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        let is_volatile = input_.iter().any(|v| v.is_volatile());
        {
            // do start graph stuff with f
            let f = self.delegate();
            let mut inner = f.access();
            if !is_volatile {
                inner.previous_functions = input_
                    .iter()
                    .map(|v| if let Some(grad_fn) = v.grad_fn() {
                             (RootKind::RootFunc(grad_fn), v.varid())
                         } else {
                             (RootKind::RootVar(v.clone()), v.varid())
                         })
                    .collect();
                inner.needs_input_grad = input_.iter().map(|v| v.requires_grad()).collect();
                inner.requires_grad = inner.needs_input_grad.iter().any(|v| *v);
            }
        }
        let v;
        {
            let mut input_tensors = input_.iter_mut().map(|v| v.data()).collect();
            v = self.forward(&mut input_tensors);
        }
        let f = self.delegate();
        let mut inner = f.access();
        let output = if is_volatile {
            let args = VariableArgs::build().volatile(true).done();
            v.into_iter().map(|t| VarKind::new_args(t, &args)).collect()
        } else {
            let args = VariableArgs::build()
                .creator(Some(f.clone()))
                .requires_grad(inner.requires_grad)
                .done();
            let mut output: VarKindList =
                v.into_iter().map(|t| VarKind::new_args(t, &args)).collect();
            for (i, v) in output.iter().enumerate() {
                inner.output_ids.insert(v.varid(), i);
            }
            if !inner.to_save.is_empty() {
                /* if a tensor was modified in place replace the old variable with the new one */
                let mut t2v = HashMap::new();
                for ref mut var in input_.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for ref mut var in &mut output.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for t in inner.to_save.iter() {
                    inner.saved_variables.push(t2v[t]);
                }
                inner.to_save.clear();
            };
            if !inner.non_differentiable.is_empty() {
                for ref mut var in &mut output {
                    if inner.non_differentiable.contains(&var.tid()) {
                        var.requires_nograd()
                    }
                }
                inner.non_differentiable.clear();
            };
            output
        };
        output
    }
}
<|endoftext|>