<commit_before>use std::io::net::ip::{Port, IpAddr};

use router::Router;
use middleware::{ MiddlewareStack, Middleware, Action };
use into_middleware::IntoMiddleware;
use into_error_handler::IntoErrorHandler;
use nickel_error::{ NickelError, ErrorWithStatusCode };
use server::Server;

use http::method::{ Method, Get, Post, Put, Delete };
use http::status::NotFound;
use request::Request;
use response::Response;

//pre defined middleware
use json_body_parser::JsonBodyParser;
use query_string::QueryStringParser;
use default_error_handler::DefaultErrorHandler;

/// Nickel is the application object. It's the surface that
/// holds all public APIs.

pub struct Nickel{
    middleware_stack: MiddlewareStack,
    server: Option<Server>,
}

impl Nickel {

    /// In order to use Nickels API one first has to create an instance.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut server = Nickel::new();
    /// ```
    pub fn new() -> Nickel {
        let mut middleware_stack = MiddlewareStack::new();

        // Hook up the default error handler by default. Users are
        // free to cancel it out from their custom error handler if
        // they don't like the default behaviour.
        middleware_stack.add_error_handler(DefaultErrorHandler);

        Nickel {
            middleware_stack: middleware_stack,
            server: None
        }
    }

    /// Registers a middleware handler which will be invoked among other middleware
    /// handlers before each request. Middleware can be stacked and is invoked in the
    /// same order it was registered.
    ///
    /// A middleware handler is nearly identical to a regular route handler with the only
    /// difference that it expects a return value of boolean. That is to indicate whether
    /// other middleware handler (if any) further down the stack should continue or if the
    /// middleware invocation should be stopped after the current handler.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn logger (req: &Request, res: &mut Response) -> Result<Action, NickelError>{
    ///     println!("logging request: {}", req.origin.request_uri);
    ///     Ok(Continue)
    /// }
    /// ```
    pub fn utilize<T: Middleware>(&mut self, handler: T){
        self.middleware_stack.add_middleware(handler);
    }

    /// Registers a handler to be used for a specific GET request.
    /// Handlers are assigned to paths and paths are allowed to contain
    /// variables and wildcards. A handler added through this API will
    /// be attached to the default router. Consider creating the router
    /// middleware manually for advanced functionality.
    ///
    /// # Example without variables and wildcards
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user");
    /// };
    /// server.get("/user", handler);
    /// ```
    /// # Example with variables
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     let text = format!("This is user: {}", request.params.get(&"userid".to_string()));
    ///     response.send(text.as_slice());
    /// };
    /// server.get("/user/:userid", handler);
    /// ```
    /// # Example with simple wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 but not /user/extended/list/4711");
    /// };
    /// server.get("/user/*/:userid", handler);
    /// ```
    /// # Example with double wildcard
    ///
    /// ```{rust,ignore}
    /// fn handler (request: Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 and also /user/extended/list/4711");
    /// };
    /// server.get("/user/**/:userid", handler);
    /// ```
    pub fn get(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Get, uri, handler);
    }

    /// Registers a handler to be used for a specific POST request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/post/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.post("/a/post/request", handler);
    /// ```
    pub fn post(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Post, uri, handler);
    }

    /// Registers a handler to be used for a specific PUT request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/put/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.put("/a/put/request", handler);
    /// ```
    pub fn put(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Put, uri, handler);
    }

    /// Registers a handler to be used for a specific DELETE request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a DELETE request to /a/delete/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.delete("/a/delete/request", handler);
    /// ```
    pub fn delete(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Delete, uri, handler);
    }

    fn register_route_with_new_router(&mut self, method: Method, uri: &str, handler: fn(request: &Request, response: &mut Response)) {
        let mut router = Router::new();
        router.add_route(method, String::from_str(uri), handler);
        self.utilize(router);
    }

    /// Registers an error handler which will be invoked among other error handler
    /// as soon as any regular handler returned an error
    ///
    /// A error handler is nearly identical to a regular middleware handler with the only
    /// difference that it takes an additional error parameter or type `NickelError.
    ///
    /// # Example
    ///
    /// ```{rust}
    /// # extern crate http;
    /// # extern crate nickel;
    /// # fn main() {
    /// use nickel::{Nickel, Request, Response, Action, Continue, Halt};
    /// use nickel::{NickelError, ErrorWithStatusCode, get_media_type};
    /// use http::status::NotFound;
    ///
    /// fn error_handler(err: &NickelError, req: &Request, response: &mut Response)
    ///                  -> Result<Action, NickelError>{
    ///    match err.kind {
    ///        ErrorWithStatusCode(NotFound) => {
    ///            response.origin.headers.content_type = get_media_type("html");
    ///            response.origin.status = NotFound;
    ///            response.send("<h1>Call the police!<h1>");
    ///            Ok(Halt)
    ///        },
    ///        _ => Ok(Continue)
    ///    }
    /// }
    ///
    /// let mut server = Nickel::new();
    /// server.handle_error(error_handler)
    /// # }
    /// ```
    pub fn handle_error(&mut self, handler: fn(err: &NickelError,
                                               req: &Request,
                                               res: &mut Response)
                                            -> Result<Action, NickelError>){
        let handler = IntoErrorHandler::from_fn(handler);
        self.middleware_stack.add_error_handler(handler);
    }

    /// Create a new middleware to serve as a router.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    ///
    /// let mut server = Nickel::new();
    /// let mut router = Nickel::router();
    ///
    /// fn foo_handler(request: &Request, response: &mut Response) {
    ///     response.send("Hi from /foo");
    /// };
    ///
    /// router.get("/foo", foo_handler);
    /// server.utilize(router);
    /// ```
    pub fn router() -> Router {
        Router::new()
    }

    /// Create a new middleware to parse JSON bodies.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// # #![feature(phase)]
    /// # #[phase(plugin)] extern crate nickel_macros;
    /// # extern crate nickel;
    /// # extern crate serialize;
    /// # use nickel::{Nickel, Request, Response};
    /// use nickel::JsonBody;
    ///
    /// # fn main() {
    /// #[deriving(Decodable, Encodable)]
    /// struct Person {
    ///     first_name: String,
    ///     last_name:  String,
    /// }
    ///
    /// let router = router! {
    ///     post "/a/post/request" => |request, response| {
    ///         let person = request.json_as::<Person>().unwrap();
    ///         let text = format!("Hello {} {}", person.first_name, person.last_name);
    ///         response.send(text);
    ///     }
    /// };
    ///
    /// let mut server = Nickel::new();
    /// // It is currently a requirement that the json_body_parser middleware
    /// // is added before any routes that require it.
    /// server.utilize(Nickel::json_body_parser());
    /// server.utilize(router);
    /// # }
    /// ```
    pub fn json_body_parser() -> JsonBodyParser {
        JsonBodyParser
    }

    /// Create a new middleware to parse the query string.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// # #![feature(phase)]
    /// # #[phase(plugin)] extern crate nickel_macros;
    /// # extern crate nickel;
    /// # use nickel::{Nickel, Request, Response};
    /// use nickel::QueryString;
    /// # fn main() {
    /// let router = router! {
    ///     get "/a/get/request" => |request, response| {
    ///         let foo = request.query("foo", "this is the default value, if foo is not present!");
    ///         response.send(foo[0].as_slice());
    ///     }
    /// };
    ///
    /// let mut server = Nickel::new();
    /// // It is currently a requirement that the query_string middleware
    /// // is added before any routes that require it.
    /// server.utilize(Nickel::query_string());
    /// server.utilize(router);
    /// # }
    /// ```
    pub fn query_string() -> QueryStringParser {
        QueryStringParser
    }

    /// Bind and listen for connections on the given host and port
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut server = Nickel::new();
    /// server.listen(Ipv4Addr(127, 0, 0, 1), 6767);
    /// ```
    pub fn listen(mut self, ip: IpAddr, port: Port) {
        fn not_found_handler(_request: &Request, _response: &mut Response) -> Result<Action, NickelError> {
            Err(NickelError::new("File Not Found", ErrorWithStatusCode(NotFound)))
        }

        self.middleware_stack.add_middleware(IntoMiddleware::from_fn(not_found_handler));
        self.server = Some(Server::new(self.middleware_stack, ip, port));
        self.server.unwrap().serve();
    }
}
<commit_msg>docs(nickel): fix docs for get<commit_after>use std::io::net::ip::{Port, IpAddr};

use router::Router;
use middleware::{ MiddlewareStack, Middleware, Action };
use into_middleware::IntoMiddleware;
use into_error_handler::IntoErrorHandler;
use nickel_error::{ NickelError, ErrorWithStatusCode };
use server::Server;

use http::method::{ Method, Get, Post, Put, Delete };
use http::status::NotFound;
use request::Request;
use response::Response;

//pre defined middleware
use json_body_parser::JsonBodyParser;
use query_string::QueryStringParser;
use default_error_handler::DefaultErrorHandler;

/// Nickel is the application object. It's the surface that
/// holds all public APIs.

pub struct Nickel{
    middleware_stack: MiddlewareStack,
    server: Option<Server>,
}

impl Nickel {

    /// In order to use Nickels API one first has to create an instance.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut server = Nickel::new();
    /// ```
    pub fn new() -> Nickel {
        let mut middleware_stack = MiddlewareStack::new();

        // Hook up the default error handler by default. Users are
        // free to cancel it out from their custom error handler if
        // they don't like the default behaviour.
        middleware_stack.add_error_handler(DefaultErrorHandler);

        Nickel {
            middleware_stack: middleware_stack,
            server: None
        }
    }

    /// Registers a middleware handler which will be invoked among other middleware
    /// handlers before each request. Middleware can be stacked and is invoked in the
    /// same order it was registered.
    ///
    /// A middleware handler is nearly identical to a regular route handler with the only
    /// difference that it expects a return value of boolean. That is to indicate whether
    /// other middleware handler (if any) further down the stack should continue or if the
    /// middleware invocation should be stopped after the current handler.
    ///
    /// # Example
    ///
    /// ```{rust,ignore}
    /// fn logger (req: &Request, res: &mut Response) -> Result<Action, NickelError>{
    ///     println!("logging request: {}", req.origin.request_uri);
    ///     Ok(Continue)
    /// }
    /// ```
    pub fn utilize<T: Middleware>(&mut self, handler: T){
        self.middleware_stack.add_middleware(handler);
    }

    /// Registers a handler to be used for a specific GET request.
    /// Handlers are assigned to paths and paths are allowed to contain
    /// variables and wildcards.
    ///
    /// A handler added through this API will
    /// be attached to the default router. Consider creating the router
    /// middleware manually for advanced functionality.
    ///
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// let mut server = Nickel::new();
    ///
    /// //  without variables or wildcards
    /// fn bare_handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches /user");
    /// };
    /// server.get("/user", bare_handler);
    ///
    /// // with variables
    /// fn var_handler(request: &Request, response: &mut Response) {
    ///     let text = format!("This is user: {}", request.param("userid"));
    ///     response.send(text.as_slice());
    /// };
    /// server.get("/user/:userid", var_handler);
    ///
    /// // with simple wildcard
    /// fn wild_handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 but not /user/extended/list/4711");
    /// };
    /// server.get("/user/*/:userid", wild_handler);
    ///
    /// // with double wildcard
    /// fn very_wild_handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches /user/list/4711 and also /user/extended/list/4711");
    /// };
    /// server.get("/user/**/:userid", very_wild_handler);
    /// ```
    pub fn get(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Get, uri, handler);
    }

    /// Registers a handler to be used for a specific POST request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/post/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.post("/a/post/request", handler);
    /// ```
    pub fn post(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Post, uri, handler);
    }

    /// Registers a handler to be used for a specific PUT request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a POST request to /a/put/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.put("/a/put/request", handler);
    /// ```
    pub fn put(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Put, uri, handler);
    }

    /// Registers a handler to be used for a specific DELETE request.
    /// A handler added through this API will be attached to the default router.
    /// Consider creating the router middleware manually for advanced functionality.
    ///
    /// Take a look at `get(...)` for a more detailed description.
    /// # Example
    ///
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    /// fn handler(request: &Request, response: &mut Response) {
    ///     response.send("This matches a DELETE request to /a/delete/request");
    /// };
    ///
    /// let mut server = Nickel::new();
    /// server.delete("/a/delete/request", handler);
    /// ```
    pub fn delete(&mut self, uri: &str, handler: fn(request: &Request, response: &mut Response)){
        self.register_route_with_new_router(Delete, uri, handler);
    }

    fn register_route_with_new_router(&mut self, method: Method, uri: &str, handler: fn(request: &Request, response: &mut Response)) {
        let mut router = Router::new();
        router.add_route(method, String::from_str(uri), handler);
        self.utilize(router);
    }

    /// Registers an error handler which will be invoked among other error handler
    /// as soon as any regular handler returned an error
    ///
    /// A error handler is nearly identical to a regular middleware handler with the only
    /// difference that it takes an additional error parameter or type `NickelError.
    ///
    /// # Example
    ///
    /// ```{rust}
    /// # extern crate http;
    /// # extern crate nickel;
    /// # fn main() {
    /// use nickel::{Nickel, Request, Response, Action, Continue, Halt};
    /// use nickel::{NickelError, ErrorWithStatusCode, get_media_type};
    /// use http::status::NotFound;
    ///
    /// fn error_handler(err: &NickelError, req: &Request, response: &mut Response)
    ///                  -> Result<Action, NickelError>{
    ///    match err.kind {
    ///        ErrorWithStatusCode(NotFound) => {
    ///            response.origin.headers.content_type = get_media_type("html");
    ///            response.origin.status = NotFound;
    ///            response.send("<h1>Call the police!<h1>");
    ///            Ok(Halt)
    ///        },
    ///        _ => Ok(Continue)
    ///    }
    /// }
    ///
    /// let mut server = Nickel::new();
    /// server.handle_error(error_handler)
    /// # }
    /// ```
    pub fn handle_error(&mut self, handler: fn(err: &NickelError,
                                               req: &Request,
                                               res: &mut Response)
                                            -> Result<Action, NickelError>){
        let handler = IntoErrorHandler::from_fn(handler);
        self.middleware_stack.add_error_handler(handler);
    }

    /// Create a new middleware to serve as a router.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// use nickel::{Nickel, Request, Response};
    ///
    /// let mut server = Nickel::new();
    /// let mut router = Nickel::router();
    ///
    /// fn foo_handler(request: &Request, response: &mut Response) {
    ///     response.send("Hi from /foo");
    /// };
    ///
    /// router.get("/foo", foo_handler);
    /// server.utilize(router);
    /// ```
    pub fn router() -> Router {
        Router::new()
    }

    /// Create a new middleware to parse JSON bodies.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// # #![feature(phase)]
    /// # #[phase(plugin)] extern crate nickel_macros;
    /// # extern crate nickel;
    /// # extern crate serialize;
    /// # use nickel::{Nickel, Request, Response};
    /// use nickel::JsonBody;
    ///
    /// # fn main() {
    /// #[deriving(Decodable, Encodable)]
    /// struct Person {
    ///     first_name: String,
    ///     last_name:  String,
    /// }
    ///
    /// let router = router! {
    ///     post "/a/post/request" => |request, response| {
    ///         let person = request.json_as::<Person>().unwrap();
    ///         let text = format!("Hello {} {}", person.first_name, person.last_name);
    ///         response.send(text);
    ///     }
    /// };
    ///
    /// let mut server = Nickel::new();
    /// // It is currently a requirement that the json_body_parser middleware
    /// // is added before any routes that require it.
    /// server.utilize(Nickel::json_body_parser());
    /// server.utilize(router);
    /// # }
    /// ```
    pub fn json_body_parser() -> JsonBodyParser {
        JsonBodyParser
    }

    /// Create a new middleware to parse the query string.
    ///
    ///
    /// # Example
    /// ```{rust}
    /// # #![feature(phase)]
    /// # #[phase(plugin)] extern crate nickel_macros;
    /// # extern crate nickel;
    /// # use nickel::{Nickel, Request, Response};
    /// use nickel::QueryString;
    /// # fn main() {
    /// let router = router! {
    ///     get "/a/get/request" => |request, response| {
    ///         let foo = request.query("foo", "this is the default value, if foo is not present!");
    ///         response.send(foo[0].as_slice());
    ///     }
    /// };
    ///
    /// let mut server = Nickel::new();
    /// // It is currently a requirement that the query_string middleware
    /// // is added before any routes that require it.
    /// server.utilize(Nickel::query_string());
    /// server.utilize(router);
    /// # }
    /// ```
    pub fn query_string() -> QueryStringParser {
        QueryStringParser
    }

    /// Bind and listen for connections on the given host and port
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut server = Nickel::new();
    /// server.listen(Ipv4Addr(127, 0, 0, 1), 6767);
    /// ```
    pub fn listen(mut self, ip: IpAddr, port: Port) {
        fn not_found_handler(_request: &Request, _response: &mut Response) -> Result<Action, NickelError> {
            Err(NickelError::new("File Not Found", ErrorWithStatusCode(NotFound)))
        }

        self.middleware_stack.add_middleware(IntoMiddleware::from_fn(not_found_handler));
        self.server = Some(Server::new(self.middleware_stack, ip, port));
        self.server.unwrap().serve();
    }
}
<|endoftext|>


<commit_before>use std::slice;

use std::marker;

pub trait ModuleStruct<'a> {
	fn init_module(&mut self);
}

struct TorchBackend {}
struct Parameter {}

pub struct Module<'a> {
	pub _name: &'a str,
	_backend : TorchBackend,
//	_buffers: HashTable<&str, Tensor>
//	_backward_hooks:
//	_forward_hooks:
	_params: Vec<&'a str>,
    _modulesp: Vec<*mut Module<'a>>,
	training: bool,
}
pub struct PtrIterMut<'a, T: 'a> {
	mod_iter: slice::IterMut<'a, *mut T>,
}

impl <'a>Iterator for PtrIterMut<'a, Module<'a>> {
	type Item = &'a mut Module<'a>;
	fn next(&mut self) -> Option<Self::Item> {
		if let Some(t) = self.mod_iter.next() {
			Some(unsafe { &mut **t as Self::Item })
		} else {
			None
		}
	}
}

impl <'a>Module<'a> {
	pub fn new() -> Module<'a> {
		Module {_name: "", _backend: TorchBackend {}, _params: Vec::new(),
		_modulesp: Vec::new(), training: true }
	}
	#[inline]
    fn as_mut_ptr(&mut self) -> *mut Module<'a> {
        self as *mut Module<'a>
    }
    pub fn add_module(&mut self, module: &mut ModIntf<'a>) {
    	self._modulesp.push(module.delegate().as_mut_ptr())

    }
	pub fn modules_iter_mut(&mut self) -> PtrIterMut<Module<'a> > {
		PtrIterMut {mod_iter: self._modulesp.iter_mut()}  //_marker: marker::PhantomData } }
	}
}

pub trait ModIntf<'a> {
	fn delegate(&mut self) -> &mut Module<'a>;
	fn forward(&mut self /* Tensor */ ) /* -> Tensor */;
}

#[derive(ModParse)]
pub struct Linear<'a> {
	delegate: Module<'a> ,
	in_features: u32,
}

impl <'a>Linear<'a> {
	pub fn new(/* args: LinearArgs */) -> Linear<'a> {
		let mut t = Linear {delegate: Module::new(), in_features: 0 };
		t.init_module();
		t
	}
}

impl <'a> ModIntf<'a> for Linear<'a>  {
	fn delegate(&mut self) -> &mut Module<'a> {
		&mut self.delegate
	}
	fn forward(&mut self/* Tensor */) /* -> Tensor */ {

	}
}

<commit_msg>start fleshing out modules functions<commit_after>use super::*;
use std::collections::HashMap;
use std::collections::hash_map;
use std::iter;
//use linked_hash_map;
//use linked_hash_map::LinkedHashMap;

pub trait ModuleStruct<'a> {
	fn init_module(&mut self);
}

struct TorchBackend {}
pub struct Parameter<'a> {
	data: &'a mut Tensor,
}
impl <'a>Parameter<'a> {
	fn valid(&self) -> bool {
		true
	}
}

pub struct Module<'a> {
	pub _name: &'a str,
	_backend : TorchBackend,
//	_buffers: HashTable<&str, Tensor>
//	_backward_hooks:
//	_forward_hooks:
    //_modulesp: LinkedHashMap<&'a str, *mut Module<'a>>,
    _params: HashMap<&'a str, *mut Parameter<'a>>,
    _modulesp: HashMap<&'a str, *mut Module<'a>>,
	training: bool,
}
pub struct PtrIterMut<'a, T: 'a> {
	//mod_iter: linked_hash_map::IterMut<'a, &'a str, *mut T>,
	mod_iter: hash_map::IterMut<'a, &'a str, *mut T>,
}

impl <'a, T>Iterator for PtrIterMut<'a, T> {
	type Item = (&'a str, &'a mut T);
	fn next(&mut self) -> Option<Self::Item> {
		if let Some((name, t)) = self.mod_iter.next() {
			Some((name, unsafe { &mut **t as &mut T}))
		} else {
			None
		}
	}
}

impl <'a>Module<'a> {
	pub fn new() -> Module<'a> {
		Module {_name: "", _backend: TorchBackend {}, _params: HashMap::new(),
		//_modulesp: LinkedHashMap::new(),
		_modulesp: HashMap::new(),
		training: true }
	}
	#[inline]
    fn as_mut_ptr(&mut self) -> *mut Module<'a> {
        self as *mut Module<'a>
    }
    pub fn add_module(&mut self, module: &mut ModIntf<'a>) {
    	let m = module.delegate();
    	self._modulesp.insert(m._name, m.as_mut_ptr());

    }
	pub fn modules_iter_mut(&mut self) -> PtrIterMut<Module<'a> > {
		PtrIterMut {mod_iter: self._modulesp.iter_mut()}
	}
	pub fn params_iter_mut(&mut self) -> PtrIterMut<Parameter<'a> > {
		PtrIterMut {mod_iter: self._params.iter_mut()}
	}
	fn _apply(&mut self, callback: fn(&mut Tensor))  {
		for (_, module) in self.modules_iter_mut() { module._apply(callback) }
		for param in self.params_iter_mut().filter_map(|(_, p)| {if p.valid() {Some(p)} else {None} }) {
			callback(param.data);
/*
			if let Some(g) = p._grad {
					callback(g.data)
			}
			*/
			/* see also _buffers */
		}
	}
	fn apply(&mut self, callback: fn(&mut Self)) {
		for (_, module) in self.modules_iter_mut() {
			module.apply(callback)
		}
		callback(self)
	}
	pub fn eval(&mut self) {
		self.train(false)
	}
	fn train(&mut self, mode: bool) {
		self.training = mode;
		for (_, module) in self.modules_iter_mut() {
			module.train(mode)
		}

	}
}

pub trait ModIntf<'a> {
	fn delegate(&mut self) -> &mut Module<'a>;
	fn forward(&mut self /* Tensor */ ) /* -> Tensor */;
}

#[derive(ModParse)]
pub struct Linear<'a> {
	delegate: Module<'a> ,
	in_features: u32,
}

impl <'a>Linear<'a> {
	pub fn new(/* args: LinearArgs */) -> Linear<'a> {
		let mut t = Linear {delegate: Module::new(), in_features: 0 };
		t.init_module();
		t
	}
}

impl <'a> ModIntf<'a> for Linear<'a>  {
	fn delegate(&mut self) -> &mut Module<'a> {
		&mut self.delegate
	}
	fn forward(&mut self/* Tensor */) /* -> Tensor */ {

	}
}

<|endoftext|>


<commit_before>use std::str::FromStr;

#[derive(Debug, PartialEq)]
pub enum Instruction {
    NOP,
    MOV { src: Operand, dst: Operand },
    SWP,
    SAV,
    ADD { addend: Operand },
    SUB { subtrahend: Operand} ,
    NEG,
    J { cond: Condition, dst: Label},
    JRO { dst: Operand },
}

static BAD_OPCODE_ERR: &'static str = "Bad opcode for # of arguments";
static NUM_ARGS_ERR: &'static str = "Wrong number of arguments";

impl FromStr for Instruction {
    type Err = &'static str;

    fn from_str(insn: &str) -> Result<Instruction, Self::Err> {
        let words: Vec<&str> = insn.split(' ').collect();

        match words.len() {
            1 => match words[0] {
                "NOP" => Ok(Instruction::NOP),
                "SWP" => Ok(Instruction::SWP),
                "SAV" => Ok(Instruction::SAV),
                "NEG" => Ok(Instruction::NEG),
                _ => Err(BAD_OPCODE_ERR),
            },

            2 => match words[0] {
                "ADD" => Operand::from_str(words[1]).map(|o| Instruction::ADD { addend: o }),
                "SUB" => Operand::from_str(words[1]).map(|o| Instruction::SUB { subtrahend: o }),
                "JMP" => Ok(Instruction::J { cond: Condition::Unconditional, dst: words[1].to_string() }),
                "JEZ" => Ok(Instruction::J { cond: Condition::Ez,            dst: words[1].to_string() }),
                "JNZ" => Ok(Instruction::J { cond: Condition::Nz,            dst: words[1].to_string() }),
                "JGZ" => Ok(Instruction::J { cond: Condition::Gz,            dst: words[1].to_string() }),
                "JLZ" => Ok(Instruction::J { cond: Condition::Lz,            dst: words[1].to_string() }),
                "JRO" => Operand::from_str(words[1]).map(|o| Instruction::JRO { dst: o }),
                _ => Err(BAD_OPCODE_ERR),
            },

            3 => match words[0] {
                "MOV" => Operand::from_str(words[1]).and_then(|s| Operand::from_str(words[2]).map(|d| (s, d)))
                            .map(|(s, d)| Instruction::MOV { src: s, dst: d }),
                _ => Err(BAD_OPCODE_ERR),
            },

            _ => Err(NUM_ARGS_ERR)
        }
    }
}

#[test]
fn instruction_from_str() {
    fn i(s: &str) -> Instruction {
        Instruction::from_str(s).unwrap()
    }
    assert_eq!(i("NOP"), Instruction::NOP);
    assert_eq!(i("SWP"), Instruction::SWP);
    assert_eq!(i("SAV"), Instruction::SAV);
    assert_eq!(i("NEG"), Instruction::NEG);
    assert_eq!(Instruction::from_str("BOGUS").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("ADD 10"), Instruction::ADD { addend: Operand::Lit(10) });
    assert_eq!(i("SUB 10"), Instruction::SUB { subtrahend: Operand::Lit(10) });
    assert_eq!(i("JMP LOC"), Instruction::J { cond: Condition::Unconditional,   dst: "LOC".to_string() });
    assert_eq!(i("JEZ LOC"), Instruction::J { cond: Condition::Ez,              dst: "LOC".to_string() });
    assert_eq!(i("JGZ LOC"), Instruction::J { cond: Condition::Gz,              dst: "LOC".to_string() });
    assert_eq!(i("JLZ LOC"), Instruction::J { cond: Condition::Lz,              dst: "LOC".to_string() });
    assert_eq!(i("JRO 1"), Instruction::JRO { dst: Operand::Lit(1) });
    assert_eq!(Instruction::from_str("JEQ LOC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("MOV UP DOWN"),    Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    //assert_eq!(i("MOV  UP  DOWN"), Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV UP ACC"),     Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::ACC });
    assert_eq!(i("MOV ACC ACC"),    Instruction::MOV { src: Operand::ACC, dst: Operand::ACC });
    assert_eq!(Instruction::from_str("MV UP ACC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(Instruction::from_str("1 2 3 4").unwrap_err(), NUM_ARGS_ERR);
}

#[derive(Debug, PartialEq)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

impl FromStr for Port {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err("bad port"),
        }
    }
}

#[test]
fn port_from_str() {
    assert_eq!(Port::from_str("UP").unwrap(),   Port::Up);
    assert_eq!(Port::from_str("DOWN").unwrap(), Port::Down);
    assert_eq!(Port::from_str("LEFT").unwrap(), Port::Left);
    assert_eq!(Port::from_str("RIGHT").unwrap(), Port::Right);
    assert_eq!(Port::from_str("OTHER").unwrap_err(), "bad port");
}

#[derive(Debug, PartialEq)]
pub enum Operand {
    Lit(i32),
    Port(Port),
    ACC,
}

impl FromStr for Operand {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "ACC" {
            return Ok(Operand::ACC);
        }

        let as_int = i32::from_str(s);
        if as_int.is_ok() {
            return Ok(Operand::Lit(as_int.unwrap()));
        }

        let as_port = Port::from_str(s);
        if as_port.is_ok() {
            return Ok(Operand::Port(as_port.unwrap()));
        }

        Err("Invalid operand")
    }
}

#[test]
fn operand_from_str() {
    assert_eq!(Operand::from_str("ACC").unwrap(), Operand::ACC);
    assert_eq!(Operand::from_str("32").unwrap(), Operand::Lit(32));
    assert_eq!(Operand::from_str("UP").unwrap(), Operand::Port(Port::Up));
    assert_eq!(Operand::from_str("FOO").unwrap_err(), "Invalid operand");
}

#[derive(Debug, PartialEq)]
pub enum Condition {
    Unconditional,
    Ez,
    Nz,
    Gz,
    Lz,
}

pub type Label = String;
<commit_msg>Instruction: fix instructions with extra spaces<commit_after>use std::str::FromStr;

#[derive(Debug, PartialEq)]
pub enum Instruction {
    NOP,
    MOV { src: Operand, dst: Operand },
    SWP,
    SAV,
    ADD { addend: Operand },
    SUB { subtrahend: Operand} ,
    NEG,
    J { cond: Condition, dst: Label},
    JRO { dst: Operand },
}

static BAD_OPCODE_ERR: &'static str = "Bad opcode for # of arguments";
static NUM_ARGS_ERR: &'static str = "Wrong number of arguments";

impl FromStr for Instruction {
    type Err = &'static str;

    fn from_str(insn: &str) -> Result<Instruction, Self::Err> {
        //Nightly: let words: Vec<&str> = insn.split_whitespace().collect();
        let words: Vec<&str> = insn.split(' ').filter(|s| *s != "").collect();

        match words.len() {
            1 => match words[0] {
                "NOP" => Ok(Instruction::NOP),
                "SWP" => Ok(Instruction::SWP),
                "SAV" => Ok(Instruction::SAV),
                "NEG" => Ok(Instruction::NEG),
                _ => Err(BAD_OPCODE_ERR),
            },

            2 => match words[0] {
                "ADD" => Operand::from_str(words[1]).map(|o| Instruction::ADD { addend: o }),
                "SUB" => Operand::from_str(words[1]).map(|o| Instruction::SUB { subtrahend: o }),
                "JMP" => Ok(Instruction::J { cond: Condition::Unconditional, dst: words[1].to_string() }),
                "JEZ" => Ok(Instruction::J { cond: Condition::Ez,            dst: words[1].to_string() }),
                "JNZ" => Ok(Instruction::J { cond: Condition::Nz,            dst: words[1].to_string() }),
                "JGZ" => Ok(Instruction::J { cond: Condition::Gz,            dst: words[1].to_string() }),
                "JLZ" => Ok(Instruction::J { cond: Condition::Lz,            dst: words[1].to_string() }),
                "JRO" => Operand::from_str(words[1]).map(|o| Instruction::JRO { dst: o }),
                _ => Err(BAD_OPCODE_ERR),
            },

            3 => match words[0] {
                "MOV" => Operand::from_str(words[1]).and_then(|s| Operand::from_str(words[2]).map(|d| (s, d)))
                            .map(|(s, d)| Instruction::MOV { src: s, dst: d }),
                _ => Err(BAD_OPCODE_ERR),
            },

            _ => Err(NUM_ARGS_ERR)
        }
    }
}

#[test]
fn instruction_from_str() {
    fn i(s: &str) -> Instruction {
        Instruction::from_str(s).unwrap()
    }
    assert_eq!(i("NOP"), Instruction::NOP);
    assert_eq!(i("SWP"), Instruction::SWP);
    assert_eq!(i("SAV"), Instruction::SAV);
    assert_eq!(i("NEG"), Instruction::NEG);
    assert_eq!(Instruction::from_str("BOGUS").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("ADD 10"), Instruction::ADD { addend: Operand::Lit(10) });
    assert_eq!(i("SUB 10"), Instruction::SUB { subtrahend: Operand::Lit(10) });
    assert_eq!(i("JMP LOC"), Instruction::J { cond: Condition::Unconditional,   dst: "LOC".to_string() });
    assert_eq!(i("JEZ LOC"), Instruction::J { cond: Condition::Ez,              dst: "LOC".to_string() });
    assert_eq!(i("JGZ LOC"), Instruction::J { cond: Condition::Gz,              dst: "LOC".to_string() });
    assert_eq!(i("JLZ LOC"), Instruction::J { cond: Condition::Lz,              dst: "LOC".to_string() });
    assert_eq!(i("JRO 1"), Instruction::JRO { dst: Operand::Lit(1) });
    assert_eq!(Instruction::from_str("JEQ LOC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(i("MOV UP DOWN"),    Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV  UP  DOWN"),  Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::Port(Port::Down) });
    assert_eq!(i("MOV UP ACC"),     Instruction::MOV { src: Operand::Port(Port::Up), dst: Operand::ACC });
    assert_eq!(i("MOV ACC ACC"),    Instruction::MOV { src: Operand::ACC, dst: Operand::ACC });
    assert_eq!(Instruction::from_str("MV UP ACC").unwrap_err(), BAD_OPCODE_ERR);

    assert_eq!(Instruction::from_str("1 2 3 4").unwrap_err(), NUM_ARGS_ERR);
}

#[derive(Debug, PartialEq)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

impl FromStr for Port {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err("bad port"),
        }
    }
}

#[test]
fn port_from_str() {
    assert_eq!(Port::from_str("UP").unwrap(),   Port::Up);
    assert_eq!(Port::from_str("DOWN").unwrap(), Port::Down);
    assert_eq!(Port::from_str("LEFT").unwrap(), Port::Left);
    assert_eq!(Port::from_str("RIGHT").unwrap(), Port::Right);
    assert_eq!(Port::from_str("OTHER").unwrap_err(), "bad port");
}

#[derive(Debug, PartialEq)]
pub enum Operand {
    Lit(i32),
    Port(Port),
    ACC,
}

impl FromStr for Operand {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "ACC" {
            return Ok(Operand::ACC);
        }

        let as_int = i32::from_str(s);
        if as_int.is_ok() {
            return Ok(Operand::Lit(as_int.unwrap()));
        }

        let as_port = Port::from_str(s);
        if as_port.is_ok() {
            return Ok(Operand::Port(as_port.unwrap()));
        }

        Err("Invalid operand")
    }
}

#[test]
fn operand_from_str() {
    assert_eq!(Operand::from_str("ACC").unwrap(), Operand::ACC);
    assert_eq!(Operand::from_str("32").unwrap(), Operand::Lit(32));
    assert_eq!(Operand::from_str("UP").unwrap(), Operand::Port(Port::Up));
    assert_eq!(Operand::from_str("FOO").unwrap_err(), "Invalid operand");
}

#[derive(Debug, PartialEq)]
pub enum Condition {
    Unconditional,
    Ez,
    Nz,
    Gz,
    Lz,
}

pub type Label = String;
<|endoftext|>


<commit_before>use image::{GenericImage, DynamicImage, Rgba};
use colors;
use ops;

pub struct Interpreter {
    codel_chooser: CodelChooser,
    direction_pointer: Direction,
    image: DynamicImage,
    current_color: Rgba<u8>,
    pub current_size: u32,
    pub stack: Vec<i32>,
}

#[derive(Debug)]
#[derive(PartialEq)]
enum Direction {
    Left,
    Right,
    Down,
    Up,
}

#[derive(Debug)]
enum CodelChooser {
    Right,
    Left,
}

impl Interpreter {
    pub fn new(image: DynamicImage) -> Interpreter {
        Interpreter {
            codel_chooser: CodelChooser::Left,
            direction_pointer: Direction::Right,
            current_color: image.get_pixel(0, 0),
            current_size: 0,
            image: image,
            stack: Vec::new()
        }
    }

    pub fn run(&mut self) {

        let mut failed_attempts = 0;
        let mut failed_white_attempts = 0;

        // These are here for debugging purposes
        let mut x = 0;
        let mut y = 0;

        loop {
            // If white just sail through until we hit another color
            if self.current_color_code() == 1 {
                let (x_next, y_next) = self.move_step(x, y);
                // If in white and hit boundary or black then step through dp
                if self.restricted(x_next, y_next) {
                    self.toggle_cc();
                    self.step_dp();
                    failed_white_attempts += 1;
                } else {
                    x = x_next;
                    y = y_next;
                    self.current_color = self.pixel_at(x, y);
                    failed_attempts = 0;
                    failed_white_attempts = 0;
                }

                if failed_white_attempts == 4 {
                    break;
                }
            } else {
                let (x_next, y_next, size) = self.next_block(x, y);
                // If we are hitting the edge or a black codel try another direction
                if self.restricted(x_next, y_next) {
                    self.update_movement(failed_attempts);
                    failed_attempts += 1;
                } else { // here we move into the next color block after processing the action
                    x = x_next;
                    y = y_next;

                    let next_color = self.pixel_at(x, y);
                    self.current_size = size;
                    self.process_action(&next_color);
                    self.current_color = next_color;
                    failed_attempts = 0;
                    failed_white_attempts = 0;
                }

                if failed_attempts == 8 {
                    break;
                }
            }

            // println!("Loop restarted");
        }

    }

    fn process_action(&mut self, next_color: &Rgba<u8>) {
        let color_code = colors::color_code(&next_color).unwrap();
        if color_code != 1 {
            let curr_color = self.current_color_code();
            let shade_diff = colors::shade_difference(curr_color, color_code);
            let hue_diff = colors::hue_difference(curr_color, color_code);
            ops::call_op(self, (hue_diff, shade_diff));
        }
    }

    // Based on the number of attempts this either toggles the CC back and forth
    // or updates the current direction of the DP
    fn update_movement(&mut self, attempts: i32) {
        if attempts % 2 == 0 {
            self.toggle_cc();
        } else {
            self.step_dp();
        }
    }

    pub fn toggle_cc(&mut self) {
        match self.codel_chooser {
            CodelChooser::Right => self.codel_chooser = CodelChooser::Left,
            CodelChooser::Left => self.codel_chooser = CodelChooser::Right,
        };
    }

    pub fn step_dp(&mut self) {
      match self.direction_pointer {
        Direction::Left => self.direction_pointer = Direction::Up,
        Direction::Right => self.direction_pointer = Direction::Down,
        Direction::Down => self.direction_pointer = Direction::Left,
        Direction::Up => self.direction_pointer = Direction::Right,
      }
    }

    pub fn step_dp_counter(&mut self) {
      match self.direction_pointer {
        Direction::Left => self.direction_pointer = Direction::Down,
        Direction::Right => self.direction_pointer = Direction::Up,
        Direction::Down => self.direction_pointer = Direction::Right,
        Direction::Up => self.direction_pointer = Direction::Left,
      }
    }

    // This method sets the current size after finding the next block
    fn next_block(&mut self, x: i32, y: i32) -> (i32, i32, u32) {
        let marked_size = self.image.width() * self.image.height();
        let mut marked = vec![false; marked_size as usize];
        let mut size = 0;
        let mut mx = x;
        let mut my = y;
        self.block_walk_recursive(&mut size, x, y, &mut mx, &mut my, &mut marked);
        let (x, y) = self.move_step(mx, my);
        (x, y, size)
    }

    fn move_step(&self, x: i32, y: i32) -> (i32, i32) {
        match self.direction_pointer {
            Direction::Left => (x-1, y),
            Direction::Up => (x, y-1),
            Direction::Right => (x+1, y),
            Direction::Down => (x, y+1),
        }
    }

    // Check if the color at a given position is different
    // than the current color
    fn current_color_eq(&self, x: i32, y: i32) -> bool {
        self.current_color == self.pixel_at(x, y)
    }

    fn pixel_at(&self, x: i32, y: i32) -> Rgba<u8> {
        self.image.get_pixel(x as u32, y as u32)
    }

    fn current_color_code(&self) -> i32 {
        match colors::color_code(&self.current_color) {
            Some(color) => color,
            None => 0,
        }
    }

    fn color_at(&self, x: i32, y: i32) -> i32 {
        match colors::color_code(&self.pixel_at(x, y)) {
            Some(color) => color,
            None => 0,
        }
    }

    fn restricted(&self, x: i32, y: i32) -> bool {
        !self.image.in_bounds(x as u32, y as u32) || self.color_at(x, y) == 0
    }

    // block walk should count the size of the current block
    // and set it on the Interpreter
    fn block_walk_recursive(
        &self,
        size: &mut u32,
        x: i32,
        y: i32,
        mx: &mut i32,
        my: &mut i32,
        marked: &mut Vec<bool>,
    ) {
        // println!("mx: {:?}, my: {:?}", mx, my);
        // println!(" x: {:?},  y: {:?}\n", x, y);
        // println!("DP: {:?}, CC: {:?}", self.direction_pointer, self.codel_chooser);
        if self.image.in_bounds(x as u32, y as u32) {
            // println!("curr: {:?},  step: {:?}\n", self.current_color_code(), self.color_at(x, y));
            let visit_index = self.marked_index(x, y) as usize;

            if self.current_color_eq(x, y) && !marked[visit_index] {
                *size += 1;
                marked[visit_index] = true;

                match self.direction_pointer {
                    Direction::Left => {
                        if x < *mx {
                            *mx = x;
                            *my = y;
                        }
                        else if x == *mx {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if y > *my {
                                        *my = y;
                                    }
                                },

                                CodelChooser::Right => {
                                    if y < *my {
                                        *my = y;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Right => {
                        if x > *mx {
                            *mx = x;
                            *my = y;
                        }
                        else if x == *mx {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if y < *my {
                                        *my = y;
                                    }
                                },

                                CodelChooser::Right => {
                                    if y > *my {
                                        *my = y;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Down => {
                        if y > *my {
                            *mx = x;
                            *my = y;
                        }
                        else if y == *my {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if x > *mx {
                                        *mx = x;
                                    }
                                },

                                CodelChooser::Right => {
                                    if x < *mx {
                                        *mx = x;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Up => {
                        if y < *my {
                            *mx = x;
                            *my = y;
                        }
                        else if y == *my {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if x < *mx {
                                        *mx = x;
                                    }
                                },

                                CodelChooser::Right => {
                                    if x > *mx {
                                        *mx = x;
                                    }
                                }
                            }
                        }
                    },
                }

                self.block_walk_recursive(size, x+1, y, mx, my, marked);
                self.block_walk_recursive(size, x-1, y, mx, my, marked);
                self.block_walk_recursive(size, x, y+1, mx, my, marked);
                self.block_walk_recursive(size, x, y-1, mx, my, marked);
            }
        }
        // print!("\n");
    }

    fn marked_index(&self, x: i32, y: i32) -> i32 {
        if x == 0 && y == 0 {
            return 0;
        }

        if x == 0 {
            return y * self.image.width() as i32;
        }

        if y == 0 {
            return x;
        }

        (x-1) * self.image.width() as i32 + (y-1)
    }
}<commit_msg>Fix index bug<commit_after>use image::{GenericImage, DynamicImage, Rgba};
use colors;
use ops;

pub struct Interpreter {
    codel_chooser: CodelChooser,
    direction_pointer: Direction,
    image: DynamicImage,
    current_color: Rgba<u8>,
    pub current_size: u32,
    pub stack: Vec<i32>,
}

#[derive(Debug)]
#[derive(PartialEq)]
enum Direction {
    Left,
    Right,
    Down,
    Up,
}

#[derive(Debug)]
enum CodelChooser {
    Right,
    Left,
}

impl Interpreter {
    pub fn new(image: DynamicImage) -> Interpreter {
        Interpreter {
            codel_chooser: CodelChooser::Left,
            direction_pointer: Direction::Right,
            current_color: image.get_pixel(0, 0),
            current_size: 0,
            image: image,
            stack: Vec::new()
        }
    }

    pub fn run(&mut self) {

        let mut failed_attempts = 0;
        let mut failed_white_attempts = 0;

        // These are here for debugging purposes
        let mut x = 0;
        let mut y = 0;

        loop {
            // If white just sail through until we hit another color
            if self.current_color_code() == 1 {
                let (x_next, y_next) = self.move_step(x, y);
                // If in white and hit boundary or black then step through dp
                if self.restricted(x_next, y_next) {
                    self.toggle_cc();
                    self.step_dp();
                    failed_white_attempts += 1;
                } else {
                    x = x_next;
                    y = y_next;
                    self.current_color = self.pixel_at(x, y);
                    failed_attempts = 0;
                    failed_white_attempts = 0;
                }

                if failed_white_attempts == 4 {
                    break;
                }
            } else {
                let (x_next, y_next, size) = self.next_block(x, y);
                // If we are hitting the edge or a black codel try another direction
                if self.restricted(x_next, y_next) {
                    self.update_movement(failed_attempts);
                    failed_attempts += 1;
                } else { // here we move into the next color block after processing the action
                    x = x_next;
                    y = y_next;

                    let next_color = self.pixel_at(x, y);
                    self.current_size = size;
                    self.process_action(&next_color);
                    self.current_color = next_color;
                    failed_attempts = 0;
                    failed_white_attempts = 0;
                }

                if failed_attempts == 8 {
                    break;
                }
            }
        }

    }

    fn process_action(&mut self, next_color: &Rgba<u8>) {
        let color_code = colors::color_code(&next_color).unwrap();
        if color_code != 1 {
            let curr_color = self.current_color_code();
            let shade_diff = colors::shade_difference(curr_color, color_code);
            let hue_diff = colors::hue_difference(curr_color, color_code);
            ops::call_op(self, (hue_diff, shade_diff));
        }
    }

    // Based on the number of attempts this either toggles the CC back and forth
    // or updates the current direction of the DP
    fn update_movement(&mut self, attempts: i32) {
        if attempts % 2 == 0 {
            self.toggle_cc();
        } else {
            self.step_dp();
        }
    }

    pub fn toggle_cc(&mut self) {
        match self.codel_chooser {
            CodelChooser::Right => self.codel_chooser = CodelChooser::Left,
            CodelChooser::Left => self.codel_chooser = CodelChooser::Right,
        };
    }

    pub fn step_dp(&mut self) {
      match self.direction_pointer {
        Direction::Left => self.direction_pointer = Direction::Up,
        Direction::Right => self.direction_pointer = Direction::Down,
        Direction::Down => self.direction_pointer = Direction::Left,
        Direction::Up => self.direction_pointer = Direction::Right,
      }
    }

    pub fn step_dp_counter(&mut self) {
      match self.direction_pointer {
        Direction::Left => self.direction_pointer = Direction::Down,
        Direction::Right => self.direction_pointer = Direction::Up,
        Direction::Down => self.direction_pointer = Direction::Right,
        Direction::Up => self.direction_pointer = Direction::Left,
      }
    }

    // This method sets the current size after finding the next block
    fn next_block(&mut self, x: i32, y: i32) -> (i32, i32, u32) {
        let marked_size = self.image.width() * self.image.height();
        let mut marked = vec![false; marked_size as usize];
        let mut size = 0;
        let mut mx = x;
        let mut my = y;
        self.block_walk_recursive(&mut size, x, y, &mut mx, &mut my, &mut marked);
        let (x, y) = self.move_step(mx, my);
        (x, y, size)
    }

    fn move_step(&self, x: i32, y: i32) -> (i32, i32) {
        match self.direction_pointer {
            Direction::Left => (x-1, y),
            Direction::Up => (x, y-1),
            Direction::Right => (x+1, y),
            Direction::Down => (x, y+1),
        }
    }

    // Check if the color at a given position is different
    // than the current color
    fn current_color_eq(&self, x: i32, y: i32) -> bool {
        self.current_color == self.pixel_at(x, y)
    }

    fn pixel_at(&self, x: i32, y: i32) -> Rgba<u8> {
        self.image.get_pixel(x as u32, y as u32)
    }

    fn current_color_code(&self) -> i32 {
        match colors::color_code(&self.current_color) {
            Some(color) => color,
            None => 0,
        }
    }

    fn color_at(&self, x: i32, y: i32) -> i32 {
        match colors::color_code(&self.pixel_at(x, y)) {
            Some(color) => color,
            None => 0,
        }
    }

    fn restricted(&self, x: i32, y: i32) -> bool {
        !self.image.in_bounds(x as u32, y as u32) || self.color_at(x, y) == 0
    }

    // block walk should count the size of the current block
    // and set it on the Interpreter
    fn block_walk_recursive(
        &self,
        size: &mut u32,
        x: i32,
        y: i32,
        mx: &mut i32,
        my: &mut i32,
        marked: &mut Vec<bool>,
    ) {
        // println!("mx: {:?}, my: {:?}", mx, my);
        // println!(" x: {:?},  y: {:?}", x, y);
        // println!("DP: {:?}, CC: {:?}", self.direction_pointer, self.codel_chooser);
        // println!("Size: {:?}", size);
        if self.image.in_bounds(x as u32, y as u32) {
            // println!("curr: {:?},  step: {:?}\n", self.current_color_code(), self.color_at(x, y));
            let visit_index = self.marked_index(x, y) as usize;
            // println!("{:?}\n", visit_index);

            if self.current_color_eq(x, y) && !marked[visit_index] {
                *size += 1;
                marked[visit_index] = true;

                match self.direction_pointer {
                    Direction::Left => {
                        if x < *mx {
                            *mx = x;
                            *my = y;
                        }
                        else if x == *mx {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if y > *my {
                                        *my = y;
                                    }
                                },

                                CodelChooser::Right => {
                                    if y < *my {
                                        *my = y;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Right => {
                        if x > *mx {
                            *mx = x;
                            *my = y;
                        }
                        else if x == *mx {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if y < *my {
                                        *my = y;
                                    }
                                },

                                CodelChooser::Right => {
                                    if y > *my {
                                        *my = y;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Down => {
                        if y > *my {
                            *mx = x;
                            *my = y;
                        }
                        else if y == *my {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if x > *mx {
                                        *mx = x;
                                    }
                                },

                                CodelChooser::Right => {
                                    if x < *mx {
                                        *mx = x;
                                    }
                                }
                            }
                        }
                    },
                    Direction::Up => {
                        if y < *my {
                            *mx = x;
                            *my = y;
                        }
                        else if y == *my {
                            match self.codel_chooser {
                                CodelChooser::Left => {
                                    if x < *mx {
                                        *mx = x;
                                    }
                                },

                                CodelChooser::Right => {
                                    if x > *mx {
                                        *mx = x;
                                    }
                                }
                            }
                        }
                    },
                }

                self.block_walk_recursive(size, x+1, y, mx, my, marked);
                self.block_walk_recursive(size, x-1, y, mx, my, marked);
                self.block_walk_recursive(size, x, y+1, mx, my, marked);
                self.block_walk_recursive(size, x, y-1, mx, my, marked);
            }
        }
        // print!("\n");
    }
    // 0, 1, 2
    // 3, 4, 5
    // 6, 7, 8
    fn marked_index(&self, x: i32, y: i32) -> i32 {
        if x == 0 && y == 0 {
            return 0;
        }

        if x == 0 {
            return y * self.image.width() as i32;
        }

        if y == 0 {
            return x;
        }

        y * self.image.width() as i32 + x
    }
}<|endoftext|>


<commit_before>// Exports of <CoreFoundation/CFNumber.h>

use libc::c_void;

use base::*;

#[repr(C)]
struct __CFBoolean;
pub type CFBooleanRef = *const __CFBoolean;

extern "C" {
    pub fn CFBooleanGetTypeID() -> CFTypeID;
    pub fn CFBooleanGetValue(boolean: CFBooleanRef) -> Boolean;
}

pub type CFNumberType = CFIndex;
pub const kCFNumberSInt8Type: CFNumberType = 1;
pub const kCFNumberSInt16Type: CFNumberType = 2;
pub const kCFNumberSInt32Type: CFNumberType = 3;
pub const kCFNumberSInt64Type: CFNumberType = 4;
pub const kCFNumberFloat32Type: CFNumberType = 5;
pub const kCFNumberFloat64Type: CFNumberType = 6;
pub const kCFNumberCharType: CFNumberType = 7;
pub const kCFNumberShortType: CFNumberType = 8;
pub const kCFNumberIntType: CFNumberType = 9;
pub const kCFNumberLongType: CFNumberType = 10;
pub const kCFNumberLongLongType: CFNumberType = 11;
pub const kCFNumberFloatType: CFNumberType = 12;
pub const kCFNumberDoubleType: CFNumberType = 13;
pub const kCFNumberCFIndexType: CFNumberType = 14;
pub const kCFNumberNSIntegerType: CFNumberType = 15;
pub const kCFNumberCGFloatType: CFNumberType = 16;
pub const kCFNumberMaxType: CFNumberType = 16;

#[repr(C)]
struct __CFNumber;
pub type CFNumberRef = *const __CFNumber;

extern "C" {
    pub fn CFNumberGetTypeID() -> CFTypeID;
    pub fn CFNumberCreate(alloc: CFAllocatorRef, theType: CFNumberType, valuePtr: *mut c_void) -> CFNumberRef;
    pub fn CFNumberGetType(number: CFNumberRef) -> CFNumberType;
    pub fn CFNumberGetByteSize(number: CFNumberRef) -> CFIndex;
    pub fn CFNumberIsFloatType(number: CFNumberRef) -> Boolean;
    pub fn CFNumberGetValue(number: CFNumberRef, theType: CFNumberType, valuePtr: *mut c_void) -> CFNumberRef;
    pub fn CFNumberCompare(number: CFNumberRef, otherNumber: CFNumberRef, context: *const c_void) -> CFComparisonResult;
}
<commit_msg>fix code review of adding CFNumber.h bindings<commit_after>// Exports of <CoreFoundation/CFNumber.h>

use libc::c_void;

use base::*;

#[repr(C)]
struct __CFBoolean;
pub type CFBooleanRef = *const __CFBoolean;

extern "C" {
    pub fn CFBooleanGetTypeID() -> CFTypeID;
    pub fn CFBooleanGetValue(boolean: CFBooleanRef) -> Boolean;
}

pub type CFNumberType = CFIndex;
pub const kCFNumberSInt8Type: CFNumberType = 1;
pub const kCFNumberSInt16Type: CFNumberType = 2;
pub const kCFNumberSInt32Type: CFNumberType = 3;
pub const kCFNumberSInt64Type: CFNumberType = 4;
pub const kCFNumberFloat32Type: CFNumberType = 5;
pub const kCFNumberFloat64Type: CFNumberType = 6;
pub const kCFNumberCharType: CFNumberType = 7;
pub const kCFNumberShortType: CFNumberType = 8;
pub const kCFNumberIntType: CFNumberType = 9;
pub const kCFNumberLongType: CFNumberType = 10;
pub const kCFNumberLongLongType: CFNumberType = 11;
pub const kCFNumberFloatType: CFNumberType = 12;
pub const kCFNumberDoubleType: CFNumberType = 13;
pub const kCFNumberCFIndexType: CFNumberType = 14;
pub const kCFNumberNSIntegerType: CFNumberType = 15;
pub const kCFNumberCGFloatType: CFNumberType = 16;
pub const kCFNumberMaxType: CFNumberType = 16;

#[repr(C)]
struct __CFNumber;
pub type CFNumberRef = *const __CFNumber;

extern "C" {
    pub fn CFNumberGetTypeID() -> CFTypeID;
    pub fn CFNumberCreate(alloc: CFAllocatorRef, theType: CFNumberType, valuePtr: *const c_void) -> CFNumberRef;
    pub fn CFNumberGetType(number: CFNumberRef) -> CFNumberType;
    pub fn CFNumberGetByteSize(number: CFNumberRef) -> CFIndex;
    pub fn CFNumberIsFloatType(number: CFNumberRef) -> Boolean;
    pub fn CFNumberGetValue(number: CFNumberRef, theType: CFNumberType, valuePtr: *mut c_void) -> Boolean;
    pub fn CFNumberCompare(number: CFNumberRef, otherNumber: CFNumberRef, context: *mut c_void) -> CFComparisonResult;
}
<|endoftext|>


<commit_before>use super::interface::DBusInterface;

extern crate core;
use self::core::ops::Deref;

extern crate dbus;
use self::dbus::{Connection, Error};

use std::collections::btree_map::BTreeMap;

pub struct DBusObject<'a> {
    ifaces: BTreeMap<String, Box<DBusInterface>>,

    conn: &'a Connection,
    path: String,
}

impl<'a> DBusObject<'a> {
    pub fn new(conn: &'a Connection, ifaces: BTreeMap<String, Box<DBusInterface>>, path: &str) -> Result<Self, Error> {
        try!(conn.register_object_path(path));

        Ok(DBusObject {
            ifaces: ifaces,

            conn: conn,
            path: path.to_string(),
        })
    }

    pub fn path(&self) -> &str {
        &self.path
    }

    pub fn get_interface(&self, interface: &str) -> Option<&DBusInterface> {
        self.ifaces.get(interface).map(|iface| {
            iface.deref()
        })
    }
}

impl<'a> Drop for DBusObject<'a> {
    fn drop(&mut self) {
        self.conn.unregister_object_path(self.path())
    }
}
<commit_msg>object: implement get_interface_mut<commit_after>use super::interface::DBusInterface;

extern crate core;
use self::core::ops::Deref;

extern crate dbus;
use self::dbus::{Connection, Error};

use std::collections::btree_map::BTreeMap;

pub struct DBusObject<'a> {
    ifaces: BTreeMap<String, Box<DBusInterface>>,

    conn: &'a Connection,
    path: String,
}

impl<'a> DBusObject<'a> {
    pub fn new(conn: &'a Connection, ifaces: BTreeMap<String, Box<DBusInterface>>, path: &str) -> Result<Self, Error> {
        try!(conn.register_object_path(path));

        Ok(DBusObject {
            ifaces: ifaces,

            conn: conn,
            path: path.to_string(),
        })
    }

    pub fn path(&self) -> &str {
        &self.path
    }

    pub fn get_interface(&self, interface: &str) -> Option<&DBusInterface> {
        self.ifaces.get(interface).map(|iface| {
            // FIXME: Why is .deref() necessary?
            iface.deref()
        })
    }

    pub fn get_interface_mut<'i>(&'i mut self, interface: &str) -> Option<&'i mut (DBusInterface + 'i)> {
        self.ifaces.get_mut(interface).map(|iface| ->
            // FIXME: Why is this not just iface.deref_mut().
            // FIXME: Why is .deref_mut() necessary?
            &'i mut DBusInterface { &mut **iface }
        )
    }
}

impl<'a> Drop for DBusObject<'a> {
    fn drop(&mut self) {
        self.conn.unregister_object_path(self.path())
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::fmt;
use std::slice;

use ::Result;
use ::Error;

use octets;
use crypto;
use rand;
use stream;

const FORM_BIT: u8 = 0x80;
const KEY_PHASE_BIT: u8 = 0x40;
const DEMUX_BIT: u8 = 0x08;

const TYPE_MASK: u8 = 0x7f;

const MAX_CID_LEN: u8 = 18;

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Type {
    Initial,
    Retry,
    Handshake,
    ZeroRTT,
    VersionNegotiation,
    Application,
}

pub fn has_long_header(b: u8) -> bool {
    b & FORM_BIT != 0
}

#[derive(Clone)]
pub struct Header {
    pub ty: Type,
    pub version: u32,
    pub flags: u8,
    pub dcid: Vec<u8>,
    pub scid: Vec<u8>,
    pub token: Option<Vec<u8>>,
}

impl Header {
    pub fn decode_long(buf: &mut [u8]) -> Result<Header> {
        let mut b = octets::Bytes::new(buf);
        Header::long_from_bytes(&mut b)
    }

    pub fn decode_short(buf: &mut [u8], dcil: usize) -> Result<Header> {
        let mut b = octets::Bytes::new(buf);
        Header::short_from_bytes(&mut b, dcil)
    }

    pub fn long_from_bytes(b: &mut octets::Bytes) -> Result<Header> {
        let first = b.get_u8()?;

        if !has_long_header(first) {
            return Err(Error::WrongForm);
        }

        let version = b.get_u32()?;

        let ty = if version == 0 {
            Type::VersionNegotiation
        } else {
            match first & TYPE_MASK {
                0x7f => Type::Initial,
                0x7e => Type::Retry,
                0x7d => Type::Handshake,
                0x7c => Type::ZeroRTT,
                _    => return Err(Error::UnknownPacket),
            }
        };

        let (dcil, scil) = match b.get_u8() {
            Ok(v) => {
                let mut dcil = v >> 4;
                let mut scil = v & 0xf;

                if dcil > MAX_CID_LEN || scil > MAX_CID_LEN {
                    return Err(Error::InvalidPacket);
                }

                if dcil > 0 {
                    dcil += 3;
                }

                if scil > 0 {
                    scil += 3;
                }

                (dcil, scil)
            },

            Err(_) => return Err(Error::BufferTooShort),
        };

        let dcid = b.get_bytes(dcil as usize)?.to_vec();
        let scid = b.get_bytes(scil as usize)?.to_vec();

        // End of invariants.

        let mut token: Option<Vec<u8>> = None;

        match ty {
            Type::Initial => {
                token = Some(b.get_bytes_with_varint_length()?.to_vec());
            },

            Type::Retry => {
                panic!("Retry not supported");
            },

            Type::VersionNegotiation => {
                panic!("Retry not supported");
            },

            _ => (),
        };

        Ok(Header {
            ty,
            flags: first & FORM_BIT,
            version,
            dcid,
            scid,
            token,
        })
    }

    pub fn long_to_bytes(hdr: &Header, out: &mut octets::Bytes) -> Result<()> {
        let ty: u8 = match hdr.ty {
                Type::Initial   => 0x7f,
                Type::Retry     => 0x7e,
                Type::Handshake => 0x7d,
                Type::ZeroRTT   => 0x7c,
                // TODO: unify handling of version negotiation
                _               => return Err(Error::UnknownPacket),
        };

        let first = FORM_BIT | ty;

        out.put_u8(first)?;

        out.put_u32(hdr.version)?;

        let mut cil: u8 = 0;

        if hdr.dcid.len() > 0 {
            cil |= ((hdr.dcid.len() - 3) as u8) << 4;
        }

        if hdr.scid.len() > 0 {
            cil |= ((hdr.scid.len() - 3) as u8) & 0xf;
        }

        out.put_u8(cil)?;

        out.put_bytes(&hdr.dcid)?;
        out.put_bytes(&hdr.scid)?;

        // Only Initial packet have a token.
        if hdr.ty == Type::Initial {
            match hdr.token {
                Some(ref v) => {
                    out.put_bytes(v)?;
                },

                None => {
                    // No token, so lemgth = 0.
                    out.put_varint(0)?;
                }
            }
        }

        Ok(())
    }

    pub fn short_from_bytes(b: &mut octets::Bytes, dcil: usize) -> Result<Header> {
        let first = b.get_u8()?;

        if has_long_header(first) {
            return Err(Error::WrongForm);
        }

        let dcid = b.get_bytes(dcil)?;

        Ok(Header {
            ty: Type::Application,
            flags: first & FORM_BIT,
            version: 0,
            dcid: dcid.to_vec(),
            scid: Vec::new(),
            token: None,
        })
    }

    pub fn short_to_bytes(hdr: &Header, out: &mut octets::Bytes) -> Result<()> {
        let mut first = rand::rand_u8();

        // Unset form bit for short header.
        first &= !FORM_BIT;

        // TODO: support key update
        first &= !KEY_PHASE_BIT;

        // "The third bit (0x20) of octet 0 is set to 1."
        first |= 0x20;

        // "The fourth bit (0x10) of octet 0 is set to 1."
        first |= 0x10;

        // Clear Google QUIC demultiplexing bit
        first &= !DEMUX_BIT;

        out.put_u8(first)?;
        out.put_bytes(&hdr.dcid)?;

        Ok(())
    }
}

impl fmt::Debug for Header {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.ty);

        if self.ty != Type::Application {
            write!(f, " vers={:x}", self.version);
        }

        let vec: Vec<String> = self.dcid.iter()
                                        .map(|b| format!("{:02x}", b))
                                        .collect();

        write!(f, " dcid={}", vec.join(""));

        if self.ty != Type::Application {
            let vec: Vec<String> = self.scid.iter()
                                            .map(|b| format!("{:02x}", b))
                                            .collect();

            write!(f, " scid={}", vec.join(""));
        }

        Ok(())
    }
}

pub fn pkt_num_len(pn: u64) -> Result<usize> {
    let len = if pn < 128 {
        1
    } else if pn < 16384 {
        2
    } else if pn < 1_073_741_824 {
        4
    } else {
        return Err(Error::InvalidPacket);
    };

    Ok(len)
}

pub fn decrypt_pkt_num(b: &mut octets::Bytes, aead: &crypto::Open)
                                                    -> Result<(u64,usize)> {
    let max_pn_len = if b.cap() < aead.nonce_len() + 4 + 4 {
        b.cap() - aead.nonce_len() - 4
    } else {
        4
    };

    let mut pn_and_sample = b.peek_bytes(max_pn_len + aead.nonce_len() + 4)?;

    let (mut ciphertext, sample) = pn_and_sample.split_at(max_pn_len).unwrap();

    let ciphertext = ciphertext.as_mut();

    // Decrypt first byte of pkt num into separate buffer to get length.
    let mut first: u8 = ciphertext[0];

    aead.xor_keystream(sample.as_ref(), slice::from_mut(&mut first))?;

    let len = if first >> 7 == 0 {
        1
    } else {
        // Map most significant 2 bits to actual pkt num length.
        match first >> 6 {
            2 => 2,
            3 => 4,
            _ => return Err(Error::InvalidPacket),
        }
    };

    // Decrypt full pkt num in-place.
    aead.xor_keystream(sample.as_ref(), &mut ciphertext[..len])?;

    let mut plaintext = Vec::with_capacity(len);
    plaintext.extend_from_slice(&ciphertext[..len]);

    // Mask the 2 most significant bits to remove the encoded length.
    if len > 1 {
        plaintext[0] &= 0x3f;
    }

    let mut b = octets::Bytes::new(&mut plaintext);

    // Extract packet number corresponding to the decoded length.
    let out = match len {
        1 => u64::from(b.get_u8()?),
        2 => u64::from(b.get_u16()?),
        4 => u64::from(b.get_u32()?),
        _ => return Err(Error::InvalidPacket),
    };

    // TODO: implement packet number truncation
    Ok((out as u64, len))
}

pub fn encode_pkt_num(pn: u64, b: &mut octets::Bytes) -> Result<()> {
    let len = pkt_num_len(pn)?;

    match len {
        1 => {
            let mut buf = b.put_u8(pn as u8)?;
            buf[0] &= !0x80;
        },

        2 => {
            let mut buf = b.put_u16(pn as u16)?;
            buf[0] &= !0xc0;
            buf[0] |= 0x80;
        },

        4 => {
            let mut buf = b.put_u32(pn as u32)?;
            buf[0] |= 0xc0;
        },

        _ => return Err(Error::InvalidPacket),
    };

    Ok(())
}

pub fn decrypt_pkt(payload: &mut [u8], pn: u64, ad: &[u8], aead: &crypto::Open)
                                                            -> Result<usize> {
    aead.open_with_u64_counter(pn, ad, payload)
}

pub fn encrypt_pkt(payload: &mut [u8], pn: u64, ad: &[u8], aead: &crypto::Seal)
                                                            -> Result<usize> {
    aead.seal_with_u64_counter(pn, ad, payload)
}

pub fn negotiate_version(hdr: &Header, out: &mut [u8]) -> Result<usize> {
    let mut b = octets::Bytes::new(out);

    let first = rand::rand_u8() | FORM_BIT;

    b.put_u8(first)?;
    b.put_u32(0)?;

    // Invert client's scid and dcid.
    let mut cil: u8 = 0;
    if hdr.scid.len() > 0 {
        cil |= ((hdr.scid.len() - 3) as u8) << 4;
    }

    if hdr.dcid.len() > 0 {
        cil |= ((hdr.dcid.len() - 3) as u8) & 0xf;
    }

    b.put_u8(cil)?;
    b.put_bytes(&hdr.scid)?;
    b.put_bytes(&hdr.dcid)?;
    b.put_u32(::VERSION_DRAFT15)?;

    Ok(b.off())
}

pub struct PktNumSpace {
    pub pkt_type: Type,

    pub expected_pkt_num: u64,

    pub last_pkt_num: u64,

    pub need_ack: Vec<u64>,

    pub crypto_level: crypto::Level,

    pub crypto_open: Option<crypto::Open>,
    pub crypto_seal: Option<crypto::Seal>,

    pub crypto_stream: stream::Stream,
}

impl PktNumSpace {
    pub fn new(ty: Type, crypto_level: crypto::Level) -> PktNumSpace {
        PktNumSpace {
            pkt_type: ty,

            expected_pkt_num: 0,

            last_pkt_num: 0,

            need_ack: Vec::new(),

            crypto_level,

            crypto_open: None,
            crypto_seal: None,

            crypto_stream: stream::Stream::new(),
        }
    }

    pub fn cipher(&self) -> crypto::Algorithm {
        match self.crypto_open {
            Some(ref v) => v.algorithm(),
            None => crypto::Algorithm::Null,
        }
    }

    pub fn overhead(&self) -> usize {
        self.crypto_seal.as_ref().unwrap().tag_len()
    }
}

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

    use std::fs::File;
    use std::io::prelude::*;

    #[test]
    fn decrypt_initial() {
        decrypt_initial_helper("test_data/initial_ngtcp2.txt");
        decrypt_initial_helper("test_data/initial_quicly.txt");
    }

    fn decrypt_initial_helper(file_name: &'static str) {
        let mut file = File::open(file_name).unwrap();

        let mut buf = Vec::new();
        file.read_to_end(&mut buf).unwrap();

        let mut b = octets::Bytes::new(&mut buf);

        let hdr = Header::long_from_bytes(&mut b).unwrap();

        let payload_len = b.get_varint().unwrap() as usize;

        let (aead_open, _) =
            crypto::derive_initial_key_material(&hdr.dcid, true).unwrap();

        let (pn, pn_len) = decrypt_pkt_num(&mut b, &aead_open).unwrap();
        b.skip(pn_len).unwrap();

        let payload_offset = b.off();

        let (header, mut payload) = b.split_at(payload_offset).unwrap();

        decrypt_pkt(payload.slice(payload_len - pn_len).unwrap(),
                    pn, header.as_ref(), &aead_open).unwrap();
    }
}
<commit_msg>packet: fix error message<commit_after>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::fmt;
use std::slice;

use ::Result;
use ::Error;

use octets;
use crypto;
use rand;
use stream;

const FORM_BIT: u8 = 0x80;
const KEY_PHASE_BIT: u8 = 0x40;
const DEMUX_BIT: u8 = 0x08;

const TYPE_MASK: u8 = 0x7f;

const MAX_CID_LEN: u8 = 18;

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Type {
    Initial,
    Retry,
    Handshake,
    ZeroRTT,
    VersionNegotiation,
    Application,
}

pub fn has_long_header(b: u8) -> bool {
    b & FORM_BIT != 0
}

#[derive(Clone)]
pub struct Header {
    pub ty: Type,
    pub version: u32,
    pub flags: u8,
    pub dcid: Vec<u8>,
    pub scid: Vec<u8>,
    pub token: Option<Vec<u8>>,
}

impl Header {
    pub fn decode_long(buf: &mut [u8]) -> Result<Header> {
        let mut b = octets::Bytes::new(buf);
        Header::long_from_bytes(&mut b)
    }

    pub fn decode_short(buf: &mut [u8], dcil: usize) -> Result<Header> {
        let mut b = octets::Bytes::new(buf);
        Header::short_from_bytes(&mut b, dcil)
    }

    pub fn long_from_bytes(b: &mut octets::Bytes) -> Result<Header> {
        let first = b.get_u8()?;

        if !has_long_header(first) {
            return Err(Error::WrongForm);
        }

        let version = b.get_u32()?;

        let ty = if version == 0 {
            Type::VersionNegotiation
        } else {
            match first & TYPE_MASK {
                0x7f => Type::Initial,
                0x7e => Type::Retry,
                0x7d => Type::Handshake,
                0x7c => Type::ZeroRTT,
                _    => return Err(Error::UnknownPacket),
            }
        };

        let (dcil, scil) = match b.get_u8() {
            Ok(v) => {
                let mut dcil = v >> 4;
                let mut scil = v & 0xf;

                if dcil > MAX_CID_LEN || scil > MAX_CID_LEN {
                    return Err(Error::InvalidPacket);
                }

                if dcil > 0 {
                    dcil += 3;
                }

                if scil > 0 {
                    scil += 3;
                }

                (dcil, scil)
            },

            Err(_) => return Err(Error::BufferTooShort),
        };

        let dcid = b.get_bytes(dcil as usize)?.to_vec();
        let scid = b.get_bytes(scil as usize)?.to_vec();

        // End of invariants.

        let mut token: Option<Vec<u8>> = None;

        match ty {
            Type::Initial => {
                token = Some(b.get_bytes_with_varint_length()?.to_vec());
            },

            Type::Retry => {
                panic!("Retry not supported");
            },

            Type::VersionNegotiation => {
                panic!("Version negotiation not supported");
            },

            _ => (),
        };

        Ok(Header {
            ty,
            flags: first & FORM_BIT,
            version,
            dcid,
            scid,
            token,
        })
    }

    pub fn long_to_bytes(hdr: &Header, out: &mut octets::Bytes) -> Result<()> {
        let ty: u8 = match hdr.ty {
                Type::Initial   => 0x7f,
                Type::Retry     => 0x7e,
                Type::Handshake => 0x7d,
                Type::ZeroRTT   => 0x7c,
                // TODO: unify handling of version negotiation
                _               => return Err(Error::UnknownPacket),
        };

        let first = FORM_BIT | ty;

        out.put_u8(first)?;

        out.put_u32(hdr.version)?;

        let mut cil: u8 = 0;

        if hdr.dcid.len() > 0 {
            cil |= ((hdr.dcid.len() - 3) as u8) << 4;
        }

        if hdr.scid.len() > 0 {
            cil |= ((hdr.scid.len() - 3) as u8) & 0xf;
        }

        out.put_u8(cil)?;

        out.put_bytes(&hdr.dcid)?;
        out.put_bytes(&hdr.scid)?;

        // Only Initial packet have a token.
        if hdr.ty == Type::Initial {
            match hdr.token {
                Some(ref v) => {
                    out.put_bytes(v)?;
                },

                None => {
                    // No token, so lemgth = 0.
                    out.put_varint(0)?;
                }
            }
        }

        Ok(())
    }

    pub fn short_from_bytes(b: &mut octets::Bytes, dcil: usize) -> Result<Header> {
        let first = b.get_u8()?;

        if has_long_header(first) {
            return Err(Error::WrongForm);
        }

        let dcid = b.get_bytes(dcil)?;

        Ok(Header {
            ty: Type::Application,
            flags: first & FORM_BIT,
            version: 0,
            dcid: dcid.to_vec(),
            scid: Vec::new(),
            token: None,
        })
    }

    pub fn short_to_bytes(hdr: &Header, out: &mut octets::Bytes) -> Result<()> {
        let mut first = rand::rand_u8();

        // Unset form bit for short header.
        first &= !FORM_BIT;

        // TODO: support key update
        first &= !KEY_PHASE_BIT;

        // "The third bit (0x20) of octet 0 is set to 1."
        first |= 0x20;

        // "The fourth bit (0x10) of octet 0 is set to 1."
        first |= 0x10;

        // Clear Google QUIC demultiplexing bit
        first &= !DEMUX_BIT;

        out.put_u8(first)?;
        out.put_bytes(&hdr.dcid)?;

        Ok(())
    }
}

impl fmt::Debug for Header {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.ty);

        if self.ty != Type::Application {
            write!(f, " vers={:x}", self.version);
        }

        let vec: Vec<String> = self.dcid.iter()
                                        .map(|b| format!("{:02x}", b))
                                        .collect();

        write!(f, " dcid={}", vec.join(""));

        if self.ty != Type::Application {
            let vec: Vec<String> = self.scid.iter()
                                            .map(|b| format!("{:02x}", b))
                                            .collect();

            write!(f, " scid={}", vec.join(""));
        }

        Ok(())
    }
}

pub fn pkt_num_len(pn: u64) -> Result<usize> {
    let len = if pn < 128 {
        1
    } else if pn < 16384 {
        2
    } else if pn < 1_073_741_824 {
        4
    } else {
        return Err(Error::InvalidPacket);
    };

    Ok(len)
}

pub fn decrypt_pkt_num(b: &mut octets::Bytes, aead: &crypto::Open)
                                                    -> Result<(u64,usize)> {
    let max_pn_len = if b.cap() < aead.nonce_len() + 4 + 4 {
        b.cap() - aead.nonce_len() - 4
    } else {
        4
    };

    let mut pn_and_sample = b.peek_bytes(max_pn_len + aead.nonce_len() + 4)?;

    let (mut ciphertext, sample) = pn_and_sample.split_at(max_pn_len).unwrap();

    let ciphertext = ciphertext.as_mut();

    // Decrypt first byte of pkt num into separate buffer to get length.
    let mut first: u8 = ciphertext[0];

    aead.xor_keystream(sample.as_ref(), slice::from_mut(&mut first))?;

    let len = if first >> 7 == 0 {
        1
    } else {
        // Map most significant 2 bits to actual pkt num length.
        match first >> 6 {
            2 => 2,
            3 => 4,
            _ => return Err(Error::InvalidPacket),
        }
    };

    // Decrypt full pkt num in-place.
    aead.xor_keystream(sample.as_ref(), &mut ciphertext[..len])?;

    let mut plaintext = Vec::with_capacity(len);
    plaintext.extend_from_slice(&ciphertext[..len]);

    // Mask the 2 most significant bits to remove the encoded length.
    if len > 1 {
        plaintext[0] &= 0x3f;
    }

    let mut b = octets::Bytes::new(&mut plaintext);

    // Extract packet number corresponding to the decoded length.
    let out = match len {
        1 => u64::from(b.get_u8()?),
        2 => u64::from(b.get_u16()?),
        4 => u64::from(b.get_u32()?),
        _ => return Err(Error::InvalidPacket),
    };

    // TODO: implement packet number truncation
    Ok((out as u64, len))
}

pub fn encode_pkt_num(pn: u64, b: &mut octets::Bytes) -> Result<()> {
    let len = pkt_num_len(pn)?;

    match len {
        1 => {
            let mut buf = b.put_u8(pn as u8)?;
            buf[0] &= !0x80;
        },

        2 => {
            let mut buf = b.put_u16(pn as u16)?;
            buf[0] &= !0xc0;
            buf[0] |= 0x80;
        },

        4 => {
            let mut buf = b.put_u32(pn as u32)?;
            buf[0] |= 0xc0;
        },

        _ => return Err(Error::InvalidPacket),
    };

    Ok(())
}

pub fn decrypt_pkt(payload: &mut [u8], pn: u64, ad: &[u8], aead: &crypto::Open)
                                                            -> Result<usize> {
    aead.open_with_u64_counter(pn, ad, payload)
}

pub fn encrypt_pkt(payload: &mut [u8], pn: u64, ad: &[u8], aead: &crypto::Seal)
                                                            -> Result<usize> {
    aead.seal_with_u64_counter(pn, ad, payload)
}

pub fn negotiate_version(hdr: &Header, out: &mut [u8]) -> Result<usize> {
    let mut b = octets::Bytes::new(out);

    let first = rand::rand_u8() | FORM_BIT;

    b.put_u8(first)?;
    b.put_u32(0)?;

    // Invert client's scid and dcid.
    let mut cil: u8 = 0;
    if hdr.scid.len() > 0 {
        cil |= ((hdr.scid.len() - 3) as u8) << 4;
    }

    if hdr.dcid.len() > 0 {
        cil |= ((hdr.dcid.len() - 3) as u8) & 0xf;
    }

    b.put_u8(cil)?;
    b.put_bytes(&hdr.scid)?;
    b.put_bytes(&hdr.dcid)?;
    b.put_u32(::VERSION_DRAFT15)?;

    Ok(b.off())
}

pub struct PktNumSpace {
    pub pkt_type: Type,

    pub expected_pkt_num: u64,

    pub last_pkt_num: u64,

    pub need_ack: Vec<u64>,

    pub crypto_level: crypto::Level,

    pub crypto_open: Option<crypto::Open>,
    pub crypto_seal: Option<crypto::Seal>,

    pub crypto_stream: stream::Stream,
}

impl PktNumSpace {
    pub fn new(ty: Type, crypto_level: crypto::Level) -> PktNumSpace {
        PktNumSpace {
            pkt_type: ty,

            expected_pkt_num: 0,

            last_pkt_num: 0,

            need_ack: Vec::new(),

            crypto_level,

            crypto_open: None,
            crypto_seal: None,

            crypto_stream: stream::Stream::new(),
        }
    }

    pub fn cipher(&self) -> crypto::Algorithm {
        match self.crypto_open {
            Some(ref v) => v.algorithm(),
            None => crypto::Algorithm::Null,
        }
    }

    pub fn overhead(&self) -> usize {
        self.crypto_seal.as_ref().unwrap().tag_len()
    }
}

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

    use std::fs::File;
    use std::io::prelude::*;

    #[test]
    fn decrypt_initial() {
        decrypt_initial_helper("test_data/initial_ngtcp2.txt");
        decrypt_initial_helper("test_data/initial_quicly.txt");
    }

    fn decrypt_initial_helper(file_name: &'static str) {
        let mut file = File::open(file_name).unwrap();

        let mut buf = Vec::new();
        file.read_to_end(&mut buf).unwrap();

        let mut b = octets::Bytes::new(&mut buf);

        let hdr = Header::long_from_bytes(&mut b).unwrap();

        let payload_len = b.get_varint().unwrap() as usize;

        let (aead_open, _) =
            crypto::derive_initial_key_material(&hdr.dcid, true).unwrap();

        let (pn, pn_len) = decrypt_pkt_num(&mut b, &aead_open).unwrap();
        b.skip(pn_len).unwrap();

        let payload_offset = b.off();

        let (header, mut payload) = b.split_at(payload_offset).unwrap();

        decrypt_pkt(payload.slice(payload_len - pn_len).unwrap(),
                    pn, header.as_ref(), &aead_open).unwrap();
    }
}
<|endoftext|>