<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write(&buf) {
                    Ok(_) => {
                        let _ = stream.flush();
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                };
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<commit_msg>fix bug for network not send all msg<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<|endoftext|>


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

#![allow(bad_style)]

use std::io;
use std::mem;
use std::net::{TcpListener, TcpStream, UdpSocket};
use std::os::windows::io::FromRawSocket;
use std::sync::{Once, ONCE_INIT};

use winapi::*;
use ws2_32::*;
use kernel32::*;

const WSA_FLAG_OVERLAPPED: DWORD = 0x01;
const HANDLE_FLAG_INHERIT: DWORD = 0x00000001;

pub mod c {
    pub use winapi::*;
    pub use ws2_32::*;

    pub use winapi::SOCKADDR as sockaddr;
}

mod impls;

#[inline]
fn init() {
    static INIT: Once = ONCE_INIT;

    INIT.call_once(|| {
        //To initialize Winsock
        let _ = UdpSocket::bind("127.0.0.1:34254");
    });
}

pub struct Socket {
    socket: SOCKET,
}

impl Socket {
    pub fn new(family: c_int, ty: c_int) -> io::Result<Socket> {
        init();
        let socket = try!(unsafe {
            match WSASocketW(family, ty, 0, 0 as *mut _, 0,
                             WSA_FLAG_OVERLAPPED) {
                INVALID_SOCKET => Err(io::Error::last_os_error()),
                n => Ok(Socket { socket: n }),
            }
        });
        try!(socket.set_no_inherit());
        Ok(socket)
    }

    pub fn raw(&self) -> SOCKET { self.socket }

    fn into_socket(self) -> SOCKET {
        let socket = self.socket;
        mem::forget(self);
        socket
    }

    pub fn into_tcp_listener(self) -> TcpListener {
        unsafe { TcpListener::from_raw_socket(self.into_socket()) }
    }

    pub fn into_tcp_stream(self) -> TcpStream {
        unsafe { TcpStream::from_raw_socket(self.into_socket()) }
    }

    pub fn into_udp_socket(self) -> UdpSocket {
        unsafe { UdpSocket::from_raw_socket(self.into_socket()) }
    }

    fn set_no_inherit(&self) -> io::Result<()> {
        ::cvt_win(unsafe {
            SetHandleInformation(self.socket as HANDLE, HANDLE_FLAG_INHERIT, 0)
        }).map(|_| ())
    }
}

impl ::FromInner for Socket {
    type Inner = SOCKET;
    fn from_inner(socket: SOCKET) -> Socket {
        Socket { socket: socket }
    }
}

impl Drop for Socket {
    fn drop(&mut self) {
        unsafe {
            let _ = closesocket(self.socket);
        }
    }
}
<commit_msg>Touch up winsock comment slightly<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(bad_style)]

use std::io;
use std::mem;
use std::net::{TcpListener, TcpStream, UdpSocket};
use std::os::windows::io::FromRawSocket;
use std::sync::{Once, ONCE_INIT};

use winapi::*;
use ws2_32::*;
use kernel32::*;

const WSA_FLAG_OVERLAPPED: DWORD = 0x01;
const HANDLE_FLAG_INHERIT: DWORD = 0x00000001;

pub mod c {
    pub use winapi::*;
    pub use ws2_32::*;

    pub use winapi::SOCKADDR as sockaddr;
}

mod impls;

fn init() {
    static INIT: Once = ONCE_INIT;

    INIT.call_once(|| {
        // Initialize winsock through the standard library by just creating a
        // dummy socket. Whether this is successful or not we drop the result as
        // libstd will be sure to have initialized winsock.
        let _ = UdpSocket::bind("127.0.0.1:34254");
    });
}

pub struct Socket {
    socket: SOCKET,
}

impl Socket {
    pub fn new(family: c_int, ty: c_int) -> io::Result<Socket> {
        init();
        let socket = try!(unsafe {
            match WSASocketW(family, ty, 0, 0 as *mut _, 0,
                             WSA_FLAG_OVERLAPPED) {
                INVALID_SOCKET => Err(io::Error::last_os_error()),
                n => Ok(Socket { socket: n }),
            }
        });
        try!(socket.set_no_inherit());
        Ok(socket)
    }

    pub fn raw(&self) -> SOCKET { self.socket }

    fn into_socket(self) -> SOCKET {
        let socket = self.socket;
        mem::forget(self);
        socket
    }

    pub fn into_tcp_listener(self) -> TcpListener {
        unsafe { TcpListener::from_raw_socket(self.into_socket()) }
    }

    pub fn into_tcp_stream(self) -> TcpStream {
        unsafe { TcpStream::from_raw_socket(self.into_socket()) }
    }

    pub fn into_udp_socket(self) -> UdpSocket {
        unsafe { UdpSocket::from_raw_socket(self.into_socket()) }
    }

    fn set_no_inherit(&self) -> io::Result<()> {
        ::cvt_win(unsafe {
            SetHandleInformation(self.socket as HANDLE, HANDLE_FLAG_INHERIT, 0)
        }).map(|_| ())
    }
}

impl ::FromInner for Socket {
    type Inner = SOCKET;
    fn from_inner(socket: SOCKET) -> Socket {
        Socket { socket: socket }
    }
}

impl Drop for Socket {
    fn drop(&mut self) {
        unsafe {
            let _ = closesocket(self.socket);
        }
    }
}
<|endoftext|>


<commit_before>//! This module implements a reference search.
//! First, the element at the cursor position must be either an `ast::Name`
//! or `ast::NameRef`. If it's a `ast::NameRef`, at the classification step we
//! try to resolve the direct tree parent of this element, otherwise we
//! already have a definition and just need to get its HIR together with
//! some information that is needed for futher steps of searching.
//! After that, we collect files that might contain references and look
//! for text occurrences of the identifier. If there's an `ast::NameRef`
//! at the index that the match starts at and its tree parent is
//! resolved to the search element definition, we get a reference.

mod classify;
mod name_definition;
mod rename;
mod search_scope;

use once_cell::unsync::Lazy;
use ra_db::{SourceDatabase, SourceDatabaseExt};
use ra_prof::profile;
use ra_syntax::{algo::find_node_at_offset, ast, AstNode, SourceFile, SyntaxNode, TextUnit};

use crate::{
    db::RootDatabase, display::ToNav, FilePosition, FileRange, NavigationTarget, RangeInfo,
};

pub(crate) use self::{
    classify::{classify_name, classify_name_ref},
    name_definition::{NameDefinition, NameKind},
    rename::rename,
};

pub use self::search_scope::SearchScope;

#[derive(Debug, Clone)]
pub struct ReferenceSearchResult {
    declaration: NavigationTarget,
    references: Vec<FileRange>,
}

impl ReferenceSearchResult {
    pub fn declaration(&self) -> &NavigationTarget {
        &self.declaration
    }

    pub fn references(&self) -> &[FileRange] {
        &self.references
    }

    /// Total number of references
    /// At least 1 since all valid references should
    /// Have a declaration
    pub fn len(&self) -> usize {
        self.references.len() + 1
    }
}

// allow turning ReferenceSearchResult into an iterator
// over FileRanges
impl IntoIterator for ReferenceSearchResult {
    type Item = FileRange;
    type IntoIter = std::vec::IntoIter<FileRange>;

    fn into_iter(mut self) -> Self::IntoIter {
        let mut v = Vec::with_capacity(self.len());
        v.push(FileRange { file_id: self.declaration.file_id(), range: self.declaration.range() });
        v.append(&mut self.references);
        v.into_iter()
    }
}

pub(crate) fn find_all_refs(
    db: &RootDatabase,
    position: FilePosition,
    search_scope: Option<SearchScope>,
) -> Option<RangeInfo<ReferenceSearchResult>> {
    let parse = db.parse(position.file_id);
    let syntax = parse.tree().syntax().clone();
    let RangeInfo { range, info: (name, def) } = find_name(db, &syntax, position)?;

    let declaration = match def.kind {
        NameKind::Macro(mac) => mac.to_nav(db),
        NameKind::Field(field) => field.to_nav(db),
        NameKind::AssocItem(assoc) => assoc.to_nav(db),
        NameKind::Def(def) => NavigationTarget::from_def(db, def)?,
        NameKind::SelfType(ref ty) => match ty.as_adt() {
            Some((adt, _)) => adt.to_nav(db),
            None => return None,
        },
        NameKind::Local(local) => local.to_nav(db),
        NameKind::GenericParam(_) => return None,
    };

    let search_scope = {
        let base = def.search_scope(db);
        match search_scope {
            None => base,
            Some(scope) => base.intersection(&scope),
        }
    };

    let references = process_definition(db, def, name, search_scope);

    Some(RangeInfo::new(range, ReferenceSearchResult { declaration, references }))
}

fn find_name<'a>(
    db: &RootDatabase,
    syntax: &SyntaxNode,
    position: FilePosition,
) -> Option<RangeInfo<(String, NameDefinition)>> {
    if let Some(name) = find_node_at_offset::<ast::Name>(&syntax, position.offset) {
        let def = classify_name(db, position.file_id, &name)?;
        let range = name.syntax().text_range();
        return Some(RangeInfo::new(range, (name.text().to_string(), def)));
    }
    let name_ref = find_node_at_offset::<ast::NameRef>(&syntax, position.offset)?;
    let def = classify_name_ref(db, position.file_id, &name_ref)?;
    let range = name_ref.syntax().text_range();
    Some(RangeInfo::new(range, (name_ref.text().to_string(), def)))
}

fn process_definition(
    db: &RootDatabase,
    def: NameDefinition,
    name: String,
    scope: SearchScope,
) -> Vec<FileRange> {
    let _p = profile("process_definition");

    let pat = name.as_str();
    let mut refs = vec![];

    for (file_id, search_range) in scope {
        let text = db.file_text(file_id);
        let parse = Lazy::new(|| SourceFile::parse(&text));

        for (idx, _) in text.match_indices(pat) {
            let offset = TextUnit::from_usize(idx);

            if let Some(name_ref) =
                find_node_at_offset::<ast::NameRef>(parse.tree().syntax(), offset)
            {
                let range = name_ref.syntax().text_range();
                if let Some(search_range) = search_range {
                    if !range.is_subrange(&search_range) {
                        continue;
                    }
                }
                if let Some(d) = classify_name_ref(db, file_id, &name_ref) {
                    if d == def {
                        refs.push(FileRange { file_id, range });
                    }
                }
            }
        }
    }
    refs
}

#[cfg(test)]
mod tests {
    use crate::{
        mock_analysis::{analysis_and_position, single_file_with_position, MockAnalysis},
        ReferenceSearchResult, SearchScope,
    };

    #[test]
    fn test_find_all_refs_for_local() {
        let code = r#"
    fn main() {
        let mut i = 1;
        let j = 1;
        i = i<|> + j;

        {
            i = 0;
        }

        i = 5;
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 5);
    }

    #[test]
    fn test_find_all_refs_for_param_inside() {
        let code = r#"
    fn foo(i : u32) -> u32 {
        i<|>
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_for_fn_param() {
        let code = r#"
    fn foo(i<|> : u32) -> u32 {
        i
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_field_name() {
        let code = r#"
            //- /lib.rs
            struct Foo {
                pub spam<|>: u32,
            }

            fn main(s: Foo) {
                let f = s.spam;
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_impl_item_name() {
        let code = r#"
            //- /lib.rs
            struct Foo;
            impl Foo {
                fn f<|>(&self) {  }
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 1);
    }

    #[test]
    fn test_find_all_refs_enum_var_name() {
        let code = r#"
            //- /lib.rs
            enum Foo {
                A,
                B<|>,
                C,
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 1);
    }

    #[test]
    fn test_find_all_refs_two_modules() {
        let code = r#"
            //- /lib.rs
            pub mod foo;
            pub mod bar;

            fn f() {
                let i = foo::Foo { n: 5 };
            }

            //- /foo.rs
            use crate::bar;

            pub struct Foo {
                pub n: u32,
            }

            fn f() {
                let i = bar::Bar { n: 5 };
            }

            //- /bar.rs
            use crate::foo;

            pub struct Bar {
                pub n: u32,
            }

            fn f() {
                let i = foo::Foo<|> { n: 5 };
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);
    }

    // `mod foo;` is not in the results because `foo` is an `ast::Name`.
    // So, there are two references: the first one is a definition of the `foo` module,
    // which is the whole `foo.rs`, and the second one is in `use foo::Foo`.
    #[test]
    fn test_find_all_refs_decl_module() {
        let code = r#"
            //- /lib.rs
            mod foo<|>;

            use foo::Foo;

            fn f() {
                let i = Foo { n: 5 };
            }

            //- /foo.rs
            pub struct Foo {
                pub n: u32,
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_super_mod_vis() {
        let code = r#"
            //- /lib.rs
            mod foo;

            //- /foo.rs
            mod some;
            use some::Foo;

            fn f() {
                let i = Foo { n: 5 };
            }

            //- /foo/some.rs
            pub(super) struct Foo<|> {
                pub n: u32,
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);
    }

    #[test]
    fn test_find_all_refs_with_scope() {
        let code = r#"
            //- /lib.rs
            mod foo;
            mod bar;

            pub fn quux<|>() {}

            //- /foo.rs
            fn f() { super::quux(); }

            //- /bar.rs
            fn f() { super::quux(); }
        "#;

        let (mock, pos) = MockAnalysis::with_files_and_position(code);
        let bar = mock.id_of("/bar.rs");
        let analysis = mock.analysis();

        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);

        let refs =
            analysis.find_all_refs(pos, Some(SearchScope::single_file(bar))).unwrap().unwrap();
        assert_eq!(refs.len(), 2);
    }

    fn get_all_refs(text: &str) -> ReferenceSearchResult {
        let (analysis, position) = single_file_with_position(text);
        analysis.find_all_refs(position, None).unwrap().unwrap()
    }
}
<commit_msg>Add a test that uses classify_name<commit_after>//! This module implements a reference search.
//! First, the element at the cursor position must be either an `ast::Name`
//! or `ast::NameRef`. If it's a `ast::NameRef`, at the classification step we
//! try to resolve the direct tree parent of this element, otherwise we
//! already have a definition and just need to get its HIR together with
//! some information that is needed for futher steps of searching.
//! After that, we collect files that might contain references and look
//! for text occurrences of the identifier. If there's an `ast::NameRef`
//! at the index that the match starts at and its tree parent is
//! resolved to the search element definition, we get a reference.

mod classify;
mod name_definition;
mod rename;
mod search_scope;

use once_cell::unsync::Lazy;
use ra_db::{SourceDatabase, SourceDatabaseExt};
use ra_prof::profile;
use ra_syntax::{algo::find_node_at_offset, ast, AstNode, SourceFile, SyntaxNode, TextUnit};

use crate::{
    db::RootDatabase, display::ToNav, FilePosition, FileRange, NavigationTarget, RangeInfo,
};

pub(crate) use self::{
    classify::{classify_name, classify_name_ref},
    name_definition::{NameDefinition, NameKind},
    rename::rename,
};

pub use self::search_scope::SearchScope;

#[derive(Debug, Clone)]
pub struct ReferenceSearchResult {
    declaration: NavigationTarget,
    references: Vec<FileRange>,
}

impl ReferenceSearchResult {
    pub fn declaration(&self) -> &NavigationTarget {
        &self.declaration
    }

    pub fn references(&self) -> &[FileRange] {
        &self.references
    }

    /// Total number of references
    /// At least 1 since all valid references should
    /// Have a declaration
    pub fn len(&self) -> usize {
        self.references.len() + 1
    }
}

// allow turning ReferenceSearchResult into an iterator
// over FileRanges
impl IntoIterator for ReferenceSearchResult {
    type Item = FileRange;
    type IntoIter = std::vec::IntoIter<FileRange>;

    fn into_iter(mut self) -> Self::IntoIter {
        let mut v = Vec::with_capacity(self.len());
        v.push(FileRange { file_id: self.declaration.file_id(), range: self.declaration.range() });
        v.append(&mut self.references);
        v.into_iter()
    }
}

pub(crate) fn find_all_refs(
    db: &RootDatabase,
    position: FilePosition,
    search_scope: Option<SearchScope>,
) -> Option<RangeInfo<ReferenceSearchResult>> {
    let parse = db.parse(position.file_id);
    let syntax = parse.tree().syntax().clone();
    let RangeInfo { range, info: (name, def) } = find_name(db, &syntax, position)?;

    let declaration = match def.kind {
        NameKind::Macro(mac) => mac.to_nav(db),
        NameKind::Field(field) => field.to_nav(db),
        NameKind::AssocItem(assoc) => assoc.to_nav(db),
        NameKind::Def(def) => NavigationTarget::from_def(db, def)?,
        NameKind::SelfType(ref ty) => match ty.as_adt() {
            Some((adt, _)) => adt.to_nav(db),
            None => return None,
        },
        NameKind::Local(local) => local.to_nav(db),
        NameKind::GenericParam(_) => return None,
    };

    let search_scope = {
        let base = def.search_scope(db);
        match search_scope {
            None => base,
            Some(scope) => base.intersection(&scope),
        }
    };

    let references = process_definition(db, def, name, search_scope);

    Some(RangeInfo::new(range, ReferenceSearchResult { declaration, references }))
}

fn find_name<'a>(
    db: &RootDatabase,
    syntax: &SyntaxNode,
    position: FilePosition,
) -> Option<RangeInfo<(String, NameDefinition)>> {
    if let Some(name) = find_node_at_offset::<ast::Name>(&syntax, position.offset) {
        let def = classify_name(db, position.file_id, &name)?;
        let range = name.syntax().text_range();
        return Some(RangeInfo::new(range, (name.text().to_string(), def)));
    }
    let name_ref = find_node_at_offset::<ast::NameRef>(&syntax, position.offset)?;
    let def = classify_name_ref(db, position.file_id, &name_ref)?;
    let range = name_ref.syntax().text_range();
    Some(RangeInfo::new(range, (name_ref.text().to_string(), def)))
}

fn process_definition(
    db: &RootDatabase,
    def: NameDefinition,
    name: String,
    scope: SearchScope,
) -> Vec<FileRange> {
    let _p = profile("process_definition");

    let pat = name.as_str();
    let mut refs = vec![];

    for (file_id, search_range) in scope {
        let text = db.file_text(file_id);
        let parse = Lazy::new(|| SourceFile::parse(&text));

        for (idx, _) in text.match_indices(pat) {
            let offset = TextUnit::from_usize(idx);

            if let Some(name_ref) =
                find_node_at_offset::<ast::NameRef>(parse.tree().syntax(), offset)
            {
                let range = name_ref.syntax().text_range();
                if let Some(search_range) = search_range {
                    if !range.is_subrange(&search_range) {
                        continue;
                    }
                }
                if let Some(d) = classify_name_ref(db, file_id, &name_ref) {
                    if d == def {
                        refs.push(FileRange { file_id, range });
                    }
                }
            }
        }
    }
    refs
}

#[cfg(test)]
mod tests {
    use crate::{
        mock_analysis::{analysis_and_position, single_file_with_position, MockAnalysis},
        ReferenceSearchResult, SearchScope,
    };

    #[test]
    fn test_find_all_refs_for_local() {
        let code = r#"
    fn main() {
        let mut i = 1;
        let j = 1;
        i = i<|> + j;

        {
            i = 0;
        }

        i = 5;
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 5);
    }

    #[test]
    fn test_find_all_refs_for_param_inside() {
        let code = r#"
    fn foo(i : u32) -> u32 {
        i<|>
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_for_fn_param() {
        let code = r#"
    fn foo(i<|> : u32) -> u32 {
        i
    }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_field_name() {
        let code = r#"
            //- /lib.rs
            struct Foo {
                pub spam<|>: u32,
            }

            fn main(s: Foo) {
                let f = s.spam;
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_impl_item_name() {
        let code = r#"
            //- /lib.rs
            struct Foo;
            impl Foo {
                fn f<|>(&self) {  }
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 1);
    }

    #[test]
    fn test_find_all_refs_enum_var_name() {
        let code = r#"
            //- /lib.rs
            enum Foo {
                A,
                B<|>,
                C,
            }
        "#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 1);
    }

    #[test]
    fn test_find_all_refs_two_modules() {
        let code = r#"
            //- /lib.rs
            pub mod foo;
            pub mod bar;

            fn f() {
                let i = foo::Foo { n: 5 };
            }

            //- /foo.rs
            use crate::bar;

            pub struct Foo {
                pub n: u32,
            }

            fn f() {
                let i = bar::Bar { n: 5 };
            }

            //- /bar.rs
            use crate::foo;

            pub struct Bar {
                pub n: u32,
            }

            fn f() {
                let i = foo::Foo<|> { n: 5 };
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);
    }

    // `mod foo;` is not in the results because `foo` is an `ast::Name`.
    // So, there are two references: the first one is a definition of the `foo` module,
    // which is the whole `foo.rs`, and the second one is in `use foo::Foo`.
    #[test]
    fn test_find_all_refs_decl_module() {
        let code = r#"
            //- /lib.rs
            mod foo<|>;

            use foo::Foo;

            fn f() {
                let i = Foo { n: 5 };
            }

            //- /foo.rs
            pub struct Foo {
                pub n: u32,
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_super_mod_vis() {
        let code = r#"
            //- /lib.rs
            mod foo;

            //- /foo.rs
            mod some;
            use some::Foo;

            fn f() {
                let i = Foo { n: 5 };
            }

            //- /foo/some.rs
            pub(super) struct Foo<|> {
                pub n: u32,
            }
        "#;

        let (analysis, pos) = analysis_and_position(code);
        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);
    }

    #[test]
    fn test_find_all_refs_with_scope() {
        let code = r#"
            //- /lib.rs
            mod foo;
            mod bar;

            pub fn quux<|>() {}

            //- /foo.rs
            fn f() { super::quux(); }

            //- /bar.rs
            fn f() { super::quux(); }
        "#;

        let (mock, pos) = MockAnalysis::with_files_and_position(code);
        let bar = mock.id_of("/bar.rs");
        let analysis = mock.analysis();

        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
        assert_eq!(refs.len(), 3);

        let refs =
            analysis.find_all_refs(pos, Some(SearchScope::single_file(bar))).unwrap().unwrap();
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_all_refs_macro_def() {
        let code = r#"
        #[macro_export]
        macro_rules! m1<|> { () => (()) }

        fn foo() {
            m1();
            m1();
        }"#;

        let refs = get_all_refs(code);
        assert_eq!(refs.len(), 3);
    }

    fn get_all_refs(text: &str) -> ReferenceSearchResult {
        let (analysis, position) = single_file_with_position(text);
        analysis.find_all_refs(position, None).unwrap().unwrap()
    }
}
<|endoftext|>


<commit_before>/*!
This crate provides a robust regular expression parser.

This crate defines two primary types:

* [`Ast`](ast/enum.Ast.html) is the abstract syntax of a regular expression.
  An abstract syntax corresponds to a *structured representation* of the
  concrete syntax of a regular expression, where the concrete syntax is the
  pattern string itself (e.g., `foo(bar)+`). Given some abstract syntax, it
  can be converted back to the original concrete syntax (modulo some details,
  like whitespace). To a first approximation, the abstract syntax is complex
  and difficult to analyze.
* [`Hir`](hir/struct.Hir.html) is the high-level intermediate representation
  ("HIR" or "high-level IR" for short) of regular expression. It corresponds to
  an intermediate state of a regular expression that sits between the abstract
  syntax and the low level compiled opcodes that are eventually responsible for
  executing a regular expression search. Given some high-level IR, it is not
  possible to produce the original concrete syntax (although it is possible to
  produce an equivalent concrete syntax, but it will likely scarcely resemble
  the original pattern). To a first approximation, the high-level IR is simple
  and easy to analyze.

These two types come with conversion routines:

* An [`ast::parse::Parser`](ast/parse/struct.Parser.html) converts concrete
  syntax (a `&str`) to an [`Ast`](ast/enum.Ast.html).
* A [`hir::translate::Translator`](hir/translate/struct.Translator.html)
  converts an [`Ast`](ast/enum.Ast.html) to a [`Hir`](hir/struct.Hir.html).

As a convenience, the above two conversion routines are combined into one via
the top-level [`Parser`](struct.Parser.html) type. This `Parser` will first
convert your pattern to an `Ast` and then convert the `Ast` to an `Hir`.

# Example

This example shows how to parse a pattern string into its HIR:

```
use regex_syntax::Parser;
use regex_syntax::hir::{self, Hir};

let hir = Parser::new().parse("a|b").unwrap();
assert_eq!(hir, Hir::alternation(vec![
    Hir::literal(hir::Literal::Unicode('a')),
    Hir::literal(hir::Literal::Unicode('b')),
]));
```

# Concrete syntax supported

The concrete syntax is documented as part of the public API of the
[`regex` crate](https://docs.rs/regex/%2A/regex/#syntax).

# Input safety

A key feature of this library is that it is safe to use with end user facing
input. This plays a significant role in the internal implementation. In
particular:

1. Parsers provide a `nest_limit` option that permits callers to control how
   deeply nested a regular expression is allowed to be. This makes it possible
   to do case analysis over an `Ast` or an `Hir` using recursion without
   worrying about stack overflow.
2. Since relying on a particular stack size is brittle, this crate goes to
   great lengths to ensure that all interactions with both the `Ast` and the
   `Hir` do not use recursion. Namely, they use constant stack space and heap
   space proportional to the size of the original pattern string (in bytes).
   This includes the type's corresponding destructors. (One exception to this
   is literal extraction, but this will eventually get fixed.)

# Error reporting

The `Display` implementations on all `Error` types exposed in this library
provide nice human readable errors that are suitable for showing to end users
in a monospace font.

# Literal extraction

This crate provides limited support for
[literal extraction from `Hir` values](hir/literal/struct.Literals.html).
Be warned that literal extraction currently uses recursion, and therefore,
stack size proportional to the size of the `Hir`.

The purpose of literal extraction is to speed up searches. That is, if you
know a regular expression must match a prefix or suffix literal, then it is
often quicker to search for instances of that literal, and then confirm or deny
the match using the full regular expression engine. These optimizations are
done automatically in the `regex` crate.

# Crate features

An important feature provided by this crate is its Unicode support. This
includes things like case folding, boolean properties, general categories,
scripts and Unicode-aware support for the Perl classes `\w`, `\s` and `\d`.
However, a downside of this support is that it requires bundling several
Unicode data tables that are substantial in size.

A fair number of use cases do not require full Unicode support. For this
reason, this crate exposes a number of features to control which Unicode
data is available.

If a regular expression attempts to use a Unicode feature that is not available
because the corresponding crate feature was disabled, then translating that
regular expression to an `Hir` will return an error. (It is still possible
construct an `Ast` for such a regular expression, since Unicode data is not
used until translation to an `Hir`.) Stated differently, enabling or disabling
any of the features below can only add or subtract from the total set of valid
regular expressions. Enabling or disabling a feature will never modify the
match semantics of a regular expression.

The following features are available:

* **unicode** -
  Enables all Unicode features. This feature is enabled by default, and will
  always cover all Unicode features, even if more are added in the future.
* **unicode-age** -
  Provide the data for the
  [Unicode `Age` property](https://www.unicode.org/reports/tr44/tr44-24.html#Character_Age).
  This makes it possible to use classes like `\p{Age:6.0}` to refer to all
  codepoints first introduced in Unicode 6.0
* **unicode-bool** -
  Provide the data for numerous Unicode boolean properties. The full list
  is not included here, but contains properties like `Alphabetic`, `Emoji`,
  `Lowercase`, `Math`, `Uppercase` and `White_Space`.
* **unicode-case** -
  Provide the data for case insensitive matching using
  [Unicode's "simple loose matches" specification](https://www.unicode.org/reports/tr18/#Simple_Loose_Matches).
* **unicode-gencat** -
  Provide the data for
  [Uncode general categories](https://www.unicode.org/reports/tr44/tr44-24.html#General_Category_Values).
  This includes, but is not limited to, `Decimal_Number`, `Letter`,
  `Math_Symbol`, `Number` and `Punctuation`.
* **unicode-perl** -
  Provide the data for supporting the Unicode-aware Perl character classes,
  corresponding to `\w`, `\s` and `\d`. This is also necessary for using
  Unicode-aware word boundary assertions. Note that if this feature is
  disabled, the `\s` and `\d` character classes are still available if the
  `unicode-bool` and `unicode-gencat` features are enabled, respectively.
* **unicode-script** -
  Provide the data for
  [Unicode scripts and script extensions](https://www.unicode.org/reports/tr24/).
  This includes, but is not limited to, `Arabic`, `Cyrillic`, `Hebrew`,
  `Latin` and `Thai`.
* **unicode-segment** -
  Provide the data necessary to provide the properties used to implement the
  [Unicode text segmentation algorithms](https://www.unicode.org/reports/tr29/).
  This enables using classes like `\p{gcb=Extend}`, `\p{wb=Katakana}` and
  `\p{sb=ATerm}`.
*/

#![deny(missing_docs)]
#![forbid(unsafe_code)]

pub use error::{Error, Result};
pub use parser::{Parser, ParserBuilder};
pub use unicode::UnicodeWordError;

pub mod ast;
mod either;
mod error;
pub mod hir;
mod parser;
mod unicode;
mod unicode_tables;
pub mod utf8;

/// Escapes all regular expression meta characters in `text`.
///
/// The string returned may be safely used as a literal in a regular
/// expression.
pub fn escape(text: &str) -> String {
    let mut quoted = String::new();
    escape_into(text, &mut quoted);
    quoted
}

/// Escapes all meta characters in `text` and writes the result into `buf`.
///
/// This will append escape characters into the given buffer. The characters
/// that are appended are safe to use as a literal in a regular expression.
pub fn escape_into(text: &str, buf: &mut String) {
    buf.reserve(text.len());
    for c in text.chars() {
        if is_meta_character(c) {
            buf.push('\\');
        }
        buf.push(c);
    }
}

/// Returns true if the give character has significance in a regex.
///
/// These are the only characters that are allowed to be escaped, with one
/// exception: an ASCII space character may be escaped when extended mode (with
/// the `x` flag) is enabld. In particular, `is_meta_character(' ')` returns
/// `false`.
///
/// Note that the set of characters for which this function returns `true` or
/// `false` is fixed and won't change in a semver compatible release.
pub fn is_meta_character(c: char) -> bool {
    match c {
        '\\' | '.' | '+' | '*' | '?' | '(' | ')' | '|' | '[' | ']' | '{'
        | '}' | '^' | '$' | '#' | '&' | '-' | '~' => true,
        _ => false,
    }
}

/// Returns true if and only if the given character is a Unicode word
/// character.
///
/// A Unicode word character is defined by
/// [UTS#18 Annex C](http://unicode.org/reports/tr18/#Compatibility_Properties).
/// In particular, a character
/// is considered a word character if it is in either of the `Alphabetic` or
/// `Join_Control` properties, or is in one of the `Decimal_Number`, `Mark`
/// or `Connector_Punctuation` general categories.
///
/// # Panics
///
/// If the `unicode-perl` feature is not enabled, then this function panics.
/// For this reason, it is recommended that callers use
/// [`try_is_word_character`](fn.try_is_word_character.html)
/// instead.
pub fn is_word_character(c: char) -> bool {
    try_is_word_character(c).expect("unicode-perl feature must be enabled")
}

/// Returns true if and only if the given character is a Unicode word
/// character.
///
/// A Unicode word character is defined by
/// [UTS#18 Annex C](http://unicode.org/reports/tr18/#Compatibility_Properties).
/// In particular, a character
/// is considered a word character if it is in either of the `Alphabetic` or
/// `Join_Control` properties, or is in one of the `Decimal_Number`, `Mark`
/// or `Connector_Punctuation` general categories.
///
/// # Errors
///
/// If the `unicode-perl` feature is not enabled, then this function always
/// returns an error.
pub fn try_is_word_character(
    c: char,
) -> std::result::Result<bool, UnicodeWordError> {
    unicode::is_word_character(c)
}

/// Returns true if and only if the given character is an ASCII word character.
///
/// An ASCII word character is defined by the following character class:
/// `[_0-9a-zA-Z]'.
pub fn is_word_byte(c: u8) -> bool {
    match c {
        b'_' | b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' => true,
        _ => false,
    }
}

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

    #[test]
    fn escape_meta() {
        assert_eq!(
            escape(r"\.+*?()|[]{}^$#&-~"),
            r"\\\.\+\*\?\(\)\|\[\]\{\}\^\$\#\&\-\~".to_string()
        );
    }

    #[test]
    fn word_byte() {
        assert!(is_word_byte(b'a'));
        assert!(!is_word_byte(b'-'));
    }

    #[test]
    #[cfg(feature = "unicode-perl")]
    fn word_char() {
        assert!(is_word_character('a'), "ASCII");
        assert!(is_word_character('à'), "Latin-1");
        assert!(is_word_character('β'), "Greek");
        assert!(is_word_character('\u{11011}'), "Brahmi (Unicode 6.0)");
        assert!(is_word_character('\u{11611}'), "Modi (Unicode 7.0)");
        assert!(is_word_character('\u{11711}'), "Ahom (Unicode 8.0)");
        assert!(is_word_character('\u{17828}'), "Tangut (Unicode 9.0)");
        assert!(is_word_character('\u{1B1B1}'), "Nushu (Unicode 10.0)");
        assert!(is_word_character('\u{16E40}'), "Medefaidrin (Unicode 11.0)");
        assert!(!is_word_character('-'));
        assert!(!is_word_character('☃'));
    }

    #[test]
    #[should_panic]
    #[cfg(not(feature = "unicode-perl"))]
    fn word_char_disabled_panic() {
        assert!(is_word_character('a'));
    }

    #[test]
    #[cfg(not(feature = "unicode-perl"))]
    fn word_char_disabled_error() {
        assert!(try_is_word_character('a').is_err());
    }
}
<commit_msg>syntax/doc: fix enabld -> enabled<commit_after>/*!
This crate provides a robust regular expression parser.

This crate defines two primary types:

* [`Ast`](ast/enum.Ast.html) is the abstract syntax of a regular expression.
  An abstract syntax corresponds to a *structured representation* of the
  concrete syntax of a regular expression, where the concrete syntax is the
  pattern string itself (e.g., `foo(bar)+`). Given some abstract syntax, it
  can be converted back to the original concrete syntax (modulo some details,
  like whitespace). To a first approximation, the abstract syntax is complex
  and difficult to analyze.
* [`Hir`](hir/struct.Hir.html) is the high-level intermediate representation
  ("HIR" or "high-level IR" for short) of regular expression. It corresponds to
  an intermediate state of a regular expression that sits between the abstract
  syntax and the low level compiled opcodes that are eventually responsible for
  executing a regular expression search. Given some high-level IR, it is not
  possible to produce the original concrete syntax (although it is possible to
  produce an equivalent concrete syntax, but it will likely scarcely resemble
  the original pattern). To a first approximation, the high-level IR is simple
  and easy to analyze.

These two types come with conversion routines:

* An [`ast::parse::Parser`](ast/parse/struct.Parser.html) converts concrete
  syntax (a `&str`) to an [`Ast`](ast/enum.Ast.html).
* A [`hir::translate::Translator`](hir/translate/struct.Translator.html)
  converts an [`Ast`](ast/enum.Ast.html) to a [`Hir`](hir/struct.Hir.html).

As a convenience, the above two conversion routines are combined into one via
the top-level [`Parser`](struct.Parser.html) type. This `Parser` will first
convert your pattern to an `Ast` and then convert the `Ast` to an `Hir`.

# Example

This example shows how to parse a pattern string into its HIR:

```
use regex_syntax::Parser;
use regex_syntax::hir::{self, Hir};

let hir = Parser::new().parse("a|b").unwrap();
assert_eq!(hir, Hir::alternation(vec![
    Hir::literal(hir::Literal::Unicode('a')),
    Hir::literal(hir::Literal::Unicode('b')),
]));
```

# Concrete syntax supported

The concrete syntax is documented as part of the public API of the
[`regex` crate](https://docs.rs/regex/%2A/regex/#syntax).

# Input safety

A key feature of this library is that it is safe to use with end user facing
input. This plays a significant role in the internal implementation. In
particular:

1. Parsers provide a `nest_limit` option that permits callers to control how
   deeply nested a regular expression is allowed to be. This makes it possible
   to do case analysis over an `Ast` or an `Hir` using recursion without
   worrying about stack overflow.
2. Since relying on a particular stack size is brittle, this crate goes to
   great lengths to ensure that all interactions with both the `Ast` and the
   `Hir` do not use recursion. Namely, they use constant stack space and heap
   space proportional to the size of the original pattern string (in bytes).
   This includes the type's corresponding destructors. (One exception to this
   is literal extraction, but this will eventually get fixed.)

# Error reporting

The `Display` implementations on all `Error` types exposed in this library
provide nice human readable errors that are suitable for showing to end users
in a monospace font.

# Literal extraction

This crate provides limited support for
[literal extraction from `Hir` values](hir/literal/struct.Literals.html).
Be warned that literal extraction currently uses recursion, and therefore,
stack size proportional to the size of the `Hir`.

The purpose of literal extraction is to speed up searches. That is, if you
know a regular expression must match a prefix or suffix literal, then it is
often quicker to search for instances of that literal, and then confirm or deny
the match using the full regular expression engine. These optimizations are
done automatically in the `regex` crate.

# Crate features

An important feature provided by this crate is its Unicode support. This
includes things like case folding, boolean properties, general categories,
scripts and Unicode-aware support for the Perl classes `\w`, `\s` and `\d`.
However, a downside of this support is that it requires bundling several
Unicode data tables that are substantial in size.

A fair number of use cases do not require full Unicode support. For this
reason, this crate exposes a number of features to control which Unicode
data is available.

If a regular expression attempts to use a Unicode feature that is not available
because the corresponding crate feature was disabled, then translating that
regular expression to an `Hir` will return an error. (It is still possible
construct an `Ast` for such a regular expression, since Unicode data is not
used until translation to an `Hir`.) Stated differently, enabling or disabling
any of the features below can only add or subtract from the total set of valid
regular expressions. Enabling or disabling a feature will never modify the
match semantics of a regular expression.

The following features are available:

* **unicode** -
  Enables all Unicode features. This feature is enabled by default, and will
  always cover all Unicode features, even if more are added in the future.
* **unicode-age** -
  Provide the data for the
  [Unicode `Age` property](https://www.unicode.org/reports/tr44/tr44-24.html#Character_Age).
  This makes it possible to use classes like `\p{Age:6.0}` to refer to all
  codepoints first introduced in Unicode 6.0
* **unicode-bool** -
  Provide the data for numerous Unicode boolean properties. The full list
  is not included here, but contains properties like `Alphabetic`, `Emoji`,
  `Lowercase`, `Math`, `Uppercase` and `White_Space`.
* **unicode-case** -
  Provide the data for case insensitive matching using
  [Unicode's "simple loose matches" specification](https://www.unicode.org/reports/tr18/#Simple_Loose_Matches).
* **unicode-gencat** -
  Provide the data for
  [Uncode general categories](https://www.unicode.org/reports/tr44/tr44-24.html#General_Category_Values).
  This includes, but is not limited to, `Decimal_Number`, `Letter`,
  `Math_Symbol`, `Number` and `Punctuation`.
* **unicode-perl** -
  Provide the data for supporting the Unicode-aware Perl character classes,
  corresponding to `\w`, `\s` and `\d`. This is also necessary for using
  Unicode-aware word boundary assertions. Note that if this feature is
  disabled, the `\s` and `\d` character classes are still available if the
  `unicode-bool` and `unicode-gencat` features are enabled, respectively.
* **unicode-script** -
  Provide the data for
  [Unicode scripts and script extensions](https://www.unicode.org/reports/tr24/).
  This includes, but is not limited to, `Arabic`, `Cyrillic`, `Hebrew`,
  `Latin` and `Thai`.
* **unicode-segment** -
  Provide the data necessary to provide the properties used to implement the
  [Unicode text segmentation algorithms](https://www.unicode.org/reports/tr29/).
  This enables using classes like `\p{gcb=Extend}`, `\p{wb=Katakana}` and
  `\p{sb=ATerm}`.
*/

#![deny(missing_docs)]
#![forbid(unsafe_code)]

pub use error::{Error, Result};
pub use parser::{Parser, ParserBuilder};
pub use unicode::UnicodeWordError;

pub mod ast;
mod either;
mod error;
pub mod hir;
mod parser;
mod unicode;
mod unicode_tables;
pub mod utf8;

/// Escapes all regular expression meta characters in `text`.
///
/// The string returned may be safely used as a literal in a regular
/// expression.
pub fn escape(text: &str) -> String {
    let mut quoted = String::new();
    escape_into(text, &mut quoted);
    quoted
}

/// Escapes all meta characters in `text` and writes the result into `buf`.
///
/// This will append escape characters into the given buffer. The characters
/// that are appended are safe to use as a literal in a regular expression.
pub fn escape_into(text: &str, buf: &mut String) {
    buf.reserve(text.len());
    for c in text.chars() {
        if is_meta_character(c) {
            buf.push('\\');
        }
        buf.push(c);
    }
}

/// Returns true if the give character has significance in a regex.
///
/// These are the only characters that are allowed to be escaped, with one
/// exception: an ASCII space character may be escaped when extended mode (with
/// the `x` flag) is enabled. In particular, `is_meta_character(' ')` returns
/// `false`.
///
/// Note that the set of characters for which this function returns `true` or
/// `false` is fixed and won't change in a semver compatible release.
pub fn is_meta_character(c: char) -> bool {
    match c {
        '\\' | '.' | '+' | '*' | '?' | '(' | ')' | '|' | '[' | ']' | '{'
        | '}' | '^' | '$' | '#' | '&' | '-' | '~' => true,
        _ => false,
    }
}

/// Returns true if and only if the given character is a Unicode word
/// character.
///
/// A Unicode word character is defined by
/// [UTS#18 Annex C](http://unicode.org/reports/tr18/#Compatibility_Properties).
/// In particular, a character
/// is considered a word character if it is in either of the `Alphabetic` or
/// `Join_Control` properties, or is in one of the `Decimal_Number`, `Mark`
/// or `Connector_Punctuation` general categories.
///
/// # Panics
///
/// If the `unicode-perl` feature is not enabled, then this function panics.
/// For this reason, it is recommended that callers use
/// [`try_is_word_character`](fn.try_is_word_character.html)
/// instead.
pub fn is_word_character(c: char) -> bool {
    try_is_word_character(c).expect("unicode-perl feature must be enabled")
}

/// Returns true if and only if the given character is a Unicode word
/// character.
///
/// A Unicode word character is defined by
/// [UTS#18 Annex C](http://unicode.org/reports/tr18/#Compatibility_Properties).
/// In particular, a character
/// is considered a word character if it is in either of the `Alphabetic` or
/// `Join_Control` properties, or is in one of the `Decimal_Number`, `Mark`
/// or `Connector_Punctuation` general categories.
///
/// # Errors
///
/// If the `unicode-perl` feature is not enabled, then this function always
/// returns an error.
pub fn try_is_word_character(
    c: char,
) -> std::result::Result<bool, UnicodeWordError> {
    unicode::is_word_character(c)
}

/// Returns true if and only if the given character is an ASCII word character.
///
/// An ASCII word character is defined by the following character class:
/// `[_0-9a-zA-Z]'.
pub fn is_word_byte(c: u8) -> bool {
    match c {
        b'_' | b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' => true,
        _ => false,
    }
}

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

    #[test]
    fn escape_meta() {
        assert_eq!(
            escape(r"\.+*?()|[]{}^$#&-~"),
            r"\\\.\+\*\?\(\)\|\[\]\{\}\^\$\#\&\-\~".to_string()
        );
    }

    #[test]
    fn word_byte() {
        assert!(is_word_byte(b'a'));
        assert!(!is_word_byte(b'-'));
    }

    #[test]
    #[cfg(feature = "unicode-perl")]
    fn word_char() {
        assert!(is_word_character('a'), "ASCII");
        assert!(is_word_character('à'), "Latin-1");
        assert!(is_word_character('β'), "Greek");
        assert!(is_word_character('\u{11011}'), "Brahmi (Unicode 6.0)");
        assert!(is_word_character('\u{11611}'), "Modi (Unicode 7.0)");
        assert!(is_word_character('\u{11711}'), "Ahom (Unicode 8.0)");
        assert!(is_word_character('\u{17828}'), "Tangut (Unicode 9.0)");
        assert!(is_word_character('\u{1B1B1}'), "Nushu (Unicode 10.0)");
        assert!(is_word_character('\u{16E40}'), "Medefaidrin (Unicode 11.0)");
        assert!(!is_word_character('-'));
        assert!(!is_word_character('☃'));
    }

    #[test]
    #[should_panic]
    #[cfg(not(feature = "unicode-perl"))]
    fn word_char_disabled_panic() {
        assert!(is_word_character('a'));
    }

    #[test]
    #[cfg(not(feature = "unicode-perl"))]
    fn word_char_disabled_error() {
        assert!(try_is_word_character('a').is_err());
    }
}
<|endoftext|>


<commit_before>use std::collections::HashSet;

use crate::assist_ctx::{Assist, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, SyntaxNode, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
pub(crate) fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

pub(crate) fn add_missing_impl_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    use SyntaxKind::{IMPL_BLOCK, ITEM_LIST, WHITESPACE};

    let node = ctx.covering_node();
    let kinds = node.ancestors().take(3).map(SyntaxNode::kind);
    // Only suggest this in `impl Foo for S { [Item...] <|> }` cursor position
    if !Iterator::eq(kinds, [WHITESPACE, ITEM_LIST, IMPL_BLOCK].iter().cloned()) {
        return None;
    }

    let impl_node = node.ancestors().find_map(ast::ImplBlock::cast)?;

    let trait_def = {
        let db = ctx.db;
        // TODO: Can we get the position of cursor itself rather than supplied range?
        let range = ctx.frange;
        let position = FilePosition { file_id: range.file_id, offset: range.range.start() };
        let resolver = hir::source_binder::resolver_for_position(db, position);

        resolve_target_trait_def(db, &resolver, impl_node)?
    };

    let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
    let def_name = |&def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

    let trait_items = trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
    let impl_items = impl_node.item_list()?.impl_items();

    let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
    let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

    let trait_fn_names = trait_fns.iter().filter_map(def_name).collect::<HashSet<_>>();
    let impl_fn_names = impl_fns.iter().filter_map(def_name).collect::<HashSet<_>>();

    let missing_fn_names = trait_fn_names.difference(&impl_fn_names).collect::<HashSet<_>>();
    let missing_fns = trait_fns
        .iter()
        .filter(|&t| def_name(t).map(|n| missing_fn_names.contains(&n)).unwrap_or(false));

    unimplemented!()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{ check_assist };

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
}

struct S;

impl Foo for S {
    <|>
}",
            "
trait Foo {
    fn foo(&self);
}

struct S;

impl Foo for S {
    fn foo(&self) {
        <|>
    }
}",
        );
    }
}
<commit_msg>Implement a simple working assist<commit_after>use std::collections::HashSet;

use crate::{Assist, AssistId, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, SyntaxNode, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
pub(crate) fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

pub(crate) fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children() {
        if child.kind() == SyntaxKind::SEMI {
            buf.push_str(" { unimplemented!() }")
        } else {
            child.text().push_to(&mut buf);
        }
    }

    buf.trim_end().to_string()
}

pub(crate) fn add_missing_impl_members(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    use SyntaxKind::{IMPL_BLOCK, ITEM_LIST, WHITESPACE};

    let node = ctx.covering_node();
    let kinds = node.ancestors().take(3).map(SyntaxNode::kind);
    // Only suggest this in `impl Foo for S { [Item...] <|> }` cursor position
    if !Iterator::eq(kinds, [WHITESPACE, ITEM_LIST, IMPL_BLOCK].iter().cloned()) {
        return None;
    }

    let impl_node = node.ancestors().find_map(ast::ImplBlock::cast)?;
    let impl_item_list = impl_node.item_list()?;

    let trait_def = {
        let db = ctx.db;
        // TODO: Can we get the position of cursor itself rather than supplied range?
        let range = ctx.frange;
        let position = FilePosition { file_id: range.file_id, offset: range.range.start() };
        let resolver = hir::source_binder::resolver_for_position(db, position);

        resolve_target_trait_def(db, &resolver, impl_node)?
    };

    let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
    let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

    let trait_items = trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
    let impl_items = impl_item_list.impl_items();

    let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
    let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

    let trait_fn_names = trait_fns.iter().cloned().filter_map(def_name).collect::<HashSet<_>>();
    let impl_fn_names = impl_fns.iter().cloned().filter_map(def_name).collect::<HashSet<_>>();

    let missing_fn_names = trait_fn_names.difference(&impl_fn_names).collect::<HashSet<_>>();
    let missing_fns: Vec<_> = trait_fns
        .iter()
        .cloned()
        .filter(|t| def_name(t).map(|n| missing_fn_names.contains(&n)).unwrap_or(false))
        .collect();

    if missing_fns.is_empty() {
        return None;
    }

    let last_whitespace_node =
        impl_item_list.syntax().children().filter_map(ast::Whitespace::cast).last()?.syntax();

    ctx.add_action(AssistId("add_impl_missing_members"), "add impl missing members", |edit| {
        let func_bodies = missing_fns.into_iter().map(build_func_body).join("\n");
        let func_bodies = String::from("\n") + &func_bodies;

        let first_impl_item = impl_item_list.impl_items().next();
        // FIXME: We should respect the indent of the first item from the item list or the indent of leading block + some default indent (4?)
        // Another approach is to not indent at all if there are no items here
        let indent = first_impl_item.and_then(|i| leading_indent(i.syntax())).unwrap_or_default();
        let func_bodies = reindent(&func_bodies, indent) + "\n";

        let changed_range = last_whitespace_node.range();
        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(changed_range.start() + replaced_text_range - TextUnit::of_str("\n"));
    });

    ctx.build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{ check_assist };

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Hailstone_sequence
// Define a struct which we can make into an interator.
struct Hailstone {
  priv current: int, // Accessible online only to itself.
       start:   int  // Accessible to everything.
}

// Define a constructor for the struct.
impl Hailstone {
  fn new(n: int) -> Hailstone {
    Hailstone { current: n, start: n }
  }
}

// Implement the hailstone iteration sequence.
impl Iterator<int> for Hailstone {
  // This gets called to fetch the next item of the iterator.
  fn next(&mut self) -> Option<int> {
    let current = self.current;    // We need to cache the current value.
    match current {
      0               => {
        // Resets the iterator.
        self.current = self.start;
        None
      },
      1               => {
        // At the end, yield 1 and roll over next time called.
        self.current = 0;
        Some(1)
      },
      x if x % 2 == 0 => {
        // Got an even.
        self.current = x / 2;
        Some(current)
      },
      x if x % 2 == 1 => {
        // Got an odd.
        self.current = (3 * x) + 1;
        Some(current)
      },
      _               => None
    }
  }
}

fn main() {
  // Find the hailstone for 27.
  let two_seven: ~[int] = Hailstone::new(27).collect();
  println!("Testing: {}, Length: {}, Values: {}...{}", two_seven[0], two_seven.len(), two_seven.slice(0,4).to_str(), two_seven.slice(two_seven.len()-4, two_seven.len()).to_str());
  // Find the longest.
  let mut biggest = range(0, 100000).map(|x| Hailstone::new(x)).max_by(|&mut x| x.len()).unwrap();
  let size = biggest.len();
  println!("Largest: {}, Size: {}", biggest.start, size);
}
<commit_msg>Fixed hailstone priv<commit_after>// Implements http://rosettacode.org/wiki/Hailstone_sequence
// Define a struct which we can make into an interator.
struct Hailstone {
  current: int, // Accessible online only to itself.
  start:   int  // Accessible to everything.
}

// Define a constructor for the struct.
impl Hailstone {
  fn new(n: int) -> Hailstone {
    Hailstone { current: n, start: n }
  }
}

// Implement the hailstone iteration sequence.
impl Iterator<int> for Hailstone {
  // This gets called to fetch the next item of the iterator.
  fn next(&mut self) -> Option<int> {
    let current = self.current;    // We need to cache the current value.
    match current {
      0               => {
        // Resets the iterator.
        self.current = self.start;
        None
      },
      1               => {
        // At the end, yield 1 and roll over next time called.
        self.current = 0;
        Some(1)
      },
      x if x % 2 == 0 => {
        // Got an even.
        self.current = x / 2;
        Some(current)
      },
      x if x % 2 == 1 => {
        // Got an odd.
        self.current = (3 * x) + 1;
        Some(current)
      },
      _               => None
    }
  }
}

fn main() {
  // Find the hailstone for 27.
  let two_seven: ~[int] = Hailstone::new(27).collect();
  println!("Testing: {}, Length: {}, Values: {}...{}", two_seven[0], two_seven.len(), two_seven.slice(0,4).to_str(), two_seven.slice(two_seven.len()-4, two_seven.len()).to_str());
  // Find the longest.
  let mut biggest = range(0, 100000).map(|x| Hailstone::new(x)).max_by(|&mut x| x.len()).unwrap();
  let size = biggest.len();
  println!("Largest: {}, Size: {}", biggest.start, size);
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Hailstone_sequence
struct Hailstone {
  priv current: int,
       start:   int
}

impl Hailstone {
  fn new(n: int) -> Hailstone {
    Hailstone { current: n, start: n }
  }
}

impl Iterator<int> for Hailstone {
  fn next(&mut self) -> Option<int> {
    let current = self.current;
    match current {
      0               => {
        self.current = self.start; // Resets the iterator.
        None
      },
      1               => {
        self.current = 0;
        Some(1)
      },
      x if x % 2 == 0 => {
        self.current = x / 2;
        Some(current)
      },
      x if x % 2 == 1 => {
        self.current = (3 * x) + 1;
        Some(current)
      },
      _               => None
    }
  }
}

fn main() {
  // Find the hailstone for 27.
  let two_seven: ~[int] = Hailstone::new(27).collect();
  println!("Testing: {}, Length: {}, Values: {}...{}", two_seven[0], two_seven.len(), two_seven.slice(0,4).to_str(), two_seven.slice(two_seven.len()-4, two_seven.len()).to_str());
  // Find the longest.
  let mut biggest = range(0, 100000).map(|x| Hailstone::new(x)).max_by(|&mut x| x.len()).unwrap();
  let size = biggest.len();
  println!("Largest: {}, Size: {}", biggest.start, size);
}
<commit_msg>Better comments on Hailstone.<commit_after>// Implements http://rosettacode.org/wiki/Hailstone_sequence
// Define a struct which we can make into an interator.
struct Hailstone {
  priv current: int, // Accessible online only to itself.
       start:   int  // Accessible to everything.
}

// Define a constructor for the struct.
impl Hailstone {
  fn new(n: int) -> Hailstone {
    Hailstone { current: n, start: n }
  }
}

// Implement the hailstone iteration sequence.
impl Iterator<int> for Hailstone {
  // This gets called to fetch the next item of the iterator.
  fn next(&mut self) -> Option<int> {
    let current = self.current;    // We need to cache the current value.
    match current {
      0               => {
        // Resets the iterator.
        self.current = self.start;
        None
      },
      1               => {
        // At the end, yield 1 and roll over next time called.
        self.current = 0;
        Some(1)
      },
      x if x % 2 == 0 => {
        // Got an even.
        self.current = x / 2;
        Some(current)
      },
      x if x % 2 == 1 => {
        // Got an odd.
        self.current = (3 * x) + 1;
        Some(current)
      },
      _               => None
    }
  }
}

fn main() {
  // Find the hailstone for 27.
  let two_seven: ~[int] = Hailstone::new(27).collect();
  println!("Testing: {}, Length: {}, Values: {}...{}", two_seven[0], two_seven.len(), two_seven.slice(0,4).to_str(), two_seven.slice(two_seven.len()-4, two_seven.len()).to_str());
  // Find the longest.
  let mut biggest = range(0, 100000).map(|x| Hailstone::new(x)).max_by(|&mut x| x.len()).unwrap();
  let size = biggest.len();
  println!("Largest: {}, Size: {}", biggest.start, size);
}
<|endoftext|>


<commit_before>// http://rosettacode.org/wiki/Towers_of_Hanoi
// not_tested

fn move(n: int, from: int, to: int, via: int) {
    if n > 0 {
        move(n - 1, from, via, to);
        println!("Move disk from pole {:d} to pole {:d}", from, to);
        move(n - 1, via, to, from);
    }
}

fn main() {
    move(4, 1, 2, 3);
}
<commit_msg>Fixing breakage due to move keywordification.<commit_after>// http://rosettacode.org/wiki/Towers_of_Hanoi
// not_tested

fn play(n: int, from: int, to: int, via: int) {
    if n > 0 {
        play(n - 1, from, via, to);
        println!("Move disk from pole {:d} to pole {:d}", from, to);
        play(n - 1, via, to, from);
    }
}

fn main() {
    play(4, 1, 2, 3);
}
<|endoftext|>