<commit_before>extern crate ndarray;
extern crate libc;
extern crate rayon;
extern crate cblas_sys;
extern crate openblas_src;
#[cfg(not(feature="blas"))]
extern crate matrixmultiply;
#[allow(unused_imports)]
use self::rayon::iter::*;
use self::libc::{c_float, c_int};
use ndarray_ext::NdArray;
use std::mem;
use std::slice;
use tensor::Tensor;

macro_rules! get_xw {
    ($op:expr, $yw:expr, $kw:expr) => {
        $op.stride_w * ($yw - 1) - $op.pad_w + ($op.dilation_w * ($kw - 1) + 1)
    };
}

macro_rules! get_xh {
    ($op:expr, $yh:expr, $kh:expr) => {
        $op.stride_h * ($yh - 1) - $op.pad_h + ($op.dilation_h * ($kh - 1) + 1)
    };
}

macro_rules! get_yw {
    ($op:expr, $xw:expr, $kw:expr) => {
        ($xw + 2 * $op.pad_w - ($op.dilation_w * ($kw - 1) + 1)) / $op.stride_w + 1
    };
}

macro_rules! get_yh {
    ($op:expr, $xh:expr, $kh:expr) => {
        ($xh + 2 * $op.pad_h - ($op.dilation_h * ($kh - 1) + 1)) / $op.stride_h + 1
    };
}

// Returns: &Vec<f32>
macro_rules! get_or_insert_cols {
    ($me:expr, $batch_size:expr, $num_elements_in_batch_c:expr) => {
        unsafe {
            let slf: &mut Self = mem::transmute($me);
            let cols: &Vec<f32> = mem::transmute(
                slf.cols.get_or_insert_with(||
                    alloc_uninitialized_buf($batch_size * $num_elements_in_batch_c)
                )
            );
            cols
        }
    };
}

pub mod conv2d;
pub mod conv2d_transpose;

#[link(name = "conv")]
#[no_mangle]
extern "C" {

    fn im2col_cpu(
        data_im: *const c_float,
        channels: c_int,
        height: c_int,
        width: c_int,
        kernel_h: c_int,
        kernel_w: c_int,
        pad_h: c_int,
        pad_w: c_int,
        stride_h: c_int,
        stride_w: c_int,
        dilation_h: c_int,
        dilation_w: c_int,
        data_col: *const c_float,
    );

    fn col2im_cpu(
        data_col: *const c_float,
        channels: c_int,
        height: c_int,
        width: c_int,
        kernel_h: c_int,
        kernel_w: c_int,
        pad_h: c_int,
        pad_w: c_int,
        stride_h: c_int,
        stride_w: c_int,
        dilation_h: c_int,
        dilation_w: c_int,
        data_im: *const c_float,
    );
}

#[inline]
fn im2col(
    data_im: &c_float,
    channels: usize,
    height: usize,
    width: usize,
    kernel_h: usize,
    kernel_w: usize,
    pad_h: usize,
    pad_w: usize,
    stride_h: usize,
    stride_w: usize,
    dilation_h: usize,
    dilation_w: usize,
    data_col: &c_float,
) {
    unsafe {
        im2col_cpu(
            data_im as *const c_float,
            channels as i32,
            height as i32,
            width as i32,
            kernel_h as i32,
            kernel_w as i32,
            pad_h as i32,
            pad_w as i32,
            stride_h as i32,
            stride_w as i32,
            dilation_h as i32,
            dilation_w as i32,
            data_col,
        )
    }
}

#[inline]
fn col2im(
    data_col: &f32,
    channels: usize,
    height: usize,
    width: usize,
    kernel_h: usize,
    kernel_w: usize,
    pad_h: usize,
    pad_w: usize,
    stride_h: usize,
    stride_w: usize,
    dilation_h: usize,
    dilation_w: usize,
    data_im: &f32,
) {
    unsafe {
        col2im_cpu(
            data_col as *const c_float,
            channels as c_int,
            height as c_int,
            width as c_int,
            kernel_h as c_int,
            kernel_w as c_int,
            pad_h as c_int,
            pad_w as c_int,
            stride_h as c_int,
            stride_w as c_int,
            dilation_h as c_int,
            dilation_w as c_int,
            data_im as *const c_float,
        )
    }
}

#[inline]
fn alloc_uninitialized_buf(size: usize) -> Vec<f32>
{
    let mut buf = Vec::with_capacity(size);
    unsafe {
        buf.set_len(size);
    }
    buf
}

#[inline]
fn sgemm(trans_a: bool, trans_b: bool,
         a: &f32, b: &f32, c: &f32,
         m: usize, n: usize, k: usize, alpha: f32, beta: f32)
{
    #[cfg(feature="blas")] {
        let m = m as i32;
        let n = n as i32;
        let k = k as i32;
        unsafe {
            cblas_sys::cblas_sgemm(
                cblas_sys::CBLAS_LAYOUT::CblasRowMajor,
                if trans_a {
                    cblas_sys::CBLAS_TRANSPOSE::CblasTrans
                } else {
                    cblas_sys::CBLAS_TRANSPOSE::CblasNoTrans
                },
                if trans_b {
                    cblas_sys::CBLAS_TRANSPOSE::CblasTrans
                } else {
                    cblas_sys::CBLAS_TRANSPOSE::CblasNoTrans
                },
                m, n, k,
                alpha,
                a as *const f32, if trans_a { m } else { k }, // lda
                b as *const f32, if trans_b { k } else { n }, // ldb
                beta,
                mem::transmute::<&f32, *mut f32>(c), n, // ldc
            );
        }
    }
    #[cfg(not(feature="blas"))] {
        let rsa = if trans_a { 1 } else { k };
        let csa = if trans_a { m } else { 1 };
        let rsb = if trans_b { 1 } else { n };
        let csb = if trans_b { k } else { 1 };
        let rsc = n;
        let csc = 1;
        unsafe {
            let c: *mut f32 = mem::transmute(c);
            matrixmultiply::sgemm(m, k, n,
                                  alpha,
                                  a as *const f32, rsa as isize, csa as isize,
                                  b as *const f32, rsb as isize, csb as isize,
                                  beta,
                                  c as *mut f32, rsc as isize, csc as isize
            )
        }
    }
}

#[test]
fn test_gemm_trans_a() {
    let a = [1., 2., 3., 4., 5., 6.];
    let b = [1., 2., 3., 4.];
    let c = [0.; 6];
    let m = 3; // row of op(a)
    let n = 2; // col of op(b)
    let k = 2; // col of op(a)
    sgemm(true, false, &a[0], &b[0], &c[0], m, n, k, 1., 0.);
    assert_eq!(&c, &[13.0, 18.0, 17.0, 24.0, 21.0, 30.0]);
}

#[test]
fn test_gemm_trans_b() {
    let a = [1., 2., 3., 4.];
    let b = [1., 2., 3., 4., 5., 6.];
    let c = [0.; 6];
    let m = 2; // row of op(a)
    let n = 3; // col of op(b)
    let k = 2; // col of op(a)
    sgemm(false, true, &a[0], &b[0], &c[0], m, n, k, 1., 0.);
    assert_eq!(&c, &[5., 11., 17., 11., 25., 39.]);
}

#[test]
fn test_conv_filter_grad()
{
    use ::op::Op;
    let op = conv2d::Conv2DFilterGrad {
        pad_h: 0,
        pad_w: 0,
        stride_h: 1,
        stride_w: 1,
        dilation_h: 1,
        dilation_w: 1,
    };

    let (kh, kw) = (2, 2);
    let (xch, ych) = (3, 2);
    let (yh, yw) = (2, 2);
    let batch_size = 2;

    let x = ::ndarray_ext::ones(&[batch_size, yh, yw, kh, kw, xch]);
    let g = ::ndarray_ext::ones(&[batch_size, ych, yh, yw]);
    let w = ::ndarray_ext::ones(&[ych, xch, kh, kw]);

    let ret = op.compute(::runtime::OpComputeContext {
        xs: vec![&x, &g, &w],
        node: &::ops::zeros(&[0]) // dummy (not used)
    });

    assert_eq!(w.shape(), ret[0].as_ref().unwrap().shape());  // (2, 3, 2, 2)
    assert_eq!(ret[0].clone().unwrap().into_raw_vec(), vec![8.; 24]);
}

#[test]
fn test_sequential_sgemm()
{
    let x = [0., 1., 2., 3.];
    let y = [0., 1., 2., 3.];
    let z = [0.; 8];

    for i in 0..2 {
        sgemm(false, false, &x[0], &y[0], &z[i * 4], 2, 2, 2, 1., 0.)
    }
    assert_eq!([2.0, 3.0, 6.0, 11.0, 2.0, 3.0, 6.0, 11.0], z);
}

#[test]
fn test_sgemm_acc()
{
    let x = [0., 1., 2., 3.];
    let y = [0., 1., 2., 3.];
    let z = [0.; 4];

    let num_iter = 3.;

    for _ in 0..num_iter as usize {
        sgemm(false, false, &x[0], &y[0], &z[0], 2, 2, 2, 1., 1.)
    }
    assert_eq!([2.*num_iter, 3.*num_iter, 6.*num_iter, 11.*num_iter], z);
}<commit_msg>Add feature flags for blas crates<commit_after>extern crate ndarray;
extern crate libc;
extern crate rayon;
#[cfg(feature="blas")]
extern crate cblas_sys;
#[cfg(feature="blas")]
extern crate openblas_src;
#[cfg(not(feature="blas"))]
extern crate matrixmultiply;
#[allow(unused_imports)]
use self::rayon::iter::*;
use self::libc::{c_float, c_int};
use ndarray_ext::NdArray;
use std::mem;
use std::slice;
use tensor::Tensor;

macro_rules! get_xw {
    ($op:expr, $yw:expr, $kw:expr) => {
        $op.stride_w * ($yw - 1) - $op.pad_w + ($op.dilation_w * ($kw - 1) + 1)
    };
}

macro_rules! get_xh {
    ($op:expr, $yh:expr, $kh:expr) => {
        $op.stride_h * ($yh - 1) - $op.pad_h + ($op.dilation_h * ($kh - 1) + 1)
    };
}

macro_rules! get_yw {
    ($op:expr, $xw:expr, $kw:expr) => {
        ($xw + 2 * $op.pad_w - ($op.dilation_w * ($kw - 1) + 1)) / $op.stride_w + 1
    };
}

macro_rules! get_yh {
    ($op:expr, $xh:expr, $kh:expr) => {
        ($xh + 2 * $op.pad_h - ($op.dilation_h * ($kh - 1) + 1)) / $op.stride_h + 1
    };
}

// Returns: &Vec<f32>
macro_rules! get_or_insert_cols {
    ($me:expr, $batch_size:expr, $num_elements_in_batch_c:expr) => {
        unsafe {
            let slf: &mut Self = mem::transmute($me);
            let cols: &Vec<f32> = mem::transmute(
                slf.cols.get_or_insert_with(||
                    alloc_uninitialized_buf($batch_size * $num_elements_in_batch_c)
                )
            );
            cols
        }
    };
}

pub mod conv2d;
pub mod conv2d_transpose;

#[link(name = "conv")]
#[no_mangle]
extern "C" {

    fn im2col_cpu(
        data_im: *const c_float,
        channels: c_int,
        height: c_int,
        width: c_int,
        kernel_h: c_int,
        kernel_w: c_int,
        pad_h: c_int,
        pad_w: c_int,
        stride_h: c_int,
        stride_w: c_int,
        dilation_h: c_int,
        dilation_w: c_int,
        data_col: *const c_float,
    );

    fn col2im_cpu(
        data_col: *const c_float,
        channels: c_int,
        height: c_int,
        width: c_int,
        kernel_h: c_int,
        kernel_w: c_int,
        pad_h: c_int,
        pad_w: c_int,
        stride_h: c_int,
        stride_w: c_int,
        dilation_h: c_int,
        dilation_w: c_int,
        data_im: *const c_float,
    );
}

#[inline]
fn im2col(
    data_im: &c_float,
    channels: usize,
    height: usize,
    width: usize,
    kernel_h: usize,
    kernel_w: usize,
    pad_h: usize,
    pad_w: usize,
    stride_h: usize,
    stride_w: usize,
    dilation_h: usize,
    dilation_w: usize,
    data_col: &c_float,
) {
    unsafe {
        im2col_cpu(
            data_im as *const c_float,
            channels as i32,
            height as i32,
            width as i32,
            kernel_h as i32,
            kernel_w as i32,
            pad_h as i32,
            pad_w as i32,
            stride_h as i32,
            stride_w as i32,
            dilation_h as i32,
            dilation_w as i32,
            data_col,
        )
    }
}

#[inline]
fn col2im(
    data_col: &f32,
    channels: usize,
    height: usize,
    width: usize,
    kernel_h: usize,
    kernel_w: usize,
    pad_h: usize,
    pad_w: usize,
    stride_h: usize,
    stride_w: usize,
    dilation_h: usize,
    dilation_w: usize,
    data_im: &f32,
) {
    unsafe {
        col2im_cpu(
            data_col as *const c_float,
            channels as c_int,
            height as c_int,
            width as c_int,
            kernel_h as c_int,
            kernel_w as c_int,
            pad_h as c_int,
            pad_w as c_int,
            stride_h as c_int,
            stride_w as c_int,
            dilation_h as c_int,
            dilation_w as c_int,
            data_im as *const c_float,
        )
    }
}

#[inline]
fn alloc_uninitialized_buf(size: usize) -> Vec<f32>
{
    let mut buf = Vec::with_capacity(size);
    unsafe {
        buf.set_len(size);
    }
    buf
}

#[inline]
fn sgemm(trans_a: bool, trans_b: bool,
         a: &f32, b: &f32, c: &f32,
         m: usize, n: usize, k: usize, alpha: f32, beta: f32)
{
    #[cfg(feature="blas")] {
        let m = m as i32;
        let n = n as i32;
        let k = k as i32;
        unsafe {
            cblas_sys::cblas_sgemm(
                cblas_sys::CBLAS_LAYOUT::CblasRowMajor,
                if trans_a {
                    cblas_sys::CBLAS_TRANSPOSE::CblasTrans
                } else {
                    cblas_sys::CBLAS_TRANSPOSE::CblasNoTrans
                },
                if trans_b {
                    cblas_sys::CBLAS_TRANSPOSE::CblasTrans
                } else {
                    cblas_sys::CBLAS_TRANSPOSE::CblasNoTrans
                },
                m, n, k,
                alpha,
                a as *const f32, if trans_a { m } else { k }, // lda
                b as *const f32, if trans_b { k } else { n }, // ldb
                beta,
                mem::transmute::<&f32, *mut f32>(c), n, // ldc
            );
        }
    }
    #[cfg(not(feature="blas"))] {
        let rsa = if trans_a { 1 } else { k };
        let csa = if trans_a { m } else { 1 };
        let rsb = if trans_b { 1 } else { n };
        let csb = if trans_b { k } else { 1 };
        let rsc = n;
        let csc = 1;
        unsafe {
            let c: *mut f32 = mem::transmute(c);
            matrixmultiply::sgemm(m, k, n,
                                  alpha,
                                  a as *const f32, rsa as isize, csa as isize,
                                  b as *const f32, rsb as isize, csb as isize,
                                  beta,
                                  c as *mut f32, rsc as isize, csc as isize
            )
        }
    }
}

#[test]
fn test_gemm_trans_a() {
    let a = [1., 2., 3., 4., 5., 6.];
    let b = [1., 2., 3., 4.];
    let c = [0.; 6];
    let m = 3; // row of op(a)
    let n = 2; // col of op(b)
    let k = 2; // col of op(a)
    sgemm(true, false, &a[0], &b[0], &c[0], m, n, k, 1., 0.);
    assert_eq!(&c, &[13.0, 18.0, 17.0, 24.0, 21.0, 30.0]);
}

#[test]
fn test_gemm_trans_b() {
    let a = [1., 2., 3., 4.];
    let b = [1., 2., 3., 4., 5., 6.];
    let c = [0.; 6];
    let m = 2; // row of op(a)
    let n = 3; // col of op(b)
    let k = 2; // col of op(a)
    sgemm(false, true, &a[0], &b[0], &c[0], m, n, k, 1., 0.);
    assert_eq!(&c, &[5., 11., 17., 11., 25., 39.]);
}

#[test]
fn test_conv_filter_grad()
{
    use ::op::Op;
    let op = conv2d::Conv2DFilterGrad {
        pad_h: 0,
        pad_w: 0,
        stride_h: 1,
        stride_w: 1,
        dilation_h: 1,
        dilation_w: 1,
    };

    let (kh, kw) = (2, 2);
    let (xch, ych) = (3, 2);
    let (yh, yw) = (2, 2);
    let batch_size = 2;

    let x = ::ndarray_ext::ones(&[batch_size, yh, yw, kh, kw, xch]);
    let g = ::ndarray_ext::ones(&[batch_size, ych, yh, yw]);
    let w = ::ndarray_ext::ones(&[ych, xch, kh, kw]);

    let ret = op.compute(::runtime::OpComputeContext {
        xs: vec![&x, &g, &w],
        node: &::ops::zeros(&[0]) // dummy (not used)
    });

    assert_eq!(w.shape(), ret[0].as_ref().unwrap().shape());  // (2, 3, 2, 2)
    assert_eq!(ret[0].clone().unwrap().into_raw_vec(), vec![8.; 24]);
}

#[test]
fn test_sequential_sgemm()
{
    let x = [0., 1., 2., 3.];
    let y = [0., 1., 2., 3.];
    let z = [0.; 8];

    for i in 0..2 {
        sgemm(false, false, &x[0], &y[0], &z[i * 4], 2, 2, 2, 1., 0.)
    }
    assert_eq!([2.0, 3.0, 6.0, 11.0, 2.0, 3.0, 6.0, 11.0], z);
}

#[test]
fn test_sgemm_acc()
{
    let x = [0., 1., 2., 3.];
    let y = [0., 1., 2., 3.];
    let z = [0.; 4];

    let num_iter = 3.;

    for _ in 0..num_iter as usize {
        sgemm(false, false, &x[0], &y[0], &z[0], 2, 2, 2, 1., 1.)
    }
    assert_eq!([2.*num_iter, 3.*num_iter, 6.*num_iter, 11.*num_iter], z);
}<|endoftext|>


<commit_before>use std::cell::{Cell, RefCell};
use std::rc::Rc;

#[derive(Debug, Clone, PartialEq)]
pub struct Token {
    category: String,
    lexeme: String,
}

fn lex(s: String) -> Vec<Token> {
    let mut tokens = vec![];

    for c in s.chars() {
        let (cat, lexeme) = match c {
            '|' => ("|", "".to_string()),
            '*' => ("*", "".to_string()),
            '+' => ("+", "".to_string()),
            '(' => ("(", "".to_string()),
            ')' => (")", "".to_string()),
            _ => ("Lit", c.to_string()),
        };

        tokens.push( Token{ category: cat.to_string(), lexeme: lexeme.to_string() } );
    }

    tokens.push( Token{ category: "EOF".to_string(), lexeme: "".to_string() } );

    tokens
}

#[derive(Debug)]
pub enum NodeCat {
    T(String),
    NT(String),
}

#[derive(Debug)]
pub struct Node {
    category: NodeCat,
    children: Vec<Rc<RefCell<Node>>>,
}

impl Node {
    fn new_nt(val: &'static str) -> Rc<RefCell<Node>> {
        Rc::new(RefCell::new(Node {
            category: NodeCat::NT(val.to_string()),
            children: vec![],
        }))
    }

    fn new_t(val: &'static str) -> Rc<RefCell<Node>> {
        Rc::new(RefCell::new(Node {
            category: NodeCat::T(val.to_string()),
            children: vec![],
        }))
    }

    fn print(&self) {
        println!("{:?}", self.category);
        for c in &self.children {
            println!("- {:?}", c.borrow().category);
        }

        for c in &self.children {
            if let NodeCat::NT(_) = c.borrow().category {
                c.borrow().print();
            }
        }
    }

    fn preorder_walk(x: Rc<RefCell<Node>>, level: usize) {
        let x = x.borrow();
        let separator = "|-- ";
        let s = format!("{}{:?}", separator, x.category);
        let space_n = level * separator.len();

        let mut space = String::new();
        for i in 0..space_n {
            space.push_str(" ");
        }
        println!("{}{}", space, s);

        for c in &x.children {
            Node::preorder_walk(c.clone(), level + 1);
        }
    }
}

#[derive(Debug)]
pub struct Parser {
    index: Cell<usize>,
    src: String,
    tokens: Vec<Token>,
    tree: Rc<RefCell<Node>>,
    focus: RefCell<Rc<RefCell<Node>>>,
}

impl Parser {
    fn new(src: String) -> Parser {
        let tokens = lex(src.clone());
        let root = Node::new_nt("Re");
        Parser {
            index: Cell::new(0),
            src: src,
            tokens: tokens,
            tree: root.clone(),
            focus: RefCell::new(root.clone()),
        }
    }

    fn next(&self) {
        self.index.set(self.index.get() + 1);
    }

    fn print(&self) {
        println!("\nParser");
        let focus = self.focus.borrow();
        println!("focus {:?}", focus.borrow().category);
        Node::preorder_walk(self.tree.clone(), 0);
    }

    fn parse(&self) -> bool {
        return self.re();
    }

    fn re(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("Re panic");

        println!("Re {:?}", token);

        match token.category.as_str() {
            // Re -> Lit Ops
            "Lit" => {
                {
                    let a = Node::new_t("Lit");
                    let new_focus = Node::new_nt("Ops");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.ops();
            },

            // Re -> ( Re ) Ops
            "(" => {
                let new_focus = Node::new_nt("Re");
                let ops_focus = Node::new_nt("Ops");
                {
                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(Node::new_t("("));
                        focus.children.push(new_focus.clone());
                        focus.children.push(Node::new_t(")"));
                        focus.children.push(ops_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }

                self.next();
                if self.re() {
                    let token = self.tokens.get(self.index.get()).expect("Re panic");
                    if token.category == ")" {
                        {
                            *self.focus.borrow_mut() = ops_focus;
                        }
                        self.next();
                        return self.ops();
                    }
                }

                return false;
            },

            "EOF" => return true,

            _ => return false,
        }
    }

    fn ops(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("Ops panic");

        println!("Ops {:?}", token);

        match token.category.as_str() {
            // Ops -> | Re
            // First
            "|" => {
                {
                    let a = Node::new_t("|");
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re();
            },
            // Ops -> * ReFinish
            // First
            "*" => {
                {
                    let a = Node::new_t("*");
                    let new_focus = Node::new_nt("ReFinish");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re_finish();
            },
            // Ops -> + ReFinish
            // First
            "+" => {
                {
                    let a = Node::new_t("+");
                    let new_focus = Node::new_nt("ReFinish");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re_finish();
            },

            // Follow
            "Lit" | "(" => {
                {
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                return self.re();
            },

            // Follow
            ")" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t("Lambda"));
                }
                return true;
            },

            // Follow
            "EOF" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t("EOF"));
                }
                return true;
            },

            _ => return false,
        }
    }

    fn re_finish(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("ReFinish panic");

        println!("ReFinish {:?}", token);

        match token.category.as_str() {
            // ReFinish -> Re
            // Follow
            "Lit" | "(" => {
                {
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                return self.re();
            },
            // Follow
            ")" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t("Lambda"));
                }
                return true;
            },
            // Follow
            "EOF" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t("EOF"));
                }
                return true;
            },

            _ => return false,
        }
    }

}

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

    #[test]
    fn lex_test() {
        let cases = vec![
            ("a", vec![("Lit", "a"), ("EOF", "")]),
            ("aa", vec![("Lit", "a"), ("Lit", "a"), ("EOF", "")]),
            ("a|b|c", vec![("Lit", "a"), ("|", ""), ("Lit", "b"), ("|", ""), ("Lit", "c"), ("EOF", "")]),
            ("(a)", vec![("(", ""), ("Lit", "a"), (")", ""), ("EOF", "")]),
            ("a(b)", vec![("Lit", "a"), ("(", ""), ("Lit", "b"), (")", ""), ("EOF", "")]),
            //"(a|b)",
            //"(a|b)*",
        ];

        for (c, e) in cases {
            let expected: Vec<Token> = e.iter()
                .map(|&(cat, lexeme)| Token { category: cat.to_string(), lexeme: lexeme.to_string()})
                .collect();
            assert_eq!(lex(c.to_string()), expected);
        }
    }

    #[test]
    fn parse_test() {

        let cases = vec![
            "a",
            "(a)",
            "(aa)",
            "a*b",
            "(a*)b",
            "(a)*b",
            "a|(cde)*a+",
            "((((aaa))))",
        ];

        let expect = vec![
            true,
            true,
            true,
            true,
            true,
            true,
            true,
            true,
        ];

        for (c, e) in cases.iter().zip(expect.iter()) {
            println!("\nCase {:?}\n", c);
            let p = Parser::new(c.to_string());
            assert_eq!(p.parse(), *e, "In {:?}", c);
            p.print();
        }
    }
}
<commit_msg>wip: to automata<commit_after>use std::cell::{Cell, RefCell};
use std::rc::Rc;

#[derive(Debug, Clone, PartialEq)]
pub struct Token {
    category: String,
    lexeme: String,
}

fn lex(s: String) -> Vec<Token> {
    let mut tokens = vec![];

    for c in s.chars() {
        let (cat, lexeme) = match c {
            '|' => ("|", "".to_string()),
            '*' => ("*", "".to_string()),
            '+' => ("+", "".to_string()),
            '(' => ("(", "".to_string()),
            ')' => (")", "".to_string()),
            _ => ("Lit", c.to_string()),
        };

        tokens.push( Token{ category: cat.to_string(), lexeme: lexeme.to_string() } );
    }

    tokens.push( Token{ category: "EOF".to_string(), lexeme: "".to_string() } );

    tokens
}

#[derive(Debug)]
pub enum NodeCat {
    T(Token),
    NT(String),
}

#[derive(Debug)]
pub struct Node {
    category: NodeCat,
    children: Vec<Rc<RefCell<Node>>>,
}

impl Node {
    fn new_nt(val: &'static str) -> Rc<RefCell<Node>> {
        Rc::new(RefCell::new(Node {
            category: NodeCat::NT(val.to_string()),
            children: vec![],
        }))
    }

    fn new_t(val: Token) -> Rc<RefCell<Node>> {
        Rc::new(RefCell::new(Node {
            category: NodeCat::T(val),
            children: vec![],
        }))
    }

    fn print(&self) {
        println!("{:?}", self.category);
        for c in &self.children {
            println!("- {:?}", c.borrow().category);
        }

        for c in &self.children {
            if let NodeCat::NT(_) = c.borrow().category {
                c.borrow().print();
            }
        }
    }

    fn preorder_walk(x: Rc<RefCell<Node>>, level: usize) {
        let x = x.borrow();
        let separator = "|-- ";
        let s = format!("{}{:?}", separator, x.category);
        let space_n = level * separator.len();

        let mut space = String::new();
        for i in 0..space_n {
            space.push_str(" ");
        }
        println!("{}{}", space, s);

        for c in &x.children {
            Node::preorder_walk(c.clone(), level + 1);
        }
    }

    fn postorder_walk(x: Rc<RefCell<Node>>) {
        let x = x.borrow();
        for c in &x.children {
            //println!("{:?}", c);
            Node::postorder_walk(c.clone());
        }

        println!("{:?}", x.category);
    }
}

#[derive(Debug)]
pub struct Parser {
    index: Cell<usize>,
    src: String,
    tokens: Vec<Token>,
    tree: Rc<RefCell<Node>>,
    focus: RefCell<Rc<RefCell<Node>>>,
}

impl Parser {
    fn new(src: String) -> Parser {
        let tokens = lex(src.clone());
        let root = Node::new_nt("Re");
        Parser {
            index: Cell::new(0),
            src: src,
            tokens: tokens,
            tree: root.clone(),
            focus: RefCell::new(root.clone()),
        }
    }

    fn next(&self) {
        self.index.set(self.index.get() + 1);
    }

    fn print(&self) {
        println!("\nParser");
        let focus = self.focus.borrow();
        println!("focus {:?}", focus.borrow().category);
        Node::preorder_walk(self.tree.clone(), 0);
    }

    fn parse(&self) -> bool {
        return self.re();
    }

    fn re(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("Re panic");

        println!("Re {:?}", token);

        match token.category.as_str() {
            // Re -> Lit Ops
            "Lit" => {
                {
                    let a = Node::new_t(token.clone());
                    let new_focus = Node::new_nt("Ops");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.ops();
            },

            // Re -> ( Re ) Ops
            "(" => {
                let new_focus = Node::new_nt("Re");
                let ops_focus = Node::new_nt("Ops");
                {
                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(Node::new_t(token.clone()));
                        focus.children.push(new_focus.clone());
                        focus.children.push(Node::new_t(Token { category: ")".to_string(), lexeme: "".to_string()}));
                        focus.children.push(ops_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }

                self.next();
                if self.re() {
                    let token = self.tokens.get(self.index.get()).expect("Re panic");
                    if token.category == ")" {
                        {
                            *self.focus.borrow_mut() = ops_focus;
                        }
                        self.next();
                        return self.ops();
                    }
                }

                return false;
            },

            "EOF" => return true,

            _ => return false,
        }
    }

    fn ops(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("Ops panic");

        println!("Ops {:?}", token);

        match token.category.as_str() {
            // Ops -> | Re
            // First
            "|" => {
                {
                    let a = Node::new_t(token.clone());
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re();
            },
            // Ops -> * ReFinish
            // First
            "*" => {
                {
                    let a = Node::new_t(token.clone());
                    let new_focus = Node::new_nt("ReFinish");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re_finish();
            },
            // Ops -> + ReFinish
            // First
            "+" => {
                {
                    let a = Node::new_t(token.clone());
                    let new_focus = Node::new_nt("ReFinish");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(a);
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                self.next();
                return self.re_finish();
            },

            // Follow
            "Lit" | "(" => {
                {
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                return self.re();
            },

            // Follow
            ")" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t(Token { category: "Lambda".to_string(), lexeme: "".to_string()}));
                }
                return true;
            },

            // Follow
            "EOF" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t(Token { category: "EOF".to_string(), lexeme: "".to_string()}));
                }
                return true;
            },

            _ => return false,
        }
    }

    fn re_finish(&self) -> bool {
        let token = self.tokens.get(self.index.get()).expect("ReFinish panic");

        println!("ReFinish {:?}", token);

        match token.category.as_str() {
            // ReFinish -> Re
            // Follow
            "Lit" | "(" => {
                {
                    let new_focus = Node::new_nt("Re");

                    {
                        let focus = self.focus.borrow();
                        let mut focus = focus.borrow_mut();
                        focus.children.push(new_focus.clone());
                    }

                    *self.focus.borrow_mut() = new_focus;
                }
                return self.re();
            },
            // Follow
            ")" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t(Token { category: "Lambda".to_string(), lexeme: "".to_string()}));
                }
                return true;
            },
            // Follow
            "EOF" => {
                {
                    let focus = self.focus.borrow();
                    let mut focus = focus.borrow_mut();
                    focus.children.push(Node::new_t(Token { category: "EOF".to_string(), lexeme: "".to_string()}));
                }
                return true;
            },

            _ => return false,
        }
    }

}

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

    #[test]
    fn lex_test() {
        let cases = vec![
            ("a", vec![("Lit", "a"), ("EOF", "")]),
            ("aa", vec![("Lit", "a"), ("Lit", "a"), ("EOF", "")]),
            ("a|b|c", vec![("Lit", "a"), ("|", ""), ("Lit", "b"), ("|", ""), ("Lit", "c"), ("EOF", "")]),
            ("(a)", vec![("(", ""), ("Lit", "a"), (")", ""), ("EOF", "")]),
            ("a(b)", vec![("Lit", "a"), ("(", ""), ("Lit", "b"), (")", ""), ("EOF", "")]),
            //"(a|b)",
            //"(a|b)*",
        ];

        for (c, e) in cases {
            let expected: Vec<Token> = e.iter()
                .map(|&(cat, lexeme)| Token { category: cat.to_string(), lexeme: lexeme.to_string()})
                .collect();
            assert_eq!(lex(c.to_string()), expected);
        }
    }

    #[test]
    fn parse_test() {

        let cases = vec![
            "a",
            "(a)",
            "(aa)",
            "a*b",
            "(a*)b",
            "(a)*b",
            "a|(cde)*a+",
            "((((aaa))))",
        ];

        let expect = vec![
            true,
            true,
            true,
            true,
            true,
            true,
            true,
            true,
        ];

        for (c, e) in cases.iter().zip(expect.iter()) {
            println!("\nCase {:?}\n", c);
            let p = Parser::new(c.to_string());
            assert_eq!(p.parse(), *e, "In {:?}", c);
            p.print();
        }
    }

    #[test]
    fn engine_test() {

        let cases = vec![
            "a",
            //"(a)",
            //"(aa)",
            //"a*b",
            //"(a*)b",
            //"(a)*b",
            //"a|(cde)*a+",
            //"((((aaa))))",
        ];

        let expect = vec![
            true,
            true,
            true,
            true,
            true,
            true,
            true,
            true,
        ];

        for (c, e) in cases.iter().zip(expect.iter()) {
            println!("\nCase {:?}\n", c);
            let p = Parser::new(c.to_string());
            p.parse();
            println!("WALK");
            Node::postorder_walk(p.tree.clone());
            println!("TREE");
            p.print();
        }
    }
}
<|endoftext|>


<commit_before>use nom::IResult;
use nom::IResult::*;
use nom::{be_u8, be_u16, ErrorKind};

use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use std::convert::AsMut;
use std::convert::From;

use network::protocol::proxy_protocol::header::*;

pub const UNKNOWN_FAMILY: u32 = 1;
pub const INCORRECT_VERSION: u32 = 2;

const PROTOCOL_SIGNATURE_V2: [u8; 12] = [0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A];

//NOTE: Only support the V2 for the moment
pub fn parse(buf: &[u8]) -> IResult<&[u8], ProxyProtocolHeader> {
  match tag!(buf, &PROTOCOL_SIGNATURE_V2) {
    Done(rest, _) => {
      match parse_v2_header(rest) {
        Done(rest, header) => Done(rest, ProxyProtocolHeader::V2(header)),
        Incomplete(i) => Incomplete(i),
        Error(e) => Error(e),
      }
    },
    Incomplete(i) => Incomplete(i),
    Error(e) => Error(e),
  }
}

//TODO: Refactor this method using cond! parser if it's possible
fn read_and_verify_ver_and_command(input: &[u8]) -> IResult<&[u8], u8> {
  match be_u8(input) {
    Done(rest, ver_and_cmd) => {
      // Verify if it's the version 2
      if (ver_and_cmd >> 4) & 0x0f == 0x02 {
        return Done(rest, ver_and_cmd)
      }
      Error(error_code!(ErrorKind::Custom(INCORRECT_VERSION)))
    },
    Error(e)      => Error(e),
    Incomplete(e) => Incomplete(e)
  }
}

fn parse_v2_header(input: &[u8]) -> IResult<&[u8], HeaderV2> {
  do_parse!(
    input,
    ver_and_cmd: read_and_verify_ver_and_command >>
    family: be_u8 >>
    len: be_u16 >>
    addr: apply!(parse_addr_v2, family) >>
    (
      HeaderV2 {
        signature: PROTOCOL_SIGNATURE_V2,
        ver_and_cmd,
        family,
        len,
        addr
      }
    )
  )
}

fn parse_addr_v2(input: &[u8], family: u8) -> IResult<&[u8], ProxyAddr> {
  // family: 4 bits for address family and 4 bits for transport protocol
  match (family >> 4) & 0x0f {
    0x00 => Done(input, ProxyAddr::AfUnspec),
    0x01 => parse_ipv4_on_v2(input),
    0x02 => parse_ipv6_on_v2(input),
    _ => Error(error_code!(ErrorKind::Custom(UNKNOWN_FAMILY))),
  }
}

//TODO: Maybe we can refactor: e.g. count_fixed!(u8, take!(2), 3)
named!(take_4_bytes, take!(4));

fn parse_ipv4_on_v2(input: &[u8]) -> IResult<&[u8], ProxyAddr> {
  do_parse!(
    input,
    src_ip: take_4_bytes >>
    dest_ip: take_4_bytes >>
    src_port: be_u16 >>
    dest_port: be_u16 >>
    (
      ProxyAddr::Ipv4Addr {
        src_addr: SocketAddrV4::new(Ipv4Addr::new(src_ip[0], src_ip[1], src_ip[2], src_ip[3]), src_port),
        dst_addr: SocketAddrV4::new(Ipv4Addr::new(dest_ip[0], dest_ip[1], dest_ip[2], dest_ip[3]), dest_port)
      }
    )
  )
}

named!(take_16_bytes, take!(16));

// Transform a slice to an array
fn clone_into_array<A, T>(slice: &[T]) -> A where A: Sized + Default + AsMut<[T]>, T: Clone {
  let mut a = Default::default();
  <A as AsMut<[T]>>::as_mut(&mut a).clone_from_slice(slice);
  a
}

fn parse_ipv6_on_v2(input: &[u8]) -> IResult<&[u8], ProxyAddr> {
  do_parse!(
    input,
    src_ip: take_16_bytes >>
    dest_ip: take_16_bytes >>
    src_port: be_u16 >>
    dest_port: be_u16 >>
    (
      ProxyAddr::Ipv6Addr {
        src_addr: SocketAddrV6::new(Ipv6Addr::from(clone_into_array::<[u8; 16], u8>(&src_ip[..])), src_port, 0 ,0),
        dst_addr: SocketAddrV6::new(Ipv6Addr::from(clone_into_array::<[u8; 16], u8>(&dest_ip[..])), dest_port, 0 ,0),
      }
    )
  )
}

#[cfg(test)]
mod test {

  use super::*;
  use std::net::{IpAddr, SocketAddr};
  use nom::Needed::Size;

  #[test]
  fn test_parse_proxy_protocol_v2_ipv4_addr_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
      0x00, 0x0C,                                                             // address sizes = 12
      0x7D, 0x19, 0x0A, 0x01,                                                 // source address
      0x0A, 0x04, 0x05, 0x08,                                                 // destination address
      0x1F, 0x90,                                                             // source port
      0x10, 0x68,                                                             // destination port
    ];

    let src_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200);
    let expected = ProxyProtocolHeader::V2(HeaderV2::new(Some((src_addr, dst_addr))));

    assert_eq!(Done(&[][..], expected), parse(input));
  }

  #[test]
  fn it_should_parse_proxy_protocol_v2_ipv6_addr_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,                         // MAGIC header
      0x20,                                                                                // Version 2 and command LOCAL
      0x21,                                                                                // family AF_UNIX with IPv6
      0x00, 0x24,                                                                                // address sizes = 36
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // source address
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, // destination address
      0x1F, 0x90,                                                                                // source port
      0x10, 0x68,                                                                                // destination port
    ];

    let src_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)), 4200);
    let expected = ProxyProtocolHeader::V2(HeaderV2::new(Some((src_addr, dst_addr))));

    assert_eq!(Done(&[][..], expected), parse(input));
  }

  #[test]
  fn it_should_parse_proxy_protocol_v2_afunspec_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x00,                                                                   // family AF_UNSPEC and transport protocol unknow
      0x00, 0x00,                                                             // address sizes = 0
    ];

    let expected = ProxyProtocolHeader::V2(HeaderV2::new(None));

    assert_eq!(Done(&[][..], expected), parse(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_v2_with_unknown_version() {
    let unknow_version = 0x30;

    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      unknow_version,                                                         // Version 2 and command LOCAL
    ];

    assert_eq!(Error(ErrorKind::Custom(INCORRECT_VERSION)), parse(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_with_unknown_family() {
    let unknow_family = 0x30;

    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      unknow_family,                                                          // family
      0x00, 0x00,                                                             // address sizes = 0
    ];

    assert_eq!(Error(ErrorKind::Custom(UNKNOWN_FAMILY)), parse(input));
  }

  #[test]
  fn it_should_not_parse_request_without_magic_header() {

    let input = &[
      0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, // INCORRECT MAGIC header
    ];

    assert_eq!(Error(ErrorKind::Tag), parse(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_v2_ipv4_addr_header_with_missing_data() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
    ];

    assert_eq!(Incomplete(Size(4)), parse(input));
  }
}<commit_msg>Refactor the proxy protocol parser following the review<commit_after>use nom::{be_u8, be_u16};

use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use std::convert::AsMut;
use std::convert::From;

use network::protocol::proxy_protocol::header::*;

const PROTOCOL_SIGNATURE_V2: [u8; 12] = [0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A];

named!(read_and_verify_ver_and_command<u8>, verify!(be_u8, |ver_and_cmd: u8| (ver_and_cmd >> 4) & 0x0f == 0x02));

named!(pub parse_v2_header<HeaderV2>,
  do_parse!(
    signature: tag!(&PROTOCOL_SIGNATURE_V2) >>
    ver_and_cmd: read_and_verify_ver_and_command >>
    family: be_u8 >>
    len: be_u16 >>
    addr: apply!(parse_addr_v2, family) >>
    (
      HeaderV2 {
        signature: PROTOCOL_SIGNATURE_V2,
        ver_and_cmd,
        family,
        len,
        addr
      }
    )
  )
);

named_args!(address_family(family: u8)<u8>, value!((family >> 4) & 0x0f));

named!(unspec<ProxyAddr>, value!(ProxyAddr::AfUnspec));

named_args!(parse_addr_v2(family: u8)<ProxyAddr>,
  // family: 4 bits for address family and 4 bits for transport protocol
  switch!(apply!(address_family, family),
    0x00 => call!(unspec) |
    0x01 => call!(parse_ipv4_on_v2) |
    0x02 => call!(parse_ipv6_on_v2)
  )
);

named!(take_4_bytes, take!(4));

named!(parse_ipv4_on_v2<ProxyAddr>,
  do_parse!(
    src_ip: take_4_bytes >>
    dest_ip: take_4_bytes >>
    src_port: be_u16 >>
    dest_port: be_u16 >>
    (
      ProxyAddr::Ipv4Addr {
        src_addr: SocketAddrV4::new(Ipv4Addr::new(src_ip[0], src_ip[1], src_ip[2], src_ip[3]), src_port),
        dst_addr: SocketAddrV4::new(Ipv4Addr::new(dest_ip[0], dest_ip[1], dest_ip[2], dest_ip[3]), dest_port)
      }
    )
  )
);

named!(take_16_bytes, take!(16));

named!(parse_ipv6_on_v2<ProxyAddr>,
  do_parse!(
    src_ip: take_16_bytes >>
    dest_ip: take_16_bytes >>
    src_port: be_u16 >>
    dest_port: be_u16 >>
    (
      ProxyAddr::Ipv6Addr {
        src_addr: SocketAddrV6::new(Ipv6Addr::from(clone_into_array::<[u8; 16], u8>(&src_ip[..])), src_port, 0 ,0),
        dst_addr: SocketAddrV6::new(Ipv6Addr::from(clone_into_array::<[u8; 16], u8>(&dest_ip[..])), dest_port, 0 ,0),
      }
    )
  )
);

// Transform a slice to an array
fn clone_into_array<A, T>(slice: &[T]) -> A where A: Sized + Default + AsMut<[T]>, T: Clone {
  let mut a = Default::default();
  <A as AsMut<[T]>>::as_mut(&mut a).clone_from_slice(slice);
  a
}

#[cfg(test)]
mod test {

  use super::*;
  use nom::ErrorKind;
  use nom::IResult;
  use nom::IResult::*;
  use std::net::{IpAddr, SocketAddr};
  use nom::Needed::Size;

  #[test]
  fn test_parse_proxy_protocol_v2_ipv4_addr_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
      0x00, 0x0C,                                                             // address sizes = 12
      0x7D, 0x19, 0x0A, 0x01,                                                 // source address
      0x0A, 0x04, 0x05, 0x08,                                                 // destination address
      0x1F, 0x90,                                                             // source port
      0x10, 0x68,                                                             // destination port
    ];

    let src_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200);
    let expected = HeaderV2::new(Some((src_addr, dst_addr)));

    assert_eq!(Done(&[][..], expected), parse_v2_header(input));
  }

  #[test]
  fn it_should_parse_proxy_protocol_v2_ipv6_addr_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,                         // MAGIC header
      0x20,                                                                                // Version 2 and command LOCAL
      0x21,                                                                                // family AF_UNIX with IPv6
      0x00, 0x24,                                                                                // address sizes = 36
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // source address
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, // destination address
      0x1F, 0x90,                                                                                // source port
      0x10, 0x68,                                                                                // destination port
    ];

    let src_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)), 4200);
    let expected = HeaderV2::new(Some((src_addr, dst_addr)));

    assert_eq!(Done(&[][..], expected), parse_v2_header(input));
  }

  #[test]
  fn it_should_parse_proxy_protocol_v2_afunspec_header() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x00,                                                                   // family AF_UNSPEC and transport protocol unknow
      0x00, 0x00,                                                             // address sizes = 0
    ];

    let expected = HeaderV2::new(None);

    assert_eq!(Done(&[][..], expected), parse_v2_header(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_v2_with_unknown_version() {
    let unknow_version = 0x30;

    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      unknow_version,                                                         // Version 2 and command LOCAL
    ];

    assert_eq!(Error(ErrorKind::Verify), parse_v2_header(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_with_unknown_family() {
    let unknow_family = 0x30;

    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      unknow_family,                                                          // family
      0x00, 0x00,                                                             // address sizes = 0
    ];

    assert_eq!(Error(ErrorKind::Switch), parse_v2_header(input));
  }

  #[test]
  fn it_should_not_parse_request_without_magic_header() {

    let input = &[
      0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, 0x0D, // INCORRECT MAGIC header
    ];

    assert_eq!(Error(ErrorKind::Tag), parse_v2_header(input));
  }

  #[test]
  fn it_should_not_parse_proxy_protocol_v2_ipv4_addr_header_with_missing_data() {
    let input = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
    ];

    assert_eq!(Incomplete(Size(16)), parse_v2_header(input));
  }
}<|endoftext|>


<commit_before>use rsh::State;
use std::collections::HashMap;

pub type Builtin = fn(State) -> State;

pub fn load() -> HashMap<String, Builtin> {
    HashMap::new()
}
<commit_msg>cd builtin now works<commit_after>use rsh::State;
use std::path::PathBuf;
use std::collections::HashMap;
use std::slice::Iter;

pub type Builtin = fn(State) -> State;

pub fn load() -> HashMap<String, Builtin> {
    let mut h = HashMap::new();

    h.insert("cd".to_string(), cd as fn(State) -> State);
    // h.insert("ls".to_string(), ls as fn(State) -> State);

    h
}

fn cd(s: State) -> State {
    match s.argv.get(1) {
        Some(x) => {
            let mut newS = s.clone();

            newS.cwd = PathBuf::from(x);
            return newS
        },
        None => s.clone(),
    }
}

//  fn ls<'a>(s: State) -> State {
//      let dirs: Iter<String>;

//      if s.argv.len() == 0 {
//          let cwd = s.cwd.to_str().unwrap_or("").to_string();
//          let mut args: &'a Vec<String>;
//          args = Vec::new();
//          args.push(cwd);

//          dirs = args.iter();
//      } else {
//          let mut args = s.argv.clone().iter();
//          args.next();
//          dirs = args;
//      }

//      for d in dirs {
//          let mut p = PathBuf::from(d);

//          // Cheking if file so we don't do extra processing
//          if p.is_file() {
//              println!("{}", p.to_str().unwrap_or(""));
//              continue;
//          }

//          // Unwrapping because we know it's a dir, not a file
//          for entry in p.read_dir().unwrap() {
//              match entry {
//                  Ok(e) => {
//                      // TODO replace this unwrap to something safer
//                      println!("{}", e.file_name().into_string().unwrap());
//                  },
//                  Err(e) => println!("Error: {}", e),

//              }
//          }
//      }

//      s
// }

<|endoftext|>


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

// This should go in either 'sys' or 'os'
#[cfg(target_arch = "x86")]
#[cfg(target_arch = "x86_64")]
#[allow(dead_assignment)]
fn cpuid(func: u32) -> (u32, u32, u32, u32) {
    let mut a = 0u32;
    let mut b = 0u32;
    let mut c = 0u32;
    let mut d = 0u32;

    unsafe {
        asm!(
        "
        mov $4, %eax;
        cpuid;
        mov %eax, $0;
        mov %ebx, $1;
        mov %ecx, $2;
        mov %edx, $3;
        "
        : "=r" (a), "=r" (b), "=r" (c), "=r" (d)
        : "r" (func)
        : "eax", "ebx", "ecx", "edx"
        )
    }

    return (a, b, c, d);
}

#[cfg(target_arch = "x86")]
#[cfg(target_arch = "x86_64")]
pub fn supports_aesni() -> bool {
    let (_, _, c, _) = cpuid(1);
    return (c & 0x02000000) != 0;
}
<commit_msg>Re-write supports_aesni() function to compile on 32-bit x86<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[cfg(target_arch = "x86")]
#[cfg(target_arch = "x86_64")]
pub fn supports_aesni() -> bool {
	let mut flags: u32;
    unsafe {
        asm!(
        "
        mov $$1, %eax;
        cpuid;
		mov %ecx, $0;
        "
        : "=r" (flags) // output
        : // input
        : "eax", "ebx", "ecx", "edx" // clobbers
        )
		// No idea why, but on 32-bit targets, the compiler complains
		// about not having enough registers. Adding in this dummy
		// section, however, seems to fix it.
        asm!("")
    }

    return (flags & 0x02000000) != 0;
}
<|endoftext|>


<commit_before>//! Implement push-based [`Write`] trait for both compressing and decompressing.
use std::io::{self, Write};

#[cfg(feature = "tokio")]
use futures::Poll;
#[cfg(feature = "tokio")]
use tokio_io::AsyncWrite;

use zstd_safe;

use dict::{DecoderDictionary, EncoderDictionary};
use stream::{raw, zio};

#[cfg(test)]
#[cfg(feature = "tokio")]
mod async_tests;

#[cfg(test)]
mod tests;

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call [`finish()`] before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
///
/// [`finish()`]: #method.finish
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: zio::Writer<W, raw::Encoder>,
}

/// A decoder that decompress and forward data to another writer.
pub struct Decoder<W: Write> {
    // output writer (decompressed data)
    writer: zio::Writer<W, raw::Decoder>,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,

    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F>(encoder: Encoder<W>, on_finish: F) -> Self
    where
        F: 'static + FnMut(io::Result<W>),
    {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder.as_ref().unwrap().get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder.as_mut().unwrap().get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder.take().unwrap().finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder.as_mut().unwrap().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder.as_mut().unwrap().flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21).
    ///
    /// A level of `0` uses zstd's default (currently `3`).
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    ///
    /// A level of `0` uses zstd's default (currently `3`).
    pub fn with_dictionary(
        writer: W,
        level: i32,
        dictionary: &[u8],
    ) -> io::Result<Self> {
        let encoder = raw::Encoder::with_dictionary(level, dictionary)?;
        let writer = zio::Writer::new(writer, encoder);
        Ok(Encoder { writer })
    }

    /// Creates a new encoder, using an existing prepared `EncoderDictionary`.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_prepared_dictionary(
        writer: W,
        dictionary: &EncoderDictionary,
    ) -> io::Result<Self> {
        let encoder = raw::Encoder::with_prepared_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, encoder);
        Ok(Encoder { writer })
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| {
            result.unwrap();
        })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>(
        self,
        f: F,
    ) -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.writer.writer()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.writer.writer_mut()
    }

    /// **Required**: Finishes the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`try_finish(self)`](#method.try_finish).
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// To get back `self` in case an error happened, use `try_finish`.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after
    ///           writing your data, consider using `auto_finish()` to get an
    ///           `AutoFinishEncoder`.
    pub fn finish(self) -> io::Result<W> {
        self.try_finish().map_err(|(_, err)| err)
    }

    /// **Required**: Attempts to finish the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`finish(self)`](#method.finish).
    ///
    /// This returns the inner writer if the finish was successful, or the
    /// object plus an error if it wasn't.
    ///
    /// `write` on this object will panic after `try_finish` has been called,
    /// even if it fails.
    pub fn try_finish(mut self) -> Result<W, (Self, io::Error)> {
        match self.writer.finish() {
            // Return the writer, because why not
            Ok(()) => Ok(self.writer.into_inner().0),
            Err(e) => Err((self, e)),
        }
    }

    /// Attemps to finish the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`finish(self)`](#method.finish).
    pub fn do_finish(&mut self) -> io::Result<()> {
        self.writer.finish()
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        zstd_safe::cstream_in_size()
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

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

#[cfg(feature = "tokio")]
impl<W: AsyncWrite> AsyncWrite for Encoder<W> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        try_nb!(self.do_finish());
        self.writer.writer_mut().shutdown()
    }
}

impl<W: Write> Decoder<W> {
    /// Creates a new decoder.
    pub fn new(writer: W) -> io::Result<Self> {
        Self::with_dictionary(writer, &[])
    }

    /// Creates a new decoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, dictionary: &[u8]) -> io::Result<Self> {
        let decoder = raw::Decoder::with_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, decoder);
        Ok(Decoder { writer })
    }

    /// Creates a new decoder, using an existing prepared `DecoderDictionary`.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_prepared_dictionary(
        writer: W,
        dictionary: &DecoderDictionary,
    ) -> io::Result<Self> {
        let decoder = raw::Decoder::with_prepared_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, decoder);
        Ok(Decoder { writer })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.writer.writer()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this decoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.writer.writer_mut()
    }

    /// Returns the inner `Write`.
    pub fn into_inner(self) -> W {
        self.writer.into_inner().0
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        zstd_safe::cstream_in_size()
    }
}

impl<W: Write> Write for Decoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

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

#[cfg(feature = "tokio")]
impl<W: AsyncWrite> AsyncWrite for Decoder<W> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.writer.writer_mut().shutdown()
    }
}

fn _assert_traits() {
    fn _assert_send<T: Send>(_: T) {}

    _assert_send(Decoder::new(Vec::new()));
    _assert_send(Encoder::new(Vec::new(), 1));
}
<commit_msg>Add `dyn` to trait objects<commit_after>//! Implement push-based [`Write`] trait for both compressing and decompressing.
use std::io::{self, Write};

#[cfg(feature = "tokio")]
use futures::Poll;
#[cfg(feature = "tokio")]
use tokio_io::AsyncWrite;

use zstd_safe;

use dict::{DecoderDictionary, EncoderDictionary};
use stream::{raw, zio};

#[cfg(test)]
#[cfg(feature = "tokio")]
mod async_tests;

#[cfg(test)]
mod tests;

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call [`finish()`] before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
///
/// [`finish()`]: #method.finish
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: zio::Writer<W, raw::Encoder>,
}

/// A decoder that decompress and forward data to another writer.
pub struct Decoder<W: Write> {
    // output writer (decompressed data)
    writer: zio::Writer<W, raw::Decoder>,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,

    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<dyn FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F>(encoder: Encoder<W>, on_finish: F) -> Self
    where
        F: 'static + FnMut(io::Result<W>),
    {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder.as_ref().unwrap().get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder.as_mut().unwrap().get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder.take().unwrap().finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder.as_mut().unwrap().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder.as_mut().unwrap().flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21).
    ///
    /// A level of `0` uses zstd's default (currently `3`).
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    ///
    /// A level of `0` uses zstd's default (currently `3`).
    pub fn with_dictionary(
        writer: W,
        level: i32,
        dictionary: &[u8],
    ) -> io::Result<Self> {
        let encoder = raw::Encoder::with_dictionary(level, dictionary)?;
        let writer = zio::Writer::new(writer, encoder);
        Ok(Encoder { writer })
    }

    /// Creates a new encoder, using an existing prepared `EncoderDictionary`.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_prepared_dictionary(
        writer: W,
        dictionary: &EncoderDictionary,
    ) -> io::Result<Self> {
        let encoder = raw::Encoder::with_prepared_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, encoder);
        Ok(Encoder { writer })
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| {
            result.unwrap();
        })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>(
        self,
        f: F,
    ) -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.writer.writer()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.writer.writer_mut()
    }

    /// **Required**: Finishes the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`try_finish(self)`](#method.try_finish).
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// To get back `self` in case an error happened, use `try_finish`.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after
    ///           writing your data, consider using `auto_finish()` to get an
    ///           `AutoFinishEncoder`.
    pub fn finish(self) -> io::Result<W> {
        self.try_finish().map_err(|(_, err)| err)
    }

    /// **Required**: Attempts to finish the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`finish(self)`](#method.finish).
    ///
    /// This returns the inner writer if the finish was successful, or the
    /// object plus an error if it wasn't.
    ///
    /// `write` on this object will panic after `try_finish` has been called,
    /// even if it fails.
    pub fn try_finish(mut self) -> Result<W, (Self, io::Error)> {
        match self.writer.finish() {
            // Return the writer, because why not
            Ok(()) => Ok(self.writer.into_inner().0),
            Err(e) => Err((self, e)),
        }
    }

    /// Attemps to finish the stream.
    ///
    /// You *need* to finish the stream when you're done writing, either with
    /// this method or with [`finish(self)`](#method.finish).
    pub fn do_finish(&mut self) -> io::Result<()> {
        self.writer.finish()
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        zstd_safe::cstream_in_size()
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

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

#[cfg(feature = "tokio")]
impl<W: AsyncWrite> AsyncWrite for Encoder<W> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        try_nb!(self.do_finish());
        self.writer.writer_mut().shutdown()
    }
}

impl<W: Write> Decoder<W> {
    /// Creates a new decoder.
    pub fn new(writer: W) -> io::Result<Self> {
        Self::with_dictionary(writer, &[])
    }

    /// Creates a new decoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, dictionary: &[u8]) -> io::Result<Self> {
        let decoder = raw::Decoder::with_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, decoder);
        Ok(Decoder { writer })
    }

    /// Creates a new decoder, using an existing prepared `DecoderDictionary`.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_prepared_dictionary(
        writer: W,
        dictionary: &DecoderDictionary,
    ) -> io::Result<Self> {
        let decoder = raw::Decoder::with_prepared_dictionary(dictionary)?;
        let writer = zio::Writer::new(writer, decoder);
        Ok(Decoder { writer })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.writer.writer()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this decoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.writer.writer_mut()
    }

    /// Returns the inner `Write`.
    pub fn into_inner(self) -> W {
        self.writer.into_inner().0
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        zstd_safe::cstream_in_size()
    }
}

impl<W: Write> Write for Decoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

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

#[cfg(feature = "tokio")]
impl<W: AsyncWrite> AsyncWrite for Decoder<W> {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.writer.writer_mut().shutdown()
    }
}

fn _assert_traits() {
    fn _assert_send<T: Send>(_: T) {}

    _assert_send(Decoder::new(Vec::new()));
    _assert_send(Encoder::new(Vec::new(), 1));
}
<|endoftext|>