<commit_before>use crate::AsArray1D;
use plotpy::{Curve, Plot, Shapes, Text};
use std::collections::HashMap;

/// Holds the id and coordinates of an item
#[derive(Clone, Debug)]
struct Item {
    id: usize,   // identification number
    x: [f64; 3], // coordinates
}

/// Holds items
#[derive(Clone, Debug)]
struct Container {
    items: Vec<Item>,
}

/// Implements a grid for fast searching entries by coordinates
///
/// # Reference
///
/// * Durand, Farias, and Pedroso (2015) Computing intersections between
///   non-compatible curves and finite elements, Computational Mechanics;
///   DOI=10.1007/s00466-015-1181-y
pub struct GridSearch {
    // constants
    ndim: usize,      // space dimension
    ndiv: [usize; 3], // number of divisions along each direction
    xmin: [f64; 3],   // min values
    xmax: [f64; 3],   // max values
    xdelta: [f64; 3], // difference between max and min
    size: [f64; 3],   // size of each container
    cf: [usize; 3],   // coefficients [1, ndiv[0], ndiv[0]*ndiv[1]] (Eq. 8)

    // auxiliary variable
    ratio: [usize; 3], // ratio = trunc(δx[i]/Δx[i]) (Eq. 8)

    // holds non-empty containers. maps container.index to container.data
    containers: HashMap<usize, Container>,
}

impl GridSearch {
    /// Creates a new GridSearch
    ///
    /// # Input
    ///
    /// * `xmin, xmax` -- min and max coordinates (len = 2 or 3 == ndim)
    /// * `ndiv` -- number of divisions along each dimension (len = 2 or 3 == ndim)
    pub fn new(xmin: &[f64], xmax: &[f64], ndiv: &[usize]) -> Result<Self, &'static str> {
        // check input
        let ndim = ndiv.len();
        if ndim < 2 || ndim > 3 {
            return Err("len(ndiv) == ndim must be 2 or 3");
        }
        if xmin.len() != ndim {
            return Err("size of xmin must equal ndim == len(ndiv)");
        }
        if xmax.len() != ndim {
            return Err("size of xmax must equal ndim == len(ndiv)");
        }

        // allocate gs
        let mut grid = GridSearch {
            ndim,
            ndiv: [0; 3],
            xmin: [0.0; 3],
            xmax: [0.0; 3],
            xdelta: [0.0; 3],
            size: [0.0; 3],
            ratio: [0; 3],
            cf: [0; 3],
            containers: HashMap::new(),
        };

        // check and compute sizes
        for i in 0..ndim {
            grid.ndiv[i] = ndiv[i];
            if ndiv[i] < 1 {
                return Err("ndiv must be at least 1");
            }
            grid.xmin[i] = xmin[i];
            grid.xmax[i] = xmax[i];
            grid.xdelta[i] = xmax[i] - xmin[i];
            if grid.xdelta[i] <= 0.0 {
                return Err("xmax must be greater than xmin");
            }
            grid.size[i] = grid.xdelta[i] / (ndiv[i] as f64);
        }

        // coefficient
        grid.cf[0] = 1;
        grid.cf[1] = grid.ndiv[0];
        grid.cf[2] = grid.ndiv[0] * grid.ndiv[1];

        // done
        Ok(grid)
    }

    /// Inserts a new item to some container in the grid
    ///
    /// # Input
    ///
    /// * `id` -- identification number for the item
    /// * `x` -- coordinates of the item (an 1D array with len = ndim)
    pub fn insert<'a, T, U>(&mut self, id: usize, x: &'a T) -> Result<(), &'static str>
    where
        T: AsArray1D<'a, U>,
        U: 'a + Into<f64>,
    {
        if x.size() != self.ndim {
            return Err("x must have len equal to ndim");
        }
        match self.container_index(x) {
            Some(index) => {
                // new item data
                let mut item = Item {
                    id,
                    x: [x.at(0).into(), x.at(1).into(), 0.0],
                };
                if self.ndim == 3 {
                    item.x[2] = x.at(2).into();
                }

                // update container
                if self.containers.contains_key(&index) {
                    let container = self.containers.get_mut(&index).unwrap();
                    container.items.push(item);

                // initialize container
                } else {
                    self.containers.insert(index, Container { items: vec![item] });
                }
            }
            None => return Err("point is outside the grid"),
        }
        Ok(())
    }

    /// Returns a drawing of this object
    pub fn plot(&self) -> Result<Plot, &'static str> {
        // create plot
        let mut plot = Plot::new();

        // draw grid
        let mut xmin = vec![0.0; self.ndim];
        let mut xmax = vec![0.0; self.ndim];
        let mut ndiv = vec![0; self.ndim];
        for i in 0..self.ndim {
            xmin[i] = self.xmin[i];
            xmax[i] = self.xmax[i];
            ndiv[i] = self.ndiv[i];
        }
        let mut shapes = Shapes::new();
        shapes
            .set_alt_text_color("#5d5d5d")
            .draw_grid(&xmin, &xmax, &ndiv, false, true)?;
        plot.add(&shapes);

        // draw items
        if self.ndim == 2 {
            let mut curve = Curve::new();
            let mut text = Text::new();
            curve
                .set_marker_style("o")
                .set_marker_color("#fab32faa")
                .set_marker_line_color("black")
                .set_marker_line_width(0.5);
            text.set_color("#cd0000");
            for container in self.containers.values() {
                for item in &container.items {
                    let txt = format!("{}", item.id);
                    curve.draw(&[item.x[0]], &[item.x[1]]);
                    text.draw(item.x[0], item.x[1], &txt);
                }
            }
            plot.add(&curve);
            plot.add(&text);
        }
        Ok(plot)
    }

    /// Calculates the container index where the point x should be located
    ///
    /// # Output
    ///
    /// * returns the index of the container or None if the point is out-of-range
    fn container_index<'a, T, U>(&mut self, x: &'a T) -> Option<usize>
    where
        T: AsArray1D<'a, U>,
        U: 'a + Into<f64>,
    {
        let mut index = 0;
        for i in 0..self.ndim {
            if x.at(i).into() < self.xmin[i] || x.at(i).into() > self.xmax[i] {
                return None;
            }
            self.ratio[i] = ((x.at(i).into() - self.xmin[i]) / self.size[i]) as usize;
            if self.ratio[i] == self.ndiv[i] {
                // the point is exactly on the max edge, thus select inner container
                self.ratio[i] -= 1; // move to the inside
            }
            index += self.ratio[i] * self.cf[i];
        }
        Some(index)
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

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

    struct TestData {
        id: usize,
        x: Vec<f64>,
        container_index: usize,
    }

    fn get_test_data() -> Vec<TestData> {
        vec![
            TestData {
                id: 100,
                x: vec![0.0, 0.2],
                container_index: 6,
            },
            TestData {
                id: 200,
                x: vec![0.2, 0.6000000000000001],
                container_index: 12,
            },
            TestData {
                id: 300,
                x: vec![0.4, 1.0],
                // in this case, the ratio is:
                // 3.0000000000000004, 2.9999999999999996
                // thus, the point falls in #13 instead of #18
                container_index: 13,
            },
            TestData {
                id: 400,
                x: vec![0.6, 1.4],
                // in this case, the ratio is:
                // 4, 3.9999999999999996
                // thus, the point falls in #19 instead of #24
                container_index: 19,
            },
            TestData {
                id: 500,
                x: vec![0.8, 1.8],
                container_index: 24,
            },
        ]
    }

    #[test]
    fn new_fails_on_wrong_input() {
        let grid = GridSearch::new(&[0.0, 0.0], &[1.0, 1.0], &[1]);
        assert_eq!(grid.err(), Some("len(ndiv) == ndim must be 2 or 3"));
        let grid = GridSearch::new(&[0.0], &[1.0, 1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("size of xmin must equal ndim == len(ndiv)"));
        let grid = GridSearch::new(&[0.0, 0.0], &[1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("size of xmax must equal ndim == len(ndiv)"));
        let grid = GridSearch::new(&[0.0, 0.0], &[0.0, 1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("xmax must be greater than xmin"));
    }

    #[test]
    fn new_works() -> Result<(), &'static str> {
        let grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        assert_eq!(grid.ndim, 2);
        assert_eq!(grid.ndiv, [5, 5, 0]);
        assert_eq!(grid.xmin, [-0.2, -0.2, 0.0]);
        assert_eq!(grid.xmax, [0.8, 1.8, 0.0]);
        assert_eq!(grid.xdelta, [1.0, 2.0, 0.0]);
        assert_eq!(grid.size, [0.2, 0.4, 0.0]);
        assert_eq!(grid.cf, [1, 5, 25]);
        assert_eq!(grid.ratio, [0, 0, 0]);
        assert_eq!(grid.containers.len(), 0);
        println!("{:.20}", grid.size[0]);
        assert_eq!(grid.size[0], 0.20000000000000001110); // TODO: check if this is a problem
        let b0 = 0.20000000000000001110_f64.to_bits();
        let b1 = 0.2_f64.to_bits();
        println!("{} =? {}", b0, b1);
        assert_eq!(b0, b1);
        Ok(())
    }

    #[test]
    fn container_index_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            let index = grid.container_index(&data.x).unwrap();
            assert_eq!(index, data.container_index);
        }
        Ok(())
    }

    #[test]
    fn insert_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            grid.insert(data.id, &data.x)?;
            let container = grid.containers.get(&data.container_index).unwrap();
            let item = container.items.iter().find(|item| item.id == data.id).unwrap();
            assert_eq!(item.id, data.id);
        }
        assert_eq!(grid.containers.len(), 5);
        Ok(())
    }

    #[test]
    fn plot_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            grid.insert(data.id, &data.x)?;
        }
        let plot = grid.plot()?;
        plot.save("/tmp/gemlab/search_grid_plot_works.svg")?;
        Ok(())
    }
}
<commit_msg>Improve test<commit_after>use crate::AsArray1D;
use plotpy::{Curve, Plot, Shapes, Text};
use std::collections::HashMap;

/// Holds the id and coordinates of an item
struct Item {
    id: usize,   // identification number
    x: [f64; 3], // coordinates
}

/// Holds items
struct Container {
    items: Vec<Item>,
}

/// Implements a grid for fast searching entries by coordinates
///
/// # Reference
///
/// * Durand, Farias, and Pedroso (2015) Computing intersections between
///   non-compatible curves and finite elements, Computational Mechanics;
///   DOI=10.1007/s00466-015-1181-y
pub struct GridSearch {
    // constants
    ndim: usize,      // space dimension
    ndiv: [usize; 3], // number of divisions along each direction
    xmin: [f64; 3],   // min values
    xmax: [f64; 3],   // max values
    xdelta: [f64; 3], // difference between max and min
    size: [f64; 3],   // size of each container
    cf: [usize; 3],   // coefficients [1, ndiv[0], ndiv[0]*ndiv[1]] (Eq. 8)

    // auxiliary variable
    ratio: [usize; 3], // ratio = trunc(δx[i]/Δx[i]) (Eq. 8)

    // holds non-empty containers. maps container.index to container.data
    containers: HashMap<usize, Container>,
}

impl GridSearch {
    /// Creates a new GridSearch
    ///
    /// # Input
    ///
    /// * `xmin, xmax` -- min and max coordinates (len = 2 or 3 == ndim)
    /// * `ndiv` -- number of divisions along each dimension (len = 2 or 3 == ndim)
    pub fn new(xmin: &[f64], xmax: &[f64], ndiv: &[usize]) -> Result<Self, &'static str> {
        // check input
        let ndim = ndiv.len();
        if ndim < 2 || ndim > 3 {
            return Err("len(ndiv) == ndim must be 2 or 3");
        }
        if xmin.len() != ndim {
            return Err("size of xmin must equal ndim == len(ndiv)");
        }
        if xmax.len() != ndim {
            return Err("size of xmax must equal ndim == len(ndiv)");
        }

        // allocate gs
        let mut grid = GridSearch {
            ndim,
            ndiv: [0; 3],
            xmin: [0.0; 3],
            xmax: [0.0; 3],
            xdelta: [0.0; 3],
            size: [0.0; 3],
            ratio: [0; 3],
            cf: [0; 3],
            containers: HashMap::new(),
        };

        // check and compute sizes
        for i in 0..ndim {
            grid.ndiv[i] = ndiv[i];
            if ndiv[i] < 1 {
                return Err("ndiv must be at least equal to 1");
            }
            grid.xmin[i] = xmin[i];
            grid.xmax[i] = xmax[i];
            grid.xdelta[i] = xmax[i] - xmin[i];
            if grid.xdelta[i] <= 0.0 {
                return Err("xmax must be greater than xmin");
            }
            grid.size[i] = grid.xdelta[i] / (ndiv[i] as f64);
        }

        // coefficient
        grid.cf[0] = 1;
        grid.cf[1] = grid.ndiv[0];
        grid.cf[2] = grid.ndiv[0] * grid.ndiv[1];

        // done
        Ok(grid)
    }

    /// Inserts a new item to some container in the grid
    ///
    /// # Input
    ///
    /// * `id` -- identification number for the item
    /// * `x` -- coordinates of the item (an 1D array with len = ndim)
    pub fn insert<'a, T, U>(&mut self, id: usize, x: &'a T) -> Result<(), &'static str>
    where
        T: AsArray1D<'a, U>,
        U: 'a + Into<f64>,
    {
        if x.size() != self.ndim {
            return Err("x must have len equal to ndim");
        }
        match self.container_index(x) {
            Some(index) => {
                // new item data
                let mut item = Item {
                    id,
                    x: [x.at(0).into(), x.at(1).into(), 0.0],
                };
                if self.ndim == 3 {
                    item.x[2] = x.at(2).into();
                }

                // update container
                if self.containers.contains_key(&index) {
                    let container = self.containers.get_mut(&index).unwrap();
                    container.items.push(item);

                // initialize container
                } else {
                    self.containers.insert(index, Container { items: vec![item] });
                }
            }
            None => return Err("point is outside the grid"),
        }
        Ok(())
    }

    /// Returns a drawing of this object
    pub fn plot(&self) -> Result<Plot, &'static str> {
        // create plot
        let mut plot = Plot::new();

        // draw grid
        let mut xmin = vec![0.0; self.ndim];
        let mut xmax = vec![0.0; self.ndim];
        let mut ndiv = vec![0; self.ndim];
        for i in 0..self.ndim {
            xmin[i] = self.xmin[i];
            xmax[i] = self.xmax[i];
            ndiv[i] = self.ndiv[i];
        }
        let mut shapes = Shapes::new();
        shapes
            .set_alt_text_color("#5d5d5d")
            .draw_grid(&xmin, &xmax, &ndiv, false, true)?;
        plot.add(&shapes);

        // draw items
        if self.ndim == 2 {
            let mut curve = Curve::new();
            let mut text = Text::new();
            curve
                .set_marker_style("o")
                .set_marker_color("#fab32faa")
                .set_marker_line_color("black")
                .set_marker_line_width(0.5);
            text.set_color("#cd0000");
            for container in self.containers.values() {
                for item in &container.items {
                    let txt = format!("{}", item.id);
                    curve.draw(&[item.x[0]], &[item.x[1]]);
                    text.draw(item.x[0], item.x[1], &txt);
                }
            }
            plot.add(&curve);
            plot.add(&text);
        }
        Ok(plot)
    }

    /// Calculates the container index where the point x should be located
    ///
    /// # Output
    ///
    /// * returns the index of the container or None if the point is out-of-range
    fn container_index<'a, T, U>(&mut self, x: &'a T) -> Option<usize>
    where
        T: AsArray1D<'a, U>,
        U: 'a + Into<f64>,
    {
        let mut index = 0;
        for i in 0..self.ndim {
            if x.at(i).into() < self.xmin[i] || x.at(i).into() > self.xmax[i] {
                return None;
            }
            self.ratio[i] = ((x.at(i).into() - self.xmin[i]) / self.size[i]) as usize;
            if self.ratio[i] == self.ndiv[i] {
                // the point is exactly on the max edge, thus select inner container
                self.ratio[i] -= 1; // move to the inside
            }
            index += self.ratio[i] * self.cf[i];
        }
        Some(index)
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

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

    struct TestData {
        id: usize,
        x: Vec<f64>,
        container_index: usize,
    }

    fn get_test_data() -> Vec<TestData> {
        vec![
            TestData {
                id: 100,
                x: vec![0.0, 0.2],
                container_index: 6,
            },
            TestData {
                id: 200,
                x: vec![0.2, 0.6000000000000001],
                container_index: 12,
            },
            TestData {
                id: 300,
                x: vec![0.4, 1.0],
                // in this case, the ratio is:
                // 3.0000000000000004, 2.9999999999999996
                // thus, the point falls in #13 instead of #18
                container_index: 13,
            },
            TestData {
                id: 400,
                x: vec![0.6, 1.4],
                // in this case, the ratio is:
                // 4, 3.9999999999999996
                // thus, the point falls in #19 instead of #24
                container_index: 19,
            },
            TestData {
                id: 500,
                x: vec![0.8, 1.8],
                container_index: 24,
            },
        ]
    }

    #[test]
    fn new_fails_on_wrong_input() {
        let grid = GridSearch::new(&[0.0, 0.0], &[1.0, 1.0], &[1]);
        assert_eq!(grid.err(), Some("len(ndiv) == ndim must be 2 or 3"));
        let grid = GridSearch::new(&[0.0], &[1.0, 1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("size of xmin must equal ndim == len(ndiv)"));
        let grid = GridSearch::new(&[0.0, 0.0], &[1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("size of xmax must equal ndim == len(ndiv)"));
        let grid = GridSearch::new(&[0.0, 0.0], &[1.0, 1.0], &[0, 1]);
        assert_eq!(grid.err(), Some("ndiv must be at least equal to 1"));
        let grid = GridSearch::new(&[0.0, 0.0], &[0.0, 1.0], &[1, 1]);
        assert_eq!(grid.err(), Some("xmax must be greater than xmin"));
    }

    #[test]
    fn new_works() -> Result<(), &'static str> {
        let grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        assert_eq!(grid.ndim, 2);
        assert_eq!(grid.ndiv, [5, 5, 0]);
        assert_eq!(grid.xmin, [-0.2, -0.2, 0.0]);
        assert_eq!(grid.xmax, [0.8, 1.8, 0.0]);
        assert_eq!(grid.xdelta, [1.0, 2.0, 0.0]);
        assert_eq!(grid.size, [0.2, 0.4, 0.0]);
        assert_eq!(grid.cf, [1, 5, 25]);
        assert_eq!(grid.ratio, [0, 0, 0]);
        assert_eq!(grid.containers.len(), 0);
        println!("{:.20}", grid.size[0]);
        assert_eq!(grid.size[0], 0.20000000000000001110); // TODO: check if this is a problem
        let b0 = 0.20000000000000001110_f64.to_bits();
        let b1 = 0.2_f64.to_bits();
        println!("{} =? {}", b0, b1);
        assert_eq!(b0, b1);
        Ok(())
    }

    #[test]
    fn container_index_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            let index = grid.container_index(&data.x).unwrap();
            assert_eq!(index, data.container_index);
        }
        Ok(())
    }

    #[test]
    fn insert_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            grid.insert(data.id, &data.x)?;
            let container = grid.containers.get(&data.container_index).unwrap();
            let item = container.items.iter().find(|item| item.id == data.id).unwrap();
            assert_eq!(item.id, data.id);
        }
        assert_eq!(grid.containers.len(), 5);
        Ok(())
    }

    #[test]
    fn plot_works() -> Result<(), &'static str> {
        let mut grid = GridSearch::new(&[-0.2, -0.2], &[0.8, 1.8], &[5, 5])?;
        for data in get_test_data() {
            grid.insert(data.id, &data.x)?;
        }
        let plot = grid.plot()?;
        plot.save("/tmp/gemlab/search_grid_plot_works.svg")?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::collections::hash_map::{HashMap, Entry};
use std::path::Path;

use url::Url;

use ffi;
use ffi::opt;
use ffi::easy::Easy;
use http::Response;
use http::body::{Body,ToBody};
use {ProgressCb,ErrCode};

use self::Method::{Get, Head, Post, Put, Patch, Delete};
use self::BodyType::{Fixed, Chunked};

const DEFAULT_TIMEOUT_MS: usize = 30_000;

pub struct Handle {
    easy: Easy,
}

impl Handle {
    pub fn new() -> Handle {
        return configure(Handle { easy: Easy::new() }
            .timeout(DEFAULT_TIMEOUT_MS)
            .connect_timeout(DEFAULT_TIMEOUT_MS));

        #[cfg(unix)]
        fn configure(handle: Handle) -> Handle {
            let probe = ::openssl::probe::probe();
            let handle = match probe.cert_file {
                Some(ref path) => handle.ssl_ca_info(path),
                None => handle,
            };
            match probe.cert_dir {
                Some(ref path) => handle.ssl_ca_path(path),
                None => handle,
            }
        }

        #[cfg(not(unix))]
        fn configure(handle: Handle) -> Handle { handle }
    }

    pub fn timeout(mut self, ms: usize) -> Handle {
        self.easy.setopt(opt::TIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn connect_timeout(mut self, ms: usize) -> Handle {
        self.easy.setopt(opt::CONNECTTIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn verbose(mut self) -> Handle {
        self.easy.setopt(opt::VERBOSE, 1).unwrap();
        self
    }

    pub fn proxy<U: ToUrl>(mut self, proxy: U) -> Handle {
        proxy.with_url_str(|s| {
            self.easy.setopt(opt::PROXY, s).unwrap();
        });

        self
    }

    pub fn ssl_ca_path(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::CAPATH, path).unwrap();
        self
    }

    pub fn ssl_ca_info(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::CAINFO, path).unwrap();
        self
    }

    pub fn cookie_jar(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::COOKIEJAR, path).unwrap();
        self
    }

    pub fn cookie_file(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::COOKIEFILE, path).unwrap();
        self
    }

    pub fn cookies(self, path: &Path) -> Handle {
        self.cookie_jar(path).cookie_file(path)
    }

    pub fn get<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Get).uri(uri)
    }

    pub fn head<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Head).uri(uri)
    }

    pub fn post<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Post).uri(uri).body(body)
    }

    pub fn put<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Put).uri(uri).body(body)
    }

    pub fn patch<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Patch).uri(uri).body(body)
    }

    pub fn delete<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Delete).uri(uri)
    }
}

#[derive(Copy, Clone)]
pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Patch,
    Delete,
    Trace,
    Connect
}

pub struct Request<'a, 'b> {
    err: Option<ErrCode>,
    handle: &'a mut Handle,
    method: Method,
    headers: HashMap<String, Vec<String>>,
    body: Option<Body<'b>>,
    body_type: Option<BodyType>,
    content_type: bool, // whether or not the content type was set
    expect_continue: bool, // whether to expect a 100 continue from the server
    progress: Option<Box<ProgressCb<'b>>>,
    follow: bool,
}

enum BodyType {
    Fixed(usize),
    Chunked,
}

impl<'a, 'b> Request<'a, 'b> {
    pub fn new(handle: &'a mut Handle, method: Method) -> Request<'a, 'b> {
        Request {
            err: None,
            handle: handle,
            method: method,
            headers: HashMap::new(),
            body: None,
            body_type: None,
            content_type: false,
            expect_continue: false,
            progress: None,
            follow: false,
        }
    }

    pub fn uri<U: ToUrl>(mut self, uri: U) -> Request<'a, 'b> {
        uri.with_url_str(|s| {
            match self.handle.easy.setopt(opt::URL, s) {
                Ok(_) => {}
                Err(e) => self.err = Some(e)
            }
        });

        self
    }

    pub fn body<B: ToBody<'b>>(mut self, body: B) -> Request<'a, 'b> {
        self.body = Some(body.to_body());
        self
    }

    pub fn content_type(mut self, ty: &str) -> Request<'a, 'b> {
        if !self.content_type {
            self.content_type = true;
            append_header(&mut self.headers, "Content-Type", ty);
        }

        self
    }

    pub fn content_length(mut self, len: usize) -> Request<'a, 'b> {
        self.body_type = Some(Fixed(len));
        self
    }

    pub fn chunked(mut self) -> Request<'a, 'b> {
        self.body_type = Some(Chunked);
        self
    }

    pub fn expect_continue(mut self) -> Request<'a, 'b> {
        self.expect_continue = true;
        self
    }

    pub fn header(mut self, name: &str, val: &str) -> Request<'a, 'b> {
        append_header(&mut self.headers, name, val);
        self
    }

    pub fn get_header(&self, name: &str) -> Option<&[String]> {
        self.headers.get(name).map(|a| &a[..])
    }

    pub fn headers<'c, 'd, I: Iterator<Item=(&'c str, &'d str)>>(mut self, hdrs: I) -> Request<'a, 'b> {
        for (name, val) in hdrs {
            append_header(&mut self.headers, name, val);
        }

        self
    }

    pub fn progress<F>(mut self, cb: F) -> Request<'a, 'b>
        where F: FnMut(usize, usize, usize, usize) + 'b
    {
        self.progress = Some(Box::new(cb) as Box<ProgressCb<'b>>);
        self
    }

    pub fn follow_redirects(mut self, follow: bool) -> Request<'a, 'b> {
        self.follow = follow;
        self
    }

    pub fn exec(self) -> Result<Response, ErrCode> {
        // Deconstruct the struct
        let Request {
            err,
            handle,
            method,
            mut headers,
            mut body,
            body_type,
            content_type,
            expect_continue,
            progress,
            follow,
            ..
        } = self;

        if follow {
            try!(handle.easy.setopt(opt::FOLLOWLOCATION, 1));
        }

        match err {
            Some(e) => return Err(e),
            None => {}
        }

        // Clear custom headers set from the previous request
        try!(handle.easy.setopt(opt::HTTPHEADER, 0));

        match method {
            Get => try!(handle.easy.setopt(opt::HTTPGET, 1)),
            Head => try!(handle.easy.setopt(opt::NOBODY, 1)),
            Post => try!(handle.easy.setopt(opt::POST, 1)),
            Put => try!(handle.easy.setopt(opt::UPLOAD, 1)),
            Patch => {
                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "PATCH"));
                try!(handle.easy.setopt(opt::UPLOAD, 1));
            },
            Delete => {
                if body.is_some() {
                    try!(handle.easy.setopt(opt::UPLOAD, 1));
                }

                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "DELETE"));
            }
            _ => unimplemented!()
        }

        match body.as_ref() {
            None => {}
            Some(body) => {
                let body_type = body_type.unwrap_or(match body.get_size() {
                    Some(len) => Fixed(len),
                    None => Chunked,
                });

                match body_type {
                    Fixed(len) => {
                        match method {
                            Post => try!(handle.easy.setopt(opt::POSTFIELDSIZE, len)),
                            Put | Patch | Delete  => try!(handle.easy.setopt(opt::INFILESIZE, len)),
                            _ => {}
                        }
                        append_header(&mut headers, "Content-Length",
                                      &len.to_string());
                    }
                    Chunked => {
                        append_header(&mut headers, "Transfer-Encoding",
                                      "chunked");
                    }

                }

                if !content_type {
                    append_header(&mut headers, "Content-Type", "application/octet-stream");
                }

                if !expect_continue {
                    append_header(&mut headers, "Expect", "");
                }
            }
        }

        let mut ffi_headers = ffi::List::new();

        if !headers.is_empty() {
            let mut buf = Vec::new();

            for (k, v) in headers.iter() {
                buf.extend(k.bytes());
                buf.extend(": ".bytes());

                for v in v.iter() {
                    buf.extend(v.bytes());
                    buf.push(0);
                    ffi_headers.push_bytes(&buf);
                    buf.truncate(k.len() + 2);
                }

                buf.truncate(0);
            }

            try!(handle.easy.setopt(opt::HTTPHEADER, &ffi_headers));
        }

        handle.easy.perform(body.as_mut(), progress)
    }
}

fn append_header(map: &mut HashMap<String, Vec<String>>, key: &str, val: &str) {
    match map.entry(key.to_string()) {
        Entry::Vacant(entry) => {
            let mut values = Vec::new();
            values.push(val.to_string());
            entry.insert(values)
        },
        Entry::Occupied(entry) => entry.into_mut()
    };
}

pub trait ToUrl{
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str);
}

impl<'a> ToUrl for &'a str {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        f(self);
    }
}

impl<'a> ToUrl for &'a Url {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        self.to_string().with_url_str(f);
    }
}

impl ToUrl for String {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        self[..].with_url_str(f);
    }
}

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

    #[test]
    fn get_header() {
        let mut h = Handle::new();
        let r = h.get("/foo").header("foo", "bar");
        assert_eq!(r.get_header("foo"), Some(&["bar".to_string()][..]));
    }
}
<commit_msg>Bind LOW_SPEED_LIMIT/LOW_SPEED_TIME options<commit_after>use std::collections::hash_map::{HashMap, Entry};
use std::path::Path;

use url::Url;

use ffi;
use ffi::opt;
use ffi::easy::Easy;
use http::Response;
use http::body::{Body,ToBody};
use {ProgressCb,ErrCode};

use self::Method::{Get, Head, Post, Put, Patch, Delete};
use self::BodyType::{Fixed, Chunked};

const DEFAULT_TIMEOUT_MS: usize = 30_000;

pub struct Handle {
    easy: Easy,
}

impl Handle {
    pub fn new() -> Handle {
        return configure(Handle { easy: Easy::new() }
            .timeout(DEFAULT_TIMEOUT_MS)
            .connect_timeout(DEFAULT_TIMEOUT_MS));

        #[cfg(unix)]
        fn configure(handle: Handle) -> Handle {
            let probe = ::openssl::probe::probe();
            let handle = match probe.cert_file {
                Some(ref path) => handle.ssl_ca_info(path),
                None => handle,
            };
            match probe.cert_dir {
                Some(ref path) => handle.ssl_ca_path(path),
                None => handle,
            }
        }

        #[cfg(not(unix))]
        fn configure(handle: Handle) -> Handle { handle }
    }

    pub fn timeout(mut self, ms: usize) -> Handle {
        self.easy.setopt(opt::TIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn connect_timeout(mut self, ms: usize) -> Handle {
        self.easy.setopt(opt::CONNECTTIMEOUT_MS, ms).unwrap();
        self
    }

    /// Set the time in seconds that the transfer speed should be below
    /// the `low_speed_limit` rate of bytes per second for the library to
    /// consider it too slow and abort.
    ///
    /// The default for this option is 0 which means that this option is
    /// disabled.
    pub fn low_speed_timeout(mut self, seconds: usize) -> Handle {
        self.easy.setopt(opt::LOW_SPEED_TIME, seconds).unwrap();
        self
    }

    /// Set the average transfer speed in bytes per second that the
    /// transfer should be below during `low_speed_timeout` seconds for
    /// libcurl to consider it to be too slow and abort.
    ///
    /// The default for this option is 0 which means that this option is
    /// disabled.
    pub fn low_speed_limit(mut self, bytes_per_second: usize) -> Handle {
        self.easy.setopt(opt::LOW_SPEED_LIMIT, bytes_per_second).unwrap();
        self
    }

    pub fn verbose(mut self) -> Handle {
        self.easy.setopt(opt::VERBOSE, 1).unwrap();
        self
    }

    pub fn proxy<U: ToUrl>(mut self, proxy: U) -> Handle {
        proxy.with_url_str(|s| {
            self.easy.setopt(opt::PROXY, s).unwrap();
        });

        self
    }

    pub fn ssl_ca_path(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::CAPATH, path).unwrap();
        self
    }

    pub fn ssl_ca_info(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::CAINFO, path).unwrap();
        self
    }

    pub fn cookie_jar(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::COOKIEJAR, path).unwrap();
        self
    }

    pub fn cookie_file(mut self, path: &Path) -> Handle {
        self.easy.setopt(opt::COOKIEFILE, path).unwrap();
        self
    }

    pub fn cookies(self, path: &Path) -> Handle {
        self.cookie_jar(path).cookie_file(path)
    }

    pub fn get<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Get).uri(uri)
    }

    pub fn head<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Head).uri(uri)
    }

    pub fn post<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Post).uri(uri).body(body)
    }

    pub fn put<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Put).uri(uri).body(body)
    }

    pub fn patch<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Patch).uri(uri).body(body)
    }

    pub fn delete<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Delete).uri(uri)
    }
}

#[derive(Copy, Clone)]
pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Patch,
    Delete,
    Trace,
    Connect
}

pub struct Request<'a, 'b> {
    err: Option<ErrCode>,
    handle: &'a mut Handle,
    method: Method,
    headers: HashMap<String, Vec<String>>,
    body: Option<Body<'b>>,
    body_type: Option<BodyType>,
    content_type: bool, // whether or not the content type was set
    expect_continue: bool, // whether to expect a 100 continue from the server
    progress: Option<Box<ProgressCb<'b>>>,
    follow: bool,
}

enum BodyType {
    Fixed(usize),
    Chunked,
}

impl<'a, 'b> Request<'a, 'b> {
    pub fn new(handle: &'a mut Handle, method: Method) -> Request<'a, 'b> {
        Request {
            err: None,
            handle: handle,
            method: method,
            headers: HashMap::new(),
            body: None,
            body_type: None,
            content_type: false,
            expect_continue: false,
            progress: None,
            follow: false,
        }
    }

    pub fn uri<U: ToUrl>(mut self, uri: U) -> Request<'a, 'b> {
        uri.with_url_str(|s| {
            match self.handle.easy.setopt(opt::URL, s) {
                Ok(_) => {}
                Err(e) => self.err = Some(e)
            }
        });

        self
    }

    pub fn body<B: ToBody<'b>>(mut self, body: B) -> Request<'a, 'b> {
        self.body = Some(body.to_body());
        self
    }

    pub fn content_type(mut self, ty: &str) -> Request<'a, 'b> {
        if !self.content_type {
            self.content_type = true;
            append_header(&mut self.headers, "Content-Type", ty);
        }

        self
    }

    pub fn content_length(mut self, len: usize) -> Request<'a, 'b> {
        self.body_type = Some(Fixed(len));
        self
    }

    pub fn chunked(mut self) -> Request<'a, 'b> {
        self.body_type = Some(Chunked);
        self
    }

    pub fn expect_continue(mut self) -> Request<'a, 'b> {
        self.expect_continue = true;
        self
    }

    pub fn header(mut self, name: &str, val: &str) -> Request<'a, 'b> {
        append_header(&mut self.headers, name, val);
        self
    }

    pub fn get_header(&self, name: &str) -> Option<&[String]> {
        self.headers.get(name).map(|a| &a[..])
    }

    pub fn headers<'c, 'd, I: Iterator<Item=(&'c str, &'d str)>>(mut self, hdrs: I) -> Request<'a, 'b> {
        for (name, val) in hdrs {
            append_header(&mut self.headers, name, val);
        }

        self
    }

    pub fn progress<F>(mut self, cb: F) -> Request<'a, 'b>
        where F: FnMut(usize, usize, usize, usize) + 'b
    {
        self.progress = Some(Box::new(cb) as Box<ProgressCb<'b>>);
        self
    }

    pub fn follow_redirects(mut self, follow: bool) -> Request<'a, 'b> {
        self.follow = follow;
        self
    }

    pub fn exec(self) -> Result<Response, ErrCode> {
        // Deconstruct the struct
        let Request {
            err,
            handle,
            method,
            mut headers,
            mut body,
            body_type,
            content_type,
            expect_continue,
            progress,
            follow,
            ..
        } = self;

        if follow {
            try!(handle.easy.setopt(opt::FOLLOWLOCATION, 1));
        }

        match err {
            Some(e) => return Err(e),
            None => {}
        }

        // Clear custom headers set from the previous request
        try!(handle.easy.setopt(opt::HTTPHEADER, 0));

        match method {
            Get => try!(handle.easy.setopt(opt::HTTPGET, 1)),
            Head => try!(handle.easy.setopt(opt::NOBODY, 1)),
            Post => try!(handle.easy.setopt(opt::POST, 1)),
            Put => try!(handle.easy.setopt(opt::UPLOAD, 1)),
            Patch => {
                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "PATCH"));
                try!(handle.easy.setopt(opt::UPLOAD, 1));
            },
            Delete => {
                if body.is_some() {
                    try!(handle.easy.setopt(opt::UPLOAD, 1));
                }

                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "DELETE"));
            }
            _ => unimplemented!()
        }

        match body.as_ref() {
            None => {}
            Some(body) => {
                let body_type = body_type.unwrap_or(match body.get_size() {
                    Some(len) => Fixed(len),
                    None => Chunked,
                });

                match body_type {
                    Fixed(len) => {
                        match method {
                            Post => try!(handle.easy.setopt(opt::POSTFIELDSIZE, len)),
                            Put | Patch | Delete  => try!(handle.easy.setopt(opt::INFILESIZE, len)),
                            _ => {}
                        }
                        append_header(&mut headers, "Content-Length",
                                      &len.to_string());
                    }
                    Chunked => {
                        append_header(&mut headers, "Transfer-Encoding",
                                      "chunked");
                    }

                }

                if !content_type {
                    append_header(&mut headers, "Content-Type", "application/octet-stream");
                }

                if !expect_continue {
                    append_header(&mut headers, "Expect", "");
                }
            }
        }

        let mut ffi_headers = ffi::List::new();

        if !headers.is_empty() {
            let mut buf = Vec::new();

            for (k, v) in headers.iter() {
                buf.extend(k.bytes());
                buf.extend(": ".bytes());

                for v in v.iter() {
                    buf.extend(v.bytes());
                    buf.push(0);
                    ffi_headers.push_bytes(&buf);
                    buf.truncate(k.len() + 2);
                }

                buf.truncate(0);
            }

            try!(handle.easy.setopt(opt::HTTPHEADER, &ffi_headers));
        }

        handle.easy.perform(body.as_mut(), progress)
    }
}

fn append_header(map: &mut HashMap<String, Vec<String>>, key: &str, val: &str) {
    match map.entry(key.to_string()) {
        Entry::Vacant(entry) => {
            let mut values = Vec::new();
            values.push(val.to_string());
            entry.insert(values)
        },
        Entry::Occupied(entry) => entry.into_mut()
    };
}

pub trait ToUrl{
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str);
}

impl<'a> ToUrl for &'a str {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        f(self);
    }
}

impl<'a> ToUrl for &'a Url {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        self.to_string().with_url_str(f);
    }
}

impl ToUrl for String {
    fn with_url_str<F>(self, f: F) where F: FnOnce(&str) {
        self[..].with_url_str(f);
    }
}

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

    #[test]
    fn get_header() {
        let mut h = Handle::new();
        let r = h.get("/foo").header("foo", "bar");
        assert_eq!(r.get_header("foo"), Some(&["bar".to_string()][..]));
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use structopt::StructOpt;

use anyhow::Result;

use strum_macros::{
    Display,
    EnumIter,
    EnumString,
};

pub mod genes;
pub mod search_xml;
pub mod sequences;
pub mod utils;

#[derive(Debug, Display, PartialEq, Eq, Hash, EnumString, EnumIter)]
#[strum(serialize_all = "kebab-case")]
pub enum Groupable {
    #[strum(ascii_case_insensitive)]
    Base,

    #[strum(ascii_case_insensitive)]
    Crs,

    #[strum(ascii_case_insensitive)]
    Feedback,

    #[strum(ascii_case_insensitive)]
    GoAnnotations,

    #[strum(ascii_case_insensitive)]
    InteractingProteins,

    #[strum(ascii_case_insensitive)]
    InteractingRnas,

    #[strum(ascii_case_insensitive)]
    RegionInfo,

    #[strum(ascii_case_insensitive)]
    Orfs,

    #[strum(ascii_case_insensitive)]
    Precompute,

    #[strum(ascii_case_insensitive)]
    QaStatus,

    #[strum(ascii_case_insensitive)]
    R2dt,

    #[strum(ascii_case_insensitive)]
    RfamHits,

    #[strum(ascii_case_insensitive)]
    SoInfo,
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum GenesCommand {
    /// This command does a join between the sequence information, locus information and
    /// gene information. Any sequence not part of a gene is ignored. It then splits
    /// the resulting merged entries by assembly so it produces a directory of files,
    /// one per assmebly. These can then be merge together to produce complete gene
    /// entries. This assumes the locus and sequence information is sorted by a common
    /// id.
    SelectAndSplit {
        /// Filename to read the locus level information from
        #[structopt(parse(from_os_str))]
        locus: PathBuf,

        /// Filename to read the selected members from
        #[structopt(parse(from_os_str))]
        sequences: PathBuf,

        /// Output directory to write to, does not need to exist
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    /// This will take all gene member information and merge them into genes. This needs
    /// enough memory to fit information on all members from an assembly into memory
    /// at once.
    MergeAssembly {
        /// A file of all gene members for a given assembly. The file does not need to be
        /// sorted.
        #[structopt(parse(from_os_str))]
        members: PathBuf,

        /// Filename to write complete gene to
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    AsXml {
        /// File of the complete genes to read
        #[structopt(parse(from_os_str))]
        genes: PathBuf,

        /// Filename to write the xml data to
        #[structopt(parse(from_os_str))]
        xml_output: PathBuf,

        /// Filename to write the counts to
        #[structopt(parse(from_os_str))]
        count_output: PathBuf,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum SequenceCommand {
    Merge {
        #[structopt(parse(from_os_str))]
        base: PathBuf,

        #[structopt(parse(from_os_str))]
        crs: PathBuf,

        #[structopt(parse(from_os_str))]
        feedback: PathBuf,

        #[structopt(parse(from_os_str))]
        go_annotations: PathBuf,

        #[structopt(parse(from_os_str))]
        interacting_proteins: PathBuf,

        #[structopt(parse(from_os_str))]
        interacting_rnas: PathBuf,

        #[structopt(parse(from_os_str))]
        precompute: PathBuf,

        #[structopt(parse(from_os_str))]
        qa_status: PathBuf,

        #[structopt(parse(from_os_str))]
        r2dt_hits: PathBuf,

        #[structopt(parse(from_os_str))]
        rfam_hits: PathBuf,

        #[structopt(parse(from_os_str))]
        orfs: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename of the SO term tree metadata.
        so_term_tree: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename to write the results to, '-' means stdout
        output: PathBuf,
    },

    Normalize {
        #[structopt(parse(from_os_str))]
        /// Filename of the accessions
        accessions: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename for the merged metadata
        metadata: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename to write the results to, '-' means stdout
        output: PathBuf,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum Subcommand {
    /// Merge each distinct type of metadata into single merged entries. This assumes that
    /// all files are sorted in the same way. Additionally, the xref file must have at
    /// least one entry for all possible urs_taxids.
    Group {
        /// Type of data to group
        #[structopt(case_insensitive = true)]
        data_type: Groupable,

        /// Filename to read the results from, '-' means stdin
        #[structopt(parse(from_os_str))]
        path: PathBuf,

        /// The maximum count of the entries.
        max_count: usize,

        /// Filename to write the results to, '-' means stdout
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    /// Set of commands dealing with building data for sequence data export
    Sequences {
        #[structopt(subcommand)]
        command: SequenceCommand,
    },

    /// Commands dealing with building data for gene export
    Genes {
        #[structopt(subcommand)]
        command: GenesCommand,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
struct Opt {
    /// Set the logging option, more is more verbose.
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: u32,

    #[structopt(subcommand)]
    command: Subcommand,
}

fn main() -> Result<()> {
    let opt = Opt::from_args();

    let level = match opt.verbose {
        0 => simplelog::LevelFilter::Warn,
        1 => simplelog::LevelFilter::Info,
        2 => simplelog::LevelFilter::Debug,
        _ => simplelog::LevelFilter::Trace,
    };
    simplelog::TermLogger::init(
        level,
        simplelog::Config::default(),
        simplelog::TerminalMode::Stderr,
    )
    .unwrap_or_else(|_| eprintln!("Failed to create logger, ignore"));

    match opt.command {
        Subcommand::Group {
            data_type,
            path,
            max_count,
            output,
        } => match data_type {
            Groupable::Base => sequences::basic::group(&path, max_count, &output)?,
            Groupable::Crs => sequences::crs::group(&path, max_count, &output)?,
            Groupable::Feedback => sequences::feedback::group(&path, max_count, &output)?,
            Groupable::GoAnnotations => sequences::go_annotation::group(&path, max_count, &output)?,
            Groupable::InteractingProteins => {
                sequences::interacting_protein::group(&path, max_count, &output)?
            },
            Groupable::InteractingRnas => {
                sequences::interacting_rna::group(&path, max_count, &output)?
            },
            Groupable::RegionInfo => genes::region::group(&path, max_count, &output)?,
            Groupable::Precompute => sequences::precompute::group(&path, max_count, &output)?,
            Groupable::QaStatus => sequences::qa_status::group(&path, max_count, &output)?,
            Groupable::R2dt=> sequences::r2dt::group(&path, max_count, &output)?,
            Groupable::RfamHits => sequences::rfam_hit::group(&path, max_count, &output)?,
            Groupable::Orfs => sequences::orf::group(&path, max_count, &output)?,
            Groupable::SoInfo => Err(anyhow::anyhow!("May not group so info"))?,
        },
        Subcommand::Sequences {
            command,
        } => match command {
            SequenceCommand::Merge {
                base,
                crs,
                feedback,
                go_annotations,
                interacting_proteins,
                interacting_rnas,
                precompute,
                qa_status,
                r2dt_hits,
                rfam_hits,
                orfs,
                so_term_tree,
                output,
            } => sequences::writers::write_merge(
                vec![
                    base,
                    crs,
                    feedback,
                    go_annotations,
                    interacting_proteins,
                    interacting_rnas,
                    precompute,
                    qa_status,
                    r2dt_hits,
                    rfam_hits,
                    orfs,
                    so_term_tree,
                ],
                &output,
            )?,
            SequenceCommand::Normalize {
                accessions,
                metadata,
                output,
            } => sequences::writers::write(&accessions, &metadata, &output)?,
        },
        Subcommand::Genes {
            command,
        } => match command {
            GenesCommand::SelectAndSplit {
                locus,
                sequences,
                output,
            } => genes::writers::write_selected_members(&locus, &sequences, &output)?,
            GenesCommand::MergeAssembly {
                members,
                output,
            } => genes::writers::write_merged_members(&members, &output)?,
            GenesCommand::AsXml {
                genes,
                xml_output,
                count_output,
            } => genes::writers::write_search_files(&genes, &xml_output, &count_output)?,
        },
    }

    Ok(())
}
<commit_msg>Simplify enum definition<commit_after>use std::path::PathBuf;
use structopt::StructOpt;

use anyhow::Result;

use strum_macros::{
    Display,
    EnumIter,
    EnumString,
};

pub mod genes;
pub mod search_xml;
pub mod sequences;
pub mod utils;

#[derive(Debug, Display, PartialEq, Eq, Hash, EnumString, EnumIter)]
#[strum(ascii_case_insensitive, serialize_all = "kebab-case")]
pub enum Groupable {
    Base,
    Crs,
    Feedback,
    GoAnnotations,
    InteractingProteins,
    InteractingRnas,
    RegionInfo,
    Orfs,
    Precompute,
    QaStatus,
    R2dt,
    RfamHits,
    SoInfo,
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum GenesCommand {
    /// This command does a join between the sequence information, locus information and
    /// gene information. Any sequence not part of a gene is ignored. It then splits
    /// the resulting merged entries by assembly so it produces a directory of files,
    /// one per assmebly. These can then be merge together to produce complete gene
    /// entries. This assumes the locus and sequence information is sorted by a common
    /// id.
    SelectAndSplit {
        /// Filename to read the locus level information from
        #[structopt(parse(from_os_str))]
        locus: PathBuf,

        /// Filename to read the selected members from
        #[structopt(parse(from_os_str))]
        sequences: PathBuf,

        /// Output directory to write to, does not need to exist
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    /// This will take all gene member information and merge them into genes. This needs
    /// enough memory to fit information on all members from an assembly into memory
    /// at once.
    MergeAssembly {
        /// A file of all gene members for a given assembly. The file does not need to be
        /// sorted.
        #[structopt(parse(from_os_str))]
        members: PathBuf,

        /// Filename to write complete gene to
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    AsXml {
        /// File of the complete genes to read
        #[structopt(parse(from_os_str))]
        genes: PathBuf,

        /// Filename to write the xml data to
        #[structopt(parse(from_os_str))]
        xml_output: PathBuf,

        /// Filename to write the counts to
        #[structopt(parse(from_os_str))]
        count_output: PathBuf,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum SequenceCommand {
    Merge {
        #[structopt(parse(from_os_str))]
        base: PathBuf,

        #[structopt(parse(from_os_str))]
        crs: PathBuf,

        #[structopt(parse(from_os_str))]
        feedback: PathBuf,

        #[structopt(parse(from_os_str))]
        go_annotations: PathBuf,

        #[structopt(parse(from_os_str))]
        interacting_proteins: PathBuf,

        #[structopt(parse(from_os_str))]
        interacting_rnas: PathBuf,

        #[structopt(parse(from_os_str))]
        precompute: PathBuf,

        #[structopt(parse(from_os_str))]
        qa_status: PathBuf,

        #[structopt(parse(from_os_str))]
        r2dt_hits: PathBuf,

        #[structopt(parse(from_os_str))]
        rfam_hits: PathBuf,

        #[structopt(parse(from_os_str))]
        orfs: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename of the SO term tree metadata.
        so_term_tree: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename to write the results to, '-' means stdout
        output: PathBuf,
    },

    Normalize {
        #[structopt(parse(from_os_str))]
        /// Filename of the accessions
        accessions: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename for the merged metadata
        metadata: PathBuf,

        #[structopt(parse(from_os_str))]
        /// Filename to write the results to, '-' means stdout
        output: PathBuf,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum Subcommand {
    /// Merge each distinct type of metadata into single merged entries. This assumes that
    /// all files are sorted in the same way. Additionally, the xref file must have at
    /// least one entry for all possible urs_taxids.
    Group {
        /// Type of data to group
        #[structopt(case_insensitive = true)]
        data_type: Groupable,

        /// Filename to read the results from, '-' means stdin
        #[structopt(parse(from_os_str))]
        path: PathBuf,

        /// The maximum count of the entries.
        max_count: usize,

        /// Filename to write the results to, '-' means stdout
        #[structopt(parse(from_os_str))]
        output: PathBuf,
    },

    /// Set of commands dealing with building data for sequence data export
    Sequences {
        #[structopt(subcommand)]
        command: SequenceCommand,
    },

    /// Commands dealing with building data for gene export
    Genes {
        #[structopt(subcommand)]
        command: GenesCommand,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
struct Opt {
    /// Set the logging option, more is more verbose.
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: u32,

    #[structopt(subcommand)]
    command: Subcommand,
}

fn main() -> Result<()> {
    let opt = Opt::from_args();

    let level = match opt.verbose {
        0 => simplelog::LevelFilter::Warn,
        1 => simplelog::LevelFilter::Info,
        2 => simplelog::LevelFilter::Debug,
        _ => simplelog::LevelFilter::Trace,
    };
    simplelog::TermLogger::init(
        level,
        simplelog::Config::default(),
        simplelog::TerminalMode::Stderr,
    )
    .unwrap_or_else(|_| eprintln!("Failed to create logger, ignore"));

    match opt.command {
        Subcommand::Group {
            data_type,
            path,
            max_count,
            output,
        } => match data_type {
            Groupable::Base => sequences::basic::group(&path, max_count, &output)?,
            Groupable::Crs => sequences::crs::group(&path, max_count, &output)?,
            Groupable::Feedback => sequences::feedback::group(&path, max_count, &output)?,
            Groupable::GoAnnotations => sequences::go_annotation::group(&path, max_count, &output)?,
            Groupable::InteractingProteins => {
                sequences::interacting_protein::group(&path, max_count, &output)?
            },
            Groupable::InteractingRnas => {
                sequences::interacting_rna::group(&path, max_count, &output)?
            },
            Groupable::RegionInfo => genes::region::group(&path, max_count, &output)?,
            Groupable::Precompute => sequences::precompute::group(&path, max_count, &output)?,
            Groupable::QaStatus => sequences::qa_status::group(&path, max_count, &output)?,
            Groupable::R2dt=> sequences::r2dt::group(&path, max_count, &output)?,
            Groupable::RfamHits => sequences::rfam_hit::group(&path, max_count, &output)?,
            Groupable::Orfs => sequences::orf::group(&path, max_count, &output)?,
            Groupable::SoInfo => Err(anyhow::anyhow!("May not group so info"))?,
        },
        Subcommand::Sequences {
            command,
        } => match command {
            SequenceCommand::Merge {
                base,
                crs,
                feedback,
                go_annotations,
                interacting_proteins,
                interacting_rnas,
                precompute,
                qa_status,
                r2dt_hits,
                rfam_hits,
                orfs,
                so_term_tree,
                output,
            } => sequences::writers::write_merge(
                vec![
                    base,
                    crs,
                    feedback,
                    go_annotations,
                    interacting_proteins,
                    interacting_rnas,
                    precompute,
                    qa_status,
                    r2dt_hits,
                    rfam_hits,
                    orfs,
                    so_term_tree,
                ],
                &output,
            )?,
            SequenceCommand::Normalize {
                accessions,
                metadata,
                output,
            } => sequences::writers::write(&accessions, &metadata, &output)?,
        },
        Subcommand::Genes {
            command,
        } => match command {
            GenesCommand::SelectAndSplit {
                locus,
                sequences,
                output,
            } => genes::writers::write_selected_members(&locus, &sequences, &output)?,
            GenesCommand::MergeAssembly {
                members,
                output,
            } => genes::writers::write_merged_members(&members, &output)?,
            GenesCommand::AsXml {
                genes,
                xml_output,
                count_output,
            } => genes::writers::write_search_files(&genes, &xml_output, &count_output)?,
        },
    }

    Ok(())
}
<|endoftext|>


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

#![feature(test)]

extern crate test;

const RAND_BENCH_N: u64 = 1000;

use test::Bencher;

use rand::distributions::{Bernoulli, Distribution, Standard};
use rand::prelude::*;
use rand_pcg::{Pcg32, Pcg64Mcg};

#[bench]
fn misc_gen_bool_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_bool(0.18);
        }
        accum
    })
}

#[bench]
fn misc_gen_bool_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        let mut p = 0.18;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_bool(p);
            p += 0.0001;
        }
        accum
    })
}

#[bench]
fn misc_gen_ratio_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_ratio(2, 3);
        }
        accum
    })
}

#[bench]
fn misc_gen_ratio_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for i in 2..(crate::RAND_BENCH_N as u32 + 2) {
            accum ^= rng.gen_ratio(i, i + 1);
        }
        accum
    })
}

#[bench]
fn misc_bernoulli_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let d = rand::distributions::Bernoulli::new(0.18).unwrap();
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.sample(d);
        }
        accum
    })
}

#[bench]
fn misc_bernoulli_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        let mut p = 0.18;
        for _ in 0..crate::RAND_BENCH_N {
            let d = Bernoulli::new(p).unwrap();
            accum ^= rng.sample(d);
            p += 0.0001;
        }
        accum
    })
}

#[bench]
fn gen_1k_iter_repeat(b: &mut Bencher) {
    use std::iter;
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u64> = iter::repeat(()).map(|()| rng.gen()).take(128).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1k_sample_iter(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u64> = Standard.sample_iter(&mut rng).take(128).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1k_gen_array(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        // max supported array length is 32!
        let v: [[u64; 32]; 4] = rng.gen();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1k_fill(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    let mut buf = [0u64; 128];
    b.iter(|| {
        rng.fill(&mut buf[..]);
        buf
    });
    b.bytes = 1024;
}
<commit_msg>Misc bench: add gen_1kb_u16_*<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(test)]

extern crate test;

const RAND_BENCH_N: u64 = 1000;

use test::Bencher;

use rand::distributions::{Bernoulli, Distribution, Standard};
use rand::prelude::*;
use rand_pcg::{Pcg32, Pcg64Mcg};

#[bench]
fn misc_gen_bool_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_bool(0.18);
        }
        accum
    })
}

#[bench]
fn misc_gen_bool_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        let mut p = 0.18;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_bool(p);
            p += 0.0001;
        }
        accum
    })
}

#[bench]
fn misc_gen_ratio_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.gen_ratio(2, 3);
        }
        accum
    })
}

#[bench]
fn misc_gen_ratio_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        for i in 2..(crate::RAND_BENCH_N as u32 + 2) {
            accum ^= rng.gen_ratio(i, i + 1);
        }
        accum
    })
}

#[bench]
fn misc_bernoulli_const(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let d = rand::distributions::Bernoulli::new(0.18).unwrap();
        let mut accum = true;
        for _ in 0..crate::RAND_BENCH_N {
            accum ^= rng.sample(d);
        }
        accum
    })
}

#[bench]
fn misc_bernoulli_var(b: &mut Bencher) {
    let mut rng = Pcg32::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let mut accum = true;
        let mut p = 0.18;
        for _ in 0..crate::RAND_BENCH_N {
            let d = Bernoulli::new(p).unwrap();
            accum ^= rng.sample(d);
            p += 0.0001;
        }
        accum
    })
}

#[bench]
fn gen_1kb_u16_iter_repeat(b: &mut Bencher) {
    use std::iter;
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u16> = iter::repeat(()).map(|()| rng.gen()).take(512).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u16_sample_iter(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u16> = Standard.sample_iter(&mut rng).take(512).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u16_gen_array(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        // max supported array length is 32!
        let v: [[u16; 32]; 16] = rng.gen();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u16_fill(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    let mut buf = [0u16; 512];
    b.iter(|| {
        rng.fill(&mut buf[..]);
        buf
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u64_iter_repeat(b: &mut Bencher) {
    use std::iter;
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u64> = iter::repeat(()).map(|()| rng.gen()).take(128).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u64_sample_iter(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        let v: Vec<u64> = Standard.sample_iter(&mut rng).take(128).collect();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u64_gen_array(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    b.iter(|| {
        // max supported array length is 32!
        let v: [[u64; 32]; 4] = rng.gen();
        v
    });
    b.bytes = 1024;
}

#[bench]
fn gen_1kb_u64_fill(b: &mut Bencher) {
    let mut rng = Pcg64Mcg::from_rng(&mut thread_rng()).unwrap();
    let mut buf = [0u64; 128];
    b.iter(|| {
        rng.fill(&mut buf[..]);
        buf
    });
    b.bytes = 1024;
}
<|endoftext|>


<commit_before>#![feature(test)]

#[macro_use]
extern crate slog;
extern crate slog_json;

extern crate test;

use std::io;
use test::Bencher;
use slog::*;

struct BlackBoxDrain;

impl Drain for BlackBoxDrain {
    fn log(&self,
           buf: &mut Vec<u8>,
           ri: &Record, o : &OwnedKeyValueNode) -> Result<()> {

        test::black_box((buf, ri, o));
        Ok(())
    }
}

struct BlackBoxWriter;

impl io::Write for BlackBoxWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        test::black_box(buf);
        Ok(buf.len())
    }

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

#[bench]
fn log_discard_empty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!());

    b.iter(|| {
        info!(log, "");
    });
}

#[bench]
fn log_discard_nonempty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        info!(log, "", "what" => "write");
    });
}

#[bench]
fn logger_clone_empty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!());

    b.iter(|| {
        log.clone()
    });
}

#[bench]
fn logger_clone_nonempty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        log.clone()
    });
}

#[bench]
fn logger_new_empty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!());

    b.iter(|| {
        log.new(o!())
    });
}

#[bench]
fn logger_new_nonempty(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        log.new(o!("what" => "write"));
    });
}

#[bench]
fn log_discard_i32val(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!());

    b.iter(|| {
        info!(log, "", "i32" => 5);
    });
}

#[bench]
fn log_discard_i32closure(b: &mut Bencher) {
    let log = BlackBoxDrain.into_logger(o!());

    b.iter(|| {
        info!(log, "", "i32" => |_:&Record|{5});
    });
}

#[bench]
fn log_stream_json_blackbox_i32val(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "",  "i32" => 5);
    });
}

#[bench]
fn log_stream_json_blackbox_i32closure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "i32" => |_:&Record|{5});
    });
}

#[bench]
fn log_stream_json_blackbox_i32pushclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "i32" => PushLazy(|_:&Record, ser : ValueSerializer|{
            ser.serialize(5)
        }));
    });
}

const LONG_STRING : &'static str = "A long string that would take some time to allocate";

#[bench]
fn log_stream_json_blackbox_strclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "str" => |_:&Record| {
            String::from(LONG_STRING)
        });
    });
}

#[bench]
fn log_stream_json_blackbox_strpushclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "str" => PushLazy(|_:&Record, ser : ValueSerializer|{
            ser.serialize(LONG_STRING)
        }));
    });
}
<commit_msg>Fix: Bench fails to compile<commit_after>#![feature(test)]

#[macro_use]
extern crate slog;
extern crate slog_json;

extern crate test;

use std::io;
use test::Bencher;
use slog::*;

struct BlackBoxDrain;

impl Drain for BlackBoxDrain {
    fn log(&self,
           buf: &mut Vec<u8>,
           ri: &Record, o : &OwnedKeyValueList) -> Result<()> {

        test::black_box((buf, ri, o));
        Ok(())
    }
}

struct BlackBoxWriter;

impl io::Write for BlackBoxWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        test::black_box(buf);
        Ok(buf.len())
    }

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

#[bench]
fn log_discard_empty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!());

    b.iter(|| {
        info!(log, "");
    });
}

#[bench]
fn log_discard_nonempty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        info!(log, "", "what" => "write");
    });
}

#[bench]
fn logger_clone_empty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!());

    b.iter(|| {
        log.clone()
    });
}

#[bench]
fn logger_clone_nonempty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        log.clone()
    });
}

#[bench]
fn logger_new_empty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!());

    b.iter(|| {
        log.new(o!())
    });
}

#[bench]
fn logger_new_nonempty(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!("build" => "123456", "id" => 123456));

    b.iter(|| {
        log.new(o!("what" => "write"));
    });
}

#[bench]
fn log_discard_i32val(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!());

    b.iter(|| {
        info!(log, "", "i32" => 5);
    });
}

#[bench]
fn log_discard_i32closure(b: &mut Bencher) {
    let log = Logger::root(BlackBoxDrain, o!());

    b.iter(|| {
        info!(log, "", "i32" => |_:&Record|{5});
    });
}

#[bench]
fn log_stream_json_blackbox_i32val(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "",  "i32" => 5);
    });
}

#[bench]
fn log_stream_json_blackbox_i32closure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "i32" => |_:&Record|{5});
    });
}

#[bench]
fn log_stream_json_blackbox_i32pushclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "i32" => PushLazy(|_:&Record, ser : ValueSerializer|{
            ser.serialize(5)
        }));
    });
}

const LONG_STRING : &'static str = "A long string that would take some time to allocate";

#[bench]
fn log_stream_json_blackbox_strclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "str" => |_:&Record| {
            String::from(LONG_STRING)
        });
    });
}

#[bench]
fn log_stream_json_blackbox_strpushclosure(b: &mut Bencher) {
    let drain = stream(BlackBoxWriter, slog_json::new());

    let log = Logger::root(drain, o!());

    b.iter(|| {
        info!(log, "", "str" => PushLazy(|_:&Record, ser : ValueSerializer|{
            ser.serialize(LONG_STRING)
        }));
    });
}
<|endoftext|>