<commit_before>//! Contains an implementation of splay trees where each node has a key/value
//! pair to be used in maps and sets. The only requirement is that the key must
//! implement the Ord trait.
//!
//! # Example
//!
//! ```rust
//! use splay::SplayMap;
//!
//! let mut map = SplayMap::new();
//! map.insert("foo", "bar");
//! map.insert("hello", "world");
//! map.insert("splay", "tree");
//!
//! for (k, v) in map.into_iter() {
//!     println!("{} => {}", k, v);
//! }
//! ```

#![feature(unsafe_destructor)]
//#![cfg_attr(not(test), deny(experimental, unstable))]

pub use set::SplaySet;
pub use map::SplayMap;

pub mod set;
pub mod map;
mod node;

#[cfg(test)]
mod test {
    use super::{SplayMap, SplaySet};

    // Lots of these are shamelessly stolen from the TreeMap tests, it'd be
    // awesome if they could share them...

    #[test]
    fn insert_simple() {
        let mut t = SplayMap::new();
        assert!(t.insert(1, 2));
        assert!(!t.insert(1, 3));
        assert!(t.insert(2, 3));
    }

    #[test]
    fn remove_simple() {
        let mut t = SplayMap::new();
        assert!(t.insert(1, 2));
        assert!(t.insert(2, 3));
        assert!(t.insert(3, 4));
        assert!(t.insert(0, 5));
        assert!(t.remove(&1));
    }

    #[test]
    fn pop_simple() {
        let mut t = SplayMap::new();
        assert!(t.insert(1, 2));
        assert!(t.insert(2, 3));
        assert!(t.insert(3, 4));
        assert!(t.insert(0, 5));
        assert_eq!(t.pop(&1), Some(2));
        assert_eq!(t.pop(&1), None);
        assert_eq!(t.pop(&0), Some(5));
    }

    #[test]
    fn find_mut_simple() {
        let mut t = SplayMap::new();
        assert!(t.insert(1, 2));
        assert!(t.insert(2, 3));
        assert!(t.insert(3, 4));
        assert!(t.insert(0, 5));

        {
            let ptr = t.find_mut(&1);
            assert!(ptr.is_some());
            let ptr = ptr.unwrap();
            assert!(*ptr == 2);
            *ptr = 4;
        }

        let ptr = t.find_mut(&1);
        assert!(ptr.is_some());
        assert!(*ptr.unwrap() == 4);
    }

    #[test]
    fn test_len() {
        let mut m = SplayMap::new();
        assert!(m.insert(3, 6));
        assert!(m.len() == 1);
        assert!(m.insert(0, 0));
        assert!(m.len() == 2);
        assert!(m.insert(4, 8));
        assert!(m.len() == 3);
        assert!(m.remove(&3));
        assert!(m.len() == 2);
        assert!(!m.remove(&5));
        assert!(m.len() == 2);
        assert!(m.insert(2, 4));
        assert!(m.len() == 3);
        assert!(m.insert(1, 2));
        assert!(m.len() == 4);
    }

    #[test]
    fn test_clear() {
        let mut m = SplayMap::new();
        m.clear();
        assert!(m.insert(5, 11));
        assert!(m.insert(12, -3));
        assert!(m.insert(19, 2));
        m.clear();
        assert!(m.find(&5).is_none());
        assert!(m.find(&12).is_none());
        assert!(m.find(&19).is_none());
        assert!(m.is_empty());
    }

    #[test]
    fn insert_replace() {
        let mut m = SplayMap::new();
        assert!(m.insert(5, 2));
        assert!(m.insert(2, 9));
        assert!(!m.insert(2, 11));
        assert!(m.find(&2).unwrap() == &11);
    }

    #[test]
    fn find_empty() {
        let m: SplayMap<i32, i32> = SplayMap::new();
        assert!(m.find(&5) == None);
    }

    #[test]
    fn find_not_found() {
        let mut m = SplayMap::new();
        assert!(m.insert(1, 2));
        assert!(m.insert(5, 3));
        assert!(m.insert(9, 3));
        assert!(m.find(&2) == None);
    }

    #[test]
    fn get_works() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        assert!(m[1] == 1);
    }

    #[test]
    fn into_iter() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        m.insert(2, 1);
        m.insert(0, 1);
        let mut cur = 0;
        for (k, v) in m.into_iter() {
            assert_eq!(k, cur);
            assert_eq!(v, 1);
            cur += 1;
        }
    }

    #[test]
    fn into_iter_backwards() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        m.insert(2, 1);
        m.insert(0, 1);
        let mut cur = 2;
        for (k, v) in m.into_iter().rev() {
            assert_eq!(k, cur);
            assert_eq!(v, 1);
            cur -= 1;
        }
    }

    #[test] #[should_fail]
    fn get_panic_works() {
        let mut m = SplayMap::new();
        m.insert(2, 2);
        m[1];
    }

    #[test]
    fn large() {
        use std::rand::random;
        let mut m = SplaySet::new();
        let mut v = Vec::new();

        for _ in range(0, 400) {
            let i: i32 = random();
            m.insert(i);
            v.push(i);
        }

        for i in v.iter() {
            assert!(m.contains(i));
        }
    }
}
<commit_msg>Fix some tests<commit_after>//! Contains an implementation of splay trees where each node has a key/value
//! pair to be used in maps and sets. The only requirement is that the key must
//! implement the Ord trait.
//!
//! # Example
//!
//! ```rust
//! use splay::SplayMap;
//!
//! let mut map = SplayMap::new();
//! map.insert("foo", "bar");
//! map.insert("hello", "world");
//! map.insert("splay", "tree");
//!
//! for (k, v) in map.into_iter() {
//!     println!("{} => {}", k, v);
//! }
//! ```

#![feature(unsafe_destructor)]
#![cfg_attr(test, allow(unstable))]
#![cfg_attr(test, deny(warnings))]

pub use self::set::SplaySet;
pub use self::map::SplayMap;

pub mod set;
pub mod map;
mod node;

#[cfg(test)]
mod test {
    use super::{SplayMap, SplaySet};

    // Lots of these are shamelessly stolen from the TreeMap tests, it'd be
    // awesome if they could share them...

    #[test]
    fn insert_simple() {
        let mut t = SplayMap::new();
        assert_eq!(t.insert(1, 2), None);
        assert_eq!(t.insert(1, 3), Some(2));
        assert_eq!(t.insert(2, 3), None);
    }

    #[test]
    fn remove_simple() {
        let mut t = SplayMap::new();
        assert_eq!(t.insert(1, 2), None);
        assert_eq!(t.insert(2, 3), None);
        assert_eq!(t.insert(3, 4), None);
        assert_eq!(t.insert(0, 5), None);
        assert_eq!(t.remove(&1), Some(2));
    }

    #[test]
    fn pop_simple() {
        let mut t = SplayMap::new();
        assert_eq!(t.insert(1, 2), None);
        assert_eq!(t.insert(2, 3), None);
        assert_eq!(t.insert(3, 4), None);
        assert_eq!(t.insert(0, 5), None);
        assert_eq!(t.remove(&1), Some(2));
        assert_eq!(t.remove(&1), None);
        assert_eq!(t.remove(&0), Some(5));
    }

    #[test]
    fn find_mut_simple() {
        let mut t = SplayMap::new();
        assert_eq!(t.insert(1, 2), None);
        assert_eq!(t.insert(2, 3), None);
        assert_eq!(t.insert(3, 4), None);
        assert_eq!(t.insert(0, 5), None);

        {
            let ptr = t.get_mut(&1);
            assert!(ptr.is_some());
            let ptr = ptr.unwrap();
            assert_eq!(*ptr, 2);
            *ptr = 4;
        }

        let ptr = t.get_mut(&1);
        assert!(ptr.is_some());
        assert_eq!(*ptr.unwrap(), 4);
    }

    #[test]
    fn test_len() {
        let mut m = SplayMap::new();
        assert_eq!(m.insert(3, 6), None);
        assert_eq!(m.len(), 1);
        assert_eq!(m.insert(0, 0), None);
        assert_eq!(m.len(), 2);
        assert_eq!(m.insert(4, 8), None);
        assert_eq!(m.len(), 3);
        assert_eq!(m.remove(&3), Some(6));
        assert_eq!(m.len(), 2);
        assert_eq!(m.remove(&5), None);
        assert_eq!(m.len(), 2);
        assert_eq!(m.insert(2, 4), None);
        assert_eq!(m.len(), 3);
        assert_eq!(m.insert(1, 2), None);
        assert_eq!(m.len(), 4);
    }

    #[test]
    fn test_clear() {
        let mut m = SplayMap::new();
        m.clear();
        assert_eq!(m.insert(5, 11), None);
        assert_eq!(m.insert(12, -3), None);
        assert_eq!(m.insert(19, 2), None);
        m.clear();
        assert_eq!(m.get(&5), None);
        assert_eq!(m.get(&12), None);
        assert_eq!(m.get(&19), None);
        assert!(m.is_empty());
    }

    #[test]
    fn insert_replace() {
        let mut m = SplayMap::new();
        assert_eq!(m.insert(5, 2), None);
        assert_eq!(m.insert(2, 9), None);
        assert_eq!(m.insert(2, 11), Some(9));
        assert_eq!(m[2], 11);
    }

    #[test]
    fn find_empty() {
        let m: SplayMap<i32, i32> = SplayMap::new();
        assert_eq!(m.get(&5), None);
    }

    #[test]
    fn find_not_found() {
        let mut m = SplayMap::new();
        assert_eq!(m.insert(1, 2), None);
        assert_eq!(m.insert(5, 3), None);
        assert_eq!(m.insert(9, 3), None);
        assert_eq!(m.get(&2), None);
    }

    #[test]
    fn get_works() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        assert_eq!(m[1], 1);
    }

    #[test]
    fn into_iter() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        m.insert(2, 1);
        m.insert(0, 1);
        let mut cur = 0;
        for (k, v) in m.into_iter() {
            assert_eq!(k, cur);
            assert_eq!(v, 1);
            cur += 1;
        }
    }

    #[test]
    fn into_iter_backwards() {
        let mut m = SplayMap::new();
        m.insert(1, 1);
        m.insert(2, 1);
        m.insert(0, 1);
        let mut cur = 2;
        for (k, v) in m.into_iter().rev() {
            assert_eq!(k, cur);
            assert_eq!(v, 1);
            cur -= 1;
        }
    }

    #[test] #[should_fail]
    fn get_panic_works() {
        let mut m = SplayMap::new();
        m.insert(2, 2);
        m[1];
    }

    #[test]
    fn large() {
        use std::rand::random;
        let mut m = SplaySet::new();
        let mut v = Vec::new();

        for _ in range(0, 400) {
            let i: i32 = random();
            m.insert(i);
            v.push(i);
        }

        for i in v.iter() {
            assert!(m.contains(i));
        }
    }
}
<|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<HttpStream>) -> 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 {
				println!("Writing {}..{}", self.write_pos, bytes.len());
				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 (l, 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)
		}));
		::std::thread::spawn(move || {
			srv.run();
 		});
		Ok(Server {
			server: Some(l),
			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>Removed log<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<HttpStream>) -> 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 (l, 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)
		}));
		::std::thread::spawn(move || {
			srv.run();
 		});
		Ok(Server {
			server: Some(l),
			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>mod index;
mod matching;
<commit_msg>Export matching module.<commit_after>mod index;
pub mod matching;
<|endoftext|>


<commit_before>#![feature(libc)]

#![allow(dead_code)]
#![allow(non_camel_case_types)]

extern crate libc;
use libc::c_int;

/**
 * Abstraction of libnl core functionality
 */
#[link(name="nl-3")]
extern "C" {
	// Exposed socket functions
	fn nl_socket_alloc() -> *const nl_sock;
	fn nl_socket_free(socket: *const nl_sock);
	fn nl_socket_get_fd(socket: *const nl_sock) -> i32;
	fn nl_socket_set_buffer_size(socket: *const nl_sock, rxbuf: c_int, txbuf: c_int) -> i32;

	fn nl_socket_set_local_port(socket: *const nl_sock, port: u32);
	fn nl_socket_get_local_port(socket: *const nl_sock) -> u32;

	fn nl_connect(socket: *const nl_sock, protocol: u32) -> i32;
	fn nl_close(socket: *const nl_sock);

	fn nl_socket_set_cb(socket: *const nl_sock, cb: *const nl_cb);
	fn nl_socket_get_cb(socket: *const nl_sock) -> nl_cb;

	// Exposed socket transceiver
	fn nl_send_simple(socket: *const nl_sock, msg_type: c_int, flags: c_int, buf: *const u8, size: c_int) -> i32;
	// Exposed msg functions
	fn nlmsg_alloc() -> *const nl_msg;
	fn nlmsg_free(msg: *const nl_msg);

}

// exposed structures - these are wrapped
struct nl_sock;
struct nl_msg;
struct nl_cb;

// RSNL datatypes wrapping the libnl data structures
pub struct socket {
	ptr: *const nl_sock
}

/* library version of nlmsghdr
 * the name is altered to differentiate
 * between the native libnl and rsnl
 */
pub struct rsnl_msghdr {
	nlmsg_len: u32,
	nlmsg_type: u16,
	nlmsg_flags: u16,
	nlmsg_seq: u32,
	nlmsg_pid: u32
}

pub struct msg {
	ptr: *const nl_msg
}

pub enum NetlinkProtocol {
	route,
	unused,
	usersock,
	firewall,
	sock_diag,
	nflog,
	xfrm,
	selinux,
	iscsi,
	audit,
	fib_lookup,
	connector,
	netfilter,
	ip6_fw,
	dnrtmsg,
	kobject_uevent,
	DMEVENTS,
	scsitransport,
	ecryptfs,
	rdma,
	crypto,
	zu= 30
}

impl socket {
	pub fn new() -> socket {
		unsafe {
			let nlptr = nl_socket_alloc();
			socket {
				ptr: nlptr
			}
		}
	}

	pub fn free(&self) {
		unsafe {
			nl_socket_free(self.ptr)
		}
	}

	pub fn set_buffer_size(&self, rxbuf: c_int, txbuf: c_int) -> i32 {
		unsafe {
			nl_socket_set_buffer_size(self.ptr, rxbuf, txbuf)
		}
	}

	pub fn connect(&self, protocol: NetlinkProtocol) -> i32 {
		unsafe { nl_connect(self.ptr, protocol as u32) }
	}

	pub fn close(&self) {
		unsafe{ nl_close(self.ptr) }
	}

	pub fn get_fd(&self) -> i32 {
		unsafe { nl_socket_get_fd(self.ptr) }
	}

	pub fn get_local_port(&self) -> u32 {
		unsafe { nl_socket_get_local_port(self.ptr) }
	}

	pub fn set_local_port(&self, port: u32) {
		unsafe { nl_socket_set_local_port(self.ptr, port) }
	}

	pub fn send_simple(&self, msg_type: c_int, flags: c_int, buf: *const u8, size: c_int) -> c_int {
		unsafe { nl_send_simple(self.ptr, msg_type, flags, buf, size) }
	}
}

impl msg {
	pub fn new() -> msg {
	unsafe {
		let nlmsg = nlmsg_alloc();
		msg {
			ptr: nlmsg
		}
	}
	}

	pub fn free(&self) {
		unsafe{ nlmsg_free(self.ptr); }
	}
}
<commit_msg>Structures to/from C FFI need to be marked as C repr<commit_after>#![feature(libc)]

#![allow(dead_code)]
#![allow(non_camel_case_types)]

extern crate libc;
use libc::c_int;

/**
 * Abstraction of libnl core functionality
 */
#[link(name="nl-3")]
extern "C" {
	// Exposed socket functions
	fn nl_socket_alloc() -> *const nl_sock;
	fn nl_socket_free(socket: *const nl_sock);
	fn nl_socket_get_fd(socket: *const nl_sock) -> i32;
	fn nl_socket_set_buffer_size(socket: *const nl_sock, rxbuf: c_int, txbuf: c_int) -> i32;

	fn nl_socket_set_local_port(socket: *const nl_sock, port: u32);
	fn nl_socket_get_local_port(socket: *const nl_sock) -> u32;

	fn nl_connect(socket: *const nl_sock, protocol: u32) -> i32;
	fn nl_close(socket: *const nl_sock);

	fn nl_socket_set_cb(socket: *const nl_sock, cb: *const nl_cb);
	fn nl_socket_get_cb(socket: *const nl_sock) -> nl_cb;

	// Exposed socket transceiver
	fn nl_send_simple(socket: *const nl_sock, msg_type: c_int, flags: c_int, buf: *const u8, size: c_int) -> i32;
	// Exposed msg functions
	fn nlmsg_alloc() -> *const nl_msg;
	fn nlmsg_free(msg: *const nl_msg);

}

// exposed structures - these are wrapped
#[repr(C)]
struct nl_sock;
#[repr(C)]
struct nl_msg;
#[repr(C)]
struct nl_cb;

// RSNL datatypes wrapping the libnl data structures
pub struct socket {
	ptr: *const nl_sock
}

/* library version of nlmsghdr
 * the name is altered to differentiate
 * between the native libnl and rsnl
 */
pub struct rsnl_msghdr {
	nlmsg_len: u32,
	nlmsg_type: u16,
	nlmsg_flags: u16,
	nlmsg_seq: u32,
	nlmsg_pid: u32
}

pub struct msg {
	ptr: *const nl_msg
}

pub enum NetlinkProtocol {
	route,
	unused,
	usersock,
	firewall,
	sock_diag,
	nflog,
	xfrm,
	selinux,
	iscsi,
	audit,
	fib_lookup,
	connector,
	netfilter,
	ip6_fw,
	dnrtmsg,
	kobject_uevent,
	DMEVENTS,
	scsitransport,
	ecryptfs,
	rdma,
	crypto,
	zu= 30
}

impl socket {
	pub fn new() -> socket {
		unsafe {
			let nlptr = nl_socket_alloc();
			socket {
				ptr: nlptr
			}
		}
	}

	pub fn free(&self) {
		unsafe {
			nl_socket_free(self.ptr)
		}
	}

	pub fn set_buffer_size(&self, rxbuf: c_int, txbuf: c_int) -> i32 {
		unsafe {
			nl_socket_set_buffer_size(self.ptr, rxbuf, txbuf)
		}
	}

	pub fn connect(&self, protocol: NetlinkProtocol) -> i32 {
		unsafe { nl_connect(self.ptr, protocol as u32) }
	}

	pub fn close(&self) {
		unsafe{ nl_close(self.ptr) }
	}

	pub fn get_fd(&self) -> i32 {
		unsafe { nl_socket_get_fd(self.ptr) }
	}

	pub fn get_local_port(&self) -> u32 {
		unsafe { nl_socket_get_local_port(self.ptr) }
	}

	pub fn set_local_port(&self, port: u32) {
		unsafe { nl_socket_set_local_port(self.ptr, port) }
	}

	pub fn send_simple(&self, msg_type: c_int, flags: c_int, buf: *const u8, size: c_int) -> c_int {
		unsafe { nl_send_simple(self.ptr, msg_type, flags, buf, size) }
	}
}

impl msg {
	pub fn new() -> msg {
	unsafe {
		let nlmsg = nlmsg_alloc();
		msg {
			ptr: nlmsg
		}
	}
	}

	pub fn free(&self) {
		unsafe{ nlmsg_free(self.ptr); }
	}
}
<|endoftext|>


<commit_before>extern crate num;
extern crate rutorch;

use rutorch::*;
use std::ops::Index;
use num::Num;
use std::cmp::max;

pub trait ArgsArray<T> {
	fn is_leaf(&self) -> bool;
	fn child(&self) -> &T;
	fn dim(&self) -> usize;
}

impl <T> ArgsArray<T> for [T]  {
	fn is_leaf(&self) -> bool { false }
	fn child(&self) ->  &T { &self[0]}
	fn dim(&self) -> usize { self.len()}
}

impl <T:Num + Sized> ArgsArray<T> for T  {
	fn is_leaf(&self) -> bool { true }
	fn child(&self) ->  &T { &self}
	fn dim(&self) -> usize { 1 }
}

//fn from(args: &ArgsArray<T>) -> Self;

trait Tensor<'a> : Index<&'a [isize]> {
	fn new() -> Self;
}

struct FloatTensor {
	t: *mut THFloatTensor,
	dims: Vec<isize>,
}

impl <'a>Tensor<'a> for FloatTensor {
	fn new() -> Self {
		unsafe { FloatTensor { t : THFloatTensor_new(), dims: Vec::new()  }}
	}
}

impl <'a> Index<&'a [isize]> for FloatTensor {
	type Output = f32;

	fn index(&self, idx: &'a [isize]) -> &f32 {
		let mut index : isize = 0;
		let lastidx = max(0, idx.len() as isize - 1) as usize;
		if idx.len() != self.dims.len() { panic!("bad dimlen")}
		for i in 0..lastidx {
			if idx[i] >= self.dims[i] { panic!("bad dimlen")}
			index += idx[i] * self.dims[i]
		}
		if idx[lastidx] >= self.dims[lastidx] { panic!("bad dimlen")}
		index += idx[lastidx];
		unsafe {&*(*(*self.t).storage).data.offset(index)}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
    #[test]
    fn args_array_works() {
    	let v = [[5, 2], [3, 4]];
    	let x = v.child();
    	let z = x.child();
    	println!("is_leaf = {} {} {:?}, {} {} {:?}, {} {} {:?}",
    		v.is_leaf(), v.dim(), v,
            x.is_leaf(), x.dim(), x,
            z.is_leaf(), z.dim(), z
            );
    }
    fn tensor_ops_work() {
		let v = [[5, 2], [3, 4]];

    }
}
<commit_msg>implement FloatStorage and sized FloatTensor allocation<commit_after>extern crate num;
extern crate rutorch;

use rutorch::*;
use std::ops::Index;
use num::Num;
use std::cmp::max;

pub trait ArgsArray<T> {
	fn is_leaf(&self) -> bool;
	fn child(&self) -> &T;
	fn dim(&self) -> usize;
}

impl <T> ArgsArray<T> for [T]  {
	fn is_leaf(&self) -> bool { false }
	fn child(&self) ->  &T { &self[0]}
	fn dim(&self) -> usize { self.len()}
}

impl <T:Num + Sized> ArgsArray<T> for T  {
	fn is_leaf(&self) -> bool { true }
	fn child(&self) ->  &T { &self}
	fn dim(&self) -> usize { 1 }
}

trait Storage {
	fn new() -> Self;
	fn sized(len: isize) -> Self;
}

struct FloatStorage {
	t: *mut THFloatStorage,
}

impl Storage for FloatStorage {
	fn new() -> Self {
		unsafe { FloatStorage  {t : THFloatStorage_new()} }
	}
	fn sized(size: isize) ->Self {
		unsafe { FloatStorage  {t : THFloatStorage_newWithSize(size)} }
	}
}

impl Index<isize> for FloatStorage {
	type Output = f32;

	fn index(&self, idx: isize) -> &f32 {
		unsafe {&*(*self.t).data.offset(idx)}
	}
}

//fn from(args: &ArgsArray<T>) -> Self;

trait Tensor<'a> : Index<&'a [isize]> {
	fn new() -> Self;
	fn sized(size: &'a [isize]) -> Self;
	fn randn(dims: &'a [isize]) -> Self;
}

struct FloatTensor {
	t: *mut THFloatTensor,
	storage: FloatStorage,
	dims: Vec<isize>,
}

impl <'a>Tensor<'a> for FloatTensor {
	fn new() -> Self {
		unsafe { FloatTensor { t : THFloatTensor_new(), storage: FloatStorage::new(), dims: Vec::new()  }}
	}
	fn sized(dims: &'a [isize]) -> Self {
		let size = dims.iter().product();
		let storage = FloatStorage::sized(size);
		let strides = vec![1; dims.len()];
		let mut t = THFloatTensor { size : dims.clone().as_ptr() as *mut ::std::os::raw::c_long,
			stride: strides.as_ptr() as *mut ::std::os::raw::c_long,
			nDimension : dims.len() as i32, storage: storage.t, storageOffset: 0, refcount: 1,
			flag: TH_TENSOR_REFCOUNTED as i8};
		FloatTensor {t : &mut t, storage: storage, dims: Vec::from(dims)}
	}
	fn randn(dims: &'a [isize]) -> Self {
		/* XXX */
		Tensor::sized(dims)
	}
}

impl <'a> Index<&'a [isize]> for FloatTensor {
	type Output = f32;

	fn index(&self, idx: &'a [isize]) -> &f32 {
		let mut index : isize = 0;
		let lastidx = max(0, idx.len() as isize - 1) as usize;
		if idx.len() != self.dims.len() { panic!("bad dimlen")}
		for i in 0..lastidx {
			if idx[i] >= self.dims[i] { panic!("bad dimlen")}
			index += idx[i] * self.dims[i]
		}
		if idx[lastidx] >= self.dims[lastidx] { panic!("bad dimlen")}
		index += idx[lastidx];
		&self.storage[index]
	}
}

#[cfg(test)]
mod tests {
	use super::*;
    #[test]
    fn args_array_works() {
    	let v = [[5, 2], [3, 4]];
    	let x = v.child();
    	let z = x.child();
    	println!("is_leaf = {} {} {:?}, {} {} {:?}, {} {} {:?}",
    		v.is_leaf(), v.dim(), v,
            x.is_leaf(), x.dim(), x,
            z.is_leaf(), z.dim(), z
            );
    }
    fn tensor_ops_work() {
		let v = [[5, 2], [3, 4]];

    }
}
<|endoftext|>


<commit_before>extern crate lmdb_rs as lmdb;

use lmdb::{FromMdbValue, ToMdbValue, MdbValue,
           DbFlags, DbHandle, Database, Environment, Cursor, MDB_val};
use lmdb::core::MdbResult;
use std::marker::PhantomData;

#[macro_export]
macro_rules! lmdb_not_found {
    (
        $x:expr
    ) => {
        match $x {
            Err(lmdb::MdbError::NotFound) => true,
            _ => false
        }
    };
}

#[macro_export]
macro_rules! impl_table {
    (
        $s:ident, $k:ty, $v:ty
    ) => {
        impl $s {
            pub fn table<'db>(db: ::lmdb::Database<'db>) -> ::lmdb::core::MdbResult<Table<'db, $k, $v>> {
                try!(Self::setup(&db));
                Ok(Table::new(db))
            }
        }
    };
}

/// Defines all neccessary information to create/open a database.
pub trait TableDef {
    fn name() -> &'static str;
    fn flags() -> DbFlags;
    fn setup(db: &Database) -> MdbResult<()>;

    fn open(env: &Environment, create: bool) -> MdbResult<DbHandle> {
        if create {
            env.create_db(Self::name(), Self::flags())
        } else {
            env.get_db(Self::name(), Self::flags())
        }
    }
}

/// Wraps a lmdb::Database, keeping type information for K and V.
pub struct Table<'db, K, V> {
    db: Database<'db>,
    k: PhantomData<K>,
    v: PhantomData<V>,
}

impl<'db, K: FromMdbValue+ToMdbValue, V: FromMdbValue+ToMdbValue> Table<'db, K, V> {
    #[inline(always)]
    pub fn new(db: Database<'db>) -> Table<'db, K, V> {
        Table {db: db, k: PhantomData, v: PhantomData}
    }

    #[inline(always)]
    pub fn set(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.set(key, value)
    }

    #[inline(always)]
    pub fn insert(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.insert(key, value)
    }

    /// Checks if the item exists.
    pub fn insert_item(&self, key: &K, value: &V) -> MdbResult<()> {
        if try!(self.has_item(key, value)) {
            return Err(::lmdb::MdbError::KeyExists);
        }
        self.set(key, value)
    }

    /// Returns true if item exists.
    pub fn has_item(&self, key: &K, value: &V) -> MdbResult<bool> {
        let mut cursor = try!(self.new_cursor());
        match cursor.to_item(key, value) {
            Ok(()) => Ok(true),
            Err(::lmdb::MdbError::NotFound) => Ok(false),
            Err(e) => Err(e),
        }
    }

    #[inline(always)]
    pub fn get(&self, key: &K) -> MdbResult<V> {
        self.db.get(key)
    }

    #[inline(always)]
    pub fn del(&self, key: &K) -> MdbResult<()> {
        self.db.del(key)
    }

    #[inline(always)]
    pub fn del_item(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.del_item(key, value)
    }

    #[inline(always)]
    pub fn new_cursor<'table>(&'table self) -> MdbResult<TypedCursor<'table, K, V>> {
        Ok(TypedCursor {cursor: try!(self.db.new_cursor()), k: PhantomData, v: PhantomData})
    }
}

/// Is a typed version of lmdb::Cursor.
pub struct TypedCursor<'table, K, V> {
    cursor: Cursor<'table>,
    k: PhantomData<K>,
    v: PhantomData<V>,
}

impl<'table, K: FromMdbValue+ToMdbValue, V: FromMdbValue+ToMdbValue> TypedCursor<'table, K, V> {
    #[inline(always)]
    pub fn to_first(&mut self) -> MdbResult<()> {
        self.cursor.to_first()
    }

    #[inline(always)]
    pub fn to_last(&mut self) -> MdbResult<()> {
        self.cursor.to_last()
    }

    #[inline(always)]
    pub fn to_key(&mut self, key: &K) -> MdbResult<()> {
        self.cursor.to_key(key)
    }

    #[inline(always)]
    pub fn to_gte_key(&mut self, key: &K) -> MdbResult<()> {
        self.cursor.to_gte_key(key)
    }

    #[inline(always)]
    pub fn to_next_key(&mut self) -> MdbResult<()> {
        self.cursor.to_next_key()
    }

    #[inline(always)]
    pub fn to_prev_key(&mut self) -> MdbResult<()> {
        self.cursor.to_prev_key()
    }

    #[inline(always)]
    pub fn to_item(&mut self, key: &K, value: &V) -> MdbResult<()> {
        self.cursor.to_item(key, value)
    }

    #[inline(always)]
    pub fn to_gte_item(&mut self, key: &K, value: &V) -> MdbResult<()> {
        self.cursor.to_gte_item(key, value)
    }

    #[inline(always)]
    pub fn to_first_item(&mut self) -> MdbResult<()> {
        self.cursor.to_first_item()
    }

    #[inline(always)]
    pub fn to_last_item(&mut self) -> MdbResult<()> {
        self.cursor.to_last_item()
    }

    #[inline(always)]
    pub fn to_next_item(&mut self) -> MdbResult<()> {
        self.cursor.to_next_item()
    }

    #[inline(always)]
    pub fn to_prev_item(&mut self) -> MdbResult<()> {
        self.cursor.to_prev_item()
    }

    #[inline(always)]
    pub fn item_count(&mut self) -> MdbResult<u64> {
        self.cursor.item_count()
    }

    #[inline(always)]
    pub fn get(&mut self) -> MdbResult<(K, V)> {
        self.cursor.get()
    }

    #[inline(always)]
    pub fn get_value(&mut self) -> MdbResult<V> {
        self.cursor.get_value()
    }

    #[inline(always)]
    pub fn get_key(&mut self) -> MdbResult<K> {
        self.cursor.get_key()
    }

    #[inline(always)]
    pub fn replace(&mut self, value: &V) -> MdbResult<()> {
        self.cursor.replace(value)
    }

    #[inline(always)]
    pub fn add_item(&mut self, value: &V) -> MdbResult<()> {
        self.cursor.add_item(value)
    }

    #[inline(always)]
    pub fn del(&mut self) -> MdbResult<()> {
        self.cursor.del()
    }

    #[inline(always)]
    pub fn del_item(&mut self) -> MdbResult<()> {
        self.cursor.del_item()
    }

    #[inline(always)]
    pub fn del_all(&mut self) -> MdbResult<()> {
        self.cursor.del_all()
    }
}

pub extern "C" fn sort<T:FromMdbValue+Ord>(lhs_val: *const MDB_val, rhs_val: *const MDB_val) -> lmdb::c_int {
    let lhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(lhs_val)});
    let rhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(rhs_val)});

    let order: i8 = unsafe { std::mem::transmute(lhs.cmp(&rhs)) };
    order as lmdb::c_int
}

pub extern "C" fn sort_reverse<T:FromMdbValue+Ord>(lhs_val: *const MDB_val, rhs_val: *const MDB_val) -> lmdb::c_int {
    let lhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(lhs_val)});
    let rhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(rhs_val)});
    let order: i8 = unsafe { std::mem::transmute(lhs.cmp(&rhs).reverse()) };
    order as lmdb::c_int
}

#[test]
fn test_simple_table() {
    use std::path::Path;
    use lmdb::core::{DbIntKey, DbAllowDups, DbAllowIntDups, DbDupFixed};

    let env = lmdb::EnvBuilder::new().max_dbs(2).autocreate_dir(true).open(&Path::new("./test/db1"), 0o777).unwrap();

    struct MyFirstTable;
    impl TableDef for MyFirstTable {
        fn name() -> &'static str { "my_first_table" }
        fn flags() -> DbFlags { DbIntKey | DbAllowDups | DbAllowIntDups | DbDupFixed }
        fn setup(db: &Database) -> MdbResult<()> {
            db.set_dupsort(sort_reverse::<u64>)
        }
    }
    impl MyFirstTable {
        fn table<'db>(db: Database<'db>) -> MdbResult<Table<'db, u64, u64>> {
            try!(Self::setup(&db));
            Ok(Table{db: db, k: PhantomData, v: PhantomData})
        }
    }

    struct MyBlobTable;
    impl TableDef for MyBlobTable {
        fn name() -> &'static str { "my_blob_table" }
        fn flags() -> DbFlags { DbIntKey }
        fn setup(_db: &Database) -> MdbResult<()> { Ok(()) }
    }
    impl MyBlobTable {
        fn table<'db>(db: Database<'db>) -> MdbResult<Table<'db, u64, &[u8]>> {
            try!(Self::setup(&db));
            Ok(Table{db: db, k: PhantomData, v: PhantomData})
        }
    }

    // A Unique(Id64, Id64 DESC) table
    let table_handle = MyFirstTable::open(&env, true).unwrap();
    let blobs_handle = MyBlobTable::open(&env, true).unwrap();

    // prepare database
    {
        let txn = env.new_transaction().unwrap();
        {
            let table = MyFirstTable::table(txn.bind(&table_handle)).unwrap();
            let blobs = MyBlobTable::table(txn.bind(&blobs_handle)).unwrap();

            let big_blob: &[u8] = b"Test";
            blobs.set(&1, &big_blob).unwrap();

            let back: &[u8] = blobs.get(&1).unwrap();
            assert_eq!(&back[..], &b"Test"[..]);

            let key = 1;
            table.set(&key, &200).unwrap();
            table.set(&key, &200).unwrap();
            table.set(&key, &100).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &1).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &3).unwrap();
            table.set(&key, &100).unwrap();
        }
        txn.commit().unwrap();
    }
    // read
    {
        let rdr = env.get_reader().unwrap();

        let table = MyFirstTable::table(rdr.bind(&table_handle)).unwrap();

        let mut cursor = table.new_cursor().unwrap();
        cursor.to_key(&1).unwrap(); //  positions on first item of key
        assert_eq!((1, 200), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 100), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 3), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 2), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 1), cursor.get().unwrap());

        assert!(lmdb_not_found!(cursor.to_next_item()));

        let mut cursor = table.new_cursor().unwrap();
        {
            assert!(lmdb_not_found!(cursor.to_item(&2, &3)));

            cursor.to_gte_item(&1, &4).unwrap();
            assert_eq!((1, 3), cursor.get().unwrap());

            cursor.to_next_item().unwrap();
            assert_eq!((1, 2), cursor.get().unwrap());

            cursor.to_next_item().unwrap();
            assert_eq!((1, 1), cursor.get().unwrap());

            assert!(lmdb_not_found!(cursor.to_next_item()));
        }
    }
}
<commit_msg>Add get_ref()<commit_after>extern crate lmdb_rs as lmdb;

use lmdb::{FromMdbValue, ToMdbValue, MdbValue,
           DbFlags, DbHandle, Database, Environment, Cursor, MDB_val};
use lmdb::core::MdbResult;
use std::marker::PhantomData;

#[macro_export]
macro_rules! lmdb_not_found {
    (
        $x:expr
    ) => {
        match $x {
            Err(lmdb::MdbError::NotFound) => true,
            _ => false
        }
    };
}

#[macro_export]
macro_rules! impl_table {
    (
        $s:ident, $k:ty, $v:ty
    ) => {
        impl $s {
            pub fn table<'db>(db: ::lmdb::Database<'db>) -> ::lmdb::core::MdbResult<Table<'db, $k, $v>> {
                try!(Self::setup(&db));
                Ok(Table::new(db))
            }
        }
    };
}

/// Defines all neccessary information to create/open a database.
pub trait TableDef {
    fn name() -> &'static str;
    fn flags() -> DbFlags;
    fn setup(db: &Database) -> MdbResult<()>;

    fn open(env: &Environment, create: bool) -> MdbResult<DbHandle> {
        if create {
            env.create_db(Self::name(), Self::flags())
        } else {
            env.get_db(Self::name(), Self::flags())
        }
    }
}

/// Wraps a lmdb::Database, keeping type information for K and V.
pub struct Table<'db, K, V> {
    db: Database<'db>,
    k: PhantomData<K>,
    v: PhantomData<V>,
}

impl<'db, K, V> Table<'db, K, V>
where K: FromMdbValue + ToMdbValue,
      V: FromMdbValue + ToMdbValue,
{
    #[inline(always)]
    pub fn new(db: Database<'db>) -> Table<'db, K, V> {
        Table {db: db, k: PhantomData, v: PhantomData}
    }

    #[inline(always)]
    pub fn set(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.set(key, value)
    }

    #[inline(always)]
    pub fn insert(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.insert(key, value)
    }

    /// Checks if the item exists.
    pub fn insert_item(&self, key: &K, value: &V) -> MdbResult<()> {
        if try!(self.has_item(key, value)) {
            return Err(::lmdb::MdbError::KeyExists);
        }
        self.set(key, value)
    }

    /// Returns true if item exists.
    pub fn has_item(&self, key: &K, value: &V) -> MdbResult<bool> {
        let mut cursor = try!(self.new_cursor());
        match cursor.to_item(key, value) {
            Ok(()) => Ok(true),
            Err(::lmdb::MdbError::NotFound) => Ok(false),
            Err(e) => Err(e),
        }
    }

    #[inline(always)]
    pub fn get(&self, key: &K) -> MdbResult<V> {
        self.db.get(key)
    }

    #[inline(always)]
    pub fn get_ref<'a, V2>(&self, key: &K) -> MdbResult<&'a V2> where &'a V2: FromMdbValue {
        self.db.get(key)
    }

    #[inline(always)]
    pub fn del(&self, key: &K) -> MdbResult<()> {
        self.db.del(key)
    }

    #[inline(always)]
    pub fn del_item(&self, key: &K, value: &V) -> MdbResult<()> {
        self.db.del_item(key, value)
    }

    #[inline(always)]
    pub fn new_cursor<'table>(&'table self) -> MdbResult<TypedCursor<'table, K, V>> {
        Ok(TypedCursor {cursor: try!(self.db.new_cursor()), k: PhantomData, v: PhantomData})
    }
}

/// Is a typed version of lmdb::Cursor.
pub struct TypedCursor<'table, K, V> {
    cursor: Cursor<'table>,
    k: PhantomData<K>,
    v: PhantomData<V>,
}

impl<'table, K: FromMdbValue+ToMdbValue, V: FromMdbValue+ToMdbValue> TypedCursor<'table, K, V> {
    #[inline(always)]
    pub fn to_first(&mut self) -> MdbResult<()> {
        self.cursor.to_first()
    }

    #[inline(always)]
    pub fn to_last(&mut self) -> MdbResult<()> {
        self.cursor.to_last()
    }

    #[inline(always)]
    pub fn to_key(&mut self, key: &K) -> MdbResult<()> {
        self.cursor.to_key(key)
    }

    #[inline(always)]
    pub fn to_gte_key(&mut self, key: &K) -> MdbResult<()> {
        self.cursor.to_gte_key(key)
    }

    #[inline(always)]
    pub fn to_next_key(&mut self) -> MdbResult<()> {
        self.cursor.to_next_key()
    }

    #[inline(always)]
    pub fn to_prev_key(&mut self) -> MdbResult<()> {
        self.cursor.to_prev_key()
    }

    #[inline(always)]
    pub fn to_item(&mut self, key: &K, value: &V) -> MdbResult<()> {
        self.cursor.to_item(key, value)
    }

    #[inline(always)]
    pub fn to_gte_item(&mut self, key: &K, value: &V) -> MdbResult<()> {
        self.cursor.to_gte_item(key, value)
    }

    #[inline(always)]
    pub fn to_first_item(&mut self) -> MdbResult<()> {
        self.cursor.to_first_item()
    }

    #[inline(always)]
    pub fn to_last_item(&mut self) -> MdbResult<()> {
        self.cursor.to_last_item()
    }

    #[inline(always)]
    pub fn to_next_item(&mut self) -> MdbResult<()> {
        self.cursor.to_next_item()
    }

    #[inline(always)]
    pub fn to_prev_item(&mut self) -> MdbResult<()> {
        self.cursor.to_prev_item()
    }

    #[inline(always)]
    pub fn item_count(&mut self) -> MdbResult<u64> {
        self.cursor.item_count()
    }

    #[inline(always)]
    pub fn get(&mut self) -> MdbResult<(K, V)> {
        self.cursor.get()
    }

    #[inline(always)]
    pub fn get_value(&mut self) -> MdbResult<V> {
        self.cursor.get_value()
    }

    #[inline(always)]
    pub fn get_key(&mut self) -> MdbResult<K> {
        self.cursor.get_key()
    }

    #[inline(always)]
    pub fn replace(&mut self, value: &V) -> MdbResult<()> {
        self.cursor.replace(value)
    }

    #[inline(always)]
    pub fn add_item(&mut self, value: &V) -> MdbResult<()> {
        self.cursor.add_item(value)
    }

    #[inline(always)]
    pub fn del(&mut self) -> MdbResult<()> {
        self.cursor.del()
    }

    #[inline(always)]
    pub fn del_item(&mut self) -> MdbResult<()> {
        self.cursor.del_item()
    }

    #[inline(always)]
    pub fn del_all(&mut self) -> MdbResult<()> {
        self.cursor.del_all()
    }
}

pub extern "C" fn sort<T:FromMdbValue+Ord>(lhs_val: *const MDB_val, rhs_val: *const MDB_val) -> lmdb::c_int {
    let lhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(lhs_val)});
    let rhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(rhs_val)});

    let order: i8 = unsafe { std::mem::transmute(lhs.cmp(&rhs)) };
    order as lmdb::c_int
}

pub extern "C" fn sort_reverse<T:FromMdbValue+Ord>(lhs_val: *const MDB_val, rhs_val: *const MDB_val) -> lmdb::c_int {
    let lhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(lhs_val)});
    let rhs = T::from_mdb_value(&unsafe{MdbValue::from_raw(rhs_val)});
    let order: i8 = unsafe { std::mem::transmute(lhs.cmp(&rhs).reverse()) };
    order as lmdb::c_int
}

#[test]
fn test_simple_table() {
    use std::path::Path;
    use lmdb::core::{DbIntKey, DbAllowDups, DbAllowIntDups, DbDupFixed};

    let env = lmdb::EnvBuilder::new().max_dbs(2).autocreate_dir(true).open(&Path::new("./test/db1"), 0o777).unwrap();

    struct MyFirstTable;
    impl TableDef for MyFirstTable {
        fn name() -> &'static str { "my_first_table" }
        fn flags() -> DbFlags { DbIntKey | DbAllowDups | DbAllowIntDups | DbDupFixed }
        fn setup(db: &Database) -> MdbResult<()> {
            db.set_dupsort(sort_reverse::<u64>)
        }
    }
    impl MyFirstTable {
        fn table<'db>(db: Database<'db>) -> MdbResult<Table<'db, u64, u64>> {
            try!(Self::setup(&db));
            Ok(Table{db: db, k: PhantomData, v: PhantomData})
        }
    }

    struct MyBlobTable;
    impl TableDef for MyBlobTable {
        fn name() -> &'static str { "my_blob_table" }
        fn flags() -> DbFlags { DbIntKey }
        fn setup(_db: &Database) -> MdbResult<()> { Ok(()) }
    }
    impl MyBlobTable {
        fn table<'db>(db: Database<'db>) -> MdbResult<Table<'db, u64, &[u8]>> {
            try!(Self::setup(&db));
            Ok(Table{db: db, k: PhantomData, v: PhantomData})
        }
    }

    // A Unique(Id64, Id64 DESC) table
    let table_handle = MyFirstTable::open(&env, true).unwrap();
    let blobs_handle = MyBlobTable::open(&env, true).unwrap();

    // prepare database
    {
        let txn = env.new_transaction().unwrap();
        {
            let table = MyFirstTable::table(txn.bind(&table_handle)).unwrap();
            let blobs = MyBlobTable::table(txn.bind(&blobs_handle)).unwrap();

            let big_blob: &[u8] = b"Test";
            blobs.set(&1, &big_blob).unwrap();

            let back: &[u8] = blobs.get(&1).unwrap();
            assert_eq!(&back[..], &b"Test"[..]);

            let key = 1;
            table.set(&key, &200).unwrap();
            table.set(&key, &200).unwrap();
            table.set(&key, &100).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &1).unwrap();
            table.set(&key, &2).unwrap();
            table.set(&key, &3).unwrap();
            table.set(&key, &100).unwrap();
        }
        txn.commit().unwrap();
    }
    // read
    {
        let rdr = env.get_reader().unwrap();

        let table = MyFirstTable::table(rdr.bind(&table_handle)).unwrap();

        let mut cursor = table.new_cursor().unwrap();
        cursor.to_key(&1).unwrap(); //  positions on first item of key
        assert_eq!((1, 200), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 100), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 3), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 2), cursor.get().unwrap());

        cursor.to_next_item().unwrap();
        assert_eq!((1, 1), cursor.get().unwrap());

        assert!(lmdb_not_found!(cursor.to_next_item()));

        let mut cursor = table.new_cursor().unwrap();
        {
            assert!(lmdb_not_found!(cursor.to_item(&2, &3)));

            cursor.to_gte_item(&1, &4).unwrap();
            assert_eq!((1, 3), cursor.get().unwrap());

            cursor.to_next_item().unwrap();
            assert_eq!((1, 2), cursor.get().unwrap());

            cursor.to_next_item().unwrap();
            assert_eq!((1, 1), cursor.get().unwrap());

            assert!(lmdb_not_found!(cursor.to_next_item()));
        }
    }
}
<|endoftext|>


<commit_before>//! A machine learning library for supervised regression trainings
//!
//! This library wants to enable its user to write training algorithms
//! independent of the model trained or the cost function tried to
//! minimize.
//! Consequently its two main traits are currently `Model` and `Cost`.
//! The two submodules `model` and `cost` provide ready to use
//! implementations of said traits.

extern crate num;

use std::iter::Iterator;
use num::{Zero, One, Num};

/// A Model is defines how to predict a target from an input
///
/// A model usually depends on several coefficents whose values
/// are derived using a training algorithm
pub trait Model : Clone{
    /// Input features
    type Input;
    /// Target type
    type Target : Num + Copy + One;

    /// Predicts a target for the inputs based on the internal coefficents
    fn predict(&self, &Self::Input) -> Self::Target;

    /// The number of internal coefficents this model depends on
    fn num_coefficents(&self) -> usize;

    /// Value predict derived by the n-th `coefficent` at `input`
    fn gradient(&self, coefficent : usize, input : &Self::Input) -> Self::Target;

    /// Mutable reference to the n-th `coefficent`
    fn coefficent(& mut self, coefficent : usize) -> & mut Self::Target;
}

/// Cost functions those value is supposed be minimized by the training algorithm
pub trait Cost{

    type Error : Num + Copy;

    /// Value of the cost function derived by the n-th coefficent at x expressed in E(x) and dY(x)/dx
    ///
    /// This method is called by SGD based training algorithm in order to
    /// determine the delta of the coefficents
    fn gradient(&self, error : Self::Error, gradient_error_by_coefficent : Self::Error) -> Self::Error;
}

/// Changes all coefficents of model based on their derivation of the cost function at features
///
/// Can be used to implement stochastic or batch gradient descent
pub fn inert_gradient_descent_step<C, M>(
    cost : &C,
    model : &mut M,
    features : &M::Input,
    truth : M::Target,
    learning_rate : M::Target,
    inertia : M::Target,
    velocity : & mut Vec<M::Target>
)
    where C : Cost, M : Model<Target=C::Error>
{
    let inv_inertia = M::Target::one() - inertia;
    let prediction = model.predict(&features);
    let error = prediction - truth;

    for ci in 0..model.num_coefficents(){

        velocity[ci] = inertia * velocity[ci] - inv_inertia * learning_rate * cost.gradient(error, model.gradient(ci, features));
        *model.coefficent(ci) = *model.coefficent(ci) + velocity[ci];
    }
}

/// Changes all coefficents of model based on their derivation of the cost function at features
///
/// Can be used to implement stochastic or batch gradient descent
pub fn gradient_descent_step<C, M>(cost : &C, model : &mut M, features : &M::Input, truth : M::Target, learning_rate : M::Target)
    where C : Cost, M : Model<Target=C::Error>
{
    let prediction = model.predict(&features);
    let error = prediction - truth;

    for ci in 0..model.num_coefficents(){
        *model.coefficent(ci) = *model.coefficent(ci) - learning_rate * cost.gradient(error, model.gradient(ci, features));
    }
}

/// Trains a model
pub fn stochastic_gradient_descent<C, M, H>(cost : &C, start : M, history : H, learning_rate : M::Target) -> M
    where C : Cost,
    M : Model<Target=C::Error>,
    H : Iterator<Item=(M::Input, M::Target)>
{

    let mut next = start.clone();
    for (features, truth) in history{

        gradient_descent_step(cost, & mut next, &features, truth, learning_rate);
    }

    next
}

/// Trains a model
pub fn inert_stochastic_gradient_descent<C, M, H>(cost : &C, start : M, history : H, learning_rate : M::Target, inertia : M::Target) -> M
    where C : Cost,
    M : Model<Target=C::Error>,
    H : Iterator<Item=(M::Input, M::Target)>
{

    let mut velocity = Vec::new();
    velocity.resize(start.num_coefficents(), M::Target::zero());
    let mut next = start.clone();
    for (features, truth) in history{

        inert_gradient_descent_step(cost, & mut next, &features, truth, learning_rate, inertia, & mut velocity);
    }

    next
}

pub mod model;
pub mod cost;
pub mod linear_algebra;

#[cfg(test)]
mod tests {

    #[test]
    fn estimate_median() {

        use model::Constant;
        use cost::LeastAbsoluteDeviation;
        use gradient_descent_step;

        let features = ();
        let history = [1f64, 3.0, 4.0, 7.0, 8.0, 11.0, 29.0]; //median is seven

        let cost = LeastAbsoluteDeviation{};
        let mut model = Constant{c : 0.0};

        let learning_rate_start = 0.4;
        let learning_rate_stop = 0.001;
        let num_steps = 200;
        let learning_rate_gradient = (learning_rate_start - learning_rate_stop) / (num_steps as f64);

        for (count_step, &truth) in history.iter().cycle().take(num_steps).enumerate(){

            let adapted_learning_rate = learning_rate_stop + learning_rate_gradient * (num_steps - count_step) as f64;
            gradient_descent_step(&cost, & mut model, &features, truth, adapted_learning_rate);
            println!("model: {:?}, learning_rate: {:?}", model, adapted_learning_rate);
        }

        assert!(model.c < 7.1);
        assert!(model.c > 6.9);
    }

    #[test]
    fn estimate_mean() {

        use model::Constant;
        use cost::LeastSquares;
        use gradient_descent_step;

        let features = ();
        let history = [1f64, 3.0, 4.0, 7.0, 8.0, 11.0, 29.0]; //mean is 9

        let cost = LeastSquares{};
        let mut model = Constant{c : 0.0};

        let learning_rate_start = 0.4;
        let learning_rate_stop = 0.001;
        let num_steps = 60;
        let learning_rate_gradient = (learning_rate_start - learning_rate_stop) / (num_steps as f64);

        for (count_step, &truth) in history.iter().cycle().take(num_steps).enumerate(){

            let adapted_learning_rate = learning_rate_stop + learning_rate_gradient * (num_steps - count_step) as f64;
            gradient_descent_step(&cost, & mut model, &features, truth, adapted_learning_rate);
            println!("model: {:?}, learning_rate: {:?}", model, adapted_learning_rate);
        }

        assert!(model.c < 9.1);
        assert!(model.c > 8.9);
    }

    #[test]
    fn linear_stochastic_gradient_descent() {

        use cost::LeastSquares;
        use model::Linear;
        use stochastic_gradient_descent;

        let history = [(0f64, 3f64), (1.0, 4.0), (2.0, 5.0)];

        let start = Linear{m : 0.0, c : 0.0};

        let learning_rate = 0.2;

        let cost = LeastSquares{};
        let model = stochastic_gradient_descent(&cost, start, history.iter().cycle().take(20).cloned(), learning_rate);

        assert!(model.m < 1.1);
        assert!(model.m > 0.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }

    #[test]
    fn linear_stochastic_gradient_descent_iter() {

        use model::Linear;
        use gradient_descent_step;
        use cost::LeastSquares;

        let history = [(0f64, 3f64), (1.0, 4.0), (2.0, 5.0)];

        let cost = LeastSquares{};
        let mut model = Linear{m : 0.0, c : 0.0};

        let learning_rate = 0.2;

        for &(features, truth) in history.iter().cycle().take(20){

            gradient_descent_step(&cost, & mut model, &features, truth, learning_rate);
            println!("model: {:?}", model);
        }

        assert!(model.m < 1.1);
        assert!(model.m > 0.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }

    #[test]
    fn linear_sgd_2d()
    {
        use cost::LeastSquares;
        use model::Linear;
        use inert_stochastic_gradient_descent;

        let history = [([0.0, 7.0], 17.0), ([1.0, 2.0], 8.0), ([2.0, -2.0], 1.0)];

        let start = Linear{m : [0.0, 0.0], c : 0.0};

        let learning_rate = 0.1;

        let cost = LeastSquares{};
        let model = inert_stochastic_gradient_descent(&cost, start, history.iter().cycle().take(15000).cloned(), learning_rate, 0.9);

        println!("{:?}", model);

        assert!(model.m[0] < 1.1);
        assert!(model.m[0] > 0.9);
        assert!(model.m[1] < 2.1);
        assert!(model.m[1] > 1.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }
}
<commit_msg>reorded traits<commit_after>//! A machine learning library for supervised regression trainings
//!
//! This library wants to enable its user to write training algorithms
//! independent of the model trained or the cost function tried to
//! minimize.
//! Consequently its two main traits are currently `Model` and `Cost`.
//! The two submodules `model` and `cost` provide ready to use
//! implementations of said traits.

extern crate num;

use std::iter::Iterator;
use num::{Zero, One, Num};

/// A Model is defines how to predict a target from an input
///
/// A model usually depends on several coefficents whose values
/// are derived using a training algorithm
pub trait Model : Clone{
    /// Input features
    type Input;
    /// Target type
    type Target : Num + One + Copy;

    /// Predicts a target for the inputs based on the internal coefficents
    fn predict(&self, &Self::Input) -> Self::Target;

    /// The number of internal coefficents this model depends on
    fn num_coefficents(&self) -> usize;

    /// Value predict derived by the n-th `coefficent` at `input`
    fn gradient(&self, coefficent : usize, input : &Self::Input) -> Self::Target;

    /// Mutable reference to the n-th `coefficent`
    fn coefficent(& mut self, coefficent : usize) -> & mut Self::Target;
}

/// Cost functions those value is supposed be minimized by the training algorithm
pub trait Cost{

    type Error : Num + Copy;

    /// Value of the cost function derived by the n-th coefficent at x expressed in E(x) and dY(x)/dx
    ///
    /// This method is called by SGD based training algorithm in order to
    /// determine the delta of the coefficents
    fn gradient(&self, error : Self::Error, gradient_error_by_coefficent : Self::Error) -> Self::Error;
}

/// Changes all coefficents of model based on their derivation of the cost function at features
///
/// Can be used to implement stochastic or batch gradient descent
pub fn inert_gradient_descent_step<C, M>(
    cost : &C,
    model : &mut M,
    features : &M::Input,
    truth : M::Target,
    learning_rate : M::Target,
    inertia : M::Target,
    velocity : & mut Vec<M::Target>
)
    where C : Cost, M : Model<Target=C::Error>
{
    let inv_inertia = M::Target::one() - inertia;
    let prediction = model.predict(&features);
    let error = prediction - truth;

    for ci in 0..model.num_coefficents(){

        velocity[ci] = inertia * velocity[ci] - inv_inertia * learning_rate * cost.gradient(error, model.gradient(ci, features));
        *model.coefficent(ci) = *model.coefficent(ci) + velocity[ci];
    }
}

/// Changes all coefficents of model based on their derivation of the cost function at features
///
/// Can be used to implement stochastic or batch gradient descent
pub fn gradient_descent_step<C, M>(cost : &C, model : &mut M, features : &M::Input, truth : M::Target, learning_rate : M::Target)
    where C : Cost, M : Model<Target=C::Error>
{
    let prediction = model.predict(&features);
    let error = prediction - truth;

    for ci in 0..model.num_coefficents(){
        *model.coefficent(ci) = *model.coefficent(ci) - learning_rate * cost.gradient(error, model.gradient(ci, features));
    }
}

/// Trains a model
pub fn stochastic_gradient_descent<C, M, H>(cost : &C, start : M, history : H, learning_rate : M::Target) -> M
    where C : Cost,
    M : Model<Target=C::Error>,
    H : Iterator<Item=(M::Input, M::Target)>
{

    let mut next = start.clone();
    for (features, truth) in history{

        gradient_descent_step(cost, & mut next, &features, truth, learning_rate);
    }

    next
}

/// Trains a model
pub fn inert_stochastic_gradient_descent<C, M, H>(
    cost : &C,
    start : M,
    history : H,
    learning_rate : M::Target,
    inertia : M::Target
) -> M
    where C : Cost,
    M : Model<Target=C::Error>,
    H : Iterator<Item=(M::Input, M::Target)>
{

    let mut velocity = Vec::new();
    velocity.resize(start.num_coefficents(), M::Target::zero());
    let mut next = start.clone();
    for (features, truth) in history{

        inert_gradient_descent_step(cost, & mut next, &features, truth, learning_rate, inertia, & mut velocity);
    }

    next
}

pub mod model;
pub mod cost;
pub mod linear_algebra;

#[cfg(test)]
mod tests {

    #[test]
    fn estimate_median() {

        use model::Constant;
        use cost::LeastAbsoluteDeviation;
        use gradient_descent_step;

        let features = ();
        let history = [1.0, 3.0, 4.0, 7.0, 8.0, 11.0, 29.0]; //median is seven

        let cost = LeastAbsoluteDeviation{};
        let mut model = Constant{c : 0.0};

        let learning_rate_start = 0.4;
        let learning_rate_stop = 0.001;
        let num_steps = 200;
        let learning_rate_gradient = (learning_rate_start - learning_rate_stop) / (num_steps as f64);

        for (count_step, &truth) in history.iter().cycle().take(num_steps).enumerate(){

            let adapted_learning_rate = learning_rate_stop + learning_rate_gradient * (num_steps - count_step) as f64;
            gradient_descent_step(&cost, & mut model, &features, truth, adapted_learning_rate);
            println!("model: {:?}, learning_rate: {:?}", model, adapted_learning_rate);
        }

        assert!(model.c < 7.1);
        assert!(model.c > 6.9);
    }

    #[test]
    fn estimate_mean() {

        use model::Constant;
        use cost::LeastSquares;
        use gradient_descent_step;

        let features = ();
        let history = [1f64, 3.0, 4.0, 7.0, 8.0, 11.0, 29.0]; //mean is 9

        let cost = LeastSquares{};
        let mut model = Constant{c : 0.0};

        let learning_rate_start = 0.4;
        let learning_rate_stop = 0.001;
        let num_steps = 60;
        let learning_rate_gradient = (learning_rate_start - learning_rate_stop) / (num_steps as f64);

        for (count_step, &truth) in history.iter().cycle().take(num_steps).enumerate(){

            let adapted_learning_rate = learning_rate_stop + learning_rate_gradient * (num_steps - count_step) as f64;
            gradient_descent_step(&cost, & mut model, &features, truth, adapted_learning_rate);
            println!("model: {:?}, learning_rate: {:?}", model, adapted_learning_rate);
        }

        assert!(model.c < 9.1);
        assert!(model.c > 8.9);
    }

    #[test]
    fn linear_stochastic_gradient_descent() {

        use cost::LeastSquares;
        use model::Linear;
        use stochastic_gradient_descent;

        let history = [(0f64, 3f64), (1.0, 4.0), (2.0, 5.0)];

        let start = Linear{m : 0.0, c : 0.0};

        let learning_rate = 0.2;

        let cost = LeastSquares{};
        let model = stochastic_gradient_descent(&cost, start, history.iter().cycle().take(20).cloned(), learning_rate);

        assert!(model.m < 1.1);
        assert!(model.m > 0.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }

    #[test]
    fn linear_stochastic_gradient_descent_iter() {

        use model::Linear;
        use gradient_descent_step;
        use cost::LeastSquares;

        let history = [(0f64, 3f64), (1.0, 4.0), (2.0, 5.0)];

        let cost = LeastSquares{};
        let mut model = Linear{m : 0.0, c : 0.0};

        let learning_rate = 0.2;

        for &(features, truth) in history.iter().cycle().take(20){

            gradient_descent_step(&cost, & mut model, &features, truth, learning_rate);
            println!("model: {:?}", model);
        }

        assert!(model.m < 1.1);
        assert!(model.m > 0.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }

    #[test]
    fn linear_sgd_2d()
    {
        use cost::LeastSquares;
        use model::Linear;
        use inert_stochastic_gradient_descent;

        let history = [([0.0, 7.0], 17.0), ([1.0, 2.0], 8.0), ([2.0, -2.0], 1.0)];

        let start = Linear{m : [0.0, 0.0], c : 0.0};

        let learning_rate = 0.1;

        let cost = LeastSquares{};
        let model = inert_stochastic_gradient_descent(&cost, start, history.iter().cycle().take(15000).cloned(), learning_rate, 0.9);

        println!("{:?}", model);

        assert!(model.m[0] < 1.1);
        assert!(model.m[0] > 0.9);
        assert!(model.m[1] < 2.1);
        assert!(model.m[1] > 1.9);
        assert!(model.c < 3.1);
        assert!(model.c > 2.9);
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate matches;
pub extern crate mime;

use std::io;

pub enum DataUrlError {
    NotADataUrl,
    NoComma,
    InvalidMimeType(mime::FromStrError),
}

pub struct DataUrl<'a> {
    mime_type: mime::Mime,
    base64: bool,
    encoded_body_plus_fragment: &'a str,
}

impl<'a> DataUrl<'a> {
    /// <https://fetch.spec.whatwg.org/#data-url-processor>
    /// but starting from a string rather than a Url, to avoid extra string copies.
    pub fn process(input: &'a str) -> Result<Self, DataUrlError> {
        use DataUrlError::*;

        let after_colon = pretend_parse_data_url(input).ok_or(NotADataUrl)?;
        let comma = after_colon.bytes().position(|byte| byte == b',').ok_or(NoComma)?;

        let (mime_type, base64) = parse_header(&after_colon[..comma]).map_err(InvalidMimeType)?;
        let encoded_body_plus_fragment = &after_colon[comma + 1..];

        Ok(DataUrl { mime_type, base64, encoded_body_plus_fragment })
    }

    pub fn mime_type(&self) -> &mime::Mime {
        &self.mime_type
    }

    /// Streaming-decode the data URL’s body to `sink`.
    ///
    /// Errors while writing to the sink are propagated.
    /// Invalid base64 causes an error with `e.kind() == ErrorKind::InvalidData`.
    /// When decoding without error, the URL’s fragment identifier is returned if it has one.
    ///
    /// The fragment identifier is represented as in the origin input.
    /// It needs to be either percent-encoded to obtain the same string as in a parsed URL,
    /// or percent-decoded to interpret it as text.
    pub fn decode_body<W>(&self, sink: W) -> io::Result<Option<&'a str>>
        where W: io::Write
    {
        if self.base64 {
            decode_with_base64(self.encoded_body_plus_fragment, sink)
        } else {
            decode_without_base64(self.encoded_body_plus_fragment, sink)
        }
    }

    pub fn decode_body_to_vec(&self) -> Result<(Vec<u8>, Option<&str>), ()> {
        let mut sink = Vec::new();
        match self.decode_body(&mut sink) {
            Ok(url_fragment) => {
                Ok((sink, url_fragment))
            }
            Err(e) => {
                // Vec::write_all never returns an error
                debug_assert!(e.kind() == io::ErrorKind::InvalidData && self.base64);

                Err(())
            }
        }
    }
}

macro_rules! require {
    ($condition: expr) => {
        if !$condition {
            return None
        }
    }
}

/// Similar to <https://url.spec.whatwg.org/#concept-basic-url-parser>
/// followed by <https://url.spec.whatwg.org/#concept-url-serializer>
///
/// * `None`: not a data URL.
///
/// * `Some(s)`: sort of the result of serialization, except:
///
///   - `data:` prefix removed
///   - The fragment is included
///   - Other components are **not** UTF-8 percent-encoded
///   - ASCII tabs and newlines in the middle are **not** removed
fn pretend_parse_data_url(input: &str) -> Option<&str> {
    // Trim C0 control or space
    let left_trimmed = input.trim_left_matches(|ch| ch <= ' ');

    let mut bytes = left_trimmed.bytes();
    {
        // Ignore ASCII tabs or newlines
        let mut iter = bytes.by_ref().filter(|&byte| !matches!(byte, b'\t' | b'\n' | b'\r'));
        require!(iter.next()?.to_ascii_lowercase() == b'd');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()?.to_ascii_lowercase() == b't');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()? == b':');
    }
    let bytes_consumed = left_trimmed.len() - bytes.len();
    let after_colon = &left_trimmed[bytes_consumed..];

    // Trim C0 control or space
    Some(after_colon.trim_right_matches(|ch| ch <= ' '))
}

fn parse_header(from_colon_to_comma: &str) -> Result<(mime::Mime, bool), mime::FromStrError> {
    let input = from_colon_to_comma.chars()
        .filter(|&c| !matches!(c, '\t' | '\n' | '\r'))  // Removed by the URL parser
        .collect::<String>();

    let input = input.trim_matches(' ');

    let (input, base64) = match without_base64_suffix(input) {
        Some(s) => (s, true),
        None => (input, false),
    };

    // FIXME: does Mime::from_str match the MIME Sniffing Standard’s parsing algorithm?
    // <https://mimesniff.spec.whatwg.org/#parse-a-mime-type>
    Ok((input.parse()?, base64))
}

/// None: no base64 suffix
fn without_base64_suffix(s: &str) -> Option<&str> {
    remove_suffix(
        remove_suffix(s, "base64", str::eq_ignore_ascii_case)?
            .trim_right_matches(' '),
        ";", str::eq
    )
}

fn remove_suffix<'a, Eq>(haystack: &'a str, needle: &str, eq: Eq) -> Option<&'a str>
    where Eq: Fn(&str, &str) -> bool
{
    let start_index = haystack.len().checked_sub(needle.len())?;
    let (before, after) = haystack.split_at(start_index);
    if eq(after, needle) {
        Some(before)
    } else {
        None
    }
}

/// This is <https://url.spec.whatwg.org/#string-percent-decode> while also:
///
/// * Ignoring ASCII tab or newlines
/// * Stopping at the first '#' (which indicates the start of the fragment)
///
/// Anything that would have been UTF-8 percent-encoded by the URL parser
/// would be percent-decoded here.
/// We skip that round-trip and pass it through unchanged.
fn decode_without_base64<W>(encoded_body_plus_fragment: &str, mut sink: W)
                            -> io::Result<Option<&str>>
    where W: io::Write
{
    let bytes = encoded_body_plus_fragment.as_bytes();
    let mut slice_start = 0;
    for (i, &byte) in bytes.iter().enumerate() {
        // We only need to look for 5 different "special" byte values.
        // For everything else we make slices as large as possible, borrowing the input,
        // in order to make fewer write_all() calls.
        if matches!(byte, b'%' | b'#' | b'\t' | b'\n' | b'\r') {
            // Write everything (if anything) "non-special" we’ve accumulated
            // before this special byte
            if i > slice_start {
                sink.write_all(&bytes[slice_start..i])?;
            }
            // Then deal with the special byte.
            match byte {
                b'%' => {
                    let l = bytes.get(i + 2).and_then(|&b| (b as char).to_digit(16));
                    let h = bytes.get(i + 1).and_then(|&b| (b as char).to_digit(16));
                    if let (Some(h), Some(l)) = (h, l) {
                        // '%' followed by two ASCII hex digits
                        let one_byte = h as u8 * 0x10 + l as u8;
                        sink.write_all(&[one_byte])?;
                        slice_start = i + 3;
                    } else {
                        // Do nothing. Leave slice_start unchanged.
                        // The % sign will be part of the next slice.
                    }
                }

                b'#' => {
                    let fragment_start = i + 1;
                    return Ok(Some(&encoded_body_plus_fragment[fragment_start..]))
                }

                // Ignore over '\t' | '\n' | '\r'
                _ => slice_start = i + 1
            }
        }
    }
    sink.write_all(&bytes[slice_start..])?;
    Ok(None)
}

/// `decode_without_base64()` composed with
/// <https://infra.spec.whatwg.org/#isomorphic-decode> composed with
/// <https://infra.spec.whatwg.org/#forgiving-base64-decode>.
fn decode_with_base64<W>(encoded_body_plus_fragment: &str, sink: W) -> io::Result<Option<&str>>
    where W: io::Write
{
    let mut decoder = Base64Decoder {
        sink,
        bit_buffer: 0,
        buffer_bit_length: 0,
        padding_symbols: 0,
    };
    let fragment = decode_without_base64(encoded_body_plus_fragment, &mut decoder)?;
    match (decoder.buffer_bit_length, decoder.padding_symbols) {
        (0, 0) => {
            // A multiple of four of alphabet symbols, and nothing else.
        }
        (12, 2) | (12, 0) => {
            // A multiple of four of alphabet symbols, followed by two more symbols,
            // optionally followed by two padding characters (which make a total multiple of four).
            let byte_buffer = [
                (decoder.bit_buffer >> 4) as u8,
            ];
            decoder.sink.write_all(&byte_buffer)?;
        }
        (18, 1) | (18, 0) => {
            // A multiple of four of alphabet symbols, followed by three more symbols,
            // optionally followed by one padding character (which make a total multiple of four).
            let byte_buffer = [
                (decoder.bit_buffer >> 10) as u8,
                (decoder.bit_buffer >> 2) as u8,
            ];
            decoder.sink.write_all(&byte_buffer)?;
        }
        _ => {
            // No other combination is acceptable
            Err(io::ErrorKind::InvalidData)?
        }
    }
    Ok(fragment)
}

struct Base64Decoder<W> {
    sink: W,
    bit_buffer: u32,
    buffer_bit_length: u8,
    padding_symbols: u8,
}

impl<W> io::Write for Base64Decoder<W> where W: io::Write {
    fn write(&mut self, _buf: &[u8]) -> io::Result<usize> { unimplemented!() }
    fn flush(&mut self) -> io::Result<()> { unimplemented!() }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        for &byte in buf.iter() {
            let value = BASE64_DECODE_TABLE[byte as usize];
            if value < 0 {
                // A character that’s not part of the alphabet

                // Remove ASCII whitespace
                // '\t' | '\n' | '\r' was already filtered by decode_without_base64()
                if byte == b' ' || byte == b'\x0C' {
                    continue
                }

                if byte == b'=' {
                    self.padding_symbols = self.padding_symbols.saturating_add(8);
                    continue
                }

                Err(io::ErrorKind::InvalidData)?
            }
            if self.padding_symbols > 0 {
                // Alphabet symbols after padding
                Err(io::ErrorKind::InvalidData)?
            }
            self.bit_buffer <<= 6;
            self.bit_buffer |= value as u32;
            if self.buffer_bit_length < 24 {
                self.buffer_bit_length += 6;
            } else {
                // We’ve accumulated four times 6 bits, which equals three times 8 bits.
                let byte_buffer = [
                    (self.bit_buffer >> 16) as u8,
                    (self.bit_buffer >> 8) as u8,
                    self.bit_buffer as u8,
                ];
                self.sink.write_all(&byte_buffer)?;
                self.buffer_bit_length = 0;
                // No need to reset self.bit_buffer,
                // since next time we’re only gonna read relevant bits.
            }
        }
        Ok(())
    }
}

/// Generated by `make_base64_decode_table.py` based on "Table 1: The Base 64 Alphabet"
/// at <https://tools.ietf.org/html/rfc4648#section-4>
///
/// Array indices are the byte value of symbols.
/// Array values are their positions in the base64 alphabet,
/// or -1 for symbols not in the alphabet.
/// The position contributes 6 bits to the decoded bytes.
const BASE64_DECODE_TABLE: [i8; 256] = [
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63,
    52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1,
    -1,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
    15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1,
    -1, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
];
<commit_msg>Invalid MIME types default to "text/plain;charset=US-ASCII"<commit_after>#[macro_use] extern crate matches;
pub extern crate mime;

use std::io;

pub enum DataUrlError {
    NotADataUrl,
    NoComma,
}

pub struct DataUrl<'a> {
    mime_type: mime::Mime,
    base64: bool,
    encoded_body_plus_fragment: &'a str,
}

impl<'a> DataUrl<'a> {
    /// <https://fetch.spec.whatwg.org/#data-url-processor>
    /// but starting from a string rather than a Url, to avoid extra string copies.
    pub fn process(input: &'a str) -> Result<Self, DataUrlError> {
        use DataUrlError::*;

        let after_colon = pretend_parse_data_url(input).ok_or(NotADataUrl)?;
        let comma = after_colon.bytes().position(|byte| byte == b',').ok_or(NoComma)?;

        let (mime_type, base64) = parse_header(&after_colon[..comma]);
        let encoded_body_plus_fragment = &after_colon[comma + 1..];

        Ok(DataUrl { mime_type, base64, encoded_body_plus_fragment })
    }

    pub fn mime_type(&self) -> &mime::Mime {
        &self.mime_type
    }

    /// Streaming-decode the data URL’s body to `sink`.
    ///
    /// Errors while writing to the sink are propagated.
    /// Invalid base64 causes an error with `e.kind() == ErrorKind::InvalidData`.
    /// When decoding without error, the URL’s fragment identifier is returned if it has one.
    ///
    /// The fragment identifier is represented as in the origin input.
    /// It needs to be either percent-encoded to obtain the same string as in a parsed URL,
    /// or percent-decoded to interpret it as text.
    pub fn decode_body<W>(&self, sink: W) -> io::Result<Option<&'a str>>
        where W: io::Write
    {
        if self.base64 {
            decode_with_base64(self.encoded_body_plus_fragment, sink)
        } else {
            decode_without_base64(self.encoded_body_plus_fragment, sink)
        }
    }

    pub fn decode_body_to_vec(&self) -> Result<(Vec<u8>, Option<&str>), ()> {
        let mut sink = Vec::new();
        match self.decode_body(&mut sink) {
            Ok(url_fragment) => {
                Ok((sink, url_fragment))
            }
            Err(e) => {
                // Vec::write_all never returns an error
                debug_assert!(e.kind() == io::ErrorKind::InvalidData && self.base64);

                Err(())
            }
        }
    }
}

macro_rules! require {
    ($condition: expr) => {
        if !$condition {
            return None
        }
    }
}

/// Similar to <https://url.spec.whatwg.org/#concept-basic-url-parser>
/// followed by <https://url.spec.whatwg.org/#concept-url-serializer>
///
/// * `None`: not a data URL.
///
/// * `Some(s)`: sort of the result of serialization, except:
///
///   - `data:` prefix removed
///   - The fragment is included
///   - Other components are **not** UTF-8 percent-encoded
///   - ASCII tabs and newlines in the middle are **not** removed
fn pretend_parse_data_url(input: &str) -> Option<&str> {
    // Trim C0 control or space
    let left_trimmed = input.trim_left_matches(|ch| ch <= ' ');

    let mut bytes = left_trimmed.bytes();
    {
        // Ignore ASCII tabs or newlines
        let mut iter = bytes.by_ref().filter(|&byte| !matches!(byte, b'\t' | b'\n' | b'\r'));
        require!(iter.next()?.to_ascii_lowercase() == b'd');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()?.to_ascii_lowercase() == b't');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()? == b':');
    }
    let bytes_consumed = left_trimmed.len() - bytes.len();
    let after_colon = &left_trimmed[bytes_consumed..];

    // Trim C0 control or space
    Some(after_colon.trim_right_matches(|ch| ch <= ' '))
}

fn parse_header(from_colon_to_comma: &str) -> (mime::Mime, bool) {
    let input = from_colon_to_comma.chars()
        .filter(|&c| !matches!(c, '\t' | '\n' | '\r'))  // Removed by the URL parser
        .collect::<String>();

    let input = input.trim_matches(' ');

    let (input, base64) = match without_base64_suffix(input) {
        Some(s) => (s, true),
        None => (input, false),
    };

    // FIXME: does Mime::from_str match the MIME Sniffing Standard’s parsing algorithm?
    // <https://mimesniff.spec.whatwg.org/#parse-a-mime-type>
    let mime_type = input.parse()
        .unwrap_or_else(|_| "text/plain;charset=US-ASCII".parse().unwrap());

    (mime_type, base64)
}

/// None: no base64 suffix
fn without_base64_suffix(s: &str) -> Option<&str> {
    remove_suffix(
        remove_suffix(s, "base64", str::eq_ignore_ascii_case)?
            .trim_right_matches(' '),
        ";", str::eq
    )
}

fn remove_suffix<'a, Eq>(haystack: &'a str, needle: &str, eq: Eq) -> Option<&'a str>
    where Eq: Fn(&str, &str) -> bool
{
    let start_index = haystack.len().checked_sub(needle.len())?;
    let (before, after) = haystack.split_at(start_index);
    if eq(after, needle) {
        Some(before)
    } else {
        None
    }
}

/// This is <https://url.spec.whatwg.org/#string-percent-decode> while also:
///
/// * Ignoring ASCII tab or newlines
/// * Stopping at the first '#' (which indicates the start of the fragment)
///
/// Anything that would have been UTF-8 percent-encoded by the URL parser
/// would be percent-decoded here.
/// We skip that round-trip and pass it through unchanged.
fn decode_without_base64<W>(encoded_body_plus_fragment: &str, mut sink: W)
                            -> io::Result<Option<&str>>
    where W: io::Write
{
    let bytes = encoded_body_plus_fragment.as_bytes();
    let mut slice_start = 0;
    for (i, &byte) in bytes.iter().enumerate() {
        // We only need to look for 5 different "special" byte values.
        // For everything else we make slices as large as possible, borrowing the input,
        // in order to make fewer write_all() calls.
        if matches!(byte, b'%' | b'#' | b'\t' | b'\n' | b'\r') {
            // Write everything (if anything) "non-special" we’ve accumulated
            // before this special byte
            if i > slice_start {
                sink.write_all(&bytes[slice_start..i])?;
            }
            // Then deal with the special byte.
            match byte {
                b'%' => {
                    let l = bytes.get(i + 2).and_then(|&b| (b as char).to_digit(16));
                    let h = bytes.get(i + 1).and_then(|&b| (b as char).to_digit(16));
                    if let (Some(h), Some(l)) = (h, l) {
                        // '%' followed by two ASCII hex digits
                        let one_byte = h as u8 * 0x10 + l as u8;
                        sink.write_all(&[one_byte])?;
                        slice_start = i + 3;
                    } else {
                        // Do nothing. Leave slice_start unchanged.
                        // The % sign will be part of the next slice.
                    }
                }

                b'#' => {
                    let fragment_start = i + 1;
                    return Ok(Some(&encoded_body_plus_fragment[fragment_start..]))
                }

                // Ignore over '\t' | '\n' | '\r'
                _ => slice_start = i + 1
            }
        }
    }
    sink.write_all(&bytes[slice_start..])?;
    Ok(None)
}

/// `decode_without_base64()` composed with
/// <https://infra.spec.whatwg.org/#isomorphic-decode> composed with
/// <https://infra.spec.whatwg.org/#forgiving-base64-decode>.
fn decode_with_base64<W>(encoded_body_plus_fragment: &str, sink: W) -> io::Result<Option<&str>>
    where W: io::Write
{
    let mut decoder = Base64Decoder {
        sink,
        bit_buffer: 0,
        buffer_bit_length: 0,
        padding_symbols: 0,
    };
    let fragment = decode_without_base64(encoded_body_plus_fragment, &mut decoder)?;
    match (decoder.buffer_bit_length, decoder.padding_symbols) {
        (0, 0) => {
            // A multiple of four of alphabet symbols, and nothing else.
        }
        (12, 2) | (12, 0) => {
            // A multiple of four of alphabet symbols, followed by two more symbols,
            // optionally followed by two padding characters (which make a total multiple of four).
            let byte_buffer = [
                (decoder.bit_buffer >> 4) as u8,
            ];
            decoder.sink.write_all(&byte_buffer)?;
        }
        (18, 1) | (18, 0) => {
            // A multiple of four of alphabet symbols, followed by three more symbols,
            // optionally followed by one padding character (which make a total multiple of four).
            let byte_buffer = [
                (decoder.bit_buffer >> 10) as u8,
                (decoder.bit_buffer >> 2) as u8,
            ];
            decoder.sink.write_all(&byte_buffer)?;
        }
        _ => {
            // No other combination is acceptable
            Err(io::ErrorKind::InvalidData)?
        }
    }
    Ok(fragment)
}

struct Base64Decoder<W> {
    sink: W,
    bit_buffer: u32,
    buffer_bit_length: u8,
    padding_symbols: u8,
}

impl<W> io::Write for Base64Decoder<W> where W: io::Write {
    fn write(&mut self, _buf: &[u8]) -> io::Result<usize> { unimplemented!() }
    fn flush(&mut self) -> io::Result<()> { unimplemented!() }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        for &byte in buf.iter() {
            let value = BASE64_DECODE_TABLE[byte as usize];
            if value < 0 {
                // A character that’s not part of the alphabet

                // Remove ASCII whitespace
                // '\t' | '\n' | '\r' was already filtered by decode_without_base64()
                if byte == b' ' || byte == b'\x0C' {
                    continue
                }

                if byte == b'=' {
                    self.padding_symbols = self.padding_symbols.saturating_add(8);
                    continue
                }

                Err(io::ErrorKind::InvalidData)?
            }
            if self.padding_symbols > 0 {
                // Alphabet symbols after padding
                Err(io::ErrorKind::InvalidData)?
            }
            self.bit_buffer <<= 6;
            self.bit_buffer |= value as u32;
            if self.buffer_bit_length < 24 {
                self.buffer_bit_length += 6;
            } else {
                // We’ve accumulated four times 6 bits, which equals three times 8 bits.
                let byte_buffer = [
                    (self.bit_buffer >> 16) as u8,
                    (self.bit_buffer >> 8) as u8,
                    self.bit_buffer as u8,
                ];
                self.sink.write_all(&byte_buffer)?;
                self.buffer_bit_length = 0;
                // No need to reset self.bit_buffer,
                // since next time we’re only gonna read relevant bits.
            }
        }
        Ok(())
    }
}

/// Generated by `make_base64_decode_table.py` based on "Table 1: The Base 64 Alphabet"
/// at <https://tools.ietf.org/html/rfc4648#section-4>
///
/// Array indices are the byte value of symbols.
/// Array values are their positions in the base64 alphabet,
/// or -1 for symbols not in the alphabet.
/// The position contributes 6 bits to the decoded bytes.
const BASE64_DECODE_TABLE: [i8; 256] = [
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63,
    52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1,
    -1,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
    15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1,
    -1, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
];
<|endoftext|>