<commit_before>use std::io::{ Write, Result as IoResult };

use exec::{Step, Message};
use eval::Expr;
use data::DataMode;
use session::ValueID;

type StateId = u32;
type CounterId = u32;

#[derive(Clone, Debug)]
pub struct Nfa {
    states: Vec<State>,
    initial: StateId,
    success: StateId,
}

impl Nfa {
    pub fn new() -> Nfa {
        let mut nfa = Nfa { states: vec![], initial: 0, success: 0 };
        nfa.initial = nfa.add_state();
        nfa.success = nfa.add_state();
        nfa
    }

    pub fn add_state(&mut self) -> StateId {
        let id = self.states.len() as StateId;
        self.states.push(State::new());
        id
    }

    pub fn add_transition(&mut self, from: StateId, to: StateId, action: Action) {
        self.states[from as usize].transitions.push(Transition { target: to, action: action});
    }

    pub fn to_graphviz(&self, f: &mut Write) -> IoResult<()> {
        try!(writeln!(f, "digraph G {{"));
        for (id, state) in self.states.iter().enumerate() {
            for transition in &state.transitions {
                try!(writeln!(f, "{} -> {} [ label=\"{:?}\" ];", id, transition.target, transition.action));
            }
        }
        try!(writeln!(f, "}}"));
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct State {
    transitions: Vec<Transition>,
}

impl State {
    fn new() -> State {
        State { transitions: vec![] }
    }
}

#[derive(Clone, Debug)]
pub struct Transition {
    target: StateId,
    action: Action,
}

#[derive(Clone, Debug)]
pub enum Action {
    Epsilon,
    Lower(Message, Option<(StateId, StateId)>),
    Upper(Message, Option<(StateId, StateId)>),

    RepeatDnInit(CounterId, Expr), // Down-evaluate count, initialize counter to zero
    RepeatDnBack(CounterId), // Increment counter, guard on counter < count
    RepeatDnExit(CounterId), // Guard on counter == count

    RepeatUpInit(CounterId), // initialize counter to zero
    RepeatUpBack(CounterId), // increment counter
    RepeatUpExit(CounterId, Expr), // Guard on up-evaluate count

    // Down-evaluate outer variables, pop first inner variables down, set counter to zero
    ForInit(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
    // guard on counter < size, push next outer variables up, pop next inner variables down
    ForBack(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
    // guard on counter == size, Up-evaluate outer variables
    ForExit(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
}

pub fn from_step_tree(s: &Step) -> Nfa {
    let mut nfa = Nfa::new();
    let initial = nfa.initial;
    let success = nfa.success;
    from_step_tree_inner(s, &mut nfa, initial, success);

    fn from_step_tree_inner(s: &Step, nfa: &mut Nfa, from: StateId, to: StateId) {
        match *s {
            Step::Nop => nfa.add_transition(from, to, Action::Epsilon),
            Step::Token(ref message) => {
                let m: Message = message.clone();
                nfa.add_transition(from, to, Action::Lower(m, None))
            }
            Step::TokenTop(ref message, box ref body) => {
                let is = nfa.add_state();
                let ie = nfa.add_state();
                from_step_tree_inner(body, nfa, is, ie);
                nfa.add_transition(from, to, Action::Upper(message.clone(), Some((is, ie))))
            }
            Step::Seq(ref steps) => {
                if let Some((last, rest)) = steps.split_last() {
                    let mut conn = from;
                    for c in rest.iter() {
                        let next = nfa.add_state();
                        from_step_tree_inner(c, nfa, conn, next);
                        conn = next;
                    }
                    from_step_tree_inner(last, nfa, conn, to);
                } else {
                    nfa.add_transition(from, to, Action::Epsilon);
                }
            }
            Step::Repeat(ref count, box ref inner, up) => {
                let start = nfa.add_state();
                let end = nfa.add_state();
                let counter = start;

                if up {
                    nfa.add_transition(from, start, Action::RepeatUpInit(counter));
                    nfa.add_transition(end, start, Action::RepeatUpBack(counter));
                    nfa.add_transition(end, to, Action::RepeatUpExit(counter, count.clone()));
                } else {
                    nfa.add_transition(from, start, Action::RepeatDnInit(counter, count.clone()));
                    nfa.add_transition(end, start, Action::RepeatDnBack(counter));
                    nfa.add_transition(end, to, Action::RepeatDnExit(counter));
                }
                from_step_tree_inner(inner, nfa, start, end);
            }
            Step::Foreach(width, ref vars, box ref inner) => {
                let start = nfa.add_state();
                let end = nfa.add_state();
                let counter = start;

                nfa.add_transition(from, start, Action::ForInit(counter, width, vars.clone()));
                nfa.add_transition(end, start, Action::ForBack(counter, width, vars.clone()));
                nfa.add_transition(end, to, Action::ForExit(counter, width, vars.clone()));

                from_step_tree_inner(inner, nfa, start, end);
            }
        }
    }

    nfa
}
<commit_msg>nfa: Include the body of upper messages directly in the NFA<commit_after>use std::io::{ Write, Result as IoResult };

use exec::{Step, Message};
use eval::Expr;
use data::DataMode;
use session::ValueID;

type StateId = u32;
type CounterId = u32;

#[derive(Clone, Debug)]
pub struct Nfa {
    states: Vec<State>,
    initial: StateId,
    success: StateId,
}

impl Nfa {
    pub fn new() -> Nfa {
        let mut nfa = Nfa { states: vec![], initial: 0, success: 0 };
        nfa.initial = nfa.add_state();
        nfa.success = nfa.add_state();
        nfa
    }

    pub fn add_state(&mut self) -> StateId {
        let id = self.states.len() as StateId;
        self.states.push(State::new());
        id
    }

    pub fn add_transition(&mut self, from: StateId, to: StateId, action: Action) {
        self.states[from as usize].transitions.push(Transition { target: to, action: action});
    }

    pub fn to_graphviz(&self, f: &mut Write) -> IoResult<()> {
        try!(writeln!(f, "digraph G {{"));
        for (id, state) in self.states.iter().enumerate() {
            for transition in &state.transitions {
                try!(writeln!(f, "{} -> {} [ label=\"{:?}\" ];", id, transition.target, transition.action));
            }
        }
        try!(writeln!(f, "}}"));
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct State {
    transitions: Vec<Transition>,
}

impl State {
    fn new() -> State {
        State { transitions: vec![] }
    }
}

#[derive(Clone, Debug)]
pub struct Transition {
    target: StateId,
    action: Action,
}

#[derive(Clone, Debug)]
pub enum Action {
    Epsilon,
    Lower(Message),
    UpperBegin(Message),
    UpperEnd(Message),

    RepeatDnInit(CounterId, Expr), // Down-evaluate count, initialize counter to zero
    RepeatDnBack(CounterId), // Increment counter, guard on counter < count
    RepeatDnExit(CounterId), // Guard on counter == count

    RepeatUpInit(CounterId), // initialize counter to zero
    RepeatUpBack(CounterId), // increment counter
    RepeatUpExit(CounterId, Expr), // Guard on up-evaluate count

    // Down-evaluate outer variables, pop first inner variables down, set counter to zero
    ForInit(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
    // guard on counter < size, push next outer variables up, pop next inner variables down
    ForBack(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
    // guard on counter == size, Up-evaluate outer variables
    ForExit(CounterId, u32, Vec<(ValueID, Expr, DataMode)>),
}

pub fn from_step_tree(s: &Step) -> Nfa {
    let mut nfa = Nfa::new();
    let initial = nfa.initial;
    let success = nfa.success;
    from_step_tree_inner(s, &mut nfa, initial, success);

    fn from_step_tree_inner(s: &Step, nfa: &mut Nfa, from: StateId, to: StateId) {
        match *s {
            Step::Nop => nfa.add_transition(from, to, Action::Epsilon),
            Step::Token(ref message) => {
                let m: Message = message.clone();
                nfa.add_transition(from, to, Action::Lower(m))
            }
            Step::TokenTop(ref message, box ref body) => {
                let is = nfa.add_state();
                let ie = nfa.add_state();
                nfa.add_transition(from, is, Action::UpperBegin(message.clone()));
                from_step_tree_inner(body, nfa, is, ie);
                nfa.add_transition(ie, to, Action::UpperEnd(message.clone()));
            }
            Step::Seq(ref steps) => {
                if let Some((last, rest)) = steps.split_last() {
                    let mut conn = from;
                    for c in rest.iter() {
                        let next = nfa.add_state();
                        from_step_tree_inner(c, nfa, conn, next);
                        conn = next;
                    }
                    from_step_tree_inner(last, nfa, conn, to);
                } else {
                    nfa.add_transition(from, to, Action::Epsilon);
                }
            }
            Step::Repeat(ref count, box ref inner, up) => {
                let start = nfa.add_state();
                let end = nfa.add_state();
                let counter = start;

                if up {
                    nfa.add_transition(from, start, Action::RepeatUpInit(counter));
                    nfa.add_transition(end, start, Action::RepeatUpBack(counter));
                    nfa.add_transition(end, to, Action::RepeatUpExit(counter, count.clone()));
                } else {
                    nfa.add_transition(from, start, Action::RepeatDnInit(counter, count.clone()));
                    nfa.add_transition(end, start, Action::RepeatDnBack(counter));
                    nfa.add_transition(end, to, Action::RepeatDnExit(counter));
                }
                from_step_tree_inner(inner, nfa, start, end);
            }
            Step::Foreach(width, ref vars, box ref inner) => {
                let start = nfa.add_state();
                let end = nfa.add_state();
                let counter = start;

                nfa.add_transition(from, start, Action::ForInit(counter, width, vars.clone()));
                nfa.add_transition(end, start, Action::ForBack(counter, width, vars.clone()));
                nfa.add_transition(end, to, Action::ForExit(counter, width, vars.clone()));

                from_step_tree_inner(inner, nfa, start, end);
            }
        }
    }

    nfa
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::Result;

/// Header used by TGA image files
#[derive(Debug)]
struct Header {
    id_length: u8, // length of ID string
    map_type: u8, // color map type
    image_type: u8, // image type code
    map_origin: u16, // starting index of map
    map_length: u16, // length of map
    map_entry_size: u8, // size of map entries in bits
    x_origin: u16, // x-origin of image
    y_origin: u16, // y-origin of image
    image_width: u16, // width of image
    image_height: u16, // height of image
    pixel_depth: u8, // bits per pixel
    image_desc: u8, // image descriptor
}

#[derive(Clone,Copy)]
pub struct TgaColor {
    pub rgba: [u8; 4],
}

impl TgaColor {
    fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(w.write_u8(self.rgba[2]));
        try!(w.write_u8(self.rgba[1]));
        try!(w.write_u8(self.rgba[0]));
        try!(w.write_u8(self.rgba[3]));
        Ok(())
    }
}

const RAW_TRUE_COLOR: u8 = 2;
const ARGB_DEPTH: u8 = 32;

impl Header {
    fn new(w: u16, h: u16) -> Header {
        Header {
            id_length: 0,
            map_type: 0,
            image_type: RAW_TRUE_COLOR,
            map_origin: 0,
            map_length: 0,
            map_entry_size: 0,
            x_origin: 0,
            y_origin: 0,
            image_width: w,
            image_height: h,
            pixel_depth: ARGB_DEPTH,
            image_desc: 0,
        }
    }

    fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(w.write_u8(self.id_length));
        try!(w.write_u8(self.map_type));
        try!(w.write_u8(self.image_type));
        try!(w.write_u16::<LittleEndian>(self.map_origin));
        try!(w.write_u16::<LittleEndian>(self.map_length));
        try!(w.write_u8(self.map_entry_size));
        try!(w.write_u16::<LittleEndian>(self.x_origin));
        try!(w.write_u16::<LittleEndian>(self.y_origin));
        try!(w.write_u16::<LittleEndian>(self.image_width));
        try!(w.write_u16::<LittleEndian>(self.image_height));
        try!(w.write_u8(self.pixel_depth));
        try!(w.write_u8(self.image_desc));
        Ok(())
    }
}

struct Matrix<T> {
    width: usize,
    height: usize,
    vec: Vec<T>,
}

pub struct TgaImage {
    m: Matrix<TgaColor>,
}

impl<T: Copy> Matrix<T> {
    pub fn set(self: &mut Self, i: usize, j: usize, t: T) -> () {
        self.vec[j * self.height + i] = t
    }
    fn get(self: &Self, i: usize, j: usize) -> T {
        self.vec[j * self.height + i]
    }
}

impl TgaImage {
    pub fn new(w: usize, h: usize) -> TgaImage {
        let c = TgaColor { rgba: [0, 0, 0, 255] };
        TgaImage {
            m: Matrix {
                width: w,
                height: h,
                vec: vec![c; w * h],
            },
        }
    }
    pub fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(Header::new(self.m.width as u16, self.m.height as u16).write(w));
        for c in &self.m.vec {
            try!(c.write(w));
        }
        Ok(())
    }
    pub fn set(self: &mut Self, i: usize, j: usize, c: TgaColor) -> () {
        self.m.set(i, j, c)
    }
}
<commit_msg>simplify the tga<commit_after>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::Result;

/// Header used by TGA image files
#[derive(Debug)]
struct Header {
    id_length: u8, // length of ID string
    map_type: u8, // color map type
    image_type: u8, // image type code
    map_origin: u16, // starting index of map
    map_length: u16, // length of map
    map_entry_size: u8, // size of map entries in bits
    x_origin: u16, // x-origin of image
    y_origin: u16, // y-origin of image
    image_width: u16, // width of image
    image_height: u16, // height of image
    pixel_depth: u8, // bits per pixel
    image_desc: u8, // image descriptor
}

#[derive(Clone,Copy)]
pub struct TgaColor {
    pub rgba: [u8; 4],
}

impl TgaColor {
    fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(w.write_u8(self.rgba[2]));
        try!(w.write_u8(self.rgba[1]));
        try!(w.write_u8(self.rgba[0]));
        try!(w.write_u8(self.rgba[3]));
        Ok(())
    }
}

const RAW_TRUE_COLOR: u8 = 2;
const ARGB_DEPTH: u8 = 32;

impl Header {
    fn new(w: u16, h: u16) -> Header {
        Header {
            id_length: 0,
            map_type: 0,
            image_type: RAW_TRUE_COLOR,
            map_origin: 0,
            map_length: 0,
            map_entry_size: 0,
            x_origin: 0,
            y_origin: 0,
            image_width: w,
            image_height: h,
            pixel_depth: ARGB_DEPTH,
            image_desc: 0,
        }
    }

    fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(w.write_u8(self.id_length));
        try!(w.write_u8(self.map_type));
        try!(w.write_u8(self.image_type));
        try!(w.write_u16::<LittleEndian>(self.map_origin));
        try!(w.write_u16::<LittleEndian>(self.map_length));
        try!(w.write_u8(self.map_entry_size));
        try!(w.write_u16::<LittleEndian>(self.x_origin));
        try!(w.write_u16::<LittleEndian>(self.y_origin));
        try!(w.write_u16::<LittleEndian>(self.image_width));
        try!(w.write_u16::<LittleEndian>(self.image_height));
        try!(w.write_u8(self.pixel_depth));
        try!(w.write_u8(self.image_desc));
        Ok(())
    }
}

pub struct TgaImage {
    width: usize,
    height: usize,
    data: Vec<TgaColor>,
}

impl TgaImage {
    pub fn new(w: usize, h: usize) -> TgaImage {
        let c = TgaColor { rgba: [0, 0, 0, 255] };
        TgaImage {
            width: w,
            height: h,
            data: vec![c; w * h],
        }
    }
    pub fn write<W: WriteBytesExt>(self: &Self, w: &mut W) -> Result<()> {
        try!(Header::new(self.width as u16, self.height as u16).write(w));
        for c in &self.data {
            try!(c.write(w));
        }
        Ok(())
    }
    pub fn set(self: &mut Self, i: usize, j: usize, t: TgaColor) -> () {
        self.data[j * self.height + i] = t
    }
    fn get(self: &Self, i: usize, j: usize) -> TgaColor {
        self.data[j * self.height + i]
    }
}
<|endoftext|>


<commit_before>#[link(
    name = "depviz",
    vers = "0.1.0",
    uuid = "aecc116d-c1bd-4847-9435-f020d565e181"
)];

#[author = "KokaKiwi <kokakiwi@kokakiwi.net>"];
#[license = "MIT"];

#[crate_type = "lib"];

extern mod extra;
extern mod syntax;

use std::os;
use extra::getopts::groups;

mod depviz;

static BRIEF_USAGE: &'static str = "Rust crate dependencies visualization.";

fn print_usage(program: &str, opts: &[groups::OptGroup])
{
    println!("Usage: {:s} [OPTIONS] FILE", program);
    println!("");
    print(groups::usage(BRIEF_USAGE, opts));
}

fn main_args(args: &[~str]) -> int
{
    let program = args[0].as_slice();

    let opts = ~[
        groups::optflag("h", "help", "Show this help and exit."),
    ];

    let matches = match groups::getopts(args, opts) {
        Ok(m) => m,
        Err(f) => fail!(f.to_err_msg()),
    };

    if matches.opts_present([~"h", ~"help"])
    {
        print_usage(program, opts);
        return 0;
    }

    if matches.free.len() <= 1
    {
        print_usage(program, opts);
        return 1;
    }

    return 0;
}

fn main()
{
    os::set_exit_status(main_args(os::args()));
}
<commit_msg>Bad, bad commit. :(<commit_after>#[link(
    name = "depviz",
    vers = "0.1.0",
    uuid = "aecc116d-c1bd-4847-9435-f020d565e181"
)];

#[author = "KokaKiwi <kokakiwi@kokakiwi.net>"];
#[license = "MIT"];

#[crate_type = "bin"];

extern mod extra;
extern mod syntax;

use std::os;
use extra::getopts::groups;

mod depviz;

static BRIEF_USAGE: &'static str = "Rust crate dependencies visualization.";

fn print_usage(program: &str, opts: &[groups::OptGroup])
{
    println!("Usage: {:s} [OPTIONS] FILE", program);
    println!("");
    print(groups::usage(BRIEF_USAGE, opts));
}

fn main_args(args: &[~str]) -> int
{
    let program = args[0].as_slice();

    let opts = ~[
        groups::optflag("h", "help", "Show this help and exit."),
    ];

    let matches = match groups::getopts(args, opts) {
        Ok(m) => m,
        Err(f) => fail!(f.to_err_msg()),
    };

    if matches.opts_present([~"h", ~"help"])
    {
        print_usage(program, opts);
        return 0;
    }

    if matches.free.len() <= 1
    {
        print_usage(program, opts);
        return 1;
    }

    return 0;
}

fn main()
{
    os::set_exit_status(main_args(os::args()));
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::{File, OpenOptions};
use std::process::{Command, Stdio};
use std::os::unix::AsRawFd;
use libc::{c_ushort, c_int, c_ulong};
use std::str;

pub struct TTY {
    file: File,
    dimensions: (usize, usize),
    original_state: String
}

pub trait IO {
    fn write(&mut self, line: &str);
    fn read(&mut self) -> Option<String>;
    fn last(&self) -> &str;
    fn lines(&self) -> Vec<String>;
    fn dimensions(&self) -> (usize, usize);
    fn reset(&self);
}

impl IO for TTY {
    fn write(&mut self, line: &str) {
        let it = format!("{}", line);
        self.file.write(it.as_bytes()).unwrap();
    }

    fn read(&mut self) -> Option<String> {
        let mut buffer = [0] ;
        let res = match self.file.read(&mut buffer) {
            Ok(c) if  c > 0 => {
                str::from_utf8(&buffer).ok().map(|x| x.to_string())
            },
            _ => None,
        };
        res
    }

    fn last(&self) -> &str {
        "fail"
    }

    fn lines(&self) -> Vec<String> {
        let mut lines: Vec<String> = Vec::new();
        lines.push("fail".to_string());
        lines
    }

    fn dimensions(&self) -> (usize, usize) {
        self.dimensions
    }

    fn reset(&self) {
        TTY::stty(&self.file, &[self.original_state.as_slice()]);
    }
}

impl TTY {
    pub fn new() -> TTY {
        let path = Path::new("/dev/tty");
        let file = OpenOptions::new().read(true).write(true).open(&path).unwrap();
        let dimension = TTY::get_window_size(&file);
        let orig_state = TTY::previous_state(&file);

        TTY::no_echo_no_escaping(&file);

        TTY {
            original_state: orig_state,
            dimensions: dimension,
            file: file,
        }
    }

    fn get_window_size(path: &File) -> (usize, usize) {
        extern {
            fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int;
        }
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(path.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            (size.cols as usize, size.rows as usize)
        } else {
            panic!("Could not read winsize from /dev/tty")
        }
    }

    fn stty(file: &File, args: &[&str]) -> Option<String> {
        let output = Command::new("stty").args(args).stdin(Stdio::capture()).output().unwrap();
        String::from_utf8(output.stdout).ok()
    }

    fn no_echo_no_escaping(file: &File) {
        TTY::stty(file, &["-echo", "-icanon"]);
    }

    fn previous_state(file: &File) -> String {
        TTY::stty(file, &["-g"]).unwrap_or("".to_string())
    }
}
<commit_msg>Use new path and remove no longer needed file param<commit_after>use std::io::prelude::*;
use std::fs::{File, OpenOptions};
use std::process::{Command, Stdio};
use std::os::unix::AsRawFd;
use libc::{c_ushort, c_int, c_ulong};
use std::str;
use std::path::Path;

pub struct TTY {
    file: File,
    dimensions: (usize, usize),
    original_state: String
}

pub trait IO {
    fn write(&mut self, line: &str);
    fn read(&mut self) -> Option<String>;
    fn last(&self) -> &str;
    fn lines(&self) -> Vec<String>;
    fn dimensions(&self) -> (usize, usize);
    fn reset(&self);
}

impl IO for TTY {
    fn write(&mut self, line: &str) {
        let it = format!("{}", line);
        self.file.write(it.as_bytes()).unwrap();
    }

    fn read(&mut self) -> Option<String> {
        let mut buffer = [0] ;
        let res = match self.file.read(&mut buffer) {
            Ok(c) if  c > 0 => {
                str::from_utf8(&buffer).ok().map(|x| x.to_string())
            },
            _ => None,
        };
        res
    }

    fn last(&self) -> &str {
        "fail"
    }

    fn lines(&self) -> Vec<String> {
        let mut lines: Vec<String> = Vec::new();
        lines.push("fail".to_string());
        lines
    }

    fn dimensions(&self) -> (usize, usize) {
        self.dimensions
    }

    fn reset(&self) {
        TTY::stty(&[self.original_state.as_slice()]);
    }
}

impl TTY {
    pub fn new() -> TTY {
        let path = Path::new("/dev/tty");
        let file = OpenOptions::new().read(true).write(true).open(&path).unwrap();
        let dimension = TTY::get_window_size(&file);
        let orig_state = TTY::previous_state();

        TTY::no_echo_no_escaping();

        TTY {
            original_state: orig_state,
            dimensions: dimension,
            file: file,
        }
    }

    fn get_window_size(path: &File) -> (usize, usize) {
        extern {
            fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int;
        }
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(path.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            (size.cols as usize, size.rows as usize)
        } else {
            panic!("Could not read winsize from /dev/tty")
        }
    }

    fn stty(args: &[&str]) -> Option<String> {
        let output = Command::new("stty").args(args).stdin(Stdio::capture()).output().unwrap();
        String::from_utf8(output.stdout).ok()
    }

    fn no_echo_no_escaping() {
        TTY::stty(&["-echo", "-icanon"]);
    }

    fn previous_state() -> String {
        TTY::stty(&["-g"]).unwrap_or("".to_string())
    }
}
<|endoftext|>


<commit_before>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

extern crate png;

use std::io::{self, Read, Write};

use color::ColorType;
use image::{ImageDecoder, ImageError, ImageResult};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not posible and
/// these are therefore readed at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PNGDecoder<R: Read> {
    colortype: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PNGDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PNGDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let decoder = png::Decoder::new_with_limits(r, limits);
        let (_, mut reader) = decoder.read_info()?;
        let colortype = reader.output_color_type().into();

        Ok(PNGDecoder { colortype, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PNGDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        let (w, h) = self.reader.info().size();
        (u64::from(w), u64::from(h))
    }

    fn colortype(&self) -> ColorType {
        self.colortype
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        // This should be slightly faster than the default implementation
        let mut data = vec![0; self.reader.output_buffer_size()];
        self.reader.next_frame(&mut data)?;
        Ok(data)
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> io::Result<()> {
        let (ct, bits) = match color {
            ColorType::Gray(bits) => Some((png::ColorType::Grayscale, bits)),
            ColorType::RGB(bits) => Some((png::ColorType::RGB, bits)),
            ColorType::Palette(bits) => Some((png::ColorType::Indexed, bits)),
            ColorType::GrayA(bits) => Some((png::ColorType::GrayscaleAlpha, bits)),
            ColorType::RGBA(bits) => Some((png::ColorType::RGBA, bits)),
            _ => None,
        }
        .and_then(|(ct, bits)| Some((ct, png::BitDepth::from_u8(bits)?)))
        // FIXME: After #1066 lands, return the ImageError directly.
        .ok_or(io::Error::new(io::ErrorKind::InvalidInput, ImageError::UnsupportedColor(color)))?;

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header()?;
        writer.write_image_data(data).map_err(|e| e.into())
    }
}

impl From<(png::ColorType, png::BitDepth)> for ColorType {
    fn from((ct, bits): (png::ColorType, png::BitDepth)) -> ColorType {
        use self::png::ColorType::*;
        let bits = bits as u8;
        match ct {
            Grayscale => ColorType::Gray(bits),
            RGB => ColorType::RGB(bits),
            Indexed => ColorType::Palette(bits),
            GrayscaleAlpha => ColorType::GrayA(bits),
            RGBA => ColorType::RGBA(bits),
        }
    }
}

impl From<ColorType> for (png::ColorType, png::BitDepth) {
    fn from(ct: ColorType) -> (png::ColorType, png::BitDepth) {
        use self::png::ColorType::*;
        let (ct, bits) = match ct {
            ColorType::Gray(bits) => (Grayscale, bits),
            ColorType::RGB(bits) => (RGB, bits),
            ColorType::Palette(bits) => (Indexed, bits),
            ColorType::GrayA(bits) => (GrayscaleAlpha, bits),
            ColorType::RGBA(bits) => (RGBA, bits),
            ColorType::BGRA(bits) => (RGBA, bits),
            ColorType::BGR(bits) => (RGB, bits),
        };
        (ct, png::BitDepth::from_u8(bits).unwrap())
    }
}

impl From<png::DecodingError> for ImageError {
    fn from(err: png::DecodingError) -> ImageError {
        use self::png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            Format(desc) => ImageError::FormatError(desc.into_owned()),
            InvalidSignature => ImageError::FormatError("invalid signature".into()),
            CrcMismatch { .. } => ImageError::FormatError("CRC error".into()),
            Other(desc) => ImageError::FormatError(desc.into_owned()),
            CorruptFlateStream => {
                ImageError::FormatError("compressed data stream corrupted".into())
            }
            LimitsExceeded => ImageError::InsufficientMemory,
        }
    }
}

#[cfg(test)]
mod tests {
    use image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PNGDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::RGB(8),
            dec.colortype(),
            "Image MUST have the RGB(8) format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }
}
<commit_msg>Fix documentation of PNGReader<commit_after>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

extern crate png;

use std::io::{self, Read, Write};

use color::ColorType;
use image::{ImageDecoder, ImageError, ImageResult};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not possible and
/// these are therefore read at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PNGDecoder<R: Read> {
    colortype: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PNGDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PNGDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let decoder = png::Decoder::new_with_limits(r, limits);
        let (_, mut reader) = decoder.read_info()?;
        let colortype = reader.output_color_type().into();

        Ok(PNGDecoder { colortype, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PNGDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        let (w, h) = self.reader.info().size();
        (u64::from(w), u64::from(h))
    }

    fn colortype(&self) -> ColorType {
        self.colortype
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        // This should be slightly faster than the default implementation
        let mut data = vec![0; self.reader.output_buffer_size()];
        self.reader.next_frame(&mut data)?;
        Ok(data)
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> io::Result<()> {
        let (ct, bits) = match color {
            ColorType::Gray(bits) => Some((png::ColorType::Grayscale, bits)),
            ColorType::RGB(bits) => Some((png::ColorType::RGB, bits)),
            ColorType::Palette(bits) => Some((png::ColorType::Indexed, bits)),
            ColorType::GrayA(bits) => Some((png::ColorType::GrayscaleAlpha, bits)),
            ColorType::RGBA(bits) => Some((png::ColorType::RGBA, bits)),
            _ => None,
        }
        .and_then(|(ct, bits)| Some((ct, png::BitDepth::from_u8(bits)?)))
        // FIXME: After #1066 lands, return the ImageError directly.
        .ok_or(io::Error::new(io::ErrorKind::InvalidInput, ImageError::UnsupportedColor(color)))?;

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header()?;
        writer.write_image_data(data).map_err(|e| e.into())
    }
}

impl From<(png::ColorType, png::BitDepth)> for ColorType {
    fn from((ct, bits): (png::ColorType, png::BitDepth)) -> ColorType {
        use self::png::ColorType::*;
        let bits = bits as u8;
        match ct {
            Grayscale => ColorType::Gray(bits),
            RGB => ColorType::RGB(bits),
            Indexed => ColorType::Palette(bits),
            GrayscaleAlpha => ColorType::GrayA(bits),
            RGBA => ColorType::RGBA(bits),
        }
    }
}

impl From<ColorType> for (png::ColorType, png::BitDepth) {
    fn from(ct: ColorType) -> (png::ColorType, png::BitDepth) {
        use self::png::ColorType::*;
        let (ct, bits) = match ct {
            ColorType::Gray(bits) => (Grayscale, bits),
            ColorType::RGB(bits) => (RGB, bits),
            ColorType::Palette(bits) => (Indexed, bits),
            ColorType::GrayA(bits) => (GrayscaleAlpha, bits),
            ColorType::RGBA(bits) => (RGBA, bits),
            ColorType::BGRA(bits) => (RGBA, bits),
            ColorType::BGR(bits) => (RGB, bits),
        };
        (ct, png::BitDepth::from_u8(bits).unwrap())
    }
}

impl From<png::DecodingError> for ImageError {
    fn from(err: png::DecodingError) -> ImageError {
        use self::png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            Format(desc) => ImageError::FormatError(desc.into_owned()),
            InvalidSignature => ImageError::FormatError("invalid signature".into()),
            CrcMismatch { .. } => ImageError::FormatError("CRC error".into()),
            Other(desc) => ImageError::FormatError(desc.into_owned()),
            CorruptFlateStream => {
                ImageError::FormatError("compressed data stream corrupted".into())
            }
            LimitsExceeded => ImageError::InsufficientMemory,
        }
    }
}

#[cfg(test)]
mod tests {
    use image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PNGDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::RGB(8),
            dec.colortype(),
            "Image MUST have the RGB(8) format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }
}
<|endoftext|>


<commit_before>use heartbeats_simple_sys::pow::*;
use std::mem;
use std::io::{self, Write};
use std::fs::File;

/// Contains the Heartbeat and its window data buffer.
pub struct HeartbeatPow {
    pub hb: heartbeat_pow_context,
    pub hbr: Vec<heartbeat_pow_record>,
    pub log: Option<File>,
}

impl HeartbeatPow {
    /// Allocate and initialize a new `Heartbeat`.
    pub fn new(window_size: usize,
               hwc_callback: Option<heartbeat_pow_window_complete>,
               mut log: Option<File>) -> Result<HeartbeatPow, &'static str> {
        let mut hbr: Vec<heartbeat_pow_record> = Vec::with_capacity(window_size);
        let hb: heartbeat_pow_context = unsafe {
            // must explicitly set size so we can read data later
            // (Rust isn't aware of native code modifying the buffer)
            hbr.set_len(window_size);
            let mut hb = mem::uninitialized();
            match heartbeat_pow_init(&mut hb,
                                     hbr.capacity() as u64,
                                     hbr.as_mut_ptr(),
                                     hwc_callback) {
                0 => hb,
                _ => return Err("Failed to initialize heartbeat")
            }
        };
        // write header to log file if there is one
        if let Some(ref mut l) = log {
            l.write_all(format!("{:6} {:6} \
                         {:11} {:11} {:11} \
                         {:15} {:15} {:20} {:20} \
                         {:15} {:15} {:15} \
                         {:15} {:15} {:15} {:15} \
                         {:15} {:15} {}\n",
                        "HB", "Tag",
                        "Global_Work", "Window_Work", "Work",
                        "Global_Time", "Window_Time", "Start_Time", "End_Time",
                        "Global_Perf", "Window_Perf", "Instant_Perf",
                        "Global_Energy", "Window_Energy", "Start_Energy", "End_Energy",
                        "Global_Pwr", "Window_Pwr", "Instant_Pwr").as_bytes()).unwrap()
        }
        Ok(HeartbeatPow { hb: hb, hbr: hbr, log: log, })
    }

    /// Issue a heartbeat
    pub fn heartbeat(&mut self,
                     tag: u64,
                     work: u64,
                     start_time: u64,
                     end_time: u64,
                     start_energy: u64,
                     end_energy: u64) {
        unsafe {
            heartbeat_pow(&mut self.hb, tag, work, start_time, end_time, start_energy, end_energy)
        }
    }

    fn write_log(r: &heartbeat_pow_record, l: &mut File) -> io::Result<usize> {
        l.write(format!("{:<6} {:<6} \
                         {:<11} {:<11} {:<11} \
                         {:<15} {:<15} {:<20} {:<20} \
                         {:<15.6} {:<15.6} {:<15.6} \
                         {:<15} {:<15} {:<15} {:<15} \
                         {:<15.6} {:<15.6} {:<.6}\n",
                        r.id, r.user_tag,
                        r.wd.global, r.wd.window, r.work,
                        r.td.global, r.td.window, r.start_time, r.end_time,
                        r.perf.global, r.perf.window, r.perf.instant,
                        r.ed.global, r.ed.window, r.start_energy, r.end_energy,
                        r.pwr.global, r.pwr.window, r.pwr.instant).as_bytes())
    }

    /// Rust-only function that logs the buffer (up to buffer_index) to a file.
    pub fn log_to_buffer_index(&mut self) -> io::Result<()> {
        match self.log {
            Some(ref mut l) => {
                for i in 0..self.hb.ws.buffer_index {
                    match HeartbeatPow::write_log(self.hbr.get(i as usize).unwrap(), l) {
                        Ok(_) => (),
                        Err(e) => return Err(e),
                    }
                }
                l.flush()
            }
            None => Ok(())
        }
    }
}

#[cfg(test)]
mod test {
    use super::HeartbeatPow;
    use std::fs::File;

    #[test]
    fn test_simple() {
        const TIME_INC: u64 = 1000000000;
        const ENERGY_INC: u64 = 1000000;
        let mut hb = HeartbeatPow::new(5, None, None).unwrap();
        let mut start_time: u64 = 0;
        let mut end_time: u64 = TIME_INC;
        let mut start_energy: u64 = 0;
        let mut end_energy: u64 = ENERGY_INC;
        for tag in 0..10 {
            hb.heartbeat(tag, 1, start_time, end_time, start_energy, end_energy);
            start_time = end_time;
            end_time += TIME_INC;
            start_energy = end_energy;
            end_energy += ENERGY_INC;
        }
    }

    #[test]
    fn test_file() {
        let mut hb = HeartbeatPow::new(5, None, Some(File::create("foo.log").unwrap())).unwrap();
        hb.heartbeat(0, 1, 0, 1000, 0, 0);
        hb.log_to_buffer_index().unwrap();
    }
}
<commit_msg>Add wrapper types<commit_after>use heartbeats_simple_sys::pow::*;
use std::mem;
use std::io::{self, Write};
use std::fs::File;

pub type HeartbeatPowContext = heartbeat_pow_context;
pub type HeartbeatPowRecord = heartbeat_pow_record;
pub type HeartbeatPowWindowComplete = heartbeat_pow_window_complete;

/// Contains the Heartbeat and its window data buffer.
pub struct HeartbeatPow {
    pub hb: HeartbeatPowContext,
    pub hbr: Vec<HeartbeatPowRecord>,
    pub log: Option<File>,
}

impl HeartbeatPow {
    /// Allocate and initialize a new `Heartbeat`.
    pub fn new(window_size: usize,
               hwc_callback: Option<HeartbeatPowWindowComplete>,
               mut log: Option<File>) -> Result<HeartbeatPow, &'static str> {
        let mut hbr: Vec<HeartbeatPowRecord> = Vec::with_capacity(window_size);
        let hb: HeartbeatPowContext = unsafe {
            // must explicitly set size so we can read data later
            // (Rust isn't aware of native code modifying the buffer)
            hbr.set_len(window_size);
            let mut hb = mem::uninitialized();
            match heartbeat_pow_init(&mut hb,
                                     hbr.capacity() as u64,
                                     hbr.as_mut_ptr(),
                                     hwc_callback) {
                0 => hb,
                _ => return Err("Failed to initialize heartbeat")
            }
        };
        // write header to log file if there is one
        if let Some(ref mut l) = log {
            l.write_all(format!("{:6} {:6} \
                         {:11} {:11} {:11} \
                         {:15} {:15} {:20} {:20} \
                         {:15} {:15} {:15} \
                         {:15} {:15} {:15} {:15} \
                         {:15} {:15} {}\n",
                        "HB", "Tag",
                        "Global_Work", "Window_Work", "Work",
                        "Global_Time", "Window_Time", "Start_Time", "End_Time",
                        "Global_Perf", "Window_Perf", "Instant_Perf",
                        "Global_Energy", "Window_Energy", "Start_Energy", "End_Energy",
                        "Global_Pwr", "Window_Pwr", "Instant_Pwr").as_bytes()).unwrap()
        }
        Ok(HeartbeatPow { hb: hb, hbr: hbr, log: log, })
    }

    /// Issue a heartbeat
    pub fn heartbeat(&mut self,
                     tag: u64,
                     work: u64,
                     start_time: u64,
                     end_time: u64,
                     start_energy: u64,
                     end_energy: u64) {
        unsafe {
            heartbeat_pow(&mut self.hb, tag, work, start_time, end_time, start_energy, end_energy)
        }
    }

    fn write_log(r: &HeartbeatPowRecord, l: &mut File) -> io::Result<usize> {
        l.write(format!("{:<6} {:<6} \
                         {:<11} {:<11} {:<11} \
                         {:<15} {:<15} {:<20} {:<20} \
                         {:<15.6} {:<15.6} {:<15.6} \
                         {:<15} {:<15} {:<15} {:<15} \
                         {:<15.6} {:<15.6} {:<.6}\n",
                        r.id, r.user_tag,
                        r.wd.global, r.wd.window, r.work,
                        r.td.global, r.td.window, r.start_time, r.end_time,
                        r.perf.global, r.perf.window, r.perf.instant,
                        r.ed.global, r.ed.window, r.start_energy, r.end_energy,
                        r.pwr.global, r.pwr.window, r.pwr.instant).as_bytes())
    }

    /// Rust-only function that logs the buffer (up to buffer_index) to a file.
    pub fn log_to_buffer_index(&mut self) -> io::Result<()> {
        match self.log {
            Some(ref mut l) => {
                for i in 0..self.hb.ws.buffer_index {
                    match HeartbeatPow::write_log(self.hbr.get(i as usize).unwrap(), l) {
                        Ok(_) => (),
                        Err(e) => return Err(e),
                    }
                }
                l.flush()
            }
            None => Ok(())
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;

    #[test]
    fn test_simple() {
        const TIME_INC: u64 = 1000000000;
        const ENERGY_INC: u64 = 1000000;
        let mut hb = HeartbeatPow::new(5, None, None).unwrap();
        let mut start_time: u64 = 0;
        let mut end_time: u64 = TIME_INC;
        let mut start_energy: u64 = 0;
        let mut end_energy: u64 = ENERGY_INC;
        for tag in 0..10 {
            hb.heartbeat(tag, 1, start_time, end_time, start_energy, end_energy);
            start_time = end_time;
            end_time += TIME_INC;
            start_energy = end_energy;
            end_energy += ENERGY_INC;
        }
    }

    #[test]
    fn test_callback() {
        static mut received_cb: bool = false;
        extern fn callback(_hb: *const HeartbeatPowContext) {
            unsafe {
                received_cb = true;
            }
        }

        let mut hb = HeartbeatPow::new(1, Some(callback), None).unwrap();
        hb.heartbeat(0, 1, 0, 1000, 0, 0);
        unsafe {
            assert!(received_cb);
        }
    }

    #[test]
    fn test_file() {
        let mut hb = HeartbeatPow::new(5, None, Some(File::create("foo.log").unwrap())).unwrap();
        hb.heartbeat(0, 1, 0, 1000, 0, 0);
        hb.log_to_buffer_index().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::collections::VecMap;
use std::rand;

use poly::*;

/// Split a secret into N shares, of which K are required to re-combine. Returns
/// a map of share IDs to share values.
pub fn split<'a, T: rand::Rng>(n: u8, k: u8, secret: &'a Vec<u8>, rng: &mut T) -> VecMap<Vec<u8>> {
    // Generate a random K-degree polynomial for each byte of the secret.
    let polys: Vec<Vec<u8>> = secret.iter().map(|b| {
        generate(k-1, *b, rng)
    }).collect();

    // Collect the evaluation of each polynomial with the share ID as the input.
    let mut shares: VecMap<Vec<u8>> = VecMap::with_capacity(n as usize);
    for id in (1..n+1) {
        shares.insert(id as usize, polys.iter().map(|p| eval(p, id)).collect());
    }
    return shares;
}

/// Combine a map of share IDs into the original secret.
///
/// N.B.: There is no way to know if this is successful or not.
pub fn combine<'a>(shares: &'a VecMap<Vec<u8>>) -> Vec<u8> {
    let mut points: Vec<Vec<(u8, u8)>> = Vec::new();
    for (id, share) in shares.iter() {
        for (i, v) in share.iter().enumerate() {
            if points.len() <= i {
                points.push(Vec::new())
            }
            points[i].push((id as u8, *v));
        }
    }
    return points.iter().map(|v| interpolate(v, 0)).collect();
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::VecMap;
    use std::rand;

    #[test]
    fn test_split() {
        let mut rng = rand::ChaChaRng::new_unseeded();
        let actual = split(5, 3, &vec![1, 2, 3, 4, 5], &mut rng);

        let mut expected: VecMap<Vec<u8>> = VecMap::new();
        expected.insert(1, vec![118, 163, 66, 80, 187]);
        expected.insert(2, vec![239, 91, 129, 172, 98]);
        expected.insert(3, vec![152, 250, 192, 248, 220]);
        expected.insert(4, vec![198, 176, 28, 79, 203]);
        expected.insert(5, vec![177, 17, 93, 27, 117]);

        assert_eq!(actual, expected)
    }

    #[test]
    fn test_combine() {
        let mut shares: VecMap<Vec<u8>> = VecMap::new();
        shares.insert(1, vec![118, 163, 66, 80, 187]);
        shares.insert(2, vec![239, 91, 129, 172, 98]);
        shares.insert(3, vec![152, 250, 192, 248, 220]);

        assert_eq!(combine(&shares), vec![1, 2, 3, 4, 5])
    }
}
<commit_msg>Simplify split into a map.<commit_after>use std::collections::VecMap;
use std::rand;

use poly::*;

/// Split a secret into N shares, of which K are required to re-combine. Returns
/// a map of share IDs to share values.
pub fn split<'a, T: rand::Rng>(n: u8, k: u8, secret: &'a Vec<u8>, rng: &mut T) -> VecMap<Vec<u8>> {
    // Generate a random K-degree polynomial for each byte of the secret.
    let polys: Vec<Vec<u8>> = secret.iter().map(|b| {
        generate(k-1, *b, rng)
    }).collect();

    // Collect the evaluation of each polynomial with the share ID as the input.
    (1..n+1).map(|id| {
        (id as usize, polys.iter().map(|p| eval(p, id)).collect())
    }).collect()
}

/// Combine a map of share IDs into the original secret.
///
/// N.B.: There is no way to know if this is successful or not.
pub fn combine<'a>(shares: &'a VecMap<Vec<u8>>) -> Vec<u8> {
    let mut points: Vec<Vec<(u8, u8)>> = Vec::new();
    for (id, share) in shares.iter() {
        for (i, v) in share.iter().enumerate() {
            if points.len() <= i {
                points.push(Vec::new())
            }
            points[i].push((id as u8, *v));
        }
    }
    return points.iter().map(|v| interpolate(v, 0)).collect();
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::VecMap;
    use std::rand;

    #[test]
    fn test_split() {
        let mut rng = rand::ChaChaRng::new_unseeded();
        let actual = split(5, 3, &vec![1, 2, 3, 4, 5], &mut rng);

        let mut expected: VecMap<Vec<u8>> = VecMap::new();
        expected.insert(1, vec![118, 163, 66, 80, 187]);
        expected.insert(2, vec![239, 91, 129, 172, 98]);
        expected.insert(3, vec![152, 250, 192, 248, 220]);
        expected.insert(4, vec![198, 176, 28, 79, 203]);
        expected.insert(5, vec![177, 17, 93, 27, 117]);

        assert_eq!(actual, expected)
    }

    #[test]
    fn test_combine() {
        let mut shares: VecMap<Vec<u8>> = VecMap::new();
        shares.insert(1, vec![118, 163, 66, 80, 187]);
        shares.insert(2, vec![239, 91, 129, 172, 98]);
        shares.insert(3, vec![152, 250, 192, 248, 220]);

        assert_eq!(combine(&shares), vec![1, 2, 3, 4, 5])
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::collections::HashMap;
use std::io::BufReader;
use std::borrow::Cow;

use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use errors::*;
use {ExcelReader, Range, DataType, CellErrorType};
use vba::VbaProject;
use cfb::Cfb;
use utils::{read_u16, read_u32, read_slice};

enum CfbWrap {
    Wb(Cfb, BufReader<File>),
    Vba(VbaProject),
}

/// A struct representing an old xls format file (CFB)
pub struct Xls {
    codepage: u16,
    sheets: HashMap<String, Range>,
    strings: Vec<String>,
    cfb: Option<CfbWrap>,
}

impl ExcelReader for Xls {

    fn new(r: File) -> Result<Self> {
        let len = try!(r.metadata()).len() as usize;
        let mut r = BufReader::new(r);
        let mut cfb = try!(Cfb::new(&mut r, len ));
        let wb = try!(cfb.get_stream("Workbook", &mut r));

        let mut xls = Xls {
            codepage: 1200,
            sheets: HashMap::new(),
            strings: Vec::new(),
            cfb: Some(CfbWrap::Wb(cfb, r)),
        };
        try!(xls.parse_workbook(&wb));
        Ok(xls)
    }

    fn has_vba(&mut self) -> bool {
        match self.cfb {
            Some(CfbWrap::Wb(ref cfb, _)) => cfb.has_directory("_VBA_PROJECT_CUR"),
            Some(CfbWrap::Vba(_)) => true,
            None => unreachable!(), // option is used to transfer ownership only
        }
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        if let Some(CfbWrap::Wb(..)) = self.cfb {
            match self.cfb.take() {
                Some(CfbWrap::Wb(cfb, mut r)) => {
                    let vba = try!(VbaProject::from_cfb(&mut r, cfb));
                    self.cfb = Some(CfbWrap::Vba(vba));
                },
                _ => unreachable!(),
            }
        }

        match self.cfb {
            Some(CfbWrap::Vba(ref v)) => Ok(Cow::Borrowed(v)),
            _ => unreachable!(),
        }
    }

    /// Parses Workbook stream, no need for the relationships variable
    fn read_sheets_names(&mut self, _: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>> {
        Ok(self.sheets.keys().map(|k| (k.to_string(), k.to_string())).collect())
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        Ok(self.strings.clone())
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        Ok(HashMap::new())
    }

    fn read_worksheet_range(&mut self, name: &str, _: &[String]) -> Result<Range> {
        match self.sheets.get(name) {
            None => Err(format!("Sheet '{}' does not exist", name).into()),
            Some(r) => Ok(r.clone()),
        }
    }
}

impl Xls {
    fn parse_workbook(&mut self, stream: &[u8]) -> Result<()> {

        let mut sheets = Vec::new();
        {
            let mut wb = stream;
            let records = RecordIter { stream: &mut wb };
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0005 {
                        return Err("Expecting Workbook BOF".into());
                    }, // BOF,
                    0x0042 => self.codepage = read_u16(&r.data), // CodePage (defines encoding)
                    0x013D => {
                        let sheet_len = r.data.len() / 2;
                        sheets.reserve(sheet_len);
                    }, // RRTabId
                    0x0085 => sheets.push(try!(parse_sheet_name(&r.data))), // BoundSheet8
                    0x00FC => self.strings = try!(parse_sst(&r.data)), // SST
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
        }

        'sh: for (pos, name) in sheets.into_iter() {
            let mut sh = &stream[pos..];
            let records = RecordIter { stream: &mut sh };
            let mut range = Range::default();
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0010 { continue 'sh; }, // BOF, worksheet
                    0x0200 => range = try!(parse_dimensions(&r.data)), // Dimensions
                    0x0203 => try!(parse_number(&r.data, &mut range)), // Number
                    0x0205 => try!(parse_bool_err(&r.data, &mut range)), // BoolErr
                    0x027E => try!(parse_rk(&r.data, &mut range)), // RK
                    0x00FD => try!(parse_label_sst(&r.data, &self.strings, &mut range)), // LabelSst
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
            self.sheets.insert(name, range);
        }
        Ok(())
    }
}

fn parse_sheet_name(r: &[u8]) -> Result<(usize, String)> {
    let pos = read_u32(r) as usize;
    let sheet = try!(parse_short_string(&r[6..]));
    Ok((pos, sheet))
}

fn parse_sst(r: &[u8]) -> Result<Vec<String>> {
    if r.len() < 8 {
        return Err("Invalid sst length".into());
    }
    let len = read_slice::<i32>(&r[4..]) as usize;
    let mut sst = Vec::with_capacity(len);
    let mut read = &mut &r[8..];
    for _ in 0..len {
        sst.push(try!(read_rich_extended_string(read)));
    }
    Ok(sst)
}

fn parse_number(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 14 {
        return Err("Invalid number length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = read_slice::<f64>(&r[6..]);
    range.set_value((row as u32, col as u32), DataType::Float(v))
}

fn parse_bool_err(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 8 {
        return Err("Invalid BoolErr length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = match r[7] {
        0x00 => DataType::Bool(r[6] != 0),
        0x01 => match r[6] {
            0x00 => DataType::Error(CellErrorType::Null),
            0x07 => DataType::Error(CellErrorType::Div0),
            0x0F => DataType::Error(CellErrorType::Value),
            0x17 => DataType::Error(CellErrorType::Ref),
            0x1D => DataType::Error(CellErrorType::Name),
            0x24 => DataType::Error(CellErrorType::Num),
            0x2A => DataType::Error(CellErrorType::NA),
            0x2B => DataType::Error(CellErrorType::GettingData),
            e => return Err(format!("Unrecognized error {:x}", e).into()),
        },
        e => return Err(format!("Unrecognized fError {:x}", e).into()),
    };
    range.set_value((row as u32, col as u32), v)
}

fn parse_rk(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid rk length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);

    let d100 = (r[6] & 1) != 0;
    let is_int = (r[6] & 2) != 0;

    let mut v = [0u8; 8];
    v[4..].copy_from_slice(&r[6..10]);
    v[0] &= 0xFC;
    let v = if is_int {
        let v = (read_slice::<i32>(&v[4..]) >> 2) as i64;
        DataType::Int( if d100 { v/100 } else { v })
    } else {
        let v = read_slice(&v);
        DataType::Float( if d100 { v/100.0 } else { v })
    };

    range.set_value((row as u32, col as u32), v)
}

fn parse_short_string(r: &[u8]) -> Result<String> {
    if r.len() < 2 {
        return Err("Invalid short string length".into());
    }
    let mut len = r[0] as usize;
    let zero_high_byte = r[1] == 0;
    let bytes = if zero_high_byte {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; len * 2];
        for (i, sce) in r[2..2 + len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        len *= 2;
        Cow::Borrowed(&r[2..2 + len])
    };
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn parse_label_sst(r: &[u8], strings: &[String], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid short string length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let i = read_u32(&r[6..]) as usize;
    range.set_value((row as u32, col as u32), DataType::String(strings[i].clone()))
}

fn parse_dimensions(r: &[u8]) -> Result<Range> {
    if r.len() != 14 {
        return Err("Invalid dimensions lengths".into());
    }
    let rw_first = read_u32(&r[0..4]);
    let rw_last = read_u32(&r[4..8]);
    let col_first = read_u16(&r[8..10]);
    let col_last = read_u16(&r[10..12]);

    let start = (rw_first, col_first as u32);
    let size = ((rw_last - rw_first) as usize, (col_last - col_first) as usize);

    Ok(Range::new(start, size))
}

fn read_rich_extended_string(r: &mut &[u8]) -> Result<String> {
    if r.len() < 3 {
        return Err("Invalid rich extended string length".into());
    }

    let mut len = read_u16(r) as usize;
    let mut start = 3;
    let flags = r[2];
    let high_byte = flags & 0x1;
    let ext_st = flags & 0x4;
    let rich_st = flags & 0x8;
    let str_len = len;

    if rich_st != 0 {
        len += 4 * read_u16(&r[start..]) as usize;
        start += 2;
    }
    if ext_st != 0 {
        len += read_slice::<i32>(&r[start..]) as usize;
        start += 4;
    }

    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; str_len * 2];
        for (i, sce) in r[start..start + str_len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        len += str_len;
        Cow::Borrowed(&r[start..start + 2 * str_len])
    };

    if r.len() < start + len {
        return Err(format!("Invalid rich extended string \
                            length: {} < {}", r.len(), start + len).into());
    }

    *r = &r[start + len..];
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

struct Record<'a> {
    typ: u16,
    data: Cow<'a, [u8]>,
}

struct RecordIter<'a> {
    stream: &'a[u8],
}

impl<'a> Iterator for RecordIter<'a> {
    type Item=Result<Record<'a>>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.stream.len() < 4 {
            return if self.stream.is_empty() {
                None
            } else {
                Some(Err("Expecting record type and length, found end of stream".into()))
            }
        }
        let t = read_u16(self.stream);
        let mut len = read_u16(&self.stream[2..]) as usize;
        if self.stream.len() < len + 4 {
            return Some(Err("Expecting record length, found end of stream".into()));
        }
        let (mut data, mut next) = self.stream.split_at(len + 4);
        self.stream = next;

        // Append next record data if it is a Continue record
        let cow = if next.len() > 4 && read_u16(next) == 0x003C {
            let mut c = data[4..].to_vec();
            while next.len() > 4 && read_u16(next) == 0x003C {
                len = read_u16(&self.stream[2..]) as usize;
                if self.stream.len() < len + 4 {
                    return Some(Err("Expecting continue record length, found end of stream".into()));
                }
                let sp = self.stream.split_at(len + 4);
                data = sp.0;
                next = sp.1;
                c.extend_from_slice(&data[4..]);
                self.stream = next;
            }
            Cow::Owned(c)
        } else {
            Cow::Borrowed(&data[4..])
        };

        Some(Ok(Record { typ: t, data: cow }))
    }
}
<commit_msg>xls return error when parsing too long strings (instead of panic)<commit_after>use std::fs::File;
use std::collections::HashMap;
use std::io::BufReader;
use std::borrow::Cow;

use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use errors::*;
use {ExcelReader, Range, DataType, CellErrorType};
use vba::VbaProject;
use cfb::Cfb;
use utils::{read_u16, read_u32, read_slice};

enum CfbWrap {
    Wb(Cfb, BufReader<File>),
    Vba(VbaProject),
}

/// A struct representing an old xls format file (CFB)
pub struct Xls {
    codepage: u16,
    sheets: HashMap<String, Range>,
    strings: Vec<String>,
    cfb: Option<CfbWrap>,
}

impl ExcelReader for Xls {

    fn new(r: File) -> Result<Self> {
        let len = try!(r.metadata()).len() as usize;
        let mut r = BufReader::new(r);
        let mut cfb = try!(Cfb::new(&mut r, len ));
        let wb = try!(cfb.get_stream("Workbook", &mut r));

        let mut xls = Xls {
            codepage: 1200,
            sheets: HashMap::new(),
            strings: Vec::new(),
            cfb: Some(CfbWrap::Wb(cfb, r)),
        };
        try!(xls.parse_workbook(&wb));
        Ok(xls)
    }

    fn has_vba(&mut self) -> bool {
        match self.cfb {
            Some(CfbWrap::Wb(ref cfb, _)) => cfb.has_directory("_VBA_PROJECT_CUR"),
            Some(CfbWrap::Vba(_)) => true,
            None => unreachable!(), // option is used to transfer ownership only
        }
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        if let Some(CfbWrap::Wb(..)) = self.cfb {
            match self.cfb.take() {
                Some(CfbWrap::Wb(cfb, mut r)) => {
                    let vba = try!(VbaProject::from_cfb(&mut r, cfb));
                    self.cfb = Some(CfbWrap::Vba(vba));
                },
                _ => unreachable!(),
            }
        }

        match self.cfb {
            Some(CfbWrap::Vba(ref v)) => Ok(Cow::Borrowed(v)),
            _ => unreachable!(),
        }
    }

    /// Parses Workbook stream, no need for the relationships variable
    fn read_sheets_names(&mut self, _: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>> {
        Ok(self.sheets.keys().map(|k| (k.to_string(), k.to_string())).collect())
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        Ok(self.strings.clone())
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        Ok(HashMap::new())
    }

    fn read_worksheet_range(&mut self, name: &str, _: &[String]) -> Result<Range> {
        match self.sheets.get(name) {
            None => Err(format!("Sheet '{}' does not exist", name).into()),
            Some(r) => Ok(r.clone()),
        }
    }
}

impl Xls {
    fn parse_workbook(&mut self, stream: &[u8]) -> Result<()> {

        let mut sheets = Vec::new();
        {
            let mut wb = stream;
            let records = RecordIter { stream: &mut wb };
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0005 {
                        return Err("Expecting Workbook BOF".into());
                    }, // BOF,
                    0x0042 => self.codepage = read_u16(&r.data), // CodePage (defines encoding)
                    0x013D => {
                        let sheet_len = r.data.len() / 2;
                        sheets.reserve(sheet_len);
                    }, // RRTabId
                    0x0085 => sheets.push(try!(parse_sheet_name(&r.data))), // BoundSheet8
                    0x00FC => self.strings = try!(parse_sst(&r.data)), // SST
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
        }

        'sh: for (pos, name) in sheets.into_iter() {
            let mut sh = &stream[pos..];
            let records = RecordIter { stream: &mut sh };
            let mut range = Range::default();
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0010 { continue 'sh; }, // BOF, worksheet
                    0x0200 => range = try!(parse_dimensions(&r.data)), // Dimensions
                    0x0203 => try!(parse_number(&r.data, &mut range)), // Number
                    0x0205 => try!(parse_bool_err(&r.data, &mut range)), // BoolErr
                    0x027E => try!(parse_rk(&r.data, &mut range)), // RK
                    0x00FD => try!(parse_label_sst(&r.data, &self.strings, &mut range)), // LabelSst
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
            self.sheets.insert(name, range);
        }
        Ok(())
    }
}

fn parse_sheet_name(r: &[u8]) -> Result<(usize, String)> {
    let pos = read_u32(r) as usize;
    let sheet = try!(parse_short_string(&r[6..]));
    Ok((pos, sheet))
}

fn parse_sst(r: &[u8]) -> Result<Vec<String>> {
    if r.len() < 8 {
        return Err("Invalid sst length".into());
    }
    let len = read_slice::<i32>(&r[4..]) as usize;
    let mut sst = Vec::with_capacity(len);
    let mut read = &mut &r[8..];
    for _ in 0..len {
        sst.push(try!(read_rich_extended_string(read)));
    }
    Ok(sst)
}

fn parse_number(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 14 {
        return Err("Invalid number length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = read_slice::<f64>(&r[6..]);
    range.set_value((row as u32, col as u32), DataType::Float(v))
}

fn parse_bool_err(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 8 {
        return Err("Invalid BoolErr length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = match r[7] {
        0x00 => DataType::Bool(r[6] != 0),
        0x01 => match r[6] {
            0x00 => DataType::Error(CellErrorType::Null),
            0x07 => DataType::Error(CellErrorType::Div0),
            0x0F => DataType::Error(CellErrorType::Value),
            0x17 => DataType::Error(CellErrorType::Ref),
            0x1D => DataType::Error(CellErrorType::Name),
            0x24 => DataType::Error(CellErrorType::Num),
            0x2A => DataType::Error(CellErrorType::NA),
            0x2B => DataType::Error(CellErrorType::GettingData),
            e => return Err(format!("Unrecognized error {:x}", e).into()),
        },
        e => return Err(format!("Unrecognized fError {:x}", e).into()),
    };
    range.set_value((row as u32, col as u32), v)
}

fn parse_rk(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid rk length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);

    let d100 = (r[6] & 1) != 0;
    let is_int = (r[6] & 2) != 0;

    let mut v = [0u8; 8];
    v[4..].copy_from_slice(&r[6..10]);
    v[0] &= 0xFC;
    let v = if is_int {
        let v = (read_slice::<i32>(&v[4..]) >> 2) as i64;
        DataType::Int( if d100 { v/100 } else { v })
    } else {
        let v = read_slice(&v);
        DataType::Float( if d100 { v/100.0 } else { v })
    };

    range.set_value((row as u32, col as u32), v)
}

fn parse_short_string(r: &[u8]) -> Result<String> {
    if r.len() < 2 {
        return Err("Invalid short string length".into());
    }
    let mut len = r[0] as usize;
    let high_byte = r[1];
    if high_byte != 0 { len *= 2; }
    if r.len() < 2 + len {
        return Err(format!("Invalid short string length: {} < {}", r.len(), 2 + len).into());
    }
    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; len * 2];
        for (i, sce) in r[2..2 + len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        len *= 2;
        Cow::Borrowed(&r[2..2 + len])
    };
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn parse_label_sst(r: &[u8], strings: &[String], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid short string length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let i = read_u32(&r[6..]) as usize;
    range.set_value((row as u32, col as u32), DataType::String(strings[i].clone()))
}

fn parse_dimensions(r: &[u8]) -> Result<Range> {
    if r.len() != 14 {
        return Err("Invalid dimensions lengths".into());
    }
    let rw_first = read_u32(&r[0..4]);
    let rw_last = read_u32(&r[4..8]);
    let col_first = read_u16(&r[8..10]);
    let col_last = read_u16(&r[10..12]);

    let start = (rw_first, col_first as u32);
    let size = ((rw_last - rw_first) as usize, (col_last - col_first) as usize);

    Ok(Range::new(start, size))
}

fn read_rich_extended_string(r: &mut &[u8]) -> Result<String> {
    if r.len() < 3 {
        return Err("Invalid rich extended string length".into());
    }

    let mut len = read_u16(r) as usize;
    let mut start = 3;
    let flags = r[2];
    let high_byte = flags & 0x1;
    let ext_st = flags & 0x4;
    let rich_st = flags & 0x8;
    let str_len = len;

    if high_byte != 0 { len += str_len; }
    if rich_st != 0 {
        len += 4 * read_u16(&r[start..]) as usize;
        start += 2;
    }
    if ext_st != 0 {
        len += read_slice::<i32>(&r[start..]) as usize;
        start += 4;
    }
    if r.len() < start + len {
        return Err(format!("Invalid rich extended string length: {} < {}",
                           r.len(), start + len).into());
    }

    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; str_len * 2];
        for (i, sce) in r[start..start + str_len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        Cow::Borrowed(&r[start..start + 2 * str_len])
    };

    *r = &r[start + len..];
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

struct Record<'a> {
    typ: u16,
    data: Cow<'a, [u8]>,
}

struct RecordIter<'a> {
    stream: &'a[u8],
}

impl<'a> Iterator for RecordIter<'a> {
    type Item=Result<Record<'a>>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.stream.len() < 4 {
            return if self.stream.is_empty() {
                None
            } else {
                Some(Err("Expecting record type and length, found end of stream".into()))
            }
        }
        let t = read_u16(self.stream);
        let mut len = read_u16(&self.stream[2..]) as usize;
        if self.stream.len() < len + 4 {
            return Some(Err("Expecting record length, found end of stream".into()));
        }
        let (mut data, mut next) = self.stream.split_at(len + 4);
        self.stream = next;

        // Append next record data if it is a Continue record
        let cow = if next.len() > 4 && read_u16(next) == 0x003C {
            let mut c = data[4..].to_vec();
            while next.len() > 4 && read_u16(next) == 0x003C {
                len = read_u16(&self.stream[2..]) as usize;
                if self.stream.len() < len + 4 {
                    return Some(Err("Expecting continue record length, found end of stream".into()));
                }
                let sp = self.stream.split_at(len + 4);
                data = sp.0;
                next = sp.1;
                c.extend_from_slice(&data[4..]);
                self.stream = next;
            }
            Cow::Owned(c)
        } else {
            Cow::Borrowed(&data[4..])
        };

        Some(Ok(Record { typ: t, data: cow }))
    }
}
<|endoftext|>


<commit_before>extern crate serde;
extern crate serde_json;

use hyper::Client;
use hyper::client::{Body, Response};
use hyper::header::Headers;
use hyper::status::StatusCode;
use jira_request::*;
use std::io::Read;

pub fn send<R: JIRARequest, T: serde::de::Deserialize>(request: R) -> Result<Vec<T>, String> {
    let client = Client::new();
    let url = format!("{}{}", &request.base_url(), &request.path());

    let mut res: Response = client.request(request.method(), &url)
        .headers(request.headers().unwrap_or(Headers::new()))
        .body(&request.body().unwrap_or("".to_string()))
        .send()
        .map_err(|e| e.to_string())?;

    let mut body = String::new();
    res.read_to_string(&mut body).map_err(|e| e.to_string())?;
    match res.status {
        StatusCode::Ok => serde_json::from_str::<Vec<T>>(&body).map_err(|e| e.to_string()),
        _ => Err(body)
    }
}<commit_msg>:art: Change the return value to the generic single type <commit_after>extern crate serde;
extern crate serde_json;

use hyper::Client;
use hyper::client::{Body, Response};
use hyper::header::Headers;
use hyper::status::StatusCode;
use jira_request::*;
use std::io::Read;

pub fn send<R: JIRARequest, T: serde::de::Deserialize>(request: R) -> Result<T, String> {
    let client = Client::new();
    let url = format!("{}{}", &request.base_url(), &request.path());

    let mut res: Response = client.request(request.method(), &url)
        .headers(request.headers().unwrap_or(Headers::new()))
        .body(&request.body().unwrap_or("".to_string()))
        .send()
        .map_err(|e| e.to_string())?;

    let mut body = String::new();
    res.read_to_string(&mut body).map_err(|e| e.to_string())?;
    match res.status {
        StatusCode::Ok => serde_json::from_str::<T>(&body).map_err(|e| e.to_string()),
        _ => Err(body)
    }
}<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The normal and derived distributions.

use crate::utils::ziggurat;
use num_traits::Float;
use crate::{ziggurat_tables, Distribution, Open01};
use rand::Rng;
use core::fmt;

/// Samples floating-point numbers according to the normal distribution
/// `N(0, 1)` (a.k.a. a standard normal, or Gaussian). This is equivalent to
/// `Normal::new(0.0, 1.0)` but faster.
///
/// See `Normal` for the general normal distribution.
///
/// Implemented via the ZIGNOR variant[^1] of the Ziggurat method.
///
/// [^1]: Jurgen A. Doornik (2005). [*An Improved Ziggurat Method to
///       Generate Normal Random Samples*](
///       https://www.doornik.com/research/ziggurat.pdf).
///       Nuffield College, Oxford
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand_distr::StandardNormal;
///
/// let val: f64 = thread_rng().sample(StandardNormal);
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct StandardNormal;

impl Distribution<f32> for StandardNormal {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        // TODO: use optimal 32-bit implementation
        let x: f64 = self.sample(rng);
        x as f32
    }
}

impl Distribution<f64> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x * x / 2.0).exp()
        }
        #[inline]
        fn zero_case<R: Rng + ?Sized>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            while -2.0 * y < x * x {
                let x_: f64 = rng.sample(Open01);
                let y_: f64 = rng.sample(Open01);

                x = x_.ln() / ziggurat_tables::ZIG_NORM_R;
                y = y_.ln();
            }

            if u < 0.0 {
                x - ziggurat_tables::ZIG_NORM_R
            } else {
                ziggurat_tables::ZIG_NORM_R - x
            }
        }

        ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F,
            pdf,
            zero_case,
        )
    }
}

/// The normal distribution `N(mean, std_dev**2)`.
///
/// This uses the ZIGNOR variant of the Ziggurat method, see [`StandardNormal`]
/// for more details.
///
/// Note that [`StandardNormal`] is an optimised implementation for mean 0, and
/// standard deviation 1.
///
/// # Example
///
/// ```
/// use rand_distr::{Normal, Distribution};
///
/// // mean 2, standard deviation 3
/// let normal = Normal::new(2.0, 3.0).unwrap();
/// let v = normal.sample(&mut rand::thread_rng());
/// println!("{} is from a N(2, 9) distribution", v)
/// ```
///
/// [`StandardNormal`]: crate::StandardNormal
#[derive(Clone, Copy, Debug)]
pub struct Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    mean: F,
    std_dev: F,
}

/// Error type returned from `Normal::new` and `LogNormal::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// `std_dev < 0` or `nan`.
    StdDevTooSmall,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Error::StdDevTooSmall => "std_dev < 0 or is NaN in normal distribution",
        })
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

impl<F> Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct a new `Normal` distribution with the given mean and
    /// standard deviation.
    #[inline]
    pub fn new(mean: F, std_dev: F) -> Result<Normal<F>, Error> {
        if !(std_dev >= F::zero()) {
            return Err(Error::StdDevTooSmall);
        }
        Ok(Normal { mean, std_dev })
    }
}

impl<F> Distribution<F> for Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        let n: F = rng.sample(StandardNormal);
        self.mean + self.std_dev * n
    }
}

/// The log-normal distribution `ln N(mean, std_dev**2)`.
///
/// If `X` is log-normal distributed, then `ln(X)` is `N(mean, std_dev**2)`
/// distributed.
///
/// # Example
///
/// ```
/// use rand_distr::{LogNormal, Distribution};
///
/// // mean 2, standard deviation 3
/// let log_normal = LogNormal::new(2.0, 3.0).unwrap();
/// let v = log_normal.sample(&mut rand::thread_rng());
/// println!("{} is from an ln N(2, 9) distribution", v)
/// ```
#[derive(Clone, Copy, Debug)]
pub struct LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    norm: Normal<F>,
}

impl<F> LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct a new `LogNormal` distribution with the given mean
    /// and standard deviation of the logarithm of the distribution.
    #[inline]
    pub fn new(mean: F, std_dev: F) -> Result<LogNormal<F>, Error> {
        if !(std_dev >= F::zero()) {
            return Err(Error::StdDevTooSmall);
        }
        Ok(LogNormal {
            norm: Normal::new(mean, std_dev).unwrap(),
        })
    }
}

impl<F> Distribution<F> for LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.norm.sample(rng).exp()
    }
}

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

    #[test]
    fn test_normal() {
        let norm = Normal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(210);
        for _ in 0..1000 {
            norm.sample(&mut rng);
        }
    }
    #[test]
    #[should_panic]
    fn test_normal_invalid_sd() {
        Normal::new(10.0, -1.0).unwrap();
    }

    #[test]
    fn test_log_normal() {
        let lnorm = LogNormal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(211);
        for _ in 0..1000 {
            lnorm.sample(&mut rng);
        }
    }
    #[test]
    #[should_panic]
    fn test_log_normal_invalid_sd() {
        LogNormal::new(10.0, -1.0).unwrap();
    }
}
<commit_msg>rand_distr: add from_mean_cv ctor to Normal and LogNormal<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The normal and derived distributions.

use crate::utils::ziggurat;
use num_traits::Float;
use crate::{ziggurat_tables, Distribution, Open01};
use rand::Rng;
use core::fmt;

/// Samples floating-point numbers according to the normal distribution
/// `N(0, 1)` (a.k.a. a standard normal, or Gaussian). This is equivalent to
/// `Normal::new(0.0, 1.0)` but faster.
///
/// See `Normal` for the general normal distribution.
///
/// Implemented via the ZIGNOR variant[^1] of the Ziggurat method.
///
/// [^1]: Jurgen A. Doornik (2005). [*An Improved Ziggurat Method to
///       Generate Normal Random Samples*](
///       https://www.doornik.com/research/ziggurat.pdf).
///       Nuffield College, Oxford
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand_distr::StandardNormal;
///
/// let val: f64 = thread_rng().sample(StandardNormal);
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct StandardNormal;

impl Distribution<f32> for StandardNormal {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        // TODO: use optimal 32-bit implementation
        let x: f64 = self.sample(rng);
        x as f32
    }
}

impl Distribution<f64> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x * x / 2.0).exp()
        }
        #[inline]
        fn zero_case<R: Rng + ?Sized>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            while -2.0 * y < x * x {
                let x_: f64 = rng.sample(Open01);
                let y_: f64 = rng.sample(Open01);

                x = x_.ln() / ziggurat_tables::ZIG_NORM_R;
                y = y_.ln();
            }

            if u < 0.0 {
                x - ziggurat_tables::ZIG_NORM_R
            } else {
                ziggurat_tables::ZIG_NORM_R - x
            }
        }

        ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F,
            pdf,
            zero_case,
        )
    }
}

/// The normal distribution `N(mean, std_dev**2)`.
///
/// This uses the ZIGNOR variant of the Ziggurat method, see [`StandardNormal`]
/// for more details.
///
/// Note that [`StandardNormal`] is an optimised implementation for mean 0, and
/// standard deviation 1.
///
/// # Example
///
/// ```
/// use rand_distr::{Normal, Distribution};
///
/// // mean 2, standard deviation 3
/// let normal = Normal::new(2.0, 3.0).unwrap();
/// let v = normal.sample(&mut rand::thread_rng());
/// println!("{} is from a N(2, 9) distribution", v)
/// ```
///
/// [`StandardNormal`]: crate::StandardNormal
#[derive(Clone, Copy, Debug)]
pub struct Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    mean: F,
    std_dev: F,
}

/// Error type returned from `Normal::new` and `LogNormal::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// The mean value is too small (log-normal samples must be positive)
    MeanTooSmall,
    /// The standard deviation or other dispersion parameter is negative,
    /// infinite, or not-a-number.
    BadVariance,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Error::MeanTooSmall => "mean < 0 or NaN in log-normal distribution",
            Error::BadVariance => "variation parameter is negative or non-finite in (log)normal distribution",
        })
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

impl<F> Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from mean and standard deviation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   standard deviation (`σ ≥ 0`)
    #[inline]
    pub fn new(mean: F, std_dev: F) -> Result<Normal<F>, Error> {
        if !std_dev.is_finite() || std_dev < F::zero() {
            return Err(Error::BadVariance);
        }
        Ok(Normal { mean, std_dev })
    }

    /// Construct, from mean and coefficient of variation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   coefficient of variation (`cv = σ / μ` where `cv ≥ 0`)
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<Normal<F>, Error> {
        if !cv.is_finite() || cv < F::zero() {
            return Err(Error::BadVariance);
        }
        let std_dev = cv * mean;
        Ok(Normal { mean, std_dev })
    }
}

impl<F> Distribution<F> for Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        let n: F = rng.sample(StandardNormal);
        self.mean + self.std_dev * n
    }
}

/// The log-normal distribution `ln N(mean, std_dev**2)`.
///
/// If `X` is log-normal distributed, then `ln(X)` is `N(mean, std_dev**2)`
/// distributed.
///
/// # Example
///
/// ```
/// use rand_distr::{LogNormal, Distribution};
///
/// // mean 2, standard deviation 3
/// let log_normal = LogNormal::new(2.0, 3.0).unwrap();
/// let v = log_normal.sample(&mut rand::thread_rng());
/// println!("{} is from an ln N(2, 9) distribution", v)
/// ```
#[derive(Clone, Copy, Debug)]
pub struct LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    norm: Normal<F>,
}

impl<F> LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from (log-space) mean and standard deviation
    ///
    /// Parameters are the "standard" log-space measures (these are the mean
    /// and standard deviation of the logarithm of samples):
    ///
    /// -   `mu` (`μ`, unrestricted) is the mean of the underlying distribution
    /// -   `sigma` (`σ ≥ 0`) is the standard deviation of the
    ///     underlying Normal distribution
    #[inline]
    pub fn new(mu: F, sigma: F) -> Result<LogNormal<F>, Error> {
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }

    /// Construct, from (linear-space) mean and coefficient of variation
    ///
    /// Parameters are linear-space measures:
    ///
    /// -   mean (`μ > 0`) is the (real) mean of the distribution
    /// -   coefficient of variation (`cv = σ / μ`, requiring `cv ≥ 0`) is a
    ///     standardized measure of dispersion
    ///
    /// As a special exception, `μ = 0, cv = 0` is allowed (samples are `-inf`).
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<LogNormal<F>, Error> {
        if cv == F::zero() {
            let mu = mean.ln();
            let norm = Normal::new(mu, F::zero()).unwrap();
            return Ok(LogNormal { norm });
        }
        if !(mean > F::zero()) {
            return Err(Error::MeanTooSmall);
        }
        if !(cv >= F::zero()) {
            return Err(Error::BadVariance);
        }

        // Using X ~ lognormal(μ, σ), CV² = Var(X) / E(X)²
        // E(X) = exp(μ + σ² / 2) = exp(μ) × exp(σ² / 2)
        // Var(X) = exp(2μ + σ²)(exp(σ²) - 1) = E(X)² × (exp(σ²) - 1)
        // but Var(X) = (CV × E(X))² so CV² = exp(σ²) - 1
        // thus σ² = log(CV² + 1)
        // and exp(μ) = E(X) / exp(σ² / 2) = E(X) / sqrt(CV² + 1)
        let a = F::one() + cv * cv; // e
        let mu = (mean / a.sqrt()).ln();
        let sigma = a.ln().sqrt();
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }
}

impl<F> Distribution<F> for LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.norm.sample(rng).exp()
    }
}

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

    #[test]
    fn test_normal() {
        let norm = Normal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(210);
        for _ in 0..1000 {
            norm.sample(&mut rng);
        }
    }
    #[test]
    fn test_normal_cv() {
        let norm = Normal::from_mean_cv(1024.0, 1.0 / 256.0).unwrap();
        assert_eq!((norm.mean, norm.std_dev), (1024.0, 4.0));
    }
    #[test]
    fn test_normal_invalid_sd() {
        assert!(Normal::new(10.0, -1.0).is_err());
        assert!(Normal::from_mean_cv(10.0, -1.0).is_err());
    }

    #[test]
    fn test_log_normal() {
        let lnorm = LogNormal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(211);
        for _ in 0..1000 {
            lnorm.sample(&mut rng);
        }
    }
    #[test]
    fn test_log_normal_cv() {
        let lnorm = LogNormal::from_mean_cv(0.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (-std::f64::INFINITY, 0.0));

        let lnorm = LogNormal::from_mean_cv(1.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (0.0, 0.0));

        let e = std::f64::consts::E;
        let lnorm = LogNormal::from_mean_cv(e.sqrt(), (e - 1.0).sqrt()).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (0.0, 1.0));

        let lnorm = LogNormal::from_mean_cv(e.powf(1.5), (e - 1.0).sqrt()).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (1.0, 1.0));
    }
    #[test]
    fn test_log_normal_invalid_sd() {
        assert!(LogNormal::from_mean_cv(-1.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(0.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(1.0, -1.0).is_err());
    }
}
<|endoftext|>