<commit_before>
extern crate rustc_serialize;
extern crate amqp;
#[macro_use]
extern crate log;
extern crate argparse;
extern crate rand;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options, Consumer};
use amqp::protocol;
use std::default::Default;
use rustc_serialize::json;
use std::slice;

// for debugging
fn listen_discovery(channel: &mut Channel) {
    fn participant_discovered(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>){
        let s = std::str::from_utf8(&body).unwrap();
        let info: ParticipantInfo = json::decode(s).unwrap();
        info!("MsgFlo participant discovered: {:?}", info);
        channel.basic_ack(deliver.delivery_tag, false);
    }

    let queue_name = "fbp";
    let consumer_name = channel.basic_consume(participant_discovered as ConsumerCallBackFn, queue_name,
                                                "", false, false, false, false, Table::new());
    info!("listening for discovery messages");
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
pub struct ParticipantPort {
    pub id: String, // port name
    pub queue: String, // the associated message queue
    // FIXME: support. Is a keyword so needs some special handling   type: String, // datatype, ex: "boolean"
    // options: queue options as specified by the message queue implementation
    // description?? TODO: standardize
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
pub struct ParticipantInfo {
    pub id: String, // unique name
    pub role: String ,// role participant has
    pub component: String, // component the participant is instance of
    pub label: Option<String>, // (optional) short human-readable description
    pub icon: Option<String>, // (optional)
    pub inports: Vec<ParticipantPort>,
    pub outports: Vec<ParticipantPort>,
}

type SendFunction = fn(String, Vec<u8>);
pub type ProcessFunction = fn(Vec<u8>) -> Result<Vec<u8>, Vec<u8>>;
pub struct Participant {
    pub info: ParticipantInfo,
    pub process: ProcessFunction,
}

struct Connection {
    session: Session,
    channel: Channel,
}

// FIXME: wrap in protocol/command/payload envelope
fn create_queue_and_send(channel: &mut Channel, queue_name: &str, payload: String) {

    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    queue_declare.expect("queue creation failed");
    let content_type = Some("application/json".to_string());
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
    debug!("sending on {}: {}", queue_name, payload);
    let res = channel.basic_publish("", queue_name, true, false, props, payload.into_bytes());
    res.expect("send on new queue failed");
}

fn send_discovery(channel: &mut Channel, info: &ParticipantInfo) {
    let queue_name = "fbp"; // TODO: use an exchange istead, requires protocol change in msgflo

    let payload = json::encode(&info).unwrap();
    create_queue_and_send(channel, queue_name, payload);
    info!("sent participant discovery: {}  {}({})", info.id, info.role, info.component);
}

struct PortConsumer {
    process: ProcessFunction,
    portname: String,
    outqueue: String, // FIXME: allow sending on any port, also multiple times
}

fn send_out(channel: &mut Channel, exchange: String, data: Vec<u8>) {

    let routing_key = "".to_string();
    let content_type = Some("application/json".to_string()); // TODO: should be parameter
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
    let s = channel.basic_publish(exchange, routing_key, true, false, props, data);
    s.expect("failed to send");
    debug!("sent output");
}

impl Consumer for PortConsumer {
    fn handle_delivery(&mut self,
                       channel: &mut Channel,
                       deliver: protocol::basic::Deliver,
                       headers: protocol::basic::BasicProperties,
                       body: Vec<u8>) {

        debug!("calling process()");
        let f = self.process;
        let res = f(body);
        debug!("process() returned");

        if res.is_ok() {
            debug!("ACKing and sending");
            send_out(channel, self.outqueue.to_string(), res.unwrap());
            let r = channel.basic_ack(deliver.delivery_tag, false);
        } else {
            error!("process() errored");
            let r = channel.basic_nack(deliver.delivery_tag, true, true);
        }
        // FIXME: send error data
    }
}

// FIXME: actually call ProcessFunction
fn setup_inport(participant: &Participant, port: &ParticipantPort, connection: &mut Connection) {
    debug!("setup inport: {}", port.queue.to_string());

    let consumer = PortConsumer {
        process: participant.process,
        portname: port.id.to_string(),
        outqueue: participant.info.outports[0].queue.to_string(),
    };

    // create
    let declare = connection.channel.queue_declare(port.queue.to_string(), false, true, false, false, false, Table::new());
    declare.expect("inport queue creation failed");

    // subscribe
    let q = port.queue.to_string();
    let cons = connection.channel.basic_consume(consumer, q,
                                                "".to_string(), false, false, false, false, Table::new());

    cons.expect("iport setup failed");

    debug!("inport setup done: {:?}, {:?}", port.id.to_string(), port.queue.to_string());
}

fn setup_outport(participant: &Participant, port: &ParticipantPort, connection: &mut Connection) {

    let exchange_type = "fanout".to_string();
    let declare = connection.channel.exchange_declare(port.queue.to_string(), exchange_type,
                                                    false, true, false, false, false, Table::new());
    declare.expect("outport setup failed");

    debug!("setup outport done: {:?}, {:?}", port.id.to_string(), port.queue.to_string());
}

fn start_participant(participant: &Participant, options: &ParticipantOptions) -> Connection {

    let mut session = Session::open_url(&options.broker).expect("Can't create AMQP session");
    let mut channel = session.open_channel(1).expect("could not open AMQP channel");

    let mut conn = Connection { session: session, channel: channel };

    // setup ports
    setup_inport(&participant, &participant.info.inports[0], &mut conn);
    setup_outport(&participant, &participant.info.outports[0], &mut conn);

    // send MsgFlo participant discovery message
    send_discovery(&mut conn.channel, &participant.info);

    return conn;
}

fn stop_participant(participant: &Participant, connection: &mut Connection) {

    let closed = connection.channel.close(200, "Bye".to_string());
    connection.session.close(200, "Good Bye".to_string());
}

use log::{LogRecord, LogLevel, LogMetadata, LogLevelFilter};

struct SimpleLogger {
    showlevel: LogLevel,
}

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.showlevel
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}

#[derive(Debug)]
struct ParticipantOptions {
    role: String,
    broker: String,
}

impl Default for ParticipantOptions {
    fn default() -> ParticipantOptions {
        use rand::{thread_rng, Rng};

        let id: String = thread_rng().gen_ascii_chars().take(5).collect();
        ParticipantOptions {
            broker: "amqp://localhost//".to_string(),
            role: "".to_string(), // TODO: allow anonymous, format!("msgflo-rust-{}", id),
        }
    }
}

fn parse(options: &mut ParticipantOptions) {

    use argparse::{StoreTrue, Store};
    let mut parser = argparse::ArgumentParser::new();

    parser.refer(&mut options.role)
        .add_option(&["--role"], Store, "Participant role name");
    parser.refer(&mut options.broker)
        .add_option(&["--broker"], Store, "Address of messaging broker")
        .envvar("MSGFLO_BROKER");

    parser.parse_args_or_exit(); // XXX: should return out
}

// XXX: seems rust-amqp makes program hangs forever if error occurs / channel is borked?
// TODO: pass port info in/out of process()
// TODO: nicer way to declare ports? ideally they are enums not stringly typed?
pub fn participant_main(p: Participant) {

    let mut options = ParticipantOptions { .. Default::default() };
    parse(&mut options);

    if options.role == "" {
        options.role = p.info.role.to_string();
    }

    let mut c = start_participant(&p, &options);
    println!("{}({}) started", &options.role, &p.info.component);
    c.channel.start_consuming();

    stop_participant(&p, &mut c);
}

<commit_msg>Remove unused logger code<commit_after>
extern crate rustc_serialize;
extern crate amqp;
#[macro_use]
extern crate log;
extern crate argparse;
extern crate rand;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options, Consumer};
use amqp::protocol;
use std::default::Default;
use rustc_serialize::json;
use std::slice;

// for debugging
fn listen_discovery(channel: &mut Channel) {
    fn participant_discovered(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>){
        let s = std::str::from_utf8(&body).unwrap();
        let info: ParticipantInfo = json::decode(s).unwrap();
        info!("MsgFlo participant discovered: {:?}", info);
        channel.basic_ack(deliver.delivery_tag, false);
    }

    let queue_name = "fbp";
    let consumer_name = channel.basic_consume(participant_discovered as ConsumerCallBackFn, queue_name,
                                                "", false, false, false, false, Table::new());
    info!("listening for discovery messages");
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
pub struct ParticipantPort {
    pub id: String, // port name
    pub queue: String, // the associated message queue
    // FIXME: support. Is a keyword so needs some special handling   type: String, // datatype, ex: "boolean"
    // options: queue options as specified by the message queue implementation
    // description?? TODO: standardize
}

#[derive(Debug, Default, RustcDecodable, RustcEncodable)]
pub struct ParticipantInfo {
    pub id: String, // unique name
    pub role: String ,// role participant has
    pub component: String, // component the participant is instance of
    pub label: Option<String>, // (optional) short human-readable description
    pub icon: Option<String>, // (optional)
    pub inports: Vec<ParticipantPort>,
    pub outports: Vec<ParticipantPort>,
}

type SendFunction = fn(String, Vec<u8>);
pub type ProcessFunction = fn(Vec<u8>) -> Result<Vec<u8>, Vec<u8>>;
pub struct Participant {
    pub info: ParticipantInfo,
    pub process: ProcessFunction,
}

struct Connection {
    session: Session,
    channel: Channel,
}

// FIXME: wrap in protocol/command/payload envelope
fn create_queue_and_send(channel: &mut Channel, queue_name: &str, payload: String) {

    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    queue_declare.expect("queue creation failed");
    let content_type = Some("application/json".to_string());
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
    debug!("sending on {}: {}", queue_name, payload);
    let res = channel.basic_publish("", queue_name, true, false, props, payload.into_bytes());
    res.expect("send on new queue failed");
}

fn send_discovery(channel: &mut Channel, info: &ParticipantInfo) {
    let queue_name = "fbp"; // TODO: use an exchange istead, requires protocol change in msgflo

    let payload = json::encode(&info).unwrap();
    create_queue_and_send(channel, queue_name, payload);
    info!("sent participant discovery: {}  {}({})", info.id, info.role, info.component);
}

struct PortConsumer {
    process: ProcessFunction,
    portname: String,
    outqueue: String, // FIXME: allow sending on any port, also multiple times
}

fn send_out(channel: &mut Channel, exchange: String, data: Vec<u8>) {

    let routing_key = "".to_string();
    let content_type = Some("application/json".to_string()); // TODO: should be parameter
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };
    let s = channel.basic_publish(exchange, routing_key, true, false, props, data);
    s.expect("failed to send");
    debug!("sent output");
}

impl Consumer for PortConsumer {
    fn handle_delivery(&mut self,
                       channel: &mut Channel,
                       deliver: protocol::basic::Deliver,
                       headers: protocol::basic::BasicProperties,
                       body: Vec<u8>) {

        debug!("calling process()");
        let f = self.process;
        let res = f(body);
        debug!("process() returned");

        if res.is_ok() {
            debug!("ACKing and sending");
            send_out(channel, self.outqueue.to_string(), res.unwrap());
            let r = channel.basic_ack(deliver.delivery_tag, false);
        } else {
            error!("process() errored");
            let r = channel.basic_nack(deliver.delivery_tag, true, true);
        }
        // FIXME: send error data
    }
}

// FIXME: actually call ProcessFunction
fn setup_inport(participant: &Participant, port: &ParticipantPort, connection: &mut Connection) {
    debug!("setup inport: {}", port.queue.to_string());

    let consumer = PortConsumer {
        process: participant.process,
        portname: port.id.to_string(),
        outqueue: participant.info.outports[0].queue.to_string(),
    };

    // create
    let declare = connection.channel.queue_declare(port.queue.to_string(), false, true, false, false, false, Table::new());
    declare.expect("inport queue creation failed");

    // subscribe
    let q = port.queue.to_string();
    let cons = connection.channel.basic_consume(consumer, q,
                                                "".to_string(), false, false, false, false, Table::new());

    cons.expect("iport setup failed");

    debug!("inport setup done: {:?}, {:?}", port.id.to_string(), port.queue.to_string());
}

fn setup_outport(participant: &Participant, port: &ParticipantPort, connection: &mut Connection) {

    let exchange_type = "fanout".to_string();
    let declare = connection.channel.exchange_declare(port.queue.to_string(), exchange_type,
                                                    false, true, false, false, false, Table::new());
    declare.expect("outport setup failed");

    debug!("setup outport done: {:?}, {:?}", port.id.to_string(), port.queue.to_string());
}

fn start_participant(participant: &Participant, options: &ParticipantOptions) -> Connection {

    let mut session = Session::open_url(&options.broker).expect("Can't create AMQP session");
    let mut channel = session.open_channel(1).expect("could not open AMQP channel");

    let mut conn = Connection { session: session, channel: channel };

    // setup ports
    setup_inport(&participant, &participant.info.inports[0], &mut conn);
    setup_outport(&participant, &participant.info.outports[0], &mut conn);

    // send MsgFlo participant discovery message
    send_discovery(&mut conn.channel, &participant.info);

    return conn;
}

fn stop_participant(participant: &Participant, connection: &mut Connection) {

    let closed = connection.channel.close(200, "Bye".to_string());
    connection.session.close(200, "Good Bye".to_string());
}

#[derive(Debug)]
struct ParticipantOptions {
    role: String,
    broker: String,
}

impl Default for ParticipantOptions {
    fn default() -> ParticipantOptions {
        use rand::{thread_rng, Rng};

        let id: String = thread_rng().gen_ascii_chars().take(5).collect();
        ParticipantOptions {
            broker: "amqp://localhost//".to_string(),
            role: "".to_string(), // TODO: allow anonymous, format!("msgflo-rust-{}", id),
        }
    }
}

fn parse(options: &mut ParticipantOptions) {

    use argparse::{StoreTrue, Store};
    let mut parser = argparse::ArgumentParser::new();

    parser.refer(&mut options.role)
        .add_option(&["--role"], Store, "Participant role name");
    parser.refer(&mut options.broker)
        .add_option(&["--broker"], Store, "Address of messaging broker")
        .envvar("MSGFLO_BROKER");

    parser.parse_args_or_exit(); // XXX: should return out
}

// XXX: seems rust-amqp makes program hangs forever if error occurs / channel is borked?
// TODO: pass port info in/out of process()
// TODO: nicer way to declare ports? ideally they are enums not stringly typed?
pub fn participant_main(p: Participant) {

    let mut options = ParticipantOptions { .. Default::default() };
    parse(&mut options);

    if options.role == "" {
        options.role = p.info.role.to_string();
    }

    let mut c = start_participant(&p, &options);
    println!("{}({}) started", &options.role, &p.info.component);
    c.channel.start_consuming();

    stop_participant(&p, &mut c);
}

<|endoftext|>


<commit_before>//! A simple but effective sign library, written in Rust.
//!
//! Ported from [nobi](https://github.com/cyx/nobi),
//! which itself is a port of [itsdangerous](http://pythonhosted.org/itsdangerous/).
//!
//! # Possible use cases
//!
//! * Creating an activation link for users
//! * Creating a password reset link
//!
//! # Basic Example:
//!
//! ```rust
//! use nobsign::Signer;
//! let signer = Signer::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // Later check the signature and get the value back
//! let unsigned = signer.unsign(&signed).unwrap();
//! ```
//!
//! # Example with timestamped signatures
//!
//! ```rust
//! use nobsign::TimestampSigner;
//! let signer = TimestampSigner::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // In your code, you can verify the expiration:
//! signer.unsign(&signed, 86400).unwrap(); // 1 day expiration
//! ```
extern crate rustc_serialize;
extern crate ring;
extern crate time;
extern crate byteorder;

// Use same EPOCH as nobi, the Ruby implementation
const EPOCH : i32 = 1293840000;

use ring::{digest, hmac};
use rustc_serialize::base64::{ToBase64, FromBase64, URL_SAFE};
use byteorder::{ByteOrder, LittleEndian};

#[derive(Debug,PartialEq)]
pub enum Error {
    BadData,
    BadSignature,
    BadTimeSignature,
    SignatureExpired,
}

pub struct Signer {
    separator: char,
    s_key: hmac::SigningKey,
    v_key: hmac::VerificationKey,
}

pub struct TimestampSigner {
    signer: Signer,
}

fn int_to_bytes(n: i32) -> [u8; 4] {
    let mut out = [0; 4];

    LittleEndian::write_i32(&mut out, n);
    out
}
fn bytes_to_int(n: &[u8]) -> i32 {
    assert!(n.len() == 4);
    LittleEndian::read_i32(n)
}

impl Signer {
    pub fn new(secret: &[u8]) -> Signer {
        static ALGORITHM: &'static digest::Algorithm = &digest::SHA1;
        let initial_key = hmac::SigningKey::new(ALGORITHM, secret);
        let derived_key = hmac::sign(&initial_key, b"nobi.Signer");

        Signer {
            separator: '.',
            s_key: hmac::SigningKey::new(ALGORITHM, derived_key.as_ref()),
            v_key: hmac::VerificationKey::new(ALGORITHM, derived_key.as_ref()),
        }
    }

    pub fn sign(&self, value: &str) -> String {
        format!("{}{}{}", value, self.separator, &self.signature(value))
    }

    pub fn unsign(&self, value: &str) -> Result<String, Error> {
        let mut splitted = value.rsplitn(2, |c| c == self.separator);

        let sig = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let sig = try!(sig.from_base64().map_err(|_| Error::BadSignature));
        try!(hmac::verify(&self.v_key, value.as_bytes(), &sig)
                .map_err(|_| Error::BadSignature));

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        let sig = hmac::sign(&self.s_key, value.as_bytes());
        sig.as_ref().to_base64(URL_SAFE)
    }
}

impl TimestampSigner {
    pub fn new(secret: &[u8]) -> TimestampSigner {
        TimestampSigner { signer: Signer::new(secret) }
    }

    fn get_timestamp(&self) -> i32 {
        let now = time::now_utc().to_timespec();
        now.sec as i32 - EPOCH
    }

    pub fn sign(&self, value: &str) -> String {
        let timestamp = self.get_timestamp();
        let timestamp = int_to_bytes(timestamp);
        let timestamp = timestamp.to_base64(URL_SAFE);
        let value = format!("{}{}{}", value, self.signer.separator, timestamp);

        format!("{}{}{}",
                value,
                self.signer.separator,
                &self.signature(&value))
    }

    pub fn unsign(&self, value: &str, max_age: u32) -> Result<String, Error> {
        let result = try!(self.signer.unsign(value));

        if !result.contains(self.signer.separator) {
            return Err(Error::BadTimeSignature);
        }

        let mut splitted = result.rsplitn(2, |c| c == self.signer.separator);

        let timestamp = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let timestamp = match timestamp.from_base64() {
            Err(_) => return Err(Error::BadTimeSignature),
            Ok(timestamp) => timestamp,
        };
        if timestamp.len() != 4 {
            return Err(Error::BadTimeSignature);
        }
        let timestamp = bytes_to_int(&timestamp);
        let age = self.get_timestamp() - timestamp;

        if age > max_age as i32 {
            return Err(Error::SignatureExpired);
        }

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        self.signer.signature(value)
    }
}

#[test]
fn signs() {
    let signer = Signer::new(b"my-key");

    let signed = signer.sign("value");
    assert_eq!("value.EWkF3-80sipsPgLQ01NuTuPb0jQ".to_owned(), signed);

    assert_eq!("value".to_owned(), signer.unsign(&signed).unwrap());
}

#[test]
fn bad_unsign() {
    let signer = Signer::new(b"my-key");

    let signed = "value.ABCDEF";
    assert_eq!(Err(Error::BadSignature), signer.unsign(&signed));
}

#[test]
fn signs_with_timestamp() {
    let signer = TimestampSigner::new(b"my-key");

    let signed = signer.sign("value");

    assert_eq!("value".to_owned(), signer.unsign(&signed, 100).unwrap());
}

#[test]
fn bad_unsign_with_timestamp() {
    let signer = TimestampSigner::new(b"my-key");

    let signed = "value.ABCDEF";
    assert_eq!(Err(Error::BadSignature), signer.unsign(&signed, 10));

    let signed = "value.EWkF3-80sipsPgLQ01NuTuPb0jQ";
    assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));

    let signed = "value.AB.HHdWpuF7QVDZ_02wvECHvtV8vIc";
    assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));
}

#[test]
fn unsign_expired() {
    use std::thread::sleep_ms;
    let signer = TimestampSigner::new(b"my-key");
    let signed = signer.sign("value");
    sleep_ms(1000);
    assert_eq!(Err(Error::SignatureExpired), signer.unsign(&signed, 0));
}

#[test]
fn with_secure_secret() {
    use ring::{digest, rand};
    let mut key = vec![0u8; digest::SHA1.digest_len];
    rand::fill_secure_random(&mut key[..]).unwrap();

    let signer = Signer::new(&key);
    let signed = signer.sign("Hello world!");

    assert_eq!("Hello world!", signer.unsign(&signed).unwrap());
}
<commit_msg>fix: Move algorithm reference to top and use it<commit_after>//! A simple but effective sign library, written in Rust.
//!
//! Ported from [nobi](https://github.com/cyx/nobi),
//! which itself is a port of [itsdangerous](http://pythonhosted.org/itsdangerous/).
//!
//! # Possible use cases
//!
//! * Creating an activation link for users
//! * Creating a password reset link
//!
//! # Basic Example:
//!
//! ```rust
//! use nobsign::Signer;
//! let signer = Signer::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // Later check the signature and get the value back
//! let unsigned = signer.unsign(&signed).unwrap();
//! ```
//!
//! # Example with timestamped signatures
//!
//! ```rust
//! use nobsign::TimestampSigner;
//! let signer = TimestampSigner::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // In your code, you can verify the expiration:
//! signer.unsign(&signed, 86400).unwrap(); // 1 day expiration
//! ```
extern crate rustc_serialize;
extern crate ring;
extern crate time;
extern crate byteorder;

// Use same EPOCH as nobi, the Ruby implementation
const EPOCH : i32 = 1293840000;

use ring::{digest, hmac};
use rustc_serialize::base64::{ToBase64, FromBase64, URL_SAFE};
use byteorder::{ByteOrder, LittleEndian};

static ALGORITHM: &'static digest::Algorithm = &digest::SHA1;

#[derive(Debug,PartialEq)]
pub enum Error {
    BadData,
    BadSignature,
    BadTimeSignature,
    SignatureExpired,
}

pub struct Signer {
    separator: char,
    s_key: hmac::SigningKey,
    v_key: hmac::VerificationKey,
}

pub struct TimestampSigner {
    signer: Signer,
}

fn int_to_bytes(n: i32) -> [u8; 4] {
    let mut out = [0; 4];

    LittleEndian::write_i32(&mut out, n);
    out
}
fn bytes_to_int(n: &[u8]) -> i32 {
    assert!(n.len() == 4);
    LittleEndian::read_i32(n)
}

impl Signer {
    pub fn new(secret: &[u8]) -> Signer {
        let initial_key = hmac::SigningKey::new(ALGORITHM, secret);
        let derived_key = hmac::sign(&initial_key, b"nobi.Signer");

        Signer {
            separator: '.',
            s_key: hmac::SigningKey::new(ALGORITHM, derived_key.as_ref()),
            v_key: hmac::VerificationKey::new(ALGORITHM, derived_key.as_ref()),
        }
    }

    pub fn sign(&self, value: &str) -> String {
        format!("{}{}{}", value, self.separator, &self.signature(value))
    }

    pub fn unsign(&self, value: &str) -> Result<String, Error> {
        let mut splitted = value.rsplitn(2, |c| c == self.separator);

        let sig = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let sig = try!(sig.from_base64().map_err(|_| Error::BadSignature));
        try!(hmac::verify(&self.v_key, value.as_bytes(), &sig)
                .map_err(|_| Error::BadSignature));

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        let sig = hmac::sign(&self.s_key, value.as_bytes());
        sig.as_ref().to_base64(URL_SAFE)
    }
}

impl TimestampSigner {
    pub fn new(secret: &[u8]) -> TimestampSigner {
        TimestampSigner { signer: Signer::new(secret) }
    }

    fn get_timestamp(&self) -> i32 {
        let now = time::now_utc().to_timespec();
        now.sec as i32 - EPOCH
    }

    pub fn sign(&self, value: &str) -> String {
        let timestamp = self.get_timestamp();
        let timestamp = int_to_bytes(timestamp);
        let timestamp = timestamp.to_base64(URL_SAFE);
        let value = format!("{}{}{}", value, self.signer.separator, timestamp);

        format!("{}{}{}",
                value,
                self.signer.separator,
                &self.signature(&value))
    }

    pub fn unsign(&self, value: &str, max_age: u32) -> Result<String, Error> {
        let result = try!(self.signer.unsign(value));

        if !result.contains(self.signer.separator) {
            return Err(Error::BadTimeSignature);
        }

        let mut splitted = result.rsplitn(2, |c| c == self.signer.separator);

        let timestamp = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let timestamp = match timestamp.from_base64() {
            Err(_) => return Err(Error::BadTimeSignature),
            Ok(timestamp) => timestamp,
        };
        if timestamp.len() != 4 {
            return Err(Error::BadTimeSignature);
        }
        let timestamp = bytes_to_int(&timestamp);
        let age = self.get_timestamp() - timestamp;

        if age > max_age as i32 {
            return Err(Error::SignatureExpired);
        }

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        self.signer.signature(value)
    }
}

#[test]
fn signs() {
    let signer = Signer::new(b"my-key");

    let signed = signer.sign("value");
    assert_eq!("value.EWkF3-80sipsPgLQ01NuTuPb0jQ".to_owned(), signed);

    assert_eq!("value".to_owned(), signer.unsign(&signed).unwrap());
}

#[test]
fn bad_unsign() {
    let signer = Signer::new(b"my-key");

    let signed = "value.ABCDEF";
    assert_eq!(Err(Error::BadSignature), signer.unsign(&signed));
}

#[test]
fn signs_with_timestamp() {
    let signer = TimestampSigner::new(b"my-key");

    let signed = signer.sign("value");

    assert_eq!("value".to_owned(), signer.unsign(&signed, 100).unwrap());
}

#[test]
fn bad_unsign_with_timestamp() {
    let signer = TimestampSigner::new(b"my-key");

    let signed = "value.ABCDEF";
    assert_eq!(Err(Error::BadSignature), signer.unsign(&signed, 10));

    let signed = "value.EWkF3-80sipsPgLQ01NuTuPb0jQ";
    assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));

    let signed = "value.AB.HHdWpuF7QVDZ_02wvECHvtV8vIc";
    assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));
}

#[test]
fn unsign_expired() {
    use std::thread::sleep_ms;
    let signer = TimestampSigner::new(b"my-key");
    let signed = signer.sign("value");
    sleep_ms(1000);
    assert_eq!(Err(Error::SignatureExpired), signer.unsign(&signed, 0));
}

#[test]
fn with_secure_secret() {
    use ring::rand;
    let mut key = vec![0u8; ALGORITHM.digest_len];
    rand::fill_secure_random(&mut key[..]).unwrap();

    let signer = Signer::new(&key);
    let signed = signer.sign("Hello world!");

    assert_eq!("Hello world!", signer.unsign(&signed).unwrap());
}
<|endoftext|>


<commit_before>#![crate_name = "msgpack"]

#![unstable = "this library is still in rapid development"]

#![feature(core)]
#![feature(collections)]
#![cfg_attr(test, feature(test))]

extern crate byteorder;
extern crate rustc_serialize as serialize;

pub const MSGPACK_VERSION : u32 = 5;

pub mod core;

pub use core::decode::{
    read_nil,
};

pub use core::decode::serialize::{
    Decoder,
};

/// Temporary. Move to error.rs module.
pub use core::{Error, ReadError, MarkerError};

#[cfg(test)]
mod bench;

/// Stage 1. Low-level decoding functions.
/// Stage 2. Value decoding functions.
/// Stage 3. Deserialization.
/// Stage 4. Low-level encoding functions.
/// Stage 5. Value encoding functions.
/// Stage 6. Serialization.
<commit_msg>[Misc] Fixed doc style.<commit_after>#![crate_name = "msgpack"]

#![unstable = "this library is still in rapid development"]

#![feature(core)]
#![feature(collections)]
#![cfg_attr(test, feature(test))]

extern crate byteorder;
extern crate rustc_serialize as serialize;

pub const MSGPACK_VERSION : u32 = 5;

pub mod core;

pub use core::decode::{
    read_nil,
};

pub use core::decode::serialize::{
    Decoder,
};

/// Temporary. Move to error.rs module.
pub use core::{Error, ReadError, MarkerError};

#[cfg(test)]
mod bench;

// Stage 1. Low-level decoding functions.
// Stage 2. Value decoding functions.
// Stage 3. Deserialization.
// Stage 4. Low-level encoding functions.
// Stage 5. Value encoding functions.
// Stage 6. Serialization.
<|endoftext|>


<commit_before>/*!
```
extern crate hertz;
extern crate mel;

fn main() {
    let sample_rate = 44100;
    let window_size = 4096;
    let power_spectrum_size = window_size / 2;
    let filter_count = 100;

    let mut calls: usize = 0;

    for (row, col, value) in mel::enumerate_mel_scaling_matrix(
        sample_rate,
        window_size,
        power_spectrum_size,
        filter_count,
    ) {
        calls += 1;
    }
    assert_eq!(calls, power_spectrum_size * filter_count);
}
```
*/

extern crate ndarray;
use ndarray::ArrayBase;

#[macro_use]
extern crate nalgebra;
use nalgebra::ApproxEq;

extern crate num;
use num::{Float, FromPrimitive};

extern crate hertz;
extern crate apodize;

extern crate itertools;
use itertools::linspace;

macro_rules! f64_from_usize {
    ($val:expr) => {
        <f64>::from_usize($val)
            .expect("type `f64` can't represent a specific value of type `usize` on this architecture.");
    }
}

macro_rules! usize_from_f64 {
    ($val:expr) => {
        <usize>::from_f64($val)
            .expect("type `usize` can't represent a specific value of type `f64` on this architecture.");
    }
}

#[inline]
pub fn hertz_from_mel(mel: f64) -> f64 {
    700. * ((10.).powf(mel / 2595.) - 1.)
}

#[inline]
pub fn mel_from_hertz(hertz: f64) -> f64 {
    2595. * (1. + hertz / 700.).log10()
}

#[test]
fn test_mel() {
    assert_approx_eq_eps!(
        549.64, mel_from_hertz(440.), 0.01);
    assert_approx_eq_eps!(
        440., hertz_from_mel(549.64), 0.01);

    let mel = 0.;
    assert_approx_eq!(mel, mel_from_hertz(hertz_from_mel(mel)));
    let mel = 100.;
    assert_approx_eq_eps!(
        mel, mel_from_hertz(hertz_from_mel(mel)), 0.0001);
    let mel = 3000.;
    assert_approx_eq_eps!(
        mel, mel_from_hertz(hertz_from_mel(mel)), 0.0001);

    let hertz = 0.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
    let hertz = 1000.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
    let hertz = 44100. / 2.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
}

/// col index changes fastest
pub struct MelScalingMatrixEnumerator<WindowIter> {
    // parameters

    /// equivalent to col count
    input_size: usize,
    /// equivalent to row count
    output_size: usize,

    mel_from_hertz: fn(f64) -> f64,
    hertz_from_mel: fn(f64) -> f64,

    max_hertz: f64,

    window_function: fn(usize) -> WindowIter,

    // state

    row_index: usize,
    col_index: usize,

    start_mels_iter: itertools::Linspace<f64>,
    end_mels_iter: itertools::Linspace<f64>,

    // this gets set anew for every row
    window_start: usize,
    window_size: usize,
    window_iter: WindowIter
}

impl<WindowIter> MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    #[inline]
    pub fn is_done(&self) -> bool {
        self.is_after_last_row() && self.is_after_last_col()
    }

    #[inline]
    pub fn is_after_last_row(&self) -> bool {
        self.output_size <= self.row_index
    }

    #[inline]
    pub fn is_after_last_col(&self) -> bool {
        self.input_size <= self.col_index
    }

    #[inline]
    pub fn index(&self) -> usize {
        self.row_index * self.input_size + self.col_index
    }

    #[inline]
    pub fn remaining(&self) -> usize {
        // TODO is this off by one ?
        self.len() - self.index()
    }
}

impl<WindowIter> Iterator for MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    type Item = (usize, usize, f64);

    #[inline]
    fn next(&mut self) -> Option<(usize, usize, f64)> {
        if self.is_done() {
            return None;
        }

        // start a new row/filter
        if self.is_after_last_col() {
            self.col_index = 0;
            self.row_index += 1;
            if self.is_after_last_row() {
                return None;
            }

            let start_mel = self.start_mels_iter.next()
                .expect("self.start_mels_iter.next()");
            let end_mel = self.end_mels_iter.next()
                .expect("self.end_mels_iter.next()");

            let hertz_from_mel = self.hertz_from_mel;
            let start_hertz = hertz_from_mel(start_mel);
            let end_hertz = hertz_from_mel(end_mel);

            // TODO maybe round or floor or ceil here
            self.window_start = usize_from_f64!(
                start_hertz /
                self.max_hertz *
                f64_from_usize!(self.output_size));

            let window_end = usize_from_f64!(
                end_hertz /
                self.max_hertz *
                f64_from_usize!(self.output_size));

            self.window_size = window_end - self.window_start;

            let window_function = self.window_function;
            self.window_iter = window_function(self.window_size);
        }

        let col = self.col_index;

        let value = if col < self.window_start {
            match self.window_iter.next() {
                Some(value) => {
                    value / f64_from_usize!(self.window_size)
                },
                None => 0.
            }
        } else {
            0.
        };

        self.col_index += 1;
        Some((self.row_index, col, value))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.remaining();
        (remaining, Some(remaining))
    }
}

impl<WindowIter> ExactSizeIterator for MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    #[inline]
    fn len(&self) -> usize {
        self.input_size * self.output_size
    }
}

/// (or many vectors at once in form of a matrix)
/// to transform that vector into mel space.
/// efficiently
/// rows are mel
/// cols are freq
/// TODO get frequency index pair for nth mel
/// TODO enumerate scaling vector
/// TODO which dimension first
pub fn enumerate_mel_scaling_matrix_base<WindowIter>(
    // TODO more descriptive argument names
    min_hertz: f64,
    max_hertz: f64,
    // TODO shape
    input_size: usize,
    output_size: usize,
    mel_from_hertz: fn(f64) -> f64,
    hertz_from_mel: fn(f64) -> f64,
    window_function: fn(usize) -> WindowIter
) -> MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    assert!(min_hertz < max_hertz);
    assert!(output_size < input_size);
    assert!(0 < output_size);
    assert!(0 < input_size);

    let min_mel = mel_from_hertz(min_hertz);
    let max_mel = mel_from_hertz(max_hertz);

    // initially start_mels_iter == end_mels_iter
    let mut start_mels_iter = linspace(min_mel, max_mel, output_size + 1);
    let mut end_mels_iter = linspace(max_mel, max_mel, output_size + 1);
    end_mels_iter.next();

    let start_mel = start_mels_iter.next().unwrap();
    let end_mel = end_mels_iter.next().unwrap();

    let start_hertz = hertz_from_mel(start_mel);
    let end_hertz = hertz_from_mel(end_mel);

    // TODO maybe round or floor or ceil here
    let start_index = usize_from_f64!(
        start_hertz / max_hertz * f64_from_usize!(output_size));

    let end_index = usize_from_f64!(
        end_hertz / max_hertz * f64_from_usize!(output_size));

    let window_size = end_index - start_index;

    let window_iter = window_function(window_size);

    MelScalingMatrixEnumerator::<WindowIter> {
        // parameters
        input_size: input_size,
        output_size: output_size,
        window_function: window_function,
        hertz_from_mel: hertz_from_mel,
        mel_from_hertz: mel_from_hertz,
        max_hertz: max_hertz,

        // state
        row_index: 0,
        col_index: 0,

        start_mels_iter: start_mels_iter,
        end_mels_iter: end_mels_iter,

        window_start: start_index,
        window_size: window_size,
        window_iter: window_iter,
    }
}

/// one filter per row
pub fn enumerate_mel_scaling_matrix(
    sample_rate: usize,
    window_size: usize,
    input_size: usize,
    output_size: usize,
) -> MelScalingMatrixEnumerator<apodize::TriangularWindowIter>
{
    enumerate_mel_scaling_matrix_base(
        hertz::rayleigh(
            f64_from_usize!(sample_rate),
            f64_from_usize!(window_size)),
        hertz::nyquist(f64_from_usize!(sample_rate)),
        input_size,
        output_size,
        mel_from_hertz,
        hertz_from_mel,
        apodize::triangular_iter
    )
}

#[test]
fn test_enumerate_mel_scaling_matrix() {
    let sample_rate = 44100;
    let window_size = 100;
    let power_spectrum_size = window_size / 2;
    let filter_count = 10;

    let mut data: ArrayBase<Vec<f64>, (usize, usize)> =
        ArrayBase::zeros((filter_count, power_spectrum_size));

    for (row, col, value) in enumerate_mel_scaling_matrix(
        sample_rate,
        window_size,
        power_spectrum_size,
        filter_count,
    ) {
        data[(row, col)] = value;
    }
    println!("{:?}", data);
}

// TODO test that identity matrix with
// output_size = input_size
// and conversion functions = identity

// TODO test in a small scale that scaling matrixes are correct
// then test it in emir directly and check correctness visually
<commit_msg>fixes - promising results, seems to be working<commit_after>/*!
```
extern crate hertz;
extern crate mel;

fn main() {
    let sample_rate = 44100;
    let window_size = 4096;
    let power_spectrum_size = window_size / 2;
    let filter_count = 100;

    let mut calls: usize = 0;

    for (row, col, value) in mel::enumerate_mel_scaling_matrix(
        sample_rate,
        window_size,
        power_spectrum_size,
        filter_count,
    ) {
        calls += 1;
    }
    assert_eq!(calls, power_spectrum_size * filter_count);
}
```
*/

extern crate ndarray;
use ndarray::ArrayBase;

#[macro_use]
extern crate nalgebra;
use nalgebra::ApproxEq;

extern crate num;
use num::{Float, FromPrimitive};

extern crate hertz;
extern crate apodize;

extern crate itertools;
use itertools::linspace;

macro_rules! f64_from_usize {
    ($val:expr) => {
        <f64>::from_usize($val)
            .expect("type `f64` can't represent a specific value of type `usize` on this architecture.");
    }
}

macro_rules! usize_from_f64 {
    ($val:expr) => {
        <usize>::from_f64($val)
            .expect("type `usize` can't represent a specific value of type `f64` on this architecture.");
    }
}

#[inline]
pub fn hertz_from_mel(mel: f64) -> f64 {
    700. * ((10.).powf(mel / 2595.) - 1.)
}

#[inline]
pub fn mel_from_hertz(hertz: f64) -> f64 {
    2595. * (1. + hertz / 700.).log10()
}

#[test]
fn test_mel() {
    assert_approx_eq_eps!(
        549.64, mel_from_hertz(440.), 0.01);
    assert_approx_eq_eps!(
        440., hertz_from_mel(549.64), 0.01);

    let mel = 0.;
    assert_approx_eq!(mel, mel_from_hertz(hertz_from_mel(mel)));
    let mel = 100.;
    assert_approx_eq_eps!(
        mel, mel_from_hertz(hertz_from_mel(mel)), 0.0001);
    let mel = 3000.;
    assert_approx_eq_eps!(
        mel, mel_from_hertz(hertz_from_mel(mel)), 0.0001);

    let hertz = 0.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
    let hertz = 1000.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
    let hertz = 44100. / 2.;
    assert_approx_eq!(hertz, hertz_from_mel(mel_from_hertz(hertz)));
}

/// col index changes fastest
pub struct MelScalingMatrixEnumerator<WindowIter> {
    // parameters

    /// equivalent to col count
    input_size: usize,
    /// equivalent to row count
    output_size: usize,

    mel_from_hertz: fn(f64) -> f64,
    hertz_from_mel: fn(f64) -> f64,

    max_hertz: f64,

    window_function: fn(usize) -> WindowIter,

    // state

    row_index: usize,
    col_index: usize,

    start_mels_iter: itertools::Linspace<f64>,
    end_mels_iter: itertools::Linspace<f64>,

    // this gets set anew for every row
    window_start: usize,
    window_size: usize,
    window_iter: WindowIter
}

impl<WindowIter> MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    #[inline]
    pub fn is_done(&self) -> bool {
        self.is_after_last_row() && self.is_after_last_col()
    }

    #[inline]
    pub fn is_after_last_row(&self) -> bool {
        self.output_size <= self.row_index
    }

    #[inline]
    pub fn is_after_last_col(&self) -> bool {
        self.input_size <= self.col_index
    }

    #[inline]
    pub fn index(&self) -> usize {
        self.row_index * self.input_size + self.col_index
    }

    #[inline]
    pub fn remaining(&self) -> usize {
        // TODO is this off by one ?
        self.len() - self.index()
    }
}

impl<WindowIter> Iterator for MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    type Item = (usize, usize, f64);

    #[inline]
    fn next(&mut self) -> Option<(usize, usize, f64)> {
        if self.is_done() {
            return None;
        }

        // start a new row/filter
        if self.is_after_last_col() {
            self.col_index = 0;
            self.row_index += 1;
            if self.is_after_last_row() {
                return None;
            }

            let start_mel = self.start_mels_iter.next()
                .expect("self.start_mels_iter.next()");
            let end_mel = self.end_mels_iter.next()
                .expect("self.end_mels_iter.next()");

            let hertz_from_mel = self.hertz_from_mel;
            let start_hertz = hertz_from_mel(start_mel);
            let end_hertz = hertz_from_mel(end_mel);

            // TODO maybe round or floor or ceil here
            self.window_start = usize_from_f64!(
                start_hertz /
                self.max_hertz *
                f64_from_usize!(self.input_size));

            let window_end = usize_from_f64!(
                end_hertz /
                self.max_hertz *
                f64_from_usize!(self.input_size));

            self.window_size = window_end - self.window_start;

            println!("window_start = {}, window_end = {}, window_size = {}", self.window_start, window_end, self.window_size);

            let window_function = self.window_function;
            self.window_iter = window_function(self.window_size);
        }

        let col = self.col_index;

        let value = if col < self.window_start {
            0.
        } else {
            match self.window_iter.next() {
                Some(value) => {
                    value / f64_from_usize!(self.window_size)
                },
                None => 0.
            }
        };

        self.col_index += 1;
        Some((self.row_index, col, value))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.remaining();
        (remaining, Some(remaining))
    }
}

impl<WindowIter> ExactSizeIterator for MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    #[inline]
    fn len(&self) -> usize {
        self.input_size * self.output_size
    }
}

/// (or many vectors at once in form of a matrix)
/// to transform that vector into mel space.
/// efficiently
/// rows are mel
/// cols are freq
/// TODO get frequency index pair for nth mel
/// TODO enumerate scaling vector
/// TODO which dimension first
pub fn enumerate_mel_scaling_matrix_base<WindowIter>(
    // TODO more descriptive argument names
    min_hertz: f64,
    max_hertz: f64,
    // TODO shape
    input_size: usize,
    output_size: usize,
    mel_from_hertz: fn(f64) -> f64,
    hertz_from_mel: fn(f64) -> f64,
    window_function: fn(usize) -> WindowIter
) -> MelScalingMatrixEnumerator<WindowIter>
    where WindowIter: Iterator<Item=f64>
{
    assert!(min_hertz < max_hertz);
    assert!(output_size < input_size);
    assert!(0 < output_size);
    assert!(0 < input_size);

    let min_mel = mel_from_hertz(min_hertz);
    let max_mel = mel_from_hertz(max_hertz);

    // initially start_mels_iter == end_mels_iter
    let mut start_mels_iter = linspace(min_mel, max_mel, output_size + 1);
    let mut end_mels_iter = linspace(min_mel, max_mel, output_size + 1);
    end_mels_iter.next();

    let start_mel = start_mels_iter.next().unwrap();
    let end_mel = end_mels_iter.next().unwrap();

    let start_hertz = hertz_from_mel(start_mel);
    let end_hertz = hertz_from_mel(end_mel);

    // TODO maybe round or floor or ceil here
    let start_index = usize_from_f64!(
        start_hertz / max_hertz * f64_from_usize!(input_size));

    let end_index = usize_from_f64!(
        end_hertz / max_hertz * f64_from_usize!(input_size));

    let window_size = end_index - start_index;

    println!("window_start = {}, window_end = {}, window_size = {}", start_index, end_index, window_size);

    let window_iter = window_function(window_size);

    MelScalingMatrixEnumerator::<WindowIter> {
        // parameters
        input_size: input_size,
        output_size: output_size,
        window_function: window_function,
        hertz_from_mel: hertz_from_mel,
        mel_from_hertz: mel_from_hertz,
        max_hertz: max_hertz,

        // state
        row_index: 0,
        col_index: 0,

        start_mels_iter: start_mels_iter,
        end_mels_iter: end_mels_iter,

        window_start: start_index,
        window_size: window_size,
        window_iter: window_iter,
    }
}

/// one filter per row
pub fn enumerate_mel_scaling_matrix(
    sample_rate: usize,
    window_size: usize,
    input_size: usize,
    output_size: usize,
) -> MelScalingMatrixEnumerator<apodize::TriangularWindowIter>
{
    enumerate_mel_scaling_matrix_base(
        hertz::rayleigh(
            f64_from_usize!(sample_rate),
            f64_from_usize!(window_size)),
        hertz::nyquist(f64_from_usize!(sample_rate)),
        input_size,
        output_size,
        mel_from_hertz,
        hertz_from_mel,
        apodize::triangular_iter
    )
}

#[test]
fn test_enumerate_mel_scaling_matrix() {
    let sample_rate = 44100;
    let window_size = 100;
    let power_spectrum_size = window_size / 2;
    let filter_count = 10;

    let mut data: ArrayBase<Vec<f64>, (usize, usize)> =
        ArrayBase::zeros((filter_count, power_spectrum_size));

    for (row, col, value) in enumerate_mel_scaling_matrix(
        sample_rate,
        window_size,
        power_spectrum_size,
        filter_count,
    ) {
        data[(row, col)] = value;
    }
    println!("{:?}", data);
}

// TODO test that identity matrix with
// output_size = input_size
// and conversion functions = identity

// TODO test in a small scale that scaling matrixes are correct
// then test it in emir directly and check correctness visually
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate chrono;
extern crate rustc_serialize;
extern crate xml as xml_rs;

pub mod binary;
pub mod xml;
mod builder;

use chrono::{DateTime, UTC};
use chrono::format::ParseError as ChronoParseError;
use std::collections::BTreeMap;
use std::io::{Read, Seek, SeekFrom};
use std::io::Error as IoError;

#[derive(Clone, Debug, PartialEq)]
pub enum Plist {
    Array(Vec<Plist>),
    Dictionary(BTreeMap<String, Plist>),
    Boolean(bool),
    Data(Vec<u8>),
    Date(DateTime<UTC>),
    Real(f64),
    Integer(i64),
    String(String),
}

use rustc_serialize::base64::{STANDARD, ToBase64};
use rustc_serialize::json::Json as RustcJson;

impl Plist {
    pub fn from_events<T>(events: T) -> Result<Plist, ()>
        where T: IntoIterator<Item = ReadResult<PlistEvent>>
    {
        let iter = events.into_iter();
        let builder = builder::Builder::new(iter);

        match builder.build() {
            Ok(plist) => Ok(plist),
            Err(_) => Err(()),
        }
    }

    pub fn into_events(self) -> Vec<PlistEvent> {
        let mut events = Vec::new();
        self.into_events_inner(&mut events);
        events
    }

    fn into_events_inner(self, events: &mut Vec<PlistEvent>) {
        match self {
            Plist::Array(array) => {
                events.push(PlistEvent::StartArray(Some(array.len() as u64)));
                for value in array.into_iter() {
                    value.into_events_inner(events);
                }
                events.push(PlistEvent::EndArray);
            }
            Plist::Dictionary(dict) => {
                events.push(PlistEvent::StartDictionary(Some(dict.len() as u64)));
                for (key, value) in dict.into_iter() {
                    events.push(PlistEvent::StringValue(key));
                    value.into_events_inner(events);
                }
                events.push(PlistEvent::EndDictionary);
            }
            Plist::Boolean(value) => events.push(PlistEvent::BooleanValue(value)),
            Plist::Data(value) => events.push(PlistEvent::DataValue(value)),
            Plist::Date(value) => events.push(PlistEvent::DateValue(value)),
            Plist::Real(value) => events.push(PlistEvent::RealValue(value)),
            Plist::Integer(value) => events.push(PlistEvent::IntegerValue(value)),
            Plist::String(value) => events.push(PlistEvent::StringValue(value)),
        }
    }

    pub fn into_rustc_serialize_json(self) -> RustcJson {
        match self {
            Plist::Array(value) => {
                RustcJson::Array(value.into_iter().map(|p| p.into_rustc_serialize_json()).collect())
            }
            Plist::Dictionary(value) => {
                RustcJson::Object(value.into_iter()
                                       .map(|(k, v)| (k, v.into_rustc_serialize_json()))
                                       .collect())
            }
            Plist::Boolean(value) => RustcJson::Boolean(value),
            Plist::Data(value) => RustcJson::String(value.to_base64(STANDARD)),
            Plist::Date(value) => RustcJson::String(value.to_rfc3339()),
            Plist::Real(value) => RustcJson::F64(value),
            Plist::Integer(value) => RustcJson::I64(value),
            Plist::String(value) => RustcJson::String(value),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum PlistEvent {
    StartPlist,
    EndPlist,

    StartArray(Option<u64>),
    EndArray,

    StartDictionary(Option<u64>),
    EndDictionary,

    BooleanValue(bool),
    DataValue(Vec<u8>),
    DateValue(DateTime<UTC>),
    IntegerValue(i64),
    RealValue(f64),
    StringValue(String),
}

pub type ReadResult<T> = Result<T, ReadError>;

#[derive(Debug)]
pub enum ReadError {
    InvalidData,
    UnexpectedEof,
    UnsupportedType,
    Io(IoError),
}

impl From<IoError> for ReadError {
    fn from(io_error: IoError) -> ReadError {
        ReadError::Io(io_error)
    }
}

impl From<ChronoParseError> for ReadError {
    fn from(_: ChronoParseError) -> ReadError {
        ReadError::InvalidData
    }
}

pub enum EventReader<R: Read + Seek> {
    Uninitialized(Option<R>),
    Xml(xml::EventReader<R>),
    Binary(binary::EventReader<R>),
}

impl<R: Read + Seek> EventReader<R> {
    pub fn new(reader: R) -> EventReader<R> {
        EventReader::Uninitialized(Some(reader))
    }

    fn is_binary(reader: &mut R) -> Result<bool, IoError> {
        try!(reader.seek(SeekFrom::Start(0)));
        let mut magic = [0; 8];
        try!(reader.read(&mut magic));
        try!(reader.seek(SeekFrom::Start(0)));

        Ok(if &magic == b"bplist00" {
            true
        } else {
            false
        })
    }
}

impl<R: Read + Seek> Iterator for EventReader<R> {
    type Item = ReadResult<PlistEvent>;

    fn next(&mut self) -> Option<ReadResult<PlistEvent>> {
        let mut reader = match *self {
            EventReader::Xml(ref mut parser) => return parser.next(),
            EventReader::Binary(ref mut parser) => return parser.next(),
            EventReader::Uninitialized(ref mut reader) => reader.take().unwrap(),
        };

        let event_reader = match EventReader::is_binary(&mut reader) {
            Ok(true) => EventReader::Binary(binary::EventReader::new(reader)),
            Ok(false) => EventReader::Xml(xml::EventReader::new(reader)),
            Err(err) => {
                ::std::mem::replace(self, EventReader::Uninitialized(Some(reader)));
                return Some(Err(ReadError::Io(err)));
            }
        };

        ::std::mem::replace(self, event_reader);

        self.next()
    }
}
<commit_msg>Hide EventReader variants<commit_after>extern crate byteorder;
extern crate chrono;
extern crate rustc_serialize;
extern crate xml as xml_rs;

pub mod binary;
pub mod xml;
mod builder;

use chrono::{DateTime, UTC};
use chrono::format::ParseError as ChronoParseError;
use std::collections::BTreeMap;
use std::io::{Read, Seek, SeekFrom};
use std::io::Error as IoError;

#[derive(Clone, Debug, PartialEq)]
pub enum Plist {
    Array(Vec<Plist>),
    Dictionary(BTreeMap<String, Plist>),
    Boolean(bool),
    Data(Vec<u8>),
    Date(DateTime<UTC>),
    Real(f64),
    Integer(i64),
    String(String),
}

use rustc_serialize::base64::{STANDARD, ToBase64};
use rustc_serialize::json::Json as RustcJson;

impl Plist {
    pub fn from_events<T>(events: T) -> Result<Plist, ()>
        where T: IntoIterator<Item = ReadResult<PlistEvent>>
    {
        let iter = events.into_iter();
        let builder = builder::Builder::new(iter);

        match builder.build() {
            Ok(plist) => Ok(plist),
            Err(_) => Err(()),
        }
    }

    pub fn into_events(self) -> Vec<PlistEvent> {
        let mut events = Vec::new();
        self.into_events_inner(&mut events);
        events
    }

    fn into_events_inner(self, events: &mut Vec<PlistEvent>) {
        match self {
            Plist::Array(array) => {
                events.push(PlistEvent::StartArray(Some(array.len() as u64)));
                for value in array.into_iter() {
                    value.into_events_inner(events);
                }
                events.push(PlistEvent::EndArray);
            }
            Plist::Dictionary(dict) => {
                events.push(PlistEvent::StartDictionary(Some(dict.len() as u64)));
                for (key, value) in dict.into_iter() {
                    events.push(PlistEvent::StringValue(key));
                    value.into_events_inner(events);
                }
                events.push(PlistEvent::EndDictionary);
            }
            Plist::Boolean(value) => events.push(PlistEvent::BooleanValue(value)),
            Plist::Data(value) => events.push(PlistEvent::DataValue(value)),
            Plist::Date(value) => events.push(PlistEvent::DateValue(value)),
            Plist::Real(value) => events.push(PlistEvent::RealValue(value)),
            Plist::Integer(value) => events.push(PlistEvent::IntegerValue(value)),
            Plist::String(value) => events.push(PlistEvent::StringValue(value)),
        }
    }

    pub fn into_rustc_serialize_json(self) -> RustcJson {
        match self {
            Plist::Array(value) => {
                RustcJson::Array(value.into_iter().map(|p| p.into_rustc_serialize_json()).collect())
            }
            Plist::Dictionary(value) => {
                RustcJson::Object(value.into_iter()
                                       .map(|(k, v)| (k, v.into_rustc_serialize_json()))
                                       .collect())
            }
            Plist::Boolean(value) => RustcJson::Boolean(value),
            Plist::Data(value) => RustcJson::String(value.to_base64(STANDARD)),
            Plist::Date(value) => RustcJson::String(value.to_rfc3339()),
            Plist::Real(value) => RustcJson::F64(value),
            Plist::Integer(value) => RustcJson::I64(value),
            Plist::String(value) => RustcJson::String(value),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum PlistEvent {
    StartPlist,
    EndPlist,

    StartArray(Option<u64>),
    EndArray,

    StartDictionary(Option<u64>),
    EndDictionary,

    BooleanValue(bool),
    DataValue(Vec<u8>),
    DateValue(DateTime<UTC>),
    IntegerValue(i64),
    RealValue(f64),
    StringValue(String),
}

pub type ReadResult<T> = Result<T, ReadError>;

#[derive(Debug)]
pub enum ReadError {
    InvalidData,
    UnexpectedEof,
    UnsupportedType,
    Io(IoError),
}

impl From<IoError> for ReadError {
    fn from(io_error: IoError) -> ReadError {
        ReadError::Io(io_error)
    }
}

impl From<ChronoParseError> for ReadError {
    fn from(_: ChronoParseError) -> ReadError {
        ReadError::InvalidData
    }
}

pub struct EventReader<R: Read + Seek>(EventReaderInner<R>);

enum EventReaderInner<R: Read + Seek> {
    Uninitialized(Option<R>),
    Xml(xml::EventReader<R>),
    Binary(binary::EventReader<R>),
}

impl<R: Read + Seek> EventReader<R> {
    pub fn new(reader: R) -> EventReader<R> {
        EventReader(EventReaderInner::Uninitialized(Some(reader)))
    }

    fn is_binary(reader: &mut R) -> Result<bool, IoError> {
        try!(reader.seek(SeekFrom::Start(0)));
        let mut magic = [0; 8];
        try!(reader.read(&mut magic));
        try!(reader.seek(SeekFrom::Start(0)));

        Ok(if &magic == b"bplist00" {
            true
        } else {
            false
        })
    }
}

impl<R: Read + Seek> Iterator for EventReader<R> {
    type Item = ReadResult<PlistEvent>;

    fn next(&mut self) -> Option<ReadResult<PlistEvent>> {
        let mut reader = match self.0 {
            EventReaderInner::Xml(ref mut parser) => return parser.next(),
            EventReaderInner::Binary(ref mut parser) => return parser.next(),
            EventReaderInner::Uninitialized(ref mut reader) => reader.take().unwrap(),
        };

        let event_reader = match EventReader::is_binary(&mut reader) {
            Ok(true) => EventReaderInner::Binary(binary::EventReader::new(reader)),
            Ok(false) => EventReaderInner::Xml(xml::EventReader::new(reader)),
            Err(err) => {
                ::std::mem::replace(&mut self.0, EventReaderInner::Uninitialized(Some(reader)));
                return Some(Err(ReadError::Io(err)));
            }
        };

        ::std::mem::replace(&mut self.0, event_reader);

        self.next()
    }
}
<|endoftext|>


<commit_before>#![allow(unstable)]

use std::ops::{Index, IndexMut};

#[derive(Copy, Clone, PartialEq)]
pub struct Matrix4([[f32; 4]; 4]);

impl Index<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index<'a>(&'a self, row: &usize) -> &'a [f32; 4] {
        &self.0[*row]
    }
}

impl IndexMut<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index_mut<'a>(&'a mut self, row: &usize) -> &'a mut [f32; 4] {
        &mut self.0[*row]
    }
}

impl Matrix4 {
    pub fn as_array(&self) -> &[[f32; 4]; 4] {
        &self.0
    }

    pub fn from_array(array: [[f32; 4]; 4]) -> Self {
        Matrix4(array)
    }

    pub fn identity() -> Self {
        Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0],
                             [0.0, 0.0, 1.0, 0.0],
                             [0.0, 0.0, 0.0, 1.0]])
    }

	/// Builds a scale 4 * 4 matrix created from 3 scalars.
	/// Input matrix multiplied by this scale matrix.
    pub fn scale(&mut self, x: f32, y: f32, z: f32) {
        self.0[0][0] *= x;
        self.0[0][1] *= x;
        self.0[0][2] *= x;
        self.0[0][3] *= x;

        self.0[1][0] *= y;
        self.0[1][1] *= y;
        self.0[1][2] *= y;
        self.0[1][3] *= y;

        self.0[2][0] *= z;
        self.0[2][1] *= z;
        self.0[2][2] *= z;
        self.0[2][3] *= z;
    }

    pub fn scale_by_vector(&mut self, v: &Vector4) {
        self.scale(v[0], v[1], v[2])
    }

	/// Builds a translation 4 * 4 matrix created from a vector of 3 components.
	/// Input matrix multiplied by this translation matrix.
    pub fn translate(&mut self, x: f32, y: f32, z: f32) {
        self.0[3][0] += self.0[0][0] * x + self.0[1][0] * y + self.0[2][0] * z;
        self.0[3][1] += self.0[0][1] * x + self.0[1][1] * y + self.0[2][1] * z;
        self.0[3][2] += self.0[0][2] * x + self.0[1][2] * y + self.0[2][2] * z;
    }

    pub fn translate_by_vector(&mut self, v: &Vector4) {
        self.translate(v[0], v[1], v[2])
    }
}

#[derive(Copy, Clone, PartialEq)]
pub struct Vector4([f32; 4]);

impl Index<usize> for Vector4 {
    type Output = f32;

    fn index<'a>(&'a self, index: &usize) -> &'a f32 {
        &self.0[*index]
    }
}

impl IndexMut<usize> for Vector4 {
    type Output = f32;

    fn index_mut<'a>(&'a mut self, index: &usize) -> &'a mut f32 {
        &mut self.0[*index]
    }
}

impl Vector4 {
    pub fn as_array(&self) -> &[f32; 4] {
        &self.0
    }

    pub fn from_array(array: [f32; 4]) -> Self {
        Vector4(array)
    }

    pub fn new(x: f32, y: f32, z: f32) -> Self {
        Vector4::from_array([x, y, z, 1.0])
    }

    pub fn from_scalar(scalar: f32) -> Self {
        Vector4::from_array([scalar, scalar, scalar, 1.0])
    }
}

#[test]
fn test_matrix_to_array() {
    let array = {
        let foo = Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                                       [0.0, 1.0, 0.0, 0.0],
                                       [0.0, 0.0, 1.0, 0.0],
                                       [0.0, 0.0, 0.0, 1.0]]);
        assert_eq!(foo.as_array()[0][0], 1.0);
        *foo.as_array()
    };
    assert_eq!(array[1][1], 1.0);
}

#[test]
fn test_scale() {
    let mut m = Matrix4::identity();
    assert_eq!(m[2][2], 1.0);
    m.scale(3.0, 3.0, 3.0);
    assert_eq!(m[2][2], 3.0);
    m.scale_by_vector(&Vector4::new(2.0, 4.0, 3.0));
    assert_eq!(m[2][2], 9.0);
}

#[test]
fn test_translate() {
    let mut m = Matrix4::identity();
    m.translate(1.0, 2.0, 3.0);
    assert_eq!(m[3][0], 1.0);
    assert_eq!(m[3][1], 2.0);
    assert_eq!(m[3][2], 3.0);
}
<commit_msg>Vector4::new takes a w param<commit_after>#![allow(unstable)]

use std::ops::{Index, IndexMut};

#[derive(Copy, Clone, PartialEq)]
pub struct Matrix4([[f32; 4]; 4]);

impl Index<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index<'a>(&'a self, row: &usize) -> &'a [f32; 4] {
        &self.0[*row]
    }
}

impl IndexMut<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index_mut<'a>(&'a mut self, row: &usize) -> &'a mut [f32; 4] {
        &mut self.0[*row]
    }
}

impl Matrix4 {
    pub fn as_array(&self) -> &[[f32; 4]; 4] {
        &self.0
    }

    pub fn from_array(array: [[f32; 4]; 4]) -> Self {
        Matrix4(array)
    }

    pub fn identity() -> Self {
        Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0],
                             [0.0, 0.0, 1.0, 0.0],
                             [0.0, 0.0, 0.0, 1.0]])
    }

	/// Builds a scale 4 * 4 matrix created from 3 scalars.
	/// Input matrix multiplied by this scale matrix.
    pub fn scale(&mut self, x: f32, y: f32, z: f32) {
        self.0[0][0] *= x;
        self.0[0][1] *= x;
        self.0[0][2] *= x;
        self.0[0][3] *= x;

        self.0[1][0] *= y;
        self.0[1][1] *= y;
        self.0[1][2] *= y;
        self.0[1][3] *= y;

        self.0[2][0] *= z;
        self.0[2][1] *= z;
        self.0[2][2] *= z;
        self.0[2][3] *= z;
    }

    pub fn scale_by_vector(&mut self, v: &Vector4) {
        self.scale(v[0], v[1], v[2])
    }

	/// Builds a translation 4 * 4 matrix created from a vector of 3 components.
	/// Input matrix multiplied by this translation matrix.
    pub fn translate(&mut self, x: f32, y: f32, z: f32) {
        self.0[3][0] += self.0[0][0] * x + self.0[1][0] * y + self.0[2][0] * z;
        self.0[3][1] += self.0[0][1] * x + self.0[1][1] * y + self.0[2][1] * z;
        self.0[3][2] += self.0[0][2] * x + self.0[1][2] * y + self.0[2][2] * z;
    }

    pub fn translate_by_vector(&mut self, v: &Vector4) {
        self.translate(v[0], v[1], v[2])
    }
}

#[derive(Copy, Clone, PartialEq)]
pub struct Vector4([f32; 4]);

impl Index<usize> for Vector4 {
    type Output = f32;

    fn index<'a>(&'a self, index: &usize) -> &'a f32 {
        &self.0[*index]
    }
}

impl IndexMut<usize> for Vector4 {
    type Output = f32;

    fn index_mut<'a>(&'a mut self, index: &usize) -> &'a mut f32 {
        &mut self.0[*index]
    }
}

impl Vector4 {
    pub fn as_array(&self) -> &[f32; 4] {
        &self.0
    }

    pub fn from_array(array: [f32; 4]) -> Self {
        Vector4(array)
    }

    pub fn new(x: f32, y: f32, z: f32, w: f32) -> Self {
        Vector4::from_array([x, y, z, w])
    }

    pub fn from_scalar(scalar: f32) -> Self {
        Vector4::from_array([scalar, scalar, scalar, 1.0])
    }
}

#[test]
fn test_matrix_to_array() {
    let array = {
        let foo = Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                                       [0.0, 1.0, 0.0, 0.0],
                                       [0.0, 0.0, 1.0, 0.0],
                                       [0.0, 0.0, 0.0, 1.0]]);
        assert_eq!(foo.as_array()[0][0], 1.0);
        *foo.as_array()
    };
    assert_eq!(array[1][1], 1.0);
}

#[test]
fn test_scale() {
    let mut m = Matrix4::identity();
    assert_eq!(m[2][2], 1.0);
    m.scale(3.0, 3.0, 3.0);
    assert_eq!(m[2][2], 3.0);
    m.scale_by_vector(&Vector4::new(2.0, 4.0, 3.0, 1.0));
    assert_eq!(m[2][2], 9.0);
}

#[test]
fn test_translate() {
    let mut m = Matrix4::identity();
    m.translate(1.0, 2.0, 3.0);
    assert_eq!(m[3][0], 1.0);
    assert_eq!(m[3][1], 2.0);
    assert_eq!(m[3][2], 3.0);
}
<|endoftext|>


<commit_before>//! A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
//! uses HashMaps internally, and so Eq and Hash are both required.
//!
//! # Examples
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::for_strings();
//! chain.feed_str("I like cats and I like dogs.");
//! println!("{}", chain.generate_str());
//! ```
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new(0u8, 255);
//! chain.feed(vec![1u8, 2, 3, 5]).feed(vec![3u8, 9, 2]);
//! println!("{}", chain.generate());
//! ```
#![experimental]
#![feature(slicing_syntax)]
#![warn(missing_docs)]

use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::hash::Hash;
use std::io::{BufferedReader, File};
use std::rand::{Rng, task_rng};
use std::rc::Rc;

/// A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
/// uses HashMaps internally, and so Eq and Hash are both required.
pub struct Chain<T: Eq + Hash> {
    map: HashMap<Rc<T>, HashMap<Rc<T>, uint>>,
    start: Rc<T>,
    end: Rc<T>,
}

impl<T: Eq + Hash> Chain<T> {
    /// Constructs a new Markov chain using the given tokens as the marked starting and ending
    /// points for generation.
    pub fn new(start: T, end: T) -> Chain<T> {
        let start = Rc::new(start);
        let end = Rc::new(end);
        Chain {
            map: {
                let mut map = HashMap::new();
                map.insert(start.clone(), HashMap::new());
                map
            },
            start: start, end: end
        }
    }

    /// Determines whether or not the chain is empty. A chain is considered empty if nothing has
    /// been fed into it.
    pub fn is_empty(&self) -> bool {
        self.map[self.start.clone()].is_empty()
    }

    /// Feeds the chain a collection of tokens. This operation is O(n) where n is the number of
    /// tokens to be fed into the chain.
    pub fn feed(&mut self, tokens: Vec<T>) -> &mut Chain<T> {
        if tokens.len() == 0 { return self }
        let mut toks = Vec::new();
        toks.push(self.start.clone());
        toks.extend(tokens.into_iter().map(|token| {
            let rc = Rc::new(token);
            if !self.map.contains_key(&rc) {
                self.map.insert(rc.clone(), HashMap::new());
            }
            rc
        }));
        toks.push(self.end.clone());
        for p in toks.windows(2) {
            self.map[p[0]].add(p[1].clone());
        }
        self
    }

    /// Generates a collection of tokens from the chain. This operation is O(mn) where m is the
    /// length of the generated collection, and n is the number of possible states from a given
    /// state.
    pub fn generate(&self) -> Vec<Rc<T>> {
        let mut ret = Vec::new();
        let mut curs = self.start.clone();
        while curs != self.end {
            curs = self.map[curs].next();
            ret.push(curs.clone());
        }
        ret.pop();
        ret
    }

    /// Generates a collection of tokens from the chain, starting with the given token. This
    /// operation is O(mn) where m is the length of the generated collection, and n is the number
    /// of possible states from a given state. This returns an empty vector if the token is not
    /// found.
    pub fn generate_from_token(&self, token: T) -> Vec<Rc<T>> {
        let token = Rc::new(token);
        if !self.map.contains_key(&token) { return Vec::new() }
        let mut ret = vec![token.clone()];
        let mut curs = token;
        while curs != self.end {
            curs = self.map[curs].next();
            ret.push(curs.clone());
        }
        ret.pop();
        ret
    }
}

impl Chain<String> {
    /// Creates a new Chain intended specifically for strings. This uses the Unicode start of text
    /// and end of text control characters as the starting and ending tokens for the chain.
    pub fn for_strings() -> Chain<String> {
        Chain::new("\u0002".into_string(), "\u0003".into_string())
    }

    /// Feeds a string of text into the chain.
    pub fn feed_str(&mut self, string: &str) -> &mut Chain<String> {
        self.feed(string.split_str(" ").map(|s| s.into_string()).collect())
    }

    /// Feeds a properly formatted file into the chain. This file should be formatted such that
    /// each line is a new sentence. Punctuation may be included if it is desired.
    pub fn feed_file(&mut self, path: &Path) -> &mut Chain<String> {
        let mut reader = BufferedReader::new(File::open(path));
        for line in reader.lines() {
            let line = line.unwrap();
            let words: Vec<_> = line.split([' ', '\t', '\n', '\r'][])
                                    .filter(|word| !word.is_empty())
                                    .collect();
            self.feed(words.iter().map(|s| s.into_string()).collect());
        }
        self
    }

    /// Generates a random string of text.
    pub fn generate_str(&self) -> String {
        let vec = self.generate();
        let mut ret = String::new();
        for s in vec.iter() {
            ret.push_str(s[]);
            ret.push_str(" ");
        }
        let len = ret.len();
        ret.truncate(len - 1);
        ret
    }

    /// Generates a random string of text starting with the desired token. This returns an empty
    /// string if the token is not found.
    pub fn generate_str_from_token(&self, string: &str) -> String {
        let vec = self.generate_from_token(string.into_string());
        let mut ret = String::new();
        for s in vec.iter() {
            ret.push_str(s[]);
            ret.push_str(" ");
        }
        let len = ret.len();
        if len > 0 {
            ret.truncate(len - 1);
        }
        ret
    }
}

/// A collection of states for the Markov chain.
trait States<T: PartialEq> {
    /// Adds a state to this states collection.
    fn add(&mut self, token: Rc<T>);
    /// Gets the next state from this collection of states.
    fn next(&self) -> Rc<T>;
}

impl<T: Eq + Hash> States<T> for HashMap<Rc<T>, uint> {
    fn add(&mut self, token: Rc<T>) {
        match self.entry(token) {
            Occupied(mut e) => *e.get_mut() += 1,
            Vacant(e) => { e.set(1); },
        }
    }

    fn next(&self) -> Rc<T> {
        let mut sum = 0;
        for &value in self.values() {
            sum += value;
        }
        let mut rng = task_rng();
        let cap = rng.gen_range(0, sum);
        sum = 0;
        for (key, &value) in self.iter() {
            sum += value;
            if sum > cap {
                return key.clone()
            }
        }
        unreachable!("The random number generator failed.")
    }
}

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

    #[test]
    fn new() {
        Chain::new(0u, 100u);
        Chain::for_strings();
    }

    #[test]
    fn is_empty() {
        let mut chain = Chain::new(0u, 100u);
        assert!(chain.is_empty());
        chain.feed(vec![1u, 2u, 3u]);
        assert!(!chain.is_empty());
    }

    #[test]
    fn feed() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
    }

    #[test]
    fn generate() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate().map_in_place(|v| *v);
        assert!([vec![3u, 5u, 10u], vec![3u, 5u, 12u], vec![5u, 10u], vec![5u, 12u]].contains(&v));
    }

    #[test]
    fn generate_from_token() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate_from_token(5u).map_in_place(|v| *v);
        assert!([vec![5u, 10u], vec![5u, 12u]].contains(&v));
    }

    #[test]
    fn generate_from_unfound_token() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate_from_token(9u).map_in_place(|v| *v);
        assert_eq!(v, vec![]);
    }

    #[test]
    fn feed_str() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats and dogs");
    }

    #[test]
    fn generate_str() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("I hate cats");
        assert!(["I like cats", "I hate cats"].contains(&chain.generate_str()[]));
    }

    #[test]
    fn generate_str_from_token() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert!(["cats", "cats are cute"].contains(&chain.generate_str_from_token("cats")[]));
    }

    #[test]
    fn generate_str_from_unfound_token() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert_eq!(chain.generate_str_from_token("test"), "");
    }
}
<commit_msg>Updated for changes to unicode escape syntax.<commit_after>//! A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
//! uses HashMaps internally, and so Eq and Hash are both required.
//!
//! # Examples
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::for_strings();
//! chain.feed_str("I like cats and I like dogs.");
//! println!("{}", chain.generate_str());
//! ```
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new(0u8, 255);
//! chain.feed(vec![1u8, 2, 3, 5]).feed(vec![3u8, 9, 2]);
//! println!("{}", chain.generate());
//! ```
#![experimental]
#![feature(slicing_syntax)]
#![warn(missing_docs)]

use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::hash::Hash;
use std::io::{BufferedReader, File};
use std::rand::{Rng, task_rng};
use std::rc::Rc;

/// A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
/// uses HashMaps internally, and so Eq and Hash are both required.
pub struct Chain<T: Eq + Hash> {
    map: HashMap<Rc<T>, HashMap<Rc<T>, uint>>,
    start: Rc<T>,
    end: Rc<T>,
}

impl<T: Eq + Hash> Chain<T> {
    /// Constructs a new Markov chain using the given tokens as the marked starting and ending
    /// points for generation.
    pub fn new(start: T, end: T) -> Chain<T> {
        let start = Rc::new(start);
        let end = Rc::new(end);
        Chain {
            map: {
                let mut map = HashMap::new();
                map.insert(start.clone(), HashMap::new());
                map
            },
            start: start, end: end
        }
    }

    /// Determines whether or not the chain is empty. A chain is considered empty if nothing has
    /// been fed into it.
    pub fn is_empty(&self) -> bool {
        self.map[self.start.clone()].is_empty()
    }

    /// Feeds the chain a collection of tokens. This operation is O(n) where n is the number of
    /// tokens to be fed into the chain.
    pub fn feed(&mut self, tokens: Vec<T>) -> &mut Chain<T> {
        if tokens.len() == 0 { return self }
        let mut toks = Vec::new();
        toks.push(self.start.clone());
        toks.extend(tokens.into_iter().map(|token| {
            let rc = Rc::new(token);
            if !self.map.contains_key(&rc) {
                self.map.insert(rc.clone(), HashMap::new());
            }
            rc
        }));
        toks.push(self.end.clone());
        for p in toks.windows(2) {
            self.map[p[0]].add(p[1].clone());
        }
        self
    }

    /// Generates a collection of tokens from the chain. This operation is O(mn) where m is the
    /// length of the generated collection, and n is the number of possible states from a given
    /// state.
    pub fn generate(&self) -> Vec<Rc<T>> {
        let mut ret = Vec::new();
        let mut curs = self.start.clone();
        while curs != self.end {
            curs = self.map[curs].next();
            ret.push(curs.clone());
        }
        ret.pop();
        ret
    }

    /// Generates a collection of tokens from the chain, starting with the given token. This
    /// operation is O(mn) where m is the length of the generated collection, and n is the number
    /// of possible states from a given state. This returns an empty vector if the token is not
    /// found.
    pub fn generate_from_token(&self, token: T) -> Vec<Rc<T>> {
        let token = Rc::new(token);
        if !self.map.contains_key(&token) { return Vec::new() }
        let mut ret = vec![token.clone()];
        let mut curs = token;
        while curs != self.end {
            curs = self.map[curs].next();
            ret.push(curs.clone());
        }
        ret.pop();
        ret
    }
}

impl Chain<String> {
    /// Creates a new Chain intended specifically for strings. This uses the Unicode start of text
    /// and end of text control characters as the starting and ending tokens for the chain.
    pub fn for_strings() -> Chain<String> {
        Chain::new("\u{0002}".into_string(), "\u{0003}".into_string())
    }

    /// Feeds a string of text into the chain.
    pub fn feed_str(&mut self, string: &str) -> &mut Chain<String> {
        self.feed(string.split_str(" ").map(|s| s.into_string()).collect())
    }

    /// Feeds a properly formatted file into the chain. This file should be formatted such that
    /// each line is a new sentence. Punctuation may be included if it is desired.
    pub fn feed_file(&mut self, path: &Path) -> &mut Chain<String> {
        let mut reader = BufferedReader::new(File::open(path));
        for line in reader.lines() {
            let line = line.unwrap();
            let words: Vec<_> = line.split([' ', '\t', '\n', '\r'][])
                                    .filter(|word| !word.is_empty())
                                    .collect();
            self.feed(words.iter().map(|s| s.into_string()).collect());
        }
        self
    }

    /// Generates a random string of text.
    pub fn generate_str(&self) -> String {
        let vec = self.generate();
        let mut ret = String::new();
        for s in vec.iter() {
            ret.push_str(s[]);
            ret.push_str(" ");
        }
        let len = ret.len();
        ret.truncate(len - 1);
        ret
    }

    /// Generates a random string of text starting with the desired token. This returns an empty
    /// string if the token is not found.
    pub fn generate_str_from_token(&self, string: &str) -> String {
        let vec = self.generate_from_token(string.into_string());
        let mut ret = String::new();
        for s in vec.iter() {
            ret.push_str(s[]);
            ret.push_str(" ");
        }
        let len = ret.len();
        if len > 0 {
            ret.truncate(len - 1);
        }
        ret
    }
}

/// A collection of states for the Markov chain.
trait States<T: PartialEq> {
    /// Adds a state to this states collection.
    fn add(&mut self, token: Rc<T>);
    /// Gets the next state from this collection of states.
    fn next(&self) -> Rc<T>;
}

impl<T: Eq + Hash> States<T> for HashMap<Rc<T>, uint> {
    fn add(&mut self, token: Rc<T>) {
        match self.entry(token) {
            Occupied(mut e) => *e.get_mut() += 1,
            Vacant(e) => { e.set(1); },
        }
    }

    fn next(&self) -> Rc<T> {
        let mut sum = 0;
        for &value in self.values() {
            sum += value;
        }
        let mut rng = task_rng();
        let cap = rng.gen_range(0, sum);
        sum = 0;
        for (key, &value) in self.iter() {
            sum += value;
            if sum > cap {
                return key.clone()
            }
        }
        unreachable!("The random number generator failed.")
    }
}

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

    #[test]
    fn new() {
        Chain::new(0u, 100u);
        Chain::for_strings();
    }

    #[test]
    fn is_empty() {
        let mut chain = Chain::new(0u, 100u);
        assert!(chain.is_empty());
        chain.feed(vec![1u, 2u, 3u]);
        assert!(!chain.is_empty());
    }

    #[test]
    fn feed() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
    }

    #[test]
    fn generate() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate().map_in_place(|v| *v);
        assert!([vec![3u, 5u, 10u], vec![3u, 5u, 12u], vec![5u, 10u], vec![5u, 12u]].contains(&v));
    }

    #[test]
    fn generate_from_token() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate_from_token(5u).map_in_place(|v| *v);
        assert!([vec![5u, 10u], vec![5u, 12u]].contains(&v));
    }

    #[test]
    fn generate_from_unfound_token() {
        let mut chain = Chain::new(0u, 100u);
        chain.feed(vec![3u, 5u, 10u]).feed(vec![5u, 12u]);
        let v = chain.generate_from_token(9u).map_in_place(|v| *v);
        assert_eq!(v, vec![]);
    }

    #[test]
    fn feed_str() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats and dogs");
    }

    #[test]
    fn generate_str() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("I hate cats");
        assert!(["I like cats", "I hate cats"].contains(&chain.generate_str()[]));
    }

    #[test]
    fn generate_str_from_token() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert!(["cats", "cats are cute"].contains(&chain.generate_str_from_token("cats")[]));
    }

    #[test]
    fn generate_str_from_unfound_token() {
        let mut chain = Chain::for_strings();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert_eq!(chain.generate_str_from_token("test"), "");
    }
}
<|endoftext|>