<commit_before>#![feature(conservative_impl_trait)]

extern crate densearray;
extern crate rng;
extern crate sharedmem;

extern crate rand;

//use io::{IoBuffer};
use rw::{ReadBuffer, ReadAccumulateBuffer, WriteBuffer, AccumulateBuffer};

use rng::xorshift::{Xorshiftplus128Rng};

use rand::{Rng};
use std::cell::{Cell};
use std::io::{Read, Write};
use std::ops::{Deref};

pub mod data;
pub mod io;
pub mod opt;
pub mod prelude;
pub mod rw;
pub mod timing;

#[derive(Clone, Copy)]
pub enum OpCapability {
  Forward,
  Backward,
  RForward,
  RBackward,
}

impl OpCapability {
  pub fn enable_backward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      _ => true,
    }
  }

  pub fn enable_r_forward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      OpCapability::Backward => false,
      _ => true,
    }
  }

  pub fn enable_r_backward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      OpCapability::Backward => false,
      OpCapability::RForward => false,
      _ => true,
    }
  }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum OpPhase {
  Inference,
  Learning,
}

#[derive(Clone, Copy, Debug)]
pub enum Regularization {
  L2(f32),
}

pub trait Operator {
  fn _next(&self) -> u64;
  fn _epoch(&self) -> u64;
}

#[derive(Clone, Default)]
pub struct OperatorNode {
  pub curr_epoch:   Cell<u64>,
  pub curr_count:   Cell<u64>,
}

impl Operator for OperatorNode {
  fn _next(&self) -> u64 {
    self.curr_epoch.set(self.curr_epoch.get() + 1);
    self.curr_count.set(0);
    self.curr_epoch.get()
  }

  fn _epoch(&self) -> u64 {
    self.curr_epoch.get()
  }
}

impl OperatorNode {
  pub fn step(&self, next_epoch: u64) {
    assert!(next_epoch >= self.curr_epoch.get());
    if next_epoch > self.curr_epoch.get() {
      self.curr_epoch.set(next_epoch);
      self.curr_count.set(0);
    }
    self.curr_count.set(self.curr_count.get() + 1);
  }

  pub fn count(&self) -> u64 {
    self.curr_count.get()
  }

  pub fn limit(&self, max_count: u64) -> bool {
    self.curr_count.get() <= max_count
  }
}

pub trait OperatorOutput {
}

pub trait NewDiffOperator<S>: Operator {
  //type Output;
  type IoBuf: ?Sized;

  //fn _op_output(&self, arm: usize) -> impl OperatorOutput;
  //fn _output(&self, arm: usize) -> Self::Output;
  fn _traverse_fwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, IoBuf=Self::IoBuf>));
  fn _traverse_bwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, IoBuf=Self::IoBuf>));
  //fn _traverse_fwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, Output=Self::Output, IoBuf=Self::IoBuf>));
  //fn _traverse_bwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, Output=Self::Output, IoBuf=Self::IoBuf>));

  fn _diff_param_sz(&self) -> usize { 0 }
  fn _nondiff_param_sz(&self) -> usize { 0 }

  fn _load_diff_param(&mut self, offset: usize, param_reader: &mut Self::IoBuf) -> usize { 0 }
  fn _load_nondiff_param(&mut self, offset: usize, param_reader: &mut Self::IoBuf) -> usize { 0 }
  fn _store_diff_param(&mut self, offset: usize, param_writer: &mut Self::IoBuf) -> usize { 0 }
  fn _store_nondiff_param(&mut self, offset: usize, param_writer: &mut Self::IoBuf) -> usize { 0 }
  fn _store_grad(&mut self, offset: usize, grad_writer: &mut Self::IoBuf) -> usize { 0 }

  fn _save_rng_state(&mut self) {}
  fn _restore_rng_state(&mut self) {}

  fn _next_iteration(&mut self) {}
  fn _load_batch(&mut self, _samples: &[S]) {}
  fn _init_param(&mut self, rng: &mut Xorshiftplus128Rng) {}
  fn _update_nondiff_param(&mut self, _iter_nr: usize) {}
  fn _reset_grad(&mut self) {}

  fn _forward(&mut self, phase: OpPhase);
  fn _backward(&mut self);
}

pub trait DiffLoss<S>: NewDiffOperator<S> {
  fn reset_loss(&mut self);
  fn store_loss(&mut self) -> f32;
  fn _store_accuracy(&mut self) -> usize { 0 }

  fn diff_param_sz(&mut self) -> usize {
    let epoch = self._next();
    let mut grad_sz = 0;
    self._traverse_bwd(epoch, &mut |op| {
      grad_sz += op._diff_param_sz();
    });
    grad_sz
  }

  fn nondiff_param_sz(&mut self) -> usize {
    let epoch = self._next();
    let mut nondiff_sz = 0;
    self._traverse_bwd(epoch, &mut |op| {
      nondiff_sz += op._nondiff_param_sz();
    });
    nondiff_sz
  }

  fn load_diff_param(&mut self, param_reader: &mut Self::IoBuf) -> usize {
    //param_reader.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._load_diff_param(offset, param_reader);
    });
    offset
  }

  fn load_nondiff_param(&mut self, param_reader: &mut Self::IoBuf) -> usize {
    //param_reader.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._load_nondiff_param(offset, param_reader);
    });
    offset
  }

  fn store_diff_param(&mut self, param_writer: &mut Self::IoBuf) -> usize {
    //param_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_diff_param(offset, param_writer);
    });
    offset
  }

  fn store_nondiff_param(&mut self, param_writer: &mut Self::IoBuf) -> usize {
    //param_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_nondiff_param(offset, param_writer);
    });
    offset
  }

  fn store_grad(&mut self, grad_writer: &mut Self::IoBuf) -> usize {
    //grad_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_grad(offset, grad_writer);
    });
    offset
  }

  fn next_iteration(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._next_iteration());
  }

  fn load_batch(&mut self, samples: &[S]) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._load_batch(samples));
  }

  fn save_rng_state(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._save_rng_state());
  }

  fn restore_rng_state(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._restore_rng_state());
  }

  fn init_param(&mut self, rng: &mut Xorshiftplus128Rng) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._init_param(rng));
  }

  fn update_nondiff_param(&mut self, iter_nr: usize) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._update_nondiff_param(iter_nr));
  }

  fn reset_grad(&mut self) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._reset_grad());
  }

  fn forward(&mut self, phase: OpPhase) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._forward(phase));
  }

  fn backward(&mut self) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._backward());
  }
}

pub trait DiffOperator<T> {
  type Output;
  type Rng: Rng;

  fn _output(&self, arm: usize) -> Self::Output;

  fn diff_param_sz(&self) -> usize { 0 }
  fn nondiff_param_sz(&self) -> usize { 0 }
  fn total_param_sz(&self) -> usize {
    self.diff_param_sz() + self.nondiff_param_sz()
  }

  fn save_rng_state(&mut self) {}
  fn restore_rng_state(&mut self) {}
  //fn store_rng_state(&mut self, _rng_state: &mut Write) {}
  //fn load_rng_state(&mut self, _rng_state: &mut Read) {}

  //fn reset_state(&mut self) {}

  fn init_param(&mut self, _rng: &mut Self::Rng) {}
  //fn reset_nondiff_param(&mut self) {}
  fn reset_grad(&mut self) {}

  fn store_param(&mut self, _param_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn store_nondiff_param(&mut self, _param_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn load_param(&mut self, _param_reader: &mut ReadBuffer<T>, _offset: usize) -> usize { 0 }
  fn load_nondiff_param(&mut self, _param_reader: &mut ReadBuffer<T>, _offset: usize) -> usize { 0 }
  fn update_param(&mut self, _alpha: f32, _beta: f32, _grad_reader: &mut ReadAccumulateBuffer<T>, _offset: usize) -> usize { 0 }

  fn update_nondiff_param(&mut self, _iter: usize) {}

  fn store_grad(&mut self, _grad_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn accumulate_grad(&mut self, _alpha: f32, _beta: f32, _grad_accum: &mut AccumulateBuffer<T>, _offset: usize) -> usize { 0 }
  //fn grad_step(&mut self, _alpha: f32, _beta: f32) {}

  fn reset_loss(&mut self) {}
  fn store_loss(&mut self) -> f32 { 0.0 }
  fn _store_accuracy(&mut self) -> usize { 0 }

  fn apply_grad_reg(&mut self, _reg: Regularization) {}

  fn forward(&mut self, phase: OpPhase);
  //fn fwd_reg(&mut self, _reg: Regularization) {}
  fn backward(&mut self);
  //fn bwd_reg(&mut self, _reg: Regularization) {}
  //fn r_forward(&mut self) { unimplemented!(); }
  //fn r_backward(&mut self) { unimplemented!(); }
}

pub trait DiffOperatorIo<T, Buf>/*: DiffOperator*/ where T: Copy {
  fn store_diff_param(&mut self, _param_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn store_nondiff_param(&mut self, _param_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn load_diff_param(&mut self, _param_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn load_nondiff_param(&mut self, _param_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn update_diff_param(&mut self, _alpha: T, _beta: T, _delta_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn store_grad(&mut self, _grad_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn accumulate_grad(&mut self, _alpha: T, _beta: T, _grad_writer: &mut Buf, _offset: usize) -> usize { 0 }
}

pub trait DiffOperatorInput<T, S>: DiffOperator<T> {
  fn as_op(&self) -> &DiffOperator<T, Output=Self::Output, Rng=Self::Rng> { unimplemented!(); }
  fn load_data(&mut self, samples: &[S]);
}

pub trait DiffOperatorOutput<T, U>: DiffOperator<T> {
  fn get_output(&mut self) -> U;
}

pub trait CheckpointFormat {
}

pub trait DiffOperatorCheckpoint<Format> where Format: CheckpointFormat {
  fn decode(reader: &mut Read) -> Result<Self, ()> where Self: Sized;
  fn encode(&mut self, writer: &mut Write) -> Result<(), ()>;
}
<commit_msg>Sketched out an extension for traversing an operator graph where the nodes are of a more specialized type, commented out for now.<commit_after>//#![feature(associated_type_defaults)]
#![feature(conservative_impl_trait)]

extern crate densearray;
extern crate rng;
extern crate sharedmem;

extern crate rand;

//use io::{IoBuffer};
use rw::{ReadBuffer, ReadAccumulateBuffer, WriteBuffer, AccumulateBuffer};

use rng::xorshift::{Xorshiftplus128Rng};

use rand::{Rng};
use std::cell::{Cell, RefCell};
use std::io::{Read, Write};
use std::ops::{Deref};
use std::rc::{Rc};

pub mod data;
pub mod io;
pub mod opt;
pub mod prelude;
pub mod rw;
pub mod timing;

#[derive(Clone, Copy)]
pub enum OpCapability {
  Forward,
  Backward,
  RForward,
  RBackward,
}

impl OpCapability {
  pub fn enable_backward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      _ => true,
    }
  }

  pub fn enable_r_forward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      OpCapability::Backward => false,
      _ => true,
    }
  }

  pub fn enable_r_backward(&self) -> bool {
    match *self {
      OpCapability::Forward => false,
      OpCapability::Backward => false,
      OpCapability::RForward => false,
      _ => true,
    }
  }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum OpPhase {
  Inference,
  Learning,
}

#[derive(Clone, Copy, Debug)]
pub enum Regularization {
  L2(f32),
}

pub trait Operator {
  fn _next(&self) -> u64;
  fn _epoch(&self) -> u64;
}

#[derive(Clone, Default)]
pub struct OperatorNode {
  pub curr_epoch:   Cell<u64>,
  pub curr_count:   Cell<u64>,
}

impl Operator for OperatorNode {
  fn _next(&self) -> u64 {
    self.curr_epoch.set(self.curr_epoch.get() + 1);
    self.curr_count.set(0);
    self.curr_epoch.get()
  }

  fn _epoch(&self) -> u64 {
    self.curr_epoch.get()
  }
}

impl OperatorNode {
  pub fn step(&self, next_epoch: u64) {
    assert!(next_epoch >= self.curr_epoch.get());
    if next_epoch > self.curr_epoch.get() {
      self.curr_epoch.set(next_epoch);
      self.curr_count.set(0);
    }
    self.curr_count.set(self.curr_count.get() + 1);
  }

  pub fn count(&self) -> u64 {
    self.curr_count.get()
  }

  pub fn limit(&self, max_count: u64) -> bool {
    self.curr_count.get() <= max_count
  }
}

/*pub trait OperatorOutput {
}*/

pub trait NewDiffOperator<S>: Operator {
  type IoBuf: ?Sized;
  //type OpRef = Rc<RefCell<NewDiffOperator<S, IoBuf=Self::IoBuf, OpRef=Self::OpRef>>>;

  fn _traverse_fwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, IoBuf=Self::IoBuf>));
  fn _traverse_bwd(&mut self, _epoch: u64, _apply: &mut FnMut(&mut NewDiffOperator<S, IoBuf=Self::IoBuf>));
  //fn _traverse_fwd_new(&self, _epoch: u64, _apply: &mut FnMut(Self::OpRef));
  //fn _traverse_bwd_new(&self, _epoch: u64, _apply: &mut FnMut(Self::OpRef));

  fn _diff_param_sz(&self) -> usize { 0 }
  fn _nondiff_param_sz(&self) -> usize { 0 }

  fn _load_diff_param(&mut self, offset: usize, param_reader: &mut Self::IoBuf) -> usize { 0 }
  fn _load_nondiff_param(&mut self, offset: usize, param_reader: &mut Self::IoBuf) -> usize { 0 }
  fn _store_diff_param(&mut self, offset: usize, param_writer: &mut Self::IoBuf) -> usize { 0 }
  fn _store_nondiff_param(&mut self, offset: usize, param_writer: &mut Self::IoBuf) -> usize { 0 }
  fn _store_grad(&mut self, offset: usize, grad_writer: &mut Self::IoBuf) -> usize { 0 }

  fn _save_rng_state(&mut self) {}
  fn _restore_rng_state(&mut self) {}

  fn _next_iteration(&mut self) {}
  fn _load_batch(&mut self, _samples: &[S]) {}
  fn _init_param(&mut self, rng: &mut Xorshiftplus128Rng) {}
  fn _update_nondiff_param(&mut self, _iter_nr: usize) {}
  fn _reset_grad(&mut self) {}

  fn _forward(&mut self, phase: OpPhase);
  fn _backward(&mut self);
}

pub trait DiffLoss<S>: NewDiffOperator<S> {
  fn reset_loss(&mut self);
  fn store_loss(&mut self) -> f32;
  fn _store_accuracy(&mut self) -> usize { 0 }

  fn diff_param_sz(&mut self) -> usize {
    let epoch = self._next();
    let mut grad_sz = 0;
    self._traverse_bwd(epoch, &mut |op| {
      grad_sz += op._diff_param_sz();
    });
    grad_sz
  }

  fn nondiff_param_sz(&mut self) -> usize {
    let epoch = self._next();
    let mut nondiff_sz = 0;
    self._traverse_bwd(epoch, &mut |op| {
      nondiff_sz += op._nondiff_param_sz();
    });
    nondiff_sz
  }

  fn load_diff_param(&mut self, param_reader: &mut Self::IoBuf) -> usize {
    //param_reader.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._load_diff_param(offset, param_reader);
    });
    offset
  }

  fn load_nondiff_param(&mut self, param_reader: &mut Self::IoBuf) -> usize {
    //param_reader.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._load_nondiff_param(offset, param_reader);
    });
    offset
  }

  fn store_diff_param(&mut self, param_writer: &mut Self::IoBuf) -> usize {
    //param_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_diff_param(offset, param_writer);
    });
    offset
  }

  fn store_nondiff_param(&mut self, param_writer: &mut Self::IoBuf) -> usize {
    //param_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_nondiff_param(offset, param_writer);
    });
    offset
  }

  fn store_grad(&mut self, grad_writer: &mut Self::IoBuf) -> usize {
    //grad_writer.reset();
    let epoch = self._next();
    let mut offset = 0;
    self._traverse_bwd(epoch, &mut |op| {
      offset += op._store_grad(offset, grad_writer);
    });
    offset
  }

  fn next_iteration(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._next_iteration());
  }

  fn load_batch(&mut self, samples: &[S]) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._load_batch(samples));
  }

  fn save_rng_state(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._save_rng_state());
  }

  fn restore_rng_state(&mut self) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._restore_rng_state());
  }

  fn init_param(&mut self, rng: &mut Xorshiftplus128Rng) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._init_param(rng));
  }

  fn update_nondiff_param(&mut self, iter_nr: usize) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._update_nondiff_param(iter_nr));
  }

  fn reset_grad(&mut self) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._reset_grad());
  }

  fn forward(&mut self, phase: OpPhase) {
    let epoch = self._next();
    self._traverse_fwd(epoch, &mut |op| op._forward(phase));
  }

  fn backward(&mut self) {
    let epoch = self._next();
    self._traverse_bwd(epoch, &mut |op| op._backward());
  }
}

pub trait DiffOperator<T> {
  type Output;
  type Rng: Rng;

  fn _output(&self, arm: usize) -> Self::Output;

  fn diff_param_sz(&self) -> usize { 0 }
  fn nondiff_param_sz(&self) -> usize { 0 }
  fn total_param_sz(&self) -> usize {
    self.diff_param_sz() + self.nondiff_param_sz()
  }

  fn save_rng_state(&mut self) {}
  fn restore_rng_state(&mut self) {}
  //fn store_rng_state(&mut self, _rng_state: &mut Write) {}
  //fn load_rng_state(&mut self, _rng_state: &mut Read) {}

  //fn reset_state(&mut self) {}

  fn init_param(&mut self, _rng: &mut Self::Rng) {}
  //fn reset_nondiff_param(&mut self) {}
  fn reset_grad(&mut self) {}

  fn store_param(&mut self, _param_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn store_nondiff_param(&mut self, _param_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn load_param(&mut self, _param_reader: &mut ReadBuffer<T>, _offset: usize) -> usize { 0 }
  fn load_nondiff_param(&mut self, _param_reader: &mut ReadBuffer<T>, _offset: usize) -> usize { 0 }
  fn update_param(&mut self, _alpha: f32, _beta: f32, _grad_reader: &mut ReadAccumulateBuffer<T>, _offset: usize) -> usize { 0 }

  fn update_nondiff_param(&mut self, _iter: usize) {}

  fn store_grad(&mut self, _grad_writer: &mut WriteBuffer<T>, _offset: usize) -> usize { 0 }
  fn accumulate_grad(&mut self, _alpha: f32, _beta: f32, _grad_accum: &mut AccumulateBuffer<T>, _offset: usize) -> usize { 0 }
  //fn grad_step(&mut self, _alpha: f32, _beta: f32) {}

  fn reset_loss(&mut self) {}
  fn store_loss(&mut self) -> f32 { 0.0 }
  fn _store_accuracy(&mut self) -> usize { 0 }

  fn apply_grad_reg(&mut self, _reg: Regularization) {}

  fn forward(&mut self, phase: OpPhase);
  //fn fwd_reg(&mut self, _reg: Regularization) {}
  fn backward(&mut self);
  //fn bwd_reg(&mut self, _reg: Regularization) {}
  //fn r_forward(&mut self) { unimplemented!(); }
  //fn r_backward(&mut self) { unimplemented!(); }
}

pub trait DiffOperatorIo<T, Buf>/*: DiffOperator*/ where T: Copy {
  fn store_diff_param(&mut self, _param_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn store_nondiff_param(&mut self, _param_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn load_diff_param(&mut self, _param_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn load_nondiff_param(&mut self, _param_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn update_diff_param(&mut self, _alpha: T, _beta: T, _delta_reader: &mut Buf, _offset: usize) -> usize { 0 }
  fn store_grad(&mut self, _grad_writer: &mut Buf, _offset: usize) -> usize { 0 }
  fn accumulate_grad(&mut self, _alpha: T, _beta: T, _grad_writer: &mut Buf, _offset: usize) -> usize { 0 }
}

pub trait DiffOperatorInput<T, S>: DiffOperator<T> {
  fn as_op(&self) -> &DiffOperator<T, Output=Self::Output, Rng=Self::Rng> { unimplemented!(); }
  fn load_data(&mut self, samples: &[S]);
}

pub trait DiffOperatorOutput<T, U>: DiffOperator<T> {
  fn get_output(&mut self) -> U;
}

pub trait CheckpointFormat {
}

pub trait DiffOperatorCheckpoint<Format> where Format: CheckpointFormat {
  fn decode(reader: &mut Read) -> Result<Self, ()> where Self: Sized;
  fn encode(&mut self, writer: &mut Write) -> Result<(), ()>;
}
<|endoftext|>


<commit_before>#![deny(missing_docs)]
#![doc(html_root_url = "http://arcnmx.github.io/pipe-rs/")]
#![cfg_attr(all(feature = "bench", test), feature(test))]

//! Synchronous in-memory pipe
//!
//! ## Example
//!
//! ```
//! use std::thread::spawn;
//! use std::io::{Read, Write};
//!
//! let (mut read, mut write) = pipe::pipe();
//!
//! let message = "Hello, world!";
//! spawn(move || write.write_all(message.as_bytes()).unwrap());
//!
//! let mut s = String::new();
//! read.read_to_string(&mut s).unwrap();
//!
//! assert_eq!(&s, message);
//! ```

#[cfg(feature="readwrite")]
extern crate readwrite;

use std::sync::mpsc::{sync_channel, SyncSender, Receiver};
use std::io::{self, Read, Write};
use std::cmp::min;

/// The `Read` end of a pipe (see `pipe()`)
pub struct PipeReader(Receiver<Vec<u8>>, Vec<u8>);

/// The `Write` end of a pipe (see `pipe()`)
#[derive(Clone)]
pub struct PipeWriter(SyncSender<Vec<u8>>);

/// Creates a synchronous memory pipe
pub fn pipe() -> (PipeReader, PipeWriter) {
    let (tx, rx) = sync_channel(0);

    (PipeReader(rx, Vec::new()), PipeWriter(tx))
}

/// Creates a pair of pipes for bidirectional communication.
#[cfg(feature="readwrite")]
pub fn socketpair() -> (readwrite::ReadWrite<PipeReader, PipeWriter>, readwrite::ReadWrite<PipeReader, PipeWriter>) {
    let (r1,w1) = pipe();
    let (r2,w2) = pipe();
    ((r1,w2).into(), (r2,w1).into())
}

impl PipeWriter {
    /// Extracts the inner `SyncSender` from the writer
    pub fn into_inner(self) -> SyncSender<Vec<u8>> {
        self.0
    }
}

impl PipeReader {
    /// Extracts the inner `Receiver` from the writer, and any pending buffered data
    pub fn into_inner(self) -> (Receiver<Vec<u8>>, Vec<u8>) {
        (self.0, self.1)
    }
}

impl Read for PipeReader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        while self.1.is_empty() {
            match self.0.recv() {
                // The only existing error from mpsc is EOF
                Err(_) => return Ok(0),
                Ok(data) => self.1 = data,
            }
        }

        let len = min(buf.len(), self.1.len());
        buf[..len].copy_from_slice(&self.1[..len]);
        self.1.drain(..len);
        Ok(len)
    }
}

impl Write for PipeWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let data = buf.to_vec();

        self.0.send(data)
            .map(|_| buf.len())
            .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "pipe reader has been dropped"))
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::thread::spawn;
    use std::io::{Read, Write};
    use super::*;

    #[test]
    fn pipe_reader() {
        let i = b"hello there";
        let mut o = Vec::with_capacity(i.len());
        let (mut r, mut w) = pipe();
        let guard = spawn(move || {
            w.write_all(&i[..5]).unwrap();
            w.write_all(&i[5..]).unwrap();
            drop(w);
        });

        r.read_to_end(&mut o).unwrap();
        assert_eq!(i, &o[..]);

        guard.join().unwrap();
    }

    #[test]
    fn pipe_writer_fail() {
        let i = b"hi";
        let (r, mut w) = pipe();
        let guard = spawn(move || {
            drop(r);
        });

        assert!(w.write_all(i).is_err());

        guard.join().unwrap();
    }

    #[test]
    fn small_reads() {
        let block_cnt = 20;
        const BLOCK: usize = 20;
        let (mut r, mut w) = pipe();
        let guard = spawn(move || {
            for _ in 0..block_cnt {
                let data = &[0; BLOCK];
                w.write_all(data).unwrap();
            }
        });

        let mut buff = [0; BLOCK / 2];
        let mut read = 0;
        while let Ok(size) = r.read(&mut buff) {
            // 0 means EOF
            if size == 0 {
                break;
            }
            read += size;
        }
        assert_eq!(block_cnt * BLOCK, read);

        guard.join().unwrap();
    }
}

#[cfg(all(test, feature = "bench"))]
mod bench {
    extern crate test;
    use std::thread::spawn;
    use self::test::Bencher;
    use super::*;

    #[bench]
    fn pipe_bench(b: &mut Bencher) {
        use std::iter::repeat;
        use std::io::{copy, sink};

        let data = repeat(0).enumerate().map(|(i, _)| i as u8).take(0x1000).collect::<Vec<_>>();
        b.bytes = data.len() as u64 * 0x100;
        b.iter(|| {
            let (mut r, mut w) = pipe();
            let data = data.clone();
            let guard = spawn(move || {
                for _ in 0..0x100 {
                    w.write_all(&data[..]).unwrap();
                }
            });

            copy(&mut r, &mut sink()).unwrap();

            guard.join().unwrap();
        });
    }
}
<commit_msg>Rename socketpair to bipipe<commit_after>#![deny(missing_docs)]
#![doc(html_root_url = "http://arcnmx.github.io/pipe-rs/")]
#![cfg_attr(all(feature = "bench", test), feature(test))]

//! Synchronous in-memory pipe
//!
//! ## Example
//!
//! ```
//! use std::thread::spawn;
//! use std::io::{Read, Write};
//!
//! let (mut read, mut write) = pipe::pipe();
//!
//! let message = "Hello, world!";
//! spawn(move || write.write_all(message.as_bytes()).unwrap());
//!
//! let mut s = String::new();
//! read.read_to_string(&mut s).unwrap();
//!
//! assert_eq!(&s, message);
//! ```

#[cfg(feature="readwrite")]
extern crate readwrite;

use std::sync::mpsc::{sync_channel, SyncSender, Receiver};
use std::io::{self, Read, Write};
use std::cmp::min;

/// The `Read` end of a pipe (see `pipe()`)
pub struct PipeReader(Receiver<Vec<u8>>, Vec<u8>);

/// The `Write` end of a pipe (see `pipe()`)
#[derive(Clone)]
pub struct PipeWriter(SyncSender<Vec<u8>>);

/// Creates a synchronous memory pipe
pub fn pipe() -> (PipeReader, PipeWriter) {
    let (tx, rx) = sync_channel(0);

    (PipeReader(rx, Vec::new()), PipeWriter(tx))
}

/// Creates a pair of pipes for bidirectional communication, a bit like UNIX's `socketpair(2)`.
#[cfg(feature="readwrite")]
pub fn bipipe() -> (readwrite::ReadWrite<PipeReader, PipeWriter>, readwrite::ReadWrite<PipeReader, PipeWriter>) {
    let (r1,w1) = pipe();
    let (r2,w2) = pipe();
    ((r1,w2).into(), (r2,w1).into())
}

impl PipeWriter {
    /// Extracts the inner `SyncSender` from the writer
    pub fn into_inner(self) -> SyncSender<Vec<u8>> {
        self.0
    }
}

impl PipeReader {
    /// Extracts the inner `Receiver` from the writer, and any pending buffered data
    pub fn into_inner(self) -> (Receiver<Vec<u8>>, Vec<u8>) {
        (self.0, self.1)
    }
}

impl Read for PipeReader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        while self.1.is_empty() {
            match self.0.recv() {
                // The only existing error from mpsc is EOF
                Err(_) => return Ok(0),
                Ok(data) => self.1 = data,
            }
        }

        let len = min(buf.len(), self.1.len());
        buf[..len].copy_from_slice(&self.1[..len]);
        self.1.drain(..len);
        Ok(len)
    }
}

impl Write for PipeWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let data = buf.to_vec();

        self.0.send(data)
            .map(|_| buf.len())
            .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "pipe reader has been dropped"))
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::thread::spawn;
    use std::io::{Read, Write};
    use super::*;

    #[test]
    fn pipe_reader() {
        let i = b"hello there";
        let mut o = Vec::with_capacity(i.len());
        let (mut r, mut w) = pipe();
        let guard = spawn(move || {
            w.write_all(&i[..5]).unwrap();
            w.write_all(&i[5..]).unwrap();
            drop(w);
        });

        r.read_to_end(&mut o).unwrap();
        assert_eq!(i, &o[..]);

        guard.join().unwrap();
    }

    #[test]
    fn pipe_writer_fail() {
        let i = b"hi";
        let (r, mut w) = pipe();
        let guard = spawn(move || {
            drop(r);
        });

        assert!(w.write_all(i).is_err());

        guard.join().unwrap();
    }

    #[test]
    fn small_reads() {
        let block_cnt = 20;
        const BLOCK: usize = 20;
        let (mut r, mut w) = pipe();
        let guard = spawn(move || {
            for _ in 0..block_cnt {
                let data = &[0; BLOCK];
                w.write_all(data).unwrap();
            }
        });

        let mut buff = [0; BLOCK / 2];
        let mut read = 0;
        while let Ok(size) = r.read(&mut buff) {
            // 0 means EOF
            if size == 0 {
                break;
            }
            read += size;
        }
        assert_eq!(block_cnt * BLOCK, read);

        guard.join().unwrap();
    }
}

#[cfg(all(test, feature = "bench"))]
mod bench {
    extern crate test;
    use std::thread::spawn;
    use self::test::Bencher;
    use super::*;

    #[bench]
    fn pipe_bench(b: &mut Bencher) {
        use std::iter::repeat;
        use std::io::{copy, sink};

        let data = repeat(0).enumerate().map(|(i, _)| i as u8).take(0x1000).collect::<Vec<_>>();
        b.bytes = data.len() as u64 * 0x100;
        b.iter(|| {
            let (mut r, mut w) = pipe();
            let data = data.clone();
            let guard = spawn(move || {
                for _ in 0..0x100 {
                    w.write_all(&data[..]).unwrap();
                }
            });

            copy(&mut r, &mut sink()).unwrap();

            guard.join().unwrap();
        });
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate byteorder;
extern crate rand;

use crypto::digest::Digest;
use crypto::sha2::Sha256;
use crypto::hmac::Hmac;
use crypto::mac::{Mac};
use crypto::aes::KeySize;
use crypto::aes_gcm::AesGcm;
use crypto::aead::{AeadEncryptor, AeadDecryptor};
use crypto::curve25519::{curve25519, curve25519_base};

use byteorder::{ByteOrder, BigEndian};
use rand::{OsRng, Rng};

const DHLEN : usize = 32;
const HASHLEN : usize = 32;
//const BLOCKLEN : usize = 64;
const MACLEN : usize = 16;
const CIPHERKEYLEN : usize = 32;

fn copy_memory(data: &[u8], out: &mut [u8]) {
    for count in 0..data.len() {out[count] = data[count];}
}

pub struct KeyPair {
    privkey: [u8; DHLEN],
    pubkey: [u8; DHLEN]
}

pub fn generate_keypair() -> KeyPair {
    let mut privkey = [0u8; 32];
    let mut rng = OsRng::new().unwrap();
    rng.fill_bytes(&mut privkey);
    privkey[0] &= 248;
    privkey[31] &= 127;
    privkey[31] |= 64;
    let pubkey = curve25519_base(&privkey);
    KeyPair{privkey : privkey, pubkey : pubkey}
}

fn dh(keypair : &KeyPair, pubkey: &[u8]) -> [u8; DHLEN] {
    curve25519(&keypair.privkey, pubkey)
}

fn hash(data: &[u8], out: &mut [u8]) {
    let mut h = Sha256::new();
    h.input(data);
    h.result(out);
}

fn hmac_hash(key: &[u8], data: &[u8], out: &mut [u8]) {
    let mut hmac = Hmac::new(Sha256::new(), key);
    hmac.input(data);
    hmac.raw_result(out);
}

fn hkdf(chaining_key: &[u8], data: &[u8], out: &mut [u8]) {
    let mut temp_key = [0u8; HASHLEN];
    let mut in2 = [0u8; HASHLEN];
    hmac_hash(chaining_key, data, &mut temp_key);
    hmac_hash(&temp_key, &[1u8], &mut out[0..HASHLEN]);
    copy_memory(out as &[u8], &mut in2);
    in2[HASHLEN] = 2;
    hmac_hash(&temp_key, &in2, &mut out[HASHLEN..2*HASHLEN]);
}

fn encrypt(key: &[u8], nonce: u64, authtext: &[u8], plaintext: &[u8], out: &mut[u8]) {
    let mut nonce_bytes = [0u8; 12];
    BigEndian::write_u64(&mut nonce_bytes[4..], nonce);
    let mut cipher = AesGcm::new(KeySize::KeySize256, key, &nonce_bytes, authtext);
    let mut tag = [0u8; MACLEN];
    cipher.encrypt(plaintext, &mut out[..plaintext.len()], &mut tag);
    copy_memory(&tag, &mut out[plaintext.len()..]);
}

fn decrypt(key: &[u8], nonce: u64, authtext: &[u8], ciphertext: &[u8], out: &mut[u8]) -> bool {
    let mut nonce_bytes = [0u8; 12];
    BigEndian::write_u64(&mut nonce_bytes[4..], nonce);
    let mut cipher = AesGcm::new(KeySize::KeySize256, key, &nonce_bytes, authtext);
    let tag = [0u8; MACLEN];
    cipher.decrypt(ciphertext, out, &tag)
}

pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss}

pub enum NoiseError {DecryptError}

pub struct CipherState {
    k: [u8; CIPHERKEYLEN],
    n: u64
}

struct SymmetricState {
    cipherstate : CipherState,
    has_key : bool,
    ck: [u8; HASHLEN],
    h: [u8; HASHLEN]
}

pub struct HandshakeState {
    symmetricstate: SymmetricState,
    s: KeyPair,
    e: KeyPair,
    rs: [u8; DHLEN],
    re: [u8; DHLEN],
}

impl CipherState {

    pub fn encrypt_and_increment(&mut self, authtext: &[u8], plaintext: &[u8], out: &mut [u8]) {
        encrypt(&self.k, self.n, authtext, plaintext, out);
        self.n += 1;
    }

    pub fn decrypt_and_increment(&mut self, authtext: &[u8], ciphertext: &[u8], out: &mut [u8]) -> bool {
        if !decrypt(&self.k, self.n, authtext, ciphertext, out) {
            self.n += 1;
            return false;
        }
        self.n += 1;
        true
    }
}

impl SymmetricState {

    fn new(handshake_name: &[u8]) -> SymmetricState {
        let mut hname = [0u8; HASHLEN];
        match handshake_name.len() {
            0 ... HASHLEN => copy_memory(handshake_name, &mut hname),
            _ => hash(handshake_name, &mut hname)
        }
        let cipherstate = CipherState{k : [0; CIPHERKEYLEN], n:0};
        SymmetricState{cipherstate: cipherstate, has_key : false, ck : hname, h: hname}
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = [0u8; 2*HASHLEN];
        hkdf(&self.ck, data, &mut hkdf_output);
        copy_memory(&hkdf_output[0..HASHLEN], &mut self.ck);
        copy_memory(&hkdf_output[HASHLEN..HASHLEN+CIPHERKEYLEN], &mut self.cipherstate.k);
        self.cipherstate.n = 0;
        self.has_key = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut buffer = [0u8; 65535];
        copy_memory(&data, &mut buffer);
        copy_memory(&self.h, &mut buffer[data.len()..data.len() + HASHLEN]);
        hash(&buffer, &mut self.h);
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_key {
            self.cipherstate.encrypt_and_increment(&self.h, plaintext, out);
            self.mix_hash(&out[..plaintext.len() + MACLEN]);
            return plaintext.len() + MACLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_key {
            if !self.cipherstate.decrypt_and_increment(&self.h, data, out) {
                return false;
            }
            self.mix_hash(data);
        } else {
            copy_memory(data, out);
            self.mix_hash(data);
        }
        true
    }

    fn split(&self) -> (CipherState, CipherState) {
        let mut hkdf_output = [0u8; 2*HASHLEN];
        hkdf(&self.ck, &[0u8; 0], &mut hkdf_output);
        let mut c1 = CipherState{k : [0; CIPHERKEYLEN], n:0};
        let mut c2 = CipherState{k : [0; CIPHERKEYLEN], n:0};
        copy_memory( &hkdf_output[0..CIPHERKEYLEN], &mut c1.k);
        copy_memory( &hkdf_output[HASHLEN..HASHLEN+CIPHERKEYLEN], &mut c2.k);
        (c1, c2)
    }

}

impl HandshakeState {

    pub fn new(handshake_name: &[u8], new_s : KeyPair, new_e : KeyPair, new_rs: [u8; 32], new_re : [u8; 32] ) -> HandshakeState {
        let symmetricstate = SymmetricState::new(handshake_name);
        HandshakeState{symmetricstate: symmetricstate, s: new_s, e: new_e, rs: new_rs, re: new_re}
    }

    pub fn write_handshake_message(&mut self,
                               descriptor: &[Token],
                               last: bool,
                               payload: &[u8],
                               out: &mut [u8]) -> Option<(CipherState, CipherState)> {
        let mut index = 0;
        for token in descriptor {
            match *token {
                Token::E => {
                    self.e = generate_keypair();
                    index += self.symmetricstate.encrypt_and_hash(&self.e.pubkey, &mut out[index..]);
                },
                Token::S => index += self.symmetricstate.encrypt_and_hash(&self.s.pubkey, &mut out[index..]),
                Token::Dhee => self.symmetricstate.mix_key(&dh(&self.e, &self.re)),
                Token::Dhes => self.symmetricstate.mix_key(&dh(&self.e, &self.rs)),
                Token::Dhse => self.symmetricstate.mix_key(&dh(&self.s, &self.re)),
                Token::Dhss => self.symmetricstate.mix_key(&dh(&self.s, &self.rs)),
            }
        }
        self.symmetricstate.encrypt_and_hash(payload, &mut out[index..]);
        match last {
            true => Some(self.symmetricstate.split()),
            false => None
        }
    }

    pub fn read_handshake_message(&mut self,
                              buffer: &[u8],
                              descriptor: &[Token],
                              last: bool,
                              out: &mut [u8]) -> Result<Option<(CipherState, CipherState)>, NoiseError> {
        let mut ptr = buffer;
        for token in descriptor {
            match *token {
                Token::E | Token::S => {
                    let data = match self.symmetricstate.has_key {
                        true =>  {let temp = &ptr[..DHLEN+16]; ptr = &ptr[DHLEN+16..]; temp}
                        false => {let temp = &ptr[..DHLEN];    ptr = &ptr[DHLEN..];    temp}
                    };
                    if !self.symmetricstate.decrypt_and_hash(data, match *token {
                            Token::E => &mut self.re,
                            Token::S => &mut self.rs,
                            _ => unreachable!()}) {
                        return Err(NoiseError::DecryptError);
                    }
                },
                Token::Dhee => self.symmetricstate.mix_key(&dh(&self.e, &self.re)),
                Token::Dhes => self.symmetricstate.mix_key(&dh(&self.s, &self.re)),
                Token::Dhse => self.symmetricstate.mix_key(&dh(&self.e, &self.rs)),
                Token::Dhss => self.symmetricstate.mix_key(&dh(&self.s, &self.rs)),
            }
        }
        if !self.symmetricstate.decrypt_and_hash(ptr, out) {
            return Err(NoiseError::DecryptError);
        }
        match last {
            true => Ok(Some(self.symmetricstate.split())),
            false => Ok(None)
        }
    }

}

#[cfg(test)]
mod tests {

    extern crate rustc_serialize;

    use super::*;
    use super::hash;
    use super::hmac_hash;
    use super::dh;
    use super::copy_memory;
    use super::encrypt;
    use self::rustc_serialize::hex::{FromHex, ToHex};

    #[test]
    fn crypto_tests() {

        // SHA256 test
        {
            let mut output = [0u8; 32];
            hash("abc".as_bytes(), &mut output);
            assert!(output.to_hex() == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
        }

        // HMAC-SHA256 test - RFC 4231
        {
            let key = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".from_hex().unwrap();
            let data = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd".from_hex().unwrap();
            let mut output = [0u8; 32];
            hmac_hash(&key, &data, &mut output);
            println!("{}, {}, {}", key.len(), data.len(), output.to_hex());
            assert!(output.to_hex() == "773ea91e36800e46854db8ebd09181a72959098b3ef8c122d9635514ced565fe");
        }

        // Curve25519 test - draft-curves-10
        {
            let scalar = "a546e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449ac4".from_hex().unwrap();
            let mut keypair = KeyPair{privkey: [0; 32], pubkey: [0; 32]};
            copy_memory(&scalar, &mut keypair.privkey);
            let public = "e6db6867583030db3594c1a424b15f7c726624ec26b3353b10a903a6d0ab1c4c".from_hex().unwrap();
            let output = dh(&keypair, &public);
            assert!(output.to_hex() == "c3da55379de9c6908e94ea4df28d084f32eccf03491c71f754b4075577a28552");
        }

        //AES256-GCM tests - gcm-spec.pdf
        {
            // Test Case 13
            let key = [0u8; 32];
            let nonce = 0u64;
            let plaintext = [0u8; 0];
            let authtext = [0u8; 0];
            let mut ciphertext = [0u8; 16];
            encrypt(&key, nonce, &authtext, &plaintext, &mut ciphertext);
            assert!(ciphertext.to_hex() == "530f8afbc74536b9a963b4f1c4cb738b");

            // Test Case 14
            let plaintext2 = [0u8; 16];
            let mut ciphertext2 = [0u8; 32];
            encrypt(&key, nonce, &authtext, &plaintext2, &mut ciphertext2);
            assert!(ciphertext2.to_hex() == "cea7403d4d606b6e074ec5d3baf39d18d0d1c8a799996bf0265b98b5d48ab919");
        }

    }
}
<commit_msg>Fix hkdf() so it doesn't crash.<commit_after>extern crate crypto;
extern crate byteorder;
extern crate rand;

use crypto::digest::Digest;
use crypto::sha2::Sha256;
use crypto::hmac::Hmac;
use crypto::mac::{Mac};
use crypto::aes::KeySize;
use crypto::aes_gcm::AesGcm;
use crypto::aead::{AeadEncryptor, AeadDecryptor};
use crypto::curve25519::{curve25519, curve25519_base};

use byteorder::{ByteOrder, BigEndian};
use rand::{OsRng, Rng};

const DHLEN : usize = 32;
const HASHLEN : usize = 32;
//const BLOCKLEN : usize = 64;
const MACLEN : usize = 16;
const CIPHERKEYLEN : usize = 32;

fn copy_memory(data: &[u8], out: &mut [u8]) {
    for count in 0..data.len() {out[count] = data[count];}
}

pub struct KeyPair {
    privkey: [u8; DHLEN],
    pubkey: [u8; DHLEN]
}

pub fn generate_keypair() -> KeyPair {
    let mut privkey = [0u8; 32];
    let mut rng = OsRng::new().unwrap();
    rng.fill_bytes(&mut privkey);
    privkey[0] &= 248;
    privkey[31] &= 127;
    privkey[31] |= 64;
    let pubkey = curve25519_base(&privkey);
    KeyPair{privkey : privkey, pubkey : pubkey}
}

fn dh(keypair : &KeyPair, pubkey: &[u8]) -> [u8; DHLEN] {
    curve25519(&keypair.privkey, pubkey)
}

fn hash(data: &[u8], out: &mut [u8]) {
    let mut h = Sha256::new();
    h.input(data);
    h.result(out);
}

fn hmac_hash(key: &[u8], data: &[u8], out: &mut [u8]) {
    let mut hmac = Hmac::new(Sha256::new(), key);
    hmac.input(data);
    hmac.raw_result(out);
}

fn hkdf(chaining_key: &[u8], data: &[u8], out: &mut [u8]) {
    let mut temp_key = [0u8; HASHLEN];
    let mut in2 = [0u8; HASHLEN+1];
    hmac_hash(chaining_key, data, &mut temp_key);
    hmac_hash(&temp_key, &[1u8], &mut out[0..HASHLEN]);
    copy_memory(&out[0..HASHLEN], &mut in2);
    in2[HASHLEN] = 2;
    hmac_hash(&temp_key, &in2, &mut out[HASHLEN..2*HASHLEN]);
}

fn encrypt(key: &[u8], nonce: u64, authtext: &[u8], plaintext: &[u8], out: &mut[u8]) {
    let mut nonce_bytes = [0u8; 12];
    BigEndian::write_u64(&mut nonce_bytes[4..], nonce);
    let mut cipher = AesGcm::new(KeySize::KeySize256, key, &nonce_bytes, authtext);
    let mut tag = [0u8; MACLEN];
    cipher.encrypt(plaintext, &mut out[..plaintext.len()], &mut tag);
    copy_memory(&tag, &mut out[plaintext.len()..]);
}

fn decrypt(key: &[u8], nonce: u64, authtext: &[u8], ciphertext: &[u8], out: &mut[u8]) -> bool {
    let mut nonce_bytes = [0u8; 12];
    BigEndian::write_u64(&mut nonce_bytes[4..], nonce);
    let mut cipher = AesGcm::new(KeySize::KeySize256, key, &nonce_bytes, authtext);
    let tag = [0u8; MACLEN];
    cipher.decrypt(ciphertext, out, &tag)
}

pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss}

pub enum NoiseError {DecryptError}

pub struct CipherState {
    k: [u8; CIPHERKEYLEN],
    n: u64
}

struct SymmetricState {
    cipherstate : CipherState,
    has_key : bool,
    ck: [u8; HASHLEN],
    h: [u8; HASHLEN]
}

pub struct HandshakeState {
    symmetricstate: SymmetricState,
    s: KeyPair,
    e: KeyPair,
    rs: [u8; DHLEN],
    re: [u8; DHLEN],
}

impl CipherState {

    pub fn encrypt_and_increment(&mut self, authtext: &[u8], plaintext: &[u8], out: &mut [u8]) {
        encrypt(&self.k, self.n, authtext, plaintext, out);
        self.n += 1;
    }

    pub fn decrypt_and_increment(&mut self, authtext: &[u8], ciphertext: &[u8], out: &mut [u8]) -> bool {
        if !decrypt(&self.k, self.n, authtext, ciphertext, out) {
            self.n += 1;
            return false;
        }
        self.n += 1;
        true
    }
}

impl SymmetricState {

    fn new(handshake_name: &[u8]) -> SymmetricState {
        let mut hname = [0u8; HASHLEN];
        match handshake_name.len() {
            0 ... HASHLEN => copy_memory(handshake_name, &mut hname),
            _ => hash(handshake_name, &mut hname)
        }
        let cipherstate = CipherState{k : [0; CIPHERKEYLEN], n:0};
        SymmetricState{cipherstate: cipherstate, has_key : false, ck : hname, h: hname}
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = [0u8; 2*HASHLEN];
        hkdf(&self.ck, data, &mut hkdf_output);
        copy_memory(&hkdf_output[0..HASHLEN], &mut self.ck);
        copy_memory(&hkdf_output[HASHLEN..HASHLEN+CIPHERKEYLEN], &mut self.cipherstate.k);
        self.cipherstate.n = 0;
        self.has_key = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut buffer = [0u8; 65535];
        copy_memory(&data, &mut buffer);
        copy_memory(&self.h, &mut buffer[data.len()..data.len() + HASHLEN]);
        hash(&buffer, &mut self.h);
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_key {
            self.cipherstate.encrypt_and_increment(&self.h, plaintext, out);
            self.mix_hash(&out[..plaintext.len() + MACLEN]);
            return plaintext.len() + MACLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_key {
            if !self.cipherstate.decrypt_and_increment(&self.h, data, out) {
                return false;
            }
            self.mix_hash(data);
        } else {
            copy_memory(data, out);
            self.mix_hash(data);
        }
        true
    }

    fn split(&self) -> (CipherState, CipherState) {
        let mut hkdf_output = [0u8; 2*HASHLEN];
        hkdf(&self.ck, &[0u8; 0], &mut hkdf_output);
        let mut c1 = CipherState{k : [0; CIPHERKEYLEN], n:0};
        let mut c2 = CipherState{k : [0; CIPHERKEYLEN], n:0};
        copy_memory( &hkdf_output[0..CIPHERKEYLEN], &mut c1.k);
        copy_memory( &hkdf_output[HASHLEN..HASHLEN+CIPHERKEYLEN], &mut c2.k);
        (c1, c2)
    }

}

impl HandshakeState {

    pub fn new(handshake_name: &[u8], new_s : KeyPair, new_e : KeyPair, new_rs: [u8; 32], new_re : [u8; 32] ) -> HandshakeState {
        let symmetricstate = SymmetricState::new(handshake_name);
        HandshakeState{symmetricstate: symmetricstate, s: new_s, e: new_e, rs: new_rs, re: new_re}
    }

    pub fn write_handshake_message(&mut self,
                               descriptor: &[Token],
                               last: bool,
                               payload: &[u8],
                               out: &mut [u8]) -> Option<(CipherState, CipherState)> {
        let mut index = 0;
        for token in descriptor {
            match *token {
                Token::E => {
                    self.e = generate_keypair();
                    index += self.symmetricstate.encrypt_and_hash(&self.e.pubkey, &mut out[index..]);
                },
                Token::S => index += self.symmetricstate.encrypt_and_hash(&self.s.pubkey, &mut out[index..]),
                Token::Dhee => self.symmetricstate.mix_key(&dh(&self.e, &self.re)),
                Token::Dhes => self.symmetricstate.mix_key(&dh(&self.e, &self.rs)),
                Token::Dhse => self.symmetricstate.mix_key(&dh(&self.s, &self.re)),
                Token::Dhss => self.symmetricstate.mix_key(&dh(&self.s, &self.rs)),
            }
        }
        self.symmetricstate.encrypt_and_hash(payload, &mut out[index..]);
        match last {
            true => Some(self.symmetricstate.split()),
            false => None
        }
    }

    pub fn read_handshake_message(&mut self,
                              buffer: &[u8],
                              descriptor: &[Token],
                              last: bool,
                              out: &mut [u8]) -> Result<Option<(CipherState, CipherState)>, NoiseError> {
        let mut ptr = buffer;
        for token in descriptor {
            match *token {
                Token::E | Token::S => {
                    let data = match self.symmetricstate.has_key {
                        true =>  {let temp = &ptr[..DHLEN+16]; ptr = &ptr[DHLEN+16..]; temp}
                        false => {let temp = &ptr[..DHLEN];    ptr = &ptr[DHLEN..];    temp}
                    };
                    if !self.symmetricstate.decrypt_and_hash(data, match *token {
                            Token::E => &mut self.re,
                            Token::S => &mut self.rs,
                            _ => unreachable!()}) {
                        return Err(NoiseError::DecryptError);
                    }
                },
                Token::Dhee => self.symmetricstate.mix_key(&dh(&self.e, &self.re)),
                Token::Dhes => self.symmetricstate.mix_key(&dh(&self.s, &self.re)),
                Token::Dhse => self.symmetricstate.mix_key(&dh(&self.e, &self.rs)),
                Token::Dhss => self.symmetricstate.mix_key(&dh(&self.s, &self.rs)),
            }
        }
        if !self.symmetricstate.decrypt_and_hash(ptr, out) {
            return Err(NoiseError::DecryptError);
        }
        match last {
            true => Ok(Some(self.symmetricstate.split())),
            false => Ok(None)
        }
    }

}

#[cfg(test)]
mod tests {

    extern crate rustc_serialize;

    use super::*;
    use super::hash;
    use super::hmac_hash;
    use super::dh;
    use super::copy_memory;
    use super::encrypt;
    use super::hkdf;
    use self::rustc_serialize::hex::{FromHex, ToHex};

    #[test]
    fn crypto_tests() {

        // SHA256 test
        {
            let mut output = [0u8; 32];
            hash("abc".as_bytes(), &mut output);
            assert!(output.to_hex() == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
        }

        // HMAC-SHA256 test - RFC 4231
        {
            let key = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".from_hex().unwrap();
            let data = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd".from_hex().unwrap();
            let mut output = [0u8; 32];
            hmac_hash(&key, &data, &mut output);
            println!("{}, {}, {}", key.len(), data.len(), output.to_hex());
            assert!(output.to_hex() == "773ea91e36800e46854db8ebd09181a72959098b3ef8c122d9635514ced565fe");
        }

        // Curve25519 test - draft-curves-10
        {
            let scalar = "a546e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449ac4".from_hex().unwrap();
            let mut keypair = KeyPair{privkey: [0; 32], pubkey: [0; 32]};
            copy_memory(&scalar, &mut keypair.privkey);
            let public = "e6db6867583030db3594c1a424b15f7c726624ec26b3353b10a903a6d0ab1c4c".from_hex().unwrap();
            let output = dh(&keypair, &public);
            assert!(output.to_hex() == "c3da55379de9c6908e94ea4df28d084f32eccf03491c71f754b4075577a28552");
        }

        //AES256-GCM tests - gcm-spec.pdf
        {
            // Test Case 13
            let key = [0u8; 32];
            let nonce = 0u64;
            let plaintext = [0u8; 0];
            let authtext = [0u8; 0];
            let mut ciphertext = [0u8; 16];
            encrypt(&key, nonce, &authtext, &plaintext, &mut ciphertext);
            assert!(ciphertext.to_hex() == "530f8afbc74536b9a963b4f1c4cb738b");

            // Test Case 14
            let plaintext2 = [0u8; 16];
            let mut ciphertext2 = [0u8; 32];
            encrypt(&key, nonce, &authtext, &plaintext2, &mut ciphertext2);
            assert!(ciphertext2.to_hex() == "cea7403d4d606b6e074ec5d3baf39d18d0d1c8a799996bf0265b98b5d48ab919");
        }

    }
}
<|endoftext|>


<commit_before>//! Simple crate that parses athenahealth Watertown's Nourish menu, and posts it to a slack channel
//! of choice.

#![warn(missing_docs)]

#[macro_use]
extern crate error_chain;

#[macro_use]
extern crate lazy_static;

extern crate chrono;
extern crate inflector;
extern crate linked_hash_map;
extern crate regex;
extern crate reqwest;
extern crate select;

use std::cmp;
use std::fmt::{self, Display, Write};
use std::io::prelude::*;

use chrono::{Datelike, Duration, NaiveDate};
use inflector::Inflector;
use linked_hash_map::LinkedHashMap;
use regex::Regex;
use reqwest::Url;
use select::document::Document;
use select::predicate::{Attr, Name};

pub mod errors;

use errors::*;

lazy_static! {
    /// Extracts the ingredients out of the text of a menu item.
    static ref INGREDIENTS_RE: Regex =
        Regex::new(r"Ingredients: (.*?)(?: - Serving Size.*)?\s*$").unwrap();

    /// This Regex matches nutritional information, to filter it out from the menu output.
    static ref NUTRITION_RE: Regex = Regex::new(r"Cal.*Fat.*Sat.*Sod.*Carbs.*Fib.*Pro").unwrap();

    /// This Regex matches the price information, to filter it out from the menu output.
    static ref PRICE_RE: Regex = Regex::new(r"^\d+\.\d+( ?/ ?(\d+\.\d+|[a-z]+))?$").unwrap();
}

/// The Nourish menu for a given date.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Menu(LinkedHashMap<String, Entry>);

/// A section of the menu.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Entry {
    /// The section header.
    pub heading: String,

    /// The food items listed under the header.
    pub items: Vec<String>,

    /// Information indicating what ingredients are in the food, including whether it's vegetarian,
    /// vegan, etc.
    pub dietary_info: Option<String>,
}

impl Display for Entry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "*{}*\n", self.heading)?;

        let items = if self.items.is_empty() {
            vec![String::from("Nothing today")]
        } else {
            self.items.clone()
        };

        for item in &items {
            let dietary_info = if let Some(ref info) = self.dietary_info {
                format!(" (_{}_)", info)
            } else {
                String::default()
            };

            writeln!(f, "• {}{}", item, dietary_info)?;
        }

        Ok(())
    }
}

impl Menu {
    /// Returns the entries of the menu.
    pub fn entries(&self) -> Vec<Entry> {
        self.0.values().cloned().collect()
    }

    /// Renders the menu as a Markdown string.
    pub fn to_markdown(&self) -> Result<String> {
        let mut output = String::default();

        for entry in self.entries() {
            writeln!(output, "{}", entry).unwrap();
        }

        if output.is_empty() {
            bail!(ErrorKind::EmptyMenu);
        }

        Ok(output)
    }
}

/// Retrieves and parses the menu for a given date.
pub fn retrieve_menu(date: &NaiveDate) -> Result<Menu> {
    let url = url_for_date(date);

    let mut res = reqwest::get(url.as_str()).map_err(|e| {
        ErrorKind::Network(e.to_string())
    })?;

    let body = if res.status().is_success() {
        let mut bytes = vec![];
        res.read_to_end(&mut bytes)?;
        String::from_utf8_lossy(&bytes).into_owned()
    } else {
        bail!(ErrorKind::Network(res.status().to_string()))
    };

    Ok(parse_menu(&body))
}

/// Returns the URL pointing to the Nourish menu for a given day. On the weekends, returns Friday's
/// menu.
pub fn url_for_date(date: &NaiveDate) -> Url {
    let days_from_monday = cmp::min(date.weekday().num_days_from_monday(), 4);
    let monday = *date - Duration::days(date.weekday().num_days_from_monday() as i64);

    Url::parse(&format!(
        "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.nsf\
        /weeklyMenuLaunch/8DURSE~{}/$file/day{}.htm",
        monday.format("%m-%d-%Y"),
        days_from_monday + 1
    )).unwrap()
}

/// Parses the menu information out of HTML.
pub fn parse_menu(html: &str) -> Menu {
    let document = Document::from(html);
    let menu_node = document.find(Attr("id", "center_text")).next().unwrap();

    let mut menu = LinkedHashMap::new();
    let mut current_heading = None;

    for node in menu_node.children() {
        let text = node.text().trim().to_owned();

        if let Some("font-weight:bold;") = node.attr("style") {
            current_heading = Some(text.to_lowercase().to_title_case());
        } else {
            if let Some(ref heading) = current_heading {
                // Skip breakfast
                if heading == "Breakfast Special" {
                    continue;
                }

                // Skip nutrition and price information.
                if NUTRITION_RE.is_match(&text) || PRICE_RE.is_match(&text) {
                    continue;
                }

                let entry = menu.entry(heading.clone()).or_insert_with(|| {
                    Entry {
                        heading: heading.to_owned(),
                        items: Vec::new(),
                        dietary_info: None,
                    }
                });

                if node.is(Name("div")) && !text.is_empty() {
                    entry.items.push(text);
                } else if let Some(caps) = INGREDIENTS_RE.captures(&text) {
                    entry.dietary_info = Some(caps[1].to_owned());
                }
            } else {
                println!("encountered entry without a heading");
                continue;
            }
        }
    }

    Menu(menu)
}

#[cfg(test)]
mod tests {
    use chrono::NaiveDate;
    use linked_hash_map::LinkedHashMap;

    use super::{Entry, Menu};

    #[test]
    fn parse_menu() {
        let html = r#"

CHEF'S SPECIAL PIZZA

        "#;

        let menu = super::parse_menu(html);
        let expected = {
            let mut menu = LinkedHashMap::new();
            menu.insert(
                String::from("Chef's Special Pizza"),
                Entry {
                    heading: String::from("Chef's Special Pizza"),
                    items: vec![],
                    dietary_info: None,
                },
            );
            Menu(menu)
        };

        assert_eq!(menu, expected);
    }

    #[test]
    fn monday() {
        let expected_url = "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.\
                            nsf/weeklyMenuLaunch/8DURSE~04-18-2016/$file/day1.htm";

        assert_eq!(
            expected_url,
            &super::url_for_date(&NaiveDate::from_ymd(2016, 4, 18)).to_string()
        );
    }

    #[test]
    fn saturday() {
        let expected_url = "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.\
                            nsf/weeklyMenuLaunch/8DURSE~04-11-2016/$file/day5.htm";

        assert_eq!(
            expected_url,
            &super::url_for_date(&NaiveDate::from_ymd(2016, 4, 16)).to_string()
        );

    }

    #[test]
    fn ingredients_regex() {
        use super::INGREDIENTS_RE;

        let test_cases = [
            ("Ingredients: Pork, Egg, Spicy", "Pork, Egg, Spicy"),
            (
                "Ingredients: Vegetarian - Contains Wheat, Dairy",
                "Vegetarian - Contains Wheat, Dairy",
            ),
            ("Ingredients: Vegan - Serving Size 12oz", "Vegan"),
        ];

        for &(text, ingredients) in &test_cases {
            let capture = INGREDIENTS_RE.captures(text).unwrap().get(1).unwrap();
            assert_eq!(capture.as_str(), ingredients);
        }
    }

    #[test]
    fn price_regex() {
        use super::PRICE_RE;

        let patterns = ["7.50", "0.25/oz", "2.15 / 2.65"];

        for pattern in &patterns {
            assert!(PRICE_RE.is_match(pattern), pattern.clone());
        }
    }
}
<commit_msg>don't print ingredients list if it's too long<commit_after>//! Simple crate that parses athenahealth Watertown's Nourish menu, and posts it to a slack channel
//! of choice.

#![warn(missing_docs)]

#[macro_use]
extern crate error_chain;

#[macro_use]
extern crate lazy_static;

extern crate chrono;
extern crate inflector;
extern crate linked_hash_map;
extern crate regex;
extern crate reqwest;
extern crate select;

use std::cmp;
use std::fmt::{self, Display, Write};
use std::io::prelude::*;

use chrono::{Datelike, Duration, NaiveDate};
use inflector::Inflector;
use linked_hash_map::LinkedHashMap;
use regex::Regex;
use reqwest::Url;
use select::document::Document;
use select::predicate::{Attr, Name};

pub mod errors;

use errors::*;

lazy_static! {
    /// Extracts the ingredients out of the text of a menu item.
    static ref INGREDIENTS_RE: Regex =
        Regex::new(r"Ingredients: (.*?)(?: - Serving Size.*)?\s*$").unwrap();

    /// This Regex matches nutritional information, to filter it out from the menu output.
    static ref NUTRITION_RE: Regex = Regex::new(r"Cal.*Fat.*Sat.*Sod.*Carbs.*Fib.*Pro").unwrap();

    /// This Regex matches the price information, to filter it out from the menu output.
    static ref PRICE_RE: Regex = Regex::new(r"^\d+\.\d+( ?/ ?(\d+\.\d+|[a-z]+))?$").unwrap();
}

/// The Nourish menu for a given date.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Menu(LinkedHashMap<String, Entry>);

/// A section of the menu.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Entry {
    /// The section header.
    pub heading: String,

    /// The food items listed under the header.
    pub items: Vec<String>,

    /// Information indicating what ingredients are in the food, including whether it's vegetarian,
    /// vegan, etc.
    pub dietary_info: Option<String>,
}

impl Display for Entry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "*{}*\n", self.heading)?;

        let items = if self.items.is_empty() {
            vec![String::from("Nothing today")]
        } else {
            self.items.clone()
        };

        for item in &items {
            let dietary_info = if let Some(ref info) = self.dietary_info {
                format!(" (_{}_)", info)
            } else {
                String::default()
            };

            writeln!(f, "• {}{}", item, dietary_info)?;
        }

        Ok(())
    }
}

impl Menu {
    /// Returns the entries of the menu.
    pub fn entries(&self) -> Vec<Entry> {
        self.0.values().cloned().collect()
    }

    /// Renders the menu as a Markdown string.
    pub fn to_markdown(&self) -> Result<String> {
        let mut output = String::default();

        for entry in self.entries() {
            writeln!(output, "{}", entry).unwrap();
        }

        if output.is_empty() {
            bail!(ErrorKind::EmptyMenu);
        }

        Ok(output)
    }
}

/// Retrieves and parses the menu for a given date.
pub fn retrieve_menu(date: &NaiveDate) -> Result<Menu> {
    let url = url_for_date(date);

    let mut res = reqwest::get(url.as_str()).map_err(|e| {
        ErrorKind::Network(e.to_string())
    })?;

    let body = if res.status().is_success() {
        let mut bytes = vec![];
        res.read_to_end(&mut bytes)?;
        String::from_utf8_lossy(&bytes).into_owned()
    } else {
        bail!(ErrorKind::Network(res.status().to_string()))
    };

    Ok(parse_menu(&body))
}

/// Returns the URL pointing to the Nourish menu for a given day. On the weekends, returns Friday's
/// menu.
pub fn url_for_date(date: &NaiveDate) -> Url {
    let days_from_monday = cmp::min(date.weekday().num_days_from_monday(), 4);
    let monday = *date - Duration::days(date.weekday().num_days_from_monday() as i64);

    Url::parse(&format!(
        "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.nsf\
        /weeklyMenuLaunch/8DURSE~{}/$file/day{}.htm",
        monday.format("%m-%d-%Y"),
        days_from_monday + 1
    )).unwrap()
}

/// Parses the menu information out of HTML.
pub fn parse_menu(html: &str) -> Menu {
    let document = Document::from(html);
    let menu_node = document.find(Attr("id", "center_text")).next().unwrap();

    let mut menu = LinkedHashMap::new();
    let mut current_heading = None;

    for node in menu_node.children() {
        let text = node.text().trim().to_owned();

        if let Some("font-weight:bold;") = node.attr("style") {
            current_heading = Some(text.to_lowercase().to_title_case());
        } else {
            if let Some(ref heading) = current_heading {
                // Skip breakfast
                if heading == "Breakfast Special" {
                    continue;
                }

                // Skip nutrition and price information.
                if NUTRITION_RE.is_match(&text) || PRICE_RE.is_match(&text) {
                    continue;
                }

                let entry = menu.entry(heading.clone()).or_insert_with(|| {
                    Entry {
                        heading: heading.to_owned(),
                        items: Vec::new(),
                        dietary_info: None,
                    }
                });

                if let Some(ref caps) = INGREDIENTS_RE.captures(&text) {
                    let ingredients = &caps[1];
                    if ingredients.chars().count() < 50 {
                        entry.dietary_info = Some(caps[1].to_owned());
                    }
                } else if node.is(Name("div")) && !text.is_empty() {
                    entry.items.push(text.to_owned());
                }
            } else {
                println!("encountered entry without a heading");
                continue;
            }
        }
    }

    Menu(menu)
}

#[cfg(test)]
mod tests {
    use chrono::NaiveDate;
    use linked_hash_map::LinkedHashMap;

    use super::{Entry, Menu};

    #[test]
    fn parse_menu() {
        let html = r#"

CHEF'S SPECIAL PIZZA

        "#;

        let menu = super::parse_menu(html);
        let expected = {
            let mut menu = LinkedHashMap::new();
            menu.insert(
                String::from("Chef's Special Pizza"),
                Entry {
                    heading: String::from("Chef's Special Pizza"),
                    items: vec![],
                    dietary_info: None,
                },
            );
            Menu(menu)
        };

        assert_eq!(menu, expected);
    }

    #[test]
    fn monday() {
        let expected_url = "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.\
                            nsf/weeklyMenuLaunch/8DURSE~04-18-2016/$file/day1.htm";

        assert_eq!(
            expected_url,
            &super::url_for_date(&NaiveDate::from_ymd(2016, 4, 18)).to_string()
        );
    }

    #[test]
    fn saturday() {
        let expected_url = "http://dining.guckenheimer.com/clients/athenahealth/fss/fss.\
                            nsf/weeklyMenuLaunch/8DURSE~04-11-2016/$file/day5.htm";

        assert_eq!(
            expected_url,
            &super::url_for_date(&NaiveDate::from_ymd(2016, 4, 16)).to_string()
        );

    }

    #[test]
    fn ingredients_regex() {
        use super::INGREDIENTS_RE;

        let test_cases = [
            ("Ingredients: Pork, Egg, Spicy", "Pork, Egg, Spicy"),
            (
                "Ingredients: Vegetarian - Contains Wheat, Dairy",
                "Vegetarian - Contains Wheat, Dairy",
            ),
            ("Ingredients: Vegan - Serving Size 12oz", "Vegan"),
        ];

        for &(text, ingredients) in &test_cases {
            let capture = INGREDIENTS_RE.captures(text).unwrap().get(1).unwrap();
            assert_eq!(capture.as_str(), ingredients);
        }
    }

    #[test]
    fn price_regex() {
        use super::PRICE_RE;

        let patterns = ["7.50", "0.25/oz", "2.15 / 2.65"];

        for pattern in &patterns {
            assert!(PRICE_RE.is_match(pattern), pattern.clone());
        }
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/*!
*resvg* is an [SVG](https://en.wikipedia.org/wiki/Scalable_Vector_Graphics) rendering library.

*resvg* can be used to render SVG files based on a
[static](http://www.w3.org/TR/SVG11/feature#SVG-static)
[SVG Full 1.1](https://www.w3.org/TR/SVG/Overview.html) subset.
In simple terms: no animations and scripting.

It can be used as a simple SVG to PNG converted.
And as an embeddable library to paint SVG on an application native canvas.
*/

#![forbid(unsafe_code)]
#![warn(missing_docs)]
//#![warn(missing_debug_implementations)]
#![warn(missing_copy_implementations)]

// For error-chain.
#![recursion_limit="128"]

extern crate svgdom;
extern crate base64;
extern crate libflate;
extern crate lyon_geom;
extern crate euclid;
extern crate ego_tree;
#[macro_use] extern crate log;
#[macro_use] extern crate error_chain;

#[cfg(feature = "cairo-backend")] pub extern crate cairo;
#[cfg(feature = "cairo-backend")] extern crate pango;
#[cfg(feature = "cairo-backend")] extern crate pangocairo;
#[cfg(feature = "cairo-backend")] extern crate image;

#[cfg(feature = "qt-backend")] pub extern crate resvg_qt as qt;

#[cfg(feature = "cairo-backend")] pub mod render_cairo;
#[cfg(feature = "qt-backend")] pub mod render_qt;

pub mod tree;
pub mod utils;
mod math;
mod convert;
mod error;
mod options;
mod preproc;
mod traits;

use std::path::{
    Path,
};

pub use error::{
    Error,
    ErrorKind,
    Result,
};
// reexport traits
pub use tree::{
    NodeExt,
};
pub use options::*;
pub use math::*;

/// Shorthand names for modules.
mod short {
    pub use svgdom::{
        LengthUnit as Unit,
        ElementId as EId,
        AttributeId as AId,
        AttributeValue as AValue,
    };
}

use preproc::{
    DEFAULT_FONT_FAMILY,
    DEFAULT_FONT_SIZE,
};

/// A generic interface for image rendering.
///
/// Instead of using backend implementation directly, you can
/// use this trait to write backend-independent code.
pub trait Render {
    /// Renders SVG to image.
    fn render_to_image(
        &self,
        rtree: &tree::RenderTree,
        opt: &Options,
    ) -> Result<Box<OutputImage>>;

    /// Renders SVG node to image.
    fn render_node_to_image(
        &self,
        node: tree::NodeRef,
        opt: &Options,
    ) -> Result<Box<OutputImage>>;

    /// Calculates node's bounding box.
    ///
    /// Note: this method can be pretty expensive.
    fn calc_node_bbox(
        &self,
        node: tree::NodeRef,
        opt: &Options,
    ) -> Option<Rect>;
}

/// A generic interface for output image.
pub trait OutputImage {
    /// Saves rendered image to the selected path.
    fn save(&self, path: &Path) -> bool;
}

/// Global library handle.
pub struct InitObject {
    #[cfg(feature = "qt-backend")]
    #[allow(dead_code)]
    handle: qt::GuiApp,
}

/// Creates a global library handle.
///
/// Must be invoked before any other `resvg` code.
///
/// Currently, handles `QGuiApplication` object which must be created
/// in order to draw text. If you don't plan to draw text - it's better to skip
/// the initialization.
///
/// Does nothing when only `cairo` backend is enabled.
///
/// Note: `QGuiApplication` initialization is pretty slow (up to 100ms).
///
/// # Example
///
/// ```
/// let _resvg = resvg::init();
///
/// // other code
/// ```
///
/// Also, take a look at `examples/minimal.rs`.
pub fn init() -> InitObject {
    InitObject {
        #[cfg(feature = "qt-backend")]
        handle: qt::GuiApp::new("resvg"),
    }
}

/// Returns default backend.
///
/// - If both backends are enabled - cairo backend will be returned.
/// - If no backends are enabled - will panic.
/// - Otherwise will return a corresponding backend.
#[allow(unreachable_code)]
pub fn default_backend() -> Box<Render> {
    #[cfg(feature = "cairo-backend")]
    {
        return Box::new(render_cairo::Backend);
    }

    #[cfg(feature = "qt-backend")]
    {
        return Box::new(render_qt::Backend);
    }

    unreachable!("at least one backend must be enabled")
}

/// Creates `RenderTree` from SVG data.
pub fn parse_rtree_from_data(
    text: &str,
    opt: &Options,
) -> Result<tree::RenderTree> {
    let mut doc = parse_svg(text)?;
    preproc::prepare_doc(&mut doc, opt)?;
    let rtree = convert::convert_doc(&doc, opt)?;

    Ok(rtree)
}

/// Creates `RenderTree` from file.
///
/// `.svg` and `.svgz` files are supported.
pub fn parse_rtree_from_file<P: AsRef<Path>>(
    path: P,
    opt: &Options,
) -> Result<tree::RenderTree> {
    let text = load_file(path.as_ref())?;
    parse_rtree_from_data(&text, opt)
}

fn load_file(path: &Path) -> Result<String> {
    use std::fs;
    use std::io::Read;

    let mut file = fs::File::open(path)?;
    let length = file.metadata()?.len() as usize;

    let ext = if let Some(ext) = Path::new(path).extension() {
        ext.to_str().map(|s| s.to_lowercase()).unwrap_or(String::new())
    } else {
        String::new()
    };

    match ext.as_str() {
        "svgz" => {
            let mut decoder = libflate::gzip::Decoder::new(&file)?;
            let mut decoded = Vec::new();
            decoder.read_to_end(&mut decoded)?;

            Ok(String::from_utf8(decoded)?)
        }
        "svg" => {
            let mut s = String::with_capacity(length + 1);
            file.read_to_string(&mut s)?;
            Ok(s)
        }
        _ => {
            Err(ErrorKind::InvalidFileExtension.into())
        }
    }
}

fn parse_svg(text: &str) -> Result<svgdom::Document> {
    let opt = svgdom::ParseOptions {
        parse_comments: false,
        parse_declarations: false,
        parse_unknown_elements: false,
        parse_unknown_attributes: false,
        parse_px_unit: false,
        skip_invalid_attributes: true,
        skip_invalid_css: true,
        skip_paint_fallback: true,
        .. svgdom::ParseOptions::default()
    };

    let doc = svgdom::Document::from_str_with_opt(&text, &opt)?;
    Ok(doc)
}
<commit_msg>Added parse_rtree_from_dom() method.<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/*!
*resvg* is an [SVG](https://en.wikipedia.org/wiki/Scalable_Vector_Graphics) rendering library.

*resvg* can be used to render SVG files based on a
[static](http://www.w3.org/TR/SVG11/feature#SVG-static)
[SVG Full 1.1](https://www.w3.org/TR/SVG/Overview.html) subset.
In simple terms: no animations and scripting.

It can be used as a simple SVG to PNG converted.
And as an embeddable library to paint SVG on an application native canvas.
*/

#![forbid(unsafe_code)]
#![warn(missing_docs)]
//#![warn(missing_debug_implementations)]
#![warn(missing_copy_implementations)]

// For error-chain.
#![recursion_limit="128"]

extern crate svgdom;
extern crate base64;
extern crate libflate;
extern crate lyon_geom;
extern crate euclid;
extern crate ego_tree;
#[macro_use] extern crate log;
#[macro_use] extern crate error_chain;

#[cfg(feature = "cairo-backend")] pub extern crate cairo;
#[cfg(feature = "cairo-backend")] extern crate pango;
#[cfg(feature = "cairo-backend")] extern crate pangocairo;
#[cfg(feature = "cairo-backend")] extern crate image;

#[cfg(feature = "qt-backend")] pub extern crate resvg_qt as qt;

#[cfg(feature = "cairo-backend")] pub mod render_cairo;
#[cfg(feature = "qt-backend")] pub mod render_qt;

pub mod tree;
pub mod utils;
mod math;
mod convert;
mod error;
mod options;
mod preproc;
mod traits;

use std::path::{
    Path,
};

pub use error::{
    Error,
    ErrorKind,
    Result,
};
// reexport traits
pub use tree::{
    NodeExt,
};
pub use options::*;
pub use math::*;

/// Shorthand names for modules.
mod short {
    pub use svgdom::{
        LengthUnit as Unit,
        ElementId as EId,
        AttributeId as AId,
        AttributeValue as AValue,
    };
}

use preproc::{
    DEFAULT_FONT_FAMILY,
    DEFAULT_FONT_SIZE,
};

/// A generic interface for image rendering.
///
/// Instead of using backend implementation directly, you can
/// use this trait to write backend-independent code.
pub trait Render {
    /// Renders SVG to image.
    fn render_to_image(
        &self,
        rtree: &tree::RenderTree,
        opt: &Options,
    ) -> Result<Box<OutputImage>>;

    /// Renders SVG node to image.
    fn render_node_to_image(
        &self,
        node: tree::NodeRef,
        opt: &Options,
    ) -> Result<Box<OutputImage>>;

    /// Calculates node's bounding box.
    ///
    /// Note: this method can be pretty expensive.
    fn calc_node_bbox(
        &self,
        node: tree::NodeRef,
        opt: &Options,
    ) -> Option<Rect>;
}

/// A generic interface for output image.
pub trait OutputImage {
    /// Saves rendered image to the selected path.
    fn save(&self, path: &Path) -> bool;
}

/// Global library handle.
pub struct InitObject {
    #[cfg(feature = "qt-backend")]
    #[allow(dead_code)]
    handle: qt::GuiApp,
}

/// Creates a global library handle.
///
/// Must be invoked before any other `resvg` code.
///
/// Currently, handles `QGuiApplication` object which must be created
/// in order to draw text. If you don't plan to draw text - it's better to skip
/// the initialization.
///
/// Does nothing when only `cairo` backend is enabled.
///
/// Note: `QGuiApplication` initialization is pretty slow (up to 100ms).
///
/// # Example
///
/// ```
/// let _resvg = resvg::init();
///
/// // other code
/// ```
///
/// Also, take a look at `examples/minimal.rs`.
pub fn init() -> InitObject {
    InitObject {
        #[cfg(feature = "qt-backend")]
        handle: qt::GuiApp::new("resvg"),
    }
}

/// Returns default backend.
///
/// - If both backends are enabled - cairo backend will be returned.
/// - If no backends are enabled - will panic.
/// - Otherwise will return a corresponding backend.
#[allow(unreachable_code)]
pub fn default_backend() -> Box<Render> {
    #[cfg(feature = "cairo-backend")]
    {
        return Box::new(render_cairo::Backend);
    }

    #[cfg(feature = "qt-backend")]
    {
        return Box::new(render_qt::Backend);
    }

    unreachable!("at least one backend must be enabled")
}

/// Creates `RenderTree` from SVG data.
pub fn parse_rtree_from_data(
    text: &str,
    opt: &Options,
) -> Result<tree::RenderTree> {
    let doc = parse_svg(text)?;
    parse_rtree_from_dom(doc, opt)
}

/// Creates `RenderTree` from file.
///
/// `.svg` and `.svgz` files are supported.
pub fn parse_rtree_from_file<P: AsRef<Path>>(
    path: P,
    opt: &Options,
) -> Result<tree::RenderTree> {
    let text = load_file(path.as_ref())?;
    parse_rtree_from_data(&text, opt)
}

/// Creates `RenderTree` from `svgdom::Document`.
pub fn parse_rtree_from_dom(
    mut doc: svgdom::Document,
    opt: &Options,
) -> Result<tree::RenderTree> {
    preproc::prepare_doc(&mut doc, opt)?;
    let rtree = convert::convert_doc(&doc, opt)?;

    Ok(rtree)
}

fn load_file(path: &Path) -> Result<String> {
    use std::fs;
    use std::io::Read;

    let mut file = fs::File::open(path)?;
    let length = file.metadata()?.len() as usize;

    let ext = if let Some(ext) = Path::new(path).extension() {
        ext.to_str().map(|s| s.to_lowercase()).unwrap_or(String::new())
    } else {
        String::new()
    };

    match ext.as_str() {
        "svgz" => {
            let mut decoder = libflate::gzip::Decoder::new(&file)?;
            let mut decoded = Vec::new();
            decoder.read_to_end(&mut decoded)?;

            Ok(String::from_utf8(decoded)?)
        }
        "svg" => {
            let mut s = String::with_capacity(length + 1);
            file.read_to_string(&mut s)?;
            Ok(s)
        }
        _ => {
            Err(ErrorKind::InvalidFileExtension.into())
        }
    }
}

fn parse_svg(text: &str) -> Result<svgdom::Document> {
    let opt = svgdom::ParseOptions {
        parse_comments: false,
        parse_declarations: false,
        parse_unknown_elements: false,
        parse_unknown_attributes: false,
        parse_px_unit: false,
        skip_invalid_attributes: true,
        skip_invalid_css: true,
        skip_paint_fallback: true,
        .. svgdom::ParseOptions::default()
    };

    let doc = svgdom::Document::from_str_with_opt(&text, &opt)?;
    Ok(doc)
}
<|endoftext|>