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

use std::sync::Arc;
use std::net::SocketAddr;

use core::core;
use core::core::hash::{Hash, Hashed};
use core::ser;
use conn::*;
use msg::*;
use rand;
use rand::Rng;
use types::*;
use util::LOGGER;

pub struct Protocol {
	adapter: Arc<NetAdapter>,
	addr: SocketAddr,
}

impl Protocol {
	pub fn new(adapter: Arc<NetAdapter>, addr: SocketAddr) -> Protocol {
		Protocol{adapter, addr}
	}
}

impl MessageHandler for Protocol {
	fn consume(&self, msg: &mut Message) -> Result<Option<(Vec<u8>, Type)>, Error> {
		let adapter = &self.adapter;

		match msg.header.msg_type {

			Type::Ping => {
				let ping: Ping = msg.body()?;
				adapter.peer_difficulty(self.addr, ping.total_difficulty, ping.height);

				let pong_bytes = ser::ser_vec(
					&Pong {
						total_difficulty: adapter.total_difficulty(),
						height: adapter.total_height(),
					}).unwrap();

				Ok(Some((pong_bytes, Type::Pong)))
			}

			Type::Pong => {
				let pong: Pong = msg.body()?;
				adapter.peer_difficulty(self.addr, pong.total_difficulty, pong.height);
				Ok(None)
			},

			Type::Transaction => {
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx);
				Ok(None)
			}

			Type::GetBlock => {
				let h: Hash = msg.body()?;
				debug!(LOGGER, "handle_payload: GetBlock {}", h);

				let bo = adapter.get_block(h);
				if let Some(b) = bo {
					let block_bytes = ser::ser_vec(&b).unwrap();
					return Ok(Some((block_bytes, Type::Block)));
				}
				Ok(None)
			}

			Type::Block => {
				let b: core::Block = msg.body()?;
				let bh = b.hash();

				debug!(LOGGER, "handle_payload: Block {}", bh);

				adapter.block_received(b, self.addr);
				Ok(None)
			}

			Type::GetCompactBlock => {
				let h: Hash = msg.body()?;
				debug!(LOGGER, "handle_payload: GetCompactBlock: {}", h);

				if let Some(b) = adapter.get_block(h) {
					let cb = b.as_compact_block();

					// serialize and send the block over in compact representation

					// if we have txs in the block send a compact block
					// but if block is empty -
					// to allow us to test all code paths, randomly choose to send
					// either the block or the compact block
					let mut rng = rand::thread_rng();

					if cb.kern_ids.is_empty() && rng.gen() {
						debug!(
							LOGGER,
							"handle_payload: GetCompactBlock: empty block, sending full block",
							);

						let block_bytes = ser::ser_vec(&b).unwrap();
						Ok(Some((block_bytes, Type::Block)))
					} else {
						let compact_block_bytes = ser::ser_vec(&cb).unwrap();
						Ok(Some((compact_block_bytes, Type::CompactBlock)))
					}
				} else {
					Ok(None)
				}
			}

			Type::CompactBlock => {
				let b: core::CompactBlock = msg.body()?;
				let bh = b.hash();
				debug!(LOGGER, "handle_payload: CompactBlock: {}", bh);

				adapter.compact_block_received(b, self.addr);
				Ok(None)
			}

			Type::GetHeaders => {
				// load headers from the locator
				let loc: Locator = msg.body()?;
				let headers = adapter.locate_headers(loc.hashes);

				// serialize and send all the headers over
				let header_bytes = ser::ser_vec(&Headers { headers: headers }).unwrap();
				return Ok(Some((header_bytes, Type::Headers)));
			}

			Type::Headers => {
				let headers: Headers = msg.body()?;
				adapter.headers_received(headers.headers, self.addr);
				Ok(None)
			}

			Type::GetPeerAddrs => {
				let get_peers: GetPeerAddrs = msg.body()?;
				let peer_addrs = adapter.find_peer_addrs(get_peers.capabilities);
				let peer_addrs_bytes = ser::ser_vec(
					&PeerAddrs {
						peers: peer_addrs.iter().map(|sa| SockAddr(*sa)).collect(),
					}).unwrap();
				return Ok(Some((peer_addrs_bytes, Type::PeerAddrs)));
			}

			Type::PeerAddrs => {
				let peer_addrs: PeerAddrs = msg.body()?;
				adapter.peer_addrs_received(peer_addrs.peers.iter().map(|pa| pa.0).collect());
				Ok(None)
			}

			_ => {
				debug!(LOGGER, "unknown message type {:?}", msg.header.msg_type);
				Ok(None)
			}
		}
	}
}
<commit_msg>fix bad merge - missing support for header msg type (#690)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::net::SocketAddr;

use core::core;
use core::core::hash::{Hash, Hashed};
use core::ser;
use conn::*;
use msg::*;
use rand;
use rand::Rng;
use types::*;
use util::LOGGER;

pub struct Protocol {
	adapter: Arc<NetAdapter>,
	addr: SocketAddr,
}

impl Protocol {
	pub fn new(adapter: Arc<NetAdapter>, addr: SocketAddr) -> Protocol {
		Protocol{adapter, addr}
	}
}

impl MessageHandler for Protocol {
	fn consume(&self, msg: &mut Message) -> Result<Option<(Vec<u8>, Type)>, Error> {
		let adapter = &self.adapter;

		match msg.header.msg_type {

			Type::Ping => {
				let ping: Ping = msg.body()?;
				adapter.peer_difficulty(self.addr, ping.total_difficulty, ping.height);

				let pong_bytes = ser::ser_vec(
					&Pong {
						total_difficulty: adapter.total_difficulty(),
						height: adapter.total_height(),
					}).unwrap();

				Ok(Some((pong_bytes, Type::Pong)))
			}

			Type::Pong => {
				let pong: Pong = msg.body()?;
				adapter.peer_difficulty(self.addr, pong.total_difficulty, pong.height);
				Ok(None)
			},

			Type::Transaction => {
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx);
				Ok(None)
			}

			Type::GetBlock => {
				let h: Hash = msg.body()?;
				debug!(LOGGER, "handle_payload: GetBlock {}", h);

				let bo = adapter.get_block(h);
				if let Some(b) = bo {
					let block_bytes = ser::ser_vec(&b).unwrap();
					return Ok(Some((block_bytes, Type::Block)));
				}
				Ok(None)
			}

			Type::Block => {
				let b: core::Block = msg.body()?;
				let bh = b.hash();

				debug!(LOGGER, "handle_payload: Block {}", bh);

				adapter.block_received(b, self.addr);
				Ok(None)
			}

			Type::GetCompactBlock => {
				let h: Hash = msg.body()?;
				debug!(LOGGER, "handle_payload: GetCompactBlock: {}", h);

				if let Some(b) = adapter.get_block(h) {
					let cb = b.as_compact_block();

					// serialize and send the block over in compact representation

					// if we have txs in the block send a compact block
					// but if block is empty -
					// to allow us to test all code paths, randomly choose to send
					// either the block or the compact block
					let mut rng = rand::thread_rng();

					if cb.kern_ids.is_empty() && rng.gen() {
						debug!(
							LOGGER,
							"handle_payload: GetCompactBlock: empty block, sending full block",
							);

						let block_bytes = ser::ser_vec(&b).unwrap();
						Ok(Some((block_bytes, Type::Block)))
					} else {
						let compact_block_bytes = ser::ser_vec(&cb).unwrap();
						Ok(Some((compact_block_bytes, Type::CompactBlock)))
					}
				} else {
					Ok(None)
				}
			}

			Type::CompactBlock => {
				let b: core::CompactBlock = msg.body()?;
				let bh = b.hash();
				debug!(LOGGER, "handle_payload: CompactBlock: {}", bh);

				adapter.compact_block_received(b, self.addr);
				Ok(None)
			}

			Type::GetHeaders => {
				// load headers from the locator
				let loc: Locator = msg.body()?;
				let headers = adapter.locate_headers(loc.hashes);

				// serialize and send all the headers over
				let header_bytes = ser::ser_vec(&Headers { headers: headers }).unwrap();
				return Ok(Some((header_bytes, Type::Headers)));
			}

			// "header first" block propagation - if we have not yet seen this block
			// we can go request it from some of our peers
			Type::Header => {
				let header: core::BlockHeader = msg.body()?;

				adapter.header_received(header, self.addr);

				// we do not return a hash here as we never request a single header
				// a header will always arrive unsolicited
				Ok(None)
			}

			Type::Headers => {
				let headers: Headers = msg.body()?;
				adapter.headers_received(headers.headers, self.addr);
				Ok(None)
			}

			Type::GetPeerAddrs => {
				let get_peers: GetPeerAddrs = msg.body()?;
				let peer_addrs = adapter.find_peer_addrs(get_peers.capabilities);
				let peer_addrs_bytes = ser::ser_vec(
					&PeerAddrs {
						peers: peer_addrs.iter().map(|sa| SockAddr(*sa)).collect(),
					}).unwrap();
				return Ok(Some((peer_addrs_bytes, Type::PeerAddrs)));
			}

			Type::PeerAddrs => {
				let peer_addrs: PeerAddrs = msg.body()?;
				adapter.peer_addrs_received(peer_addrs.peers.iter().map(|pa| pa.0).collect());
				Ok(None)
			}

			_ => {
				debug!(LOGGER, "unknown message type {:?}", msg.header.msg_type);
				Ok(None)
			}
		}
	}
}
<|endoftext|>


<commit_before>extern crate cobalt;

use std::net;
use std::thread;
use std::collections::HashMap;
use cobalt::{Client, Connection, ConnectionID, Config, Handler, Server};

#[test]
fn test_client_connection_failure() {

    let config = Config::default();
    let mut handler = MockClientHandler::new();
    let mut client = Client::new(config);
    client.connect(&mut handler, "127.0.0.1:12345").unwrap();

    assert_eq!(handler.connect_calls, 1);
    assert!(handler.tick_connection_calls > 0);
    assert_eq!(handler.close_calls, 1);

    assert_eq!(handler.connection_calls, 0);
    assert_eq!(handler.connection_failed_calls, 1);
    assert_eq!(handler.connection_congestion_state_calls, 0);
    assert_eq!(handler.connection_packet_lost_calls, 0);
    assert_eq!(handler.connection_lost_calls, 0);

}

#[test]
fn test_server_bind_and_shutdown() {

    let config = Config::default();
    let mut handler = MockServerHandler::new(0, false);
    let mut server = Server::new(config);
    server.bind(&mut handler, "127.0.0.1:0").unwrap();

    assert_eq!(handler.bind_calls, 1);
    assert!(handler.tick_connections_calls > 0);
    assert_eq!(handler.shutdown_calls, 1);

    assert_eq!(handler.connection_calls, 0);
    assert_eq!(handler.connection_failed_calls, 0);
    assert_eq!(handler.connection_congestion_state_calls, 0);
    assert_eq!(handler.connection_packet_lost_calls, 0);
    assert_eq!(handler.connection_lost_calls, 0);

}

#[test]
fn test_server_client_connection_lost() {

    // Get a free local socket and then drop it for quick re-use
    // this is not 100% safe but we cannot easily get the locally bound server
    // address after bind() has been called
    let address: Option<net::SocketAddr> = {
        Some(net::UdpSocket::bind("127.0.0.1:0").unwrap().local_addr().unwrap())
    };

    // Setup Test Server
    let server_address = address.clone();
    let server_thread = thread::spawn(move|| {

        let config = Config::default();
        let mut server_handler = MockServerHandler::new(35, false);
        let mut server = Server::new(config);
        server.bind(&mut server_handler, server_address.unwrap()).unwrap();

        assert_eq!(server_handler.bind_calls, 1);
        assert!(server_handler.tick_connections_calls > 0);
        assert_eq!(server_handler.shutdown_calls, 1);

        assert_eq!(server_handler.connection_calls, 1);
        assert_eq!(server_handler.connection_failed_calls, 0);
        assert_eq!(server_handler.connection_congestion_state_calls, 0);
        assert_eq!(server_handler.connection_packet_lost_calls, 0);
        assert_eq!(server_handler.connection_lost_calls, 0);
        assert_eq!(server_handler.connection_closed_calls, 0);

    });

    let config = Config::default();
    let mut client_handler = MockClientHandler::new();
    let mut client = Client::new(config);
    client.connect(&mut client_handler, address.unwrap()).unwrap();

    server_thread.join().unwrap();
    assert_eq!(client_handler.connect_calls, 1);
    assert!(client_handler.tick_connection_calls > 0);
    assert_eq!(client_handler.close_calls, 1);

    assert_eq!(client_handler.connection_calls, 1);
    assert_eq!(client_handler.connection_failed_calls, 0);
    assert_eq!(client_handler.connection_congestion_state_calls, 0);
    // This is somewhat random and depends on how excatly the two threads
    // interact
    // assert_eq!(client_handler.connection_packet_lost_calls, 0);
    assert_eq!(client_handler.connection_lost_calls, 1);
    assert_eq!(client_handler.connection_closed_calls, 0);

}

#[test]
fn test_server_client_connection_close() {

    // Get a free local socket and then drop it for quick re-use
    // this is not 100% safe but we cannot easily get the locally bound server
    // address after bind() has been called
    let address: Option<net::SocketAddr> = {
        Some(net::UdpSocket::bind("127.0.0.1:0").unwrap().local_addr().unwrap())
    };

    // Setup Test Server
    let server_address = address.clone();
    let server_thread = thread::spawn(move|| {

        let config = Config::default();
        let mut server_handler = MockServerHandler::new(35, true);
        let mut server = Server::new(config);
        server.bind(&mut server_handler, server_address.unwrap()).unwrap();

        assert_eq!(server_handler.connection_closed_calls, 1);
        assert_eq!(server_handler.closed_by_remote, false);

    });

    let config = Config::default();
    let mut client_handler = MockClientHandler::new();
    let mut client = Client::new(config);
    client.connect(&mut client_handler, address.unwrap()).unwrap();

    server_thread.join().unwrap();
    assert_eq!(client_handler.connection_closed_calls, 1);
    assert_eq!(client_handler.closed_by_remote, true);

}

// Client Mock ----------------------------------------------------------------
pub struct MockClientHandler {
    pub connect_calls: u32,
    pub tick_connection_calls: u32,
    pub close_calls: u32,

    pub connection_calls: u32,
    pub connection_failed_calls: u32,
    pub connection_congestion_state_calls: u32,
    pub connection_packet_lost_calls: u32,
    pub connection_lost_calls: u32,
    pub connection_closed_calls: u32,
    pub closed_by_remote: bool
}

impl MockClientHandler {
    pub fn new() -> MockClientHandler {
        MockClientHandler {
            connect_calls: 0,
            tick_connection_calls: 0,
            close_calls: 0,

            connection_calls: 0,
            connection_failed_calls: 0,
            connection_congestion_state_calls: 0,
            connection_packet_lost_calls: 0,
            connection_lost_calls: 0,
            connection_closed_calls: 0,
            closed_by_remote: false
        }
    }
}

impl Handler<Client> for MockClientHandler {

    fn connect(&mut self, _: &mut Client) {
        self.connect_calls += 1;
    }

    fn tick_connection(&mut self, _: &mut Client, _: &mut Connection) {
        self.tick_connection_calls += 1;
    }

    fn close(&mut self, _: &mut Client) {
        self.close_calls += 1;
    }

    fn connection(&mut self, _: &mut Client, _: &mut Connection) {
        self.connection_calls += 1;
    }

    fn connection_failed(&mut self, client: &mut Client, _: &mut Connection) {
        self.connection_failed_calls += 1;
        client.close().unwrap();
    }

    fn connection_packet_lost(
        &mut self, _: &mut Client, _: &mut Connection, _: &[u8]
    ) {
        self.connection_packet_lost_calls += 1;
    }

    fn connection_congestion_state(&mut self, _: &mut Client, _: &mut Connection, _: bool) {
        self.connection_congestion_state_calls += 1;
    }

    fn connection_lost(&mut self, client: &mut Client, _: &mut Connection) {
        self.connection_lost_calls += 1;
        client.close().unwrap();
    }

    fn connection_closed(&mut self, client: &mut Client, _: &mut Connection, by_remote: bool) {
        self.connection_closed_calls += 1;
        self.closed_by_remote = by_remote;
        client.close().unwrap();
    }

}

// Server Mock ----------------------------------------------------------------
pub struct MockServerHandler {

    shutdown_ticks: u32,
    close_connection: bool,

    pub shutdown_calls: u32,
    pub tick_connections_calls: u32,
    pub bind_calls: u32,

    pub connection_calls: u32,
    pub connection_failed_calls: u32,
    pub connection_congestion_state_calls: u32,
    pub connection_packet_lost_calls: u32,
    pub connection_lost_calls: u32,
    pub connection_closed_calls: u32,
    pub closed_by_remote: bool
}

impl MockServerHandler {
    pub fn new(shutdown_ticks: u32, close_connection: bool) -> MockServerHandler {
        MockServerHandler {
            shutdown_ticks: shutdown_ticks,
            close_connection: close_connection,

            shutdown_calls: 0,
            tick_connections_calls: 0,
            bind_calls: 0,

            connection_calls: 0,
            connection_failed_calls: 0,
            connection_congestion_state_calls: 0,
            connection_packet_lost_calls: 0,
            connection_lost_calls: 0,
            connection_closed_calls: 0,
            closed_by_remote: false
        }
    }
}

impl Handler<Server> for MockServerHandler {

    fn bind(&mut self, _: &mut Server) {
        self.bind_calls += 1;
    }

    fn tick_connections(
        &mut self, server: &mut Server,
        connections: &mut HashMap<ConnectionID, Connection>
    ) {

        // Ensure hashmap and connection object have the same id
        for (id, conn) in connections.iter() {
            assert_eq!(*id, conn.id());
        }

        self.tick_connections_calls += 1;

        if self.tick_connections_calls > self.shutdown_ticks {
            if self.close_connection {
                for (_, conn) in connections.iter_mut() {
                    conn.close();
                }

            } else {
                server.shutdown().unwrap();
            }
        }

    }

    fn shutdown(&mut self, _: &mut Server) {
        self.shutdown_calls += 1;
    }

    fn connection(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_calls += 1;
    }

    fn connection_failed(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_failed_calls += 1;
    }

    fn connection_packet_lost(
        &mut self, _: &mut Server, _: &mut Connection, _: &[u8]
    ) {
        self.connection_packet_lost_calls += 1;
    }

    fn connection_congestion_state(&mut self, _: &mut Server, _: &mut Connection, _: bool) {
        self.connection_congestion_state_calls += 1;
    }

    fn connection_lost(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_lost_calls += 1;
    }

    fn connection_closed(&mut self, _: &mut Server, _: &mut Connection, by_remote: bool) {
        self.connection_closed_calls += 1;
        self.closed_by_remote = by_remote;
        self.tick_connections_calls = 0;
        self.shutdown_ticks = 5;
        self.close_connection = false;
    }

}

<commit_msg>Further stablilization of connection_close tests.<commit_after>extern crate cobalt;

use std::net;
use std::thread;
use std::time::Duration;
use std::collections::HashMap;
use cobalt::{Client, Connection, ConnectionID, Config, Handler, Server};

#[test]
fn test_client_connection_failure() {

    let config = Config::default();
    let mut handler = MockClientHandler::new();
    let mut client = Client::new(config);
    client.connect(&mut handler, "127.0.0.1:12345").unwrap();

    assert_eq!(handler.connect_calls, 1);
    assert!(handler.tick_connection_calls > 0);
    assert_eq!(handler.close_calls, 1);

    assert_eq!(handler.connection_calls, 0);
    assert_eq!(handler.connection_failed_calls, 1);
    assert_eq!(handler.connection_congestion_state_calls, 0);
    assert_eq!(handler.connection_packet_lost_calls, 0);
    assert_eq!(handler.connection_lost_calls, 0);

}

#[test]
fn test_server_bind_and_shutdown() {

    let config = Config::default();
    let mut handler = MockServerHandler::new(0, false);
    let mut server = Server::new(config);
    server.bind(&mut handler, "127.0.0.1:0").unwrap();

    assert_eq!(handler.bind_calls, 1);
    assert!(handler.tick_connections_calls > 0);
    assert_eq!(handler.shutdown_calls, 1);

    assert_eq!(handler.connection_calls, 0);
    assert_eq!(handler.connection_failed_calls, 0);
    assert_eq!(handler.connection_congestion_state_calls, 0);
    assert_eq!(handler.connection_packet_lost_calls, 0);
    assert_eq!(handler.connection_lost_calls, 0);

}

#[test]
fn test_server_client_connection_lost() {

    // Get a free local socket and then drop it for quick re-use
    // this is not 100% safe but we cannot easily get the locally bound server
    // address after bind() has been called
    let address: Option<net::SocketAddr> = {
        Some(net::UdpSocket::bind("127.0.0.1:0").unwrap().local_addr().unwrap())
    };

    // Setup Test Server
    let server_address = address.clone();
    let server_thread = thread::spawn(move|| {

        let config = Config::default();
        let mut server_handler = MockServerHandler::new(35, false);
        let mut server = Server::new(config);
        server.bind(&mut server_handler, server_address.unwrap()).unwrap();

        assert_eq!(server_handler.bind_calls, 1);
        assert!(server_handler.tick_connections_calls > 0);
        assert_eq!(server_handler.shutdown_calls, 1);

        assert_eq!(server_handler.connection_calls, 1);
        assert_eq!(server_handler.connection_failed_calls, 0);
        assert_eq!(server_handler.connection_congestion_state_calls, 0);
        assert_eq!(server_handler.connection_packet_lost_calls, 0);
        assert_eq!(server_handler.connection_lost_calls, 0);
        assert_eq!(server_handler.connection_closed_calls, 0);

    });

    let config = Config::default();
    let mut client_handler = MockClientHandler::new();
    let mut client = Client::new(config);
    client.connect(&mut client_handler, address.unwrap()).unwrap();

    server_thread.join().unwrap();
    assert_eq!(client_handler.connect_calls, 1);
    assert!(client_handler.tick_connection_calls > 0);
    assert_eq!(client_handler.close_calls, 1);

    assert_eq!(client_handler.connection_calls, 1);
    assert_eq!(client_handler.connection_failed_calls, 0);
    assert_eq!(client_handler.connection_congestion_state_calls, 0);
    // This is somewhat random and depends on how excatly the two threads
    // interact
    // assert_eq!(client_handler.connection_packet_lost_calls, 0);
    assert_eq!(client_handler.connection_lost_calls, 1);
    assert_eq!(client_handler.connection_closed_calls, 0);

}

#[test]
fn test_server_client_connection_close() {

    // Get a free local socket and then drop it for quick re-use
    // this is not 100% safe but we cannot easily get the locally bound server
    // address after bind() has been called
    let address: Option<net::SocketAddr> = {
        Some(net::UdpSocket::bind("127.0.0.1:0").unwrap().local_addr().unwrap())
    };

    // Setup Test Server
    let server_address = address.clone();
    let server_thread = thread::spawn(move|| {

        let config = Config {
            connection_init_threshold: 1000,
            .. Default::default()
        };

        let mut server_handler = MockServerHandler::new(5, true);
        let mut server = Server::new(config);
        server.bind(&mut server_handler, server_address.unwrap()).unwrap();

        assert_eq!(server_handler.connection_lost_calls, 0);
        assert_eq!(server_handler.connection_calls, 1);
        assert_eq!(server_handler.connection_closed_calls, 1);
        assert_eq!(server_handler.closed_by_remote, false);

    });

    let config = Config {
        connection_init_threshold: 1000,
        .. Default::default()
    };

    // TODO cause might be connection marked as dropped by client and
    // an attempted reconnect?
    // client might not receive the closure frame in time?
    //
    // TODO another cause might be that the server thinks the connection is
    // already lost before having send the closure packets? this might explain
    // why the client reconnects afterwards

    let mut client_handler = MockClientHandler::new();
    let mut client = Client::new(config);
    thread::sleep(Duration::from_millis(50));
    client.connect(&mut client_handler, address.unwrap()).unwrap();

    server_thread.join().unwrap();
    assert_eq!(client_handler.connection_lost_calls, 0);
    assert_eq!(client_handler.connection_calls, 1);
    assert_eq!(client_handler.connection_closed_calls, 1);
    assert_eq!(client_handler.closed_by_remote, true);

}

// Client Mock ----------------------------------------------------------------
pub struct MockClientHandler {
    pub connect_calls: u32,
    pub tick_connection_calls: u32,
    pub close_calls: u32,

    pub connection_calls: u32,
    pub connection_failed_calls: u32,
    pub connection_congestion_state_calls: u32,
    pub connection_packet_lost_calls: u32,
    pub connection_lost_calls: u32,
    pub connection_closed_calls: u32,
    pub closed_by_remote: bool
}

impl MockClientHandler {
    pub fn new() -> MockClientHandler {
        MockClientHandler {
            connect_calls: 0,
            tick_connection_calls: 0,
            close_calls: 0,

            connection_calls: 0,
            connection_failed_calls: 0,
            connection_congestion_state_calls: 0,
            connection_packet_lost_calls: 0,
            connection_lost_calls: 0,
            connection_closed_calls: 0,
            closed_by_remote: false
        }
    }
}

impl Handler<Client> for MockClientHandler {

    fn connect(&mut self, _: &mut Client) {
        self.connect_calls += 1;
    }

    fn tick_connection(&mut self, _: &mut Client, _: &mut Connection) {
        self.tick_connection_calls += 1;
    }

    fn close(&mut self, _: &mut Client) {
        self.close_calls += 1;
    }

    fn connection(&mut self, _: &mut Client, _: &mut Connection) {
        self.connection_calls += 1;
    }

    fn connection_failed(&mut self, client: &mut Client, _: &mut Connection) {
        self.connection_failed_calls += 1;
        client.close().unwrap();
    }

    fn connection_packet_lost(
        &mut self, _: &mut Client, _: &mut Connection, _: &[u8]
    ) {
        self.connection_packet_lost_calls += 1;
    }

    fn connection_congestion_state(&mut self, _: &mut Client, _: &mut Connection, _: bool) {
        self.connection_congestion_state_calls += 1;
    }

    fn connection_lost(&mut self, client: &mut Client, _: &mut Connection) {
        self.connection_lost_calls += 1;
        client.close().unwrap();
    }

    fn connection_closed(&mut self, client: &mut Client, _: &mut Connection, by_remote: bool) {
        self.connection_closed_calls += 1;
        self.closed_by_remote = by_remote;
        client.close().unwrap();
    }

}

// Server Mock ----------------------------------------------------------------
pub struct MockServerHandler {

    shutdown_ticks: u32,
    close_connection: bool,

    pub shutdown_calls: u32,
    pub tick_connections_calls: u32,
    pub bind_calls: u32,

    pub connection_calls: u32,
    pub connection_failed_calls: u32,
    pub connection_congestion_state_calls: u32,
    pub connection_packet_lost_calls: u32,
    pub connection_lost_calls: u32,
    pub connection_closed_calls: u32,
    pub closed_by_remote: bool
}

impl MockServerHandler {
    pub fn new(shutdown_ticks: u32, close_connection: bool) -> MockServerHandler {
        MockServerHandler {
            shutdown_ticks: shutdown_ticks,
            close_connection: close_connection,

            shutdown_calls: 0,
            tick_connections_calls: 0,
            bind_calls: 0,

            connection_calls: 0,
            connection_failed_calls: 0,
            connection_congestion_state_calls: 0,
            connection_packet_lost_calls: 0,
            connection_lost_calls: 0,
            connection_closed_calls: 0,
            closed_by_remote: false
        }
    }
}

impl Handler<Server> for MockServerHandler {

    fn bind(&mut self, _: &mut Server) {
        self.bind_calls += 1;
    }

    fn tick_connections(
        &mut self, server: &mut Server,
        connections: &mut HashMap<ConnectionID, Connection>
    ) {

        // Ensure hashmap and connection object have the same id
        for (id, conn) in connections.iter() {
            assert_eq!(*id, conn.id());
        }

        // only advance until initial connection close
        // when there is a connection
        if connections.len() == 0 && self.close_connection {
            self.tick_connections_calls = 0;
        }

        self.tick_connections_calls += 1;

        if self.tick_connections_calls == self.shutdown_ticks {

            // close the client connection the first time around
            if self.close_connection {
                for (_, conn) in connections.iter_mut() {
                    conn.close();
                }

            // shutdown the server
            } else {
                server.shutdown().unwrap();
            }
        }

    }

    fn shutdown(&mut self, _: &mut Server) {
        self.shutdown_calls += 1;
    }

    fn connection(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_calls += 1;
    }

    fn connection_failed(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_failed_calls += 1;
    }

    fn connection_packet_lost(
        &mut self, _: &mut Server, _: &mut Connection, _: &[u8]
    ) {
        self.connection_packet_lost_calls += 1;
    }

    fn connection_congestion_state(&mut self, _: &mut Server, _: &mut Connection, _: bool) {
        self.connection_congestion_state_calls += 1;
    }

    fn connection_lost(&mut self, _: &mut Server, _: &mut Connection) {
        self.connection_lost_calls += 1;
    }

    fn connection_closed(&mut self, _: &mut Server, _: &mut Connection, by_remote: bool) {
        self.connection_closed_calls += 1;
        self.closed_by_remote = by_remote;
        self.tick_connections_calls = 0;
        self.shutdown_ticks = 5;
        self.close_connection = false;
    }

}

<|endoftext|>


<commit_before>#![feature(core)]
extern crate simple_parallel;

#[test]
fn unordered_map_probabilistic_out_of_ordering() {
    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |_: usize| ();
    let iter = simple_parallel::unordered_map(0..N, &f);

    // see if there are any elements where the order they come out of
    // the original iterator is different to the order in which they
    // are yielded here.
    assert!(iter.enumerate().any(|(yield_order, (true_order, ()))| yield_order != true_order));
}

#[test]
fn map_in_order() {
    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |i: usize| i;
    let iter = simple_parallel::map(0..N, &f);

    assert!(iter.enumerate().all(|(yield_order, true_order)| yield_order == true_order));
}

#[test]
fn pool_unordered() {
    let mut pool = simple_parallel::Pool::new(8);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |_: usize| {};
    let iter = pool.unordered_map(0..N, &f);

    // see if there are any elements where the order they come out of
    // the original iterator is different to the order in which they
    // are yielded here.
    let v: Vec<_> = iter.enumerate().collect();
    assert_eq!(v.len(), N);
    assert!(v.into_iter().any(|(yield_order, (true_order, ()))| yield_order != true_order));
}
#[test]
fn pool_unordered2() {
    let mut pool = simple_parallel::Pool::new(2);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |_: usize| {};
    let iter = pool.unordered_map2(0..N, &f);

    // see if there are any elements where the order they come out of
    // the original iterator is different to the order in which they
    // are yielded here.
    let v: Vec<_> = iter.enumerate().collect();
    assert_eq!(v.len(), N);
    assert!(v.into_iter().any(|(yield_order, (true_order, ()))| yield_order != true_order));
}

#[test]
fn pool_map_in_order() {
    let mut pool = simple_parallel::Pool::new(8);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |i: usize| i;
    let iter = pool.map(0..N, &f);

    assert!(std::iter::order::eq(iter, 0..N));
}
#[test]
fn pool_map2_in_order() {
    let mut pool = simple_parallel::Pool::new(8);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |i: usize| i;
    let iter = pool.map2(0..N, &f);

    assert!(std::iter::order::eq(iter, 0..N));
}
<commit_msg>Remove invalid tests.<commit_after>#![feature(core)]
extern crate simple_parallel;

#[test]
fn unordered_map_probabilistic_out_of_ordering() {
    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |_: usize| ();
    let iter = simple_parallel::unordered_map(0..N, &f);

    // see if there are any elements where the order they come out of
    // the original iterator is different to the order in which they
    // are yielded here.
    assert!(iter.enumerate().any(|(yield_order, (true_order, ()))| yield_order != true_order));
}

#[test]
fn map_in_order() {
    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |i: usize| i;
    let iter = simple_parallel::map(0..N, &f);

    assert!(iter.enumerate().all(|(yield_order, true_order)| yield_order == true_order));
}

#[test]
fn pool_unordered() {
    let mut pool = simple_parallel::Pool::new(8);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |_: usize| {};
    let iter = pool.unordered_map(0..N, &f);

    // see if there are any elements where the order they come out of
    // the original iterator is different to the order in which they
    // are yielded here.
    let v: Vec<_> = iter.enumerate().collect();
    assert_eq!(v.len(), N);
    assert!(v.into_iter().any(|(yield_order, (true_order, ()))| yield_order != true_order));
}

#[test]
fn pool_map_in_order() {
    let mut pool = simple_parallel::Pool::new(8);

    // with this many elements, its (hopefully) likely that the
    // threads won't execute sequentially.
    const N: usize = 1000;

    let f = |i: usize| i;
    let iter = pool.map(0..N, &f);

    assert!(std::iter::order::eq(iter, 0..N));
}
<|endoftext|>


<commit_before>extern crate time;
extern crate rand;
extern crate timely;
extern crate timely_sort;

use std::collections::HashMap;
use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::{InputHandle, ProbeHandle};

use timely::dataflow::operators::{LoopVariable, ConnectLoop, Probe};
use timely::dataflow::operators::generic::Operator;
use timely::dataflow::channels::pact::Exchange;

fn main() {

    timely::execute_from_args(std::env::args().skip(3), move |worker| {

        let mut input = InputHandle::new();
        let mut probe = ProbeHandle::new();

        worker.dataflow(|scope| {

            // create a new input, into which we can push edge changes.
            let edge_stream = input.to_stream(scope);

            // create a new feedback stream, which will be changes to ranks.
            let (handle, rank_stream) = scope.loop_variable(usize::max_value(), 1);

            // bring edges and ranks together!
            let changes = edge_stream.binary_frontier(
                &rank_stream,
                Exchange::new(|x: &((usize, usize), i64)| (x.0).0 as u64),
                Exchange::new(|x: &(usize, i64)| x.0 as u64),
                "PageRank",
                |_capability| {

                    // where we stash out-of-order data.
                    let mut edge_stash = HashMap::new();
                    let mut rank_stash = HashMap::new();

                    // lists of edges, ranks, and changes.
                    let mut edges = Vec::new();
                    let mut ranks = Vec::new();
                    let mut delta = Vec::new();

                    let timer = ::std::time::Instant::now();

                    move |input1, input2, output| {

                        // hold on to edge changes until it is time.
                        input1.for_each(|time, data| {
                            edge_stash.entry(time).or_insert(Vec::new()).extend(data.drain(..));
                        });

                        // hold on to rank changes until it is time.
                        input2.for_each(|time, data| {
                            rank_stash.entry(time).or_insert(Vec::new()).extend(data.drain(..));
                        });

                        let frontiers = &[input1.frontier(), input2.frontier()];

                        for (time, edge_changes) in edge_stash.iter_mut() {
                            if frontiers.iter().all(|f| !f.less_equal(time)) {

                                let mut session = output.session(time);

                                compact(edge_changes);

                                for ((src, dst), diff) in edge_changes.drain(..) {

                                    // 0. ensure enough state allocated
                                    while edges.len() <= src { edges.push(Vec::new()); }
                                    while ranks.len() <= src { ranks.push(1_000); }

                                    // 1. subtract previous distribution.
                                    allocate(ranks[src], &edges[src][..], &mut delta);
                                    for x in delta.iter_mut() { x.1 *= -1; }

                                    // 2. update edges.
                                    edges[src].push((dst, diff));
                                    compact(&mut edges[src]);

                                    // 3. re-distribute allocations.
                                    allocate(ranks[src], &edges[src][..], &mut delta);

                                    // 4. compact down and send cumulative changes.
                                    compact(&mut delta);
                                    for (dst, diff) in delta.drain(..) {
                                        session.give((dst, diff));
                                    }
                                }
                            }
                        }

                        edge_stash.retain(|_key, val| !val.is_empty());

                        for (time, rank_changes) in rank_stash.iter_mut() {
                            if frontiers.iter().all(|f| !f.less_equal(time)) {

                                let mut session = output.session(time);

                                compact(rank_changes);

                                let mut cnt = 0;
                                let mut sum = 0;
                                let mut max = 0;

                                for (src, diff) in rank_changes.drain(..) {

                                    cnt += 1;
                                    sum += diff.abs();
                                    max = if max < diff.abs() { diff.abs() } else { max };

                                    // 0. ensure enough state allocated
                                    while edges.len() <= src { edges.push(Vec::new()); }
                                    while ranks.len() <= src { ranks.push(1_000); }

                                    // 1. subtract previous distribution.
                                    allocate(ranks[src], &edges[src][..], &mut delta);
                                    for x in delta.iter_mut() { x.1 *= -1; }

                                    // 2. update ranks.
                                    ranks[src] += diff;

                                    // 3. re-distribute allocations.
                                    allocate(ranks[src], &edges[src][..], &mut delta);

                                    // 4. compact down and send cumulative changes.
                                    compact(&mut delta);
                                    for (dst, diff) in delta.drain(..) {
                                        session.give((dst, diff));
                                    }
                                }

                                println!("{:?}:\t{:?}\t{}\t{}\t{}", timer.elapsed(), time.time(), cnt, sum, max);
                            }
                        }

                        rank_stash.retain(|_key, val| !val.is_empty());

                    }
                }
            );

            changes
                .probe_with(&mut probe)
                .connect_loop(handle);

        });

        let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
        let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

        let seed: &[_] = &[1, 2, 3, worker.index()];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);
        let mut rng2: StdRng = SeedableRng::from_seed(seed);

        for _ in 0 .. edges / worker.peers() {
            input.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
        }

        input.advance_to(1);

        while probe.less_than(input.time()) {
            worker.step();
        }

        for i in 1 .. 1000 {
            input.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
            input.send(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)), -1));
            input.advance_to(1 + i);
            while probe.less_than(input.time()) {
                worker.step();
            }
        }

    }).unwrap(); // asserts error-free execution;
}

fn compact<T: Ord>(list: &mut Vec<(T, i64)>) {
    if !list.is_empty() {
        list.sort_by(|x,y| x.0.cmp(&y.0));
        for i in 0 .. list.len() - 1 {
            if list[i].0 == list[i+1].0 {
                list[i+1].1 += list[i].1;
                list[i].1 = 0;
            }
        }
        list.retain(|x| x.1 != 0);
    }
}

// this method allocates some rank between elements of `edges`.
fn allocate(rank: i64, edges: &[(usize, i64)], send: &mut Vec<(usize, i64)>) {
    if !edges.is_empty() {
        assert!(rank >= 0);
        assert!(edges.iter().all(|x| x.1 > 0));

        let degree: i64 = edges.iter().map(|x| x.1 as i64).sum();
        let share = ((rank * 5) / 6) / degree;
        for i in 0 .. edges.len() {
            if (i as i64) < (share % (edges.len() as i64)) {
                send.push((edges[i].0, edges[i].1 * (share + 1)));
            }
            else {
                send.push((edges[i].0, edges[i].1 * share));
            }
        }
    }
}<commit_msg>bug fixes<commit_after>extern crate time;
extern crate rand;
extern crate timely;
extern crate timely_sort;

use std::collections::HashMap;
use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::{InputHandle, ProbeHandle};

use timely::dataflow::operators::{LoopVariable, ConnectLoop, Probe};
use timely::dataflow::operators::generic::Operator;
use timely::dataflow::channels::pact::Exchange;

fn main() {

    timely::execute_from_args(std::env::args().skip(3), move |worker| {

        let mut input = InputHandle::new();
        let mut probe = ProbeHandle::new();

        worker.dataflow(|scope| {

            // create a new input, into which we can push edge changes.
            let edge_stream = input.to_stream(scope);

            // create a new feedback stream, which will be changes to ranks.
            let (handle, rank_stream) = scope.loop_variable(usize::max_value(), 1);

            // bring edges and ranks together!
            let changes = edge_stream.binary_frontier(
                &rank_stream,
                Exchange::new(|x: &((usize, usize), i64)| (x.0).0 as u64),
                Exchange::new(|x: &(usize, i64)| x.0 as u64),
                "PageRank",
                |_capability| {

                    // where we stash out-of-order data.
                    let mut edge_stash = HashMap::new();
                    let mut rank_stash = HashMap::new();

                    // lists of edges, ranks, and changes.
                    let mut edges = Vec::new();
                    let mut ranks = Vec::new();
                    let mut diffs = Vec::new(); // for received but un-acted upon deltas.
                    let mut delta = Vec::new();

                    let timer = ::std::time::Instant::now();

                    move |input1, input2, output| {

                        // hold on to edge changes until it is time.
                        input1.for_each(|time, data| {
                            edge_stash.entry(time).or_insert(Vec::new()).extend(data.drain(..));
                        });

                        // hold on to rank changes until it is time.
                        input2.for_each(|time, data| {
                            rank_stash.entry(time).or_insert(Vec::new()).extend(data.drain(..));
                        });

                        let frontiers = &[input1.frontier(), input2.frontier()];

                        for (time, edge_changes) in edge_stash.iter_mut() {
                            if frontiers.iter().all(|f| !f.less_equal(time)) {

                                let mut session = output.session(time);

                                compact(edge_changes);

                                for ((src, dst), diff) in edge_changes.drain(..) {

                                    // 0. ensure enough state allocated
                                    while edges.len() <= src { edges.push(Vec::new()); }
                                    while ranks.len() <= src { ranks.push(1_000); }
                                    while diffs.len() <= src { diffs.push(0); }

                                    // 1. subtract previous distribution.
                                    allocate(ranks[src], &edges[src][..], &mut delta);
                                    for x in delta.iter_mut() { x.1 *= -1; }

                                    // 2. update edges.
                                    edges[src].push((dst, diff));
                                    compact(&mut edges[src]);

                                    // 3. re-distribute allocations.
                                    allocate(ranks[src], &edges[src][..], &mut delta);

                                    // 4. compact down and send cumulative changes.
                                    compact(&mut delta);
                                    for (dst, diff) in delta.drain(..) {
                                        session.give((dst, diff));
                                    }
                                }
                            }
                        }

                        edge_stash.retain(|_key, val| !val.is_empty());

                        for (time, rank_changes) in rank_stash.iter_mut() {
                            if frontiers.iter().all(|f| !f.less_equal(time)) {

                                let mut session = output.session(time);

                                compact(rank_changes);

                                let mut cnt = 0;
                                let mut sum = 0;
                                let mut max = 0;

                                for (src, diff) in rank_changes.drain(..) {

                                    cnt += 1;
                                    sum += diff.abs();
                                    max = if max < diff.abs() { diff.abs() } else { max };

                                    // 0. ensure enough state allocated
                                    while edges.len() <= src { edges.push(Vec::new()); }
                                    while ranks.len() <= src { ranks.push(1_000); }
                                    while diffs.len() <= src { diffs.push(0); }

                                    // 1. subtract previous distribution.
                                    allocate(ranks[src], &edges[src][..], &mut delta);
                                    for x in delta.iter_mut() { x.1 *= -1; }

                                    // 2. update ranks.
                                    diffs[src] += diff;
                                    if diffs[src].abs() >= 6 {
                                        ranks[src] += diffs[src];
                                        diffs[src] = 0;
                                    }

                                    // 3. re-distribute allocations.
                                    allocate(ranks[src], &edges[src][..], &mut delta);

                                    // 4. compact down and send cumulative changes.
                                    compact(&mut delta);
                                    for (dst, diff) in delta.drain(..) {
                                        session.give((dst, diff));
                                    }
                                }

                                println!("{:?}:\t{:?}\t{}\t{}\t{}", timer.elapsed(), time.time(), cnt, sum, max);
                            }
                        }

                        rank_stash.retain(|_key, val| !val.is_empty());

                    }
                }
            );

            changes
                .probe_with(&mut probe)
                .connect_loop(handle);

        });

        let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
        let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

        let seed: &[_] = &[1, 2, 3, worker.index()];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);
        let mut rng2: StdRng = SeedableRng::from_seed(seed);

        for _ in 0 .. edges / worker.peers() {
            input.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
        }

        input.advance_to(1);

        while probe.less_than(input.time()) {
            worker.step();
        }

        for i in 1 .. 1000 {
            input.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
            input.send(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)), -1));
            input.advance_to(1 + i);
            while probe.less_than(input.time()) {
                worker.step();
            }
        }

    }).unwrap(); // asserts error-free execution;
}

fn compact<T: Ord>(list: &mut Vec<(T, i64)>) {
    if !list.is_empty() {
        list.sort_by(|x,y| x.0.cmp(&y.0));
        for i in 0 .. list.len() - 1 {
            if list[i].0 == list[i+1].0 {
                list[i+1].1 += list[i].1;
                list[i].1 = 0;
            }
        }
        list.retain(|x| x.1 != 0);
    }
}

// this method allocates some rank between elements of `edges`.
fn allocate(rank: i64, edges: &[(usize, i64)], send: &mut Vec<(usize, i64)>) {
    if !edges.is_empty() {
        assert!(rank >= 0);
        assert!(edges.iter().all(|x| x.1 > 0));

        let distribute = (rank * 5) / 6;
        let degree = edges.len() as i64;
        let share = distribute / degree;
        for i in 0 .. edges.len() {
            if (i as i64) < (distribute % (edges.len() as i64)) {
                send.push((edges[i].0, edges[i].1 * (share + 1)));
            }
            else {
                send.push((edges[i].0, edges[i].1 * share));
            }
        }
    }
}<|endoftext|>


<commit_before>// .env parsing
extern crate dotenv;
// command line argument parsing
extern crate clap;
// logging macros
#[macro_use]
extern crate log;
// console logging output
extern crate fern;
extern crate chrono;
// sockets
extern crate tokio_core;
extern crate futures;
extern crate websocket;
// Screeps API
extern crate screeps_api;
// json pretty printing
extern crate serde_json;

use futures::{Sink, Stream, Future, future, stream};

use websocket::OwnedMessage;

use screeps_api::{TokenStorage, RoomName};
use screeps_api::websocket::{Channel, ChannelUpdate, SockjsMessage, ScreepsMessage};

/// Set up dotenv and retrieve a specific variable, informatively panicking if it does not exist.
fn env(var: &str) -> String {
    dotenv::dotenv().ok();
    match ::std::env::var(var) {
        Ok(value) => value,
        Err(e) => panic!("must have `{}` defined (err: {:?})", var, e),
    }
}

fn setup_logging(verbosity: u64) {
    let log_level = match verbosity {
        0 => log::LogLevelFilter::Warn,
        1 => log::LogLevelFilter::Info,
        _ => log::LogLevelFilter::Debug,
    };
    fern::Dispatch::new()
        .level(log_level)
        .level_for("rustls", log::LogLevelFilter::Warn)
        .level_for("hyper", log::LogLevelFilter::Warn)
        .format(|out, message, record| {
            let now = chrono::Local::now();

            out.finish(format_args!("[{}][{}] {}: {}",
                                    now.format("%H:%M:%S"),
                                    record.level(),
                                    record.target(),
                                    message));
        })
        .chain(std::io::stdout())
        .apply()
        // ignore errors
        .unwrap_or(());
}

#[derive(Clone, Debug)]
struct Config {
    cpu: bool,
    messages: bool,
    credits: bool,
    console: bool,
    rooms: Vec<RoomName>,
    map_view: Vec<RoomName>,
}

impl Config {
    fn new<'a>(args: &'a clap::ArgMatches) -> Result<Self, screeps_api::data::room_name::RoomNameParseError<'a>> {
        Ok(Config {
            cpu: args.is_present("cpu"),
            messages: args.is_present("messages"),
            credits: args.is_present("credits"),
            console: args.is_present("console"),
            rooms: args.values_of("room")
                .map(|it| it.map(|v| RoomName::new(v)).collect::<Result<_, _>>())
                .unwrap_or_else(|| Ok(Vec::new()))?,
            map_view: args.values_of("map-view")
                .map(|it| it.map(|v| RoomName::new(v)).collect::<Result<_, _>>())
                .unwrap_or_else(|| Ok(Vec::new()))?,
        })
    }

    fn subscribe_with(&self, id: &str) -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        use screeps_api::websocket::subscribe;

        let mut messages = Vec::with_capacity(1 + self.cpu as usize + self.messages as usize + self.credits as usize +
                                              self.console as usize +
                                              self.rooms.len() +
                                              self.map_view.len());

        messages.push(subscribe(&Channel::ServerMessages));

        if self.cpu {
            messages.push(subscribe(&Channel::user_cpu(id)));
        }

        if self.messages {
            messages.push(subscribe(&Channel::user_messages(id)));
            messages.push(subscribe(&Channel::user_conversation(id, "57fb16b6e4dd183b746435b0")));
        }

        if self.credits {
            messages.push(subscribe(&Channel::user_credits(id)));
        }

        if self.console {
            messages.push(subscribe(&Channel::user_console(id)));
        }

        for room_name in &self.rooms {
            messages.push(subscribe(&Channel::room_detail(*room_name)));
        }

        for room_name in &self.map_view {
            messages.push(subscribe(&Channel::room_map_view(*room_name)));
        }

        Box::new(stream::iter(messages.into_iter().map(|string| Ok(OwnedMessage::Text(string)))))
    }
}

fn setup() -> Config {
    let cmd_arguments = clap::App::new("ws-debug")
        .arg(clap::Arg::with_name("verbose")
            .short("v")
            .long("verbose")
            .multiple(true)
            .help("Enables verbose logging"))
        .arg(clap::Arg::with_name("cpu")
            .short("p")
            .long("cpu")
            .help("Subscribe to user cpu and memory updates"))
        .arg(clap::Arg::with_name("credits")
            .short("c")
            .long("credits")
            .help("Subscribe to per-tick user credit updates"))
        .arg(clap::Arg::with_name("console")
            .short("o")
            .long("console")
            .help("Subscribe to user console messages"))
        .arg(clap::Arg::with_name("messages")
            .short("e")
            .long("messages")
            .help("Subscribe to user message alerts"))
        .arg(clap::Arg::with_name("room")
            .short("r")
            .long("room")
            .value_name("ROOM_NAME")
            .help("Subscribes to a room")
            .takes_value(true)
            .multiple(true))
        .arg(clap::Arg::with_name("map-view")
            .short("m")
            .long("map-view")
            .value_name("ROOM_NAME")
            .help("Subscribes to a map-view room")
            .takes_value(true)
            .multiple(true))
        .get_matches();

    setup_logging(cmd_arguments.occurrences_of("verbose"));

    match Config::new(&cmd_arguments) {
        Ok(v) => v,
        Err(e) => {
            clap::Error {
                    message: e.to_string(),
                    kind: clap::ErrorKind::InvalidValue,
                    info: None,
                }
                .exit()
        }
    }
}

fn main() {
    let config = setup();

    let token_storage = screeps_api::RcTokenStorage::default();

    let mut client = screeps_api::SyncConfig::new().unwrap().tokens(token_storage.clone()).build().unwrap();

    // Login using the API client - this will storage the auth token in token_storage.
    client.login(env("SCREEPS_API_USERNAME"), env("SCREEPS_API_PASSWORD")).expect("failed to login");

    let my_info = client.my_info().unwrap();

    info!("Logged in as {}, attempting to connect to stream.",
          my_info.username);

    let mut core = tokio_core::reactor::Core::new().expect("expected IO core to start up without issue.");

    let handle = core.handle();

    let url = screeps_api::websocket::default_url();

    let connection = websocket::ClientBuilder::from_url(&url).async_connect_secure(None, &handle);

    core.run(connection.then(|result| {
            let (client, _) = result.expect("expected successful connection to official server, found error");

            let (sink, stream) = client.split();

            sink.send(OwnedMessage::Text(screeps_api::websocket::authenticate(&token_storage.take_token().unwrap())))
                .and_then(|sink| {
                    let handler = Handler::new(token_storage, my_info, config);

                    sink.send_all((stream.and_then(move |data| future::ok(handler.handle_data(data)))
                        .or_else(|err| {
                            warn!("error occurred: {}", err);

                            future::ok::<_, websocket::WebSocketError>(Box::new(stream::empty()) as
                                                                       Box<Stream<Item = OwnedMessage,
                                                                                Error = websocket::WebSocketError>>)
                        })
                        .flatten()))
                })

        }))
        .expect("expected websocket connection to complete successfully, but an error occurred");
}

struct Handler<T>
    where T: screeps_api::TokenStorage
{
    tokens: T,
    info: screeps_api::MyInfo,
    config: Config,
}

impl<T> Handler<T>
    where T: screeps_api::TokenStorage
{
    fn new(tokens: T, info: screeps_api::MyInfo, config: Config) -> Self {
        Handler {
            tokens: tokens,
            info: info,
            config: config,
        }
    }

    fn handle_data(&self, data: OwnedMessage) -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        match data {
            OwnedMessage::Text(string) => {
                let data = SockjsMessage::parse(&string)
                    .expect("expected a correct SockJS message, found a parse error.");

                match data {
                    SockjsMessage::Open => debug!("SockJS connection opened"),
                    SockjsMessage::Heartbeat => debug!("SockJS heartbeat."),
                    SockjsMessage::Close { .. } => debug!("SockJS close"),
                    SockjsMessage::Message(message) => {
                        return Box::new(self.handle_parsed_message(message));
                    }
                    SockjsMessage::Messages(messages) => {
                        let results = messages.into_iter()
                            .map(|message| Ok(self.handle_parsed_message(message)))
                            .collect::<Vec<_>>();

                        return Box::new(stream::iter::<_, _, websocket::WebSocketError>(results).flatten());
                    }

                }
            }
            OwnedMessage::Binary(data) => warn!("ignoring binary data from websocket: {:?}", data),
            OwnedMessage::Close(data) => warn!("connection closing: {:?}", data),
            OwnedMessage::Ping(data) => return Box::new(stream::once(Ok(OwnedMessage::Pong(data)))),
            OwnedMessage::Pong(_) => (),
        }

        Box::new(stream::empty())
    }

    fn handle_parsed_message(&self,
                             message: screeps_api::websocket::parsing::ScreepsMessage)
                             -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        match message {
            ScreepsMessage::AuthFailed => panic!("authentication with stored token failed!"),
            ScreepsMessage::AuthOk { new_token } => {
                info!("authentication succeeded, now authenticated as {}.",
                      self.info.username);
                // return the token to the token storage in case we need it again - we won't in this example
                // program, but this is a good practice
                //
                // TODO: find an efficient way to do this automatically in the handler.
                self.tokens.return_token(new_token);

                return Box::new(self.config
                    .subscribe_with(&self.info.user_id)
                    .chain(stream::futures_unordered(vec![future::lazy(|| {
                        warn!("subscribed!");
                        future::ok::<_, websocket::WebSocketError>(stream::empty())
                    })])
                        .flatten())) as
                       Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>>;
            }
            ScreepsMessage::ChannelUpdate { update } => {
                self.handle_update(update);
            }
            ScreepsMessage::ServerProtocol { protocol } => {
                info!("server protocol: {}", protocol);
            }
            ScreepsMessage::ServerTime { time } => {
                info!("server time: {}", time);
            }
            ScreepsMessage::ServerPackage { package } => {
                info!("server package: {}", package);
            }
            ScreepsMessage::Other(other) => {
                warn!("ScreepsMessage::Other: {}", other);
            }
        }

        Box::new(stream::empty())
    }

    fn handle_update(&self, update: ChannelUpdate) {
        match update {
            ChannelUpdate::UserCpu { user_id, update } => info!("CPU: [{}] {:#?}", user_id, update),
            ChannelUpdate::RoomMapView { room_name, update } => {
                info!("Map View: [{}] {:?}", room_name, update);
            }
            ChannelUpdate::RoomDetail { room_name, update } => {
                info!("Room Detail: [{}] {:?}", room_name, update);
            }
            ChannelUpdate::NoRoomDetail { room_name } => {
                info!("Room Skipped: {}", room_name);
            }
            ChannelUpdate::UserConsole { user_id, update } => {
                info!("Console: [{}] {:#?}", user_id, update);
            }
            ChannelUpdate::UserCredits { user_id, update } => {
                info!("Credits: [{}] {}", user_id, update);
            }
            ChannelUpdate::UserMessage { user_id, update } => {
                info!("New message: [{}] {:#?}", user_id, update);
            }
            ChannelUpdate::UserConversation { user_id, target_user_id, update } => {
                info!("Conversation update: [{}->{}] {:#?}",
                      user_id,
                      target_user_id,
                      update);
            }
            ChannelUpdate::Other { channel, update } => {
                warn!("ChannelUpdate::Other: {}\n{}",
                      channel,
                      serde_json::to_string_pretty(&update).expect("failed to serialize json string"));
            }
        }
    }
}
<commit_msg>Output room information at INFO level in ws-debug for easy room ws debugging.<commit_after>// .env parsing
extern crate dotenv;
// command line argument parsing
extern crate clap;
// logging macros
#[macro_use]
extern crate log;
// console logging output
extern crate fern;
extern crate chrono;
// sockets
extern crate tokio_core;
extern crate futures;
extern crate websocket;
// Screeps API
extern crate screeps_api;
// json pretty printing
extern crate serde_json;

use futures::{Sink, Stream, Future, future, stream};

use websocket::OwnedMessage;

use screeps_api::{TokenStorage, RoomName};
use screeps_api::websocket::{Channel, ChannelUpdate, SockjsMessage, ScreepsMessage};

/// Set up dotenv and retrieve a specific variable, informatively panicking if it does not exist.
fn env(var: &str) -> String {
    dotenv::dotenv().ok();
    match ::std::env::var(var) {
        Ok(value) => value,
        Err(e) => panic!("must have `{}` defined (err: {:?})", var, e),
    }
}

fn setup_logging(verbosity: u64) {
    let log_level = match verbosity {
        0 => log::LogLevelFilter::Warn,
        1 => log::LogLevelFilter::Info,
        _ => log::LogLevelFilter::Debug,
    };
    fern::Dispatch::new()
        .level(log_level)
        .level_for("rustls", log::LogLevelFilter::Warn)
        .level_for("hyper", log::LogLevelFilter::Warn)
        .format(|out, message, record| {
            let now = chrono::Local::now();

            out.finish(format_args!("[{}][{}] {}: {}",
                                    now.format("%H:%M:%S"),
                                    record.level(),
                                    record.target(),
                                    message));
        })
        .chain(std::io::stdout())
        .apply()
        // ignore errors
        .unwrap_or(());
}

#[derive(Clone, Debug)]
struct Config {
    cpu: bool,
    messages: bool,
    credits: bool,
    console: bool,
    rooms: Vec<RoomName>,
    map_view: Vec<RoomName>,
}

impl Config {
    fn new<'a>(args: &'a clap::ArgMatches) -> Result<Self, screeps_api::data::room_name::RoomNameParseError<'a>> {
        Ok(Config {
            cpu: args.is_present("cpu"),
            messages: args.is_present("messages"),
            credits: args.is_present("credits"),
            console: args.is_present("console"),
            rooms: args.values_of("room")
                .map(|it| it.map(|v| RoomName::new(v)).collect::<Result<_, _>>())
                .unwrap_or_else(|| Ok(Vec::new()))?,
            map_view: args.values_of("map-view")
                .map(|it| it.map(|v| RoomName::new(v)).collect::<Result<_, _>>())
                .unwrap_or_else(|| Ok(Vec::new()))?,
        })
    }

    fn subscribe_with(&self, id: &str) -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        use screeps_api::websocket::subscribe;

        let mut messages = Vec::with_capacity(1 + self.cpu as usize + self.messages as usize + self.credits as usize +
                                              self.console as usize +
                                              self.rooms.len() +
                                              self.map_view.len());

        messages.push(subscribe(&Channel::ServerMessages));

        if self.cpu {
            messages.push(subscribe(&Channel::user_cpu(id)));
        }

        if self.messages {
            messages.push(subscribe(&Channel::user_messages(id)));
            messages.push(subscribe(&Channel::user_conversation(id, "57fb16b6e4dd183b746435b0")));
        }

        if self.credits {
            messages.push(subscribe(&Channel::user_credits(id)));
        }

        if self.console {
            messages.push(subscribe(&Channel::user_console(id)));
        }

        for room_name in &self.rooms {
            messages.push(subscribe(&Channel::room_detail(*room_name)));
        }

        for room_name in &self.map_view {
            messages.push(subscribe(&Channel::room_map_view(*room_name)));
        }

        Box::new(stream::iter(messages.into_iter().map(|string| Ok(OwnedMessage::Text(string)))))
    }
}

fn setup() -> Config {
    let cmd_arguments = clap::App::new("ws-debug")
        .arg(clap::Arg::with_name("verbose")
            .short("v")
            .long("verbose")
            .multiple(true)
            .help("Enables verbose logging"))
        .arg(clap::Arg::with_name("cpu")
            .short("p")
            .long("cpu")
            .help("Subscribe to user cpu and memory updates"))
        .arg(clap::Arg::with_name("credits")
            .short("c")
            .long("credits")
            .help("Subscribe to per-tick user credit updates"))
        .arg(clap::Arg::with_name("console")
            .short("o")
            .long("console")
            .help("Subscribe to user console messages"))
        .arg(clap::Arg::with_name("messages")
            .short("e")
            .long("messages")
            .help("Subscribe to user message alerts"))
        .arg(clap::Arg::with_name("room")
            .short("r")
            .long("room")
            .value_name("ROOM_NAME")
            .help("Subscribes to a room")
            .takes_value(true)
            .multiple(true))
        .arg(clap::Arg::with_name("map-view")
            .short("m")
            .long("map-view")
            .value_name("ROOM_NAME")
            .help("Subscribes to a map-view room")
            .takes_value(true)
            .multiple(true))
        .get_matches();

    setup_logging(cmd_arguments.occurrences_of("verbose"));

    match Config::new(&cmd_arguments) {
        Ok(v) => v,
        Err(e) => {
            clap::Error {
                    message: e.to_string(),
                    kind: clap::ErrorKind::InvalidValue,
                    info: None,
                }
                .exit()
        }
    }
}

fn main() {
    let config = setup();

    let token_storage = screeps_api::RcTokenStorage::default();

    let mut client = screeps_api::SyncConfig::new().unwrap().tokens(token_storage.clone()).build().unwrap();

    // Login using the API client - this will storage the auth token in token_storage.
    client.login(env("SCREEPS_API_USERNAME"), env("SCREEPS_API_PASSWORD")).expect("failed to login");

    let my_info = client.my_info().unwrap();

    info!("Logged in as {}, attempting to connect to stream.",
          my_info.username);

    let mut core = tokio_core::reactor::Core::new().expect("expected IO core to start up without issue.");

    let handle = core.handle();

    let url = screeps_api::websocket::default_url();

    let connection = websocket::ClientBuilder::from_url(&url).async_connect_secure(None, &handle);

    core.run(connection.then(|result| {
            let (client, _) = result.expect("expected successful connection to official server, found error");

            let (sink, stream) = client.split();

            sink.send(OwnedMessage::Text(screeps_api::websocket::authenticate(&token_storage.take_token().unwrap())))
                .and_then(|sink| {
                    let handler = Handler::new(token_storage, my_info, config);

                    sink.send_all((stream.and_then(move |data| future::ok(handler.handle_data(data)))
                        .or_else(|err| {
                            warn!("error occurred: {}", err);

                            future::ok::<_, websocket::WebSocketError>(Box::new(stream::empty()) as
                                                                       Box<Stream<Item = OwnedMessage,
                                                                                Error = websocket::WebSocketError>>)
                        })
                        .flatten()))
                })

        }))
        .expect("expected websocket connection to complete successfully, but an error occurred");
}

struct Handler<T>
    where T: screeps_api::TokenStorage
{
    tokens: T,
    info: screeps_api::MyInfo,
    config: Config,
}

impl<T> Handler<T>
    where T: screeps_api::TokenStorage
{
    fn new(tokens: T, info: screeps_api::MyInfo, config: Config) -> Self {
        Handler {
            tokens: tokens,
            info: info,
            config: config,
        }
    }

    fn handle_data(&self, data: OwnedMessage) -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        match data {
            OwnedMessage::Text(string) => {
                let data = SockjsMessage::parse(&string)
                    .expect("expected a correct SockJS message, found a parse error.");

                match data {
                    SockjsMessage::Open => debug!("SockJS connection opened"),
                    SockjsMessage::Heartbeat => debug!("SockJS heartbeat."),
                    SockjsMessage::Close { .. } => debug!("SockJS close"),
                    SockjsMessage::Message(message) => {
                        return Box::new(self.handle_parsed_message(message));
                    }
                    SockjsMessage::Messages(messages) => {
                        let results = messages.into_iter()
                            .map(|message| Ok(self.handle_parsed_message(message)))
                            .collect::<Vec<_>>();

                        return Box::new(stream::iter::<_, _, websocket::WebSocketError>(results).flatten());
                    }

                }
            }
            OwnedMessage::Binary(data) => warn!("ignoring binary data from websocket: {:?}", data),
            OwnedMessage::Close(data) => warn!("connection closing: {:?}", data),
            OwnedMessage::Ping(data) => return Box::new(stream::once(Ok(OwnedMessage::Pong(data)))),
            OwnedMessage::Pong(_) => (),
        }

        Box::new(stream::empty())
    }

    fn handle_parsed_message(&self,
                             message: screeps_api::websocket::parsing::ScreepsMessage)
                             -> Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>> {
        match message {
            ScreepsMessage::AuthFailed => panic!("authentication with stored token failed!"),
            ScreepsMessage::AuthOk { new_token } => {
                info!("authentication succeeded, now authenticated as {}.",
                      self.info.username);
                // return the token to the token storage in case we need it again - we won't in this example
                // program, but this is a good practice
                //
                // TODO: find an efficient way to do this automatically in the handler.
                self.tokens.return_token(new_token);

                return Box::new(self.config
                    .subscribe_with(&self.info.user_id)
                    .chain(stream::futures_unordered(vec![future::lazy(|| {
                        warn!("subscribed!");
                        future::ok::<_, websocket::WebSocketError>(stream::empty())
                    })])
                        .flatten())) as
                       Box<Stream<Item = OwnedMessage, Error = websocket::WebSocketError>>;
            }
            ScreepsMessage::ChannelUpdate { update } => {
                self.handle_update(update);
            }
            ScreepsMessage::ServerProtocol { protocol } => {
                info!("server protocol: {}", protocol);
            }
            ScreepsMessage::ServerTime { time } => {
                info!("server time: {}", time);
            }
            ScreepsMessage::ServerPackage { package } => {
                info!("server package: {}", package);
            }
            ScreepsMessage::Other(other) => {
                warn!("ScreepsMessage::Other: {}", other);
            }
        }

        Box::new(stream::empty())
    }

    fn handle_update(&self, update: ChannelUpdate) {
        match update {
            ChannelUpdate::UserCpu { user_id, update } => info!("CPU: [{}] {:#?}", user_id, update),
            ChannelUpdate::RoomMapView { room_name, update } => {
                info!("Map View: [{}] {:?}", room_name, update);
            }
            ChannelUpdate::RoomDetail { room_name, update } => {
                debug!("Room Detail: [{}] {:?}", room_name, update);
                info!("Room {}: {}",
                      room_name,
                      serde_json::to_string_pretty(&serde_json::Value::Object(update.objects
                              .iter()
                              .cloned()
                              .collect()))
                          .expect("expected to_string to succeed on plain map."));
            }
            ChannelUpdate::NoRoomDetail { room_name } => {
                info!("Room Skipped: {}", room_name);
            }
            ChannelUpdate::UserConsole { user_id, update } => {
                info!("Console: [{}] {:#?}", user_id, update);
            }
            ChannelUpdate::UserCredits { user_id, update } => {
                info!("Credits: [{}] {}", user_id, update);
            }
            ChannelUpdate::UserMessage { user_id, update } => {
                info!("New message: [{}] {:#?}", user_id, update);
            }
            ChannelUpdate::UserConversation { user_id, target_user_id, update } => {
                info!("Conversation update: [{}->{}] {:#?}",
                      user_id,
                      target_user_id,
                      update);
            }
            ChannelUpdate::Other { channel, update } => {
                warn!("ChannelUpdate::Other: {}\n{}",
                      channel,
                      serde_json::to_string_pretty(&update).expect("failed to serialize json string"));
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate rustfmt;
extern crate syntex_syntax as syntax;

use syntax::ast;
use syntax::ext::base::{ExtCtxt, DummyMacroLoader};
use syntax::ext::expand;
use syntax::ext::expand::{ExpansionConfig, MacroExpander};
use syntax::codemap::{CodeMap, Span, ExpnInfo, NO_EXPANSION};
use syntax::errors::Handler;
use syntax::errors::emitter::{ColorConfig};
use syntax::fold::{self, Folder};
use syntax::parse::{self, ParseSess};
use syntax::ptr::{self, P};
use syntax::util::small_vector::SmallVector;

use rustfmt::filemap::FileMap;
use rustfmt::config::{Config, WriteMode};
use rustfmt::modules::list_files;
use rustfmt::visitor::FmtVisitor;

use std::collections::HashMap;
use std::io::stdout;
use std::env;
use std::path::Path;
use std::rc::Rc;

// Small macro to simplify setting the full-expansion closures to the identity closure.
macro_rules! set_expander_fns {
    ($expander:ident,
        $( $expand:ident ),*) => {{
        $( $expander.$expand = Rc::new(Box::new(|_, node| node )); )*
    }}
}

struct ExpandData<'a> {
    config: Config,
    cx: ExtCtxt<'a>,
    krates: Vec<ast::Crate>,
    index: usize,
    span_map: HashMap<Span, Span>,
}

impl<'a> ExpandData<'a> {
    fn new(sess: &'a ParseSess, loader: &'a mut DummyMacroLoader) -> ExpandData<'a> {
        let args: Vec<String> = env::args().collect();
        if args.len() < 2 {
            panic!("Please supply a filepath to parse.")
        }
        if args.len() > 2 {
            panic!("Too many arguments. Please supply a single filepath.")
        }
        let filename = args[1].clone();

        let mut config = Config::default();
        config.write_mode = WriteMode::Overwrite;

        let ex_cfg = ExpansionConfig::default(filename.clone());
        let mut krates = vec!();
        krates.push(parse::parse_crate_from_file(&Path::new(&filename),
                                                 Vec::new(), sess).unwrap());
        let ecx = ExtCtxt::new(sess,
                               krates[0].config.clone(),
                               ex_cfg,
                               loader);
        ExpandData {
            config: config,
            cx: ecx,
            krates: krates,
            index: 0,
            span_map: HashMap::new(),
        }
    }

    fn insert(&mut self, span: Span) {
        if span.expn_id == NO_EXPANSION {
            return;
        }

        let key_sp = Span {
            lo: span.lo,
            hi: span.hi,
            expn_id: NO_EXPANSION
        };
        let callsite = self.cx.codemap().with_expn_info(span.expn_id,
                                                        |ei| ei.map(|ei| ei.call_site.clone()));
        if callsite.is_none() {
            panic!("Callsite not found!");
        }
        let mut callsite = callsite.unwrap();

        if !self.span_map.contains_key(&callsite) {
            self.span_map.insert(key_sp, span);
            return;
        }

        let callee = self.cx.codemap().with_expn_info(span.expn_id,
                                                      |ei| ei.map(|ei| ei.callee.clone()));
        if callee.is_none() {
            panic!("Callee not found!");
        }
        let callee = callee.unwrap();

        callsite = self.span_map.get(&callsite).unwrap().clone();
        let info = ExpnInfo {
            call_site: callsite,
            callee: callee
        };
        let new_id = self.cx.codemap().record_expansion(info);
        self.span_map.insert(key_sp, Span { expn_id: new_id, .. span });
    }

    fn get(&mut self, span: Span) -> Span {
        let key_sp = Span { expn_id: NO_EXPANSION, .. span };
        return self.span_map.get(&key_sp).unwrap_or(&span).clone();
    }

    fn expand_crate(&mut self) {
        let mut krate = self.krates[self.index].clone();
        {
            let mut expander = MacroExpander::new(&mut self.cx);

            set_expander_fns!(expander,
                                expand_pat,
                                expand_type,
                                expand_expr,
                                expand_stmt,
                                expand_item,
                                expand_impl_item,
                                expand_opt_expr);

            krate = expand::expand_crate(&mut expander,
                                         Vec::new(),
                                         krate).0;
        }

        krate = self.fold_crate(krate);

        self.krates.push(krate);
        self.index += 1;
    }

    fn write_file(&self) {
        let mut fm = FileMap::new();
        for (path, module) in list_files(&self.krates[self.index], self.cx.codemap()) {
            let path = path.to_str().unwrap();
            let mut visitor = FmtVisitor::from_codemap(self.cx.parse_sess, &self.config);
            visitor.format_separate_mod(module);
            fm.push((path.to_owned(), visitor.buffer));
        }
        let out = &mut stdout();

        for (ref filename, ref text) in fm {
            let prefix = Path::new(filename).file_stem().and_then(|stem| stem.to_str()).unwrap_or("");
            let parent = Path::new(filename).parent().and_then(|path| path.to_str()).unwrap_or("");
            let file = format!("{}/{}Output{}.rs", parent, prefix, self.index);
            let res = rustfmt::filemap::write_file(text, &file, out, &self.config);
            if let Err(_) = res {
                print!("Error writing to file {}", file);
            }
        }
    }
}

//Walk over AST of expanded crate to patch up spans
impl<'a> Folder for ExpandData<'a> {
    fn fold_pat(&mut self, pat: P<ast::Pat>) -> P<ast::Pat> {
        println!("Folding pat");
        if pat.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_pat(pat, self);
        }

        self.insert(pat.span);
        fold::noop_fold_pat(pat.map(|elt| ast::Pat { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_ty(&mut self, ty: P<ast::Ty>) -> P<ast::Ty> {
        println!("Folding ty");
        if ty.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_ty(ty, self);
        }

        self.insert(ty.span);
        fold::noop_fold_ty(ty.map(|elt| ast::Ty { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_expr(&mut self, expr: P<ast::Expr>) -> P<ast::Expr> {
        println!("Folding expr");
        if expr.span.expn_id == NO_EXPANSION {
            return ptr::P(fold::noop_fold_expr(expr.unwrap(), self));
        }

        self.insert(expr.span);
        ptr::P(fold::noop_fold_expr(expr.map(|elt| ast::Expr { span: self.get(elt.span), .. elt }).unwrap(), self))
    }

    fn fold_opt_expr(&mut self, opt: P<ast::Expr>) -> Option<P<ast::Expr>> {
        println!("Folding optexpr");
        if opt.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_opt_expr(opt, self);
        }

        self.insert(opt.span);
        fold::noop_fold_opt_expr(opt.map(|elt| ast::Expr { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_item(&mut self, item: P<ast::Item>) -> SmallVector<P<ast::Item>> {
        println!("Folding item");
        if item.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_item(item, self);
        }

        self.insert(item.span);
        return fold::noop_fold_item(item.map(|elt| ast::Item { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_stmt(&mut self, stmt: ast::Stmt) -> SmallVector<ast::Stmt> {
        println!("Folding statemet");
        if stmt.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_stmt(stmt, self);
        }
        println!("Span: {}", self.cx.codemap().span_to_expanded_string(stmt.span.clone()));
        println!("Got Span: {}", self.cx.codemap().span_to_expanded_string(self.get(stmt.span.clone())));

        self.insert(stmt.span);
        return fold::noop_fold_stmt(ast::Stmt { span: self.get(stmt.span), .. stmt }, self)
    }

    fn fold_impl_item(&mut self, item: ast::ImplItem) -> SmallVector<ast::ImplItem> {
        println!("Folding implitem");
        if item.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_impl_item(item, self);
        }

        self.insert(item.span);
        return fold::noop_fold_impl_item(ast::ImplItem { span: self.get(item.span), .. item }, self)
    }

    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
        fold::noop_fold_mac(mac, self)
    }
}

// Struct for checking if expansion is required.
// (Checking if AST contains macros)
struct MacChecker {
    has_mac: bool,
}

impl MacChecker {

    fn new() -> MacChecker {
        MacChecker { has_mac: false }
    }

    fn check_finished(&mut self, data: &ExpandData) -> bool {
        self.has_mac = false;
        self.fold_crate(data.krates[data.index].clone());
        !self.has_mac
    }
}

impl Folder for MacChecker {
    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
        self.has_mac = true;
        mac //No need to expand further
    }
}

// Given some filepath, repeatedly expand and write output until no further expansion possible
fn main() {
    let codemap = Rc::new(CodeMap::new());
    let tty_handler = Handler::with_tty_emitter(ColorConfig::Auto,
                                                None,
                                                true,
                                                false,
                                                codemap.clone());
    let session = ParseSess::with_span_handler(tty_handler, codemap.clone());
    let mut loader = DummyMacroLoader;
    let mut data = ExpandData::new(&session, &mut loader);
    data.write_file();
    while !MacChecker::new().check_finished(&data) {
        data.expand_crate();
        data.write_file();
    }
}<commit_msg>Removed debug code, added trait item support<commit_after>extern crate rustfmt;
extern crate syntex_syntax as syntax;

use syntax::ast;
use syntax::ext::base::{ExtCtxt, DummyMacroLoader};
use syntax::ext::expand;
use syntax::ext::expand::{ExpansionConfig, MacroExpander};
use syntax::codemap::{CodeMap, Span, ExpnInfo, NO_EXPANSION};
use syntax::errors::Handler;
use syntax::errors::emitter::{ColorConfig};
use syntax::fold::{self, Folder};
use syntax::parse::{self, ParseSess};
use syntax::ptr::{self, P};
use syntax::util::small_vector::SmallVector;

use rustfmt::filemap::FileMap;
use rustfmt::config::{Config, WriteMode};
use rustfmt::modules::list_files;
use rustfmt::visitor::FmtVisitor;

use std::collections::HashMap;
use std::io::stdout;
use std::env;
use std::path::Path;
use std::rc::Rc;

// Small macro to simplify setting the full-expansion closures to the identity closure.
macro_rules! set_expander_fns {
    ($expander:ident,
        $( $expand:ident ),*) => {{
        $( $expander.$expand = Rc::new(Box::new(|_, node| node )); )*
    }}
}

struct ExpandData<'a> {
    config: Config,
    cx: ExtCtxt<'a>,
    krates: Vec<ast::Crate>,
    index: usize,
    span_map: HashMap<Span, Span>,
}

impl<'a> ExpandData<'a> {
    fn new(sess: &'a ParseSess, loader: &'a mut DummyMacroLoader) -> ExpandData<'a> {
        let args: Vec<String> = env::args().collect();
        if args.len() < 2 {
            panic!("Please supply a filepath to parse.")
        }
        if args.len() > 2 {
            panic!("Too many arguments. Please supply a single filepath.")
        }
        let filename = args[1].clone();

        let mut config = Config::default();
        config.write_mode = WriteMode::Overwrite;

        let ex_cfg = ExpansionConfig::default(filename.clone());
        let mut krates = vec!();
        krates.push(parse::parse_crate_from_file(&Path::new(&filename),
                                                 Vec::new(), sess).unwrap());
        let ecx = ExtCtxt::new(sess,
                               krates[0].config.clone(),
                               ex_cfg,
                               loader);
        ExpandData {
            config: config,
            cx: ecx,
            krates: krates,
            index: 0,
            span_map: HashMap::new(),
        }
    }

    fn insert(&mut self, span: Span) {
        if span.expn_id == NO_EXPANSION {
            return;
        }

        let key_sp = Span {
            lo: span.lo,
            hi: span.hi,
            expn_id: NO_EXPANSION
        };
        let callsite = self.cx.codemap().with_expn_info(span.expn_id,
                                                        |ei| ei.map(|ei| ei.call_site.clone()));
        if callsite.is_none() {
            panic!("Callsite not found!");
        }
        let mut callsite = callsite.unwrap();

        if !self.span_map.contains_key(&callsite) {
            self.span_map.insert(key_sp, span);
            return;
        }

        let callee = self.cx.codemap().with_expn_info(span.expn_id,
                                                      |ei| ei.map(|ei| ei.callee.clone()));
        if callee.is_none() {
            panic!("Callee not found!");
        }
        let callee = callee.unwrap();

        callsite = self.span_map.get(&callsite).unwrap().clone();
        let info = ExpnInfo {
            call_site: callsite,
            callee: callee
        };
        let new_id = self.cx.codemap().record_expansion(info);
        self.span_map.insert(key_sp, Span { expn_id: new_id, .. span });
    }

    fn get(&mut self, span: Span) -> Span {
        let key_sp = Span { expn_id: NO_EXPANSION, .. span };
        return self.span_map.get(&key_sp).unwrap_or(&span).clone();
    }

    fn expand_crate(&mut self) {
        let mut krate = self.krates[self.index].clone();
        {
            let mut expander = MacroExpander::new(&mut self.cx);

            set_expander_fns!(expander,
                                expand_pat,
                                expand_type,
                                expand_expr,
                                expand_stmt,
                                expand_item,
                                expand_impl_item,
                                expand_trait_item,
                                expand_opt_expr);

            krate = expand::expand_crate_with_expander(&mut expander,
                                                       Vec::new(),
                                                       krate).0;
        }

        krate = self.fold_crate(krate);

        self.krates.push(krate);
        self.index += 1;
    }

    fn write_file(&self) {
        let mut fm = FileMap::new();
        for (path, module) in list_files(&self.krates[self.index], self.cx.codemap()) {
            let path = path.to_str().unwrap();
            let mut visitor = FmtVisitor::from_codemap(self.cx.parse_sess, &self.config);
            visitor.format_separate_mod(module);
            fm.push((path.to_owned(), visitor.buffer));
        }
        let out = &mut stdout();

        for (ref filename, ref text) in fm {
            let prefix = Path::new(filename).file_stem()
                         .and_then(|stem| stem.to_str()).unwrap_or("");
            let parent = Path::new(filename).parent()
                         .and_then(|path| path.to_str()).unwrap_or("");
            let file = format!("{}/{}Output{}.rs", parent, prefix, self.index);
            let res = rustfmt::filemap::write_file(text, &file, out, &self.config);
            if let Err(_) = res {
                print!("Error writing to file {}", file);
            }
        }
    }
}

//Walk over AST of expanded crate to patch up spans
impl<'a> Folder for ExpandData<'a> {
    fn fold_pat(&mut self, pat: P<ast::Pat>) -> P<ast::Pat> {
        if pat.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_pat(pat, self);
        }

        self.insert(pat.span);
        fold::noop_fold_pat(pat.map(|elt| ast::Pat { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_ty(&mut self, ty: P<ast::Ty>) -> P<ast::Ty> {
        if ty.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_ty(ty, self);
        }

        self.insert(ty.span);
        fold::noop_fold_ty(ty.map(|elt| ast::Ty { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_expr(&mut self, expr: P<ast::Expr>) -> P<ast::Expr> {
        if expr.span.expn_id == NO_EXPANSION {
            return ptr::P(fold::noop_fold_expr(expr.unwrap(), self));
        }

        self.insert(expr.span);
        ptr::P(fold::noop_fold_expr(expr.map(|elt| ast::Expr { span: self.get(elt.span), .. elt })
                                        .unwrap(), self))
    }

    fn fold_opt_expr(&mut self, opt: P<ast::Expr>) -> Option<P<ast::Expr>> {
        if opt.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_opt_expr(opt, self);
        }

        self.insert(opt.span);
        fold::noop_fold_opt_expr(opt.map(|elt| ast::Expr { span: self.get(elt.span), .. elt }),
                                 self)
    }

    fn fold_item(&mut self, item: P<ast::Item>) -> SmallVector<P<ast::Item>> {
        if item.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_item(item, self);
        }

        self.insert(item.span);
        return fold::noop_fold_item(item.map(|elt| ast::Item { span: self.get(elt.span), .. elt }), self)
    }

    fn fold_stmt(&mut self, stmt: ast::Stmt) -> SmallVector<ast::Stmt> {
        if stmt.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_stmt(stmt, self);
        }

        self.insert(stmt.span);
        return fold::noop_fold_stmt(ast::Stmt { span: self.get(stmt.span), .. stmt }, self)
    }

    fn fold_impl_item(&mut self, item: ast::ImplItem) -> SmallVector<ast::ImplItem> {
        if item.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_impl_item(item, self);
        }

        self.insert(item.span);
        return fold::noop_fold_impl_item(ast::ImplItem { span: self.get(item.span), .. item }, self)
    }

    fn fold_trait_item(&mut self, item: ast::TraitItem) -> SmallVector<ast::TraitItem> {
        if item.span.expn_id == NO_EXPANSION {
            return fold::noop_fold_trait_item(item, self);
        }

        self.insert(item.span);
        return fold::noop_fold_trait_item(ast::TraitItem { span: self.get(item.span), .. item }, self)
    }

    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
        fold::noop_fold_mac(mac, self)
    }
}

// Struct for checking if expansion is required.
// (Checking if AST contains macros)
struct MacChecker {
    has_mac: bool,
}

impl MacChecker {

    fn new() -> MacChecker {
        MacChecker { has_mac: false }
    }

    fn check_finished(&mut self, data: &ExpandData) -> bool {
        self.has_mac = false;
        self.fold_crate(data.krates[data.index].clone());
        !self.has_mac
    }
}

impl Folder for MacChecker {
    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
        self.has_mac = true;
        mac //No need to expand further
    }
}

// Given some filepath, repeatedly expand and write output until no further expansion possible
fn main() {
    let codemap = Rc::new(CodeMap::new());
    let tty_handler = Handler::with_tty_emitter(ColorConfig::Auto,
                                                None,
                                                true,
                                                false,
                                                codemap.clone());
    let session = ParseSess::with_span_handler(tty_handler, codemap.clone());
    let mut loader = DummyMacroLoader;
    let mut data = ExpandData::new(&session, &mut loader);
    data.write_file();
    while !MacChecker::new().check_finished(&data) {
        data.expand_crate();
        data.write_file();
    }
}<|endoftext|>