<commit_before>use Module;
use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Gain {
    LV,
    HV,
}

impl ::std::convert::Into<String> for Gain {
    fn into(self) -> String {
        let s = match self {
            Gain::LV => "LV",
            Gain::HV => "HV",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Gain {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "LV" => Ok(Gain::LV),
            "HV" => Ok(Gain::HV),
            gain => Err(format!("Unknow gain '{}'", gain)),
        }
    }
}

impl ::std::fmt::Display for Gain {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Gain::LV => "LV",
            &Gain::HV => "HV",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    IN1,
    IN2,
}

impl ::std::convert::Into<String> for Source {
    fn into(self) -> String {
        let s = match self {
            Source::IN1 => "SOUR1",
            Source::IN2 => "SOUR2",
        };

        String::from(s)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Decimation {
    DEC_1,
    DEC_8,
    DEC_64,
    DEC_1024,
    DEC_8192,
    DEC_65536,
}

impl ::std::convert::Into<String> for Decimation {
    fn into(self) -> String {
        let s = match self {
            Decimation::DEC_1 => "1",
            Decimation::DEC_8 => "8",
            Decimation::DEC_64 => "64",
            Decimation::DEC_1024 => "1024",
            Decimation::DEC_8192 => "8192",
            Decimation::DEC_65536 => "65536",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Decimation {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "1" => Ok(Decimation::DEC_1),
            "8" => Ok(Decimation::DEC_8),
            "64" => Ok(Decimation::DEC_64),
            "1024" => Ok(Decimation::DEC_1024),
            "8192" => Ok(Decimation::DEC_8192),
            "65536" => Ok(Decimation::DEC_65536),
            decimation => Err(format!("Unknow decimation '{}'", decimation)),
        }
    }
}

impl ::std::convert::Into<SamplingRate> for Decimation {
    fn into(self) -> SamplingRate {
        match self {
            Decimation::DEC_1 => SamplingRate::RATE_125MHz,
            Decimation::DEC_8 => SamplingRate::RATE_15_6MHz,
            Decimation::DEC_64 => SamplingRate::RATE_1_9MHz,
            Decimation::DEC_1024 => SamplingRate::RATE_103_8kHz,
            Decimation::DEC_8192 => SamplingRate::RATE_15_2kHz,
            Decimation::DEC_65536 => SamplingRate::RATE_1_9kHz,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SamplingRate {
    RATE_125MHz,
    RATE_15_6MHz,
    RATE_1_9MHz,
    RATE_103_8kHz,
    RATE_15_2kHz,
    RATE_1_9kHz,
}

impl SamplingRate {
    pub fn get_buffer_duration(self) -> ::std::time::Duration {
        let (s, ns) = match self {
            SamplingRate::RATE_125MHz => (0, 131_072),
            SamplingRate::RATE_15_6MHz => (0, 1_049_000),
            SamplingRate::RATE_1_9MHz => (0, 8_389_000),
            SamplingRate::RATE_103_8kHz => (0, 134_218_000),
            SamplingRate::RATE_15_2kHz => (1, 740_000_000),
            SamplingRate::RATE_1_9kHz => (8, 590_000_000),
        };

        ::std::time::Duration::new(s, ns)
    }
}

impl ::std::fmt::Display for SamplingRate {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &SamplingRate::RATE_125MHz => "125 MHz",
            &SamplingRate::RATE_15_6MHz => "15.6 MHz",
            &SamplingRate::RATE_1_9MHz => "1.9 MHz",
            &SamplingRate::RATE_103_8kHz => "103.8 kHz",
            &SamplingRate::RATE_15_2kHz => "15.2 kHz",
            &SamplingRate::RATE_1_9kHz => "1.9 kHz",
        };

        write!(f, "{}", display)
    }
}

impl ::std::convert::Into<Decimation> for SamplingRate {
    fn into(self) -> Decimation {
        match self {
            SamplingRate::RATE_125MHz => Decimation::DEC_1,
            SamplingRate::RATE_15_6MHz => Decimation::DEC_8,
            SamplingRate::RATE_1_9MHz => Decimation::DEC_64,
            SamplingRate::RATE_103_8kHz => Decimation::DEC_1024,
            SamplingRate::RATE_15_2kHz => Decimation::DEC_8192,
            SamplingRate::RATE_1_9kHz => Decimation::DEC_65536,
        }
    }
}

impl ::std::convert::Into<String> for SamplingRate {
    fn into(self) -> String {
        let s = match self {
            SamplingRate::RATE_125MHz => "125MHz",
            SamplingRate::RATE_15_6MHz => "15_6MHz",
            SamplingRate::RATE_1_9MHz => "1_9MHz",
            SamplingRate::RATE_103_8kHz => "103_8kHz",
            SamplingRate::RATE_15_2kHz => "15_2kHz",
            SamplingRate::RATE_1_9kHz => "1_9kHz",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for SamplingRate {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "125000000 Hz" => Ok(SamplingRate::RATE_125MHz),
            "15600000 Hz" => Ok(SamplingRate::RATE_15_6MHz),
            "1900000 Hz" => Ok(SamplingRate::RATE_1_9MHz),
            "103800 Hz" => Ok(SamplingRate::RATE_103_8kHz),
            "15200 Hz" => Ok(SamplingRate::RATE_15_2kHz),
            "1900 Hz" => Ok(SamplingRate::RATE_1_9kHz),
            rate => Err(format!("Unknow sampling rate {}", rate)),
        }
    }
}

#[derive(Clone)]
pub struct Acquire {
    socket: ::std::cell::RefCell<Socket>,
    started: bool,
}

impl ::Module for Acquire {
    fn get_socket<'a>(&'a self) -> ::std::cell::RefMut<'a, ::socket::Socket> {
        self.socket.borrow_mut()
    }
}

impl Acquire {
    pub fn new(socket: Socket) -> Self {
        Acquire {
            socket: ::std::cell::RefCell::new(socket),
            started: false,
        }
    }

    /**
     * Starts acquisition.
     */
    pub fn start(&mut self) {
        self.send("ACQ:START");
        self.started = true;
    }

    /**
     * Stops acquisition.
     */
    pub fn stop(&mut self) {
        self.send("ACQ:STOP");
        self.started = false;
    }

    pub fn is_started(&self) -> bool {
        self.started
    }

    /**
     * Stops acquisition and sets all parameters to default values.
     */
    pub fn reset(&self) {
        self.send("ACQ:RST");
    }

    /**
     * Set decimation factor.
     */
    pub fn set_decimation(&self, decimation: Decimation) {
        self.send(format!("ACQ:DEC {}", Into::<String>::into(decimation)));
    }

    /**
     * Get decimation factor.
     */
    pub fn get_decimation(&self) -> Result<Decimation, String> {
        self.send("ACQ:DEC?");

        self.receive()
            .parse()
    }

    /**
     * Get sampling rate.
     *
     * # Panics
     *
     * Calling this command makes buffer overflow.
     * See https://github.com/RedPitaya/RedPitaya/pull/110
     */
    pub fn get_sampling_rate(&self) -> Result<SamplingRate, String> {
        self.send("ACQ:SRAT?");

        self.receive()
            .parse()
    }

    /**
     * Enable averaging.
     */
    pub fn enable_average(&self) {
        self.send("ACQ:AVG ON");
    }

    /**
     * Disable averaging.
     */
    pub fn disable_average(&self) {
        self.send("ACQ:AVG OFF");
    }

    /**
     * Get averaging status.
     */
    pub fn is_average_enabled(&self) -> bool {
        self.send("ACQ:AVG?");

        let message = self.receive();

        match message.as_str() {
            "ON" => true,
            _ => false,
        }
    }

    /**
     * Set gain settings to HIGH or LOW.
     *
     * This gain is referring to jumper settings on Red Pitaya fast analog inputs.
     */
    pub fn set_gain(&self, source: Source, gain: Gain) {
        self.send(format!("ACQ:{}:GAIN {}", Into::<String>::into(source), Into::<String>::into(gain)));
    }

    /**
     * Get gain settings to HIGH or LOW.
     */
    pub fn get_gain(&self, source: Source) -> Result<Gain, String> {
        self.send(format!("ACQ:{}:GAIN?", Into::<String>::into(source)));

        self.receive()
            .parse()
    }
}

#[cfg(test)]
mod test {
    macro_rules! acquire_assert {
        ($f:ident, $e:expr) => {
            let (rx, acquire) = create_acquire();

            acquire.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_sampling_rate_get_buffer_duration() {
        let duration = ::std::time::Duration::new(8, 590_000_000);

        assert_eq!(duration, ::acquire::SamplingRate::RATE_1_9kHz.get_buffer_duration());
    }

    #[test]
    fn test_start() {
        let (rx, mut acquire) = create_acquire();

        acquire.start();
        assert_eq!("ACQ:START\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut acquire) = create_acquire();

        acquire.stop();
        assert_eq!("ACQ:STOP\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_is_started() {
        let (_, mut acquire) = create_acquire();

        assert_eq!(acquire.is_started(), false);
        acquire.start();
        assert_eq!(acquire.is_started(), true);
        acquire.stop();
        assert_eq!(acquire.is_started(), false);
    }

    #[test]
    fn test_reset() {
        acquire_assert!(reset, "ACQ:RST\r\n");
    }

    #[test]
    fn test_set_decimation() {
        let (rx, acquire) = create_acquire();

        acquire.set_decimation(::acquire::Decimation::DEC_8);
        assert_eq!("ACQ:DEC 8\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_decimation() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_decimation(), Ok(::acquire::Decimation::DEC_1));
    }

    #[test]
    fn test_get_sampling_rate() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_sampling_rate(), Ok(::acquire::SamplingRate::RATE_125MHz));
    }

    #[test]
    fn test_enable_average() {
        acquire_assert!(enable_average, "ACQ:AVG ON\r\n");
    }

    #[test]
    fn test_disable_average() {
        acquire_assert!(disable_average, "ACQ:AVG OFF\r\n");
    }

    #[test]
    fn test_is_average_enabled() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.is_average_enabled(), true);
    }

    #[test]
    fn test_set_gain() {
        let (rx, acquire) = create_acquire();

        acquire.set_gain(::acquire::Source::IN1, ::acquire::Gain::LV);
        assert_eq!("ACQ:SOUR1:GAIN LV\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_gain() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_gain(::acquire::Source::IN1), Ok(::acquire::Gain::HV));
    }

    fn create_acquire() -> (::std::sync::mpsc::Receiver<String>, ::acquire::Acquire) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::acquire::Acquire::new(socket))
    }
}
<commit_msg>[acquire] Remove internal stat<commit_after>use Module;
use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Gain {
    LV,
    HV,
}

impl ::std::convert::Into<String> for Gain {
    fn into(self) -> String {
        let s = match self {
            Gain::LV => "LV",
            Gain::HV => "HV",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Gain {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "LV" => Ok(Gain::LV),
            "HV" => Ok(Gain::HV),
            gain => Err(format!("Unknow gain '{}'", gain)),
        }
    }
}

impl ::std::fmt::Display for Gain {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Gain::LV => "LV",
            &Gain::HV => "HV",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    IN1,
    IN2,
}

impl ::std::convert::Into<String> for Source {
    fn into(self) -> String {
        let s = match self {
            Source::IN1 => "SOUR1",
            Source::IN2 => "SOUR2",
        };

        String::from(s)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Decimation {
    DEC_1,
    DEC_8,
    DEC_64,
    DEC_1024,
    DEC_8192,
    DEC_65536,
}

impl ::std::convert::Into<String> for Decimation {
    fn into(self) -> String {
        let s = match self {
            Decimation::DEC_1 => "1",
            Decimation::DEC_8 => "8",
            Decimation::DEC_64 => "64",
            Decimation::DEC_1024 => "1024",
            Decimation::DEC_8192 => "8192",
            Decimation::DEC_65536 => "65536",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for Decimation {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "1" => Ok(Decimation::DEC_1),
            "8" => Ok(Decimation::DEC_8),
            "64" => Ok(Decimation::DEC_64),
            "1024" => Ok(Decimation::DEC_1024),
            "8192" => Ok(Decimation::DEC_8192),
            "65536" => Ok(Decimation::DEC_65536),
            decimation => Err(format!("Unknow decimation '{}'", decimation)),
        }
    }
}

impl ::std::convert::Into<SamplingRate> for Decimation {
    fn into(self) -> SamplingRate {
        match self {
            Decimation::DEC_1 => SamplingRate::RATE_125MHz,
            Decimation::DEC_8 => SamplingRate::RATE_15_6MHz,
            Decimation::DEC_64 => SamplingRate::RATE_1_9MHz,
            Decimation::DEC_1024 => SamplingRate::RATE_103_8kHz,
            Decimation::DEC_8192 => SamplingRate::RATE_15_2kHz,
            Decimation::DEC_65536 => SamplingRate::RATE_1_9kHz,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SamplingRate {
    RATE_125MHz,
    RATE_15_6MHz,
    RATE_1_9MHz,
    RATE_103_8kHz,
    RATE_15_2kHz,
    RATE_1_9kHz,
}

impl SamplingRate {
    pub fn get_buffer_duration(self) -> ::std::time::Duration {
        let (s, ns) = match self {
            SamplingRate::RATE_125MHz => (0, 131_072),
            SamplingRate::RATE_15_6MHz => (0, 1_049_000),
            SamplingRate::RATE_1_9MHz => (0, 8_389_000),
            SamplingRate::RATE_103_8kHz => (0, 134_218_000),
            SamplingRate::RATE_15_2kHz => (1, 740_000_000),
            SamplingRate::RATE_1_9kHz => (8, 590_000_000),
        };

        ::std::time::Duration::new(s, ns)
    }
}

impl ::std::fmt::Display for SamplingRate {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &SamplingRate::RATE_125MHz => "125 MHz",
            &SamplingRate::RATE_15_6MHz => "15.6 MHz",
            &SamplingRate::RATE_1_9MHz => "1.9 MHz",
            &SamplingRate::RATE_103_8kHz => "103.8 kHz",
            &SamplingRate::RATE_15_2kHz => "15.2 kHz",
            &SamplingRate::RATE_1_9kHz => "1.9 kHz",
        };

        write!(f, "{}", display)
    }
}

impl ::std::convert::Into<Decimation> for SamplingRate {
    fn into(self) -> Decimation {
        match self {
            SamplingRate::RATE_125MHz => Decimation::DEC_1,
            SamplingRate::RATE_15_6MHz => Decimation::DEC_8,
            SamplingRate::RATE_1_9MHz => Decimation::DEC_64,
            SamplingRate::RATE_103_8kHz => Decimation::DEC_1024,
            SamplingRate::RATE_15_2kHz => Decimation::DEC_8192,
            SamplingRate::RATE_1_9kHz => Decimation::DEC_65536,
        }
    }
}

impl ::std::convert::Into<String> for SamplingRate {
    fn into(self) -> String {
        let s = match self {
            SamplingRate::RATE_125MHz => "125MHz",
            SamplingRate::RATE_15_6MHz => "15_6MHz",
            SamplingRate::RATE_1_9MHz => "1_9MHz",
            SamplingRate::RATE_103_8kHz => "103_8kHz",
            SamplingRate::RATE_15_2kHz => "15_2kHz",
            SamplingRate::RATE_1_9kHz => "1_9kHz",
        };

        String::from(s)
    }
}

impl ::std::str::FromStr for SamplingRate {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "125000000 Hz" => Ok(SamplingRate::RATE_125MHz),
            "15600000 Hz" => Ok(SamplingRate::RATE_15_6MHz),
            "1900000 Hz" => Ok(SamplingRate::RATE_1_9MHz),
            "103800 Hz" => Ok(SamplingRate::RATE_103_8kHz),
            "15200 Hz" => Ok(SamplingRate::RATE_15_2kHz),
            "1900 Hz" => Ok(SamplingRate::RATE_1_9kHz),
            rate => Err(format!("Unknow sampling rate {}", rate)),
        }
    }
}

#[derive(Clone)]
pub struct Acquire {
    socket: ::std::cell::RefCell<Socket>,
}

impl ::Module for Acquire {
    fn get_socket<'a>(&'a self) -> ::std::cell::RefMut<'a, ::socket::Socket> {
        self.socket.borrow_mut()
    }
}

impl Acquire {
    pub fn new(socket: Socket) -> Self {
        Acquire {
            socket: ::std::cell::RefCell::new(socket),
        }
    }

    /**
     * Starts acquisition.
     */
    pub fn start(&mut self) {
        self.send("ACQ:START");
    }

    /**
     * Stops acquisition.
     */
    pub fn stop(&mut self) {
        self.send("ACQ:STOP");
    }

    /**
     * Stops acquisition and sets all parameters to default values.
     */
    pub fn reset(&self) {
        self.send("ACQ:RST");
    }

    /**
     * Set decimation factor.
     */
    pub fn set_decimation(&self, decimation: Decimation) {
        self.send(format!("ACQ:DEC {}", Into::<String>::into(decimation)));
    }

    /**
     * Get decimation factor.
     */
    pub fn get_decimation(&self) -> Result<Decimation, String> {
        self.send("ACQ:DEC?");

        self.receive()
            .parse()
    }

    /**
     * Get sampling rate.
     *
     * # Panics
     *
     * Calling this command makes buffer overflow.
     * See https://github.com/RedPitaya/RedPitaya/pull/110
     */
    pub fn get_sampling_rate(&self) -> Result<SamplingRate, String> {
        self.send("ACQ:SRAT?");

        self.receive()
            .parse()
    }

    /**
     * Enable averaging.
     */
    pub fn enable_average(&self) {
        self.send("ACQ:AVG ON");
    }

    /**
     * Disable averaging.
     */
    pub fn disable_average(&self) {
        self.send("ACQ:AVG OFF");
    }

    /**
     * Get averaging status.
     */
    pub fn is_average_enabled(&self) -> bool {
        self.send("ACQ:AVG?");

        let message = self.receive();

        match message.as_str() {
            "ON" => true,
            _ => false,
        }
    }

    /**
     * Set gain settings to HIGH or LOW.
     *
     * This gain is referring to jumper settings on Red Pitaya fast analog inputs.
     */
    pub fn set_gain(&self, source: Source, gain: Gain) {
        self.send(format!("ACQ:{}:GAIN {}", Into::<String>::into(source), Into::<String>::into(gain)));
    }

    /**
     * Get gain settings to HIGH or LOW.
     */
    pub fn get_gain(&self, source: Source) -> Result<Gain, String> {
        self.send(format!("ACQ:{}:GAIN?", Into::<String>::into(source)));

        self.receive()
            .parse()
    }
}

#[cfg(test)]
mod test {
    macro_rules! acquire_assert {
        ($f:ident, $e:expr) => {
            let (rx, acquire) = create_acquire();

            acquire.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_sampling_rate_get_buffer_duration() {
        let duration = ::std::time::Duration::new(8, 590_000_000);

        assert_eq!(duration, ::acquire::SamplingRate::RATE_1_9kHz.get_buffer_duration());
    }

    #[test]
    fn test_start() {
        let (rx, mut acquire) = create_acquire();

        acquire.start();
        assert_eq!("ACQ:START\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut acquire) = create_acquire();

        acquire.stop();
        assert_eq!("ACQ:STOP\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_reset() {
        acquire_assert!(reset, "ACQ:RST\r\n");
    }

    #[test]
    fn test_set_decimation() {
        let (rx, acquire) = create_acquire();

        acquire.set_decimation(::acquire::Decimation::DEC_8);
        assert_eq!("ACQ:DEC 8\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_decimation() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_decimation(), Ok(::acquire::Decimation::DEC_1));
    }

    #[test]
    fn test_get_sampling_rate() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_sampling_rate(), Ok(::acquire::SamplingRate::RATE_125MHz));
    }

    #[test]
    fn test_enable_average() {
        acquire_assert!(enable_average, "ACQ:AVG ON\r\n");
    }

    #[test]
    fn test_disable_average() {
        acquire_assert!(disable_average, "ACQ:AVG OFF\r\n");
    }

    #[test]
    fn test_is_average_enabled() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.is_average_enabled(), true);
    }

    #[test]
    fn test_set_gain() {
        let (rx, acquire) = create_acquire();

        acquire.set_gain(::acquire::Source::IN1, ::acquire::Gain::LV);
        assert_eq!("ACQ:SOUR1:GAIN LV\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_gain() {
        let (_, acquire) = create_acquire();

        assert_eq!(acquire.get_gain(::acquire::Source::IN1), Ok(::acquire::Gain::HV));
    }

    fn create_acquire() -> (::std::sync::mpsc::Receiver<String>, ::acquire::Acquire) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::acquire::Acquire::new(socket))
    }
}
<|endoftext|>


<commit_before>// functions to ask the user for data, with crate:spinner

use std::io::stdin;

use regex::Regex;
use ansi_term::Colour::*;

/// Ask the user for a Yes/No answer. Optionally provide a default value. If none is provided, this
/// keeps loop{}ing
pub fn ask_bool(s: &str, default: Option<bool>) -> bool {
    lazy_static! {
        static ref R_YES: Regex = Regex::new(r"^[Yy]$").unwrap();
        static ref R_NO: Regex  = Regex::new(r"^[Nn]$").unwrap();
    }

    loop {
        ask_question(s, false);
        if match default { Some(s) => s, _ => true } {
            println!(" [Yn]: ");
        } else {
            println!(" [yN]: ");
        }

        let mut s = String::new();
        let _     = stdin().read_line(&mut s);

        if R_YES.is_match(&s[..]) {
            return true
        } else if R_NO.is_match(&s[..]) {
            return false
        } else {
            if default.is_some() {
                return default.unwrap();
            }

            // else again...
        }
    }
}

pub fn ask_uint(s: &str) -> u64 {
    unimplemented!()
}

pub fn ask_string(s: &str) -> String {
    unimplemented!()
}

pub fn ask_enum<E: From<String>>(s: &str) -> E {
    unimplemented!()
}

/// Helper function to print a imag question string. The `question` argument may not contain a
/// trailing questionmark.
///
/// The `nl` parameter can be used to configure whether a newline character should be printed
pub fn ask_question(question: &str, nl: bool) {
    if nl {
        println!("[imag]: {}?", Yellow.paint(question));
    } else {
        print!("[imag]: {}?", Yellow.paint(question));
    }
}

<commit_msg>Move real implementation to helper, so we can test.<commit_after>// functions to ask the user for data, with crate:spinner

use std::io::stdin;
use std::io::BufRead;
use std::io::BufReader;

use regex::Regex;
use ansi_term::Colour::*;

/// Ask the user for a Yes/No answer. Optionally provide a default value. If none is provided, this
/// keeps loop{}ing
pub fn ask_bool(s: &str, default: Option<bool>) -> bool {
    ask_bool_(s, default, &mut BufReader::new(stdin()))
}

fn ask_bool_<R: BufRead>(s: &str, default: Option<bool>, input: &mut R) -> bool {
    lazy_static! {
        static ref R_YES: Regex = Regex::new(r"^[Yy]$").unwrap();
        static ref R_NO: Regex  = Regex::new(r"^[Nn]$").unwrap();
    }

    loop {
        ask_question(s, false);
        if match default { Some(s) => s, _ => true } {
            println!(" [Yn]: ");
        } else {
            println!(" [yN]: ");
        }

        let mut s = String::new();
        let _     = input.read_line(&mut s);

        if R_YES.is_match(&s[..]) {
            return true
        } else if R_NO.is_match(&s[..]) {
            return false
        } else {
            if default.is_some() {
                return default.unwrap();
            }

            // else again...
        }
    }
}

pub fn ask_uint(s: &str) -> u64 {
    unimplemented!()
}

pub fn ask_string(s: &str) -> String {
    unimplemented!()
}

pub fn ask_enum<E: From<String>>(s: &str) -> E {
    unimplemented!()
}

/// Helper function to print a imag question string. The `question` argument may not contain a
/// trailing questionmark.
///
/// The `nl` parameter can be used to configure whether a newline character should be printed
pub fn ask_question(question: &str, nl: bool) {
    if nl {
        println!("[imag]: {}?", Yellow.paint(question));
    } else {
        print!("[imag]: {}?", Yellow.paint(question));
    }
}

#[cfg(test)]
mod test {
    use std::io::BufReader;

    use super::ask_bool_;

    #[test]
    fn test_ask_bool_nodefault_yes() {
        let question = "Is this true";
        let default  = None;
        let answers  = "\n\n\n\n\ny";

        assert!(ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_nodefault_no() {
        let question = "Is this true";
        let default  = None;
        let answers  = "n";

        assert!(false == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_no() {
        let question = "Is this true";
        let default  = Some(false);
        let answers  = "n";

        assert!(false == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_yes() {
        let question = "Is this true";
        let default  = Some(true);
        let answers  = "y";

        assert!(true == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_yes_answer_no() {
        let question = "Is this true";
        let default  = Some(true);
        let answers  = "n";

        assert!(false == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_no_answer_yes() {
        let question = "Is this true";
        let default  = Some(false);
        let answers  = "y";

        assert!(true == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_no_without_answer() {
        let question = "Is this true";
        let default  = Some(false);
        let answers  = "\n";

        assert!(false == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

    #[test]
    fn test_ask_bool_default_yes_without_answer() {
        let question = "Is this true";
        let default  = Some(true);
        let answers  = "\n";

        assert!(true == ask_bool_(question, default, &mut BufReader::new(answers.as_bytes())));
    }

}
<|endoftext|>


<commit_before>use rustc_serialize::{Encodable, Decodable};
use std;
use nix::unistd::gethostname;
use super::master::{self, Master, MasterResult};
use super::slave::Slave;
use super::error::ServerError;
use super::value::Topic;
use super::naming::Resolver;
use tcpros::{Client, Message, PublisherStream, ServicePair};
use rosxmlrpc::serde::XmlRpcValue;

pub struct Ros {
    master: Master,
    slave: Slave,
    hostname: String,
    resolver: Resolver,
    name: String,
}

impl Ros {
    pub fn new(name: &str) -> Result<Ros, ServerError> {
        let namespace = std::env::var("ROS_NAMESPACE").unwrap_or(String::from(""));
        Ros::new_raw(&namespace, name)
    }

    pub fn new_raw(namespace: &str, name: &str) -> Result<Ros, ServerError> {
        let master_uri = std::env::var("ROS_MASTER_URI")
            .unwrap_or(String::from("http://localhost:11311/"));

        let mut hostname = [0u8; 50];
        gethostname(&mut hostname)?;
        let hostname =
            hostname.into_iter().take_while(|&v| *v != 0u8).map(|v| *v).collect::<Vec<_>>();
        let hostname = String::from_utf8(hostname)?;

        let name = format!("{}/{}", namespace, name);
        let resolver = Resolver::new(&name)?;

        let slave = Slave::new(&master_uri, &hostname, 0, &name)?;
        let master = Master::new(&master_uri, &name, &slave.uri());
        Ok(Ros {
            master: master,
            slave: slave,
            hostname: hostname,
            resolver: resolver,
            name: name,
        })
    }

    pub fn map(&mut self, source: &str, destination: &str) -> Result<(), ServerError> {
        self.resolver.map(source, destination).map_err(|v| ServerError::Naming(v))
    }

    pub fn node_uri(&self) -> &str {
        return self.slave.uri();
    }

    pub fn param<'a, 'b>(&'a self, name: &'b str) -> Option<Parameter<'a>> {
        self.resolver.translate(name).ok().map(|v| {
            Parameter {
                master: &self.master,
                name: v,
            }
        })
    }

    pub fn parameters(&self) -> MasterResult<Vec<String>> {
        self.master.get_param_names()
    }

    pub fn state(&self) -> MasterResult<master::SystemState> {
        self.master.get_system_state()
    }

    pub fn topics(&self) -> MasterResult<Vec<Topic>> {
        self.master.get_topic_types()
    }

    pub fn client<T: ServicePair>(&self, service: &str) -> Result<Client<T>, ServerError> {
        let name = self.resolver.translate(service)?;
        let uri = self.master.lookup_service(&name)?;
        Ok(Client::new(&self.name, &uri, &name))
    }

    pub fn service<T, F>(&mut self, service: &str, handler: F) -> Result<(), ServerError>
        where T: ServicePair,
              F: Fn(T::Request) -> T::Response + Send + Sync + 'static
    {
        let name = self.resolver.translate(service)?;
        let api = self.slave.add_service::<T, F>(&self.hostname, &name, handler)?;

        if let Err(err) = self.master.register_service(&name, &api) {
            self.slave.remove_service(&name);
            self.master.unregister_service(&name, &api)?;
            Err(ServerError::from(err))
        } else {
            Ok(())
        }
    }

    pub fn subscribe<T, F>(&mut self, topic: &str, callback: F) -> Result<(), ServerError>
        where T: Message,
              F: Fn(T) -> () + Send + 'static
    {
        let name = self.resolver.translate(topic)?;
        self.slave.add_subscription::<T, F>(&name, callback)?;

        match self.master.register_subscriber(&name, &T::msg_type()) {
            Ok(publishers) => {
                if let Err(err) = self.slave
                    .add_publishers_to_subscription(&name, publishers.into_iter()) {
                    error!("Failed to subscribe to all publishers of topic '{}': {}",
                           name,
                           err);
                }
                Ok(())
            }
            Err(err) => {
                self.slave.remove_subscription(&name);
                self.master.unregister_subscriber(&name)?;
                Err(ServerError::from(err))
            }
        }
    }

    pub fn publish<T>(&mut self, topic: &str) -> Result<PublisherStream<T>, ServerError>
        where T: Message
    {
        let name = self.resolver.translate(topic)?;
        let stream = self.slave.add_publication::<T>(&self.hostname, &name)?;
        match self.master.register_publisher(&name, &T::msg_type()) {
            Ok(_) => Ok(stream),
            Err(error) => {
                error!("Failed to register publisher for topic '{}': {}",
                       name,
                       error);
                self.slave.remove_publication(&name);
                self.master.unregister_publisher(&name)?;
                Err(ServerError::from(error))
            }
        }
    }
}

pub struct Parameter<'a> {
    master: &'a Master,
    name: String,
}

impl<'a> Parameter<'a> {
    pub fn get<T: Decodable>(&self) -> MasterResult<T> {
        self.master.get_param::<T>(&self.name)
    }

    pub fn get_raw(&self) -> MasterResult<XmlRpcValue> {
        self.master.get_param_any(&self.name)
    }

    pub fn set<T: Encodable>(&self, value: &T) -> MasterResult<()> {
        self.master.set_param::<T>(&self.name, value).and(Ok(()))
    }

    pub fn delete(&self) -> MasterResult<()> {
        self.master.delete_param(&self.name).and(Ok(()))
    }

    pub fn exists(&self) -> MasterResult<bool> {
        self.master.has_param(&self.name)
    }

    pub fn search(&self) -> MasterResult<String> {
        self.master.search_param(&self.name)
    }
}
<commit_msg>Clear trailing namespace backlashes<commit_after>use rustc_serialize::{Encodable, Decodable};
use std;
use nix::unistd::gethostname;
use super::master::{self, Master, MasterResult};
use super::slave::Slave;
use super::error::ServerError;
use super::value::Topic;
use super::naming::Resolver;
use tcpros::{Client, Message, PublisherStream, ServicePair};
use rosxmlrpc::serde::XmlRpcValue;

pub struct Ros {
    master: Master,
    slave: Slave,
    hostname: String,
    resolver: Resolver,
    name: String,
}

impl Ros {
    pub fn new(name: &str) -> Result<Ros, ServerError> {
        let namespace = std::env::var("ROS_NAMESPACE").unwrap_or(String::from(""));
        Ros::new_raw(&namespace, name)
    }

    pub fn new_raw(namespace: &str, name: &str) -> Result<Ros, ServerError> {
        let namespace = namespace.trim_right_matches("/");
        let master_uri = std::env::var("ROS_MASTER_URI")
            .unwrap_or(String::from("http://localhost:11311/"));

        let mut hostname = [0u8; 50];
        gethostname(&mut hostname)?;
        let hostname =
            hostname.into_iter().take_while(|&v| *v != 0u8).map(|v| *v).collect::<Vec<_>>();
        let hostname = String::from_utf8(hostname)?;

        let name = format!("{}/{}", namespace, name);
        let resolver = Resolver::new(&name)?;

        let slave = Slave::new(&master_uri, &hostname, 0, &name)?;
        let master = Master::new(&master_uri, &name, &slave.uri());
        Ok(Ros {
            master: master,
            slave: slave,
            hostname: hostname,
            resolver: resolver,
            name: name,
        })
    }

    pub fn map(&mut self, source: &str, destination: &str) -> Result<(), ServerError> {
        self.resolver.map(source, destination).map_err(|v| ServerError::Naming(v))
    }

    pub fn node_uri(&self) -> &str {
        return self.slave.uri();
    }

    pub fn param<'a, 'b>(&'a self, name: &'b str) -> Option<Parameter<'a>> {
        self.resolver.translate(name).ok().map(|v| {
            Parameter {
                master: &self.master,
                name: v,
            }
        })
    }

    pub fn parameters(&self) -> MasterResult<Vec<String>> {
        self.master.get_param_names()
    }

    pub fn state(&self) -> MasterResult<master::SystemState> {
        self.master.get_system_state()
    }

    pub fn topics(&self) -> MasterResult<Vec<Topic>> {
        self.master.get_topic_types()
    }

    pub fn client<T: ServicePair>(&self, service: &str) -> Result<Client<T>, ServerError> {
        let name = self.resolver.translate(service)?;
        let uri = self.master.lookup_service(&name)?;
        Ok(Client::new(&self.name, &uri, &name))
    }

    pub fn service<T, F>(&mut self, service: &str, handler: F) -> Result<(), ServerError>
        where T: ServicePair,
              F: Fn(T::Request) -> T::Response + Send + Sync + 'static
    {
        let name = self.resolver.translate(service)?;
        let api = self.slave.add_service::<T, F>(&self.hostname, &name, handler)?;

        if let Err(err) = self.master.register_service(&name, &api) {
            self.slave.remove_service(&name);
            self.master.unregister_service(&name, &api)?;
            Err(ServerError::from(err))
        } else {
            Ok(())
        }
    }

    pub fn subscribe<T, F>(&mut self, topic: &str, callback: F) -> Result<(), ServerError>
        where T: Message,
              F: Fn(T) -> () + Send + 'static
    {
        let name = self.resolver.translate(topic)?;
        self.slave.add_subscription::<T, F>(&name, callback)?;

        match self.master.register_subscriber(&name, &T::msg_type()) {
            Ok(publishers) => {
                if let Err(err) = self.slave
                    .add_publishers_to_subscription(&name, publishers.into_iter()) {
                    error!("Failed to subscribe to all publishers of topic '{}': {}",
                           name,
                           err);
                }
                Ok(())
            }
            Err(err) => {
                self.slave.remove_subscription(&name);
                self.master.unregister_subscriber(&name)?;
                Err(ServerError::from(err))
            }
        }
    }

    pub fn publish<T>(&mut self, topic: &str) -> Result<PublisherStream<T>, ServerError>
        where T: Message
    {
        let name = self.resolver.translate(topic)?;
        let stream = self.slave.add_publication::<T>(&self.hostname, &name)?;
        match self.master.register_publisher(&name, &T::msg_type()) {
            Ok(_) => Ok(stream),
            Err(error) => {
                error!("Failed to register publisher for topic '{}': {}",
                       name,
                       error);
                self.slave.remove_publication(&name);
                self.master.unregister_publisher(&name)?;
                Err(ServerError::from(error))
            }
        }
    }
}

pub struct Parameter<'a> {
    master: &'a Master,
    name: String,
}

impl<'a> Parameter<'a> {
    pub fn get<T: Decodable>(&self) -> MasterResult<T> {
        self.master.get_param::<T>(&self.name)
    }

    pub fn get_raw(&self) -> MasterResult<XmlRpcValue> {
        self.master.get_param_any(&self.name)
    }

    pub fn set<T: Encodable>(&self, value: &T) -> MasterResult<()> {
        self.master.set_param::<T>(&self.name, value).and(Ok(()))
    }

    pub fn delete(&self) -> MasterResult<()> {
        self.master.delete_param(&self.name).and(Ok(()))
    }

    pub fn exists(&self) -> MasterResult<bool> {
        self.master.has_param(&self.name)
    }

    pub fn search(&self) -> MasterResult<String> {
        self.master.search_param(&self.name)
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;

extern crate serde;

extern crate toml;
extern crate clap;

use clap::{App, Arg, ArgMatches};
use std::path::Path;
use std::fs::{create_dir, File};
use std::io::prelude::*;

#[derive(Serialize, Deserialize)]
struct Config {
    data_dir: Option<String>,
}

const CONFIG_PATH: &'static str = "/etc/blochs/";
const CONFIG_FILE_NAME: &'static str = "server.toml";
const DEFAULT_DATA_DIR: &'static str = "/var/lib/blochs/";

fn main() {
    let options = get_provided_options();

    let config_path = Path::new(CONFIG_PATH);
    ensure_dir_exists(&config_path);

    let config_file_path = config_path.join(CONFIG_FILE_NAME);
    let mut read_config_file = match File::open(config_file_path.as_path()) {
        Ok(file) => file,
        Err(err) => panic!("Could not open config file {:?}: {}", config_file_path.display(), err),
    };

    let mut actual_content = String::new();
    match read_config_file.read_to_string(&mut actual_content) {
        Ok(_) => {},
        Err(err) => panic!("Could not open config file {:?}: {}", config_file_path.display(), err),
    };

    let mut config_values: Config = toml::from_str(&actual_content).unwrap();

    let config_or_default_data_dir = config_values.data_dir.unwrap_or(DEFAULT_DATA_DIR.to_string());

    let new_data_dir = options.value_of("data.dir").unwrap_or(&config_or_default_data_dir);

    config_values.data_dir = Some(new_data_dir.to_string());

    let new_config_content = toml::to_string(&config_values).unwrap();

    let mut write_config_file = match File::create(config_file_path.as_path()) {
        Ok(file) => file,
        Err(err) => panic!("Could not open config file {:?}: {}", config_file_path.display(), err),
    };

    match write_config_file.write_all(new_config_content.as_bytes()) {
        Ok(_) => println!("New config saved at {}:\n\n{}", config_file_path.display(), new_config_content),
        Err(why) => panic!("Could not write config file {:?}: {}", config_file_path.display(), why),
    };
}

fn ensure_dir_exists(path: &Path) {
    if !path.exists() {
        match create_dir(path) {
            Err(err) => panic!("Could not create directory under {}: {}", path.display(), err),
            Ok(_) => {}
        };
    }
}

fn get_provided_options<'a>() -> ArgMatches<'a> {
    return App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::with_name("data.dir")
            .long("data-dir")
            .value_name("DIR")
            .help(&format!("Sets where database data will be stored (default {})", DEFAULT_DATA_DIR))
        ).get_matches();
}
<commit_msg>extract file reading to function<commit_after>#[macro_use]
extern crate serde_derive;

extern crate serde;

extern crate toml;
extern crate clap;

use clap::{App, Arg, ArgMatches};
use std::path::Path;
use std::fs::{create_dir, File};
use std::io::prelude::*;

#[derive(Serialize, Deserialize)]
struct Config {
    data_dir: Option<String>,
}

const CONFIG_PATH: &'static str = "/etc/blochs/";
const CONFIG_FILE_NAME: &'static str = "server.toml";
const DEFAULT_DATA_DIR: &'static str = "/var/lib/blochs/";

fn main() {
    let options = get_provided_options();

    let config_path = Path::new(CONFIG_PATH);
    ensure_dir_exists(&config_path);

    let config_file_path = config_path.join(CONFIG_FILE_NAME);

    let actual_content = get_file_content(&config_file_path);

    let mut config_values: Config = toml::from_str(&actual_content).unwrap();

    let config_or_default_data_dir = config_values.data_dir.unwrap_or(DEFAULT_DATA_DIR.to_string());

    let new_data_dir = options.value_of("data.dir").unwrap_or(&config_or_default_data_dir);

    config_values.data_dir = Some(new_data_dir.to_string());

    let new_config_content = toml::to_string(&config_values).unwrap();

    let mut write_config_file = match File::create(config_file_path.as_path()) {
        Ok(file) => file,
        Err(err) => panic!("Could not open config file {:?}: {}", config_file_path.display(), err),
    };

    match write_config_file.write_all(new_config_content.as_bytes()) {
        Ok(_) => println!("New config saved at {}:\n\n{}", config_file_path.display(), new_config_content),
        Err(why) => panic!("Could not write config file {:?}: {}", config_file_path.display(), why),
    };
}

fn get_file_content(path: &Path) -> String {
    let mut content = String::new();
    if path.exists() {
        let mut file = match File::open(&path) {
            Ok(file) => file,
            Err(err) => panic!("Could not open file {:?}: {}", path.display(), err),
        };

        match file.read_to_string(&mut content) {
            Ok(_) => {},
            Err(err) => panic!("Could not read file {:?}: {}", path.display(), err),
        };
    }
    return content;
}

fn ensure_dir_exists(path: &Path) {
    if !path.exists() {
        match create_dir(path) {
            Err(err) => panic!("Could not create directory under {}: {}", path.display(), err),
            Ok(_) => {}
        };
    }
}

fn get_provided_options<'a>() -> ArgMatches<'a> {
    return App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::with_name("data.dir")
            .long("data-dir")
            .value_name("DIR")
            .help(&format!("Sets where database data will be stored (default {})", DEFAULT_DATA_DIR))
        ).get_matches();
}
<|endoftext|>


<commit_before>use gc_error::GcError;
use types::binding::Binding;
use types::js_var::{JsPtrEnum, JsVar};

pub trait Backend {
    fn alloc(&mut self, var: JsVar, ptr: Option<JsPtrEnum>) -> Result<Binding, GcError>;
    fn load(&mut self, bnd: &Binding) -> Result<(JsVar, Option<JsPtrEnum>), GcError>;
    fn store(&mut self, var: JsVar, ptr: Option<JsPtrEnum>) -> Result<(), GcError>;
}
<commit_msg>add `get_allocator` method to Backend<commit_after>use std::cell::RefCell;
use std::rc::Rc;

use gc_error::GcError;
use types::allocator::Allocator;
use types::binding::Binding;
use types::js_var::{JsPtrEnum, JsVar};

pub trait Backend {
    fn alloc(&mut self, var: JsVar, ptr: Option<JsPtrEnum>) -> Result<Binding, GcError>;
    fn load(&mut self, bnd: &Binding) -> Result<(JsVar, Option<JsPtrEnum>), GcError>;
    fn store(&mut self, var: JsVar, ptr: Option<JsPtrEnum>) -> Result<(), GcError>;
    fn get_allocator(&self) -> Rc<RefCell<Allocator<Error=GcError>>>;
}
<|endoftext|>


<commit_before>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;
use std::str;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait BAMRead {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct BAMReader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMReader {
    /// Create a new BAMReader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: &P) -> Self {
        let f = bgzf_open(path, b"r");
        let header = unsafe { htslib::bam_hdr_read(f) };
        BAMReader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl BAMRead for BAMReader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for BAMReader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// A BAM writer.
pub struct BAMWriter {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMWriter {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");
        let header_record = unsafe {
            let header_string = header.to_bytes();
            println!("{}", str::from_utf8(&header_string).unwrap());
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }
        BAMWriter { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }
        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }
        BAMWriter { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for BAMWriter {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: BAMReader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use super::record::*;
    use super::header::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = BAMReader::new(&"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));
        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }

    #[test]
    fn test_write() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bampath = tmp.path().join("test.bam");
        {
            let mut bam = BAMWriter::new(
                &bampath,
                Header::new().push_record(
                    HeaderRecord::new(b"SQ").push_tag(b"SN", &"chr1")
                                            .push_tag(b"LN", &15072423)
                )
            );
            let mut rec = record::Record::new();
            rec.set(qname, &cigar, seq, qual);
            rec.push_aux(b"NM", &Aux::Integer(15));
            bam.write(&mut rec).ok().expect("Failed to write record.");
        }
        {
            let bam = BAMReader::new(&bampath);
            let mut rec = record::Record::new();
            bam.read(&mut rec).ok().expect("Failed to read record.");
            assert_eq!(rec.qname(), qname);
            assert_eq!(rec.cigar(), cigar);
            assert_eq!(rec.seq().as_bytes(), seq);
            assert_eq!(rec.qual(), qual);
            assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<commit_msg>Formatting.<commit_after>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;
use std::str;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait BAMRead {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct BAMReader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMReader {
    /// Create a new BAMReader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: &P) -> Self {
        let f = bgzf_open(path, b"r");
        let header = unsafe { htslib::bam_hdr_read(f) };
        BAMReader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl BAMRead for BAMReader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for BAMReader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// A BAM writer.
pub struct BAMWriter {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMWriter {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");
        let header_record = unsafe {
            let header_string = header.to_bytes();
            println!("{}", str::from_utf8(&header_string).unwrap());
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }
        BAMWriter { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }
        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }
        BAMWriter { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for BAMWriter {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: BAMReader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use super::record::*;
    use super::header::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = BAMReader::new(&"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));

        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }

    #[test]
    fn test_write() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bampath = tmp.path().join("test.bam");
        {
            let mut bam = BAMWriter::new(
                &bampath,
                Header::new().push_record(
                    HeaderRecord::new(b"SQ").push_tag(b"SN", &"chr1")
                                            .push_tag(b"LN", &15072423)
                )
            );

            let mut rec = record::Record::new();
            rec.set(qname, &cigar, seq, qual);
            rec.push_aux(b"NM", &Aux::Integer(15));

            bam.write(&mut rec).ok().expect("Failed to write record.");
        }
        {
            let bam = BAMReader::new(&bampath);

            let mut rec = record::Record::new();
            bam.read(&mut rec).ok().expect("Failed to read record.");

            assert_eq!(rec.qname(), qname);
            assert_eq!(rec.cigar(), cigar);
            assert_eq!(rec.seq().as_bytes(), seq);
            assert_eq!(rec.qual(), qual);
            assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<|endoftext|>


<commit_before>//! Randomization of big integers

// Some of the tests of non-RNG-based functionality are randomized using the
// RNG-based functionality, so the RNG-based functionality needs to be enabled
// for tests.
use rand::Rng;

use BigInt;
use BigUint;
use Sign::*;
use big_digit::BigDigit;

use traits::Zero;
use integer::Integer;

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint;

    /// Generate a random BigInt of the given bit size.
    fn gen_bigint(&mut self, bit_size: usize) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    /// Generate a random `BigInt` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

#[cfg(any(feature = "rand", test))]
impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint {
        use super::big_digit::BITS;
        let (digits, rem) = bit_size.div_rem(&BITS);
        let mut data = Vec::with_capacity(digits + 1);
        for _ in 0..digits {
            data.push(self.gen());
        }
        if rem > 0 {
            let final_digit: BigDigit = self.gen();
            data.push(final_digit >> (BITS - rem));
        }
        BigUint::new(data)
    }

    fn gen_bigint(&mut self, bit_size: usize) -> BigInt {
        // Generate a random BigUint...
        let biguint = self.gen_biguint(bit_size);
        // ...and then randomly assign it a Sign...
        let sign = if biguint.is_zero() {
            // ...except that if the BigUint is zero, we need to try
            // again with probability 0.5. This is because otherwise,
            // the probability of generating a zero BigInt would be
            // double that of any other number.
            if self.gen() {
                return self.gen_bigint(bit_size);
            } else {
                NoSign
            }
        } else if self.gen() {
            Plus
        } else {
            Minus
        };
        BigInt::from_biguint(sign, biguint)
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        assert!(!bound.is_zero());
        let bits = bound.bits();
        loop {
            let n = self.gen_biguint(bits);
            if n < *bound {
                return n;
            }
        }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        assert!(*lbound < *ubound);
        return lbound + self.gen_biguint_below(&(ubound - lbound));
    }

    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt {
        assert!(*lbound < *ubound);
        let delta = (ubound - lbound).to_biguint().unwrap();
        return lbound + BigInt::from(self.gen_biguint_below(&delta))
    }
}
<commit_msg>loop gen_bigint retries<commit_after>//! Randomization of big integers

// Some of the tests of non-RNG-based functionality are randomized using the
// RNG-based functionality, so the RNG-based functionality needs to be enabled
// for tests.
use rand::Rng;

use BigInt;
use BigUint;
use Sign::*;
use big_digit::BigDigit;

use traits::Zero;
use integer::Integer;

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint;

    /// Generate a random BigInt of the given bit size.
    fn gen_bigint(&mut self, bit_size: usize) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    /// Generate a random `BigInt` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

#[cfg(any(feature = "rand", test))]
impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint {
        use super::big_digit::BITS;
        let (digits, rem) = bit_size.div_rem(&BITS);
        let mut data = Vec::with_capacity(digits + 1);
        for _ in 0..digits {
            data.push(self.gen());
        }
        if rem > 0 {
            let final_digit: BigDigit = self.gen();
            data.push(final_digit >> (BITS - rem));
        }
        BigUint::new(data)
    }

    fn gen_bigint(&mut self, bit_size: usize) -> BigInt {
        loop {
            // Generate a random BigUint...
            let biguint = self.gen_biguint(bit_size);
            // ...and then randomly assign it a Sign...
            let sign = if biguint.is_zero() {
                // ...except that if the BigUint is zero, we need to try
                // again with probability 0.5. This is because otherwise,
                // the probability of generating a zero BigInt would be
                // double that of any other number.
                if self.gen() {
                    continue;
                } else {
                    NoSign
                }
            } else if self.gen() {
                Plus
            } else {
                Minus
            };
            return BigInt::from_biguint(sign, biguint);
        }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        assert!(!bound.is_zero());
        let bits = bound.bits();
        loop {
            let n = self.gen_biguint(bits);
            if n < *bound {
                return n;
            }
        }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        assert!(*lbound < *ubound);
        return lbound + self.gen_biguint_below(&(ubound - lbound));
    }

    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt {
        assert!(*lbound < *ubound);
        let delta = (ubound - lbound).to_biguint().unwrap();
        return lbound + BigInt::from(self.gen_biguint_below(&delta))
    }
}
<|endoftext|>


<commit_before>use std::env;

use cargo::core::Workspace;
use cargo::ops::{self, MessageFormat, Packages};
use cargo::util::{CliResult, Config};
use cargo::util::important_paths::{find_root_manifest_for_wd};

#[derive(Deserialize)]
pub struct Options {
    flag_target: Option<String>,
    flag_features: Vec<String>,
    flag_all_features: bool,
    flag_jobs: Option<u32>,
    flag_manifest_path: Option<String>,
    flag_no_default_features: bool,
    flag_no_deps: bool,
    flag_open: bool,
    flag_release: bool,
    flag_verbose: u32,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    flag_message_format: MessageFormat,
    flag_package: Vec<String>,
    flag_lib: bool,
    flag_bin: Vec<String>,
    flag_bins: bool,
    flag_frozen: bool,
    flag_locked: bool,
    flag_all: bool,
    #[serde(rename = "flag_Z")]
    flag_z: Vec<String>,
}

pub const USAGE: &'static str = "
Build a package's documentation

Usage:
    cargo doc [options]

Options:
    -h, --help                   Print this message
    --open                       Opens the docs in a browser after the operation
    -p SPEC, --package SPEC ...  Package to document
    --all                        Document all packages in the workspace
    --no-deps                    Don't build documentation for dependencies
    -j N, --jobs N               Number of parallel jobs, defaults to # of CPUs
    --lib                        Document only this package's library
    --bin NAME                   Document only the specified binary
    --bins                       Document all binaries
    --release                    Build artifacts in release mode, with optimizations
    --features FEATURES          Space-separated list of features to also build
    --all-features               Build all available features
    --no-default-features        Do not build the `default` feature
    --target TRIPLE              Build for the target triple
    --manifest-path PATH         Path to the manifest to document
    -v, --verbose ...            Use verbose output (-vv very verbose/build.rs output)
    -q, --quiet                  No output printed to stdout
    --color WHEN                 Coloring: auto, always, never
    --message-format FMT         Error format: human, json [default: human]
    --frozen                     Require Cargo.lock and cache are up to date
    --locked                     Require Cargo.lock is up to date
    -Z FLAG ...                  Unstable (nightly-only) flags to Cargo

By default the documentation for the local package and all dependencies is
built. The output is all placed in `target/doc` in rustdoc's usual format.

All packages in the workspace are documented if the `--all` flag is supplied. The
`--all` flag is automatically assumed for a virtual manifest.
Note that `--exclude` has to be specified in conjunction with the `--all` flag.

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be documented. If it is not given, then the
current package is documented. For more information on SPEC and its format, see
the `cargo help pkgid` command.
";

pub fn execute(options: Options, config: &mut Config) -> CliResult {
    debug!("executing; cmd=cargo-check; args={:?}",
           env::args().collect::<Vec<_>>());

    config.configure(options.flag_verbose,
                     options.flag_quiet,
                     &options.flag_color,
                     options.flag_frozen,
                     options.flag_locked,
                     &options.flag_z)?;

    let root = find_root_manifest_for_wd(options.flag_manifest_path, config.cwd())?;
    let ws = Workspace::new(&root, config)?;

    let spec = if options.flag_all || (ws.is_virtual() && options.flag_package.is_empty()) {
        Packages::All
    } else {
        Packages::Packages(&options.flag_package)
    };

    let empty = Vec::new();
    let doc_opts = ops::DocOptions {
        open_result: options.flag_open,
        compile_opts: ops::CompileOptions {
            config: config,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|t| &t[..]),
            features: &options.flag_features,
            all_features: options.flag_all_features,
            no_default_features: options.flag_no_default_features,
            spec: spec,
            filter: ops::CompileFilter::new(options.flag_lib,
                                            &options.flag_bin, options.flag_bins,
                                            &empty, false,
                                            &empty, false,
                                            &empty, false,
                                            false),
            message_format: options.flag_message_format,
            release: options.flag_release,
            mode: ops::CompileMode::Doc {
                deps: !options.flag_no_deps,
            },
            target_rustc_args: None,
            target_rustdoc_args: None,
        },
    };

    ops::doc(&ws, &doc_opts)?;
    Ok(())
}
<commit_msg>Auto merge of #5081 - matklad:document-not-all-the-things, r=alexcrichton<commit_after>use std::env;

use cargo::core::Workspace;
use cargo::ops::{self, MessageFormat, Packages};
use cargo::util::{CliResult, Config};
use cargo::util::important_paths::{find_root_manifest_for_wd};

#[derive(Deserialize)]
pub struct Options {
    flag_target: Option<String>,
    flag_features: Vec<String>,
    flag_all_features: bool,
    flag_jobs: Option<u32>,
    flag_manifest_path: Option<String>,
    flag_no_default_features: bool,
    flag_no_deps: bool,
    flag_open: bool,
    flag_release: bool,
    flag_verbose: u32,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    flag_message_format: MessageFormat,
    flag_package: Vec<String>,
    flag_lib: bool,
    flag_bin: Vec<String>,
    flag_bins: bool,
    flag_frozen: bool,
    flag_locked: bool,
    flag_all: bool,
    flag_exclude: Vec<String>,
    #[serde(rename = "flag_Z")]
    flag_z: Vec<String>,
}

pub const USAGE: &'static str = "
Build a package's documentation

Usage:
    cargo doc [options]

Options:
    -h, --help                   Print this message
    --open                       Opens the docs in a browser after the operation
    -p SPEC, --package SPEC ...  Package to document
    --all                        Document all packages in the workspace
    --exclude SPEC ...           Exclude packages from the build
    --no-deps                    Don't build documentation for dependencies
    -j N, --jobs N               Number of parallel jobs, defaults to # of CPUs
    --lib                        Document only this package's library
    --bin NAME                   Document only the specified binary
    --bins                       Document all binaries
    --release                    Build artifacts in release mode, with optimizations
    --features FEATURES          Space-separated list of features to also build
    --all-features               Build all available features
    --no-default-features        Do not build the `default` feature
    --target TRIPLE              Build for the target triple
    --manifest-path PATH         Path to the manifest to document
    -v, --verbose ...            Use verbose output (-vv very verbose/build.rs output)
    -q, --quiet                  No output printed to stdout
    --color WHEN                 Coloring: auto, always, never
    --message-format FMT         Error format: human, json [default: human]
    --frozen                     Require Cargo.lock and cache are up to date
    --locked                     Require Cargo.lock is up to date
    -Z FLAG ...                  Unstable (nightly-only) flags to Cargo

By default the documentation for the local package and all dependencies is
built. The output is all placed in `target/doc` in rustdoc's usual format.

All packages in the workspace are documented if the `--all` flag is supplied. The
`--all` flag is automatically assumed for a virtual manifest.
Note that `--exclude` has to be specified in conjunction with the `--all` flag.

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be documented. If it is not given, then the
current package is documented. For more information on SPEC and its format, see
the `cargo help pkgid` command.
";

pub fn execute(options: Options, config: &mut Config) -> CliResult {
    debug!("executing; cmd=cargo-check; args={:?}",
           env::args().collect::<Vec<_>>());

    config.configure(options.flag_verbose,
                     options.flag_quiet,
                     &options.flag_color,
                     options.flag_frozen,
                     options.flag_locked,
                     &options.flag_z)?;

    let root = find_root_manifest_for_wd(options.flag_manifest_path, config.cwd())?;
    let ws = Workspace::new(&root, config)?;

    let spec = Packages::from_flags(options.flag_all,
                                    &options.flag_exclude,
                                    &options.flag_package)?;

    let empty = Vec::new();
    let doc_opts = ops::DocOptions {
        open_result: options.flag_open,
        compile_opts: ops::CompileOptions {
            config: config,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|t| &t[..]),
            features: &options.flag_features,
            all_features: options.flag_all_features,
            no_default_features: options.flag_no_default_features,
            spec: spec,
            filter: ops::CompileFilter::new(options.flag_lib,
                                            &options.flag_bin, options.flag_bins,
                                            &empty, false,
                                            &empty, false,
                                            &empty, false,
                                            false),
            message_format: options.flag_message_format,
            release: options.flag_release,
            mode: ops::CompileMode::Doc {
                deps: !options.flag_no_deps,
            },
            target_rustc_args: None,
            target_rustdoc_args: None,
        },
    };

    ops::doc(&ws, &doc_opts)?;
    Ok(())
}
<|endoftext|>


<commit_before>fn main() {

}
<commit_msg>implement pwd<commit_after>use std::env;

fn main() {
	println!("{}", env::current_dir().unwrap().display())
}
<|endoftext|>


<commit_before>
extern crate lapack;

use self::lapack::fortran::*;

pub trait LapackBinding: Sized {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32);
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32);
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32);
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self;
}

impl LapackBinding for f64 {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32) {
        dsyev(jobz, uplo, n, a, lda, w, work, lwork, info);
    }
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32) {
        dgetrf(m, n, a, lda, ipiv, info);
    }
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32) {
        dgetri(n, a, lda, ipiv, work, lwork, info);
    }
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self {
        dlange(norm, m, n, a, lda, work)
    }
}

impl LapackBinding for f32 {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32) {
        ssyev(jobz, uplo, n, a, lda, w, work, lwork, info);
    }
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32) {
        sgetrf(m, n, a, lda, ipiv, info);
    }
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32) {
        sgetri(n, a, lda, ipiv, work, lwork, info);
    }
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self {
        slange(norm, m, n, a, lda, work)
    }
}
<commit_msg>Add _gesvd<commit_after>
extern crate lapack;

use self::lapack::fortran::*;

pub trait LapackBinding: Sized {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32);
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32);
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32);
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self;
    fn _gesvd(jobu: u8,
              jobvt: u8,
              m: i32,
              n: i32,
              a: &mut [Self],
              lda: i32,
              s: &mut [Self],
              u: &mut [Self],
              ldu: i32,
              vt: &mut [Self],
              ldvt: i32,
              work: &mut [Self],
              lwork: i32,
              info: &mut i32);
}

impl LapackBinding for f64 {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32) {
        dsyev(jobz, uplo, n, a, lda, w, work, lwork, info);
    }
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32) {
        dgetrf(m, n, a, lda, ipiv, info);
    }
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32) {
        dgetri(n, a, lda, ipiv, work, lwork, info);
    }
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self {
        dlange(norm, m, n, a, lda, work)
    }
    fn _gesvd(jobu: u8,
              jobvt: u8,
              m: i32,
              n: i32,
              a: &mut [Self],
              lda: i32,
              s: &mut [Self],
              u: &mut [Self],
              ldu: i32,
              vt: &mut [Self],
              ldvt: i32,
              work: &mut [Self],
              lwork: i32,
              info: &mut i32) {
        dgesvd(jobu,
               jobvt,
               m,
               n,
               a,
               lda,
               s,
               u,
               ldu,
               vt,
               ldvt,
               work,
               lwork,
               info);
    }
}

impl LapackBinding for f32 {
    fn _syev(jobz: u8,
             uplo: u8,
             n: i32,
             a: &mut Vec<Self>,
             lda: i32,
             w: &mut Vec<Self>,
             work: &mut Vec<Self>,
             lwork: i32,
             info: &mut i32) {
        ssyev(jobz, uplo, n, a, lda, w, work, lwork, info);
    }
    fn _getrf(m: i32, n: i32, a: &mut Vec<Self>, lda: i32, ipiv: &mut [i32], info: &mut i32) {
        sgetrf(m, n, a, lda, ipiv, info);
    }
    fn _getri(n: i32,
              a: &mut Vec<Self>,
              lda: i32,
              ipiv: &[i32],
              work: &mut Vec<Self>,
              lwork: i32,
              info: &mut i32) {
        sgetri(n, a, lda, ipiv, work, lwork, info);
    }
    fn _lange(norm: u8, m: i32, n: i32, a: &Vec<Self>, lda: i32, work: &mut Vec<Self>) -> Self {
        slange(norm, m, n, a, lda, work)
    }
    fn _gesvd(jobu: u8,
              jobvt: u8,
              m: i32,
              n: i32,
              a: &mut [Self],
              lda: i32,
              s: &mut [Self],
              u: &mut [Self],
              ldu: i32,
              vt: &mut [Self],
              ldvt: i32,
              work: &mut [Self],
              lwork: i32,
              info: &mut i32) {
        sgesvd(jobu,
               jobvt,
               m,
               n,
               a,
               lda,
               s,
               u,
               ldu,
               vt,
               ldvt,
               work,
               lwork,
               info);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use super::metric::{Metric, MetricKind};
use clock_ticks;

/// Buckets are the primary internal storage type.
///
/// Each bucket contains a set of hashmaps containing
/// each set of metrics received by clients.
///
pub struct Buckets {
    counters: HashMap<String, f64>,
    gauges: HashMap<String, f64>,
    timers: HashMap<String, Vec<f64>>,

    server_start_time: u64,
    last_message: u64,
    bad_messages: usize,
    total_messages: usize,
}

impl Buckets {
    /// Create a new Buckets
    ///
    /// ```
    /// let bucket = Buckets::new();
    /// ```
    pub fn new() -> Buckets {
        Buckets {
            counters: HashMap::new(),
            gauges: HashMap::new(),
            timers: HashMap::new(),
            bad_messages: 0,
            total_messages: 0,
            last_message: clock_ticks::precise_time_ms(),
            server_start_time: clock_ticks::precise_time_ms(),
        }
    }

    /// Adds a metric to the bucket storage.
    ///
    /// # Examples
    ///
    /// ```
    /// use super::metric;
    /// use std::str::FromStr;
    ///
    /// let metric = metric::Metric::FromStr("foo:1|c");
    /// let mut bucket = Buckets::new();
    /// bucket.add(metric);
    /// ```
    pub fn add(&mut self, value: &Metric) {
        let name = value.name.to_owned();
        match value.kind {
            MetricKind::Counter(rate) => {
                let counter = self.counters.entry(name).or_insert(0.0);
                *counter = *counter + value.value * (1.0 / rate);
            },
            MetricKind::Gauge => {
                self.gauges.insert(name, value.value);
            },
            MetricKind::Timer => {
                let slot = self.timers.entry(name).or_insert(Vec::new());
                slot.push(value.value);
            },
        }
        self.last_message = clock_ticks::precise_time_ms();
    }

    /// Increment the bad message count by one.
    ///
    pub fn add_bad_message(&mut self) {
        self.bad_messages += 1
    }

    /// Get the count of bad messages
    pub fn bad_messages(&self) -> usize {
        self.bad_messages
    }
}

//
// Tests
//
#[cfg(test)]
mod test {
    use super::*;
    use super::super::metric::{Metric, MetricKind};

    #[test]
    fn test_bad_messages() {
        let mut buckets = Buckets::new();
        buckets.add_bad_message();
        assert_eq!(1, buckets.bad_messages());

        buckets.add_bad_message();
        assert_eq!(2, buckets.bad_messages());
    }

    #[test]
    fn test_add_increments_last_message_timer() {
        let mut buckets = Buckets::new();
        // duff value to ensure it changes.
        let original = 10;
        buckets.last_message = original;

        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(1.0));
        buckets.add(&metric);
        assert!(buckets.last_message > original);
    }

    #[test]
    fn test_add_counter_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(1.0));
        buckets.add(&metric);

        assert!(buckets.counters.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&1.0), buckets.counters.get("some.metric"));

        // Increment counter
        buckets.add(&metric);
        assert_eq!(Some(&2.0), buckets.counters.get("some.metric"));
    }

    #[test]
    fn test_add_counter_metric_sampled() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(0.1));

        buckets.add(&metric);
        assert_eq!(Some(&10.0), buckets.counters.get("some.metric"));

        let metric_two = Metric::new("some.metric", 1.0, MetricKind::Counter(0.5));
        buckets.add(&metric_two);
        assert_eq!(Some(&12.0), buckets.counters.get("some.metric"));
    }

    #[test]
    fn test_add_gauge_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 11.5, MetricKind::Gauge);
        buckets.add(&metric);
        assert!(buckets.gauges.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&11.5), buckets.gauges.get("some.metric"));
    }

    #[test]
    fn test_add_timer_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 11.5, MetricKind::Timer);
        buckets.add(&metric);
        assert!(buckets.timers.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&vec![11.5]), buckets.timers.get("some.metric"));

        let metric_two = Metric::new("some.metric", 99.5, MetricKind::Timer);
        buckets.add(&metric_two);

        let metric_three = Metric::new("other.metric", 811.5, MetricKind::Timer);
        buckets.add(&metric_three);
        assert!(buckets.timers.contains_key("some.metric"),
                "Should contain the metric key");
        assert!(buckets.timers.contains_key("other.metric"),
                "Should contain the metric key");

        assert_eq!(Some(&vec![11.5, 99.5]), buckets.timers.get("some.metric"));
        assert_eq!(Some(&vec![811.5]), buckets.timers.get("other.metric"));
    }
}
<commit_msg>Track total messages sent.<commit_after>use std::collections::HashMap;
use super::metric::{Metric, MetricKind};
use clock_ticks;

/// Buckets are the primary internal storage type.
///
/// Each bucket contains a set of hashmaps containing
/// each set of metrics received by clients.
///
pub struct Buckets {
    counters: HashMap<String, f64>,
    gauges: HashMap<String, f64>,
    timers: HashMap<String, Vec<f64>>,

    server_start_time: u64,
    last_message: u64,
    bad_messages: usize,
    total_messages: usize,
}

impl Buckets {
    /// Create a new Buckets
    ///
    /// ```
    /// let bucket = Buckets::new();
    /// ```
    pub fn new() -> Buckets {
        Buckets {
            counters: HashMap::new(),
            gauges: HashMap::new(),
            timers: HashMap::new(),
            bad_messages: 0,
            total_messages: 0,
            last_message: clock_ticks::precise_time_ms(),
            server_start_time: clock_ticks::precise_time_ms(),
        }
    }

    /// Adds a metric to the bucket storage.
    ///
    /// # Examples
    ///
    /// ```
    /// use super::metric;
    /// use std::str::FromStr;
    ///
    /// let metric = metric::Metric::FromStr("foo:1|c");
    /// let mut bucket = Buckets::new();
    /// bucket.add(metric);
    /// ```
    pub fn add(&mut self, value: &Metric) {
        let name = value.name.to_owned();
        match value.kind {
            MetricKind::Counter(rate) => {
                let counter = self.counters.entry(name).or_insert(0.0);
                *counter = *counter + value.value * (1.0 / rate);
            },
            MetricKind::Gauge => {
                self.gauges.insert(name, value.value);
            },
            MetricKind::Timer => {
                let slot = self.timers.entry(name).or_insert(Vec::new());
                slot.push(value.value);
            },
        }
        self.last_message = clock_ticks::precise_time_ms();
        self.total_messages += 1;
    }

    /// Increment the bad message count by one.
    ///
    pub fn add_bad_message(&mut self) {
        self.bad_messages += 1
    }

    /// Get the count of bad messages
    pub fn bad_messages(&self) -> usize {
        self.bad_messages
    }
}

//
// Tests
//
#[cfg(test)]
mod test {
    use super::*;
    use super::super::metric::{Metric, MetricKind};

    #[test]
    fn test_bad_messages() {
        let mut buckets = Buckets::new();
        buckets.add_bad_message();
        assert_eq!(1, buckets.bad_messages());

        buckets.add_bad_message();
        assert_eq!(2, buckets.bad_messages());
    }

    #[test]
    fn test_add_increments_total_messages() {
        let mut buckets = Buckets::new();
        // duff value to ensure it changes.
        let original = 10;
        buckets.last_message = original;

        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(1.0));
        buckets.add(&metric);
        assert!(buckets.last_message > original);
    }

    #[test]
    fn test_add_increments_last_message_timer() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(1.0));
        buckets.add(&metric);
        assert_eq!(1, buckets.total_messages);

        buckets.add(&metric);
        assert_eq!(2, buckets.total_messages);
    }

    #[test]
    fn test_add_counter_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(1.0));
        buckets.add(&metric);

        assert!(buckets.counters.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&1.0), buckets.counters.get("some.metric"));

        // Increment counter
        buckets.add(&metric);
        assert_eq!(Some(&2.0), buckets.counters.get("some.metric"));
    }

    #[test]
    fn test_add_counter_metric_sampled() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 1.0, MetricKind::Counter(0.1));

        buckets.add(&metric);
        assert_eq!(Some(&10.0), buckets.counters.get("some.metric"));

        let metric_two = Metric::new("some.metric", 1.0, MetricKind::Counter(0.5));
        buckets.add(&metric_two);
        assert_eq!(Some(&12.0), buckets.counters.get("some.metric"));
    }

    #[test]
    fn test_add_gauge_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 11.5, MetricKind::Gauge);
        buckets.add(&metric);
        assert!(buckets.gauges.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&11.5), buckets.gauges.get("some.metric"));
    }

    #[test]
    fn test_add_timer_metric() {
        let mut buckets = Buckets::new();
        let metric = Metric::new("some.metric", 11.5, MetricKind::Timer);
        buckets.add(&metric);
        assert!(buckets.timers.contains_key("some.metric"),
                "Should contain the metric key");
        assert_eq!(Some(&vec![11.5]), buckets.timers.get("some.metric"));

        let metric_two = Metric::new("some.metric", 99.5, MetricKind::Timer);
        buckets.add(&metric_two);

        let metric_three = Metric::new("other.metric", 811.5, MetricKind::Timer);
        buckets.add(&metric_three);
        assert!(buckets.timers.contains_key("some.metric"),
                "Should contain the metric key");
        assert!(buckets.timers.contains_key("other.metric"),
                "Should contain the metric key");

        assert_eq!(Some(&vec![11.5, 99.5]), buckets.timers.get("some.metric"));
        assert_eq!(Some(&vec![811.5]), buckets.timers.get("other.metric"));
    }
}
<|endoftext|>