<commit_before>//! call rank and take command of slack with rust at your helm

#[macro_use]
extern crate log;
extern crate hyper;
extern crate url;
extern crate regex;
extern crate serde;
extern crate serde_json;

mod rep;

pub use rep::{Response, ResponseBuilder};
use regex::{Captures as RegexCaptures, Regex};

use hyper::Client;
use hyper::server::{Handler as HyperHandler, Request, Response as HyperResponse};
use std::collections::HashMap;
use std::io::Read;

const DEFAULT_RESPONSE: &'static [u8] =  b"ok";

fn params<R: Read>(read: &mut R) -> HashMap<String, String> {
    let mut buffer = String::new();
    read.read_to_string(&mut buffer).unwrap();
    let mut params = HashMap::new();
    for (k, v) in url::form_urlencoded::parse(buffer.as_bytes()) {
        params.insert(k, v);
    }
    params
}

pub type Captures<'a> = RegexCaptures<'a>;

/// Responds to commands after some deferred time
/// This allows command handlers to asyncronously
/// response to commands
pub trait Responder: Sync + Send {
    fn respond(&self, response: Response) -> ();
}

pub struct DefaultResponder {
    response_url: String,
}

impl Responder for DefaultResponder {
    fn respond(&self, response: Response) {
        let client = Client::new();
        let _ = client.post(&self.response_url[..])
                      .body(serde_json::to_string(&response).unwrap().as_bytes())
                      .send();
    }
}

pub struct DiscardResponder;

impl Responder for DiscardResponder {
    fn respond(&self, _: Response) {}
}

/// Handles matched commands
pub trait Handler: Sync + Send {
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response>;
}

impl<F> Handler for F
    where F: Fn(&Command, &Option<Captures>, Box<Responder>) -> Option<Response>,
          F: Send + Sync
{
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response> {
        self(cmd, caps, responder)
    }
}

pub struct TokenValidator<H: Handler + 'static> {
    handler: H,
    token: String,
}

impl<H: Handler + 'static> Handler for TokenValidator<H> {
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response> {
        if cmd.token == self.token {
            self.handler.handle(cmd, caps, responder)
        } else {
            None
        }
    }
}

/// A matcher matches commands
pub trait Matcher: Send + Sync {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool);
}

impl<F> Matcher for F
    where F: Fn(&Command) -> (Option<Captures>, bool),
          F: Send + Sync
{
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        self(cmd)
    }
}

/// A direct command matcher
///
pub struct MatchCommand(String);

impl Matcher for MatchCommand {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        (None, cmd.command == self.0)
    }
}

/// A regex pattern matcher
///
pub struct MatchRegex(Regex);

impl Matcher for MatchRegex {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        if self.0.is_match(cmd.command.as_ref()) {
            (self.0.captures(cmd.command.as_ref()), true)
        } else {
            (None, false)
        }
    }
}

/// A binding between a command handler and which commands its handles
pub struct Route {
    handler: Box<Handler>,
    matcher: Box<Matcher>,
}

///
/// A command de-multiplexor
///
#[derive(Default)]
pub struct Mux {
    routes: Vec<Box<Route>>,
}

impl Mux {
    pub fn new() -> Mux {
        Mux { ..Default::default() }
    }

    pub fn command<C, T, H>(&mut self, cmd: C, token: T, handler: H)
        where C: Into<String>,
              T: Into<String>,
              H: Handler + 'static
    {
        self.matching(MatchCommand(cmd.into()),
                      TokenValidator {
                          handler: handler,
                          token: token.into(),
                      })
    }

    pub fn matching<M, H>(&mut self, matcher: M, handler: H)
        where M: Matcher + 'static,
              H: Handler + 'static
    {
        let route = Route {
            handler: Box::new(handler),
            matcher: Box::new(matcher),
        };
        self.route(route)
    }

    ///
    pub fn route(&mut self, route: Route) {
        self.routes.push(Box::new(route));
    }

    pub fn handler<'a>(&self, cmd: &'a Command) -> Option<(Option<Captures<'a>>, &Box<Handler>)> {
        for r in self.routes.iter() {
            if let (captures, true) = r.matcher.matches(cmd) {
                return Some((captures, &r.handler));
            }
        }
        None
    }

    pub fn handle(&self, cmd: &Command) -> Option<Response> {
        // set up responder
        let responder = DefaultResponder { response_url: cmd.response_url.clone() };
        // handle cmd
        if let &Some((ref captures, handler)) = &self.handler(&cmd) {
            info!("attempting to handle cmd ${:?}", cmd);
            handler.handle(&cmd, &captures, Box::new(responder))
        } else {
            None
        }
    }
}

#[derive(Default, Debug, PartialEq)]
pub struct Command {
    pub token: String,
    pub team_id: String,
    pub team_domain: String,
    pub channel_id: String,
    pub channel_name: String,
    pub user_id: String,
    pub user_name: String,
    pub command: String,
    pub text: String,
    pub response_url: String,
}

impl Command {
    pub fn from_params(params: HashMap<String, String>) -> Option<Command> {
        if let (Some(token),
                Some(team_id),
                Some(team_domain),
                Some(channel_id),
                Some(channel_name),
                Some(user_id),
                Some(user_name),
                Some(command),
                Some(text),
                Some(response_url)) = (params.get("token"),
                                       params.get("team_id"),
                                       params.get("team_domain"),
                                       params.get("channel_id"),
                                       params.get("channel_name"),
                                       params.get("user_id"),
                                       params.get("user_name"),
                                       params.get("command"),
                                       params.get("text"),
                                       params.get("response_url")) {
            Some(Command {
                token: token.clone(),
                team_id: team_id.clone(),
                team_domain: team_domain.clone(),
                channel_id: channel_id.clone(),
                channel_name: channel_name.clone(),
                user_id: user_id.clone(),
                user_name: user_name.clone(),
                command: command.clone(),
                text: text.clone(),
                response_url: response_url.clone(),
            })
        } else {
            None
        }
    }
}

impl HyperHandler for Mux {
    // https://api.slack.com/slash-commands
    fn handle(&self, req: Request, res: HyperResponse) {
        let (_, _, _, _, _, mut body) = req.deconstruct();
        // parse params
        let params = params(&mut body);
        // parse cmd
        if let Some(cmd) = Command::from_params(params) {
            info!("rec cmd {:?}", cmd);
            let write = |bytes: &[u8]| {
                let _ = res.send(bytes);
            };
            if let Some(resp) = self.handle(&cmd) {
                match serde_json::to_string(&resp) {
                    Ok(payload) => write(payload.as_bytes()),
                    _ => write(DEFAULT_RESPONSE),
                };
            } else {
                write(DEFAULT_RESPONSE)
            };
        } else {
            error!("rec invalid cmd");
            let _ = res.send(DEFAULT_RESPONSE);
        }
        ()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::regex::Regex;
    use std::collections::HashMap;

    #[test]
    fn matches_commands() {
        let cmd = Command { command: "/test".to_owned(), ..Default::default() };
        let (_, matched) = MatchCommand("/test".to_owned()).matches(&cmd);
        assert!(matched)
    }

    #[test]
    fn matches_regexes() {
        let cmd = Command { command: "/test".to_owned(), ..Default::default() };
        let (_, matched) = MatchRegex(Regex::new("(test)").unwrap()).matches(&cmd);
        assert!(matched)
    }

    #[test]
    fn extracts_commands() {
        let mut params = HashMap::new();
        params.insert("token".to_owned(), "test_token".to_owned());
        params.insert("team_id".to_owned(), "test_team".to_owned());
        params.insert("team_domain".to_owned(), "test_team_domain".to_owned());
        params.insert("channel_id".to_owned(), "test_channel_id".to_owned());
        params.insert("channel_name".to_owned(), "test_channel_name".to_owned());
        params.insert("user_id".to_owned(), "test_user_id".to_owned());
        params.insert("user_name".to_owned(), "test_user_name".to_owned());
        params.insert("command".to_owned(), "test_command".to_owned());
        params.insert("text".to_owned(), "test_text".to_owned());
        params.insert("response_url".to_owned(), "test_response_url".to_owned());
        match Command::from_params(params) {
            Some(cmd) => assert_eq!(cmd, Command {
                token: "test_token".to_owned(),
                team_id: "test_team".to_owned(),
                team_domain: "test_team_domain".to_owned(),
                channel_id: "test_channel_id".to_owned(),
                channel_name: "test_channel_name".to_owned(),
                user_id: "test_user_id".to_owned(),
                user_name: "test_user_name".to_owned(),
                command: "test_command".to_owned(),
                text: "test_text".to_owned(),
                response_url: "test_response_url".to_owned()
            }),
            _ => assert!(false)
        }
    }
}
<commit_msg>write out content type<commit_after>//! call rank and take command of slack with rust at your helm

#[macro_use]
extern crate log;
extern crate hyper;
extern crate url;
extern crate regex;
extern crate serde;
extern crate serde_json;

mod rep;

pub use rep::{Response, ResponseBuilder};
use regex::{Captures as RegexCaptures, Regex};

use hyper::Client;
use hyper::header::ContentType;
use hyper::server::{Handler as HyperHandler, Request, Response as HyperResponse};
use std::collections::HashMap;
use std::io::Read;

const DEFAULT_RESPONSE: &'static [u8] =  b"ok";

fn params<R: Read>(read: &mut R) -> HashMap<String, String> {
    let mut buffer = String::new();
    read.read_to_string(&mut buffer).unwrap();
    let mut params = HashMap::new();
    for (k, v) in url::form_urlencoded::parse(buffer.as_bytes()) {
        params.insert(k, v);
    }
    params
}

pub type Captures<'a> = RegexCaptures<'a>;

/// Responds to commands after some deferred time
/// This allows command handlers to asyncronously
/// response to commands
pub trait Responder: Sync + Send {
    fn respond(&self, response: Response) -> ();
}

pub struct DefaultResponder {
    response_url: String,
}

impl Responder for DefaultResponder {
    fn respond(&self, response: Response) {
        let client = Client::new();
        let _ = client.post(&self.response_url[..])
                      .body(serde_json::to_string(&response).unwrap().as_bytes())
                      .send();
    }
}

pub struct DiscardResponder;

impl Responder for DiscardResponder {
    fn respond(&self, _: Response) {}
}

/// Handles matched commands
pub trait Handler: Sync + Send {
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response>;
}

impl<F> Handler for F
    where F: Fn(&Command, &Option<Captures>, Box<Responder>) -> Option<Response>,
          F: Send + Sync
{
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response> {
        self(cmd, caps, responder)
    }
}

pub struct TokenValidator<H: Handler + 'static> {
    handler: H,
    token: String,
}

impl<H: Handler + 'static> Handler for TokenValidator<H> {
    fn handle(&self,
              cmd: &Command,
              caps: &Option<Captures>,
              responder: Box<Responder>)
              -> Option<Response> {
        if cmd.token == self.token {
            self.handler.handle(cmd, caps, responder)
        } else {
            None
        }
    }
}

/// A matcher matches commands
pub trait Matcher: Send + Sync {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool);
}

impl<F> Matcher for F
    where F: Fn(&Command) -> (Option<Captures>, bool),
          F: Send + Sync
{
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        self(cmd)
    }
}

/// A direct command matcher
///
pub struct MatchCommand(String);

impl Matcher for MatchCommand {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        (None, cmd.command == self.0)
    }
}

/// A regex pattern matcher
///
pub struct MatchRegex(Regex);

impl Matcher for MatchRegex {
    fn matches<'a>(&self, cmd: &'a Command) -> (Option<Captures<'a>>, bool) {
        if self.0.is_match(cmd.command.as_ref()) {
            (self.0.captures(cmd.command.as_ref()), true)
        } else {
            (None, false)
        }
    }
}

/// A binding between a command handler and which commands its handles
pub struct Route {
    handler: Box<Handler>,
    matcher: Box<Matcher>,
}

///
/// A command de-multiplexor
///
#[derive(Default)]
pub struct Mux {
    routes: Vec<Box<Route>>,
}

impl Mux {
    pub fn new() -> Mux {
        Mux { ..Default::default() }
    }

    pub fn command<C, T, H>(&mut self, cmd: C, token: T, handler: H)
        where C: Into<String>,
              T: Into<String>,
              H: Handler + 'static
    {
        self.matching(MatchCommand(cmd.into()),
                      TokenValidator {
                          handler: handler,
                          token: token.into(),
                      })
    }

    pub fn matching<M, H>(&mut self, matcher: M, handler: H)
        where M: Matcher + 'static,
              H: Handler + 'static
    {
        let route = Route {
            handler: Box::new(handler),
            matcher: Box::new(matcher),
        };
        self.route(route)
    }

    ///
    pub fn route(&mut self, route: Route) {
        self.routes.push(Box::new(route));
    }

    pub fn handler<'a>(&self, cmd: &'a Command) -> Option<(Option<Captures<'a>>, &Box<Handler>)> {
        for r in self.routes.iter() {
            if let (captures, true) = r.matcher.matches(cmd) {
                return Some((captures, &r.handler));
            }
        }
        None
    }

    pub fn handle(&self, cmd: &Command) -> Option<Response> {
        // set up responder
        let responder = DefaultResponder { response_url: cmd.response_url.clone() };
        // handle cmd
        if let &Some((ref captures, handler)) = &self.handler(&cmd) {
            info!("attempting to handle cmd ${:?}", cmd);
            handler.handle(&cmd, &captures, Box::new(responder))
        } else {
            None
        }
    }
}

#[derive(Default, Debug, PartialEq)]
pub struct Command {
    pub token: String,
    pub team_id: String,
    pub team_domain: String,
    pub channel_id: String,
    pub channel_name: String,
    pub user_id: String,
    pub user_name: String,
    pub command: String,
    pub text: String,
    pub response_url: String,
}

impl Command {
    pub fn from_params(params: HashMap<String, String>) -> Option<Command> {
        if let (Some(token),
                Some(team_id),
                Some(team_domain),
                Some(channel_id),
                Some(channel_name),
                Some(user_id),
                Some(user_name),
                Some(command),
                Some(text),
                Some(response_url)) = (params.get("token"),
                                       params.get("team_id"),
                                       params.get("team_domain"),
                                       params.get("channel_id"),
                                       params.get("channel_name"),
                                       params.get("user_id"),
                                       params.get("user_name"),
                                       params.get("command"),
                                       params.get("text"),
                                       params.get("response_url")) {
            Some(Command {
                token: token.clone(),
                team_id: team_id.clone(),
                team_domain: team_domain.clone(),
                channel_id: channel_id.clone(),
                channel_name: channel_name.clone(),
                user_id: user_id.clone(),
                user_name: user_name.clone(),
                command: command.clone(),
                text: text.clone(),
                response_url: response_url.clone(),
            })
        } else {
            None
        }
    }
}

impl HyperHandler for Mux {
    // https://api.slack.com/slash-commands
    fn handle(&self, req: Request, mut res: HyperResponse) {
        let (_, _, _, _, _, mut body) = req.deconstruct();

        // parse params
        let params = params(&mut body);
        // parse cmd
        if let Some(cmd) = Command::from_params(params) {
            info!("rec cmd {:?}", cmd);
            let write = |bytes: &[u8], content_type: ContentType| {
                res.headers_mut().set(content_type);
                let _ = res.send(bytes);
            };
            if let Some(resp) = self.handle(&cmd) {
                match serde_json::to_string(&resp) {
                    Ok(payload) => write(payload.as_bytes(), ContentType::json()),
                    _ => write(DEFAULT_RESPONSE, ContentType::plaintext()),
                };
            } else {
                write(DEFAULT_RESPONSE, ContentType::plaintext())
            };
        } else {
            error!("rec invalid cmd");
            let _ = res.send(DEFAULT_RESPONSE);
        }
        ()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::regex::Regex;
    use std::collections::HashMap;

    #[test]
    fn matches_commands() {
        let cmd = Command { command: "/test".to_owned(), ..Default::default() };
        let (_, matched) = MatchCommand("/test".to_owned()).matches(&cmd);
        assert!(matched)
    }

    #[test]
    fn matches_regexes() {
        let cmd = Command { command: "/test".to_owned(), ..Default::default() };
        let (_, matched) = MatchRegex(Regex::new("(test)").unwrap()).matches(&cmd);
        assert!(matched)
    }

    #[test]
    fn extracts_commands() {
        let mut params = HashMap::new();
        params.insert("token".to_owned(), "test_token".to_owned());
        params.insert("team_id".to_owned(), "test_team".to_owned());
        params.insert("team_domain".to_owned(), "test_team_domain".to_owned());
        params.insert("channel_id".to_owned(), "test_channel_id".to_owned());
        params.insert("channel_name".to_owned(), "test_channel_name".to_owned());
        params.insert("user_id".to_owned(), "test_user_id".to_owned());
        params.insert("user_name".to_owned(), "test_user_name".to_owned());
        params.insert("command".to_owned(), "test_command".to_owned());
        params.insert("text".to_owned(), "test_text".to_owned());
        params.insert("response_url".to_owned(), "test_response_url".to_owned());
        match Command::from_params(params) {
            Some(cmd) => assert_eq!(cmd, Command {
                token: "test_token".to_owned(),
                team_id: "test_team".to_owned(),
                team_domain: "test_team_domain".to_owned(),
                channel_id: "test_channel_id".to_owned(),
                channel_name: "test_channel_name".to_owned(),
                user_id: "test_user_id".to_owned(),
                user_name: "test_user_name".to_owned(),
                command: "test_command".to_owned(),
                text: "test_text".to_owned(),
                response_url: "test_response_url".to_owned()
            }),
            _ => assert!(false)
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;
extern crate byteorder;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;

use daemonize::{Daemonize};
use std::process::Command;
use std::net::{UdpSocket, SocketAddr, SocketAddrV4};
use std::str;
use std::env;
use std::fmt;
use std::str::FromStr;
use std::fs::{OpenOptions, File};
use std::path::Path;
use std::ffi::OsStr;
use std::io::{Cursor, Error, Seek, SeekFrom, stderr, Read, Write};
use udt::{UdtSocket, UdtOpts, SocketType, SocketFamily};
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::{NONCEBYTES, Key, Nonce};
use rustc_serialize::hex::{FromHex, ToHex};

macro_rules! overprint {
    ($fmt: expr) => {
        print!(concat!("\x1b[2K\r", $fmt));
        std::io::stdout().flush().unwrap();
    };
    ($fmt:expr, $($arg:tt)*) => {
        print!(concat!("\x1b[2K\r", $fmt) , $($arg)*);
        std::io::stdout().flush().unwrap();
    };
}

pub struct Server<'a> {
    filename: &'a str,
    sock: UdtSocket,
    key: Key,
}

pub struct Client<'a> {
    remote_host: &'a str,
    port_range: PortRange,
    remote_path: &'a str,
}

pub struct PortRange {
    start: u16,
    end: u16
}

enum ShoopErrKind {
    Severed,
    Fatal,
}
struct ShoopErr {
    kind: ShoopErrKind,
    msg: Option<String>,
    finished: u64,
}

impl ShoopErr {
    pub fn new(kind: ShoopErrKind, msg: &str, finished: u64) -> ShoopErr {
        ShoopErr { kind: kind, msg: Some(String::from(msg)), finished: finished }
    }

    pub fn from(err: Error, finished: u64) -> ShoopErr {
        ShoopErr { kind: ShoopErrKind::Severed, msg: Some(format!("{:?}", err)), finished: finished }
    }
}

impl PortRange {
    fn new(start: u16, end: u16) -> PortRange {
        PortRange{ start: start, end: end }
    }

    pub fn from(s: &str) -> Result<PortRange, &str> {
        let sections: Vec<&str> = s.split("-").collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"")
        }
        let (start, end) = (sections[0].parse::<u16>(),
                            sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("invalid port range");
        }
        Ok(PortRange::new(start.unwrap(), end.unwrap()))
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}

impl<'a> Server<'a> {
    pub fn new(port_range : PortRange,
               filename   : &str)
        -> Server
    {
        let sshconnstr = env::var("SSH_CONNECTION").unwrap_or(String::from("0.0.0.0 0 0.0.0.0 22")).trim().to_owned();
        let sshconn: Vec<&str> = sshconnstr.split(" ").collect();
        let ip = sshconn[2];
        let port = Server::get_open_port(port_range).unwrap();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());
        let daemonize = Daemonize::new();
        match daemonize.start() {
            Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            Err(_) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
        }

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        sock.bind(SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap())).unwrap();
        Server { sock: sock, key: key, filename: filename }
    }

    pub fn start(&self) {
        self.sock.listen(1).unwrap();

        let (stream, _) = self.sock.accept().unwrap();

        loop {
            if let Ok(starthdr) = stream.recvmsg(9) {
                let version = starthdr[0];
                let mut rdr = Cursor::new(starthdr);
                rdr.set_position(1);
                let offset = rdr.read_u64::<LittleEndian>().unwrap();
                if version == 0x00 {
                    match self.send_file(stream, offset) {
                        Ok(_) => {
                            break;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Severed, msg: _, finished: _}) => {
                            continue;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Fatal, msg, finished: _}) => {
                            panic!("{:?}", msg);
                        }
                    }
                } else {
                    panic!("Unrecognized version.");
                }
            } else {
                panic!("Failed to receive version byte from client.");
            }
        }
    }

    fn send_file(&self, stream: UdtSocket, offset: u64) -> Result<(), ShoopErr> {
        let mut f = File::open(self.filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();

        let mut remaining = 0u64;
        let mut buf = vec![0; 1024 * 1024];
        loop {
            match try!(f.read(&mut buf).map_err(|e| ShoopErr::from(e, 0u64))) {
                0    => { break },
                read => { remaining += read as u64 },
            }
        }

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(remaining).unwrap();
        match stream.sendmsg(&wtr[..]) {
            Ok(0) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "failed to write filesize header before timeout", remaining))
            },
            Err(e) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), remaining))
            }
            _ => {}
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(offset)).unwrap();
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    let nonce = secretbox::gen_nonce();
                    let Nonce(noncebytes) = nonce;
                    let mut hdr = vec![0u8; 1 + NONCEBYTES];
                    hdr[0] = NONCEBYTES as u8;
                    hdr[1..].clone_from_slice(&noncebytes);

                    let mut sealed = secretbox::seal(&payload[0..read], &nonce, &self.key);
                    let mut msg = Vec::with_capacity(hdr.len() + sealed.len());
                    msg.extend_from_slice(&hdr);
                    msg.append(&mut sealed);
                    match stream.sendmsg(&msg[..]) {
                        Ok(_) =>  { }
                        Err(e) => {
                            stream.close().expect("Error closing stream");
                            panic!("{:?}", e);
                        }
                    }
                },
                Err(e) => {
                    stream.close().expect("Error closing stream");
                    panic!("{:?}", e);
                }
            }
        }

        stream.close().expect("Error closing stream.");
        Ok(())
    }

    fn get_open_port(range: PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }
}

impl<'a> Client<'a> {
    pub fn new(remote_ssh_host : &'a str,
               port_range      : PortRange,
               remote_path     : &'a str)
        -> Client<'a>
    {
        Client{ remote_host: remote_ssh_host, port_range: port_range, remote_path: remote_path }
    }

    pub fn start(&self) {
        let cmd = format!("shoop -s '{}' -p {}", self.remote_path, self.port_range);
        // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

        overprint!(" - establishing SSH session...");
        assert!(Client::command_exists("ssh"), "`ssh` is required!");
        let output = Command::new("ssh")
                             .arg(self.remote_host.to_owned())
                             .arg(cmd)
                             .output()
                             .unwrap_or_else(|e| {
                                 panic!("failed to execute process: {}", e);
                             });
        let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
        let info: Vec<&str> = infostring.split(" ").collect();
        if info.len() != 5 {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
        overprint!(" - opening UDT connection...");
        if magic != "shoop" || version != "0" {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let mut keybytes = [0u8; 32];
        keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
        let key = Key(keybytes);

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap());
        match sock.connect(addr) {
           Ok(()) => {
               overprint!(" - connection opened, shakin' hands, makin' frands");
           },
           Err(e) => {
               panic!("errrrrrrr connecting to {}:{} - {:?}", ip, port, e);
           }
        }

        let mut offset = 0u64;
        loop {
            let mut wtr = vec![];
            wtr.push(0);
            wtr.write_u64::<LittleEndian>(offset).unwrap();
            sock.sendmsg(&wtr[..]).unwrap();

            match sock.recvmsg(8) {
               Ok(msg) => {
                   if msg.len() == 0 {
                       panic!("failed to get filesize from server, probable timeout.");
                   }
                   let mut rdr = Cursor::new(msg);
                   let filesize = rdr.read_u64::<LittleEndian>().unwrap();
                   let filename = Path::new(&self.remote_path).file_name().unwrap_or(OsStr::new("outfile")).to_str().unwrap_or("outfile");
                   overprint!(" + downloading {} ({:.1}MB)\n", filename, (filesize as f64)/(1024f64*1024f64));
                   match self.recv_file(sock, filesize, filename, &key, offset) {
                        Ok(_) => {
                            break;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Severed, msg: _, finished}) => {
                            offset = finished;
                            continue;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Fatal, msg, finished: _}) => {
                            panic!("{:?}", msg);
                        }
                   }
               }
               Err(e) => {
                   panic!("{:?}", e);
               }
            }
        }
    }

    fn command_exists(command: &str) -> bool {
        match Command::new("which").arg(command).output() {
            Ok(output) => output.status.success(),
            Err(_)     => false
        }
    }

    fn recv_file(&self, sock: UdtSocket, filesize: u64, filename: &str, key: &Key, offset: u64) -> Result<(), ShoopErr> {
        let mut f = OpenOptions::new().write(true).create(true).truncate(false).open(filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();
        let mut ts = time::precise_time_ns();
        let mut total = 0u64;
        loop {
            let buf = try!(sock.recvmsg(8192)
                               .map_err(|e| ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), total)));
            if buf.len() < 1 {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "empty msg", total));
            }
            let noncelen = buf[0] as usize;
            if noncelen != NONCEBYTES {
                return Err(ShoopErr::new(ShoopErrKind::Fatal, "nonce length not recognized", total));
            }
            if buf.len() < (1 + noncelen) {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "msg not long enough to contain nonce", total));
            }
            let mut noncebytes = [0u8; NONCEBYTES];
            noncebytes.copy_from_slice(&buf[1..1+noncelen]);
            let nonce = Nonce(noncebytes);

            let unsealed = try!(secretbox::open(&buf[1+noncelen..], &nonce, &key).map_err(|_| ShoopErr::new(ShoopErrKind::Fatal, "failed to decrypt", total)));

            total += unsealed.len() as u64;
            f.write_all(&unsealed[..]).unwrap();
            if time::precise_time_ns() > ts + 100_000_000 {
                overprint!("   {:.1}M / {:.1}M ({:.1}%)", (total as f64)/(1024f64*1024f64), (filesize as f64)/(1024f64*1024f64), (total as f64) / (filesize as f64) * 100f64);
                ts = time::precise_time_ns();
            }
            if total >= filesize {
                println!("\nEOF");
                break;
            }
        }
        let _ = sock.close();
        Ok(())
    }
}
<commit_msg>fix some resumption code<commit_after>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;
extern crate byteorder;
extern crate time;
extern crate sodiumoxide;
extern crate rustc_serialize;

use daemonize::{Daemonize};
use std::process::Command;
use std::net::{UdpSocket, SocketAddr, SocketAddrV4};
use std::str;
use std::env;
use std::fmt;
use std::str::FromStr;
use std::fs::{OpenOptions, File};
use std::path::Path;
use std::ffi::OsStr;
use std::io::{Cursor, Error, Seek, SeekFrom, stderr, Read, Write};
use udt::{UdtSocket, UdtOpts, SocketType, SocketFamily};
use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian};
use sodiumoxide::crypto::secretbox;
use sodiumoxide::crypto::secretbox::xsalsa20poly1305::{NONCEBYTES, Key, Nonce};
use rustc_serialize::hex::{FromHex, ToHex};

macro_rules! overprint {
    ($fmt: expr) => {
        print!(concat!("\x1b[2K\r", $fmt));
        std::io::stdout().flush().unwrap();
    };
    ($fmt:expr, $($arg:tt)*) => {
        print!(concat!("\x1b[2K\r", $fmt) , $($arg)*);
        std::io::stdout().flush().unwrap();
    };
}

pub struct Server<'a> {
    filename: &'a str,
    sock: UdtSocket,
    key: Key,
}

pub struct Client<'a> {
    remote_host: &'a str,
    port_range: PortRange,
    remote_path: &'a str,
}

pub struct PortRange {
    start: u16,
    end: u16
}

enum ShoopErrKind {
    Severed,
    Fatal,
}
struct ShoopErr {
    kind: ShoopErrKind,
    msg: Option<String>,
    finished: u64,
}

impl ShoopErr {
    pub fn new(kind: ShoopErrKind, msg: &str, finished: u64) -> ShoopErr {
        ShoopErr { kind: kind, msg: Some(String::from(msg)), finished: finished }
    }

    pub fn from(err: Error, finished: u64) -> ShoopErr {
        ShoopErr { kind: ShoopErrKind::Severed, msg: Some(format!("{:?}", err)), finished: finished }
    }
}

impl PortRange {
    fn new(start: u16, end: u16) -> PortRange {
        PortRange{ start: start, end: end }
    }

    pub fn from(s: &str) -> Result<PortRange, &str> {
        let sections: Vec<&str> = s.split("-").collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"")
        }
        let (start, end) = (sections[0].parse::<u16>(),
                            sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("invalid port range");
        }
        Ok(PortRange::new(start.unwrap(), end.unwrap()))
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}

impl<'a> Server<'a> {
    pub fn new(port_range : PortRange,
               filename   : &str)
        -> Server
    {
        let sshconnstr = env::var("SSH_CONNECTION").unwrap_or(String::from("0.0.0.0 0 0.0.0.0 22")).trim().to_owned();
        let sshconn: Vec<&str> = sshconnstr.split(" ").collect();
        let ip = sshconn[2];
        let port = Server::get_open_port(port_range).unwrap();
        let key = secretbox::gen_key();
        let Key(keybytes) = key;
        println!("shoop 0 {} {} {}", ip, port, keybytes.to_hex());
        let daemonize = Daemonize::new();
        match daemonize.start() {
            Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            Err(_) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
        }

        udt::init();
        let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
        sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
        sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
        sock.bind(SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap())).unwrap();
        Server { sock: sock, key: key, filename: filename }
    }

    pub fn start(&self) {
        self.sock.listen(1).unwrap();

        loop {
            let (stream, _) = self.sock.accept().unwrap();
            if let Ok(starthdr) = stream.recvmsg(9) {
                let version = starthdr[0];
                let mut rdr = Cursor::new(starthdr);
                rdr.set_position(1);
                let offset = rdr.read_u64::<LittleEndian>().unwrap();
                if version == 0x00 {
                    match self.send_file(stream, offset) {
                        Ok(_) => {
                            let _ = stream.close();
                            break;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Severed, msg: _, finished: _}) => {
                            let _ = stream.close();
                            continue;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Fatal, msg, finished: _}) => {
                            panic!("{:?}", msg);
                        }
                    }
                } else {
                    panic!("Unrecognized version.");
                }
            } else {
                panic!("Failed to receive version byte from client.");
            }
        }
    }

    fn send_file(&self, stream: UdtSocket, offset: u64) -> Result<(), ShoopErr> {
        let mut f = File::open(self.filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();

        let mut remaining = 0u64;
        let mut buf = vec![0; 1024 * 1024];
        loop {
            match try!(f.read(&mut buf).map_err(|e| ShoopErr::from(e, 0u64))) {
                0    => { break },
                read => { remaining += read as u64 },
            }
        }

        let mut wtr = vec![];
        wtr.write_u64::<LittleEndian>(remaining).unwrap();
        match stream.sendmsg(&wtr[..]) {
            Ok(0) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "failed to write filesize header before timeout", remaining))
            },
            Err(e) => {
                return Err(ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), remaining))
            }
            _ => {}
        }
        let mut payload = vec![0; 1300];
        f.seek(SeekFrom::Start(offset)).unwrap();
        loop {
            match f.read(&mut payload) {
                Ok(0) => {
                    break;
                }
                Ok(read) => {
                    let nonce = secretbox::gen_nonce();
                    let Nonce(noncebytes) = nonce;
                    let mut hdr = vec![0u8; 1 + NONCEBYTES];
                    hdr[0] = NONCEBYTES as u8;
                    hdr[1..].clone_from_slice(&noncebytes);

                    let mut sealed = secretbox::seal(&payload[0..read], &nonce, &self.key);
                    let mut msg = Vec::with_capacity(hdr.len() + sealed.len());
                    msg.extend_from_slice(&hdr);
                    msg.append(&mut sealed);
                    match stream.sendmsg(&msg[..]) {
                        Ok(_) =>  { }
                        Err(e) => {
                            stream.close().expect("Error closing stream");
                            panic!("{:?}", e);
                        }
                    }
                },
                Err(e) => {
                    stream.close().expect("Error closing stream");
                    panic!("{:?}", e);
                }
            }
        }

        stream.close().expect("Error closing stream.");
        Ok(())
    }

    fn get_open_port(range: PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }
}

impl<'a> Client<'a> {
    pub fn new(remote_ssh_host : &'a str,
               port_range      : PortRange,
               remote_path     : &'a str)
        -> Client<'a>
    {
        Client{ remote_host: remote_ssh_host, port_range: port_range, remote_path: remote_path }
    }

    pub fn start(&self) {
        let cmd = format!("shoop -s '{}' -p {}", self.remote_path, self.port_range);
        // println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

        overprint!(" - establishing SSH session...");
        assert!(Client::command_exists("ssh"), "`ssh` is required!");
        let output = Command::new("ssh")
                             .arg(self.remote_host.to_owned())
                             .arg(cmd)
                             .output()
                             .unwrap_or_else(|e| {
                                 panic!("failed to execute process: {}", e);
                             });
        let infostring = String::from_utf8_lossy(&output.stdout).to_owned().trim().to_owned();
        let info: Vec<&str> = infostring.split(" ").collect();
        if info.len() != 5 {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let (magic, version, ip, port, keyhex) = (info[0], info[1], info[2], info[3], info[4]);
        overprint!(" - opening UDT connection...");
        if magic != "shoop" || version != "0" {
            panic!("Unexpected response from server. Are you suuuuure shoop is setup on the server?");
        }

        let mut keybytes = [0u8; 32];
        keybytes.copy_from_slice(&keyhex.from_hex().unwrap()[..]);
        let key = Key(keybytes);

        udt::init();
        let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::from_str(&format!("{}:{}", ip, port)[..]).unwrap());

        let mut offset = 0u64;
        loop {
            let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
            sock.setsockopt(UdtOpts::UDP_RCVBUF, 1024000i32).unwrap();
            sock.setsockopt(UdtOpts::UDP_SNDBUF, 1024000i32).unwrap();
            match sock.connect(addr) {
               Ok(()) => {
                   overprint!(" - connection opened, shakin' hands, makin' frands");
               },
               Err(e) => {
                   panic!("errrrrrrr connecting to {}:{} - {:?}", ip, port, e);
               }
            }
            let mut wtr = vec![];
            wtr.push(0);
            wtr.write_u64::<LittleEndian>(offset).unwrap();
            match sock.sendmsg(&wtr[..]) {
                Err(_) => { sock.close().unwrap(); continue; }
                _      => {}
            }

            match sock.recvmsg(8) {
               Ok(msg) => {
                   if msg.len() == 0 {
                       panic!("failed to get filesize from server, probable timeout.");
                   }
                   let mut rdr = Cursor::new(msg);
                   let filesize = rdr.read_u64::<LittleEndian>().unwrap();
                   let filename = Path::new(&self.remote_path).file_name().unwrap_or(OsStr::new("outfile")).to_str().unwrap_or("outfile");
                   overprint!(" + downloading {} ({:.1}MB)\n", filename, (filesize as f64)/(1024f64*1024f64));
                   match self.recv_file(sock, filesize, filename, &key, offset) {
                        Ok(_) => {
                            break;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Severed, msg: _, finished}) => {
                            println!(" * [[SEVERED]]");
                            offset = finished;
                        }
                        Err(ShoopErr{ kind: ShoopErrKind::Fatal, msg, finished: _}) => {
                            panic!("{:?}", msg);
                        }
                   }
               }
               Err(_) => {}
            }
            let _ = sock.close();
        }
    }

    fn command_exists(command: &str) -> bool {
        match Command::new("which").arg(command).output() {
            Ok(output) => output.status.success(),
            Err(_)     => false
        }
    }

    fn recv_file(&self, sock: UdtSocket, filesize: u64, filename: &str, key: &Key, offset: u64) -> Result<(), ShoopErr> {
        let mut f = OpenOptions::new().write(true).create(true).truncate(false).open(filename).unwrap();
        f.seek(SeekFrom::Start(offset)).unwrap();
        let mut ts = time::precise_time_ns();
        let mut total = 0u64;
        loop {
            let buf = try!(sock.recvmsg(8192)
                               .map_err(|e| ShoopErr::new(ShoopErrKind::Severed, &format!("{:?}", e), total)));
            if buf.len() < 1 {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "empty msg", total));
            }
            let noncelen = buf[0] as usize;
            if noncelen != NONCEBYTES {
                return Err(ShoopErr::new(ShoopErrKind::Fatal, "nonce length not recognized", total));
            }
            if buf.len() < (1 + noncelen) {
                return Err(ShoopErr::new(ShoopErrKind::Severed, "msg not long enough to contain nonce", total));
            }
            let mut noncebytes = [0u8; NONCEBYTES];
            noncebytes.copy_from_slice(&buf[1..1+noncelen]);
            let nonce = Nonce(noncebytes);

            let unsealed = try!(secretbox::open(&buf[1+noncelen..], &nonce, &key).map_err(|_| ShoopErr::new(ShoopErrKind::Fatal, "failed to decrypt", total)));

            total += unsealed.len() as u64;
            f.write_all(&unsealed[..]).unwrap();
            if time::precise_time_ns() > ts + 100_000_000 {
                overprint!("   {:.1}M / {:.1}M ({:.1}%)", (total as f64)/(1024f64*1024f64), (filesize as f64)/(1024f64*1024f64), (total as f64) / (filesize as f64) * 100f64);
                ts = time::precise_time_ns();
            }
            if total >= filesize {
                println!("\nEOF");
                break;
            }
        }
        let _ = sock.close();
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! **rosc** is an implementation of the [OSC 1.0](http://opensoundcontrol.org/spec-1_0) protocol in pure Rust.
//!

extern crate byteorder;
extern crate regex;
extern crate nom;

/// Crate specific error types.
mod errors;
/// OSC data types, see [OSC 1.0 specification](http://opensoundcontrol.org/spec-1_0) for details.
mod types;

pub use crate::errors::*;
pub use crate::types::*;

/// Provides a decoding method for OSC packets.
pub mod decoder;
/// Encodes an `OscPacket` to a byte vector.
pub mod encoder;
/// Address checking and matching methods
pub mod address;<commit_msg>Update link to Open Sound Control 1.0 spec<commit_after>//! **rosc** is an implementation of the [OSC 1.0](http://opensoundcontrol.org/spec-1_0) protocol in pure Rust.
//!

extern crate byteorder;
extern crate regex;
extern crate nom;

/// Crate specific error types.
mod errors;
/// OSC data types, see [OSC 1.0 specification](https://opensoundcontrol.stanford.edu/spec-1_0.html) for details.
mod types;

pub use crate::errors::*;
pub use crate::types::*;

/// Provides a decoding method for OSC packets.
pub mod decoder;
/// Encodes an `OscPacket` to a byte vector.
pub mod encoder;
/// Address checking and matching methods
pub mod address;
<|endoftext|>


<commit_before>//! Find discrete returns in full waveform LiDAR data.
//!
//! # Why is this library called `peakbag`?
//!
//! [Peak bagging](https://en.wikipedia.org/wiki/Peak_bagging) is when you try to summit a bunch of
//! mountains, just to say you summited a bunch of mountains. While the practice of peak bagging
//! can correlate with actual appreciation for the out-of-doors and an adventuresome spirit, a peak
//! bagging attitude is neither necessary nor sufficient for a good time outside.
//!
//! Some use "peak bagger" as a derisive term for someone who likes taking a selfie on top of a
//! mountain more than just spending time outside.
//!
//! This library finds peaks in waveforms, so `peakbag` seemed as good of a name as any.

#![deny(box_pointers, fat_ptr_transmutes, missing_copy_implementations, missing_debug_implementations, missing_docs, trivial_casts, unsafe_code, unused_extern_crates, unused_import_braces, unused_qualifications, variant_size_differences)]

#[macro_use]
extern crate log;
extern crate num;

use std::f64;
use std::fmt;

use num::traits::{ToPrimitive, Unsigned};

/// Detects peaks in full waveform data.
///
/// This is a convenience method that wraps calls to `PeakDetector::new` and
/// `PeakDetector::detect_peaks`.
///
/// # Examples
///
/// ```
/// let ref data = [1u32, 2, 3, 4, 3, 2, 1];
/// let peaks = peakbag::detect_peaks(data, 3, 0, 5);
/// assert_eq!(1, peaks.len());
/// assert_eq!(4, peaks[0].amplitude);
/// assert_eq!(3, peaks[0].index);
/// ```
pub fn detect_peaks<T>(data: &[T], width: usize, floor: T, ceiling: T) -> Vec<Peak<T>>
    where T: Copy + fmt::Display + PartialOrd + ToPrimitive + Unsigned
{
    let detector = PeakDetector::new(width, floor, ceiling);
    detector.detect_peaks(data)
}

/// Configurable struct for detecting peaks.
///
/// This structure allow for fine-grained adjustment of the peak detection procedures.
#[derive(Clone, Copy, Debug)]
pub struct PeakDetector<T: Copy> {
    width: usize,
    floor: T,
    ceiling: T,
    saturation: Option<T>,
    max_kurtosis: f64,
    min_height_above_background: f64,
}

impl<T: Copy> PeakDetector<T> {
    /// Creates a new peak detector.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(1, 2, 3);
    /// ```
    pub fn new(width: usize, floor: T, ceiling: T) -> PeakDetector<T> {
        PeakDetector {
            width: width,
            floor: floor,
            ceiling: ceiling,
            saturation: None,
            max_kurtosis: f64::MAX,
            min_height_above_background: f64::MIN,
        }
    }

    /// Sets the saturation level for this peak detector.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let mut detector = PeakDetector::new(1, 2, 3).saturation(3);
    /// ```
    pub fn saturation(mut self, saturation: T) -> PeakDetector<T> {
        self.saturation = Some(saturation);
        self
    }

    /// Sets the minimum allowable height above background for a peak.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let mut detector = PeakDetector::new(1, 2, 3).min_height_above_background(4.0);
    /// ```
    pub fn min_height_above_background(mut self,
                                       min_height_above_background: f64)
                                       -> PeakDetector<T> {
        self.min_height_above_background = min_height_above_background;
        self
    }

    /// Sets the maximum allowable kurtosis for a peak.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(1, 2, 3).max_kurtosis(4.0);
    /// ```
    pub fn max_kurtosis(mut self, max_kurtosis: f64) -> PeakDetector<T> {
        self.max_kurtosis = max_kurtosis;
        self
    }
}

impl<T> PeakDetector<T> where T: Copy + fmt::Display + PartialOrd + ToPrimitive + Unsigned {
    /// Detects peaks in full waveform data.
    ///
    /// # Panics
    ///
    /// Panics if a sample value cannot be converted to an `i64` and `f64`.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(0, 8, 2);
    /// let peaks = detector.detect_peaks(&[1u32, 2, 3, 2, 1]);
    /// ```
    pub fn detect_peaks(&self, data: &[T]) -> Vec<Peak<T>> {
        let mut state = State::Ascending(0);
        let mut peaks = Vec::new();
        for (i, &sample) in data.iter().enumerate() {
            if i == 0 {
                // We assume the first sample has a slope of zero, which means it can start a leading
                // edge.
                state = State::Ascending(1);
                continue;
            }
            let slope = sample.to_i64().unwrap() - data[i - 1].to_i64().unwrap();
            if let Some(s) = self.saturation {
                if sample == s {
                    state = State::Saturated;
                }
            }
            match state {
                State::Ascending(n) => {
                    if slope < 0 {
                        state = State::Ascending(0);
                    } else if n == self.width {
                        state = State::Descending(0, i);
                    } else {
                        state = State::Ascending(n + 1);
                    }
                }
                State::Descending(n, index) => {
                    if slope > 0 {
                        if n == 0 {
                            state = State::Descending(0, i);
                        } else {
                            state = State::Ascending(1);
                        }
                    } else if n + 1 == self.width {
                        let amplitude = data[index];
                        if amplitude > self.floor && amplitude <= self.ceiling {
                            let (mean, rms, kurtosis) = self.peak_stats(&data, index);
                            let height_above_background = self.height_above_background(&data,
                                                                                index);
                            if height_above_background >= self.min_height_above_background &&
                               kurtosis < self.max_kurtosis {
                                peaks.push(Peak {
                                    index: index,
                                    amplitude: amplitude,
                                    mean: mean,
                                    rms: rms,
                                    kurtosis: kurtosis,
                                    height_above_background: height_above_background,
                                });
                            }
                        }
                        state = State::Ascending(0);
                    } else {
                        state = State::Descending(n + 1, index);
                    }
                }
                State::Saturated => {
                    // We know if we're below the floor, we have a negative slope, since we must
                    // have just gone below the floor.
                    if sample <= self.floor {
                        state = State::Ascending(0);
                    }
                }
            };
            debug!("({}) sample={}, slope={}, state={:?}",
                   i,
                   sample,
                   slope,
                   state);
        }
        peaks
    }

    fn peak_stats(&self, data: &[T], index: usize) -> (f64, f64, f64) {
        let mut values = 0u64;
        let mut values2 = 0u64;
        let mut nvalues = 0usize;
        for &sample in data.iter().skip(index - self.width).take(self.width * 2 + 1) {
            let sample = sample.to_u64().unwrap();
            values += sample;
            values2 += sample * sample;
            nvalues += 1;
        }
        let mean = values as f64 / nvalues as f64;
        let rms = (values2 as f64 / nvalues as f64 - (values as f64 / nvalues as f64).powi(2))
                      .sqrt();
        let mut kurtosis = 0f64;
        for &sample in data.iter().skip(index - self.width).take(self.width * 2 + 1) {
            let sample = sample.to_u64().unwrap();
            let temp = (sample as f64 - mean) / rms;
            kurtosis += temp.powi(4);
        }
        kurtosis = kurtosis / nvalues as f64 - 3.0;
        (mean, rms, kurtosis)
    }

    fn height_above_background(&self, data: &[T], index: usize) -> f64 {
        let slope: f64 = (data[index + self.width].to_f64().unwrap() - data[index - self.width].to_f64().unwrap()) /
                         (2.0 * self.width as f64);
        let intercept = data[index + self.width].to_f64().unwrap() -
                        slope * (index - self.width) as f64;
        data[index].to_f64().unwrap() - (slope * index as f64 + intercept)
    }
}

#[derive(Debug)]
enum State {
    Ascending(usize),
    Descending(usize, usize),
    Saturated,
}

/// A peak in the waveform data.
#[derive(Clone, Copy, Debug)]
pub struct Peak<T: Copy> {
    /// The raw intensity of the peak.
    pub amplitude: T,
    /// The index of the peak in the sample data.
    pub index: usize,
    /// The mean intensity value of the peak.
    pub mean: f64,
    /// The rms error of the peak from that mean.
    pub rms: f64,
    /// The kurtosis of the peak.
    pub kurtosis: f64,
    /// The height of the peak above a background level, as defined by the first and last points in
    /// the peak.
    pub height_above_background: f64,
}

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

    #[test]
    fn level_peak() {
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 4, 4, 4], 3, 0, 5);
        assert_eq!(1, peaks.len());
        assert_eq!(3, peaks[0].index);
    }

    #[test]
    fn floor() {
        // Must be above floor, below or equal to ceil
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 4, 5);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn ceiling() {
        // Must be above floor, below or equal to ceil
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 0, 3);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn saturation() {
        let detector = PeakDetector::new(3, 1, 8).saturation(8);
        let peaks = detector.detect_peaks(&[5u32, 6, 7, 8, 7, 6, 5]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn stats() {
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 0, 5);
        let ref peak = peaks[0];
        assert_eq!(2.2857142857142856, peak.mean);
        assert_eq!(1.0301575072754257, peak.rms);
        assert_eq!(-1.143491124260356, peak.kurtosis);
        assert_eq!(3.0, peak.height_above_background);
    }

    #[test]
    fn min_height_above_background() {
        let detector = PeakDetector::new(3, 1, 8).min_height_above_background(4.0);
        let peaks = detector.detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn peak_kurtosis() {
        let detector = PeakDetector::new(3, 1, 8).max_kurtosis(-2.0);
        let peaks = detector.detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn should_be_one_peak() {
        let ref data = [2u16, 2, 1, 2, 5, 18, 51, 107, 166, 195, 176, 125, 70, 34, 14, 7, 5, 4, 5,
                        4, 3, 1, 0, 0];
        let peaks = detect_peaks(data, 2, 15, 255);
        assert_eq!(1, peaks.len());
        assert_eq!(195, peaks[0].amplitude);
    }

    #[test]
    fn fix_arithmatic_overflow() {
        let ref data = [3u16, 3, 2, 1, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 4, 15, 50, 119, 194, 243, 255, 255, 217, 158, 97, 54, 30, 23, 20, 20, 18, 19, 20, 19, 17, 15, 13, 11, 12, 11, 10, 10, 10, 9, 10, 10, 10, 9, 11, 14, 27, 47, 66, 69, 54, 34, 19, 12, 9, 8, 7, 8, 12, 18, 28, 35, 34, 26, 16, 11, 15, 33, 64, 94, 105, 89, 58, 30, 14, 9, 7, 7, 7, 9, 9, 9, 7, 6, 5, 5, 6, 6, 6, 6, 5, 5, 4, 3, 4, 4];
        let _ = detect_peaks(data, 3, 15, 250);
    }
}
<commit_msg>Fix arithmatic error in background height<commit_after>//! Find discrete returns in full waveform LiDAR data.
//!
//! # Why is this library called `peakbag`?
//!
//! [Peak bagging](https://en.wikipedia.org/wiki/Peak_bagging) is when you try to summit a bunch of
//! mountains, just to say you summited a bunch of mountains. While the practice of peak bagging
//! can correlate with actual appreciation for the out-of-doors and an adventuresome spirit, a peak
//! bagging attitude is neither necessary nor sufficient for a good time outside.
//!
//! Some use "peak bagger" as a derisive term for someone who likes taking a selfie on top of a
//! mountain more than just spending time outside.
//!
//! This library finds peaks in waveforms, so `peakbag` seemed as good of a name as any.

#![deny(box_pointers, fat_ptr_transmutes, missing_copy_implementations, missing_debug_implementations, missing_docs, trivial_casts, unsafe_code, unused_extern_crates, unused_import_braces, unused_qualifications, variant_size_differences)]

#[macro_use]
extern crate log;
extern crate num;

use std::f64;
use std::fmt;

use num::traits::{ToPrimitive, Unsigned};

/// Detects peaks in full waveform data.
///
/// This is a convenience method that wraps calls to `PeakDetector::new` and
/// `PeakDetector::detect_peaks`.
///
/// # Examples
///
/// ```
/// let ref data = [1u32, 2, 3, 4, 3, 2, 1];
/// let peaks = peakbag::detect_peaks(data, 3, 0, 5);
/// assert_eq!(1, peaks.len());
/// assert_eq!(4, peaks[0].amplitude);
/// assert_eq!(3, peaks[0].index);
/// ```
pub fn detect_peaks<T>(data: &[T], width: usize, floor: T, ceiling: T) -> Vec<Peak<T>>
    where T: Copy + fmt::Display + PartialOrd + ToPrimitive + Unsigned
{
    let detector = PeakDetector::new(width, floor, ceiling);
    detector.detect_peaks(data)
}

/// Configurable struct for detecting peaks.
///
/// This structure allow for fine-grained adjustment of the peak detection procedures.
#[derive(Clone, Copy, Debug)]
pub struct PeakDetector<T: Copy> {
    width: usize,
    floor: T,
    ceiling: T,
    saturation: Option<T>,
    max_kurtosis: f64,
    min_height_above_background: f64,
}

impl<T: Copy> PeakDetector<T> {
    /// Creates a new peak detector.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(1, 2, 3);
    /// ```
    pub fn new(width: usize, floor: T, ceiling: T) -> PeakDetector<T> {
        PeakDetector {
            width: width,
            floor: floor,
            ceiling: ceiling,
            saturation: None,
            max_kurtosis: f64::MAX,
            min_height_above_background: f64::MIN,
        }
    }

    /// Sets the saturation level for this peak detector.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let mut detector = PeakDetector::new(1, 2, 3).saturation(3);
    /// ```
    pub fn saturation(mut self, saturation: T) -> PeakDetector<T> {
        self.saturation = Some(saturation);
        self
    }

    /// Sets the minimum allowable height above background for a peak.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let mut detector = PeakDetector::new(1, 2, 3).min_height_above_background(4.0);
    /// ```
    pub fn min_height_above_background(mut self,
                                       min_height_above_background: f64)
                                       -> PeakDetector<T> {
        self.min_height_above_background = min_height_above_background;
        self
    }

    /// Sets the maximum allowable kurtosis for a peak.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(1, 2, 3).max_kurtosis(4.0);
    /// ```
    pub fn max_kurtosis(mut self, max_kurtosis: f64) -> PeakDetector<T> {
        self.max_kurtosis = max_kurtosis;
        self
    }
}

impl<T> PeakDetector<T> where T: Copy + fmt::Display + PartialOrd + ToPrimitive + Unsigned {
    /// Detects peaks in full waveform data.
    ///
    /// # Panics
    ///
    /// Panics if a sample value cannot be converted to an `i64` and `f64`.
    ///
    /// # Examples
    ///
    /// ```
    /// use peakbag::PeakDetector;
    /// let detector = PeakDetector::new(0, 8, 2);
    /// let peaks = detector.detect_peaks(&[1u32, 2, 3, 2, 1]);
    /// ```
    pub fn detect_peaks(&self, data: &[T]) -> Vec<Peak<T>> {
        let mut state = State::Ascending(0);
        let mut peaks = Vec::new();
        for (i, &sample) in data.iter().enumerate() {
            if i == 0 {
                // We assume the first sample has a slope of zero, which means it can start a leading
                // edge.
                state = State::Ascending(1);
                continue;
            }
            let slope = sample.to_i64().unwrap() - data[i - 1].to_i64().unwrap();
            if let Some(s) = self.saturation {
                if sample == s {
                    state = State::Saturated;
                }
            }
            match state {
                State::Ascending(n) => {
                    if slope < 0 {
                        state = State::Ascending(0);
                    } else if n == self.width {
                        state = State::Descending(0, i);
                    } else {
                        state = State::Ascending(n + 1);
                    }
                }
                State::Descending(n, index) => {
                    if slope > 0 {
                        if n == 0 {
                            state = State::Descending(0, i);
                        } else {
                            state = State::Ascending(1);
                        }
                    } else if n + 1 == self.width {
                        let amplitude = data[index];
                        if amplitude > self.floor && amplitude <= self.ceiling {
                            let (mean, rms, kurtosis) = self.peak_stats(&data, index);
                            let height_above_background = self.height_above_background(&data,
                                                                                index);
                            if height_above_background >= self.min_height_above_background &&
                               kurtosis < self.max_kurtosis {
                                peaks.push(Peak {
                                    index: index,
                                    amplitude: amplitude,
                                    mean: mean,
                                    rms: rms,
                                    kurtosis: kurtosis,
                                    height_above_background: height_above_background,
                                });
                            }
                        }
                        state = State::Ascending(0);
                    } else {
                        state = State::Descending(n + 1, index);
                    }
                }
                State::Saturated => {
                    // We know if we're below the floor, we have a negative slope, since we must
                    // have just gone below the floor.
                    if sample <= self.floor {
                        state = State::Ascending(0);
                    }
                }
            };
            debug!("({}) sample={}, slope={}, state={:?}",
                   i,
                   sample,
                   slope,
                   state);
        }
        peaks
    }

    fn peak_stats(&self, data: &[T], index: usize) -> (f64, f64, f64) {
        let mut values = 0u64;
        let mut values2 = 0u64;
        let mut nvalues = 0usize;
        for &sample in data.iter().skip(index - self.width).take(self.width * 2 + 1) {
            let sample = sample.to_u64().unwrap();
            values += sample;
            values2 += sample * sample;
            nvalues += 1;
        }
        let mean = values as f64 / nvalues as f64;
        let rms = (values2 as f64 / nvalues as f64 - (values as f64 / nvalues as f64).powi(2))
                      .sqrt();
        let mut kurtosis = 0f64;
        for &sample in data.iter().skip(index - self.width).take(self.width * 2 + 1) {
            let sample = sample.to_u64().unwrap();
            let temp = (sample as f64 - mean) / rms;
            kurtosis += temp.powi(4);
        }
        kurtosis = kurtosis / nvalues as f64 - 3.0;
        (mean, rms, kurtosis)
    }

    fn height_above_background(&self, data: &[T], index: usize) -> f64 {
        let slope: f64 = (data[index + self.width].to_f64().unwrap() - data[index - self.width].to_f64().unwrap()) /
                         (2.0 * self.width as f64);
        let intercept = data[index - self.width].to_f64().unwrap() -
                        slope * (index - self.width) as f64;
        data[index].to_f64().unwrap() - (slope * index as f64 + intercept)
    }
}

#[derive(Debug)]
enum State {
    Ascending(usize),
    Descending(usize, usize),
    Saturated,
}

/// A peak in the waveform data.
#[derive(Clone, Copy, Debug)]
pub struct Peak<T: Copy> {
    /// The raw intensity of the peak.
    pub amplitude: T,
    /// The index of the peak in the sample data.
    pub index: usize,
    /// The mean intensity value of the peak.
    pub mean: f64,
    /// The rms error of the peak from that mean.
    pub rms: f64,
    /// The kurtosis of the peak.
    pub kurtosis: f64,
    /// The height of the peak above a background level, as defined by the first and last points in
    /// the peak.
    pub height_above_background: f64,
}

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

    #[test]
    fn level_peak() {
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 4, 4, 4], 3, 0, 5);
        assert_eq!(1, peaks.len());
        assert_eq!(3, peaks[0].index);
    }

    #[test]
    fn floor() {
        // Must be above floor, below or equal to ceil
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 4, 5);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn ceiling() {
        // Must be above floor, below or equal to ceil
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 0, 3);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn saturation() {
        let detector = PeakDetector::new(3, 1, 8).saturation(8);
        let peaks = detector.detect_peaks(&[5u32, 6, 7, 8, 7, 6, 5]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn stats() {
        let peaks = detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1], 3, 0, 5);
        let ref peak = peaks[0];
        assert_eq!(2.2857142857142856, peak.mean);
        assert_eq!(1.0301575072754257, peak.rms);
        assert_eq!(-1.143491124260356, peak.kurtosis);
        assert_eq!(3.0, peak.height_above_background);
    }

    #[test]
    fn min_height_above_background() {
        let detector = PeakDetector::new(3, 1, 8).min_height_above_background(4.0);
        let peaks = detector.detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn peak_kurtosis() {
        let detector = PeakDetector::new(3, 1, 8).max_kurtosis(-2.0);
        let peaks = detector.detect_peaks(&[1u32, 2, 3, 4, 3, 2, 1]);
        assert_eq!(0, peaks.len());
    }

    #[test]
    fn should_be_one_peak() {
        let ref data = [2u16, 2, 1, 2, 5, 18, 51, 107, 166, 195, 176, 125, 70, 34, 14, 7, 5, 4, 5,
                        4, 3, 1, 0, 0];
        let peaks = detect_peaks(data, 2, 15, 255);
        assert_eq!(1, peaks.len());
        assert_eq!(195, peaks[0].amplitude);
    }

    #[test]
    fn fix_arithmatic_overflow() {
        let ref data = [3u16, 3, 2, 1, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 4, 15, 50, 119, 194, 243, 255, 255, 217, 158, 97, 54, 30, 23, 20, 20, 18, 19, 20, 19, 17, 15, 13, 11, 12, 11, 10, 10, 10, 9, 10, 10, 10, 9, 11, 14, 27, 47, 66, 69, 54, 34, 19, 12, 9, 8, 7, 8, 12, 18, 28, 35, 34, 26, 16, 11, 15, 33, 64, 94, 105, 89, 58, 30, 14, 9, 7, 7, 7, 9, 9, 9, 7, 6, 5, 5, 6, 6, 6, 6, 5, 5, 4, 3, 4, 4];
        let _ = detect_peaks(data, 3, 15, 250);
    }

    #[test]
    fn reference_pulse_not_one_peak() {
        let ref data = [2u16, 1, 2, 2, 2, 1, 4, 10, 32, 80, 140, 188, 188, 149, 93, 47, 21, 9, 4,
                        4, 5, 5, 3, 1];
        let detector = PeakDetector::new(2, 15, 255).min_height_above_background(5.0);
        let peaks = detector.detect_peaks(data);
        assert_eq!(1, peaks.len());
    }
}
<|endoftext|>


<commit_before>//! A parser for OpenType fonts.
//!
//! ## Example
//!
//! ```
//! extern crate opentype;
//! extern crate truetype;
//!
//! use opentype::File;
//! use truetype::NamingTable;
//!
//! # fn main() {
//! let path = "SourceSerifPro-Regular.otf";
//! # let path = "tests/fixtures/SourceSerifPro-Regular.otf";
//! let file = File::open(path).unwrap();
//! let font = &file.fonts[0];
//!
//! assert_eq!(font.font_header.as_ref().unwrap().units_per_em, 1000);
//! assert_eq!(font.horizontal_header.as_ref().unwrap().ascender, 918);
//! let strings = match font.naming_table {
//!     Some(NamingTable::Format0(ref table)) => table.strings().unwrap(),
//!     _ => unreachable!(),
//! };
//! assert_eq!(&strings[1], "Source Serif Pro");
//! assert_eq!(&strings[9], "Frank Grießhammer");
//! # }
//! ```

extern crate postscript;
extern crate truetype;

/// An error.
pub type Error = std::io::Error;

/// A result.
pub type Result<T> = std::result::Result<T, Error>;

macro_rules! raise(
    ($message:expr) => (return Err(::Error::new(::std::io::ErrorKind::Other, $message)));
);

mod file;
mod font;

pub use file::File;
pub use font::Font;
<commit_msg>Simplify the declaration of Result<commit_after>//! A parser for OpenType fonts.
//!
//! ## Example
//!
//! ```
//! extern crate opentype;
//! extern crate truetype;
//!
//! use opentype::File;
//! use truetype::NamingTable;
//!
//! # fn main() {
//! let path = "SourceSerifPro-Regular.otf";
//! # let path = "tests/fixtures/SourceSerifPro-Regular.otf";
//! let file = File::open(path).unwrap();
//! let font = &file.fonts[0];
//!
//! assert_eq!(font.font_header.as_ref().unwrap().units_per_em, 1000);
//! assert_eq!(font.horizontal_header.as_ref().unwrap().ascender, 918);
//! let strings = match font.naming_table {
//!     Some(NamingTable::Format0(ref table)) => table.strings().unwrap(),
//!     _ => unreachable!(),
//! };
//! assert_eq!(&strings[1], "Source Serif Pro");
//! assert_eq!(&strings[9], "Frank Grießhammer");
//! # }
//! ```

extern crate postscript;
extern crate truetype;

/// An error.
pub type Error = std::io::Error;

/// A result.
pub type Result<T> = std::io::Result<T>;

macro_rules! raise(
    ($message:expr) => (return Err(::Error::new(::std::io::ErrorKind::Other, $message)));
);

mod file;
mod font;

pub use file::File;
pub use font::Font;
<|endoftext|>


<commit_before>pub mod dna_base;
pub mod base4;

/// Represents the bases used in the genome string.
/// For example the bases of the DNA are adenine (A),
/// thymine (T), guanine (G) and cytosine (C).
pub trait Base: Sized + PartialEq + Eq + Copy + Clone {
    /// Returns the "successor" base, wrapping around. Used
    /// to produce the gene product.
    fn succ(self) -> Self;

    fn from_char(c: char) -> Option<Self>;
}

fn locate_promoter<'a, T: Eq>(s: &'a [T], promoter: &[T]) -> Option<usize> {
    let plen = promoter.len();
    assert!(plen > 0);

    if s.len() < plen {
        return None;
    }

    for (i, window) in s.windows(plen).enumerate() {
        if window == promoter {
            return Some(i);
        }
    }

    return None;
}

#[derive(Debug)]
pub struct Gene<'a, B: Base + 'a> {
    pub regulatory_region: &'a [B],
    pub gene: &'a [B],
}

impl<'a, B: Base + 'a> Gene<'a, B> {
    pub fn gene_product(&self) -> Vec<B> {
        self.gene.iter().map(|b| b.succ()).collect()
    }

    pub fn find_product_in_regulatory_region(&self, product: &[B]) -> bool {
        assert!(product.len() > 0);
        for window in self.regulatory_region.windows(product.len()) {
            if window == product {
                return true;
            }
        }
        return false;
    }
}

pub struct GeneIterator<'a, 'b, B: Base + 'a + 'b> {
    start_pos: usize,
    // Genes have fixed length
    length_of_gene: usize,
    sequence: &'a [B],
    promoter: &'b [B],
}

impl<'a, 'b, B: Base + 'a + 'b> Iterator for GeneIterator<'a, 'b, B> {
    type Item = Gene<'a, B>;

    fn next(&mut self) -> Option<Self::Item> {
        match locate_promoter(&self.sequence[self.start_pos..], self.promoter) {
            Some(pos) => {
                let reg_start = self.start_pos;
                let reg_end = reg_start + pos;
                let gene_start = reg_end + self.promoter.len();
                let gene_end = gene_start + self.length_of_gene;

                // gene is not complete
                if gene_end > self.sequence.len() {
                    return None;
                }

                let gene = Gene {
                    regulatory_region: &self.sequence[reg_start..reg_end],
                    gene: &self.sequence[gene_start..gene_end],
                };
                self.start_pos = gene_end;
                return Some(gene);
            }
            None => {
                return None;
            }
        }
    }
}

/// A Genome is a string of Base
pub struct Genome<B: Base> {
    pub genome: Vec<B>,
}

// Convert genome into sections, i.e. Split at the promoter.
impl<B: Base> Genome<B> {
    pub fn iter_genes<'a, 'b>(&'a self,
                              promoter: &'b [B],
                              length_of_gene: usize)
                              -> GeneIterator<'a, 'b, B> {
        GeneIterator {
            start_pos: 0,
            length_of_gene: length_of_gene,
            sequence: &self.genome,
            promoter: promoter,
        }
    }
}
<commit_msg>Simplify the code<commit_after>pub mod dna_base;
pub mod base4;

/// Represents the bases used in the genome string.
/// For example the bases of the DNA are adenine (A),
/// thymine (T), guanine (G) and cytosine (C).
pub trait Base: Sized + PartialEq + Eq + Copy + Clone {
    /// Returns the "successor" base, wrapping around. Used
    /// to produce the gene product.
    fn succ(self) -> Self;

    fn from_char(c: char) -> Option<Self>;
}

// Locate ```substr``` in ```s```.
fn locate_substr<T: Eq>(s: &[T], substr: &[T]) -> Option<usize> {
    assert!(substr.len() > 0);

    if s.len() < substr.len() {
        return None;
    }

    for (i, window) in s.windows(substr.len()).enumerate() {
        if window == substr {
            return Some(i);
        }
    }

    return None;
}

#[derive(Debug)]
pub struct Gene<'a, B: Base + 'a> {
    pub regulatory_region: &'a [B],
    pub gene: &'a [B],
}

impl<'a, B: Base + 'a> Gene<'a, B> {
    pub fn gene_product(&self) -> Vec<B> {
        self.gene.iter().map(|b| b.succ()).collect()
    }

    pub fn find_product_in_regulatory_region(&self, product: &[B]) -> bool {
        locate_substr(self.regulatory_region, product).is_some()
    }
}

pub struct GeneIterator<'a, 'b, B: Base + 'a + 'b> {
    start_pos: usize,
    // Genes have fixed length
    length_of_gene: usize,
    sequence: &'a [B],
    promoter: &'b [B],
}

impl<'a, 'b, B: Base + 'a + 'b> Iterator for GeneIterator<'a, 'b, B> {
    type Item = Gene<'a, B>;

    fn next(&mut self) -> Option<Self::Item> {
        match locate_substr(&self.sequence[self.start_pos..], self.promoter) {
            Some(pos) => {
                let reg_start = self.start_pos;
                let reg_end = reg_start + pos;
                let gene_start = reg_end + self.promoter.len();
                let gene_end = gene_start + self.length_of_gene;

                // gene is not complete
                if gene_end > self.sequence.len() {
                    return None;
                }

                let gene = Gene {
                    regulatory_region: &self.sequence[reg_start..reg_end],
                    gene: &self.sequence[gene_start..gene_end],
                };
                self.start_pos = gene_end;
                return Some(gene);
            }
            None => {
                return None;
            }
        }
    }
}

/// A Genome is a string of Base
pub struct Genome<B: Base> {
    pub genome: Vec<B>,
}

// Convert genome into sections, i.e. Split at the promoter.
impl<B: Base> Genome<B> {
    pub fn iter_genes<'a, 'b>(&'a self,
                              promoter: &'b [B],
                              length_of_gene: usize)
                              -> GeneIterator<'a, 'b, B> {
        GeneIterator {
            start_pos: 0,
            length_of_gene: length_of_gene,
            sequence: &self.genome,
            promoter: promoter,
        }
    }
}
<|endoftext|>