<commit_before>extern crate num;
use self::num::Float;
use super::*;

pub trait Ball<D> {
    fn radius(&self) -> D;

    fn space<V>(&self) -> D
        where V: Vector<D>, D: Float
    {
        V::space_ball(self.radius())
    }
}

pub struct Box<V> {
    pub origin: V,
    pub offset: V,
}

impl<V> Box<V> {
    pub fn new(origin: V, offset: V) -> Self {
        Box{
            origin: origin,
            offset: offset,
        }
    }

    //Compute the amount of space contained in the box
    pub fn space<D>(&self) -> D
        where V: Vector<D>, D: Float
    {
        self.offset.space_box()
    }
}

impl<V> Clone for Box<V>
    where V: Clone
{
    fn clone(&self) -> Self {
        Box{
            origin: self.origin.clone(),
            offset: self.offset.clone(),
        }
    }
}
<commit_msg>add Toroid trait for toroidal space<commit_after>extern crate num;
extern crate nalgebra as na;
use self::num::{Float, FromPrimitive};
use super::*;

/// Anything that implements this trait permits its space to be infinitely wrapping.
pub trait Toroid<V> {
    /// Wrap a delta vector between two positions inside of the toroidal space
    fn wrap_delta(&self, delta: V) -> V;

    /// Wrap a position to keep it inside of the space
    fn wrap_position(&self, pos: V) -> V;
}

pub trait Ball<D> {
    fn radius(&self) -> D;

    fn space<V>(&self) -> D
        where V: Vector<D>, D: Float
    {
        V::space_ball(self.radius())
    }
}

/// A Box with a center at origin and one of the corners created by offset
///
/// The box is aligned so that the face normals point along each axis.
pub struct Box<V> {
    pub origin: V,
    pub offset: V,
}

impl<V> Box<V> {
    pub fn new(origin: V, offset: V) -> Self {
        Box{
            origin: origin,
            offset: offset,
        }
    }

    //Compute the amount of space contained in the box
    pub fn space<D>(&self) -> D
        where V: Vector<D>, D: Float
    {
        self.offset.space_box()
    }
}

impl<D> Toroid<Cartesian1<D>> for Box<Cartesian1<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: Cartesian1<D>) -> Cartesian1<D> {
        Cartesian1{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
        }
    }

    fn wrap_position(&self, pos: Cartesian1<D>) -> Cartesian1<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<D> Toroid<Cartesian2<D>> for Box<Cartesian2<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: Cartesian2<D>) -> Cartesian2<D> {
        Cartesian2{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
            y: delta.y % (D::from_u32(2u32).unwrap() * self.offset.y),
        }
    }

    fn wrap_position(&self, pos: Cartesian2<D>) -> Cartesian2<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<D> Toroid<Cartesian3<D>> for Box<Cartesian3<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: Cartesian3<D>) -> Cartesian3<D> {
        Cartesian3{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
            y: delta.y % (D::from_u32(2u32).unwrap() * self.offset.y),
            z: delta.z % (D::from_u32(2u32).unwrap() * self.offset.z),
        }
    }

    fn wrap_position(&self, pos: Cartesian3<D>) -> Cartesian3<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<D> Toroid<na::Vec1<D>> for Box<na::Vec1<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: na::Vec1<D>) -> na::Vec1<D> {
        na::Vec1{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
        }
    }

    fn wrap_position(&self, pos: na::Vec1<D>) -> na::Vec1<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<D> Toroid<na::Vec2<D>> for Box<na::Vec2<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: na::Vec2<D>) -> na::Vec2<D> {
        na::Vec2{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
            y: delta.y % (D::from_u32(2u32).unwrap() * self.offset.y),
        }
    }

    fn wrap_position(&self, pos: na::Vec2<D>) -> na::Vec2<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<D> Toroid<na::Vec3<D>> for Box<na::Vec3<D>>
    where D: Float + FromPrimitive
{
    fn wrap_delta(&self, delta: na::Vec3<D>) -> na::Vec3<D> {
        na::Vec3{
            x: delta.x % (D::from_u32(2u32).unwrap() * self.offset.x),
            y: delta.y % (D::from_u32(2u32).unwrap() * self.offset.y),
            z: delta.z % (D::from_u32(2u32).unwrap() * self.offset.z),
        }
    }

    fn wrap_position(&self, pos: na::Vec3<D>) -> na::Vec3<D> {
        self.wrap_delta(pos - self.origin) + self.origin
    }
}

impl<V> Clone for Box<V>
    where V: Clone
{
    fn clone(&self) -> Self {
        Box{
            origin: self.origin.clone(),
            offset: self.offset.clone(),
        }
    }
}
<|endoftext|>


<commit_before>use std::error;

use super::value::Value;

#[inline]
pub fn op_add(lhs: Value, rhs: Value) -> Result<Value, Box<dyn error::Error>> {
    match (&lhs, &rhs) {
        (Value::Integer(lhs), Value::Integer(rhs)) => {
            Ok(Value::Integer(*lhs + *rhs))
        },
        _ => panic!("Cannot add {:?} and {:?}", lhs, rhs)
    }
}

#[inline]
pub fn op_less_than(lhs: Value, rhs: Value) -> Result<Value, Box<dyn error::Error>> {
    match (&lhs, &rhs) {
        (Value::Integer(lhs), Value::Integer(rhs)) => Ok(Value::Boolean(*lhs < *rhs)),
        _ => panic!("Cannot less-than {:?} and {:?}", lhs, rhs)
    }
}

#[inline]
pub fn is_truthy(value: Value) -> Result<bool, Box<dyn error::Error>> {
    match &value {
        Value::Boolean(value) => Ok(*value),
        Value::Integer(value) => Ok(*value != 0),
        _ => Ok(true),
    }
}
<commit_msg>Revise truthiness/falsiness rule<commit_after>use std::error;

use super::value::Value;

#[inline]
pub fn op_add(lhs: Value, rhs: Value) -> Result<Value, Box<dyn error::Error>> {
    match (&lhs, &rhs) {
        (Value::Integer(lhs), Value::Integer(rhs)) => {
            Ok(Value::Integer(*lhs + *rhs))
        },
        _ => panic!("Cannot add {:?} and {:?}", lhs, rhs)
    }
}

#[inline]
pub fn op_less_than(lhs: Value, rhs: Value) -> Result<Value, Box<dyn error::Error>> {
    match (&lhs, &rhs) {
        (Value::Integer(lhs), Value::Integer(rhs)) => Ok(Value::Boolean(*lhs < *rhs)),
        _ => panic!("Cannot less-than {:?} and {:?}", lhs, rhs)
    }
}

/// Only `null` and `false` are falsy. Everything else is truthy.
#[inline]
pub fn is_truthy(value: Value) -> Result<bool, Box<dyn error::Error>> {
    match &value {
        Value::Null => Ok(false),
        Value::Boolean(value) => Ok(*value),
        _ => Ok(true),
    }
}
<|endoftext|>


<commit_before>extern crate shiplift;

use shiplift::Docker;

fn main() {
  let mut docker = Docker::new();
  for e in docker.events().get().unwrap() {
    println!("{:?}", e);
  }
}
<commit_msg>just print info for now<commit_after>extern crate shiplift;

use shiplift::Docker;

fn main() {
  let mut docker = Docker::new();
  println!("{:?}", docker.info().unwrap());
  //for e in docker.events().get().unwrap() {
  //  println!("{:?}", e);
  //}
}
<|endoftext|>


<commit_before>use unicode_width::UnicodeWidthStr;

use crate::buffer::Buffer;
use crate::layout::Rect;
use crate::style::Style;
use crate::symbols::line;
use crate::widgets::{Block, Widget};

/// A widget to display available tabs in a multiple panels context.
///
/// # Examples
///
/// ```
/// # use tui::widgets::{Block, Borders, Tabs};
/// # use tui::style::{Style, Color};
/// # use tui::symbols::{DOT};
/// # fn main() {
/// Tabs::default()
///     .block(Block::default().title("Tabs").borders(Borders::ALL))
///     .titles(&["Tab1", "Tab2", "Tab3", "Tab4"])
///     .style(Style::default().fg(Color::White))
///     .highlight_style(Style::default().fg(Color::Yellow));
///     .divider(DOT);
/// # }
/// ```
pub struct Tabs<'a, T>
where
    T: AsRef<str> + 'a,
{
    /// A block to wrap this widget in if necessary
    block: Option<Block<'a>>,
    /// One title for each tab
    titles: &'a [T],
    /// The index of the selected tabs
    selected: usize,
    /// The style used to draw the text
    style: Style,
    /// The style used to display the selected item
    highlight_style: Style,
    /// Tab divider
    divider: &'a str,
}

impl<'a, T> Default for Tabs<'a, T>
where
    T: AsRef<str>,
{
    fn default() -> Tabs<'a, T> {
        Tabs {
            block: None,
            titles: &[],
            selected: 0,
            style: Default::default(),
            highlight_style: Default::default(),
            divider: line::VERTICAL,
        }
    }
}

impl<'a, T> Tabs<'a, T>
where
    T: AsRef<str>,
{
    pub fn block(mut self, block: Block<'a>) -> Tabs<'a, T> {
        self.block = Some(block);
        self
    }

    pub fn titles(mut self, titles: &'a [T]) -> Tabs<'a, T> {
        self.titles = titles;
        self
    }

    pub fn select(mut self, selected: usize) -> Tabs<'a, T> {
        self.selected = selected;
        self
    }

    pub fn style(mut self, style: Style) -> Tabs<'a, T> {
        self.style = style;
        self
    }

    pub fn highlight_style(mut self, style: Style) -> Tabs<'a, T> {
        self.highlight_style = style;
        self
    }

    pub fn divider(mut self, divider: &'a str) -> Tabs<'a, T> {
        self.divider = divider;
        self
    }
}

impl<'a, T> Widget for Tabs<'a, T>
where
    T: AsRef<str>,
{
    fn draw(&mut self, area: Rect, buf: &mut Buffer) {
        let tabs_area = match self.block {
            Some(ref mut b) => {
                b.draw(area, buf);
                b.inner(area)
            }
            None => area,
        };

        if tabs_area.height < 1 {
            return;
        }

        self.background(&tabs_area, buf, self.style.bg);

        let mut x = tabs_area.left();
        let titles_length = self.titles.len();
        let divider_width = self.divider.chars().count() as u16;
        for (title, style, last_title) in self.titles.iter().enumerate().map(|(i, t)| {
            let lt = i + 1 == titles_length;
            if i == self.selected {
                (t, self.highlight_style, lt)
            } else {
                (t, self.style, lt)
            }
        }) {
            x += 1;
            if x > tabs_area.right() {
                break;
            } else {
                buf.set_string(x, tabs_area.top(), title.as_ref(), style);
                x += title.as_ref().width() as u16 + 1;
                if x >= tabs_area.right() || last_title {
                    break;
                } else {
                    buf.get_mut(x, tabs_area.top())
                        .set_symbol(self.divider)
                        .set_fg(self.style.fg)
                        .set_bg(self.style.bg);
                    x += divider_width;
                }
            }
        }
    }
}
<commit_msg>Fix documented example to fix doc-tests on CI<commit_after>use unicode_width::UnicodeWidthStr;

use crate::buffer::Buffer;
use crate::layout::Rect;
use crate::style::Style;
use crate::symbols::line;
use crate::widgets::{Block, Widget};

/// A widget to display available tabs in a multiple panels context.
///
/// # Examples
///
/// ```
/// # use tui::widgets::{Block, Borders, Tabs};
/// # use tui::style::{Style, Color};
/// # use tui::symbols::{DOT};
/// # fn main() {
/// Tabs::default()
///     .block(Block::default().title("Tabs").borders(Borders::ALL))
///     .titles(&["Tab1", "Tab2", "Tab3", "Tab4"])
///     .style(Style::default().fg(Color::White))
///     .highlight_style(Style::default().fg(Color::Yellow))
///     .divider(DOT);
/// # }
/// ```
pub struct Tabs<'a, T>
where
    T: AsRef<str> + 'a,
{
    /// A block to wrap this widget in if necessary
    block: Option<Block<'a>>,
    /// One title for each tab
    titles: &'a [T],
    /// The index of the selected tabs
    selected: usize,
    /// The style used to draw the text
    style: Style,
    /// The style used to display the selected item
    highlight_style: Style,
    /// Tab divider
    divider: &'a str,
}

impl<'a, T> Default for Tabs<'a, T>
where
    T: AsRef<str>,
{
    fn default() -> Tabs<'a, T> {
        Tabs {
            block: None,
            titles: &[],
            selected: 0,
            style: Default::default(),
            highlight_style: Default::default(),
            divider: line::VERTICAL,
        }
    }
}

impl<'a, T> Tabs<'a, T>
where
    T: AsRef<str>,
{
    pub fn block(mut self, block: Block<'a>) -> Tabs<'a, T> {
        self.block = Some(block);
        self
    }

    pub fn titles(mut self, titles: &'a [T]) -> Tabs<'a, T> {
        self.titles = titles;
        self
    }

    pub fn select(mut self, selected: usize) -> Tabs<'a, T> {
        self.selected = selected;
        self
    }

    pub fn style(mut self, style: Style) -> Tabs<'a, T> {
        self.style = style;
        self
    }

    pub fn highlight_style(mut self, style: Style) -> Tabs<'a, T> {
        self.highlight_style = style;
        self
    }

    pub fn divider(mut self, divider: &'a str) -> Tabs<'a, T> {
        self.divider = divider;
        self
    }
}

impl<'a, T> Widget for Tabs<'a, T>
where
    T: AsRef<str>,
{
    fn draw(&mut self, area: Rect, buf: &mut Buffer) {
        let tabs_area = match self.block {
            Some(ref mut b) => {
                b.draw(area, buf);
                b.inner(area)
            }
            None => area,
        };

        if tabs_area.height < 1 {
            return;
        }

        self.background(&tabs_area, buf, self.style.bg);

        let mut x = tabs_area.left();
        let titles_length = self.titles.len();
        let divider_width = self.divider.chars().count() as u16;
        for (title, style, last_title) in self.titles.iter().enumerate().map(|(i, t)| {
            let lt = i + 1 == titles_length;
            if i == self.selected {
                (t, self.highlight_style, lt)
            } else {
                (t, self.style, lt)
            }
        }) {
            x += 1;
            if x > tabs_area.right() {
                break;
            } else {
                buf.set_string(x, tabs_area.top(), title.as_ref(), style);
                x += title.as_ref().width() as u16 + 1;
                if x >= tabs_area.right() || last_title {
                    break;
                } else {
                    buf.get_mut(x, tabs_area.top())
                        .set_symbol(self.divider)
                        .set_fg(self.style.fg)
                        .set_bg(self.style.bg);
                    x += divider_width;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::OpenOptions;
use std::io::{Read, Write, BufReader, BufRead, Seek, SeekFrom};
use std::path::PathBuf;

use errors::DecodeError;
use crc32;
use constants::{NUL, CR, LF, ESCAPE, SPACE, DEFAULT_LINE_SIZE};

#[derive(Default, Debug)]
struct MetaData {
    name: Option<String>,
    line_length: Option<u16>,
    size: Option<usize>,
    crc32: Option<u32>,
    part: Option<u32>,
    begin: Option<usize>,
    end: Option<usize>,
}

/// Decodes the input file in a new output file.
///
/// If ok, returns the path of the decoded file.
///
/// # Example
/// ```rust,no_run
/// yenc::ydecode_file("test2.bin.yenc", "test2.bin");
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn ydecode_file(input_filename: &str, output_path: &str) -> Result<String, DecodeError> {
    let mut input_file = OpenOptions::new().read(true).open(input_filename)?;
    ydecode_stream(&mut input_file, output_path)
}

/// Decodes the data from a stream in a new output file.
///
/// Writes the output to a file with the filename from the header line, and places it in the
/// output path. The path of the output file is returned.
pub fn ydecode_stream<R>(read_stream: &mut R, output_path: &str) -> Result<String, DecodeError>
where
    R: Read,
{
    let mut rdr = BufReader::new(read_stream);
    let mut output_pathbuf = PathBuf::new();
    output_pathbuf.push(output_path);

    let mut checksum = crc32::Crc32::new();
    let mut yenc_block_found = false;
    let mut metadata: MetaData = Default::default();

    while !yenc_block_found {
        let mut line_buf = Vec::<u8>::with_capacity(2 * DEFAULT_LINE_SIZE as usize);
        let length = rdr.read_until(LF, &mut line_buf)?;
        if length == 0 {
            break;
        }
        if line_buf.starts_with(b"=ybegin ") {
            yenc_block_found = true;
            // parse header line and determine output filename
            metadata = parse_header_line(&line_buf, 8)?;
            output_pathbuf.push(metadata.name.unwrap().to_string().trim());
        }
    }

    if yenc_block_found {
        let mut output_file = OpenOptions::new().create(true).write(true).open(
            output_pathbuf
                .as_path(),
        )?;

        let mut footer_found = false;
        while !footer_found {
            let mut line_buf = Vec::<u8>::with_capacity(2 * DEFAULT_LINE_SIZE as usize);
            let length = rdr.read_until(LF, &mut line_buf)?;
            if length == 0 {
                break;
            }
            if line_buf.starts_with(b"=ypart ") {
                let part_metadata = parse_header_line(&line_buf, 7)?;
                metadata.part = part_metadata.part;
                metadata.begin = part_metadata.begin;
                metadata.end = part_metadata.end;
                if let Some(begin) = metadata.begin {
                    output_file.seek(SeekFrom::Start((begin - 1) as u64))?;
                }
            } else if line_buf.starts_with(b"=yend ") {
                footer_found = true;
                let mm = parse_header_line(&line_buf, 6)?;
                metadata.size = mm.size;
                metadata.crc32 = mm.crc32;
            } else {
                let decoded = ydecode_buffer(&line_buf[0..length])?;
                checksum.update_with_slice(decoded.as_slice());
                output_file.write_all(decoded.as_slice())?;
            }
        }
        if footer_found {
            if let Some(expected_size) = metadata.size {
                if expected_size != checksum.num_bytes {
                    return Err(DecodeError::IncompleteData {
                        expected_size: expected_size,
                        actual_size: checksum.num_bytes,
                    });
                }
            }
            if let Some(expected_crc) = metadata.crc32 {
                if expected_crc != checksum.crc {
                    return Err(DecodeError::InvalidChecksum);
                }
            }
        } else if let Some(expected_size) = metadata.size {
            if expected_size != checksum.num_bytes {
                return Err(DecodeError::IncompleteData {
                    expected_size: expected_size,
                    actual_size: checksum.num_bytes,
                });
            }
        }

    }
    Ok(output_pathbuf.to_str().unwrap().to_string())
}

/// Decode the yEncoded byte slice into a vector of bytes.
///
/// Carriage Return (CR) and Line Feed (LF) are ignored.
pub fn ydecode_buffer(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    let mut output = Vec::<u8>::with_capacity((input.len() as f64 * 1.02) as usize);
    let mut iter = input.iter();
    while let Some(byte) = iter.next() {
        let mut byte = *byte;
        match byte {
            NUL | CR | LF => {
                // for now, just continue
                continue;
            }
            ESCAPE => {
                match iter.next() {
                    Some(b) => {
                        byte = b.overflowing_sub(64).0;
                    }
                    None => {
                        // for now, just continue
                        continue;
                    }
                }
            }
            _ => {}
        }
        output.push(byte.overflowing_sub(42).0);
    }
    Ok(output)
}

fn parse_header_line(line_buf: &[u8], offset: usize) -> Result<MetaData, DecodeError> {
    #[derive(Debug)]
    enum State {
        Keyword,
        Value,
        End,
    };

    let mut metadata: MetaData = Default::default();
    let mut state = State::Keyword;

    let mut keyword: &[u8] = &[];
    let mut keyword_start_idx: Option<usize> = None;
    let mut value: &[u8] = &[];
    let mut value_start_idx: Option<usize> = None;

    let header_line = String::from_utf8_lossy(line_buf).to_string();
    for (i, &c) in line_buf[offset..].iter().enumerate() {
        let position = i + offset;
        match state {
            State::End => unreachable!(),
            State::Keyword => {
                match c {
                    b'a'...b'z' | b'0'...b'9' => {
                        if keyword_start_idx.is_none() {
                            keyword_start_idx = Some(position);
                        }
                        keyword = &line_buf[keyword_start_idx.unwrap()..position + 1];
                    }
                    b'=' => {
                        if keyword.is_empty() || !is_known_keyword(keyword) {
                            return Err(DecodeError::InvalidHeader {
                                line: header_line,
                                position: position,
                            });
                        } else {
                            state = State::Value;
                        }
                    }
                    CR | LF => {}
                    _ => {
                        return Err(DecodeError::InvalidHeader {
                            line: header_line,
                            position: position,
                        });
                    }
                }
            }
            State::Value => {
                match keyword {
                    b"name" => {
                        match c {
                            CR => {}
                            LF => {
                                state = State::End;
                                metadata.name = Some(String::from_utf8_lossy(value).to_string());
                            }
                            _ => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                        }
                    }
                    b"size" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.size =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"begin" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.begin =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"end" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE | LF | CR => {
                                metadata.end =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"line" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.line_length =
                                    Some(
                                        u16::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"part" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.part =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"crc32" | b"pcrc32" => {
                        match c {
                            b'0'...b'9' | b'A'...b'F' | b'a'...b'f' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                state = State::Keyword;
                                metadata.crc32 =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 16)
                                            .unwrap(),
                                    );
                                value_start_idx = None;
                            }
                            LF => {
                                state = State::End;
                                metadata.crc32 =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 16)
                                            .unwrap(),
                                    );
                                value_start_idx = None;
                            }
                            CR => {}
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    _ => unreachable!(),
                }
            }
        };
    }
    Ok(metadata)
}

fn is_known_keyword(keyword_slice: &[u8]) -> bool {
    keyword_slice == b"name" || keyword_slice == b"line" || keyword_slice == b"size" ||
        keyword_slice == b"part" ||
        keyword_slice == b"begin" || keyword_slice == b"end" || keyword_slice == b"pcrc32" ||
        keyword_slice == b"crc32"
}

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

    #[test]
    fn decode_invalid() {
        assert!(ydecode_buffer(&[b'=']).unwrap().is_empty());
    }

    #[test]
    fn decode_valid_ff() {
        assert_eq!(&vec![0xff - 0x2A], &ydecode_buffer(&[0xff]).unwrap());
    }

    #[test]
    fn decode_valid_01() {
        assert_eq!(&vec![0xff - 0x28], &ydecode_buffer(&[0x01]).unwrap());
    }

    #[test]
    fn decode_valid_esc_ff() {
        assert_eq!(
            &vec![0xff - 0x40 - 0x2A],
            &ydecode_buffer(&[b'=', 0xff]).unwrap()
        );
    }

    #[test]
    fn decode_valid_esc_01() {
        assert_eq!(
            &vec![0xff - 0x40 - 0x2A + 2],
            &ydecode_buffer(&[b'=', 0x01]).unwrap()
        );
    }

    #[test]
    fn parse_valid_header_begin() {
        let parse_result = parse_header_line(
            b"=ybegin part=1 line=128 size=189463 name=CatOnKeyboardInSpace001.jpg\n",
            8,
        );
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.part, Some(1));
        assert_eq!(metadata.size, Some(189463));
        assert_eq!(metadata.line_length, Some(128));
        assert_eq!(
            metadata.name,
            Some("CatOnKeyboardInSpace001.jpg".to_string())
        );
    }

    #[test]
    fn parse_valid_header_part() {
        let parse_result = parse_header_line(b"=ypart begin=1 end=189463\n", 7);
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.begin, Some(1));
        assert_eq!(metadata.end, Some(189463));
    }

    #[test]
    fn parse_valid_footer_end() {
        let parse_result = parse_header_line(b"=yend size=26624 part=1 pcrc32=ae052b48\n", 6);
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.part, Some(1));
        assert_eq!(metadata.size, Some(26624));
        assert_eq!(metadata.crc32, Some(0xae052b48));
    }
}
<commit_msg>Use match instead of complex if.<commit_after>use std::fs::OpenOptions;
use std::io::{Read, Write, BufReader, BufRead, Seek, SeekFrom};
use std::path::PathBuf;

use errors::DecodeError;
use crc32;
use constants::{NUL, CR, LF, ESCAPE, SPACE, DEFAULT_LINE_SIZE};

#[derive(Default, Debug)]
struct MetaData {
    name: Option<String>,
    line_length: Option<u16>,
    size: Option<usize>,
    crc32: Option<u32>,
    part: Option<u32>,
    begin: Option<usize>,
    end: Option<usize>,
}

/// Decodes the input file in a new output file.
///
/// If ok, returns the path of the decoded file.
///
/// # Example
/// ```rust,no_run
/// yenc::ydecode_file("test2.bin.yenc", "test2.bin");
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn ydecode_file(input_filename: &str, output_path: &str) -> Result<String, DecodeError> {
    let mut input_file = OpenOptions::new().read(true).open(input_filename)?;
    ydecode_stream(&mut input_file, output_path)
}

/// Decodes the data from a stream in a new output file.
///
/// Writes the output to a file with the filename from the header line, and places it in the
/// output path. The path of the output file is returned.
pub fn ydecode_stream<R>(read_stream: &mut R, output_path: &str) -> Result<String, DecodeError>
where
    R: Read,
{
    let mut rdr = BufReader::new(read_stream);
    let mut output_pathbuf = PathBuf::new();
    output_pathbuf.push(output_path);

    let mut checksum = crc32::Crc32::new();
    let mut yenc_block_found = false;
    let mut metadata: MetaData = Default::default();

    while !yenc_block_found {
        let mut line_buf = Vec::<u8>::with_capacity(2 * DEFAULT_LINE_SIZE as usize);
        let length = rdr.read_until(LF, &mut line_buf)?;
        if length == 0 {
            break;
        }
        if line_buf.starts_with(b"=ybegin ") {
            yenc_block_found = true;
            // parse header line and determine output filename
            metadata = parse_header_line(&line_buf, 8)?;
            output_pathbuf.push(metadata.name.unwrap().to_string().trim());
        }
    }

    if yenc_block_found {
        let mut output_file = OpenOptions::new().create(true).write(true).open(
            output_pathbuf
                .as_path(),
        )?;

        let mut footer_found = false;
        while !footer_found {
            let mut line_buf = Vec::<u8>::with_capacity(2 * DEFAULT_LINE_SIZE as usize);
            let length = rdr.read_until(LF, &mut line_buf)?;
            if length == 0 {
                break;
            }
            if line_buf.starts_with(b"=ypart ") {
                let part_metadata = parse_header_line(&line_buf, 7)?;
                metadata.part = part_metadata.part;
                metadata.begin = part_metadata.begin;
                metadata.end = part_metadata.end;
                if let Some(begin) = metadata.begin {
                    output_file.seek(SeekFrom::Start((begin - 1) as u64))?;
                }
            } else if line_buf.starts_with(b"=yend ") {
                footer_found = true;
                let mm = parse_header_line(&line_buf, 6)?;
                metadata.size = mm.size;
                metadata.crc32 = mm.crc32;
            } else {
                let decoded = ydecode_buffer(&line_buf[0..length])?;
                checksum.update_with_slice(decoded.as_slice());
                output_file.write_all(decoded.as_slice())?;
            }
        }
        if footer_found {
            if let Some(expected_size) = metadata.size {
                if expected_size != checksum.num_bytes {
                    return Err(DecodeError::IncompleteData {
                        expected_size: expected_size,
                        actual_size: checksum.num_bytes,
                    });
                }
            }
            if let Some(expected_crc) = metadata.crc32 {
                if expected_crc != checksum.crc {
                    return Err(DecodeError::InvalidChecksum);
                }
            }
        } else if let Some(expected_size) = metadata.size {
            if expected_size != checksum.num_bytes {
                return Err(DecodeError::IncompleteData {
                    expected_size: expected_size,
                    actual_size: checksum.num_bytes,
                });
            }
        }

    }
    Ok(output_pathbuf.to_str().unwrap().to_string())
}

/// Decode the yEncoded byte slice into a vector of bytes.
///
/// Carriage Return (CR) and Line Feed (LF) are ignored.
pub fn ydecode_buffer(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    let mut output = Vec::<u8>::with_capacity((input.len() as f64 * 1.02) as usize);
    let mut iter = input.iter();
    while let Some(byte) = iter.next() {
        let mut byte = *byte;
        match byte {
            NUL | CR | LF => {
                // for now, just continue
                continue;
            }
            ESCAPE => {
                match iter.next() {
                    Some(b) => {
                        byte = b.overflowing_sub(64).0;
                    }
                    None => {
                        // for now, just continue
                        continue;
                    }
                }
            }
            _ => {}
        }
        output.push(byte.overflowing_sub(42).0);
    }
    Ok(output)
}

fn parse_header_line(line_buf: &[u8], offset: usize) -> Result<MetaData, DecodeError> {
    #[derive(Debug)]
    enum State {
        Keyword,
        Value,
        End,
    };

    let mut metadata: MetaData = Default::default();
    let mut state = State::Keyword;

    let mut keyword: &[u8] = &[];
    let mut keyword_start_idx: Option<usize> = None;
    let mut value: &[u8] = &[];
    let mut value_start_idx: Option<usize> = None;

    let header_line = String::from_utf8_lossy(line_buf).to_string();
    for (i, &c) in line_buf[offset..].iter().enumerate() {
        let position = i + offset;
        match state {
            State::End => unreachable!(),
            State::Keyword => {
                match c {
                    b'a'...b'z' | b'0'...b'9' => {
                        if keyword_start_idx.is_none() {
                            keyword_start_idx = Some(position);
                        }
                        keyword = &line_buf[keyword_start_idx.unwrap()..position + 1];
                    }
                    b'=' => {
                        if keyword.is_empty() || !is_known_keyword(keyword) {
                            return Err(DecodeError::InvalidHeader {
                                line: header_line,
                                position: position,
                            });
                        } else {
                            state = State::Value;
                        }
                    }
                    CR | LF => {}
                    _ => {
                        return Err(DecodeError::InvalidHeader {
                            line: header_line,
                            position: position,
                        });
                    }
                }
            }
            State::Value => {
                match keyword {
                    b"name" => {
                        match c {
                            CR => {}
                            LF => {
                                state = State::End;
                                metadata.name = Some(String::from_utf8_lossy(value).to_string());
                            }
                            _ => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                        }
                    }
                    b"size" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.size =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"begin" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.begin =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"end" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE | LF | CR => {
                                metadata.end =
                                    Some(
                                        usize::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"line" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.line_length =
                                    Some(
                                        u16::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"part" => {
                        match c {
                            b'0'...b'9' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                metadata.part =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 10)
                                            .unwrap(),
                                    );
                                state = State::Keyword;
                                keyword_start_idx = None;
                                value_start_idx = None;
                            }
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    b"crc32" | b"pcrc32" => {
                        match c {
                            b'0'...b'9' | b'A'...b'F' | b'a'...b'f' => {
                                if value_start_idx.is_none() {
                                    value_start_idx = Some(position);
                                }
                                value = &line_buf[value_start_idx.unwrap()..position + 1];
                            }
                            SPACE => {
                                state = State::Keyword;
                                metadata.crc32 =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 16)
                                            .unwrap(),
                                    );
                                value_start_idx = None;
                            }
                            LF => {
                                state = State::End;
                                metadata.crc32 =
                                    Some(
                                        u32::from_str_radix(&String::from_utf8_lossy(value), 16)
                                            .unwrap(),
                                    );
                                value_start_idx = None;
                            }
                            CR => {}
                            _ => {
                                return Err(DecodeError::InvalidHeader {
                                    line: header_line,
                                    position: position,
                                });
                            }
                        }
                    }
                    _ => unreachable!(),
                }
            }
        };
    }
    Ok(metadata)
}

fn is_known_keyword(keyword_slice: &[u8]) -> bool {
    match keyword_slice {
        b"begin" | b"crc32" | b"end" | b"line"| b"name" | b"part" | b"pcrc32" | b"size" => true,
        _ => false
    }
}

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

    #[test]
    fn decode_invalid() {
        assert!(ydecode_buffer(&[b'=']).unwrap().is_empty());
    }

    #[test]
    fn decode_valid_ff() {
        assert_eq!(&vec![0xff - 0x2A], &ydecode_buffer(&[0xff]).unwrap());
    }

    #[test]
    fn decode_valid_01() {
        assert_eq!(&vec![0xff - 0x28], &ydecode_buffer(&[0x01]).unwrap());
    }

    #[test]
    fn decode_valid_esc_ff() {
        assert_eq!(
            &vec![0xff - 0x40 - 0x2A],
            &ydecode_buffer(&[b'=', 0xff]).unwrap()
        );
    }

    #[test]
    fn decode_valid_esc_01() {
        assert_eq!(
            &vec![0xff - 0x40 - 0x2A + 2],
            &ydecode_buffer(&[b'=', 0x01]).unwrap()
        );
    }

    #[test]
    fn parse_valid_header_begin() {
        let parse_result = parse_header_line(
            b"=ybegin part=1 line=128 size=189463 name=CatOnKeyboardInSpace001.jpg\n",
            8,
        );
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.part, Some(1));
        assert_eq!(metadata.size, Some(189463));
        assert_eq!(metadata.line_length, Some(128));
        assert_eq!(
            metadata.name,
            Some("CatOnKeyboardInSpace001.jpg".to_string())
        );
    }

    #[test]
    fn parse_valid_header_part() {
        let parse_result = parse_header_line(b"=ypart begin=1 end=189463\n", 7);
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.begin, Some(1));
        assert_eq!(metadata.end, Some(189463));
    }

    #[test]
    fn parse_valid_footer_end() {
        let parse_result = parse_header_line(b"=yend size=26624 part=1 pcrc32=ae052b48\n", 6);
        assert!(parse_result.is_ok());
        let metadata = parse_result.unwrap();
        assert_eq!(metadata.part, Some(1));
        assert_eq!(metadata.size, Some(26624));
        assert_eq!(metadata.crc32, Some(0xae052b48));
    }
}
<|endoftext|>


<commit_before>use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{Hash, Hasher};

use bit_vec::BitVec;

struct HashIter<'a, T>
where
    T: 'a,
{
    m: usize,
    k: usize,
    i: usize,
    obj: &'a T,
}

impl<'a, T> HashIter<'a, T>
where
    T: 'a,
{
    fn new(m: usize, k: usize, obj: &'a T) -> HashIter<'a, T> {
        HashIter {
            m: m,
            k: k,
            i: 0,
            obj: obj,
        }
    }
}

impl<'a, T> Iterator for HashIter<'a, T>
where
    T: 'a + Hash,
{
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        if self.i < self.k {
            let mut hasher = DefaultHasher::new();
            hasher.write_usize(self.i);
            self.obj.hash(&mut hasher);
            let x = (hasher.finish() as usize) % self.m;

            self.i += 1;

            Some(x)
        } else {
            None
        }
    }
}

/// Simple implementation of a [BloomFilter](https://en.wikipedia.org/wiki/Bloom_filter)
#[derive(Clone)]
pub struct BloomFilter {
    bv: BitVec,
    k: usize,
}

impl BloomFilter {
    /// Create new, empty BloomFilter with internal parameters.
    ///
    /// - `k` is the number of hash functions
    /// - `m` is the number of bits used to store state
    pub fn with_params(m: usize, k: usize) -> BloomFilter {
        BloomFilter {
            bv: BitVec::from_elem(m, false),
            k: k,
        }
    }

    /// Create new, empty BloomFilter with given properties.
    ///
    /// - `n` number of unique elements the BloomFilter is expected to hold, must be `> 0`
    /// - `p` false positive property when querying the BloomFilter after adding `n` unique
    ///   elements, must be `> 0` and `< 1`
    ///
    /// Panics if the parameters are not in range.
    pub fn with_properties(n: usize, p: f64) -> BloomFilter {
        assert!(n > 0);
        assert!(p > 0.);
        assert!(p < 1.);

        let k = (-p.log2()) as usize;
        let ln2 = (2f64).ln();
        let m = (-((n as f64) * p.ln()) / (ln2 * ln2)) as usize;

        BloomFilter::with_params(m, k)
    }

    /// Get `k` (number of hash functions).
    pub fn get_k(&self) -> usize {
        self.k
    }

    /// Get `m` (number of stored bits).
    pub fn get_m(&self) -> usize {
        self.bv.len()
    }

    /// Add new element to the BloomFilter.
    ///
    /// If the same element is added multiple times or if an element results in the same hash
    /// signature, this method does not have any effect.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            self.bv.set(pos, true);
        }
    }

    /// Guess if the given element was added to the BloomFilter.
    pub fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            if !self.bv.get(pos).unwrap() {
                return false;
            }
        }
        true
    }

    /// Clear state of the BloomFilter, so that it behaves like a fresh one.
    pub fn clear(&mut self) {
        self.bv.clear()
    }

    /// Check whether the BloomFilter is empty.
    pub fn is_empty(&self) -> bool {
        !self.bv.iter().any(|x| x)
    }

    /// Add the entire content of another bloomfilter to this BloomFilter.
    ///
    /// The result is the same as adding all elements added to `other` to `self` in the first
    /// place.
    ///
    /// Panics if `k` and `m` of the two BloomFilters are not identical.
    pub fn union(&mut self, other: &BloomFilter) {
        assert!(self.k == other.k);
        assert!(self.bv.len() == other.bv.len());
        self.bv.union(&other.bv);
    }

    /// Guess the number of unique elements added to the BloomFilter.
    pub fn guess_n(&self) -> usize {
        let m = self.bv.len() as f64;
        let k = self.k as f64;
        let x = self.bv.iter().filter(|x| *x).count() as f64;

        (-m / k * (1. - x / m).ln()) as usize
    }
}

impl fmt::Debug for BloomFilter {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "BloomFilter {{ m: {}, k: {} }}", self.bv.len(), self.k)
    }
}

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

    #[test]
    fn getter() {
        let bf = BloomFilter::with_params(100, 2);
        assert_eq!(bf.get_k(), 2);
        assert_eq!(bf.get_m(), 100);
    }

    #[test]
    fn empty() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(!bf.query(&1));
    }

    #[test]
    fn add() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        assert!(bf.query(&1));
        assert!(!bf.query(&2));
    }

    #[test]
    fn clear() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        bf.clear();
        assert!(!bf.query(&1));
    }

    #[test]
    fn is_empty() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert!(bf.is_empty());

        bf.add(&1);
        assert!(!bf.is_empty());

        bf.clear();
        assert!(bf.is_empty());
    }

    #[test]
    fn clone() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);

        let bf2 = bf1.clone();
        bf1.add(&2);
        assert!(bf2.query(&1));
        assert!(!bf2.query(&2));
    }

    #[test]
    fn union() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);
        assert!(bf1.query(&1));
        assert!(!bf1.query(&2));
        assert!(!bf1.query(&3));

        let mut bf2 = BloomFilter::with_params(100, 2);
        bf2.add(&2);
        assert!(!bf2.query(&1));
        assert!(bf2.query(&2));
        assert!(!bf2.query(&3));

        bf1.union(&bf2);
        assert!(bf1.query(&1));
        assert!(bf1.query(&2));
        assert!(!bf1.query(&3));
    }

    #[test]
    fn with_properties() {
        let bf = BloomFilter::with_properties(1000, 0.1);
        assert_eq!(bf.get_k(), 3);
        assert_eq!(bf.get_m(), 4792);
    }

    #[test]
    fn guess_n() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert_eq!(bf.guess_n(), 0);

        bf.add(&1);
        assert_eq!(bf.guess_n(), 1);

        bf.add(&1);
        assert_eq!(bf.guess_n(), 1);

        bf.add(&2);
        assert_eq!(bf.guess_n(), 2);
    }

    #[test]
    fn debug() {
        let bf = BloomFilter::with_params(100, 2);
        assert_eq!(format!("{:?}", bf), "BloomFilter { m: 100, k: 2 }");
    }
}
<commit_msg>bloomfilter: use more assert_eq<commit_after>use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{Hash, Hasher};

use bit_vec::BitVec;

struct HashIter<'a, T>
where
    T: 'a,
{
    m: usize,
    k: usize,
    i: usize,
    obj: &'a T,
}

impl<'a, T> HashIter<'a, T>
where
    T: 'a,
{
    fn new(m: usize, k: usize, obj: &'a T) -> HashIter<'a, T> {
        HashIter {
            m: m,
            k: k,
            i: 0,
            obj: obj,
        }
    }
}

impl<'a, T> Iterator for HashIter<'a, T>
where
    T: 'a + Hash,
{
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        if self.i < self.k {
            let mut hasher = DefaultHasher::new();
            hasher.write_usize(self.i);
            self.obj.hash(&mut hasher);
            let x = (hasher.finish() as usize) % self.m;

            self.i += 1;

            Some(x)
        } else {
            None
        }
    }
}

/// Simple implementation of a [BloomFilter](https://en.wikipedia.org/wiki/Bloom_filter)
#[derive(Clone)]
pub struct BloomFilter {
    bv: BitVec,
    k: usize,
}

impl BloomFilter {
    /// Create new, empty BloomFilter with internal parameters.
    ///
    /// - `k` is the number of hash functions
    /// - `m` is the number of bits used to store state
    pub fn with_params(m: usize, k: usize) -> BloomFilter {
        BloomFilter {
            bv: BitVec::from_elem(m, false),
            k: k,
        }
    }

    /// Create new, empty BloomFilter with given properties.
    ///
    /// - `n` number of unique elements the BloomFilter is expected to hold, must be `> 0`
    /// - `p` false positive property when querying the BloomFilter after adding `n` unique
    ///   elements, must be `> 0` and `< 1`
    ///
    /// Panics if the parameters are not in range.
    pub fn with_properties(n: usize, p: f64) -> BloomFilter {
        assert!(n > 0);
        assert!(p > 0.);
        assert!(p < 1.);

        let k = (-p.log2()) as usize;
        let ln2 = (2f64).ln();
        let m = (-((n as f64) * p.ln()) / (ln2 * ln2)) as usize;

        BloomFilter::with_params(m, k)
    }

    /// Get `k` (number of hash functions).
    pub fn get_k(&self) -> usize {
        self.k
    }

    /// Get `m` (number of stored bits).
    pub fn get_m(&self) -> usize {
        self.bv.len()
    }

    /// Add new element to the BloomFilter.
    ///
    /// If the same element is added multiple times or if an element results in the same hash
    /// signature, this method does not have any effect.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            self.bv.set(pos, true);
        }
    }

    /// Guess if the given element was added to the BloomFilter.
    pub fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            if !self.bv.get(pos).unwrap() {
                return false;
            }
        }
        true
    }

    /// Clear state of the BloomFilter, so that it behaves like a fresh one.
    pub fn clear(&mut self) {
        self.bv.clear()
    }

    /// Check whether the BloomFilter is empty.
    pub fn is_empty(&self) -> bool {
        !self.bv.iter().any(|x| x)
    }

    /// Add the entire content of another bloomfilter to this BloomFilter.
    ///
    /// The result is the same as adding all elements added to `other` to `self` in the first
    /// place.
    ///
    /// Panics if `k` and `m` of the two BloomFilters are not identical.
    pub fn union(&mut self, other: &BloomFilter) {
        assert_eq!(self.k, other.k);
        assert_eq!(self.bv.len(), other.bv.len());

        self.bv.union(&other.bv);
    }

    /// Guess the number of unique elements added to the BloomFilter.
    pub fn guess_n(&self) -> usize {
        let m = self.bv.len() as f64;
        let k = self.k as f64;
        let x = self.bv.iter().filter(|x| *x).count() as f64;

        (-m / k * (1. - x / m).ln()) as usize
    }
}

impl fmt::Debug for BloomFilter {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "BloomFilter {{ m: {}, k: {} }}", self.bv.len(), self.k)
    }
}

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

    #[test]
    fn getter() {
        let bf = BloomFilter::with_params(100, 2);
        assert_eq!(bf.get_k(), 2);
        assert_eq!(bf.get_m(), 100);
    }

    #[test]
    fn empty() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(!bf.query(&1));
    }

    #[test]
    fn add() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        assert!(bf.query(&1));
        assert!(!bf.query(&2));
    }

    #[test]
    fn clear() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        bf.clear();
        assert!(!bf.query(&1));
    }

    #[test]
    fn is_empty() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert!(bf.is_empty());

        bf.add(&1);
        assert!(!bf.is_empty());

        bf.clear();
        assert!(bf.is_empty());
    }

    #[test]
    fn clone() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);

        let bf2 = bf1.clone();
        bf1.add(&2);
        assert!(bf2.query(&1));
        assert!(!bf2.query(&2));
    }

    #[test]
    fn union() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);
        assert!(bf1.query(&1));
        assert!(!bf1.query(&2));
        assert!(!bf1.query(&3));

        let mut bf2 = BloomFilter::with_params(100, 2);
        bf2.add(&2);
        assert!(!bf2.query(&1));
        assert!(bf2.query(&2));
        assert!(!bf2.query(&3));

        bf1.union(&bf2);
        assert!(bf1.query(&1));
        assert!(bf1.query(&2));
        assert!(!bf1.query(&3));
    }

    #[test]
    fn with_properties() {
        let bf = BloomFilter::with_properties(1000, 0.1);
        assert_eq!(bf.get_k(), 3);
        assert_eq!(bf.get_m(), 4792);
    }

    #[test]
    fn guess_n() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert_eq!(bf.guess_n(), 0);

        bf.add(&1);
        assert_eq!(bf.guess_n(), 1);

        bf.add(&1);
        assert_eq!(bf.guess_n(), 1);

        bf.add(&2);
        assert_eq!(bf.guess_n(), 2);
    }

    #[test]
    fn debug() {
        let bf = BloomFilter::with_params(100, 2);
        assert_eq!(format!("{:?}", bf), "BloomFilter { m: 100, k: 2 }");
    }
}
<|endoftext|>


<commit_before>use consts::*;
use proto::*;
use io::ZkIo;
use listeners::{ListenerSet, Subscription};
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::io::Read;
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>,
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf,
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>,
    listeners: ListenerSet<ZkState>,
}

impl ZooKeeper {
    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
        where F: FnOnce() + Send + 'static
    {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_| ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
        where W: Watcher + 'static
    {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher, chroot.clone());
        let listeners = ListenerSet::<ZkState>::new();
        let listeners1 = listeners.clone();
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender(), listeners1);
        let sender = io.sender();

        try!(Self::zk_thread("event", move || watch.run().unwrap()));
        try!(Self::zk_thread("io", move || io.run().unwrap()));

        Ok(ZooKeeper {
            chroot: chroot,
            xid: AtomicIsize::new(1),
            io: sender,
            listeners: listeners,
        })
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start),
                }
            }
            None => (None, connect_string.len()),
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.trim().to_socket_addrs() {
                Ok(mut addrs) => {
                    match addrs.nth(0) {
                        Some(addr) => addr,
                        None => return Err(ZkError::BadArguments),
                    }
                }
                Err(_) => return Err(ZkError::BadArguments),
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self,
                                             opcode: OpCode,
                                             xid: i32,
                                             req: Req,
                                             watch: Option<Watch>)
                                             -> ZkResult<Resp> {
        let rh = RequestHeader {
            xid: xid,
            opcode: opcode,
        };
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_| ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest {
            opcode: opcode,
            data: buf,
            listener: Some(resp_tx),
            watch: watch,
        };

        try!(self.io.send(request).map_err(|err| {
            warn!("error sending request: {:?}", err);
            ZkError::ConnectionLoss
        }));

        let mut response = try!(resp_rx.recv().map_err(|err| {
            warn!("error receiving response: {:?}", err);
            ZkError::ConnectionLoss
        }));

        match response.header.err {
            0 => {
                Ok(try!(ReadFrom::read_from(&mut response.data)
                            .map_err(|_| ZkError::MarshallingError)))
            }
            e => Err(FromPrimitive::from_i32(e).unwrap()),
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => {
                if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
            }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => {
                match path {
                    "/" => Ok(chroot.clone()),
                    path => Ok(chroot.clone() + try!(Self::validate_path(path))),
                }
            }
            None => Ok(try!(Self::validate_path(path)).to_owned()),
        }
    }

    fn cut_chroot(&self, path: String) -> String {
        if let Some(ref chroot) = self.chroot {
            path[chroot.len()..].to_owned()
        } else {
            path
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest {
            typ: 0,
            scheme: scheme.to_owned(),
            auth: auth,
        };

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self,
                  path: &str,
                  data: Vec<u8>,
                  acl: Vec<Acl>,
                  mode: CreateMode)
                  -> ZkResult<String> {
        let req = CreateRequest {
            path: try!(self.path(path)),
            data: data,
            acl: acl,
            flags: mode as i32,
        };

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(self.cut_chroot(response.path))
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest {
            path: try!(self.path(path)),
            version: version,
        };

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e),
        }
    }

    pub fn exists_w<W: Watcher + 'static>(&self,
                                                               path: &str,
                                                               watcher: W)
                                                               -> ZkResult<Stat> {
        let req = ExistsRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Exist,
            watcher: Box::new(watcher),
        };

        let response: ExistsResponse = try!(self.request(OpCode::Exists,
                                                         self.xid(),
                                                         req,
                                                         Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest { path: try!(self.path(path)) };

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: Watcher + 'static>(&self,
                                                                     path: &str,
                                                                     watcher: W)
                                                                     -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Child,
            watcher: Box::new(watcher),
        };

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren,
                                                              self.xid(),
                                                              req,
                                                              Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren,
                                                              self.xid(),
                                                              req,
                                                              None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: Watcher + 'static>(&self,
                                                                 path: &str,
                                                                 watcher: W)
                                                                 -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Data,
            watcher: Box::new(watcher),
        };

        let response: GetDataResponse = try!(self.request(OpCode::GetData,
                                                          self.xid(),
                                                          req,
                                                          Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest {
            path: try!(self.path(path)),
            acl: acl,
            version: version,
        };

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest {
            path: try!(self.path(path)),
            data: data,
            version: version,
        };

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn add_listener<Listener: Fn(ZkState) + Send + 'static>(&self,
                                                                listener: Listener)
                                                                -> Subscription {
        self.listeners.subscribe(listener)
    }

    pub fn remove_listener(&self, sub: Subscription) {
        self.listeners.unsubscribe(sub);
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        if let Err(err) = self.close() {
            info!("error closing zookeeper connection in drop: {:?}", err);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky on Travis, it works on my Linux box though.
    #[test]
    #[cfg(target_os = "macos")]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos")
                                  .ok()
                                  .expect("Parse 1");
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, Some("/mesos".to_owned()));

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").ok().expect("Parse 2");
        assert_eq!(addrs,
                   vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").ok().expect("Parse 3");
        assert_eq!(addrs,
                   vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<commit_msg>Remove unused import<commit_after>use consts::*;
use proto::*;
use io::ZkIo;
use listeners::{ListenerSet, Subscription};
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>,
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf,
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>,
    listeners: ListenerSet<ZkState>,
}

impl ZooKeeper {
    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
        where F: FnOnce() + Send + 'static
    {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_| ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
        where W: Watcher + 'static
    {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher, chroot.clone());
        let listeners = ListenerSet::<ZkState>::new();
        let listeners1 = listeners.clone();
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender(), listeners1);
        let sender = io.sender();

        try!(Self::zk_thread("event", move || watch.run().unwrap()));
        try!(Self::zk_thread("io", move || io.run().unwrap()));

        Ok(ZooKeeper {
            chroot: chroot,
            xid: AtomicIsize::new(1),
            io: sender,
            listeners: listeners,
        })
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start),
                }
            }
            None => (None, connect_string.len()),
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.trim().to_socket_addrs() {
                Ok(mut addrs) => {
                    match addrs.nth(0) {
                        Some(addr) => addr,
                        None => return Err(ZkError::BadArguments),
                    }
                }
                Err(_) => return Err(ZkError::BadArguments),
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self,
                                             opcode: OpCode,
                                             xid: i32,
                                             req: Req,
                                             watch: Option<Watch>)
                                             -> ZkResult<Resp> {
        let rh = RequestHeader {
            xid: xid,
            opcode: opcode,
        };
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_| ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest {
            opcode: opcode,
            data: buf,
            listener: Some(resp_tx),
            watch: watch,
        };

        try!(self.io.send(request).map_err(|err| {
            warn!("error sending request: {:?}", err);
            ZkError::ConnectionLoss
        }));

        let mut response = try!(resp_rx.recv().map_err(|err| {
            warn!("error receiving response: {:?}", err);
            ZkError::ConnectionLoss
        }));

        match response.header.err {
            0 => {
                Ok(try!(ReadFrom::read_from(&mut response.data)
                            .map_err(|_| ZkError::MarshallingError)))
            }
            e => Err(FromPrimitive::from_i32(e).unwrap()),
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => {
                if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
            }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => {
                match path {
                    "/" => Ok(chroot.clone()),
                    path => Ok(chroot.clone() + try!(Self::validate_path(path))),
                }
            }
            None => Ok(try!(Self::validate_path(path)).to_owned()),
        }
    }

    fn cut_chroot(&self, path: String) -> String {
        if let Some(ref chroot) = self.chroot {
            path[chroot.len()..].to_owned()
        } else {
            path
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest {
            typ: 0,
            scheme: scheme.to_owned(),
            auth: auth,
        };

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self,
                  path: &str,
                  data: Vec<u8>,
                  acl: Vec<Acl>,
                  mode: CreateMode)
                  -> ZkResult<String> {
        let req = CreateRequest {
            path: try!(self.path(path)),
            data: data,
            acl: acl,
            flags: mode as i32,
        };

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(self.cut_chroot(response.path))
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest {
            path: try!(self.path(path)),
            version: version,
        };

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e),
        }
    }

    pub fn exists_w<W: Watcher + 'static>(&self,
                                            path: &str,
                                            watcher: W)
                                            -> ZkResult<Stat> {
        let req = ExistsRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Exist,
            watcher: Box::new(watcher),
        };

        let response: ExistsResponse = try!(self.request(OpCode::Exists,
                                                         self.xid(),
                                                         req,
                                                         Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest { path: try!(self.path(path)) };

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: Watcher + 'static>(&self,
                                                path: &str,
                                                watcher: W)
                                                -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Child,
            watcher: Box::new(watcher),
        };

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren,
                                                              self.xid(),
                                                              req,
                                                              Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren,
                                                              self.xid(),
                                                              req,
                                                              None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest {
            path: try!(self.path(path)),
            watch: watch,
        };

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: Watcher + 'static>(&self,
                                                                 path: &str,
                                                                 watcher: W)
                                                                 -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest {
            path: try!(self.path(path)),
            watch: true,
        };

        let watch = Watch {
            path: path.to_owned(),
            watch_type: WatchType::Data,
            watcher: Box::new(watcher),
        };

        let response: GetDataResponse = try!(self.request(OpCode::GetData,
                                                          self.xid(),
                                                          req,
                                                          Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest {
            path: try!(self.path(path)),
            acl: acl,
            version: version,
        };

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest {
            path: try!(self.path(path)),
            data: data,
            version: version,
        };

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn add_listener<Listener: Fn(ZkState) + Send + 'static>(&self,
                                                                listener: Listener)
                                                                -> Subscription {
        self.listeners.subscribe(listener)
    }

    pub fn remove_listener(&self, sub: Subscription) {
        self.listeners.unsubscribe(sub);
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        if let Err(err) = self.close() {
            info!("error closing zookeeper connection in drop: {:?}", err);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky on Travis, it works on my Linux box though.
    #[test]
    #[cfg(target_os = "macos")]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos")
                                  .ok()
                                  .expect("Parse 1");
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, Some("/mesos".to_owned()));

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").ok().expect("Parse 2");
        assert_eq!(addrs,
                   vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").ok().expect("Parse 3");
        assert_eq!(addrs,
                   vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
                                                         2181,
                                                         0,
                                                         0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<|endoftext|>