<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{HashMap, HashSet};
use std::sync::{mpsc, Arc, RwLock};
use std::path::Path;

use grpc::EnvBuilder;
use tempdir::TempDir;

use tikv::config::TiKvConfig;
use tikv::server::{Server, ServerTransport};
use tikv::server::{create_raft_storage, Config, Node, PdStoreAddrResolver, RaftClient};
use tikv::server::resolve::{self, Task as ResolveTask};
use tikv::server::transport::ServerRaftStoreRouter;
use tikv::server::readpool::ReadPool;
use tikv::raftstore::{store, Result};
use tikv::raftstore::store::{Callback, Engines, Msg as StoreMsg, SnapManager};
use tikv::raftstore::coprocessor::CoprocessorHost;
use tikv::server::transport::RaftStoreRouter;
use tikv::util::transport::SendCh;
use tikv::util::security::SecurityManager;
use tikv::util::worker::{FutureWorker, Worker};
use tikv::storage::{self, Engine};
use tikv::import::{ImportSSTService, SSTImporter};
use kvproto::raft_serverpb::{self, RaftMessage};
use kvproto::raft_cmdpb::*;

use super::pd::TestPdClient;
use super::transport_simulate::*;
use super::cluster::{Cluster, Simulator};
use super::util::create_test_engine;

type SimulateStoreTransport = SimulateTransport<StoreMsg, ServerRaftStoreRouter>;
type SimulateServerTransport =
    SimulateTransport<RaftMessage, ServerTransport<SimulateStoreTransport, PdStoreAddrResolver>>;

struct ServerMeta {
    node: Node<TestPdClient>,
    server: Server<SimulateStoreTransport, PdStoreAddrResolver>,
    router: SimulateStoreTransport,
    sim_trans: SimulateServerTransport,
    store_ch: SendCh<StoreMsg>,
    worker: Worker<ResolveTask>,
}

pub struct ServerCluster {
    metas: HashMap<u64, ServerMeta>,
    addrs: HashMap<u64, String>,
    pub storages: HashMap<u64, Box<Engine>>,
    snap_paths: HashMap<u64, TempDir>,
    pd_client: Arc<TestPdClient>,
    raft_client: RaftClient,
}

impl ServerCluster {
    pub fn new(pd_client: Arc<TestPdClient>) -> ServerCluster {
        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(1)
                .name_prefix(thd_name!("server-cluster"))
                .build(),
        );
        let security_mgr = Arc::new(SecurityManager::new(&Default::default()).unwrap());
        ServerCluster {
            metas: HashMap::new(),
            addrs: HashMap::new(),
            pd_client: pd_client,
            storages: HashMap::new(),
            snap_paths: HashMap::new(),
            raft_client: RaftClient::new(env, Arc::new(Config::default()), security_mgr),
        }
    }

    pub fn get_addr(&self, node_id: u64) -> &str {
        &self.addrs[&node_id]
    }
}

impl Simulator for ServerCluster {
    #[allow(useless_format)]
    fn run_node(
        &mut self,
        node_id: u64,
        mut cfg: TiKvConfig,
        engines: Option<Engines>,
    ) -> (u64, Engines, Option<TempDir>) {
        assert!(node_id == 0 || !self.metas.contains_key(&node_id));

        let (tmp_str, tmp) = if node_id == 0 || !self.snap_paths.contains_key(&node_id) {
            let p = TempDir::new("test_cluster").unwrap();
            (p.path().to_str().unwrap().to_owned(), Some(p))
        } else {
            let p = self.snap_paths[&node_id].path().to_str().unwrap();
            (p.to_owned(), None)
        };

        // Now we cache the store address, so here we should re-use last
        // listening address for the same store.
        if let Some(addr) = self.addrs.get(&node_id) {
            cfg.server.addr = addr.clone();
        }

        // Initialize raftstore channels.
        let mut event_loop = store::create_event_loop(&cfg.raft_store).unwrap();
        let store_sendch = SendCh::new(event_loop.channel(), "raftstore");
        let (snap_status_sender, snap_status_receiver) = mpsc::channel();
        let raft_router = ServerRaftStoreRouter::new(store_sendch.clone(), snap_status_sender);
        let sim_router = SimulateTransport::new(raft_router);

        // Create engine
        let (engines, path) = create_test_engine(engines, store_sendch.clone(), &cfg);

        // Create storage.
        let pd_worker = FutureWorker::new("test future worker");
        let read_pool = ReadPool::new("readpool", &cfg.readpool.storage, || {
            || storage::ReadPoolContext::new(pd_worker.scheduler())
        });
        let mut store = create_raft_storage(sim_router.clone(), &cfg.storage, read_pool).unwrap();
        store.start(&cfg.storage).unwrap();
        self.storages.insert(node_id, store.get_engine());

        // Create import service.
        let importer = {
            let dir = Path::new(engines.kv_engine.path()).join("import-sst");
            Arc::new(SSTImporter::new(dir).unwrap())
        };
        let import_service = ImportSSTService::new(
            cfg.import.clone(),
            sim_router.clone(),
            Arc::clone(&importer),
        );

        // Create pd client, snapshot manager, server.
        let (worker, resolver) = resolve::new_resolver(Arc::clone(&self.pd_client)).unwrap();
        let snap_mgr = SnapManager::new(tmp_str, Some(store_sendch));
        let pd_worker = FutureWorker::new("test-pd-worker");
        let server_cfg = Arc::new(cfg.server.clone());
        let security_mgr = Arc::new(SecurityManager::new(&cfg.security).unwrap());
        let mut server = Server::new(
            &server_cfg,
            &security_mgr,
            cfg.coprocessor.region_split_size.0 as usize,
            store.clone(),
            sim_router.clone(),
            resolver,
            snap_mgr.clone(),
            pd_worker.scheduler(),
            Some(engines.clone()),
            Some(import_service),
        ).unwrap();
        let addr = server.listening_addr();
        cfg.server.addr = format!("{}", addr);
        let trans = server.transport();
        let simulate_trans = SimulateTransport::new(trans.clone());
        let server_cfg = Arc::new(cfg.server.clone());

        // Create node.
        let mut node = Node::new(
            &mut event_loop,
            &cfg.server,
            &cfg.raft_store,
            Arc::clone(&self.pd_client),
        );

        // Create coprocessor.
        let coprocessor_host = CoprocessorHost::new(cfg.coprocessor, node.get_sendch());

        node.start(
            event_loop,
            engines.clone(),
            simulate_trans.clone(),
            snap_mgr.clone(),
            snap_status_receiver,
            pd_worker,
            coprocessor_host,
            importer,
        ).unwrap();
        assert!(node_id == 0 || node_id == node.id());
        let node_id = node.id();
        if let Some(tmp) = tmp {
            self.snap_paths.insert(node_id, tmp);
        }

        server.start(server_cfg, security_mgr).unwrap();

        self.metas.insert(
            node_id,
            ServerMeta {
                store_ch: node.get_sendch(),
                node: node,
                server: server,
                router: sim_router,
                sim_trans: simulate_trans,
                worker: worker,
            },
        );
        self.addrs.insert(node_id, format!("{}", addr));

        (node_id, engines, path)
    }

    fn get_snap_dir(&self, node_id: u64) -> String {
        self.snap_paths[&node_id]
            .path()
            .to_str()
            .unwrap()
            .to_owned()
    }

    fn stop_node(&mut self, node_id: u64) {
        if let Some(mut meta) = self.metas.remove(&node_id) {
            meta.server.stop().unwrap();
            meta.node.stop().unwrap();
            meta.worker.stop().unwrap().join().unwrap();
        }
    }

    fn get_node_ids(&self) -> HashSet<u64> {
        self.metas.keys().cloned().collect()
    }

    fn async_command_on_node(
        &self,
        node_id: u64,
        request: RaftCmdRequest,
        cb: Callback,
    ) -> Result<()> {
        let router = match self.metas.get(&node_id) {
            None => return Err(box_err!("missing sender for store {}", node_id)),
            Some(meta) => meta.router.clone(),
        };
        router.send_command(request, cb)
    }

    fn send_raft_msg(&mut self, raft_msg: raft_serverpb::RaftMessage) -> Result<()> {
        let store_id = raft_msg.get_to_peer().get_store_id();
        let addr = self.get_addr(store_id).to_owned();
        self.raft_client.send(store_id, &addr, raft_msg).unwrap();
        self.raft_client.flush();
        Ok(())
    }

    fn add_send_filter(&mut self, node_id: u64, filter: SendFilter) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .sim_trans
            .add_filter(filter);
    }

    fn clear_send_filters(&mut self, node_id: u64) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .sim_trans
            .clear_filters();
    }

    fn add_recv_filter(&mut self, node_id: u64, filter: RecvFilter) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .router
            .add_filter(filter);
    }

    fn clear_recv_filters(&mut self, node_id: u64) {
        self.metas.get_mut(&node_id).unwrap().router.clear_filters();
    }

    fn get_store_sendch(&self, node_id: u64) -> Option<SendCh<StoreMsg>> {
        self.metas.get(&node_id).map(|m| m.store_ch.clone())
    }
}

pub fn new_server_cluster(id: u64, count: usize) -> Cluster<ServerCluster> {
    let pd_client = Arc::new(TestPdClient::new(id));
    let sim = Arc::new(RwLock::new(ServerCluster::new(Arc::clone(&pd_client))));
    Cluster::new(id, count, sim, pd_client)
}
<commit_msg>tests: run_node should retry when it meets gRPC bind fail (#2930)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{HashMap, HashSet};
use std::sync::{mpsc, Arc, RwLock};
use std::path::Path;
use std::time::Duration;
use std::thread;

use grpc::{EnvBuilder, Error as GrpcError};
use tempdir::TempDir;

use tikv::config::TiKvConfig;
use tikv::server::{create_raft_storage, Config, Error, Node, PdStoreAddrResolver, RaftClient,
                   Server, ServerTransport};
use tikv::server::resolve::{self, Task as ResolveTask};
use tikv::server::transport::ServerRaftStoreRouter;
use tikv::server::readpool::ReadPool;
use tikv::raftstore::{store, Result};
use tikv::raftstore::store::{Callback, Engines, Msg as StoreMsg, SnapManager};
use tikv::raftstore::coprocessor::CoprocessorHost;
use tikv::server::transport::RaftStoreRouter;
use tikv::util::transport::SendCh;
use tikv::util::security::SecurityManager;
use tikv::util::worker::{FutureWorker, Worker};
use tikv::storage::{self, Engine};
use tikv::import::{ImportSSTService, SSTImporter};
use kvproto::raft_serverpb::{self, RaftMessage};
use kvproto::raft_cmdpb::*;

use super::pd::TestPdClient;
use super::transport_simulate::*;
use super::cluster::{Cluster, Simulator};
use super::util::create_test_engine;

type SimulateStoreTransport = SimulateTransport<StoreMsg, ServerRaftStoreRouter>;
type SimulateServerTransport =
    SimulateTransport<RaftMessage, ServerTransport<SimulateStoreTransport, PdStoreAddrResolver>>;

struct ServerMeta {
    node: Node<TestPdClient>,
    server: Server<SimulateStoreTransport, PdStoreAddrResolver>,
    router: SimulateStoreTransport,
    sim_trans: SimulateServerTransport,
    store_ch: SendCh<StoreMsg>,
    worker: Worker<ResolveTask>,
}

pub struct ServerCluster {
    metas: HashMap<u64, ServerMeta>,
    addrs: HashMap<u64, String>,
    pub storages: HashMap<u64, Box<Engine>>,
    snap_paths: HashMap<u64, TempDir>,
    pd_client: Arc<TestPdClient>,
    raft_client: RaftClient,
}

impl ServerCluster {
    pub fn new(pd_client: Arc<TestPdClient>) -> ServerCluster {
        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(1)
                .name_prefix(thd_name!("server-cluster"))
                .build(),
        );
        let security_mgr = Arc::new(SecurityManager::new(&Default::default()).unwrap());
        ServerCluster {
            metas: HashMap::new(),
            addrs: HashMap::new(),
            pd_client: pd_client,
            storages: HashMap::new(),
            snap_paths: HashMap::new(),
            raft_client: RaftClient::new(env, Arc::new(Config::default()), security_mgr),
        }
    }

    pub fn get_addr(&self, node_id: u64) -> &str {
        &self.addrs[&node_id]
    }
}

impl Simulator for ServerCluster {
    #[allow(useless_format)]
    fn run_node(
        &mut self,
        node_id: u64,
        mut cfg: TiKvConfig,
        engines: Option<Engines>,
    ) -> (u64, Engines, Option<TempDir>) {
        assert!(node_id == 0 || !self.metas.contains_key(&node_id));

        let (tmp_str, tmp) = if node_id == 0 || !self.snap_paths.contains_key(&node_id) {
            let p = TempDir::new("test_cluster").unwrap();
            (p.path().to_str().unwrap().to_owned(), Some(p))
        } else {
            let p = self.snap_paths[&node_id].path().to_str().unwrap();
            (p.to_owned(), None)
        };

        // Now we cache the store address, so here we should re-use last
        // listening address for the same store.
        if let Some(addr) = self.addrs.get(&node_id) {
            cfg.server.addr = addr.clone();
        }

        // Initialize raftstore channels.
        let mut event_loop = store::create_event_loop(&cfg.raft_store).unwrap();
        let store_sendch = SendCh::new(event_loop.channel(), "raftstore");
        let (snap_status_sender, snap_status_receiver) = mpsc::channel();
        let raft_router = ServerRaftStoreRouter::new(store_sendch.clone(), snap_status_sender);
        let sim_router = SimulateTransport::new(raft_router);

        // Create engine
        let (engines, path) = create_test_engine(engines, store_sendch.clone(), &cfg);

        // Create storage.
        let pd_worker = FutureWorker::new("test future worker");
        let read_pool = ReadPool::new("readpool", &cfg.readpool.storage, || {
            || storage::ReadPoolContext::new(pd_worker.scheduler())
        });
        let mut store = create_raft_storage(sim_router.clone(), &cfg.storage, read_pool).unwrap();
        store.start(&cfg.storage).unwrap();
        self.storages.insert(node_id, store.get_engine());

        // Create import service.
        let importer = {
            let dir = Path::new(engines.kv_engine.path()).join("import-sst");
            Arc::new(SSTImporter::new(dir).unwrap())
        };
        let import_service = ImportSSTService::new(
            cfg.import.clone(),
            sim_router.clone(),
            Arc::clone(&importer),
        );

        // Create pd client, snapshot manager, server.
        let (worker, resolver) = resolve::new_resolver(Arc::clone(&self.pd_client)).unwrap();
        let snap_mgr = SnapManager::new(tmp_str, Some(store_sendch));
        let pd_worker = FutureWorker::new("test-pd-worker");
        let server_cfg = Arc::new(cfg.server.clone());
        let security_mgr = Arc::new(SecurityManager::new(&cfg.security).unwrap());
        let mut server = None;
        for _ in 0..100 {
            server = Some(Server::new(
                &server_cfg,
                &security_mgr,
                cfg.coprocessor.region_split_size.0 as usize,
                store.clone(),
                sim_router.clone(),
                resolver.clone(),
                snap_mgr.clone(),
                pd_worker.scheduler(),
                Some(engines.clone()),
                Some(import_service.clone()),
            ));
            match server {
                Some(Ok(_)) => break,
                Some(Err(Error::Grpc(GrpcError::BindFail(ref addr, ref port)))) => {
                    // Servers may meet the error, when we restart them.
                    debug!("fail to create a server: bind fail {:?}", (addr, port));
                    thread::sleep(Duration::from_millis(100));
                    continue;
                }
                Some(Err(ref e)) => panic!("fail to create a server: {:?}", e),
                None => unreachable!(),
            }
        }
        let mut server = server.unwrap().unwrap();
        let addr = server.listening_addr();
        cfg.server.addr = format!("{}", addr);
        let trans = server.transport();
        let simulate_trans = SimulateTransport::new(trans.clone());
        let server_cfg = Arc::new(cfg.server.clone());

        // Create node.
        let mut node = Node::new(
            &mut event_loop,
            &cfg.server,
            &cfg.raft_store,
            Arc::clone(&self.pd_client),
        );

        // Create coprocessor.
        let coprocessor_host = CoprocessorHost::new(cfg.coprocessor, node.get_sendch());

        node.start(
            event_loop,
            engines.clone(),
            simulate_trans.clone(),
            snap_mgr.clone(),
            snap_status_receiver,
            pd_worker,
            coprocessor_host,
            importer,
        ).unwrap();
        assert!(node_id == 0 || node_id == node.id());
        let node_id = node.id();
        if let Some(tmp) = tmp {
            self.snap_paths.insert(node_id, tmp);
        }

        server.start(server_cfg, security_mgr).unwrap();

        self.metas.insert(
            node_id,
            ServerMeta {
                store_ch: node.get_sendch(),
                node: node,
                server: server,
                router: sim_router,
                sim_trans: simulate_trans,
                worker: worker,
            },
        );
        self.addrs.insert(node_id, format!("{}", addr));

        (node_id, engines, path)
    }

    fn get_snap_dir(&self, node_id: u64) -> String {
        self.snap_paths[&node_id]
            .path()
            .to_str()
            .unwrap()
            .to_owned()
    }

    fn stop_node(&mut self, node_id: u64) {
        if let Some(mut meta) = self.metas.remove(&node_id) {
            meta.server.stop().unwrap();
            meta.node.stop().unwrap();
            meta.worker.stop().unwrap().join().unwrap();
        }
    }

    fn get_node_ids(&self) -> HashSet<u64> {
        self.metas.keys().cloned().collect()
    }

    fn async_command_on_node(
        &self,
        node_id: u64,
        request: RaftCmdRequest,
        cb: Callback,
    ) -> Result<()> {
        let router = match self.metas.get(&node_id) {
            None => return Err(box_err!("missing sender for store {}", node_id)),
            Some(meta) => meta.router.clone(),
        };
        router.send_command(request, cb)
    }

    fn send_raft_msg(&mut self, raft_msg: raft_serverpb::RaftMessage) -> Result<()> {
        let store_id = raft_msg.get_to_peer().get_store_id();
        let addr = self.get_addr(store_id).to_owned();
        self.raft_client.send(store_id, &addr, raft_msg).unwrap();
        self.raft_client.flush();
        Ok(())
    }

    fn add_send_filter(&mut self, node_id: u64, filter: SendFilter) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .sim_trans
            .add_filter(filter);
    }

    fn clear_send_filters(&mut self, node_id: u64) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .sim_trans
            .clear_filters();
    }

    fn add_recv_filter(&mut self, node_id: u64, filter: RecvFilter) {
        self.metas
            .get_mut(&node_id)
            .unwrap()
            .router
            .add_filter(filter);
    }

    fn clear_recv_filters(&mut self, node_id: u64) {
        self.metas.get_mut(&node_id).unwrap().router.clear_filters();
    }

    fn get_store_sendch(&self, node_id: u64) -> Option<SendCh<StoreMsg>> {
        self.metas.get(&node_id).map(|m| m.store_ch.clone())
    }
}

pub fn new_server_cluster(id: u64, count: usize) -> Cluster<ServerCluster> {
    let pd_client = Arc::new(TestPdClient::new(id));
    let sim = Arc::new(RwLock::new(ServerCluster::new(Arc::clone(&pd_client))));
    Cluster::new(id, count, sim, pd_client)
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate serde_json;
extern crate reqwest;
use self::iron::prelude::*;
use self::iron::{status, Request, Response, IronResult};

use util::{EnvVar, read_env_var};
use util::telegram::{register, unregister, send};
use model::telegram;

use std::io::Read;

pub fn telegram(req: &mut Request) -> IronResult<Response> {
    let mut body = String::new();
    req.body
        .read_to_string(&mut body)
        .map_err(|e| IronError::new(e, (status::BadRequest, "Error reading request")))?;
    let update: telegram::Update = serde_json::from_str(&body).unwrap();
    if let Some(msg) = update.message {
        if let Some(txt) = msg.text {
            let is_private = msg.chat.chat_type == "private";
            let identifier = format!("@{}", read_env_var(&EnvVar::Name));
            let has_identifier = txt.find(&identifier).is_some();
            let txt = strip_identifier(&txt);
            if is_private || has_identifier {
                let id = msg.chat.id;
                match txt.as_str() {
                    "/start" => respond_start(id),
                    "/help" => respond_help(id),
                    "/unsubscribe" => respond_unsubscribe(id),
                    "/delays" => respond_delays(id),
                    _ => {
                        if is_private {
                            respond_unknown(id)
                        } else {
                            Ok(())
                        }
                    }
                }?;
            }
        }
    }
    Ok(Response::with((status::Ok, "ok")))
}

fn strip_identifier(msg: &str) -> String {
    let identifier = format!("@{}", read_env_var(&EnvVar::Name));
    if let Some(pos) = msg.find(&identifier) {
        return msg[..pos].to_owned();
    }
    msg.to_owned()
}

fn respond_start(chat_id: i32) -> IronResult<()> {
    let new_registration = register(chat_id).unwrap();
    let msg = if new_registration {
        "Successfully registered!"
    } else {
        "This chat has already been registered"
    };
    send(chat_id, msg)?;
    Ok(())
}

fn respond_help(chat_id: i32) -> IronResult<()> {
    send(chat_id,
         "Available commands:\n\
        /start: Subscribes this chat to be notified of SBB delays\n\
        /help: Shows this window")?;
    Ok(())
}

fn respond_unsubscribe(chat_id: i32) -> IronResult<()> {
    let unsubscribed = unregister(chat_id).unwrap();
    let msg = if unsubscribed {
        "Successfully unsubscribed from the delay notifications.\n\
        Use /start to subscribe again"
    } else {
        "This chat was not subscribed to be notified of delays in the first place, silly!\n\
        Use /start if you want to subscribe"
    };
    send(chat_id, msg)?;
    Ok(())
}

fn respond_delays(chat_id: i32) -> IronResult<()> {
    Ok(())
}

fn respond_unknown(chat_id: i32) -> IronResult<()> {
    send(chat_id, "Unknown command. Try using /help")?;
    Ok(())
}
<commit_msg>Add help description for unsubscribe<commit_after>extern crate iron;
extern crate serde_json;
extern crate reqwest;
use self::iron::prelude::*;
use self::iron::{status, Request, Response, IronResult};

use util::{EnvVar, read_env_var};
use util::telegram::{register, unregister, send};
use model::telegram;

use std::io::Read;

pub fn telegram(req: &mut Request) -> IronResult<Response> {
    let mut body = String::new();
    req.body
        .read_to_string(&mut body)
        .map_err(|e| IronError::new(e, (status::BadRequest, "Error reading request")))?;
    let update: telegram::Update = serde_json::from_str(&body).unwrap();
    if let Some(msg) = update.message {
        if let Some(txt) = msg.text {
            let is_private = msg.chat.chat_type == "private";
            let identifier = format!("@{}", read_env_var(&EnvVar::Name));
            let has_identifier = txt.find(&identifier).is_some();
            let txt = strip_identifier(&txt);
            if is_private || has_identifier {
                let id = msg.chat.id;
                match txt.as_str() {
                    "/start" => respond_start(id),
                    "/help" => respond_help(id),
                    "/unsubscribe" => respond_unsubscribe(id),
                    "/delays" => respond_delays(id),
                    _ => {
                        if is_private {
                            respond_unknown(id)
                        } else {
                            Ok(())
                        }
                    }
                }?;
            }
        }
    }
    Ok(Response::with((status::Ok, "ok")))
}

fn strip_identifier(msg: &str) -> String {
    let identifier = format!("@{}", read_env_var(&EnvVar::Name));
    if let Some(pos) = msg.find(&identifier) {
        return msg[..pos].to_owned();
    }
    msg.to_owned()
}

fn respond_start(chat_id: i32) -> IronResult<()> {
    let new_registration = register(chat_id).unwrap();
    let msg = if new_registration {
        "Successfully registered!"
    } else {
        "This chat has already been registered"
    };
    send(chat_id, msg)?;
    Ok(())
}

fn respond_help(chat_id: i32) -> IronResult<()> {
    send(chat_id,
         "Available commands:\n\
        /start: Subscribes this chat to be notified of SBB delays\n\
        /unsubscribe - Unsubscribes this chat from the delay notifications\n\
        /help: Shows this window")?;
    Ok(())
}

fn respond_unsubscribe(chat_id: i32) -> IronResult<()> {
    let unsubscribed = unregister(chat_id).unwrap();
    let msg = if unsubscribed {
        "Successfully unsubscribed from the delay notifications.\n\
        Use /start to subscribe again"
    } else {
        "This chat was not subscribed to be notified of delays in the first place, silly!\n\
        Use /start if you want to subscribe"
    };
    send(chat_id, msg)?;
    Ok(())
}

fn respond_delays(chat_id: i32) -> IronResult<()> {
    Ok(())
}

fn respond_unknown(chat_id: i32) -> IronResult<()> {
    send(chat_id, "Unknown command. Try using /help")?;
    Ok(())
}
<|endoftext|>


<commit_before>pub type BoxedScene<EventT, RendererT, EngineDataT> = Box<Scene<EventT, RendererT, EngineDataT> + 'static>;
pub type SceneChangeCallback<EventT, RendererT, EngineDataT> = Box<Fn(&mut RendererT, &mut EngineDataT) -> BoxedScene<EventT, RendererT, EngineDataT>>;

pub trait Scene<EventT, RendererT, EngineDataT> {
    fn render(&self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64);
    fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64);
    fn think(&mut self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) -> Option<SceneChangeEvent<EventT, RendererT, EngineDataT>>;
}

pub enum SceneChangeEvent<EventT, RendererT, EngineDataT> {
    PushScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    SwapScene(SceneChangeCallback<EventT, RendererT, EngineDataT>),
    PopScene,
}

pub struct SceneStack<EventT, RendererT, EngineDataT> {
    scenes: Vec<BoxedScene<EventT, RendererT, EngineDataT>>
}

impl<EventT, RendererT, EngineDataT> SceneStack<EventT, RendererT, EngineDataT> {

    pub fn new() -> SceneStack<EngineDataT, EventT, RendererT> {
        SceneStack {
            scenes: vec![]
        }
    }

    pub fn is_empty(&self) -> bool {
        self.scenes.len() == 0
    }

    pub fn push(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) {
        self.scenes.push(scene)
    }

    pub fn swap(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT>) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        let old_scene = self.scenes.pop();
        self.scenes.push(scene);
        old_scene
    }

    pub fn pop(&mut self) -> Option<BoxedScene<EventT, RendererT, EngineDataT>> {
        self.scenes.pop()
    }

    pub fn render(&mut self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(scene) => {
                scene.render(renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        }
    }

    pub fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(mut scene) => {
                scene.handle_event(event, renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        };
    }

    pub fn think(&mut self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        let event = match maybe_last_scene {
            Some(mut scene) => {
                let event = scene.think(renderer, engine_data, tick);
                self.scenes.push(scene);
                event
            },
            None => None
        };

        match event {
            Some(SceneChangeEvent::PushScene(callback)) => {
                self.push(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::SwapScene(callback)) => {
                self.swap(callback(renderer, engine_data));
            },
            Some(SceneChangeEvent::PopScene) => {
                self.pop();
            },
            None => ()
        };
    }
}
<commit_msg>Moving back towards non-closure-passing scene changing due to lifetime complexity<commit_after>pub type BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT> = Box<Scene<EventT, RendererT, EngineDataT, SceneChangeDataT> + 'static>;
pub type SceneChangeCallback<EventT, RendererT, EngineDataT, SceneChangeDataT> = Fn(&SceneChangeDataT, &mut RendererT, &mut EngineDataT) -> BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>;

pub trait Scene<EventT, RendererT, EngineDataT, SceneChangeDataT> {
    fn render(&self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64);
    fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64);
    fn think(&mut self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) -> Option<SceneChangeEvent<SceneChangeDataT>>;
}

pub enum SceneChangeEvent<SceneChangeDataT> {
    PushScene(SceneChangeDataT),
    SwapScene(SceneChangeDataT),
    PopScene,
}

pub struct SceneStack<EventT, RendererT, EngineDataT, SceneChangeDataT> {
    scenes: Vec<BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>>
}

impl<EventT, RendererT, EngineDataT, SceneChangeDataT> SceneStack<EventT, RendererT, EngineDataT, SceneChangeDataT> {

    pub fn new() -> SceneStack<EngineDataT, EventT, RendererT, SceneChangeDataT> {
        SceneStack {
            scenes: vec![]
        }
    }

    pub fn is_empty(&self) -> bool {
        self.scenes.len() == 0
    }

    pub fn push(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>) {
        self.scenes.push(scene)
    }

    pub fn swap(&mut self, scene: BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>) -> Option<BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>> {
        let old_scene = self.scenes.pop();
        self.scenes.push(scene);
        old_scene
    }

    pub fn pop(&mut self) -> Option<BoxedScene<EventT, RendererT, EngineDataT, SceneChangeDataT>> {
        self.scenes.pop()
    }

    pub fn render(&mut self, renderer: &mut RendererT, engine_data: &EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(scene) => {
                scene.render(renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        }
    }

    pub fn handle_event(&mut self, event: &EventT, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64) {
        let maybe_last_scene = self.scenes.pop();
        match maybe_last_scene {
            Some(mut scene) => {
                scene.handle_event(event, renderer, engine_data, tick);
                self.scenes.push(scene);
            },
            None => ()
        };
    }

    pub fn think(&mut self, renderer: &mut RendererT, engine_data: &mut EngineDataT, tick: u64, scene_change_handler: &SceneChangeCallback<EventT, RendererT, EngineDataT, SceneChangeDataT>) {
        let maybe_last_scene = self.scenes.pop();
        let event = match maybe_last_scene {
            Some(mut scene) => {
                let event = scene.think(renderer, engine_data, tick);
                self.scenes.push(scene);
                event
            },
            None => None
        };

        match event {
            Some(SceneChangeEvent::PushScene(ref data)) => {
                self.push(scene_change_handler(data, renderer, engine_data));
            },
            Some(SceneChangeEvent::SwapScene(ref data)) => {
                self.swap(scene_change_handler(data, renderer, engine_data));
            },
            Some(SceneChangeEvent::PopScene) => {
                self.pop();
            },
            None => ()
        };
    }
}
<|endoftext|>


<commit_before>
use core::str::StrExt;
use regex;

pub fn score(choice: &String, query: &String) -> f32 {
    if choice.len() == 0 {
        0.0
    } else if query.len() == 0 {
        1.0
    } else {
        if compute_match_length(choice, query) > 0 {
            1.0 / choice.len() as f32
        } else {
            0.0
        }
    }
}

// Find the length of the shortest substring matching the given characters.
fn compute_match_length(string: &String, query: &String) -> usize {
    // FIXME (performance). Regexp is not reused.
    let re_string = &make_query_regex(query);
    let re = match regex::Regex::new(re_string.as_slice()) {
        Ok(re)   => re,
        Err(err) => panic!("{}", err.msg)
    };
    match re.captures(string.as_slice()) {
        Some(caps) => caps.at(0).unwrap_or("").char_len(),
        None       => 0
    }
}

// Creates a regex for performing a case-insensitive non-greedy fuzzy match.
// Turns "abc" into "(?i)a.*?b.*?c.*?".
fn make_query_regex(query: &String) -> String {
    let mut parts = query
        .chars()
        .map(|ch| regex::quote(ch.to_string().as_slice()))
        .collect::<Vec<String>>();

    parts.insert(0, "(?i)".to_string());

    parts.connect(".*?")
}

#[cfg(test)]
mod test {

    use score::score;
    use core::str::StrExt;

    pub fn do_score(choice: &str, query: &str) -> f32 {
       let choice_string = choice.to_string();
       let query_string = query.to_string();
       score(&choice_string,  &query_string)
    }

    #[test]
    fn scores_zero_when_the_choice_is_empty() {
        assert_eq!(do_score("", ""), 0.0);
    }

    #[test]
    fn scores_one_when_the_query_is_empty() {
        assert_eq!(do_score("a", ""), 1.0);
    }

    #[test]
    fn scores_zero_when_the_query_is_longer_than_the_choice() {
        assert_eq!(do_score("short", "longer"), 0.0);
    }

    #[test]
    fn scores_zero_when_only_a_prefix_of_the_query_matches() {
        assert_eq!(do_score("ab", "ac"), 0.0);
    }

    #[test]
    fn scores_greater_than_zero_when_it_matches() {
        assert!(do_score("a", "a") > 0.0);
        assert!(do_score("ab", "a") > 0.0);
        assert!(do_score("ba", "a") > 0.0);
        assert!(do_score("bab", "a") > 0.0);
        assert!(do_score("babababab", "aaaa") > 0.0);
    }

    #[test]
    fn scores_one_normalized_to_length_when_the_query_equals_the_choice() {
        assert_eq!(do_score("a", "a"), 1.0);
        assert_eq!(do_score("ab", "ab"), 0.5f32);
        assert_eq!(do_score("a long string", "a long string"),
            1.0 / ("a long string".char_len() as f32));

        assert_eq!(do_score("spec/search_spec.rb", "sear"),
            1.0 / ("spec/search_spec.rb".char_len() as f32));
    }

    // Character matching

    #[test]
    fn it_matches_punctuation() {
        assert!(do_score("/! symbols $^", "/!$^") > 0.0);
    }

    #[test]
    fn it_is_case_insensitive() {
        assert_eq!(do_score("a", "A"), 1.0);
        assert_eq!(do_score("A", "a"), 1.0);
    }

    #[test]
    fn it_doesnt_match_when_the_same_letter_is_repeated_in_the_choice() {
        assert_eq!(do_score("a", "aa"), 0.0);
    }

    #[test]
    fn it_scores_higher_for_better_matches() {
        assert!(do_score("selecta.gemspec", "asp")
                > do_score("algorithm4_spec.rb", "asp"));
        assert!(do_score("README.md", "em")
                > do_score("benchmark.rb", "em"));
        assert!(do_score("search.rb", "sear")
                > do_score("spec/search_spec.rb", "sear"));
    }
}

<commit_msg>Better variable name<commit_after>
use core::str::StrExt;
use regex;

pub fn score(choice: &String, query: &String) -> f32 {
    if choice.len() == 0 {
        0.0
    } else if query.len() == 0 {
        1.0
    } else {
        if compute_match_length(choice, query) > 0 {
            1.0 / choice.len() as f32
        } else {
            0.0
        }
    }
}

// Find the length of the shortest substring matching the given characters.
fn compute_match_length(choice: &String, query: &String) -> usize {
    // FIXME (performance). Regexp is not reused.
    let re_string = &make_query_regex(query);
    let re = match regex::Regex::new(re_string.as_slice()) {
        Ok(re)   => re,
        Err(err) => panic!("{}", err.msg)
    };
    match re.captures(choice.as_slice()) {
        Some(caps) => caps.at(0).unwrap_or("").char_len(),
        None       => 0
    }
}

// Creates a regex for performing a case-insensitive non-greedy fuzzy match.
// Turns "abc" into "(?i)a.*?b.*?c.*?".
fn make_query_regex(query: &String) -> String {
    let mut parts = query
        .chars()
        .map(|ch| regex::quote(ch.to_string().as_slice()))
        .collect::<Vec<String>>();

    parts.insert(0, "(?i)".to_string());

    parts.connect(".*?")
}

#[cfg(test)]
mod test {

    use score::score;
    use core::str::StrExt;

    pub fn do_score(choice: &str, query: &str) -> f32 {
       let choice_string = choice.to_string();
       let query_string = query.to_string();
       score(&choice_string,  &query_string)
    }

    #[test]
    fn scores_zero_when_the_choice_is_empty() {
        assert_eq!(do_score("", ""), 0.0);
    }

    #[test]
    fn scores_one_when_the_query_is_empty() {
        assert_eq!(do_score("a", ""), 1.0);
    }

    #[test]
    fn scores_zero_when_the_query_is_longer_than_the_choice() {
        assert_eq!(do_score("short", "longer"), 0.0);
    }

    #[test]
    fn scores_zero_when_only_a_prefix_of_the_query_matches() {
        assert_eq!(do_score("ab", "ac"), 0.0);
    }

    #[test]
    fn scores_greater_than_zero_when_it_matches() {
        assert!(do_score("a", "a") > 0.0);
        assert!(do_score("ab", "a") > 0.0);
        assert!(do_score("ba", "a") > 0.0);
        assert!(do_score("bab", "a") > 0.0);
        assert!(do_score("babababab", "aaaa") > 0.0);
    }

    #[test]
    fn scores_one_normalized_to_length_when_the_query_equals_the_choice() {
        assert_eq!(do_score("a", "a"), 1.0);
        assert_eq!(do_score("ab", "ab"), 0.5f32);
        assert_eq!(do_score("a long string", "a long string"),
            1.0 / ("a long string".char_len() as f32));

        assert_eq!(do_score("spec/search_spec.rb", "sear"),
            1.0 / ("spec/search_spec.rb".char_len() as f32));
    }

    // Character matching

    #[test]
    fn it_matches_punctuation() {
        assert!(do_score("/! symbols $^", "/!$^") > 0.0);
    }

    #[test]
    fn it_is_case_insensitive() {
        assert_eq!(do_score("a", "A"), 1.0);
        assert_eq!(do_score("A", "a"), 1.0);
    }

    #[test]
    fn it_doesnt_match_when_the_same_letter_is_repeated_in_the_choice() {
        assert_eq!(do_score("a", "aa"), 0.0);
    }

    #[test]
    fn it_scores_higher_for_better_matches() {
        assert!(do_score("selecta.gemspec", "asp")
                > do_score("algorithm4_spec.rb", "asp"));
        assert!(do_score("README.md", "em")
                > do_score("benchmark.rb", "em"));
        assert!(do_score("search.rb", "sear")
                > do_score("spec/search_spec.rb", "sear"));
    }
}

<|endoftext|>


<commit_before>pub struct Space {
    state: SpaceState,
    position: [u8; 2],
}

impl Space {
    pub fn new(pos: [u8; 2]) -> Space {
        Space {
            state: SpaceState::Unchecked,
            position: pos,
        }
    }

    /// Sets this space as having been checked, and whether it was hit.
    ///
    /// # Errors
    ///
    /// Returns an error if the space's state is not `SpaceState::Unchecked`.
    pub fn set_checked(&mut self, hit: bool) -> Result<(), &'static str> {
        if self.state != SpaceState::Unchecked {
            Err("tried to check an already checked space")
        } else {
            self.state = SpaceState::Checked(hit);
            Ok(())
        }
    }

    pub fn is_unchecked(&self) -> bool {
        self.state == SpaceState::Unchecked
    }

    pub fn is_empty(&self) -> bool {
        self.state == SpaceState::Checked(false)
    }

    pub fn is_hit(&self) -> bool {
        self.state == SpaceState::Checked(true)
    }

    pub fn pos(&self) -> &[u8; 2] {
        &self.position
    }
}

#[derive(Clone, Copy, PartialEq)]
pub enum SpaceState {
    Unchecked,
    Checked(bool),
}
<commit_msg>Add Space tests<commit_after>pub struct Space {
    state: SpaceState,
    position: [u8; 2],
}

impl Space {
    pub fn new(pos: [u8; 2]) -> Space {
        Space {
            state: SpaceState::Unchecked,
            position: pos,
        }
    }

    /// Sets this space as having been checked, and whether it was hit.
    ///
    /// # Errors
    ///
    /// Returns an error if the space's state is not `SpaceState::Unchecked`.
    pub fn set_checked(&mut self, hit: bool) -> Result<(), &'static str> {
        if self.state != SpaceState::Unchecked {
            Err("tried to check an already checked space")
        } else {
            self.state = SpaceState::Checked(hit);
            Ok(())
        }
    }

    pub fn is_unchecked(&self) -> bool {
        self.state == SpaceState::Unchecked
    }

    pub fn is_empty(&self) -> bool {
        self.state == SpaceState::Checked(false)
    }

    pub fn is_hit(&self) -> bool {
        self.state == SpaceState::Checked(true)
    }

    pub fn pos(&self) -> &[u8; 2] {
        &self.position
    }
}

#[derive(Clone, Copy, PartialEq)]
pub enum SpaceState {
    Unchecked,
    Checked(bool),
}

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

    #[test]
    fn set_checked() {
        let mut space = Space::new([0, 0]);
        assert!(space.set_checked(false).is_ok());
        assert!(space.set_checked(false).is_err());

        space = Space::new([0, 0]);
        assert!(space.set_checked(true).is_ok());
        assert!(space.set_checked(true).is_err());
    }

    #[test]
    fn is_unchecked() {
        let mut space = Space::new([0, 0]);
        assert!(space.is_unchecked());
        assert!(space.set_checked(false).is_ok());
        assert!(!space.is_unchecked());

        space = Space::new([0, 0]);
        assert!(space.set_checked(true).is_ok());
        assert!(!space.is_unchecked());
    }

    #[test]
    fn is_empty() {
        let mut space = Space::new([0, 0]);
        assert!(!space.is_empty());
        assert!(space.set_checked(false).is_ok());
        assert!(space.is_empty());

        space = Space::new([0, 0]);
        assert!(space.set_checked(true).is_ok());
        assert!(!space.is_empty());
    }

    #[test]
    fn is_hit() {
        let mut space = Space::new([0, 0]);
        assert!(!space.is_hit());
        assert!(space.set_checked(true).is_ok());
        assert!(space.is_hit());

        space = Space::new([0, 0]);
        assert!(space.set_checked(false).is_ok());
        assert!(!space.is_hit());
    }

    #[test]
    fn pos() {
        let space = Space::new([0, 0]);
        assert_eq!(space.pos(), &[0, 0]);
    }
}
<|endoftext|>


<commit_before>//! Module for the dynamic state of the game.
//!
//! This module defines the `State` type and its entourage. The data needed to
//! run an RBattle game is split into three categories:
//!
//! - `State` holds all the significant varying state of a game. What changes as the game proceeds
//!   are the amount of goop in each node, and the outflows the user has chosen;
//!   these are the `State`.
//!
//! - The positions of nodes and their boundaries and the locations of goop
//!   sources are all fixed, and are part of the `Map`.
//!
//! - Interface elements enter hover and active states as the user mouses
//!   around, but those states are ephemeral; they are part of the `Mouse` type.
//!   When the user actually completes an interaction with an interface element,
//!   only then is the `State` affected.

use graph::Node;
use map::Map;
use square::SquareGrid;
use std::sync::Arc;
use rand::{Rng, SeedableRng, XorShiftRng};
use xorshift::XorShift128Plus;

use std::iter::repeat;

/// The complete state of an RBattle game board.
#[derive(Clone)]
pub struct State {
    /// The map being played on.
    pub map: Arc<Map>,

    /// Which nodes are occupied, and which are vacant. Indexed by node id.
    pub nodes: Vec<Option<Occupied>>,

    /// The random number generator used to drive the goop flow algorithm.
    rng: XorShift128Plus
}

/// A player id number.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct Player(pub usize);

/// The state of a node that is occupied by some player.
#[derive(Debug, Clone)]
pub struct Occupied {
    /// The player who controls this node.
    pub player: Player,

    /// Which neighbors of this node it sends goop out to.
    pub outflows: Vec<Node>,

    /// How much goop this node holds. Ranges from 0 to MAX_GOOP.
    pub goop: usize,

}

pub const MAX_GOOP: usize = 15;

/// Return a pair of mutable references to the `i`'th and `j`'th elements of
/// `slice`, where `i != j`.
fn index_mut_pair<T>(slice: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
    if i < j {
        let (left, right) = slice.split_at_mut(j);
        (&mut left[i], &mut right[0])
    } else if j < i {
        let (left, right) = slice.split_at_mut(i);
        (&mut right[0], &mut left[j])
    } else {
        panic!("can't borrow two mutable references to the same element");
    }
}

impl State {
    pub fn new(params: GameParameters) -> State {
        let graph = SquareGrid::new(params.board.0, params.board.1);
        let map = Arc::new(Map::new(graph, params.sources, params.colors));

        let mut nodes: Vec<Option<Occupied>> = repeat(None).take(map.graph.nodes()).collect();
        // Ensure that each source is occupied by its player.
        for (player, &source) in map.sources.iter().enumerate() {
            nodes[source] = Some(Occupied {
                player: Player(player),
                outflows: vec![],
                goop: 0
            });
        }

        const SEED: [u64; 2] = [0xcd9d5eaaf04bc9a7, 0x4602cc7098d01ef9];
        State { map, nodes, rng: XorShift128Plus::new(SEED) }
    }

    /// Let one unit of goop flow through each outflow.
    ///
    /// There are algorithms for finding the flow through a graph precisely, but
    /// we need something simpler here. We just visit every outgoing edge in a
    /// random order, and propagate a unit of goop if the destination permits
    /// it.
    fn flow(&mut self) {
        // Build a vector of (from, to) pairs.
        let mut outflow_list = Vec::new();
        for node in 0..self.map.graph.nodes() {
            if let &Some(ref occupied) = &self.nodes[node] {
                for &outflow in &occupied.outflows {
                    outflow_list.push((node, outflow))
                }
            }
        }

        // Put the pairs in a random order.
        self.rng.shuffle(&mut outflow_list);

        while let Some((from_index, to_index)) = outflow_list.pop() {
            match index_mut_pair(&mut self.nodes, from_index, to_index) {
                // We shouldn't have generated a pair for an empty source, and
                // when we clear a node we're supposed to remove pairs from
                // `outflow` that originate there.
                (&mut None, _) => panic!("outflow from empty node"),

                // Source has no goop. No effect.
                (&mut Some(Occupied { goop: 0, .. }), _) => (),

                // Goop flowing into an unoccupied node. New player claims ownership.
                (&mut Some(Occupied { player, ref mut goop, .. }),
                 &mut ref mut to @ None) => {
                    *goop -= 1;
                    *to = Some(Occupied { player, outflows: vec![], goop: 1 });
                },

                // Goop flowing into a node occupied by the same player.
                (&mut Some(Occupied { player: from_player, goop: ref mut from_goop, .. }),
                 &mut Some(Occupied { player: to_player,   goop: ref mut to_goop, .. }))
                    if from_player == to_player =>
                {
                    if *from_goop > 0 && *to_goop < MAX_GOOP {
                        *from_goop -= 1;
                        *to_goop += 1;
                    }
                }

                // Goop flowing into a node occupied by another player, but
                // doesn't clear it. All outflow from destination stopped.
                (&mut Some(Occupied { goop: ref mut from_goop, .. }),
                 &mut Some(Occupied { outflows: ref mut to_outflows,
                                       goop:     ref mut to_goop, .. }))
                    if *to_goop > 1 =>
                {
                    *from_goop -= 1;
                    *to_goop -= 1;
                    to_outflows.clear();
                    // Since all outflow from `to` is cancelled, remove any
                    // pending pairs from `outflow_list`.
                    outflow_list.retain(|&(from, _)| from != to_index);
                },

                // Goop flowing into an occupied node, succeeds in clearing it.
                (&mut Some(Occupied { player, goop: ref mut from_goop, .. }),
                 &mut ref mut to) => {
                    *from_goop -= 1;
                    *to = Some(Occupied { player, outflows: vec![], goop: 0 });
                    // Since all outflow from `to` is cancelled, remove any
                    // pending pairs from `outflow_list`.
                    outflow_list.retain(|&(from, _)| from != to_index);
                }
            }
        }
    }

    /// Let sources generate new goop.
    fn generate_goop(&mut self) {
        for &source in &self.map.sources {
            match &mut self.nodes[source] {
                &mut None => panic!("source nodes should always be occupied by someone"),
                &mut Some(Occupied { ref mut goop, .. }) => {
                    if *goop < MAX_GOOP {
                        *goop += 1;
                    }
                }
            }
        }
    }

    /// Advance `self` to the next state.
    pub fn advance(&mut self) {
        self.flow();
        self.generate_goop();
    }

    /// Apply `action` to this state.
    pub fn take_action(&mut self, action: Action) {
        println!("take_action({:?})", action);
        match action {
            Action::ToggleOutflow { player, from, to } => {
                match &mut self.nodes[from] {
                    // This node is empty. Don't change it.
                    &mut None => (),

                    // Some other player owns this node. Do nothing.
                    &mut Some(Occupied { player: p, .. }) if p != player => (),

                    // We own this node. Toggle the given outflow.
                    &mut Some(Occupied { ref mut outflows, .. }) => {
                        if outflows.contains(&to) {
                            outflows.retain(|&dest| dest != to);
                        } else {
                            outflows.push(to);
                        }
                    }
                }
            }
        }
    }
}

/// Actions that can be taken on a `State`.
#[derive(Debug)]
pub enum Action {
    /// The `player` has requested to toggle the outflow
    /// from `from` to `to`.
    ToggleOutflow { player: Player, from: Node, to: Node },
}

/// A set of parameters that can be used to initialize a game.
pub struct GameParameters {
    /// The dimensions of the board.
    pub board: (usize, usize),

    /// The position of the sources on the board. The number of players is the
    /// length of this vector.
    pub sources: Vec<Node>,

    /// The color assigned to each player, as an RGB triplet. This must be the
    /// same length as `sources`.
    pub colors: Vec<(u8, u8, u8)>
}

<commit_msg>Factor out `simulate_flow()` and add a few tests.<commit_after>//! Module for the dynamic state of the game.
//!
//! This module defines the `State` type and its entourage. The data needed to
//! run an RBattle game is split into three categories:
//!
//! - `State` holds all the significant varying state of a game. What changes as the game proceeds
//!   are the amount of goop in each node, and the outflows the user has chosen;
//!   these are the `State`.
//!
//! - The positions of nodes and their boundaries and the locations of goop
//!   sources are all fixed, and are part of the `Map`.
//!
//! - Interface elements enter hover and active states as the user mouses
//!   around, but those states are ephemeral; they are part of the `Mouse` type.
//!   When the user actually completes an interaction with an interface element,
//!   only then is the `State` affected.

use graph::Node;
use map::Map;
use square::SquareGrid;
use std::sync::Arc;
use rand::{Rng, SeedableRng, XorShiftRng};
use xorshift::XorShift128Plus;

use std::iter::repeat;

/// The complete state of an RBattle game board.
#[derive(Clone)]
pub struct State {
    /// The map being played on.
    pub map: Arc<Map>,

    /// Which nodes are occupied, and which are vacant. Indexed by node id.
    pub nodes: Vec<Option<Occupied>>,

    /// The random number generator used to drive the goop flow algorithm.
    rng: XorShift128Plus
}

/// A player id number.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct Player(pub usize);

/// The state of a node that is occupied by some player.
#[derive(Debug, Clone, PartialEq)]
pub struct Occupied {
    /// The player who controls this node.
    pub player: Player,

    /// Which neighbors of this node it sends goop out to.
    pub outflows: Vec<Node>,

    /// How much goop this node holds. Ranges from 0 to MAX_GOOP.
    pub goop: usize,
}

pub const MAX_GOOP: usize = 15;

/// Return a pair of mutable references to the `i`'th and `j`'th elements of
/// `slice`, where `i != j`.
fn index_mut_pair<T>(slice: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
    if i < j {
        let (left, right) = slice.split_at_mut(j);
        (&mut left[i], &mut right[0])
    } else if j < i {
        let (left, right) = slice.split_at_mut(i);
        (&mut right[0], &mut left[j])
    } else {
        panic!("can't borrow two mutable references to the same element");
    }
}

impl State {
    pub fn new(params: GameParameters) -> State {
        let graph = SquareGrid::new(params.board.0, params.board.1);
        let map = Arc::new(Map::new(graph, params.sources, params.colors));

        let mut nodes: Vec<Option<Occupied>> = repeat(None).take(map.graph.nodes()).collect();
        // Ensure that each source is occupied by its player.
        for (player, &source) in map.sources.iter().enumerate() {
            nodes[source] = Some(Occupied {
                player: Player(player),
                outflows: vec![],
                goop: 0
            });
        }

        const SEED: [u64; 2] = [0xcd9d5eaaf04bc9a7, 0x4602cc7098d01ef9];
        State { map, nodes, rng: XorShift128Plus::new(SEED) }
    }

    /// Let one unit of goop flow through each outflow.
    ///
    /// There are algorithms for finding the flow through a graph precisely, but
    /// we need something simpler here. We just visit every outgoing edge in a
    /// random order, and propagate a unit of goop if the destination permits
    /// it.
    fn flow(&mut self) {
        // Build a vector of (from, to) pairs.
        let mut outflow_list = Vec::new();
        for node in 0..self.map.graph.nodes() {
            if let &Some(ref occupied) = &self.nodes[node] {
                for &outflow in &occupied.outflows {
                    outflow_list.push((node, outflow))
                }
            }
        }

        // Put the pairs in a random order.
        self.rng.shuffle(&mut outflow_list);

        while let Some((from_index, to_index)) = outflow_list.pop() {
            let (from_node, to_node) = index_mut_pair(&mut self.nodes, from_index, to_index);
            let attacked = simulate_flow(from_node, to_node);

            if attacked {
                // `to_node` is being attacked. Disregard any outflows from it this turn.
                outflow_list.retain(|&(from, _)| from != to_index);
            }
        }
    }

    /// Let sources generate new goop.
    fn generate_goop(&mut self) {
        for &source in &self.map.sources {
            match &mut self.nodes[source] {
                &mut None => panic!("source nodes should always be occupied by someone"),
                &mut Some(Occupied { ref mut goop, .. }) => {
                    if *goop < MAX_GOOP {
                        *goop += 1;
                    }
                }
            }
        }
    }

    /// Advance `self` to the next state.
    pub fn advance(&mut self) {
        self.flow();
        self.generate_goop();
    }

    /// Apply `action` to this state.
    pub fn take_action(&mut self, action: Action) {
        println!("take_action({:?})", action);
        match action {
            Action::ToggleOutflow { player, from, to } => {
                match &mut self.nodes[from] {
                    // This node is empty. Don't change it.
                    &mut None => (),

                    // Some other player owns this node. Do nothing.
                    &mut Some(Occupied { player: p, .. }) if p != player => (),

                    // We own this node. Toggle the given outflow.
                    &mut Some(Occupied { ref mut outflows, .. }) => {
                        if outflows.contains(&to) {
                            outflows.retain(|&dest| dest != to);
                        } else {
                            outflows.push(to);
                        }
                    }
                }
            }
        }
    }
}

/// Simulate goop flow from a given cell `from_node` to another cell, `to_node`.
///
/// This only simulates flow in that particular direction;
/// if two neighboring cells are both trying to ooze into each other,
/// then the State::flow() function will helpfully call this twice,
/// once for each direction.
///
/// See the tests to learn the detailed rules of this function's behavior.
///
/// Return true if `from_node` attacked `to_node`—that is, if the two nodes are
/// occupied by two different players, and any goop flowed. (The caller needs to
/// know about this, because in this case it must stop `to_node`'s outflows.)
///
fn simulate_flow(from_node: &mut Option<Occupied>, to_node: &mut Option<Occupied>) -> bool {
    match (from_node, to_node) {
        // We shouldn't have generated a pair for an empty source, and
        // when we clear a node we're supposed to remove pairs from
        // `outflow` that originate there.
        (&mut None, _) => panic!("outflow from empty node"),

        // Source has no goop. No effect.
        (&mut Some(Occupied { goop: 0, .. }), _) => false,

        // Goop flowing into an unoccupied node. New player claims ownership.
        (&mut Some(Occupied { player, ref mut goop, .. }),
         &mut ref mut to @ None) => {
            *goop -= 1;
            *to = Some(Occupied { player, outflows: vec![], goop: 1 });
            false
        },

        // Goop flowing into a node occupied by the same player.
        (&mut Some(Occupied { player: from_player, goop: ref mut from_goop, .. }),
         &mut Some(Occupied { player: to_player,   goop: ref mut to_goop, .. }))
            if from_player == to_player =>
        {
            if *from_goop > 0 && *to_goop < MAX_GOOP {
                *from_goop -= 1;
                *to_goop += 1;
            }
            false
        }

        // Goop flowing into a node occupied by another player, but
        // doesn't clear it. All outflow from destination stopped.
        (&mut Some(Occupied { goop: ref mut from_goop, .. }),
         &mut Some(Occupied { outflows: ref mut to_outflows,
                               goop:     ref mut to_goop, .. }))
            if *to_goop > 1 =>
        {
            *from_goop -= 1;
            *to_goop -= 1;
            to_outflows.clear();
            true
        },

        // Goop flowing into an occupied node, succeeds in clearing it.
        (&mut Some(Occupied { player, goop: ref mut from_goop, .. }),
         &mut ref mut to) => {
            *from_goop -= 1;
            *to = Some(Occupied { player, outflows: vec![], goop: 0 });
            true
        }
    }
}

#[test]
fn test_flow_into_unoccupied_cell() {
    // The kingdom of Florin is invading the kingdom of Guilder and flooding it with goop.
    let mut florin = Some(Occupied { player: Player(1), outflows: vec![2], goop: 15 });
    let mut guilder = None;

    // This isn't considered an attack, since Guilder was completely unoccupied.
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);

    // One unit of goop flowed. The newly occupied territory now belongs to player 1.
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2], goop: 14 }));
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![], goop: 1 }));
}

#[test]
fn test_flow_empties_cell() {
    // Same scenario, except this time there is just one unit of goop in Florin.
    let mut florin = Some(Occupied { player: Player(1), outflows: vec![2], goop: 1 });
    let mut guilder = None;

    // As above, this isn't considered an attack.
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);

    // One unit of goop flowed, leaving Florin emptied of goop (but still
    // considered occupied by player 1).
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![], goop: 1 }));

    // In the next tick of the game, no more goop flows, because Florin is now empty.
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);

    // The state after this second simulated step is therefore exactly the same as before.
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![], goop: 1 }));
}

#[test]
fn test_no_flow_from_empty_cell() {
    // Florin, alas, has ambitions of conquest but no goop to carry them out with.
    let mut florin = Some(Occupied { player: Player(1), outflows: vec![2, 3, 4], goop: 0 });

    // Florin can try to flow into a cell that's never been occupied, but since
    // Florin has no goop, the cell does *not* become occupied.
    let mut zolot = None;
    assert_eq!(simulate_flow(&mut florin, &mut zolot), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3, 4], goop: 0 }));
    assert_eq!(zolot, None);

    // Nor can Florin attack a cell occupied by an opponent with goop.
    let mut guilder = Some(Occupied { player: Player(2), outflows: vec![5, 6], goop: 1 });
    // The attempt does not count as an attack, since no goop flowed.
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3, 4], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(2), outflows: vec![5, 6], goop: 1 }));

    // Same deal even if Guilder also has no goop.
    guilder = Some(Occupied { player: Player(2), outflows: vec![5, 6], goop: 0 });
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3, 4], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(2), outflows: vec![5, 6], goop: 0 }));

    // Nor can Florin reinforce Guilder if they happen to be occupied by the same player.
    guilder = Some(Occupied { player: Player(1), outflows: vec![5, 6], goop: 0 });
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3, 4], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![5, 6], goop: 0 }));
}

#[test]
fn test_friendly_flow() {
    // Florin is sending goop into the friendly neighboring province of Krugerrand.
    let mut florin     = Some(Occupied { player: Player(1), outflows: vec![2], goop: 8 });
    let mut krugerrand = Some(Occupied { player: Player(1), outflows: vec![3], goop: 0 });
    // This isn't an attack, since the same player occupies both.
    assert_eq!(simulate_flow(&mut florin, &mut krugerrand), false);
    // A unit of goop actually flowed.
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2], goop: 7 }));
    // Krugerrand received the goop, and its outflows are unaffected.
    assert_eq!(krugerrand, Some(Occupied { player: Player(1), outflows: vec![3], goop: 1 }));

    // It works even if the destination already has more goop than Florin.
    krugerrand.as_mut().unwrap().goop = 9;
    assert_eq!(simulate_flow(&mut florin, &mut krugerrand), false);
    assert_eq!(florin.unwrap().goop, 6);
    assert_eq!(krugerrand.unwrap().goop, 10);
}

#[test]
fn test_friendly_flow_empties_cell() {
    // Florin can send its only unit of goop to a neighboring empty or nonempty friendly cell.
    let mut florin = Some(Occupied { player: Player(1), outflows: vec![2, 3], goop: 1 });
    let mut guilder = Some(Occupied { player: Player(1), outflows: vec![4], goop: 0 });
    // This isn't an attack, since the same player occupies both.
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);
    // A unit of goop actually flowed.
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3], goop: 0 }));
    // Guilder received the goop, and its outflows are unaffected.
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![4], goop: 1 }));

    // Restore Florin's 1 unit of goop and try again.
    florin.as_mut().unwrap().goop = 1;
    assert_eq!(simulate_flow(&mut florin, &mut guilder), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3], goop: 0 }));
    assert_eq!(guilder, Some(Occupied { player: Player(1), outflows: vec![4], goop: 2 }));
}

#[test]
fn test_friendly_flow_max_goop() {
    // Florin can't pump any more goop into its friendly neighbor Pfennig,
    // which has the maximum amount already.
    let mut florin  = Some(Occupied { player: Player(1), outflows: vec![2, 3], goop: 3 });
    let mut pfennig = Some(Occupied { player: Player(1), outflows: vec![4], goop: MAX_GOOP });

    assert_eq!(simulate_flow(&mut florin, &mut pfennig), false);
    assert_eq!(florin, Some(Occupied { player: Player(1), outflows: vec![2, 3], goop: 3 }));
    assert_eq!(pfennig, Some(Occupied { player: Player(1), outflows: vec![4], goop: MAX_GOOP }));
}

#[test]
#[should_panic]
fn test_flow_from_unoccupied_cell() {
    // This should never happen, because the simulator won't pass an unoccupied
    // cell to simulate_flow().
    let mut zolot = None;
    let mut zorkmid = None;
    simulate_flow(&mut zolot, &mut zorkmid);
}

/// Actions that can be taken on a `State`.
#[derive(Debug)]
pub enum Action {
    /// The `player` has requested to toggle the outflow
    /// from `from` to `to`.
    ToggleOutflow { player: Player, from: Node, to: Node },
}

/// A set of parameters that can be used to initialize a game.
pub struct GameParameters {
    /// The dimensions of the board.
    pub board: (usize, usize),

    /// The position of the sources on the board. The number of players is the
    /// length of this vector.
    pub sources: Vec<Node>,

    /// The color assigned to each player, as an RGB triplet. This must be the
    /// same length as `sources`.
    pub colors: Vec<(u8, u8, u8)>
}

<|endoftext|>


<commit_before>extern crate stemmer;
extern crate unicode_normalization;
extern crate unicode_segmentation;

use std::iter::Peekable;

use self::stemmer::Stemmer;
use self::unicode_normalization::UnicodeNormalization;
use self::unicode_segmentation::UnicodeSegmentation;

pub struct Stems<'a> {
    words: Peekable<unicode_segmentation::UWordBoundIndices<'a>>,
    stemmer: Stemmer,
}

#[derive(Debug, PartialEq)]
pub struct StemmedWord {
    // Where the stemmed word starts
    pub stemmed_offset: usize,
    // Where the suffix starts
    pub suffix_offset: usize,
    // The stemmed word
    pub stemmed: String,
    // The difference between the stemmed word and the original lowercased one. It can be
    // used to recontruct the original word (for exact match searches)
    pub suffix: String,
}

impl<'a> Stems<'a> {
    pub fn new(text: &str) -> Stems {
        Stems{
            words: text.split_word_bound_indices().peekable(),
            stemmer: Stemmer::new("english").unwrap(),
        }
    }

    /// Return the *byte* length of the common prefix between two strings
    fn common_prefix_len(aa: &str, bb: &str) -> usize {
        aa.chars()
            .zip(bb.chars())
            .take_while(|&(a, b)| a == b)
            .fold(0, |acc, (a, _)| acc + a.len_utf8())
    }
}

impl<'a> Iterator for Stems<'a> {
    type Item = StemmedWord;

    fn next(&mut self) -> Option<StemmedWord> {
        let mut word_to_stem = String::new();
        let mut stemmed_offset = 0;
        let mut normalized = String::new();
        loop {
            match self.words.peek() {
                Some(&(pos, word)) => {
                    normalized = word.nfkc().collect::<String>();
                    if word.chars().next().unwrap().is_alphabetic() {
                        stemmed_offset = pos;
                        break;
                    } else {
                        word_to_stem.push_str(&normalized);
                        self.words.next();
                    }
                },
                None => {
                    if word_to_stem.is_empty() {
                        return None;
                    } else {
                        break;
                    }
                },
            }
        }

        if !word_to_stem.is_empty() {
            // we found the begining of the string is not a stemmable word.
            // Return the accumulated string as the stemmed word
            return Some(StemmedWord {
                            stemmed_offset: 0,
                            suffix_offset: word_to_stem.len(),
                            stemmed: word_to_stem,
                            suffix: String::new(),
                });
        }
        // normalized contains our stemmable word. advance the iter since we only peeked.
        self.words.next();
        word_to_stem = normalized;
        let mut suffix = word_to_stem.clone();
        loop {
            // loop through all non-alphabetic chars and add to suffix (which means the suffix starts
            // before the stemmed word)
            match self.words.peek() {
                Some(&(_pos, word)) => {
                    normalized = word.nfkc().collect::<String>();
                    if normalized.chars().next().unwrap().is_alphabetic() {
                        break;
                    } else {
                        suffix.push_str(&normalized);
                        self.words.next();
                    }
                },
                None => break,
            }
        }
        let stemmed = self.stemmer.stem(&word_to_stem.to_lowercase());
        let prefix_len = Stems::common_prefix_len(&stemmed, &suffix);
        Some(StemmedWord {
            stemmed_offset: stemmed_offset,
            suffix_offset: stemmed_offset + prefix_len,
            stemmed: stemmed,
            suffix: (&suffix[prefix_len..]).to_string(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{StemmedWord, Stems};

    #[test]
    fn test_stems_lowercase() {
        let input = "These words deeply test smoothly that stemming";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 5,
                          stemmed: String::from("these"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 6, suffix_offset: 10,
                          stemmed: String::from("word"), suffix: String::from("s") },
            // "deeply" stems to "deepli"
            StemmedWord { stemmed_offset: 12, suffix_offset: 17,
                          stemmed: String::from("deepli"), suffix: String::from("y") },
            StemmedWord { stemmed_offset: 19, suffix_offset: 23,
                          stemmed: String::from("test"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 24, suffix_offset: 30,
                          stemmed: String::from("smooth"), suffix: String::from("ly") },
            StemmedWord { stemmed_offset: 33, suffix_offset: 37,
                          stemmed: String::from("that"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 38, suffix_offset: 42,
                          stemmed: String::from("stem"), suffix: String::from("ming") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_mixedcase() {
        let input = "THEse Words deeplY test smOOthly that stemmING";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 5,
                          stemmed: String::from("these"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 6, suffix_offset: 10,
                          stemmed: String::from("word"), suffix: String::from("s") },
            // "deeply" stems to "deepli"
            StemmedWord { stemmed_offset: 12, suffix_offset: 17,
                          stemmed: String::from("deepli"), suffix: String::from("y") },
            StemmedWord { stemmed_offset: 19, suffix_offset: 23,
                          stemmed: String::from("test"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 24, suffix_offset: 30,
                          stemmed: String::from("smooth"), suffix: String::from("ly") },
            StemmedWord { stemmed_offset: 33, suffix_offset: 37,
                          stemmed: String::from("that"), suffix: String::from("") },
            StemmedWord { stemmed_offset: 38, suffix_offset: 42,
                          stemmed: String::from("stem"), suffix: String::from("ming") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_nonchars() {
        let input = "  @#$!== \t+-";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        assert_eq!(result.len(), 0);
    }

    #[test]
    fn test_stems_some_nonchars() {
        let input = "@!?   Let's seeing...";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 6, suffix_offset: 9,
                          stemmed: String::from("let"), suffix: String::from("'s") },
            StemmedWord { stemmed_offset: 12, suffix_offset: 15,
                          stemmed: String::from("see"), suffix: String::from("ing") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_unicode() {
        let input = "Ünicöde stemming";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 8,
                          stemmed: String::from("ünicöd"), suffix: String::from("e") },
            StemmedWord { stemmed_offset: 10, suffix_offset: 14,
                          stemmed: String::from("stem"), suffix: String::from("ming") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_unicode_lowercase_has_more_bytes() {
        let input = "İ";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 3,
                          stemmed: String::from("i̇"), suffix: String::from("") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    // Some characters don't have a corresponding uppercase chararacter with the same byte size,
    // but when those uppercase letters are transformed to lowercase, they are transformed to the
    // same number of bytes combining several characters (although there is a single character for
    // it).
    // Here's an example in Python:
    //
    //     >>> u"Ρ̓".lower() == u"ῤ"
    //     False
    //     >>> u"Ρ̓".lower() == u"ῤ".lower()
    //     False
    //     >>> u"Ρ̓".lower() == u"ῤ".upper().lower()
    //     True
    // Or Rust:
    //
    //    let upper = "Ρ̓";
    //    let lower = "ῤ";
    //    println!("lower({}) == {}: {}", upper, lower, upper.to_lowercase() == lower);
    //    println!("lower({}) == lower(upper({})): {}", upper, lower, upper.to_lowercase() == lower.to_uppercase().to_lowercase());
    //    lower(Ρ̓) == ῤ: false
    //    lower(Ρ̓) == lower(upper(ῤ)): true
    #[test]
    fn test_stems_unicode_lowercase_has_less_bytes() {
        // The input is: Ρ̓ῤῤ (11 bytes), lowercases is ῤῤῤ (9 bytes)
        let input = "\u{03A1}\u{0313}\u{03C1}\u{0313}\u{1FE4}";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 9,
                          stemmed: String::from("\u{1FE4}\u{1FE4}\u{1FE4}"), suffix: String::from("") },
           ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_common_prefix_len() {
        let tests = vec![
            ("a", "a", 1),
            ("ab", "a", 1),
            ("a", "ab", 1),
            ("ab", "ab", 2),
            ("a", "b", 0),
            ("b", "a", 0),
            ("ab", "cd", 0),
            ("ab", "bc", 0),
            ("abc", "abd", 2),
            ("ac", "abcd", 1),
            (" a", "a", 0),
            ("a", "a ", 1),
            ("xyzabc", "xyz", 3),
            ("xyz", "xyzabc", 3),
            ("öxyz", "öx", 3),
            ];
        for (aa, bb, expected) in tests {
            let prefix_len = Stems::common_prefix_len(aa, bb);
            assert_eq!(prefix_len, expected);
        }
    }
}
<commit_msg>Fixed broken tests in stems.rs<commit_after>extern crate stemmer;
extern crate unicode_normalization;
extern crate unicode_segmentation;

use std::iter::Peekable;

use self::stemmer::Stemmer;
use self::unicode_normalization::UnicodeNormalization;
use self::unicode_segmentation::UnicodeSegmentation;

pub struct Stems<'a> {
    words: Peekable<unicode_segmentation::UWordBoundIndices<'a>>,
    stemmer: Stemmer,
}

#[derive(Debug, PartialEq)]
pub struct StemmedWord {
    // Where the stemmed word starts
    pub stemmed_offset: usize,
    // Where the suffix starts
    pub suffix_offset: usize,
    // The stemmed word
    pub stemmed: String,
    // The difference between the stemmed word and the original lowercased one. It can be
    // used to recontruct the original word (for exact match searches)
    pub suffix: String,
}

impl<'a> Stems<'a> {
    pub fn new(text: &str) -> Stems {
        Stems{
            words: text.split_word_bound_indices().peekable(),
            stemmer: Stemmer::new("english").unwrap(),
        }
    }

    /// Return the *byte* length of the common prefix between two strings
    fn common_prefix_len(aa: &str, bb: &str) -> usize {
        aa.chars()
            .zip(bb.chars())
            .take_while(|&(a, b)| a == b)
            .fold(0, |acc, (a, _)| acc + a.len_utf8())
    }
}

impl<'a> Iterator for Stems<'a> {
    type Item = StemmedWord;

    fn next(&mut self) -> Option<StemmedWord> {
        let mut word_to_stem = String::new();
        let mut stemmed_offset = 0;
        let mut normalized = String::new();
        loop {
            match self.words.peek() {
                Some(&(pos, word)) => {
                    normalized = word.nfkc().collect::<String>();
                    if word.chars().next().unwrap().is_alphabetic() {
                        stemmed_offset = pos;
                        break;
                    } else {
                        word_to_stem.push_str(&normalized);
                        self.words.next();
                    }
                },
                None => {
                    if word_to_stem.is_empty() {
                        return None;
                    } else {
                        break;
                    }
                },
            }
        }

        if !word_to_stem.is_empty() {
            // we found the begining of the string is not a stemmable word.
            // Return the accumulated string as the stemmed word
            return Some(StemmedWord {
                            stemmed_offset: 0,
                            suffix_offset: word_to_stem.len(),
                            stemmed: word_to_stem,
                            suffix: String::new(),
                });
        }
        // normalized contains our stemmable word. advance the iter since we only peeked.
        self.words.next();
        word_to_stem = normalized;
        let mut suffix = word_to_stem.clone();
        loop {
            // loop through all non-alphabetic chars and add to suffix (which means the suffix starts
            // before the stemmed word)
            match self.words.peek() {
                Some(&(_pos, word)) => {
                    normalized = word.nfkc().collect::<String>();
                    if normalized.chars().next().unwrap().is_alphabetic() {
                        break;
                    } else {
                        suffix.push_str(&normalized);
                        self.words.next();
                    }
                },
                None => break,
            }
        }
        let stemmed = self.stemmer.stem(&word_to_stem.to_lowercase());
        let prefix_len = Stems::common_prefix_len(&stemmed, &suffix);
        Some(StemmedWord {
            stemmed_offset: stemmed_offset,
            suffix_offset: stemmed_offset + prefix_len,
            stemmed: stemmed,
            suffix: (&suffix[prefix_len..]).to_string(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{StemmedWord, Stems};

    #[test]
    fn test_stems_mixedcase() {
        let input = "THEse Words deeplY test smOOthly that stemmING";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 0,
                          stemmed: String::from("these"), suffix: String::from("THEse ") },
            StemmedWord { stemmed_offset: 6, suffix_offset: 6,
                          stemmed: String::from("word"), suffix: String::from("Words ") },
            // "deeply" stems to "deepli"
            StemmedWord { stemmed_offset: 12, suffix_offset: 17,
                          stemmed: String::from("deepli"), suffix: String::from("Y ") },
            StemmedWord { stemmed_offset: 19, suffix_offset: 23,
                          stemmed: String::from("test"), suffix: String::from(" ") },
            StemmedWord { stemmed_offset: 24, suffix_offset: 26,
                          stemmed: String::from("smooth"), suffix: String::from("OOthly ") },
            StemmedWord { stemmed_offset: 33, suffix_offset: 37,
                          stemmed: String::from("that"), suffix: String::from(" ") },
            StemmedWord { stemmed_offset: 38, suffix_offset: 42,
                          stemmed: String::from("stem"), suffix: String::from("mING") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_nonchars() {
        let input = "  @#$!== \t+-";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        assert_eq!(result, vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 12,
                          stemmed: String::from("  @#$!== \t+-"), suffix: String::from("") },
            ]);
    }

    #[test]
    fn test_stems_some_nonchars() {
        let input = "@!?   Let's seeing...";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 6,
                          stemmed: String::from("@!?   "), suffix: String::from("") },
            StemmedWord { stemmed_offset: 6, suffix_offset: 6,
                          stemmed: String::from("let"), suffix: String::from("Let's ") },
            StemmedWord { stemmed_offset: 12, suffix_offset: 15,
                          stemmed: String::from("see"), suffix: String::from("ing...") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_unicode() {
        let input = "Ünicöde stemming";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 0,
                          stemmed: String::from("ünicöd"), suffix: String::from("Ünicöde ") },
            StemmedWord { stemmed_offset: 10, suffix_offset: 14,
                          stemmed: String::from("stem"), suffix: String::from("ming") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_stems_unicode_lowercase_has_more_bytes() {
        let input = "İ";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 0,
                          stemmed: String::from("i̇"), suffix: String::from("İ") },
            ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    // Some characters don't have a corresponding uppercase chararacter with the same byte size,
    // but when those uppercase letters are transformed to lowercase, they are transformed to the
    // same number of bytes combining several characters (although there is a single character for
    // it).
    // Here's an example in Python:
    //
    //     >>> u"Ρ̓".lower() == u"ῤ"
    //     False
    //     >>> u"Ρ̓".lower() == u"ῤ".lower()
    //     False
    //     >>> u"Ρ̓".lower() == u"ῤ".upper().lower()
    //     True
    // Or Rust:
    //
    //    let upper = "Ρ̓";
    //    let lower = "ῤ";
    //    println!("lower({}) == {}: {}", upper, lower, upper.to_lowercase() == lower);
    //    println!("lower({}) == lower(upper({})): {}", upper, lower, upper.to_lowercase() == lower.to_uppercase().to_lowercase());
    //    lower(Ρ̓) == ῤ: false
    //    lower(Ρ̓) == lower(upper(ῤ)): true
    #[test]
    fn test_stems_unicode_lowercase_has_less_bytes() {
        // The input is: Ρ̓ῤῤ (11 bytes), lowercases is ῤῤῤ (9 bytes)
        let input = "\u{03A1}\u{0313}\u{03C1}\u{0313}\u{1FE4}";
        let result = Stems::new(input).collect::<Vec<StemmedWord>>();
        let expected = vec![
            StemmedWord { stemmed_offset: 0, suffix_offset: 0,
                          stemmed: String::from("\u{1FE4}\u{1FE4}\u{1FE4}"),
                          suffix: String::from("\u{03A1}\u{0313}\u{03C1}\u{0313}\u{03C1}") },
           ];
        assert_eq!(result.len(), expected.len());
        for (stem, expected_stem) in result.iter().zip(expected.iter()) {
            assert_eq!(stem, expected_stem);
        }
    }

    #[test]
    fn test_common_prefix_len() {
        let tests = vec![
            ("a", "a", 1),
            ("ab", "a", 1),
            ("a", "ab", 1),
            ("ab", "ab", 2),
            ("a", "b", 0),
            ("b", "a", 0),
            ("ab", "cd", 0),
            ("ab", "bc", 0),
            ("abc", "abd", 2),
            ("ac", "abcd", 1),
            (" a", "a", 0),
            ("a", "a ", 1),
            ("xyzabc", "xyz", 3),
            ("xyz", "xyzabc", 3),
            ("öxyz", "öx", 3),
            ];
        for (aa, bb, expected) in tests {
            let prefix_len = Stems::common_prefix_len(aa, bb);
            assert_eq!(prefix_len, expected);
        }
    }
}
<|endoftext|>