<commit_before>#![feature(globs)]

extern crate serialize;

extern crate conduit;
extern crate "conduit-middleware" as middleware;
extern crate "conduit-utils" as utils;

use std::fmt;
use std::fmt::{Show, Formatter};
use serialize::{Decodable, json};

use conduit::Request;
use utils::RequestDelegator;
use middleware::Middleware;

pub struct BodyReader<T>;

pub trait JsonDecodable : Decodable<json::Decoder, json::DecoderError> {}
impl<T: Decodable<json::Decoder, json::DecoderError>> JsonDecodable for T {}

impl<T: JsonDecodable + 'static> Middleware for BodyReader<T> {
    fn before(&self, req: &mut Request) -> Result<(), Box<Show + 'static>> {
        let json: T = try!(decode::<T>(req.body()).map_err(|err| {
            box format!("Couldn't parse JSON: {}", show(&*err)) as Box<Show>
        }));

        req.mut_extensions().insert(json);
        Ok(())
    }
}

// Hack around the lack of impl Show for Box<Show>
struct Shower<'a> {
    inner: &'a Show + 'a
}

impl<'a> Show for Shower<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.inner.fmt(f)
    }
}

fn show<'a>(s: &'a Show) -> Shower<'a> {
    Shower { inner: s }
}

fn decode<T: JsonDecodable + 'static>(reader: &mut Reader) -> Result<T, Box<Show>> {
    let j = try!(json::from_reader(reader).map_err(|e| box e as Box<Show>));
    let mut decoder = json::Decoder::new(j);
    Decodable::decode(&mut decoder).map_err(|e| box e as Box<Show>)
}

pub fn json_params<'a, T: JsonDecodable + 'static>(req: &'a Request) -> Option<&'a T> {
    req.extensions().find::<T>()
}

#[cfg(test)]
mod tests {
    extern crate "conduit-test" as conduit_test;

    use {json_params, BodyReader};

    use std::collections::HashMap;
    use std::io::MemReader;
    use serialize::json;

    use conduit;
    use conduit::{Request, Response, Handler};
    use middleware::MiddlewareBuilder;

    #[deriving(PartialEq, Decodable, Encodable, Show)]
    struct Person {
        name: String,
        location: String
    }

    fn handler(req: &mut Request) -> Result<Response, ()> {
        let person = json_params::<Person>(req);
        let out = person.map(|p| json::encode(p)).expect("No JSON");

        Ok(Response {
            status: (200, "OK"),
            headers: HashMap::new(),
            body: box MemReader::new(out.into_bytes()) as Box<Reader + Send>
        })
    }

    #[test]
    fn test_body_params() {
        let mut req = conduit_test::MockRequest::new(conduit::Get, "/");
        req.with_body(r#"{ "name": "Alex Crichton", "location": "San Francisco" }"#);

        let mut middleware = MiddlewareBuilder::new(handler);
        middleware.add(BodyReader::<Person>);

        let mut res = middleware.call(&mut req).ok().expect("No response");
        let person = super::decode::<Person>(&mut *res.body).ok().expect("No JSON response");
        assert_eq!(person, Person {
            name: "Alex Crichton".to_string(),
            location: "San Francisco".to_string()
        });
    }
}
<commit_msg>Update to rust master<commit_after>#![feature(globs)]

extern crate serialize;

extern crate conduit;
extern crate "conduit-middleware" as middleware;
extern crate "conduit-utils" as utils;

use std::fmt;
use std::fmt::{Show, Formatter};
use serialize::{Decodable, json};

use conduit::Request;
use utils::RequestDelegator;
use middleware::Middleware;

pub struct BodyReader<T>;

pub trait JsonDecodable : Decodable<json::Decoder, json::DecoderError> {}
impl<T: Decodable<json::Decoder, json::DecoderError>> JsonDecodable for T {}

impl<T: JsonDecodable + 'static> Middleware for BodyReader<T> {
    fn before(&self, req: &mut Request) -> Result<(), Box<Show + 'static>> {
        let json: T = try!(decode::<T>(req.body()).map_err(|err| {
            box format!("Couldn't parse JSON: {}", show(&*err)) as Box<Show>
        }));

        req.mut_extensions().insert(json);
        Ok(())
    }
}

// Hack around the lack of impl Show for Box<Show>
struct Shower<'a> {
    inner: &'a (Show + 'a),
}

impl<'a> Show for Shower<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.inner.fmt(f)
    }
}

fn show<'a>(s: &'a Show) -> Shower<'a> {
    Shower { inner: s }
}

fn decode<T: JsonDecodable + 'static>(reader: &mut Reader) -> Result<T, Box<Show>> {
    let j = try!(json::from_reader(reader).map_err(|e| box e as Box<Show>));
    let mut decoder = json::Decoder::new(j);
    Decodable::decode(&mut decoder).map_err(|e| box e as Box<Show>)
}

pub fn json_params<'a, T: JsonDecodable + 'static>(req: &'a Request) -> Option<&'a T> {
    req.extensions().find::<T>()
}

#[cfg(test)]
mod tests {
    extern crate "conduit-test" as conduit_test;

    use {json_params, BodyReader};

    use std::collections::HashMap;
    use std::io::MemReader;
    use serialize::json;

    use conduit::{Request, Response, Handler, Method};
    use middleware::MiddlewareBuilder;

    #[deriving(PartialEq, Decodable, Encodable, Show)]
    struct Person {
        name: String,
        location: String
    }

    fn handler(req: &mut Request) -> Result<Response, ()> {
        let person = json_params::<Person>(req);
        let out = person.map(|p| json::encode(p)).expect("No JSON");

        Ok(Response {
            status: (200, "OK"),
            headers: HashMap::new(),
            body: box MemReader::new(out.into_bytes()) as Box<Reader + Send>
        })
    }

    #[test]
    fn test_body_params() {
        let mut req = conduit_test::MockRequest::new(Method::Get, "/");
        req.with_body(r#"{ "name": "Alex Crichton", "location": "San Francisco" }"#);

        let mut middleware = MiddlewareBuilder::new(handler);
        middleware.add(BodyReader::<Person>);

        let mut res = middleware.call(&mut req).ok().expect("No response");
        let person = super::decode::<Person>(&mut *res.body).ok().expect("No JSON response");
        assert_eq!(person, Person {
            name: "Alex Crichton".to_string(),
            location: "San Francisco".to_string()
        });
    }
}
<|endoftext|>


<commit_before>//! jsonrpc http server.
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_http_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_http_server::*;
//!
//! struct SayHello;
//! impl MethodCommand for SayHello {
//! 	fn execute(&self, _params: Params) -> Result<Value, Error> {
//! 		Ok(Value::String("hello".to_string()))
//! 	}
//! }
//!
//! fn main() {
//! 	let io = IoHandler::new();
//! 	io.add_method("say_hello", SayHello);
//! 	let _server = Server::start(&"127.0.0.1:3030".parse().unwrap(), Arc::new(io), vec![AccessControlAllowOrigin::Null]);
//! }
//! ```

extern crate hyper;
extern crate unicase;
extern crate jsonrpc_core as jsonrpc;

mod cors;

use std::ops::Deref;
use std::sync::{Arc, Mutex};
use std::io::{Read, Write};
use std::net::SocketAddr;
use hyper::header::{Headers, Allow, ContentType, AccessControlAllowHeaders};
use hyper::method::Method;
use hyper::mime;
use hyper::server::{Request, Response};
use hyper::{Next, Encoder, Decoder};
use hyper::net::HttpStream;
use unicase::UniCase;
use self::jsonrpc::{IoHandler};

pub use hyper::header::AccessControlAllowOrigin;

pub type ServerResult = Result<Server, RpcServerError>;

/// RPC Server startup error
#[derive(Debug)]
pub enum RpcServerError {
	IoError(std::io::Error),
	Other(hyper::error::Error),
}

impl From<hyper::error::Error> for RpcServerError {
	fn from(err: hyper::error::Error) -> Self {
		match err {
			hyper::error::Error::Io(e) => RpcServerError::IoError(e),
			e => RpcServerError::Other(e)
		}
	}
}

/// PanicHandling function
pub struct PanicHandler {
	pub handler: Arc<Mutex<Option<Box<Fn() -> () + Send + 'static>>>>
}

/// jsonrpc http request handler.
pub struct ServerHandler {
	panic_handler: PanicHandler,
	jsonrpc_handler: Arc<IoHandler>,
	cors_domains: Vec<AccessControlAllowOrigin>,
	request: String,
	origin: Option<String>,
	response: Option<String>,
	error_code:  hyper::status::StatusCode,
	write_pos: usize,
}

impl Drop for ServerHandler {
	fn drop(&mut self) {
		if ::std::thread::panicking() {
			let handler = self.panic_handler.handler.lock().unwrap();
			if let Some(ref h) = *handler.deref() {
				h();
			}
		}
	}
}

impl ServerHandler {
	/// Create new request handler.
	pub fn new(jsonrpc_handler: Arc<IoHandler>, cors_domains: Vec<AccessControlAllowOrigin>, panic_handler: PanicHandler) -> Self {
		ServerHandler {
			panic_handler: panic_handler,
			jsonrpc_handler: jsonrpc_handler,
			cors_domains: cors_domains,
			request: String::new(),
			origin: None,
			response: None,
			error_code: hyper::status::StatusCode::MethodNotAllowed,
			write_pos: 0,
		}
	}

	fn response_headers(&self, origin: &Option<String>) -> Headers {
		let mut headers = Headers::new();
		headers.set(
			Allow(vec![
				Method::Options, Method::Post
			])
		);
		headers.set(ContentType::json());
		headers.set(
			AccessControlAllowHeaders(vec![
				UniCase("origin".to_owned()),
				UniCase("content-type".to_owned()),
				UniCase("accept".to_owned()),
			])
		);

		if let Some(cors_domain) = cors::get_cors_header(&self.cors_domains, origin) {
			headers.set(cors_domain);
		}

		headers
	}
}

impl hyper::server::Handler<HttpStream> for ServerHandler {
	fn on_request(&mut self, request: Request) -> Next {
		// Read origin
		self.origin = cors::read_origin(&request);

		match *request.method() {
			// Don't validate content type on options
			Method::Options => {
				self.response = Some(String::new());
				Next::write()
			},
			Method::Post => {
				// Validate the ContentType header
				// to prevent Cross-Origin XHRs with text/plain
				let content_type = request.headers().get::<ContentType>();
				if let Some(&ContentType(mime::Mime(mime::TopLevel::Application, mime::SubLevel::Json, _))) = content_type {
					Next::read()
				} else {
					self.error_code = hyper::status::StatusCode::UnsupportedMediaType;
					// Just return error
					Next::write()
				}
			},
			_ => Next::write(),
		}
	}

	/// This event occurs each time the `Request` is ready to be read from.
	fn on_request_readable(&mut self, decoder: &mut Decoder<HttpStream>) -> Next {
		match decoder.read_to_string(&mut self.request) {
			Ok(0) => {
				self.response = self.jsonrpc_handler.handle_request(&self.request);
				match self.response {
					Some(ref mut r) => r.push('\n'),
					_ => {
						self.error_code = hyper::status::StatusCode::BadRequest;
					}
				}
				Next::write()
			}
			Ok(_) => {
				Next::read()
			}
			Err(e) => match e.kind() {
				::std::io::ErrorKind::WouldBlock => Next::read(),
				_ => {
					//trace!("Read error: {}", e);
					Next::end()
				}
			}
		}
	}

		/// This event occurs after the first time this handled signals `Next::write()`.
	fn on_response(&mut self, response: &mut Response) -> Next {
		*response.headers_mut() = self.response_headers(&self.origin);
		if self.response.is_none() {
			response.set_status(self.error_code);
		}
		Next::write()
	}

		/// This event occurs each time the `Response` is ready to be written to.
		fn on_response_writable(&mut self, encoder: &mut Encoder<HttpStream>) -> Next {
		if let Some(ref response) = self.response {
			let bytes = response.as_bytes();
			if bytes.len() == self.write_pos {
				Next::end()
			} else {
				match encoder.write(&bytes[self.write_pos ..]) {
					Ok(0) => {
						Next::write()
					}
					Ok(bytes) => {
						self.write_pos += bytes;
						Next::write()
					}
					Err(e) => match e.kind() {
						::std::io::ErrorKind::WouldBlock => Next::write(),
						_ => {
							//trace!("Write error: {}", e);
							Next::end()
						}
					}
				}
			}
		} else {
			Next::end()
		}
	}
}

/// jsonrpc http server.
///
/// ```no_run
/// extern crate jsonrpc_core;
/// extern crate jsonrpc_http_server;
///
/// use std::sync::Arc;
/// use jsonrpc_core::*;
/// use jsonrpc_http_server::*;
///
/// struct SayHello;
/// impl MethodCommand for SayHello {
/// 	fn execute(&self, _params: Params) -> Result<Value, Error> {
/// 		Ok(Value::String("hello".to_string()))
/// 	}
/// }
///
/// fn main() {
/// 	let io = IoHandler::new();
/// 	io.add_method("say_hello", SayHello);
/// 	let _server = Server::start(&"127.0.0.1:3030".parse().unwrap(), Arc::new(io), vec![AccessControlAllowOrigin::Null]);
/// }
/// ```
pub struct Server {
	server: Option<hyper::server::Listening>,
	panic_handler: Arc<Mutex<Option<Box<Fn() -> () + Send>>>>
}

impl Server {
	pub fn start(addr: &SocketAddr, jsonrpc_handler: Arc<IoHandler>, cors_domains: Vec<AccessControlAllowOrigin>) -> ServerResult {
		let panic_handler = Arc::new(Mutex::new(None));
		let panic_for_server = panic_handler.clone();
		let srv = try!(try!(hyper::Server::http(addr)).handle(move |_| {
			let handler = PanicHandler { handler: panic_for_server.clone() };
			ServerHandler::new(jsonrpc_handler.clone(), cors_domains.clone(), handler)
		}));
		Ok(Server {
			server: Some(srv),
			panic_handler: panic_handler,
		})
	}

	pub fn set_panic_handler<F>(&self, handler: F)
		where F : Fn() -> () + Send + 'static {
		*self.panic_handler.lock().unwrap() = Some(Box::new(handler));
	}
}

impl Drop for Server {
	fn drop(&mut self) {
		self.server.take().unwrap().close()
	}
}
<commit_msg>informative message when requiring json content type<commit_after>//! jsonrpc http server.
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_http_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_http_server::*;
//!
//! struct SayHello;
//! impl MethodCommand for SayHello {
//! 	fn execute(&self, _params: Params) -> Result<Value, Error> {
//! 		Ok(Value::String("hello".to_string()))
//! 	}
//! }
//!
//! fn main() {
//! 	let io = IoHandler::new();
//! 	io.add_method("say_hello", SayHello);
//! 	let _server = Server::start(&"127.0.0.1:3030".parse().unwrap(), Arc::new(io), vec![AccessControlAllowOrigin::Null]);
//! }
//! ```

extern crate hyper;
extern crate unicase;
extern crate jsonrpc_core as jsonrpc;

mod cors;

use std::ops::Deref;
use std::sync::{Arc, Mutex};
use std::io::{Read, Write};
use std::net::SocketAddr;
use hyper::header::{Headers, Allow, ContentType, AccessControlAllowHeaders};
use hyper::method::Method;
use hyper::mime;
use hyper::server::{Request, Response};
use hyper::{Next, Encoder, Decoder};
use hyper::net::HttpStream;
use unicase::UniCase;
use self::jsonrpc::{IoHandler};

pub use hyper::header::AccessControlAllowOrigin;

pub type ServerResult = Result<Server, RpcServerError>;

/// RPC Server startup error
#[derive(Debug)]
pub enum RpcServerError {
	IoError(std::io::Error),
	Other(hyper::error::Error),
}

impl From<hyper::error::Error> for RpcServerError {
	fn from(err: hyper::error::Error) -> Self {
		match err {
			hyper::error::Error::Io(e) => RpcServerError::IoError(e),
			e => RpcServerError::Other(e)
		}
	}
}

/// PanicHandling function
pub struct PanicHandler {
	pub handler: Arc<Mutex<Option<Box<Fn() -> () + Send + 'static>>>>
}

/// jsonrpc http request handler.
pub struct ServerHandler {
	panic_handler: PanicHandler,
	jsonrpc_handler: Arc<IoHandler>,
	cors_domains: Vec<AccessControlAllowOrigin>,
	request: String,
	origin: Option<String>,
	response: Option<String>,
	error_code:  hyper::status::StatusCode,
	write_pos: usize,
}

impl Drop for ServerHandler {
	fn drop(&mut self) {
		if ::std::thread::panicking() {
			let handler = self.panic_handler.handler.lock().unwrap();
			if let Some(ref h) = *handler.deref() {
				h();
			}
		}
	}
}

impl ServerHandler {
	/// Create new request handler.
	pub fn new(jsonrpc_handler: Arc<IoHandler>, cors_domains: Vec<AccessControlAllowOrigin>, panic_handler: PanicHandler) -> Self {
		ServerHandler {
			panic_handler: panic_handler,
			jsonrpc_handler: jsonrpc_handler,
			cors_domains: cors_domains,
			request: String::new(),
			origin: None,
			response: None,
			error_code: hyper::status::StatusCode::MethodNotAllowed,
			write_pos: 0,
		}
	}

	fn response_headers(&self, origin: &Option<String>) -> Headers {
		let mut headers = Headers::new();
		headers.set(
			Allow(vec![
				Method::Options, Method::Post
			])
		);
		headers.set(ContentType::json());
		headers.set(
			AccessControlAllowHeaders(vec![
				UniCase("origin".to_owned()),
				UniCase("content-type".to_owned()),
				UniCase("accept".to_owned()),
			])
		);

		if let Some(cors_domain) = cors::get_cors_header(&self.cors_domains, origin) {
			headers.set(cors_domain);
		}

		headers
	}
}

impl hyper::server::Handler<HttpStream> for ServerHandler {
	fn on_request(&mut self, request: Request) -> Next {
		// Read origin
		self.origin = cors::read_origin(&request);

		match *request.method() {
			// Don't validate content type on options
			Method::Options => {
				self.response = Some(String::new());
				Next::write()
			},
			Method::Post => {
				// Validate the ContentType header
				// to prevent Cross-Origin XHRs with text/plain
				let content_type = request.headers().get::<ContentType>();
				if let Some(&ContentType(mime::Mime(mime::TopLevel::Application, mime::SubLevel::Json, _))) = content_type {
					Next::read()
				} else {
					self.error_code = hyper::status::StatusCode::UnsupportedMediaType;
					// Just return error
					Next::write()
				}
			},
			_ => Next::write(),
		}
	}

	/// This event occurs each time the `Request` is ready to be read from.
	fn on_request_readable(&mut self, decoder: &mut Decoder<HttpStream>) -> Next {
		match decoder.read_to_string(&mut self.request) {
			Ok(0) => {
				self.response = self.jsonrpc_handler.handle_request(&self.request);
				match self.response {
					Some(ref mut r) => r.push('\n'),
					_ => {
						self.error_code = hyper::status::StatusCode::BadRequest;
					}
				}
				Next::write()
			}
			Ok(_) => {
				Next::read()
			}
			Err(e) => match e.kind() {
				::std::io::ErrorKind::WouldBlock => Next::read(),
				_ => {
					//trace!("Read error: {}", e);
					Next::end()
				}
			}
		}
	}

	/// This event occurs after the first time this handled signals `Next::write()`.
	fn on_response(&mut self, response: &mut Response) -> Next {
		*response.headers_mut() = self.response_headers(&self.origin);
		if self.response.is_none() {
			response.set_status(self.error_code);
		}
		Next::write()
	}

	/// This event occurs each time the `Response` is ready to be written to.
	fn on_response_writable(&mut self, encoder: &mut Encoder<HttpStream>) -> Next {
		if self.error_code == hyper::status::StatusCode::UnsupportedMediaType {
			self.response = Some("Content-Type: application/json required\n".into());
		}
		if let Some(ref response) = self.response {
			let bytes = response.as_bytes();
			if bytes.len() == self.write_pos {
				Next::end()
			} else {
				match encoder.write(&bytes[self.write_pos..]) {
					Ok(0) => {
						Next::write()
					}
					Ok(bytes) => {
						self.write_pos += bytes;
						Next::write()
					}
					Err(e) => match e.kind() {
						::std::io::ErrorKind::WouldBlock => Next::write(),
						_ => {
							//trace!("Write error: {}", e);
							Next::end()
						}
					}
				}
			}
		} else {
			Next::end()
		}
	}
}

/// jsonrpc http server.
///
/// ```no_run
/// extern crate jsonrpc_core;
/// extern crate jsonrpc_http_server;
///
/// use std::sync::Arc;
/// use jsonrpc_core::*;
/// use jsonrpc_http_server::*;
///
/// struct SayHello;
/// impl MethodCommand for SayHello {
/// 	fn execute(&self, _params: Params) -> Result<Value, Error> {
/// 		Ok(Value::String("hello".to_string()))
/// 	}
/// }
///
/// fn main() {
/// 	let io = IoHandler::new();
/// 	io.add_method("say_hello", SayHello);
/// 	let _server = Server::start(&"127.0.0.1:3030".parse().unwrap(), Arc::new(io), vec![AccessControlAllowOrigin::Null]);
/// }
/// ```
pub struct Server {
	server: Option<hyper::server::Listening>,
	panic_handler: Arc<Mutex<Option<Box<Fn() -> () + Send>>>>
}

impl Server {
	pub fn start(addr: &SocketAddr, jsonrpc_handler: Arc<IoHandler>, cors_domains: Vec<AccessControlAllowOrigin>) -> ServerResult {
		let panic_handler = Arc::new(Mutex::new(None));
		let panic_for_server = panic_handler.clone();
		let srv = try!(try!(hyper::Server::http(addr)).handle(move |_| {
			let handler = PanicHandler { handler: panic_for_server.clone() };
			ServerHandler::new(jsonrpc_handler.clone(), cors_domains.clone(), handler)
		}));
		Ok(Server {
			server: Some(srv),
			panic_handler: panic_handler,
		})
	}

	pub fn set_panic_handler<F>(&self, handler: F)
		where F : Fn() -> () + Send + 'static {
		*self.panic_handler.lock().unwrap() = Some(Box::new(handler));
	}
}

impl Drop for Server {
	fn drop(&mut self) {
		self.server.take().unwrap().close()
	}
}
<|endoftext|>


<commit_before>extern crate websocket;

use std;
use std::sync::{Arc, Mutex};
use std::thread;
use std::sync::mpsc::channel;

use websocket::{Message, Sender, Receiver};
use websocket::message::Type;
use websocket::client::request::Url;

use hyper::header::{Cookie, CookieJar};

use rustc_serialize::json;
use rustc_serialize::json::Json;

use time;

// making a move
// out {"t":"move","d":{"from":"e2","to":"e4","b":1}}
// in {t: "ack"}
// in {"v":1,"t":"move","d":{"uci":"e2e4","san":"e4","fen":"rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR","ply":1,"clock":{"white":172800,"black":172800},"dests":{"b8":"a6c6","g8":"f6h6","h7":"h6h5","d7":"d6d5","g7":"g6g5","c7":"c6c5","f7":"f6f5","b7":"b6b5","e7":"e6e5","a7":"a6a5"},"crazyhouse":{"pockets":[{},{}]}}}
// other player moved
// in {"v":2,"t":"move","d":{"uci":"g7g5","san":"g5","fen":"rnbqkbnr/pppppp1p/8/6p1/4P3/8/PPPP1PPP/RNBQKBNR","ply":2,"clock":{"white":172800,"black":172800},"dests":{"a2":"a3a4","g1":"f3h3e2","d1":"e2f3g4h5","e1":"e2","d2":"d3d4","b1":"a3c3","e4":"e5","f1":"e2d3c4b5a6","h2":"h3h4","b2":"b3b4","f2":"f3f4","c2":"c3c4","g2":"g3g4"},"crazyhouse":{"pockets":[{},{}]}}}

#[derive(RustcEncodable)]
pub struct PingPacket {
    t: String,
    v: i64
}

#[derive(RustcEncodable)]
pub struct MovePacket {
    t: String, // this should be hardcoded
    d: Dest,
}

#[derive(RustcEncodable)]
pub struct Dest {
    pub from: String,
    pub to: String,
    // pub: i8, // bool, but 0 or 1
}

pub fn connect(jar: &CookieJar<'static>, base_url: String, sri: String, pov: Arc<Mutex<super::Pov>>) {

    let mut url;
    {
        let pov = pov.lock().unwrap();
        let version = match pov.player.version {
            Some(v) => v,
            None => 0
        };
        url = Url::parse(&format!("ws://{}{}?sri={}&version={}", base_url, pov.url.socket.clone(), sri, version)).unwrap();
    }
    let version = Arc::new(Mutex::new(0));

    //println!("Connecting to {}", url);

    // TODO: this unwrap fails when url is wrong, port for example
    let mut request = websocket::Client::connect(url).unwrap();
    request.headers.set(Cookie::from_cookie_jar(&jar));

    let response = request.send().unwrap();

    match response.validate() {
        Ok(()) => (),
        Err(e) => {
            error!("reponse validation error: {:?}", e);
            debug!("reponse status: {:?}, headers: {:?}", response.status, response.headers);
            return;
        }
    };

    //println!("Successfully connected");

    let (mut sender, mut receiver) = response.begin().split();

    let (tx, rx) = channel();

    let tx_1 = tx.clone();

    let _send_loop = thread::spawn(move || {
        loop {
            let message: Message = match rx.recv() {
                Ok(m) => m,
                Err(e) => {
                    //println!("Send Loop: {:?}", e);
                    return;
                }
            };
            match message.opcode {
                Type::Close => {
                    let _ = sender.send_message(&message);
                    return;
                }
                _ => (),
            }
            match sender.send_message(&message) {
                Ok(()) => (),
                Err(e) => {
                    //println!("Send Loop: {:?}", e);
                    let _ = sender.send_message(&Message::close());
                    return;
                }
            }
        }
    });

    let pov_1 = pov.clone();
    let _receive_loop = thread::spawn(move || {

        let handle = |obj: json::Object| {
            let mut pov = pov_1.lock().unwrap();
            match obj.get("v") {
                Some(&Json::U64(v)) => {
                    pov.player.version = Some(v as i64);
                },
                _ => ()
            }
            match obj.get("t") {
                Some(&Json::String(ref t)) if t == "move" => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    let fen = d.get("fen").unwrap().as_string().unwrap();
                    pov.game.fen = fen.to_string();
                    //cli::render_fen(fen);
                },
                Some(&Json::String(ref t)) if t == "clock" => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    match pov.clock.as_mut() {
                        Some(clock) => {
                            clock.white = d.get("white").unwrap().as_f64().unwrap();
                            clock.black = d.get("black").unwrap().as_f64().unwrap();
                        },
                        None => ()
                    };
                },
                Some(&Json::String(ref t)) if t == "end" => {
                    let _ = tx_1.send(Message::close());
                    // exit
                },
                //Some(&Json::String(ref t)) => println!("unhandled: {:?}", obj),
                _ => ()
            }
        };

        for message in receiver.incoming_messages() {
            let message: Message = match message {
                Ok(m) => m,
                Err(e) => {
                    //trace!("Receive Loop: {:?}", e);
                    let _ = tx_1.send(Message::close());
                    return;
                }
            };
            match message.opcode {
                Type::Close => {
                    let _ = tx_1.send(Message::close());
                    return;
                }
                Type::Ping => match tx_1.send(Message::pong(message.payload)) {
                    Ok(()) => (),
                    Err(e) => {
                        //println!("Receive Loop: {:?}", e);
                        return;
                    }
                },
                Type::Text => {
                    //println!("Receive Loop: {:?}", payload);
                    let json = Json::from_str(std::str::from_utf8(&message.payload.into_owned()).unwrap()).unwrap();
                    if json.is_object() {
                        let obj = json.as_object().unwrap();
                        let t = obj.get("t");
                        //println!("  type: {:?}", t);
                        //println!("full: {:?}", obj);
                        match t {
                            Some(&Json::String(ref t)) if t == "n" => { // pong
                            },
                            Some(&Json::String(ref t)) if t == "b" => { // batch
                                for item in obj.get("d").unwrap().as_array().unwrap().iter() {
                                   let obj = item.as_object().unwrap();
                                   handle(obj.clone());
                                }
                            }
                            _ => {
                                handle(obj.clone());
                            }
                        }
                    }
                }
                _ => ()//println!("Receive Loop: {:?}", message),
            }
        }
    });

    let tx_2 = tx.clone();

    // ping loop
    thread::spawn(move || {
        loop {
            std::thread::sleep_ms(1000);

            let pov = pov.lock().unwrap();
            match pov.player.version {
                Some(version) => {
                    let ping_packet = PingPacket {
                        t: "p".to_string(),
                        v: version
                    };

                    let message = Message::text(json::encode(&ping_packet).unwrap());

                    match tx_2.send(message) {
                        Ok(()) => (),
                        Err(e) => {
                            //println!("Main Loop: {:?}", e);
                        }
                    }
                },
                None => () // TODO: error
            };
        }
    });

    /*loop {
        let mut input = String::new();

        stdin().read_line(&mut input).unwrap();

        let trimmed = input.trim();

        let message = match trimmed {
            "/close" => {
                let _ = tx.send(Message::close(None));
                break;
            }
            "/ping" => Message::Ping(b"{t: 'p', v: 300}".to_vec()),
            _ => Message::Text(trimmed.to_string()),
        };

        match tx.send(message) {
            Ok(()) => (),
            Err(e) => {
                //println!("Main Loop: {:?}", e);
                break;
            }
        }
    }*/

    //let _ = send_loop.join();
    //let _ = receive_loop.join();
}

<commit_msg>logging and cleaned ping loop<commit_after>extern crate websocket;

use std;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;

use websocket::{Message, Sender, Receiver};
use websocket::message::Type;
use websocket::client::request::Url;

use hyper::header::{Cookie, CookieJar};

use rustc_serialize::json;
use rustc_serialize::json::Json;

use time;

// making a move
// out {"t":"move","d":{"from":"e2","to":"e4","b":1}}
// in {t: "ack"}
// in {"v":1,"t":"move","d":{"uci":"e2e4","san":"e4","fen":"rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR","ply":1,"clock":{"white":172800,"black":172800},"dests":{"b8":"a6c6","g8":"f6h6","h7":"h6h5","d7":"d6d5","g7":"g6g5","c7":"c6c5","f7":"f6f5","b7":"b6b5","e7":"e6e5","a7":"a6a5"},"crazyhouse":{"pockets":[{},{}]}}}
// other player moved
// in {"v":2,"t":"move","d":{"uci":"g7g5","san":"g5","fen":"rnbqkbnr/pppppp1p/8/6p1/4P3/8/PPPP1PPP/RNBQKBNR","ply":2,"clock":{"white":172800,"black":172800},"dests":{"a2":"a3a4","g1":"f3h3e2","d1":"e2f3g4h5","e1":"e2","d2":"d3d4","b1":"a3c3","e4":"e5","f1":"e2d3c4b5a6","h2":"h3h4","b2":"b3b4","f2":"f3f4","c2":"c3c4","g2":"g3g4"},"crazyhouse":{"pockets":[{},{}]}}}

#[derive(RustcEncodable)]
struct PingPacket {
    t: String,
    v: i64
}

impl PingPacket {
    pub fn new(version: i64) -> PingPacket {
        PingPacket {
            t: "p".to_string(),
            v: version
        }
    }

    pub fn to_message(&self) -> Message<'static> {
        Message::text(json::encode(self).unwrap())
    }
}

#[derive(RustcEncodable)]
pub struct MovePacket {
    t: String, // this should be hardcoded
    d: Dest,
}

#[derive(RustcEncodable)]
pub struct Dest {
    pub from: String,
    pub to: String,
    // pub: i8, // bool, but 0 or 1
}

pub fn connect(jar: &CookieJar<'static>, base_url: String, sri: String, pov: Arc<Mutex<super::Pov>>) {

    let mut url;
    {
        let pov = pov.lock().unwrap();
        let version = match pov.player.version {
            Some(v) => v,
            None => 0
        };
        url = Url::parse(&format!("ws://{}{}?sri={}&version={}", base_url, pov.url.socket.clone(), sri, version)).unwrap();
    }
    let version = Arc::new(Mutex::new(0));

    // TODO: this unwrap fails when url is wrong, port for example
    let mut request = websocket::Client::connect(url).unwrap();
    request.headers.set(Cookie::from_cookie_jar(&jar));

    let response = request.send().unwrap();

    match response.validate() {
        Ok(()) => (),
        Err(e) => {
            error!("reponse validation error: {:?}", e);
            debug!("reponse status: {:?}, headers: {:?}", response.status, response.headers);
            return;
        }
    };

    //println!("Successfully connected");

    let (mut sender, mut receiver) = response.begin().split();

    let (tx, rx) = mpsc::channel();

    let tx_1 = tx.clone();

    let _send_loop = thread::spawn(move || {
        loop {
            let message: Message = match rx.recv() {
                Ok(m) => m,
                Err(e) => {
                    //println!("Send Loop: {:?}", e);
                    return;
                }
            };
            trace!("Transmit raw: {:?}", message);
            debug!("Transmit msg: {}",
                   &std::str::from_utf8(&message.payload).unwrap());
            match message.opcode {
                Type::Close => {
                    let _ = sender.send_message(&message);
                    return;
                }
                _ => (),
            }
            match sender.send_message(&message) {
                Ok(()) => (),
                Err(e) => {
                    //println!("Send Loop: {:?}", e);
                    let _ = sender.send_message(&Message::close());
                    return;
                }
            }
        }
    });

    let pov_1 = pov.clone();
    let _receive_loop = thread::spawn(move || {

        let handle = |obj: json::Object| {
            let mut pov = pov_1.lock().unwrap();
            match obj.get("v") {
                Some(&Json::U64(v)) => {
                    pov.player.version = Some(v as i64);
                },
                _ => ()
            }
            match obj.get("t") {
                Some(&Json::String(ref t)) if t == "move" => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    let fen = d.get("fen").unwrap().as_string().unwrap();
                    pov.game.fen = fen.to_string();
                    //cli::render_fen(fen);
                },
                Some(&Json::String(ref t)) if t == "clock" => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    match pov.clock.as_mut() {
                        Some(clock) => {
                            clock.white = d.get("white").unwrap().as_f64().unwrap();
                            clock.black = d.get("black").unwrap().as_f64().unwrap();
                        },
                        None => ()
                    };
                },
                Some(&Json::String(ref t)) if t == "end" => {
                    let _ = tx_1.send(Message::close());
                    // exit
                },
                //Some(&Json::String(ref t)) => println!("unhandled: {:?}", obj),
                _ => ()
            }
        };

        for message in receiver.incoming_messages() {
            let message: Message = match message {
                Ok(m) => m,
                Err(e) => {
                    warn!("Receive loop error: {:?}", e);
                    let _ = tx_1.send(Message::close());
                    return;
                }
            };
            match message.opcode {
                Type::Close => {
                    let _ = tx_1.send(Message::close());
                    return;
                }
                Type::Ping => match tx_1.send(Message::pong(message.payload)) {
                    Ok(()) => (),
                    Err(e) => {
                        warn!("Could not send pong: {:?}", e);
                        return;
                    }
                },
                Type::Text => {
                    let str_payload = std::str::from_utf8(&message.payload).unwrap();
                    trace!("Received raw: {:?}", &message.payload);
                    trace!("Received str: {}", str_payload);
                    let json = Json::from_str(str_payload).unwrap();
                    debug!("Received obj: {:?}", json);
                    if json.is_object() {
                        let obj = json.as_object().unwrap();
                        let t = obj.get("t");
                        match t {
                            Some(&Json::String(ref t)) if t == "n" => { // pong
                            },
                            Some(&Json::String(ref t)) if t == "b" => { // batch
                                for item in obj.get("d").unwrap().as_array().unwrap().iter() {
                                   let obj = item.as_object().unwrap();
                                   handle(obj.clone());
                                }
                            }
                            _ => {
                                handle(obj.clone());
                            }
                        }
                    }
                }
                _ => debug!("Unhandled message: {:?}", message),
            }
        }
    });

    let tx_2 = tx.clone();

    thread::spawn(move || loop {
        std::thread::sleep_ms(1000);

        pov.lock().ok()
            .and_then(|p| p.player.version )
            .map(|v| PingPacket::new(v).to_message())
            .map(|p| tx_2.send(p).unwrap());
    });

    //let _ = send_loop.join();
    //let _ = receive_loop.join();
}

<|endoftext|>


<commit_before>use gameboy::{Interconnect, Interrupt, Irq, Memory};
use gameboy::frame::{Color, Frame};

const VRAM_SIZE: usize = 0x4000;
const SPRITE_DATA_SIZE: usize = 0xA0;

pub struct GpuStat {
    pub coincidence_interrupt_enabled: bool,
    pub OAM_interrupt_enabled: bool,
    pub VBlank_interrupt_enabled: bool,
    pub HBlank_interrupt_enabled: bool,
}

impl GpuStat {
    pub fn new() -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: false,
            OAM_interrupt_enabled: false,
            VBlank_interrupt_enabled: false,
            HBlank_interrupt_enabled: false,
        }
    }

    pub fn from_u8(b: u8) -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: b & 0x40 == 0x40,
            OAM_interrupt_enabled: b & 0x20 == 0x20,
            VBlank_interrupt_enabled: b & 0x10 == 0x10,
            HBlank_interrupt_enabled: b & 0x08 == 0x08,
        }
    }

    pub fn to_u8(&self, gpu: &Gpu) -> u8 {
        (if self.coincidence_interrupt_enabled {
            0x40
        } else {
            0
        }) | (if self.OAM_interrupt_enabled { 0x20 } else { 0 }) |
        (if self.VBlank_interrupt_enabled {
            0x10
        } else {
            0
        }) |
        (if self.HBlank_interrupt_enabled {
            0x08
        } else {
            0
        }) | (if gpu.ly == gpu.lyc { 0x04 } else { 0 }) | gpu.mode.to_u8()
    }
}

pub enum GpuMode {
    HBlank,
    VBlank,
    SearchingRam,
    TransferringData,
}

impl GpuMode {
    fn cycles(&self, scroll_x: u8) -> isize {
        let scroll_adjust = match scroll_x % 0x08 {
            5...7 => 2,
            1...4 => 1,
            _ => 0,
        };
        match *self {
            GpuMode::SearchingRam => 21,
            GpuMode::TransferringData => 43 + scroll_adjust,
            GpuMode::HBlank => 50 - scroll_adjust,
            GpuMode::VBlank => 114,
        }
    }

    fn to_u8(&self) -> u8 {
        match *self {
            GpuMode::HBlank => 0x00,
            GpuMode::VBlank => 0x01,
            GpuMode::SearchingRam => 0x02,
            GpuMode::TransferringData => 0x03,
        }
    }
}

pub struct Gpu {
    pub enabled: bool,
    pub ram: Memory,
    pub sprite_data: Memory,
    control_register: u8,
    stat: GpuStat,
    scroll_y: u8,
    scroll_x: u8,
    window_y: u8,
    window_x: u8,
    pub ly: u8,
    lyc: u8,
    bg_palette: u8,
    palette0: u8,
    palette1: u8,

    cycles: isize,

    pub frame: Frame,
    mode: GpuMode,

    counter: u8,
    tile_base: usize,
    background_base: usize,
}

impl Gpu {
    pub fn new() -> Gpu {
        Gpu {
            enabled: true,
            ram: Memory::new(VRAM_SIZE),
            sprite_data: Memory::new(SPRITE_DATA_SIZE),
            control_register: 0x00,
            stat: GpuStat::new(),
            scroll_y: 0x00,
            scroll_x: 0x00,
            window_y: 0x00,
            window_x: 0x00,
            ly: 0x00,
            lyc: 0x00,
            bg_palette: 0x00,
            palette0: 0x00,
            palette1: 0x00,
            cycles: 0x00,
            frame: Frame::new(),
            mode: GpuMode::HBlank,
            counter: 0,
            tile_base: 0x00,
            background_base: 0xC00,
        }
    }

    pub fn reset(&mut self) {
        *self = Gpu::new();
    }

    pub fn step(&mut self, irq: &mut Irq, cycles: usize) -> Result<(), String> {
        let cycles = cycles as isize;

        if !self.enabled {
            return Ok(());
        }

        self.cycles += cycles;

        match self.mode {
            GpuMode::HBlank => {
                self.ly = self.ly + 0x01;
                if self.cycles > 0xCC {
                    self.cycles = 0x00;
                    if self.ly > 0x90 {
                        self.switch_mode(GpuMode::VBlank, irq);
                    } else {
                        self.switch_mode(GpuMode::SearchingRam, irq);
                    }
                    self.check_coincidence(irq);
                }
            }
            GpuMode::TransferringData => {
                if self.cycles > 0xAC {
                    self.render_scanline();

                    self.cycles = 0x00;
                    self.switch_mode(GpuMode::HBlank, irq);
                }
            }
            GpuMode::SearchingRam => {
                if self.cycles > 0x50 {
                    self.cycles = 0x00;
                    self.switch_mode(GpuMode::TransferringData, irq);
                }
            }
            GpuMode::VBlank => {
                irq.request(Interrupt::Vblank);
                self.ly += 0x01;
                if self.cycles > 0x1C8 {
                    self.cycles = 0x00;
                    if self.ly > 0x99 {
                        self.ly = 0;
                        self.switch_mode(GpuMode::SearchingRam, irq);
                    }
                    self.check_coincidence(irq);
                }
            }
        }

        Ok(())
    }

    fn render_scanline(&mut self) {
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        if line >= 0x90 {
            return;
        }
        self.clear_scanline(line);
        self.render_background(line);
        // self.render_sprites();
    }

    fn clear_scanline(&mut self, line: usize) {
        for i in 0..160 {
            self.frame.pixels[line * 160 + i] = Color::new(0, 0xFF, 0xFF, 0xFF);
        }
    }

    fn render_background(&mut self, line: usize) {
        let background_map_base_address = self.background_base;
        let tile_base_address = self.tile_base;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..160 {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number = self.ram[background_map_base_address +
                                           (bg_map_row * 0x20 + bg_map_col)];
            let t = if tile_base_address == 0x00 {
                raw_tile_number as usize
            } else {
                128 + ((raw_tile_number as i8 as i16) + 128) as usize
            };

            let line_offset = (line % 0x08) << 0x01;
            let tile_data_start = tile_base_address + (t * 0x10) + line_offset;
            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.frame.pixels[line as usize * 160 + i as usize] = c;
        }
    }

    fn render_sprites(&mut self) {
        let background_map_base_address = self.background_base;
        let tile_base_address = self.tile_base;
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..160 {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number = self.ram[background_map_base_address +
                                           (bg_map_row * 0x20 + bg_map_col)];
            let t = if tile_base_address == 0x00 {
                raw_tile_number as usize
            } else {
                128 + ((raw_tile_number as i8 as i16) + 128) as usize
            };

            let line_offset = (line % 0x08) << 0x01;
            let tile_data_start = tile_base_address + (t * 0x10) + line_offset;
            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.frame.pixels[line as usize * 160 + i as usize] = c;
        }
    }

    fn get_background_color_for_byte(&self, b: u8) -> Color {
        let palette_index = match b {
            0x00 => self.bg_palette & 0x0003,
            0x01 => (self.bg_palette & 0x000C) >> 0x02,
            0x02 => (self.bg_palette & 0x0030) >> 0x04,
            0x03 => (self.bg_palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        match addr {
            0x40 => self.control_register,
            0x41 => self.stat.to_u8(self),
            0x42 => self.scroll_y,
            0x43 => self.scroll_x,
            0x4A => self.window_y,
            0x4B => self.window_x,
            0x44 => self.ly,
            0x45 => self.lyc,
            0x47 => self.bg_palette,
            0x48 => self.palette0,
            0x49 => self.palette1,
            _ => panic!("tried to read GPU memory that is not mapped"),
        }
    }

    pub fn write_u8(&mut self, addr: u16, val: u8) {
        match addr {
            0x40 => {
                self.control_register = val;
                self.enabled = self.control_register & 0x80 == 0x80;
                self.tile_base = if self.control_register & 0x10 == 0x10 {
                    0x00
                } else {
                    0x800
                };
                self.background_base = if self.control_register & 0x08 == 0x08 {
                    0x1C00
                } else {
                    0x1800
                };
            }
            0x41 => self.stat = GpuStat::from_u8(val),
            0x42 => self.scroll_y = val,
            0x43 => self.scroll_x = val,
            0x44 => (),
            0x45 => {
                self.lyc = val;
            }
            0x47 => self.bg_palette = val,
            0x48 => self.palette0 = val,
            0x49 => self.palette1 = val,
            0x4A => self.window_y = val,
            0x4B => self.window_x = val,
            _ => panic!("tried to write GPU memory that is not mapped: {:04}", addr),
        }
    }

    fn check_coincidence(&mut self, irq: &mut Irq) {
        // If LY == LYC then set the coincidence flag
        if self.lyc == self.ly && self.stat.coincidence_interrupt_enabled {
            irq.request(Interrupt::Lcd);
        }
    }

    fn switch_mode(&mut self, mode: GpuMode, irq: &mut Irq) {
        self.cycles += mode.cycles(self.scroll_x);
        match mode {
            GpuMode::SearchingRam => {
                if self.stat.OAM_interrupt_enabled {
                    irq.request(Interrupt::OAM);
                }
            }
            _ => (),
        }
        self.mode = mode;
    }
}
<commit_msg>Fix the logo scrolling issue!<commit_after>use gameboy::{Interconnect, Interrupt, Irq, Memory};
use gameboy::frame::{Color, Frame};

const VRAM_SIZE: usize = 0x4000;
const SPRITE_DATA_SIZE: usize = 0xA0;

pub struct GpuStat {
    pub coincidence_interrupt_enabled: bool,
    pub OAM_interrupt_enabled: bool,
    pub VBlank_interrupt_enabled: bool,
    pub HBlank_interrupt_enabled: bool,
}

impl GpuStat {
    pub fn new() -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: false,
            OAM_interrupt_enabled: false,
            VBlank_interrupt_enabled: false,
            HBlank_interrupt_enabled: false,
        }
    }

    pub fn from_u8(b: u8) -> GpuStat {
        GpuStat {
            coincidence_interrupt_enabled: b & 0x40 == 0x40,
            OAM_interrupt_enabled: b & 0x20 == 0x20,
            VBlank_interrupt_enabled: b & 0x10 == 0x10,
            HBlank_interrupt_enabled: b & 0x08 == 0x08,
        }
    }

    pub fn to_u8(&self, gpu: &Gpu) -> u8 {
        (if self.coincidence_interrupt_enabled {
             0x40
         } else {
             0
         }) | (if self.OAM_interrupt_enabled { 0x20 } else { 0 }) |
            (if self.VBlank_interrupt_enabled {
                 0x10
             } else {
                 0
             }) |
            (if self.HBlank_interrupt_enabled {
                 0x08
             } else {
                 0
             }) | (if gpu.ly == gpu.lyc { 0x04 } else { 0 }) | gpu.mode.to_u8()
    }
}

pub enum GpuMode {
    HBlank,
    VBlank,
    SearchingRam,
    TransferringData,
}

impl GpuMode {
    fn cycles(&self, scroll_x: u8) -> isize {
        let scroll_adjust = match scroll_x % 0x08 {
            5...7 => 2,
            1...4 => 1,
            _ => 0,
        };
        match *self {
            GpuMode::SearchingRam => 21,
            GpuMode::TransferringData => 43 + scroll_adjust,
            GpuMode::HBlank => 50 - scroll_adjust,
            GpuMode::VBlank => 114,
        }
    }

    fn to_u8(&self) -> u8 {
        match *self {
            GpuMode::HBlank => 0x00,
            GpuMode::VBlank => 0x01,
            GpuMode::SearchingRam => 0x02,
            GpuMode::TransferringData => 0x03,
        }
    }
}

pub struct Gpu {
    pub enabled: bool,
    pub ram: Memory,
    pub sprite_data: Memory,
    control_register: u8,
    stat: GpuStat,
    scroll_y: u8,
    scroll_x: u8,
    window_y: u8,
    window_x: u8,
    pub ly: u8,
    lyc: u8,
    bg_palette: u8,
    palette0: u8,
    palette1: u8,

    cycles: isize,

    pub frame: Frame,
    mode: GpuMode,

    counter: u8,
    tile_base: usize,
    background_base: usize,
}

impl Gpu {
    pub fn new() -> Gpu {
        Gpu {
            enabled: true,
            ram: Memory::new(VRAM_SIZE),
            sprite_data: Memory::new(SPRITE_DATA_SIZE),
            control_register: 0x00,
            stat: GpuStat::new(),
            scroll_y: 0x00,
            scroll_x: 0x00,
            window_y: 0x00,
            window_x: 0x00,
            ly: 0x00,
            lyc: 0x00,
            bg_palette: 0x00,
            palette0: 0x00,
            palette1: 0x00,
            cycles: 0x00,
            frame: Frame::new(),
            mode: GpuMode::HBlank,
            counter: 0,
            tile_base: 0x00,
            background_base: 0xC00,
        }
    }

    pub fn reset(&mut self) {
        *self = Gpu::new();
    }

    pub fn step(&mut self, irq: &mut Irq, cycles: usize) -> Result<(), String> {
        let cycles = cycles as isize;

        if !self.enabled {
            return Ok(());
        }

        self.cycles += cycles;

        match self.mode {
            GpuMode::HBlank => {
                self.ly = self.ly + 0x01;
                if self.cycles > 0xCC {
                    self.cycles = 0x00;
                    if self.ly > 0x90 {
                        self.switch_mode(GpuMode::VBlank, irq);
                    } else {
                        self.switch_mode(GpuMode::SearchingRam, irq);
                    }
                    self.check_coincidence(irq);
                }
            }
            GpuMode::TransferringData => {
                if self.cycles > 0xAC {
                    self.render_scanline();

                    self.cycles = 0x00;
                    self.switch_mode(GpuMode::HBlank, irq);
                }
            }
            GpuMode::SearchingRam => {
                if self.cycles > 0x50 {
                    self.cycles = 0x00;
                    self.switch_mode(GpuMode::TransferringData, irq);
                }
            }
            GpuMode::VBlank => {
                irq.request(Interrupt::Vblank);
                self.ly += 0x01;
                if self.cycles > 0x1C8 {
                    self.cycles = 0x00;
                    if self.ly > 0x99 {
                        self.ly = 0;
                        self.switch_mode(GpuMode::SearchingRam, irq);
                    }
                    self.check_coincidence(irq);
                }
            }
        }

        Ok(())
    }

    fn render_scanline(&mut self) {
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        if line >= 0x90 {
            return;
        }
        self.clear_scanline(line);
        self.render_background(line);
        // self.render_sprites();
    }

    fn clear_scanline(&mut self, line: usize) {
        for i in 0..160 {
            self.frame.pixels[line * 160 + i] = Color::new(0, 0xFF, 0xFF, 0xFF);
        }
    }

    fn render_background(&mut self, line: usize) {
        let background_map_base_address = self.background_base;
        let tile_base_address = self.tile_base;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..160 {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number =
                self.ram[background_map_base_address + (bg_map_row * 0x20 + bg_map_col)];
            let t = if tile_base_address == 0x00 {
                raw_tile_number as usize
            } else {
                128 + ((raw_tile_number as i8 as i16) + 128) as usize
            };

            let line_offset = (line % 0x08) << 0x01;
            let tile_data_start = tile_base_address + (t * 0x10) + line_offset;
            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.frame.pixels[self.ly as usize * 160 + i as usize] = c;
        }
    }

    fn render_sprites(&mut self) {
        let background_map_base_address = self.background_base;
        let tile_base_address = self.tile_base;
        let line = self.ly.wrapping_add(self.scroll_y) as usize;
        let bg_map_row = (line / 0x08) as usize;
        for i in 0..160 {
            let x = (i as u8).wrapping_add(self.scroll_x);
            let bg_map_col = (x / 8) as usize;
            let raw_tile_number =
                self.ram[background_map_base_address + (bg_map_row * 0x20 + bg_map_col)];
            let t = if tile_base_address == 0x00 {
                raw_tile_number as usize
            } else {
                128 + ((raw_tile_number as i8 as i16) + 128) as usize
            };

            let line_offset = (line % 0x08) << 0x01;
            let tile_data_start = tile_base_address + (t * 0x10) + line_offset;
            let x_shift = (x % 8).wrapping_sub(0x07).wrapping_mul(0xFF);
            let tile_data1 = (self.ram[tile_data_start] >> x_shift) & 0x01;
            let tile_data2 = (self.ram[tile_data_start + 0x01] >> x_shift) & 0x01;
            let total_row_data = (tile_data2 << 1) | tile_data1;
            let color_value = total_row_data;
            let c = self.get_background_color_for_byte(color_value as u8);
            self.frame.pixels[line as usize * 160 + i as usize] = c;
        }
    }

    fn get_background_color_for_byte(&self, b: u8) -> Color {
        let palette_index = match b {
            0x00 => self.bg_palette & 0x0003,
            0x01 => (self.bg_palette & 0x000C) >> 0x02,
            0x02 => (self.bg_palette & 0x0030) >> 0x04,
            0x03 => (self.bg_palette & 0x00C0) >> 0x06,
            _ => panic!("err: invalid pixel color value found"),
        };

        Color::from_dmg_byte(palette_index as u8)
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        match addr {
            0x40 => self.control_register,
            0x41 => self.stat.to_u8(self),
            0x42 => self.scroll_y,
            0x43 => self.scroll_x,
            0x4A => self.window_y,
            0x4B => self.window_x,
            0x44 => self.ly,
            0x45 => self.lyc,
            0x47 => self.bg_palette,
            0x48 => self.palette0,
            0x49 => self.palette1,
            _ => panic!("tried to read GPU memory that is not mapped"),
        }
    }

    pub fn write_u8(&mut self, addr: u16, val: u8) {
        match addr {
            0x40 => {
                self.control_register = val;
                self.enabled = self.control_register & 0x80 == 0x80;
                self.tile_base = if self.control_register & 0x10 == 0x10 {
                    0x00
                } else {
                    0x800
                };
                self.background_base = if self.control_register & 0x08 == 0x08 {
                    0x1C00
                } else {
                    0x1800
                };
            }
            0x41 => self.stat = GpuStat::from_u8(val),
            0x42 => self.scroll_y = val,
            0x43 => self.scroll_x = val,
            0x44 => (),
            0x45 => {
                self.lyc = val;
            }
            0x47 => self.bg_palette = val,
            0x48 => self.palette0 = val,
            0x49 => self.palette1 = val,
            0x4A => self.window_y = val,
            0x4B => self.window_x = val,
            _ => panic!("tried to write GPU memory that is not mapped: {:04}", addr),
        }
    }

    fn check_coincidence(&mut self, irq: &mut Irq) {
        // If LY == LYC then set the coincidence flag
        if self.lyc == self.ly && self.stat.coincidence_interrupt_enabled {
            irq.request(Interrupt::Lcd);
        }
    }

    fn switch_mode(&mut self, mode: GpuMode, irq: &mut Irq) {
        self.cycles += mode.cycles(self.scroll_x);
        match mode {
            GpuMode::SearchingRam => {
                if self.stat.OAM_interrupt_enabled {
                    irq.request(Interrupt::OAM);
                }
            }
            _ => (),
        }
        self.mode = mode;
    }
}
<|endoftext|>


<commit_before>
use buffer::{ImageBuffer, Pixel, ArrayLike};
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use std::old_io;
use std::old_io::IoResult;
use std::num::Int;

use color::{Rgb, Rgba};
use utils::lzw;
use utils::bitstream::LsbWriter;

use super::{Extension, Block, DisposalMethod};

#[derive(Debug, Copy)]
#[allow(unused_qualifications)]
/// The color mode the encoder will use to encode the image.
pub enum ColorMode {
    /// Image will be encoded in multiple frames if more than 256 colors are present
    TrueColor,
    /// Number of colors will be reduced
    Indexed(u8),
}
use self::ColorMode::{TrueColor, Indexed};

/// A GIF encoder.
///
/// Encodes the image either in true color using indexed colors.
/// If the mode is set to TrueColor the image is split into multiple frames
/// when the number of colors including transparent color exceeds 256.
/// Pixels with an alpha value != 1.0 will be set to alpha = 0.
pub struct Encoder<Image> {
    image: Image,
    bg_color: Option<Rgb<u8>>,
    color_mode: ColorMode
}

const TRANSPARENT: Rgba<u8> = Rgba([0, 0, 0, 0]);

impl<Container> Encoder<ImageBuffer<Rgba<u8>, Container>>
where Container: ArrayLike<u8> {
    /// Creates a new GIF encoder
    pub fn new(image: ImageBuffer<Rgba<u8>, Container>,
               bg_color: Option<Rgb<u8>>,
               color_mode: ColorMode,
              ) -> Encoder<ImageBuffer<Rgba<u8>, Container>> {
        Encoder {
            image: image,
            bg_color: bg_color,
            color_mode: color_mode
        }
    }

    /// Encodes the image
    pub fn encode<W: Writer>(&mut self, w: &mut W) -> IoResult<()> {
        // Header
        try!(w.write_all(b"GIF89a"));
        // Logical screen descriptor
        let height = self.image.height();
        let width = self.image.width();
        if width > <u16 as Int>::max_value() as u32 ||
           height > <u16 as Int>::max_value() as u32 {
            return Err(old_io::IoError{
                kind: old_io::InvalidInput,
                desc: "Image dimensions are to large for the gif format.",
                detail: None
            })
        }
        try!(w.write_le_u16(width as u16));
        try!(w.write_le_u16(height as u16));

        let (hist, transparent) = self.histogram();
        let num_colors = hist.len();

        try!(self.write_global_table(w, &hist[]));

        match self.color_mode {
            // Global color table has been written, just write the image data now
            TrueColor if num_colors <= 256 => {
                // NOTE: copy paste any changes to case: Indexed(n) if n as usize <= num_colors
                try!(self.write_control_ext(w, 0, transparent));
                try!(self.write_descriptor(w, None));
                try!(self.write_image_simple(w, &hist[], transparent));
            },
            TrueColor => {
                try!(self.write_true_color(w, hist, transparent));
            },
            Indexed(n) if n as usize <= num_colors => {
                // NOTE: copy paste from case: TrueColor if num_colors <= 256
                try!(self.write_control_ext(w, 0, transparent));
                try!(self.write_descriptor(w, None));
                try!(self.write_image_simple(w, &hist[], transparent));
            }
            Indexed(_) => {
                panic!("Indexed(n) unimplemented")
            }

        }
        w.write_u8(Block::Trailer as u8)
    }

    /// Returns a histogram of the colors in the image
    fn histogram(&self) -> (Vec<(Rgba<u8>, usize)>, Option<usize>){
        let mut hist: HashMap<_, usize> = HashMap::new();
        // Add bg_color to table
        if let Some(bg_color) = self.bg_color {
            let _ = hist.insert(bg_color.to_rgba(), 0);
        }
        for p in self.image.pixels() {
            let mut p = *p;
            // Replace transparent pixels with black, alpha = 0
            if let Some(bg_color) = self.bg_color {
                p.blend(&bg_color.to_rgba())
            } else if p[3] != 255 {
                p = TRANSPARENT;
            }
            match hist.entry(p) {
                Occupied(mut entry) => {
                    let val = entry.get_mut();
                    *val = *val + 1;
                },
                Vacant(entry) => {
                    entry.insert(1);
                }
            }
        }
        let mut colors: Vec<(Rgba<u8>, usize)> = hist.into_iter().collect();
        // Sort color map with decreasing color freqs
        colors.sort_by(|a, b| b.1.cmp(&a.1));
        // Add bg_color to table
        let transparent = colors.iter().position(|x| x.0 == TRANSPARENT);
        (colors, transparent)
    }

    /// Write the global color table and the corresponding flags
    fn write_global_table<W: Writer>(&mut self,
                                     w: &mut W,
                                     hist: &[(Rgba<u8>, usize)]
                                    ) -> IoResult<()>
    {
        let num_colors = hist.len();
        let mut flags = 0;
        flags |= 1 << 4; // sort flag
        let n = flag_n(num_colors);
        flags |= n << 4; // 2^(n+1) colors
        let (global_table, bg_index) = if num_colors <= 256 {
            flags |= 1 << 7; // glocal table flag
            flags |= n;
            let mut bg_index = 0;
            let mut table = Vec::with_capacity(3*(2 << n));
            for (i, &(color, _)) in hist.iter().enumerate() {
                let channels = &color.channels()[..3];
                if let Some(bg_color) = self.bg_color {
                    if channels == bg_color.channels() {
                        bg_index = i;
                    }
                }
                table.push_all(channels)
            }
            // Waste some space as of gif spec
            for _ in 0..((2 << n) - num_colors) {
                table.push_all(&[0, 0, 0]);
            }
            (Some(table), bg_index as u8)
        } else if let Some(bg_color) = self.bg_color {
            flags |= 1 << 7; // glocal table flag
            let mut table = Vec::with_capacity(6);
            table.push_all(&bg_color.channels()[..3]);
            table.push_all(&[0, 0, 0]);
            (Some(table), 0)
        } else {
            (None, 0)
        };
        try!(w.write_u8(flags));
        try!(w.write_u8(bg_index));
        try!(w.write_u8(0)); // aspect ration, disregard
        if let Some(global_table) = global_table {
            try!(w.write_all(&global_table[]));
        }
        Ok(())
    }

    /// Writes the graphics control extension
    fn write_control_ext<W: Writer>(&mut self,
                                    w: &mut W,
                                    delay: u16,
                                    transparent: Option<usize>
                                   ) -> IoResult<()>
    {
        try!(w.write_u8(Block::Extension as u8));
        try!(w.write_u8(Extension::Control as u8));
        try!(w.write_u8(4)); // size
        let mut field = 0;
        field |= (DisposalMethod::None as u8) << 2;
        let idx = if let Some(idx) = transparent {
            field |=  1;
            idx as u8
        } else {
            0
        };
        try!(w.write_u8(field));
        try!(w.write_le_u16(delay));
        try!(w.write_u8(idx));
        w.write_u8(0)
    }

    /// Writes the image descriptor
    fn write_descriptor<W: Writer>(&mut self,
                                   w: &mut W,
                                   table_len: Option<usize>
                                  ) -> IoResult<()>
    {
        try!(w.write_u8(Block::Image as u8));
        try!(w.write_le_u16(0)); // left
        try!(w.write_le_u16(0)); // top
        let height = self.image.height();
        let width = self.image.width();
        try!(w.write_le_u16(width as u16));
        try!(w.write_le_u16(height as u16));
        if let Some(len) = table_len {
            w.write_u8(flag_n(len) | 0x80)
        } else {
            w.write_u8(0)
        }
    }

    /// Writes and compresses the indexed data
    fn write_indices<W: Writer>(&mut self,
                                w: &mut W,
                                indices: &[u8],
                               ) -> IoResult<()>
    {
        let code_size = match flag_n(indices.len()) + 1 {
            1 => 2,
            n => n
        };
        let mut encoded_data = Vec::new();
        try!(lzw::encode(indices, LsbWriter::new(&mut encoded_data), code_size));
        try!(w.write_u8(code_size));
        for chunk in encoded_data.chunks(255) {
            try!(w.write_u8((chunk.len()) as u8));
            try!(w.write_all(chunk));
        }
        w.write_u8(0) // block terminator

    }

    /// Writes the image to the file assuming that every pixel is in the color table
    /// If not, the index of the transparent pixel is written
    fn write_image_simple<W: Writer>(&mut self,
                                     w: &mut W,
                                     hist: &[(Rgba<u8>, usize)],
                                     transparent: Option<usize>,
                                    ) -> IoResult<()>
    {
        let t_idx = match transparent { Some(i) => i as u8, None => 0 };
        let data: Vec<u8> = self.image.pixels().map(|p| {
            if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
                idx as u8
            } else {
                t_idx
            }
        }).collect();
        self.write_indices(w, &data[])

    }

    /// Writes the image as a true color image by splitting the colors
    /// over several frames
    fn write_true_color<W: Writer>(&mut self,
                                   w: &mut W,
                                   hist: Vec<(Rgba<u8>, usize)>,
                                   transparent: Option<usize>
                                  ) -> IoResult<()>
    {
        try!(self.write_nab(w, 1));
        let mut hist = hist;
        // Remove transparent idx
        if let Some(transparent) = transparent {
            let _ = hist.swap_remove(transparent);
        }
        let transparent = Some(0);
        // Calculating the indices is expensive so just do it once
        let indices: Vec<u32> = self.image.pixels().map(|p| {
            if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
                idx as u32
            } else {
                0
            }
        }).collect();
        let mut chunk_indices = Vec::with_capacity(self.image.as_slice().len());
        // starting from 1 since we inject the transparent color in every frame
        for (j, chunk) in hist.chunks(255).enumerate() {
            // Now we can reuse the indices for every subimage
            // if the substract j*255 from each an replace it with 0 if it gets negative
            chunk_indices.clear();
            for &idx in indices.iter() {
                let i: i64 = idx as i64 - j as i64*255;
                chunk_indices.push(if i < 0 || i > 255 { 0 } else { i as u8 + 1 })
            }
            // Write local color table
            let num_local_colors = chunk.len() + 1;
            let n = flag_n(num_local_colors);
            try!(self.write_control_ext(w, 0, transparent));
            try!(self.write_descriptor(w, Some(num_local_colors)));
            // Inject transparent color
            try!(w.write_all(&TRANSPARENT.channels()[..3]));
            for &(color, _) in chunk.iter() {
                try!(w.write_all(&color.channels()[..3]));
            }
            // waste some space as of gif spec
            for _ in 0..((2 << n) - num_local_colors) {
                try!(w.write_all(&[0, 0, 0]));
            }
            // write indices
            try!(self.write_indices(w, &chunk_indices[]));
        }
        Ok(())
    }

    /// Writes the netscape application block to set the number `n` of repetitions
    fn write_nab<W: Writer>(&mut self, w: &mut W, n: u16) -> IoResult<()> {
        try!(w.write_u8(Block::Extension as u8));
        try!(w.write_u8(Extension::Application as u8));
        try!(w.write_u8(0x0B)); // size
        try!(w.write_all(b"NETSCAPE2.0"));
        try!(w.write_u8(0x03)); // sub-block size
        try!(w.write_u8(0x01)); // sub-block id
        try!(w.write_le_u16(n));
        w.write_u8(0) // terminator
    }
}

// Color table len converted to flag bits
fn flag_n(size: usize) -> u8 {
    match size {
        0  ...2   => 0,
        3  ...4   => 1,
        5  ...8   => 2,
        7  ...16  => 3,
        17 ...32  => 4,
        33 ...64  => 5,
        65 ...128 => 6,
        129...256 => 7,
        _ => 7
    }
}<commit_msg>Do not write NAB since we don’t want to repeat and 0 would mean infinity!<commit_after>
use buffer::{ImageBuffer, Pixel, ArrayLike};
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use std::old_io;
use std::old_io::IoResult;
use std::num::Int;

use color::{Rgb, Rgba};
use utils::lzw;
use utils::bitstream::LsbWriter;

use super::{Extension, Block, DisposalMethod};

#[derive(Debug, Copy)]
#[allow(unused_qualifications)]
/// The color mode the encoder will use to encode the image.
pub enum ColorMode {
    /// Image will be encoded in multiple frames if more than 256 colors are present
    TrueColor,
    /// Number of colors will be reduced
    Indexed(u8),
}
use self::ColorMode::{TrueColor, Indexed};

/// A GIF encoder.
///
/// Encodes the image either in true color using indexed colors.
/// If the mode is set to TrueColor the image is split into multiple frames
/// when the number of colors including transparent color exceeds 256.
/// Pixels with an alpha value != 1.0 will be set to alpha = 0.
pub struct Encoder<Image> {
    image: Image,
    bg_color: Option<Rgb<u8>>,
    color_mode: ColorMode
}

const TRANSPARENT: Rgba<u8> = Rgba([0, 0, 0, 0]);

impl<Container> Encoder<ImageBuffer<Rgba<u8>, Container>>
where Container: ArrayLike<u8> {
    /// Creates a new GIF encoder
    pub fn new(image: ImageBuffer<Rgba<u8>, Container>,
               bg_color: Option<Rgb<u8>>,
               color_mode: ColorMode,
              ) -> Encoder<ImageBuffer<Rgba<u8>, Container>> {
        Encoder {
            image: image,
            bg_color: bg_color,
            color_mode: color_mode
        }
    }

    /// Encodes the image
    pub fn encode<W: Writer>(&mut self, w: &mut W) -> IoResult<()> {
        // Header
        try!(w.write_all(b"GIF89a"));
        // Logical screen descriptor
        let height = self.image.height();
        let width = self.image.width();
        if width > <u16 as Int>::max_value() as u32 ||
           height > <u16 as Int>::max_value() as u32 {
            return Err(old_io::IoError{
                kind: old_io::InvalidInput,
                desc: "Image dimensions are to large for the gif format.",
                detail: None
            })
        }
        try!(w.write_le_u16(width as u16));
        try!(w.write_le_u16(height as u16));

        let (hist, transparent) = self.histogram();
        let num_colors = hist.len();

        try!(self.write_global_table(w, &hist[]));

        match self.color_mode {
            // Global color table has been written, just write the image data now
            TrueColor if num_colors <= 256 => {
                // NOTE: copy paste any changes to case: Indexed(n) if n as usize <= num_colors
                try!(self.write_control_ext(w, 0, transparent));
                try!(self.write_descriptor(w, None));
                try!(self.write_image_simple(w, &hist[], transparent));
            },
            TrueColor => {
                try!(self.write_true_color(w, hist, transparent));
            },
            Indexed(n) if n as usize <= num_colors => {
                // NOTE: copy paste from case: TrueColor if num_colors <= 256
                try!(self.write_control_ext(w, 0, transparent));
                try!(self.write_descriptor(w, None));
                try!(self.write_image_simple(w, &hist[], transparent));
            }
            Indexed(_) => {
                panic!("Indexed(n) unimplemented")
            }

        }
        w.write_u8(Block::Trailer as u8)
    }

    /// Returns a histogram of the colors in the image
    fn histogram(&self) -> (Vec<(Rgba<u8>, usize)>, Option<usize>){
        let mut hist: HashMap<_, usize> = HashMap::new();
        // Add bg_color to table
        if let Some(bg_color) = self.bg_color {
            let _ = hist.insert(bg_color.to_rgba(), 0);
        }
        for p in self.image.pixels() {
            let mut p = *p;
            // Replace transparent pixels with black, alpha = 0
            if let Some(bg_color) = self.bg_color {
                p.blend(&bg_color.to_rgba())
            // do not use 255 as a limit there could be rounding error etc.
            // thus take 250 to give it some security margin.
            } else if p[3] < 250 {
                p = TRANSPARENT;
            }
            match hist.entry(p) {
                Occupied(mut entry) => {
                    let val = entry.get_mut();
                    *val = *val + 1;
                },
                Vacant(entry) => {
                    entry.insert(1);
                }
            }
        }
        let mut colors: Vec<(Rgba<u8>, usize)> = hist.into_iter().collect();
        // Sort color map with decreasing color freqs
        colors.sort_by(|a, b| b.1.cmp(&a.1));
        // Add bg_color to table
        let transparent = colors.iter().position(|x| x.0 == TRANSPARENT);
        (colors, transparent)
    }

    /// Write the global color table and the corresponding flags
    fn write_global_table<W: Writer>(&mut self,
                                     w: &mut W,
                                     hist: &[(Rgba<u8>, usize)]
                                    ) -> IoResult<()>
    {
        let num_colors = hist.len();
        let mut flags = 0;
        flags |= 1 << 4; // sort flag
        let n = flag_n(num_colors);
        flags |= n << 4; // 2^(n+1) colors
        let (global_table, bg_index) = if num_colors <= 256 {
            flags |= 1 << 7; // glocal table flag
            flags |= n;
            let mut bg_index = 0;
            let mut table = Vec::with_capacity(3*(2 << n));
            for (i, &(color, _)) in hist.iter().enumerate() {
                let channels = &color.channels()[..3];
                if let Some(bg_color) = self.bg_color {
                    if channels == bg_color.channels() {
                        bg_index = i;
                    }
                }
                table.push_all(channels)
            }
            // Waste some space as of gif spec
            for _ in 0..((2 << n) - num_colors) {
                table.push_all(&[0, 0, 0]);
            }
            (Some(table), bg_index as u8)
        } else if let Some(bg_color) = self.bg_color {
            flags |= 1 << 7; // glocal table flag
            let mut table = Vec::with_capacity(6);
            table.push_all(&bg_color.channels()[..3]);
            table.push_all(&[0, 0, 0]);
            (Some(table), 0)
        } else {
            (None, 0)
        };
        try!(w.write_u8(flags));
        try!(w.write_u8(bg_index));
        try!(w.write_u8(0)); // aspect ration, disregard
        if let Some(global_table) = global_table {
            try!(w.write_all(&global_table[]));
        }
        Ok(())
    }

    /// Writes the graphics control extension
    fn write_control_ext<W: Writer>(&mut self,
                                    w: &mut W,
                                    delay: u16,
                                    transparent: Option<usize>
                                   ) -> IoResult<()>
    {
        try!(w.write_u8(Block::Extension as u8));
        try!(w.write_u8(Extension::Control as u8));
        try!(w.write_u8(4)); // size
        let mut field = 0;
        field |= (DisposalMethod::None as u8) << 2;
        let idx = if let Some(idx) = transparent {
            field |=  1;
            idx as u8
        } else {
            0
        };
        try!(w.write_u8(field));
        try!(w.write_le_u16(delay));
        try!(w.write_u8(idx));
        w.write_u8(0)
    }

    /// Writes the image descriptor
    fn write_descriptor<W: Writer>(&mut self,
                                   w: &mut W,
                                   table_len: Option<usize>
                                  ) -> IoResult<()>
    {
        try!(w.write_u8(Block::Image as u8));
        try!(w.write_le_u16(0)); // left
        try!(w.write_le_u16(0)); // top
        let height = self.image.height();
        let width = self.image.width();
        try!(w.write_le_u16(width as u16));
        try!(w.write_le_u16(height as u16));
        if let Some(len) = table_len {
            w.write_u8(flag_n(len) | 0x80)
        } else {
            w.write_u8(0)
        }
    }

    /// Writes and compresses the indexed data
    fn write_indices<W: Writer>(&mut self,
                                w: &mut W,
                                indices: &[u8],
                               ) -> IoResult<()>
    {
        let code_size = match flag_n(indices.len()) + 1 {
            1 => 2,
            n => n
        };
        let mut encoded_data = Vec::new();
        try!(lzw::encode(indices, LsbWriter::new(&mut encoded_data), code_size));
        try!(w.write_u8(code_size));
        for chunk in encoded_data.chunks(255) {
            try!(w.write_u8((chunk.len()) as u8));
            try!(w.write_all(chunk));
        }
        w.write_u8(0) // block terminator

    }

    /// Writes the image to the file assuming that every pixel is in the color table
    /// If not, the index of the transparent pixel is written
    fn write_image_simple<W: Writer>(&mut self,
                                     w: &mut W,
                                     hist: &[(Rgba<u8>, usize)],
                                     transparent: Option<usize>,
                                    ) -> IoResult<()>
    {
        let t_idx = match transparent { Some(i) => i as u8, None => 0 };
        let data: Vec<u8> = self.image.pixels().map(|p| {
            if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
                idx as u8
            } else {
                t_idx
            }
        }).collect();
        self.write_indices(w, &data[])

    }

    /// Writes the image as a true color image by splitting the colors
    /// over several frames
    fn write_true_color<W: Writer>(&mut self,
                                   w: &mut W,
                                   hist: Vec<(Rgba<u8>, usize)>,
                                   transparent: Option<usize>
                                  ) -> IoResult<()>
    {
        let mut hist = hist;
        // Remove transparent idx
        if let Some(transparent) = transparent {
            let _ = hist.swap_remove(transparent);
        }
        let transparent = Some(0);
        // Calculating the indices is expensive so just do it once
        let indices: Vec<u32> = self.image.pixels().map(|p| {
            if let Some(idx) = hist.iter().position(|x| x.0 == *p) {
                idx as u32
            } else {
                0
            }
        }).collect();
        let mut chunk_indices = Vec::with_capacity(self.image.as_slice().len());
        // starting from 1 since we inject the transparent color in every frame
        for (j, chunk) in hist.chunks(255).enumerate() {
            // Now we can reuse the indices for every subimage
            // if the substract j*255 from each an replace it with 0 if it gets negative
            chunk_indices.clear();
            for &idx in indices.iter() {
                let i: i64 = idx as i64 - j as i64*255;
                chunk_indices.push(if i < 0 || i > 255 { 0 } else { i as u8 + 1 })
            }
            // Write local color table
            let num_local_colors = chunk.len() + 1;
            let n = flag_n(num_local_colors);
            try!(self.write_control_ext(w, 0, transparent));
            try!(self.write_descriptor(w, Some(num_local_colors)));
            // Inject transparent color
            try!(w.write_all(&TRANSPARENT.channels()[..3]));
            for &(color, _) in chunk.iter() {
                try!(w.write_all(&color.channels()[..3]));
            }
            // waste some space as of gif spec
            for _ in 0..((2 << n) - num_local_colors) {
                try!(w.write_all(&[0, 0, 0]));
            }
            // write indices
            try!(self.write_indices(w, &chunk_indices[]));
        }
        Ok(())
    }

    /// Writes the netscape application block to set the number `n` of repetitions
    #[allow(dead_code)]
    fn write_nab<W: Writer>(&mut self, w: &mut W, n: u16) -> IoResult<()> {
        try!(w.write_u8(Block::Extension as u8));
        try!(w.write_u8(Extension::Application as u8));
        try!(w.write_u8(0x0B)); // size
        try!(w.write_all(b"NETSCAPE2.0"));
        try!(w.write_u8(0x03)); // sub-block size
        try!(w.write_u8(0x01)); // sub-block id
        try!(w.write_le_u16(n));
        w.write_u8(0) // terminator
    }
}

// Color table len converted to flag bits
fn flag_n(size: usize) -> u8 {
    match size {
        0  ...2   => 0,
        3  ...4   => 1,
        5  ...8   => 2,
        7  ...16  => 3,
        17 ...32  => 4,
        33 ...64  => 5,
        65 ...128 => 6,
        129...256 => 7,
        _ => 7
    }
}<|endoftext|>


<commit_before>extern crate xml;

use std::fs::File;
use std::io::BufReader;
use std::str::FromStr;
use std::num;

use xml::reader::{EventReader, XmlEvent};
use xml::name::OwnedName;
use xml::attribute::OwnedAttribute;

pub type Coordinate = f64;

#[derive(Debug, PartialEq)]
pub struct Bounds {
    pub minlat: Coordinate,
    pub minlon: Coordinate,
    pub maxlat: Coordinate,
    pub maxlon: Coordinate,
}

#[derive(Debug, PartialEq)]
pub struct Node {
    pub id: i64,
    pub lat: Coordinate,
    pub lon: Coordinate,
}

#[derive(Debug)]
pub struct OSM {
    pub bounds: Option<Bounds>,
    pub nodes: Vec<Node>
}

impl OSM {
    fn empty() -> OSM {
        OSM {
            bounds: None,
            nodes: Vec::new()
        }
    }

    pub fn parse(file: File) -> Option<OSM> {
        let mut osm = OSM::empty();

        let reader = BufReader::new(file);
        let parser = EventReader::new(reader);

        for element in parser {
            match element {
                Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                    handle_element(&mut osm, name, attributes);
                },
                _ => { }
            }
        }

        Some(osm)
    }
}

#[derive(Debug)]
enum ErrorKind {
    BoundsMissing(AttributeError),
    IdMissing(AttributeError),
    CoordinateMissing(AttributeError),
    UnknownElement
}

#[derive(Debug)]
enum AttributeError {
    ParseFloat(num::ParseFloatError),
    ParseInt(num::ParseIntError),
    Missing
}

impl From<num::ParseFloatError> for AttributeError {
    fn from(err: num::ParseFloatError) -> AttributeError {
        AttributeError::ParseFloat(err)
    }
}

impl From<num::ParseIntError> for AttributeError {
    fn from(err: num::ParseIntError) -> AttributeError {
        AttributeError::ParseInt(err)
    }
}

#[allow(dead_code)]
enum ElementType {
    Bounds,
    Node,
    Way,
    Relation
}

impl FromStr for ElementType {
    type Err = ErrorKind;

    fn from_str(s: &str) -> Result<ElementType, ErrorKind> {
        let downcased = s.to_lowercase();

        if downcased == "bounds" {
            return Ok(ElementType::Bounds);
        }

        if downcased == "node" {
            return Ok(ElementType::Node);
        }

        if downcased == "way" {
            return Ok(ElementType::Way);
        }

        if downcased == "relation" {
            return Ok(ElementType::Relation);
        }

        Err(ErrorKind::UnknownElement)
    }
}

fn handle_element(osm: &mut OSM, name: OwnedName, attrs: Vec<OwnedAttribute>) {
    match ElementType::from_str(&name.local_name) {
        Ok(element) => {
            match element {
                ElementType::Bounds => set_bounds(osm, &attrs),
                ElementType::Node => insert_node(osm, &attrs),
                _ => ()
            }
        },
        Err(_) => ()
    }
}

fn set_bounds(osm: &mut OSM, attrs: &Vec<OwnedAttribute>) {
    match parse_bounds(&attrs) {
        Ok(bounds) => osm.bounds = Some(bounds),
        Err(_) => osm.bounds = None
    }
}

fn insert_node(osm: &mut OSM, attrs: &Vec<OwnedAttribute>) {
    match parse_node(&attrs) {
        Ok(node) => osm.nodes.push(node),
        Err(_) => ()
    }
}

fn parse_bounds(attrs: &Vec<OwnedAttribute>) -> Result<Bounds, ErrorKind> {
    let minlat = try!(find_attribute("minlat", attrs).map_err(ErrorKind::BoundsMissing));
    let minlon = try!(find_attribute("minlon", attrs).map_err(ErrorKind::BoundsMissing));
    let maxlat = try!(find_attribute("maxlat", attrs).map_err(ErrorKind::BoundsMissing));
    let maxlon = try!(find_attribute("maxlon", attrs).map_err(ErrorKind::BoundsMissing));

    Ok(Bounds {
        minlat: minlat,
        minlon: minlon,
        maxlat: maxlat,
        maxlon: maxlon,
    })
}

fn parse_node(attrs: &Vec<OwnedAttribute>) -> Result<Node, ErrorKind> {
    let id = try!(find_attribute("id", attrs).map_err(ErrorKind::IdMissing));
    let lon = try!(find_attribute("lon", attrs).map_err(ErrorKind::CoordinateMissing));
    let lat = try!(find_attribute("lat", attrs).map_err(ErrorKind::CoordinateMissing));

    Ok(Node {
        id: id,
        lon: lon,
        lat: lat
    })
}

fn find_attribute<T>( name: &str, attrs: &Vec<OwnedAttribute>) -> Result<T, AttributeError>
    where AttributeError: std::convert::From<<T as std::str::FromStr>::Err>,
          T: FromStr
{
    let attr = attrs.iter().find(|attr| attr.name.local_name == name);
    match attr {
        Some(a) => {
            let val = try!(a.value.parse::<T>());
            Ok(val)
        },
        None => Err(AttributeError::Missing)
    }
}
<commit_msg>Introduce Tag type<commit_after>extern crate xml;

use std::fs::File;
use std::io::BufReader;
use std::str::FromStr;
use std::num;

use xml::reader::{EventReader, XmlEvent};
use xml::name::OwnedName;
use xml::attribute::OwnedAttribute;

pub type Coordinate = f64;

#[derive(Debug, PartialEq)]
pub struct Tag {
    pub key: String,
    pub val: String
}

#[derive(Debug, PartialEq)]
pub struct Bounds {
    pub minlat: Coordinate,
    pub minlon: Coordinate,
    pub maxlat: Coordinate,
    pub maxlon: Coordinate,
}

#[derive(Debug, PartialEq)]
pub struct Node {
    pub id: i64,
    pub lat: Coordinate,
    pub lon: Coordinate,
    pub tags: Vec<Tag>
}

#[derive(Debug)]
pub struct OSM {
    pub bounds: Option<Bounds>,
    pub nodes: Vec<Node>
}

impl OSM {
    fn empty() -> OSM {
        OSM {
            bounds: None,
            nodes: Vec::new(),
        }
    }

    pub fn parse(file: File) -> Option<OSM> {
        let mut osm = OSM::empty();

        let reader = BufReader::new(file);
        let parser = EventReader::new(reader);

        for element in parser {
            match element {
                Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                    handle_element(&mut osm, name, attributes);
                },
                _ => { }
            }
        }

        Some(osm)
    }
}

#[derive(Debug)]
enum ErrorKind {
    BoundsMissing(AttributeError),
    IdMissing(AttributeError),
    CoordinateMissing(AttributeError),
    UnknownElement
}

#[derive(Debug)]
enum AttributeError {
    ParseFloat(num::ParseFloatError),
    ParseInt(num::ParseIntError),
    Missing
}

impl From<num::ParseFloatError> for AttributeError {
    fn from(err: num::ParseFloatError) -> AttributeError {
        AttributeError::ParseFloat(err)
    }
}

impl From<num::ParseIntError> for AttributeError {
    fn from(err: num::ParseIntError) -> AttributeError {
        AttributeError::ParseInt(err)
    }
}

#[allow(dead_code)]
enum ElementType {
    Bounds,
    Node,
    Way,
    Relation
}

impl FromStr for ElementType {
    type Err = ErrorKind;

    fn from_str(s: &str) -> Result<ElementType, ErrorKind> {
        let downcased = s.to_lowercase();

        if downcased == "bounds" {
            return Ok(ElementType::Bounds);
        }

        if downcased == "node" {
            return Ok(ElementType::Node);
        }

        if downcased == "way" {
            return Ok(ElementType::Way);
        }

        if downcased == "relation" {
            return Ok(ElementType::Relation);
        }

        Err(ErrorKind::UnknownElement)
    }
}

fn handle_element(osm: &mut OSM, name: OwnedName, attrs: Vec<OwnedAttribute>) {
    match ElementType::from_str(&name.local_name) {
        Ok(element) => {
            match element {
                ElementType::Bounds => set_bounds(osm, &attrs),
                ElementType::Node => insert_node(osm, &attrs),
                _ => ()
            }
        },
        Err(_) => ()
    }
}

fn set_bounds(osm: &mut OSM, attrs: &Vec<OwnedAttribute>) {
    match parse_bounds(&attrs) {
        Ok(bounds) => osm.bounds = Some(bounds),
        Err(_) => osm.bounds = None
    }
}

fn insert_node(osm: &mut OSM, attrs: &Vec<OwnedAttribute>) {
    match parse_node(&attrs) {
        Ok(node) => osm.nodes.push(node),
        Err(_) => ()
    }
}

fn parse_bounds(attrs: &Vec<OwnedAttribute>) -> Result<Bounds, ErrorKind> {
    let minlat = try!(find_attribute("minlat", attrs).map_err(ErrorKind::BoundsMissing));
    let minlon = try!(find_attribute("minlon", attrs).map_err(ErrorKind::BoundsMissing));
    let maxlat = try!(find_attribute("maxlat", attrs).map_err(ErrorKind::BoundsMissing));
    let maxlon = try!(find_attribute("maxlon", attrs).map_err(ErrorKind::BoundsMissing));

    Ok(Bounds {
        minlat: minlat,
        minlon: minlon,
        maxlat: maxlat,
        maxlon: maxlon,
    })
}

fn parse_node(attrs: &Vec<OwnedAttribute>) -> Result<Node, ErrorKind> {
    let id = try!(find_attribute("id", attrs).map_err(ErrorKind::IdMissing));
    let lon = try!(find_attribute("lon", attrs).map_err(ErrorKind::CoordinateMissing));
    let lat = try!(find_attribute("lat", attrs).map_err(ErrorKind::CoordinateMissing));

    Ok(Node {
        id: id,
        lon: lon,
        lat: lat,
        tags: Vec::new()
    })
}

fn find_attribute<T>( name: &str, attrs: &Vec<OwnedAttribute>) -> Result<T, AttributeError>
    where AttributeError: std::convert::From<<T as std::str::FromStr>::Err>,
          T: FromStr
{
    let attr = attrs.iter().find(|attr| attr.name.local_name == name);
    match attr {
        Some(a) => {
            let val = try!(a.value.parse::<T>());
            Ok(val)
        },
        None => Err(AttributeError::Missing)
    }
}
<|endoftext|>